wip

api_engine.py

@@ -54,7 +54,9 @@ class MatrixGameEngine:
         self.frame_width = getattr(args, 'frame_width', 640)
         self.frame_height = getattr(args, 'frame_height', 352)
         self.fps = getattr(args, 'fps', 16)
-        self.max_num_output_frames = getattr(args, 'max_num_output_frames', 90)  # Reduced for real-time
+        # For streaming sessions: following original Matrix-Game-2 default
+        # 360 frames at 16 FPS = 22.5 seconds, divisible by num_frame_per_block (3)
+        self.max_num_output_frames = getattr(args, 'max_num_output_frames', 360)  # 360 = 3 * 120, ~23 sec session
         self.seed = getattr(args, 'seed', 0)
         self.config_path = getattr(args, 'config_path', 'configs/inference_yaml/inference_universal.yaml')
         self.checkpoint_path = getattr(args, 'checkpoint_path', '')
@@ -328,13 +330,22 @@ class MatrixGameEngine:
             if mouse_condition is None:
                 mouse_condition = [[0, 0]]
             
-            # Generate conditions for multiple frames (for streaming)
-            # 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)
+            # Generate conditions following the original Matrix-Game V2 pattern
+            # For streaming, use fewer frames for real-time performance
+            max_num_output_frames = self.max_num_output_frames  # Use from constructor
             
-            # Create condition tensors
-            keyboard_tensor = torch.tensor(keyboard_condition * num_frames, dtype=self.weight_dtype).unsqueeze(0).to(self.device)
-            mouse_tensor = torch.tensor(mouse_condition * num_frames, dtype=self.weight_dtype).unsqueeze(0).to(self.device)
+            # Calculate condition frames: (max_num_output_frames - 1) * 4 + 1
+            # This follows the pattern from the original implementation
+            condition_num_frames = (max_num_output_frames - 1) * 4 + 1
+            
+            logger.info(f"Using {max_num_output_frames} output frames -> {condition_num_frames} condition frames")
+            
+            # Create condition tensors with the correct length
+            keyboard_tensor = torch.tensor(keyboard_condition * condition_num_frames, dtype=self.weight_dtype).unsqueeze(0).to(self.device)
+            mouse_tensor = torch.tensor(mouse_condition * condition_num_frames, dtype=self.weight_dtype).unsqueeze(0).to(self.device)
+            
+            logger.debug(f"Keyboard tensor shape: {keyboard_tensor.shape}")
+            logger.debug(f"Mouse tensor shape: {mouse_tensor.shape}")
             
             # Build conditional dict
             cond_concat = torch.cat([scene_data['mask'][:, :4], scene_data['latent']], dim=1)
@@ -348,10 +359,21 @@ class MatrixGameEngine:
             if mode in ['universal', 'gta_drive']:
                 conditional_dict['mouse_cond'] = mouse_tensor
             
-            # Generate noise for the frames
+            # Generate noise for the frames (following original implementation)
             sampled_noise = torch.randn(
-                [1, 16, num_frames, 44, 80], device=self.device, dtype=self.weight_dtype
+                [1, 16, max_num_output_frames, 44, 80], device=self.device, dtype=self.weight_dtype
             )
+            logger.debug(f"Generated noise shape: {sampled_noise.shape}")
+            
+            pipeline_block_size = getattr(self.pipeline, 'num_frame_per_block', 3)
+            logger.debug(f"Pipeline num_frame_per_block: {pipeline_block_size}")
+            
+            # Validate that max_num_output_frames is divisible by num_frame_per_block
+            if max_num_output_frames % pipeline_block_size != 0:
+                logger.error(f"Frame count mismatch: {max_num_output_frames} frames not divisible by block size {pipeline_block_size}")
+                raise ValueError(f"max_num_output_frames ({max_num_output_frames}) must be divisible by num_frame_per_block ({pipeline_block_size})")
+            else:
+                logger.debug(f"Frame count validation passed: {max_num_output_frames} frames / {pipeline_block_size} blocks = {max_num_output_frames // pipeline_block_size} blocks")
             
             # Generate frames with streaming pipeline
             with torch.no_grad():
@@ -386,13 +408,14 @@ class MatrixGameEngine:
                 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
+                # Decode first frame from latent (following original pattern)
                 if outputs is not None and len(outputs) > 0:
                     decode_start = time.time()
                     logger.debug("Starting VAE decoding...")
+                    logger.debug(f"Output tensor shape: {outputs.shape}")
                     
-                    # Extract first frame
-                    frame_latent = outputs[0:1, :, 0:1]  # Get first frame
+                    # Extract first frame from the generated sequence
+                    frame_latent = outputs[0:1, :, 0:1]  # Get first frame [batch=1, channels, frames=1, h, w]
                     logger.debug(f"Frame latent shape: {frame_latent.shape}")
                     
                     decoded = self.pipeline.vae_decoder.decode(frame_latent)

api_server.py

@@ -50,6 +50,13 @@ class GameSession:
         self.keyboard_state = [0, 0, 0, 0, 0, 0]  # forward, back, left, right, jump, attack
         self.mouse_state = [0, 0]  # x, y
         
+        # Game session state
+        self.current_frame = 0
+        self.max_frames = 0  # Will be set when engine is available
+        self.session_duration = 0  # Total session duration in seconds
+        self.session_start_time = None
+        self.session_active = False
+        
         self.background_tasks = []
         
     async def start(self):
@@ -57,6 +64,13 @@ class GameSession:
         self.background_tasks = [
             asyncio.create_task(self._process_action_queue()),
         ]
+        
+        # Initialize session parameters from engine
+        if self.game_manager.engine:
+            self.max_frames = self.game_manager.engine.max_num_output_frames
+            self.session_duration = self.max_frames / self.game_manager.engine.fps
+            logger.info(f"Session initialized - Max frames: {self.max_frames}, Duration: {self.session_duration:.1f}s")
+        
         logger.info(f"Started game session for user {self.user_id}")
         
     async def stop(self):
@@ -99,6 +113,8 @@ class GameSession:
                     result = await self._handle_mouse_input(data)
                 elif action_type == 'change_scene':
                     result = await self._handle_scene_change(data)
+                elif action_type == 'reset_session':
+                    result = await self._handle_session_reset(data)
                 else:
                     result = {
                         'action': action_type,
@@ -137,15 +153,36 @@ class GameSession:
                 'error': 'Stream already active'
             }
         
+        # Check if session has reached max frames
+        if self.current_frame >= self.max_frames:
+            return {
+                'action': 'start_stream',
+                'requestId': data.get('requestId'),
+                'success': False,
+                'error': 'Session has ended. Please reset to start a new session.',
+                'sessionEnded': True
+            }
+        
         fps = data.get('fps', 16)
         self.is_streaming = True
+        self.session_active = True
+        if self.session_start_time is None:
+            self.session_start_time = time.time()
+            logger.info(f"Game session started for user {self.user_id}")
+            
         self.stream_task = asyncio.create_task(self._stream_frames(fps))
         
         return {
             'action': 'start_stream',
             'requestId': data.get('requestId'),
             'success': True,
-            'message': f'Streaming started at {fps} FPS'
+            'message': f'Streaming started at {fps} FPS',
+            'sessionInfo': {
+                'maxFrames': self.max_frames,
+                'sessionDuration': self.session_duration,
+                'currentFrame': self.current_frame,
+                'remainingFrames': self.max_frames - self.current_frame
+            }
         }
     
     async def _handle_stop_stream(self, data: Dict) -> Dict:
@@ -237,6 +274,42 @@ class GameSession:
             'scene': scene_name
         }
     
+    async def _handle_session_reset(self, data: Dict) -> Dict:
+        """Handle session reset requests"""
+        # Reset session state
+        self.current_frame = 0
+        self.session_start_time = None
+        self.session_active = False
+        
+        # Stop current streaming if active
+        if self.is_streaming:
+            self.is_streaming = False
+            if self.stream_task:
+                self.stream_task.cancel()
+                try:
+                    await self.stream_task
+                except asyncio.CancelledError:
+                    pass
+                self.stream_task = None
+        
+        # Reset input states
+        self.keyboard_state = [0, 0, 0, 0, 0, 0]
+        self.mouse_state = [0, 0]
+        
+        logger.info(f"Session reset for user {self.user_id}")
+        
+        return {
+            'action': 'reset_session',
+            'requestId': data.get('requestId'),
+            'success': True,
+            'message': 'Session reset successfully',
+            'sessionInfo': {
+                'maxFrames': self.max_frames,
+                'sessionDuration': self.session_duration,
+                'currentFrame': self.current_frame
+            }
+        }
+    
     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
@@ -245,14 +318,20 @@ class GameSession:
         logger.info(f"Starting frame streaming for user {self.user_id} at {fps} FPS (interval: {frame_interval:.3f}s)")
         
         try:
-            while self.is_streaming:
+            while self.is_streaming and self.current_frame < self.max_frames:
                 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]}")
+                # Calculate session progress
+                progress_percent = (self.current_frame / self.max_frames) * 100 if self.max_frames > 0 else 0
+                remaining_frames = self.max_frames - self.current_frame
+                elapsed_time = time.time() - self.session_start_time if self.session_start_time else 0
+                
+                logger.debug(f"Session progress: {self.current_frame}/{self.max_frames} ({progress_percent:.1f}%) - "
+                           f"KB: {keyboard_condition[0]}, Mouse: {mouse_condition[0]}")
                 
                 # Check if engine is available
                 if not self.game_manager.engine:
@@ -287,7 +366,7 @@ class GameSession:
                 frame_base64 = base64.b64encode(frame_bytes).decode('utf-8')
                 encode_time = time.time() - encode_start
                 
-                # Send frame to client
+                # Send frame to client with session progress
                 send_start = time.time()
                 await self.ws.send_json({
                     'action': 'frame',
@@ -295,10 +374,21 @@ class GameSession:
                     'timestamp': time.time(),
                     'frameNumber': frame_count,
                     'generationTime': f"{generation_time:.3f}s",
-                    'frameSize': len(frame_bytes)
+                    'frameSize': len(frame_bytes),
+                    'sessionProgress': {
+                        'currentFrame': self.current_frame,
+                        'maxFrames': self.max_frames,
+                        'progress': progress_percent,
+                        'remainingFrames': remaining_frames,
+                        'elapsedTime': elapsed_time,
+                        'remainingTime': (remaining_frames / fps) if fps > 0 else 0
+                    }
                 })
                 send_time = time.time() - send_start
                 
+                # Increment session frame counter
+                self.current_frame += 1
+                
                 # Calculate total time and performance metrics
                 total_time = time.time() - stream_start_time
                 sleep_time = max(0, frame_interval - total_time)
@@ -312,6 +402,24 @@ class GameSession:
                 
                 frame_count += 1
                 await asyncio.sleep(sleep_time)
+            
+            # Handle session end
+            if self.current_frame >= self.max_frames:
+                logger.info(f"Game session ended for user {self.user_id} - {self.current_frame}/{self.max_frames} frames completed")
+                self.is_streaming = False
+                self.session_active = False
+                
+                # Send session end notification
+                await self.ws.send_json({
+                    'action': 'session_ended',
+                    'message': 'Game session completed!',
+                    'sessionStats': {
+                        'totalFrames': self.current_frame,
+                        'maxFrames': self.max_frames,
+                        'elapsedTime': elapsed_time,
+                        'sessionDuration': self.session_duration
+                    }
+                })
                 
         except asyncio.CancelledError:
             logger.info(f"Frame streaming cancelled for user {self.user_id}")

client/client.js

@@ -8,10 +8,23 @@ let lastFrameTime = 0;
 let frameCount = 0;
 let fpsUpdateInterval = null;
 
+// Session state
+let sessionInfo = {
+    maxFrames: 0,
+    sessionDuration: 0,
+    currentFrame: 0,
+    remainingFrames: 0,
+    elapsedTime: 0,
+    remainingTime: 0,
+    sessionStartTime: null,
+    sessionEnded: false
+};
+
 // DOM Elements
 const connectBtn = document.getElementById('connect-btn');
 const startStreamBtn = document.getElementById('start-stream-btn');
 const stopStreamBtn = document.getElementById('stop-stream-btn');
+const resetSessionBtn = document.getElementById('reset-session-btn');
 const sceneSelect = document.getElementById('scene-select');
 const gameCanvas = document.getElementById('game-canvas');
 const connectionLog = document.getElementById('connection-log');
@@ -19,6 +32,15 @@ const mousePosition = document.getElementById('mouse-position');
 const fpsCounter = document.getElementById('fps-counter');
 const mouseTrackingArea = document.getElementById('mouse-tracking-area');
 
+// Session progress elements
+const sessionProgress = document.getElementById('session-progress');
+const progressFill = document.getElementById('progress-fill');
+const currentFrameEl = document.getElementById('current-frame');
+const maxFramesEl = document.getElementById('max-frames');
+const elapsedTimeEl = document.getElementById('elapsed-time');
+const sessionDurationEl = document.getElementById('session-duration');
+const remainingTimeEl = document.getElementById('remaining-time');
+
 // Pointer Lock API support check
 const pointerLockSupported = 'pointerLockElement' in document || 
                             'mozPointerLockElement' in document || 
@@ -167,6 +189,8 @@ function setupWebSocketHandlers() {
         logMessage('WebSocket connection established');
         connectBtn.textContent = 'Disconnect';
         startStreamBtn.disabled = false;
+        stopStreamBtn.disabled = true;
+        resetSessionBtn.disabled = false;
         sceneSelect.disabled = false;
     };
     
@@ -225,14 +249,28 @@ function setupWebSocketHandlers() {
             case 'start_stream':
                 if (message.success) {
                     isStreaming = true;
+                    sessionInfo.sessionEnded = false;
                     startStreamBtn.disabled = true;
                     stopStreamBtn.disabled = false;
                     logMessage(`Streaming started: ${message.message}`);
                     
+                    // Update session info if provided
+                    if (message.sessionInfo) {
+                        updateSessionProgress(message.sessionInfo);
+                    }
+                    
                     // Start FPS counter
                     startFpsCounter();
                 } else {
                     logMessage(`Error starting stream: ${message.error}`);
+                    
+                    // Check if session ended
+                    if (message.sessionEnded) {
+                        sessionInfo.sessionEnded = true;
+                        startStreamBtn.disabled = true;
+                        stopStreamBtn.disabled = true;
+                        logMessage('Session has ended. Please reset to start a new session.');
+                    }
                 }
                 break;
                 
@@ -280,6 +318,56 @@ function setupWebSocketHandlers() {
                 }
                 break;
                 
+            case 'session_ended':
+                logMessage(`🏁 ${message.message}`);
+                isStreaming = false;
+                sessionInfo.sessionEnded = true;
+                startStreamBtn.disabled = true;
+                stopStreamBtn.disabled = true;
+                
+                // Show session stats if available
+                if (message.sessionStats) {
+                    const stats = message.sessionStats;
+                    logMessage(`Session Stats: ${stats.totalFrames}/${stats.maxFrames} frames, ${formatTime(stats.elapsedTime)} elapsed`);
+                }
+                
+                // Stop FPS counter
+                stopFpsCounter();
+                break;
+                
+            case 'reset_session':
+                if (message.success) {
+                    logMessage(`🔄 ${message.message}`);
+                    
+                    // Reset session state
+                    sessionInfo = {
+                        maxFrames: 0,
+                        sessionDuration: 0,
+                        currentFrame: 0,
+                        remainingFrames: 0,
+                        elapsedTime: 0,
+                        remainingTime: 0,
+                        sessionStartTime: null,
+                        sessionEnded: false
+                    };
+                    
+                    // Update UI
+                    isStreaming = false;
+                    startStreamBtn.disabled = false;
+                    stopStreamBtn.disabled = true;
+                    sessionProgress.style.display = 'none';
+                    
+                    // Update session info if provided
+                    if (message.sessionInfo) {
+                        updateSessionProgress(message.sessionInfo);
+                    }
+                    
+                    stopFpsCounter();
+                } else {
+                    logMessage(`Error resetting session: ${message.error}`);
+                }
+                break;
+                
             default:
                 logMessage(`Received message: ${JSON.stringify(message)}`);
         }
@@ -367,6 +455,47 @@ function sendChangeScene(scene) {
     }
 }
 
+// Reset session
+function sendResetSession() {
+    if (socket && socket.readyState === WebSocket.OPEN) {
+        socket.send(JSON.stringify({
+            action: 'reset_session',
+            requestId: generateRequestId()
+        }));
+    }
+}
+
+// Update session progress display
+function updateSessionProgress(progress) {
+    if (progress) {
+        sessionInfo = { ...sessionInfo, ...progress };
+        
+        // Show progress UI if session is active
+        if (sessionInfo.maxFrames > 0) {
+            sessionProgress.style.display = 'block';
+            
+            // Update progress bar
+            const progressPercent = (sessionInfo.currentFrame / sessionInfo.maxFrames) * 100;
+            progressFill.style.width = `${progressPercent}%`;
+            
+            // Update text displays
+            currentFrameEl.textContent = sessionInfo.currentFrame;
+            maxFramesEl.textContent = sessionInfo.maxFrames;
+            elapsedTimeEl.textContent = formatTime(sessionInfo.elapsedTime);
+            sessionDurationEl.textContent = formatTime(sessionInfo.sessionDuration);
+            remainingTimeEl.textContent = formatTime(sessionInfo.remainingTime);
+        }
+    }
+}
+
+// Format time in MM:SS format
+function formatTime(seconds) {
+    if (isNaN(seconds) || seconds < 0) return '0:00';
+    const minutes = Math.floor(seconds / 60);
+    const secs = Math.floor(seconds % 60);
+    return `${minutes}:${secs.toString().padStart(2, '0')}`;
+}
+
 // Process incoming frame
 function processFrame(message) {
     // Update FPS calculation
@@ -380,6 +509,11 @@ function processFrame(message) {
     if (message.frameData) {
         gameCanvas.src = `data:image/jpeg;base64,${message.frameData}`;
     }
+    
+    // Update session progress if available
+    if (message.sessionProgress) {
+        updateSessionProgress(message.sessionProgress);
+    }
 }
 
 // Generate a random request ID
@@ -430,8 +564,12 @@ function resetUI() {
     connectBtn.textContent = 'Connect';
     startStreamBtn.disabled = true;
     stopStreamBtn.disabled = true;
+    resetSessionBtn.disabled = true;
     sceneSelect.disabled = true;
     
+    // Hide session progress
+    sessionProgress.style.display = 'none';
+    
     // Reset key indicators
     for (const key in keyElements) {
         keyElements[key].classList.remove('active');
@@ -442,6 +580,18 @@ function resetUI() {
     
     // Reset streaming state
     isStreaming = false;
+    
+    // Reset session state
+    sessionInfo = {
+        maxFrames: 0,
+        sessionDuration: 0,
+        currentFrame: 0,
+        remainingFrames: 0,
+        elapsedTime: 0,
+        remainingTime: 0,
+        sessionStartTime: null,
+        sessionEnded: false
+    };
 }
 
 // Event Listeners
@@ -455,6 +605,7 @@ connectBtn.addEventListener('click', () => {
 
 startStreamBtn.addEventListener('click', sendStartStream);
 stopStreamBtn.addEventListener('click', sendStopStream);
+resetSessionBtn.addEventListener('click', sendResetSession);
 
 sceneSelect.addEventListener('change', () => {
     sendChangeScene(sceneSelect.value);

client/index.html

@@ -257,12 +257,21 @@
                 <img id="game-canvas" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" alt="Game Frame">
                 <div id="mouse-position">Mouse: 0.00, 0.00</div>
                 <div class="fps-counter" id="fps-counter">FPS: 0</div>
+                <div id="session-progress" style="position: absolute; top: 50px; left: 10px; background-color: rgba(0,0,0,0.8); color: #4CAF50; padding: 10px; border-radius: 5px; font-family: monospace; z-index: 20; display: none;">
+                    <div>Session Progress: <span id="progress-bar" style="display: inline-block; width: 200px; height: 10px; background-color: #333; border-radius: 5px; margin-left: 10px; position: relative;">
+                        <div id="progress-fill" style="height: 100%; background-color: #4CAF50; border-radius: 5px; width: 0%; transition: width 0.3s;"></div>
+                    </span></div>
+                    <div style="margin-top: 5px;">Frame: <span id="current-frame">0</span>/<span id="max-frames">0</span></div>
+                    <div>Time: <span id="elapsed-time">0:00</span> / <span id="session-duration">0:00</span></div>
+                    <div>Remaining: <span id="remaining-time">0:00</span></div>
+                </div>
             </div>
             
             <div class="controls">
                 <button id="connect-btn">Connect</button>
                 <button id="start-stream-btn" disabled>Start Stream</button>
                 <button id="stop-stream-btn" disabled>Stop Stream</button>
+                <button id="reset-session-btn" disabled style="background-color: #FF5722;">Reset Session</button>
                 <select id="scene-select" disabled>
                     <option value="universal">Universal Mode</option>
                     <option value="gta_drive">GTA Drive Mode</option>