let's try again

api_engine.py

@@ -102,9 +102,12 @@ class MatrixGameEngine:
             raise RuntimeError(error_msg)
         
         try:
+            init_start = time.time()
+            logger.info("Starting Matrix-Game V2 model initialization...")
             self._init_models()
             self.model_loaded = True
-            logger.info("Matrix-Game V2 models loaded successfully")
+            init_time = time.time() - init_start
+            logger.info(f"Matrix-Game V2 models loaded successfully in {init_time:.2f} seconds")
         except Exception as e:
             error_msg = f"Failed to initialize Matrix-Game V2 models: {str(e)}"
             logger.error(error_msg)
@@ -122,11 +125,16 @@ class MatrixGameEngine:
             logger.debug(f"Configuration loaded: {self.config}")
             
             # Initialize generator
-            generator = WanDiffusionWrapper(
-                **getattr(self.config, "model_kwargs", {}), is_causal=True)
+            logger.info("Initializing WAN Diffusion generator...")
+            model_kwargs = getattr(self.config, "model_kwargs", {})
+            logger.debug(f"Model kwargs: {model_kwargs}")
+            generator = WanDiffusionWrapper(**model_kwargs, is_causal=True)
+            logger.info("WAN Diffusion generator initialized successfully")
             
             # Initialize VAE decoder
+            logger.info("Initializing VAE decoder...")
             current_vae_decoder = VAEDecoderWrapper()
+            logger.info("VAE decoder initialized successfully")
             
             # Check if model exists locally, if not download from Hugging Face
             if not os.path.exists(self.pretrained_model_path) or not os.path.exists(os.path.join(self.pretrained_model_path, "Wan2.1_VAE.pth")):
@@ -272,7 +280,15 @@ class MatrixGameEngine:
             logger.error(error_msg)
             raise RuntimeError(error_msg)
         
+        generation_start_time = time.time()
+        
         try:
+            # Log GPU memory usage if CUDA is available
+            if torch.cuda.is_available():
+                gpu_memory_allocated = torch.cuda.memory_allocated() / 1024**3  # GB
+                gpu_memory_reserved = torch.cuda.memory_reserved() / 1024**3   # GB
+                logger.debug(f"GPU Memory - Allocated: {gpu_memory_allocated:.2f}GB, Reserved: {gpu_memory_reserved:.2f}GB")
+            
             logger.debug(f"Starting frame generation for scene: {scene_name}")
             
             # Map scene name to mode
@@ -313,7 +329,8 @@ class MatrixGameEngine:
                 mouse_condition = [[0, 0]]
             
             # Generate conditions for multiple frames (for streaming)
-            num_frames = 5  # Generate 5 frames at a time for smoother playback
+            # Must be divisible by num_frame_per_block (which is 3)
+            num_frames = 3  # Generate 3 frames at a time (matches num_frame_per_block)
             
             # Create condition tensors
             keyboard_tensor = torch.tensor(keyboard_condition * num_frames, dtype=self.weight_dtype).unsqueeze(0).to(self.device)
@@ -338,13 +355,21 @@ class MatrixGameEngine:
             
             # Generate frames with streaming pipeline
             with torch.no_grad():
-                logger.debug(f"Starting inference with mode: {mode}")
+                logger.info(f"Starting inference - Frame #{self.frame_count}, Mode: {mode}, Scene: {scene_name}")
                 logger.debug(f"Conditional dict keys: {list(conditional_dict.keys())}")
                 logger.debug(f"Noise shape: {sampled_noise.shape}")
                 
+                # Log tensor shapes for debugging
+                for key, tensor in conditional_dict.items():
+                    if hasattr(tensor, 'shape'):
+                        logger.debug(f"  {key}: {tensor.shape} ({tensor.dtype})")
+                
                 # Set seed for reproducibility
                 set_seed(self.seed + self.frame_count)
                 
+                inference_start = time.time()
+                logger.debug("Starting pipeline.inference()...")
+                
                 # Use inference method for single batch generation
                 outputs = self.pipeline.inference(
                     noise=sampled_noise,
@@ -355,17 +380,37 @@ class MatrixGameEngine:
                     mode=mode
                 )
                 
-                logger.debug(f"Inference completed, outputs type: {type(outputs)}")
+                inference_time = time.time() - inference_start
+                logger.info(f"Inference completed in {inference_time:.2f}s, outputs type: {type(outputs)}")
+                
+                if outputs is not None:
+                    logger.debug(f"Output tensor shape: {outputs.shape if hasattr(outputs, 'shape') else 'No shape attr'}")
                 
                 # Decode first frame from latent
                 if outputs is not None and len(outputs) > 0:
+                    decode_start = time.time()
+                    logger.debug("Starting VAE decoding...")
+                    
                     # Extract first frame
                     frame_latent = outputs[0:1, :, 0:1]  # Get first frame
+                    logger.debug(f"Frame latent shape: {frame_latent.shape}")
+                    
                     decoded = self.pipeline.vae_decoder.decode(frame_latent)
+                    decode_time = time.time() - decode_start
+                    logger.debug(f"VAE decoding completed in {decode_time:.3f}s")
                     
                     # Convert to numpy
+                    logger.debug(f"Decoded tensor shape: {decoded.shape}")
                     frame = decoded[0, :, 0].permute(1, 2, 0).cpu().numpy()
+                    logger.debug(f"Frame numpy shape: {frame.shape}, dtype: {frame.dtype}")
+                    
+                    # Normalize to [0, 255]
+                    frame_min, frame_max = frame.min(), frame.max()
+                    logger.debug(f"Frame value range before normalization: [{frame_min:.3f}, {frame_max:.3f}]")
+                    
                     frame = ((frame + 1) * 127.5).clip(0, 255).astype(np.uint8)
+                    logger.debug(f"Frame shape after normalization: {frame.shape}, dtype: {frame.dtype}")
+                    
                 else:
                     # Generation failed
                     error_msg = "Failed to generate frame: No output from model"
@@ -373,6 +418,8 @@ class MatrixGameEngine:
                     raise RuntimeError(error_msg)
                 
             self.frame_count += 1
+            total_generation_time = time.time() - generation_start_time
+            logger.info(f"Frame generation complete - Total time: {total_generation_time:.3f}s, Frame #{self.frame_count}")
             
         except Exception as e:
             error_msg = f"Error generating frame with Matrix-Game V2 model: {str(e)}"

api_server.py

@@ -240,15 +240,20 @@ class GameSession:
     async def _stream_frames(self, fps: int):
         """Stream frames to the client at the specified FPS"""
         frame_interval = 1.0 / fps  # Time between frames in seconds
+        frame_count = 0
+        
+        logger.info(f"Starting frame streaming for user {self.user_id} at {fps} FPS (interval: {frame_interval:.3f}s)")
         
         try:
             while self.is_streaming:
-                start_time = time.time()
+                stream_start_time = time.time()
                 
                 # Generate frame based on current keyboard and mouse state
                 keyboard_condition = [self.keyboard_state]
                 mouse_condition = [self.mouse_state]
                 
+                logger.debug(f"Frame #{frame_count} - KB: {keyboard_condition[0]}, Mouse: {mouse_condition[0]}")
+                
                 # Check if engine is available
                 if not self.game_manager.engine:
                     error_msg = f"Engine not available: {self.game_manager.engine_error}"
@@ -260,10 +265,13 @@ class GameSession:
                     return
                 
                 try:
+                    generation_start = time.time()
                     # Use the engine to generate the next frame
                     frame_bytes = self.game_manager.engine.generate_frame(
                         self.current_scene, keyboard_condition, mouse_condition
                     )
+                    generation_time = time.time() - generation_start
+                    
                 except Exception as e:
                     error_msg = f"Failed to generate frame: {str(e)}"
                     logger.error(error_msg)
@@ -275,18 +283,34 @@ class GameSession:
                     return
                 
                 # Encode as base64 for sending in JSON
+                encode_start = time.time()
                 frame_base64 = base64.b64encode(frame_bytes).decode('utf-8')
+                encode_time = time.time() - encode_start
                 
                 # Send frame to client
+                send_start = time.time()
                 await self.ws.send_json({
                     'action': 'frame',
                     'frameData': frame_base64,
-                    'timestamp': time.time()
+                    'timestamp': time.time(),
+                    'frameNumber': frame_count,
+                    'generationTime': f"{generation_time:.3f}s",
+                    'frameSize': len(frame_bytes)
                 })
+                send_time = time.time() - send_start
+                
+                # Calculate total time and performance metrics
+                total_time = time.time() - stream_start_time
+                sleep_time = max(0, frame_interval - total_time)
+                
+                # Log performance info every 10 frames or if frame takes too long
+                if frame_count % 10 == 0 or total_time > frame_interval * 1.5:
+                    logger.info(f"Frame #{frame_count} - Gen: {generation_time:.3f}s, "
+                              f"Encode: {encode_time:.3f}s, Send: {send_time:.3f}s, "
+                              f"Total: {total_time:.3f}s, Size: {len(frame_bytes)} bytes, "
+                              f"Sleep: {sleep_time:.3f}s")
                 
-                # Calculate sleep time to maintain FPS
-                elapsed = time.time() - start_time
-                sleep_time = max(0, frame_interval - elapsed)
+                frame_count += 1
                 await asyncio.sleep(sleep_time)
                 
         except asyncio.CancelledError: