Update

app.py

@@ -0,0 +1,435 @@
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "requests<3",
+#   "pillow",
+#   "opencv-python",
+#   "pyboy",
+#   "huggingface-hub",
+#   "gradio",
+#   "numpy",
+#   "nitrogen @ git+https://github.com/MineDojo/NitroGen.git@main",
+# ]
+# [tool.uv]
+# exclude-newer = "2025-12-22T00:00:00Z"
+# ///
+"""
+Unified Gradio app for NitroGen Pokemon Red player with real-time streaming
+Combines model inference and PyBoy gameplay in a single interface
+"""
+import gradio as gr
+from pathlib import Path
+import cv2
+import numpy as np
+from PIL import Image
+from pyboy import PyBoy
+from pyboy.utils import WindowEvent
+import time
+import tempfile
+import requests
+from huggingface_hub import HfFileSystem
+
+from nitrogen.inference_session import InferenceSession
+from nitrogen.shared import PATH_REPO, BUTTON_ACTION_TOKENS
+
+ROM_URL = "https://github.com/hxh-robb/pokemon-roms/raw/refs/heads/master/ROM/Pokemon%20-%20Red%20Version%20(USA,%20Europe).gb"
+STATE_PATH = "./init.state"
+
+# Game Boy button mapping
+GB_BUTTONS = {
+    "A": WindowEvent.PRESS_BUTTON_A,
+    "B": WindowEvent.PRESS_BUTTON_B,
+    "START": WindowEvent.PRESS_BUTTON_START,
+    "SELECT": WindowEvent.PRESS_BUTTON_SELECT,
+    "UP": WindowEvent.PRESS_ARROW_UP,
+    "DOWN": WindowEvent.PRESS_ARROW_DOWN,
+    "LEFT": WindowEvent.PRESS_ARROW_LEFT,
+    "RIGHT": WindowEvent.PRESS_ARROW_RIGHT,
+}
+
+GB_BUTTONS_RELEASE = {
+    "A": WindowEvent.RELEASE_BUTTON_A,
+    "B": WindowEvent.RELEASE_BUTTON_B,
+    "START": WindowEvent.RELEASE_BUTTON_START,
+    "SELECT": WindowEvent.RELEASE_BUTTON_SELECT,
+    "UP": WindowEvent.RELEASE_ARROW_UP,
+    "DOWN": WindowEvent.RELEASE_ARROW_DOWN,
+    "LEFT": WindowEvent.RELEASE_ARROW_LEFT,
+    "RIGHT": WindowEvent.RELEASE_ARROW_RIGHT,
+}
+
+def preprocess_img(frame):
+    """Convert Game Boy frame to 256x256 RGB PIL Image for model input"""
+    if isinstance(frame, Image.Image):
+        frame = np.array(frame)
+    
+    if len(frame.shape) == 2:
+        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
+    elif frame.shape[2] == 4:
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)
+    
+    frame_resized = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_AREA)
+    return Image.fromarray(frame_resized)
+
+def gamepad_to_gameboy_buttons(pred, button_threshold=0.5, joystick_threshold=0.3):
+    """Convert model's gamepad prediction to Game Boy button presses"""
+    j_left, j_right, buttons = pred["j_left"], pred["j_right"], pred["buttons"]
+    pressed_buttons = []
+    
+    if len(buttons) == 0:
+        return pressed_buttons
+    
+    button_vals = buttons[0]
+    
+    if len(button_vals) < len(BUTTON_ACTION_TOKENS):
+        return pressed_buttons
+    
+    # D-Pad mapping (indices 1-4)
+    if button_vals[1] > button_threshold:
+        pressed_buttons.append("DOWN")
+    if button_vals[2] > button_threshold:
+        pressed_buttons.append("LEFT")
+    if button_vals[3] > button_threshold:
+        pressed_buttons.append("RIGHT")
+    if button_vals[4] > button_threshold:
+        pressed_buttons.append("UP")
+    
+    # Joystick fallback if no D-pad pressed
+    if not any(b in pressed_buttons for b in ["UP", "DOWN", "LEFT", "RIGHT"]):
+        if len(j_left) > 0:
+            xl, yl = j_left[0]
+            if abs(xl) > joystick_threshold or abs(yl) > joystick_threshold:
+                if abs(xl) > abs(yl):
+                    if xl > joystick_threshold:
+                        pressed_buttons.append("RIGHT")
+                    elif xl < -joystick_threshold:
+                        pressed_buttons.append("LEFT")
+                else:
+                    if yl > joystick_threshold:
+                        pressed_buttons.append("DOWN")
+                    elif yl < -joystick_threshold:
+                        pressed_buttons.append("UP")
+    
+    # Action buttons
+    if button_vals[18] > button_threshold:  # SOUTH -> A
+        pressed_buttons.append("A")
+    if button_vals[5] > button_threshold:   # EAST -> B
+        pressed_buttons.append("B")
+    if button_vals[19] > button_threshold:  # START
+        pressed_buttons.append("START")
+    if button_vals[0] > button_threshold:   # BACK -> SELECT
+        pressed_buttons.append("SELECT")
+    
+    # Alternative mappings
+    if button_vals[10] > button_threshold and "A" not in pressed_buttons:  # NORTH -> A
+        pressed_buttons.append("A")
+    if button_vals[20] > button_threshold and "B" not in pressed_buttons:  # WEST -> B
+        pressed_buttons.append("B")
+    if button_vals[7] > button_threshold and "A" not in pressed_buttons:   # LEFT_SHOULDER -> A
+        pressed_buttons.append("A")
+    if button_vals[14] > button_threshold and "B" not in pressed_buttons:  # RIGHT_SHOULDER -> B
+        pressed_buttons.append("B")
+    
+    return pressed_buttons
+
+def play_pokemon(
+    cfg_scale: float,
+    context_length: int,
+    max_steps: int,
+    frame_skip: int,
+    button_threshold: float,
+    display_every: int,
+    update_delay: float
+):
+    """Generator that yields frames while playing Pokemon Red"""
+    
+    # Download ROM from URL
+    yield None, "⏳ Downloading ROM file...", None
+    try:
+        temp_dir = Path(tempfile.gettempdir())
+        rom_path = temp_dir / "PokemonRed.gb"
+        
+        # Download ROM if not already cached
+        if not rom_path.exists():
+            response = requests.get(ROM_URL, stream=True)
+            response.raise_for_status()
+            
+            with open(rom_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            
+            yield None, "✅ ROM downloaded successfully", None
+            time.sleep(0.5)
+        else:
+            yield None, "✅ Using cached ROM", None
+            time.sleep(0.3)
+    except Exception as e:
+        yield None, f"❌ Error downloading ROM: {str(e)}", None
+        return
+    
+    # Download checkpoint from HuggingFace using HfFileSystem
+    yield None, "⏳ Downloading checkpoint from nvidia/NitroGen...", None
+    try:
+        ckpt_path = temp_dir / "ng.pt"
+        
+        # Download checkpoint from HuggingFace Hub if not already cached
+        if not ckpt_path.exists():
+            hffs = HfFileSystem()
+            hffs.get_file("nvidia/NitroGen/ng.pt", str(ckpt_path))
+            
+            if not ckpt_path.exists():
+                yield None, "❌ Failed to download checkpoint from HuggingFace", None
+                return
+                
+            yield None, "✅ Checkpoint downloaded successfully", None
+            time.sleep(0.5)
+        else:
+            yield None, "✅ Using cached checkpoint", None
+            time.sleep(0.3)
+    except Exception as e:
+        yield None, f"❌ Error downloading checkpoint: {str(e)}", None
+        return
+    
+    # Initialize inference session
+    yield None, "⏳ Initializing inference session...", None
+    session = InferenceSession.from_ckpt(
+        str(ckpt_path), 
+        cfg_scale=cfg_scale, 
+        context_length=context_length
+    )
+    session.reset()
+    
+    # Initialize PyBoy
+    pyboy = PyBoy(str(rom_path), window="null")
+    pyboy.set_emulation_speed(0)  # Unlimited speed
+    
+    # Load save state if it exists
+    state_path = Path(STATE_PATH)
+    if state_path.exists():
+        with open(state_path, "rb") as f:
+            pyboy.load_state(f)
+        yield None, f"✅ Loaded save state: {STATE_PATH}", None
+        time.sleep(0.3)
+    else:
+        yield None, f"⚠️ Save state not found: {STATE_PATH} (starting fresh)", None
+        time.sleep(0.3)
+    
+    # Display settings
+    width, height = 640, 576
+    step_count = 0
+    
+    # Button timing: Press button briefly (4 frames), then release and wait
+    # This prevents holding buttons for too long (which would cause repeated movement)
+    # E.g., with frame_skip=16: press DOWN for 4 frames, release, wait 12 frames
+    # Result: Character moves 1 tile down, not 16 tiles
+    button_hold_frames = 4
+    
+    try:
+        while step_count < max_steps:
+            # Get screen and predict
+            screen = pyboy.screen.image
+            obs_processed = preprocess_img(screen)
+            pred = session.predict(obs_processed)
+            
+            # Convert to Game Boy buttons
+            buttons_to_press = gamepad_to_gameboy_buttons(pred, button_threshold)
+            
+            # Press buttons
+            for btn in buttons_to_press:
+                pyboy.send_input(GB_BUTTONS[btn])
+            
+            # Hold buttons for a few frames (so action registers)
+            pyboy.tick(button_hold_frames, render=False)
+            
+            # Release buttons
+            for btn in buttons_to_press:
+                pyboy.send_input(GB_BUTTONS_RELEASE[btn])
+            
+            # Tick remaining frames to complete the frame_skip cycle
+            remaining_frames = frame_skip - button_hold_frames
+            if remaining_frames > 1:
+                pyboy.tick(remaining_frames - 1, render=False)
+            if remaining_frames > 0:
+                pyboy.tick()  # Final tick with render
+            else:
+                pyboy.tick()  # Render at least once
+            
+            # Yield display update at specified frequency
+            if step_count % display_every == 0:
+                # Get frame (lightweight - no text overlay)
+                screen_np = pyboy.screen.ndarray
+                if screen_np.shape[2] == 4:
+                    screen_np = screen_np[:, :, :3]
+                
+                # Simple resize
+                frame_display = cv2.resize(
+                    screen_np, 
+                    (width, height), 
+                    interpolation=cv2.INTER_NEAREST
+                )
+                
+                # Create action info
+                action_info = f"**Step {step_count}/{max_steps}**\n\n"
+                action_info += f"🎮 **Buttons:** {', '.join(buttons_to_press) if buttons_to_press else 'None'}\n\n"
+                action_info += f"⚡ **Speed:** {frame_skip}x frame skip\n\n"
+                action_info += f"📊 **Progress:** {step_count/max_steps*100:.1f}%"
+                
+                # Create stats info
+                stats_info = f"**Inference Details**\n\n"
+                if len(pred.get("buttons", [])) > 0:
+                    button_vals = pred["buttons"][0]
+                    active_buttons = [
+                        f"{BUTTON_ACTION_TOKENS[i]}: {button_vals[i]:.2f}" 
+                        for i in range(min(len(button_vals), len(BUTTON_ACTION_TOKENS))) 
+                        if button_vals[i] > button_threshold
+                    ]
+                    if active_buttons:
+                        stats_info += "**Active Predictions:**\n"
+                        stats_info += "\n".join(f"- {btn}" for btn in active_buttons[:5])
+                    else:
+                        stats_info += "No buttons above threshold"
+                
+                # Yield frame and info (no encoding overhead)
+                yield frame_display, action_info, stats_info
+                
+                # Delay to allow Gradio to load images properly
+                time.sleep(update_delay)
+            
+            step_count += 1
+            
+    finally:
+        # Stop emulator
+        pyboy.stop()
+
+# Create Gradio interface
+with gr.Blocks(title="NitroGen Pokemon Red Player") as app:
+    gr.Markdown("# 🎮 NitroGen Pokemon Red Player")
+    gr.Markdown("Stream Pokemon Red gameplay powered by NitroGen AI model")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 🤖 Model Settings")
+            gr.Markdown("**Model:** nvidia/NitroGen (ng.pt) - automatically downloaded from HuggingFace Hub")
+            gr.Markdown("**ROM:** Automatically downloaded from configured URL")
+            gr.Markdown(f"**Save State:** {STATE_PATH}")
+            cfg_input = gr.Slider(
+                label="CFG Scale",
+                minimum=0.0,
+                maximum=3.0,
+                value=1.0,
+                step=0.1,
+                info="Classifier-free guidance scale"
+            )
+            ctx_input = gr.Slider(
+                label="Context Length",
+                minimum=1,
+                maximum=32,
+                value=1,
+                step=1,
+                info="Number of past frames to use"
+            )
+            
+            gr.Markdown("### ⚙️ Playback Settings")
+            max_steps_input = gr.Slider(
+                label="Max Steps",
+                minimum=100,
+                maximum=10000,
+                value=1000,
+                step=100,
+                info="Maximum inference steps"
+            )
+            frame_skip_input = gr.Slider(
+                label="Frame Skip",
+                minimum=1,
+                maximum=64,
+                value=16,
+                step=1,
+                info="Emulator frames per inference"
+            )
+            button_threshold_input = gr.Slider(
+                label="Button Threshold",
+                minimum=0.0,
+                maximum=1.0,
+                value=0.5,
+                step=0.05,
+                info="Threshold for button activation"
+            )
+            display_every_input = gr.Slider(
+                label="Display Every N Steps",
+                minimum=1,
+                maximum=10,
+                value=1,
+                step=1,
+                info="Update display frequency (1=every step, higher=faster but less frequent)"
+            )
+            update_delay_input = gr.Slider(
+                label="Update Delay (seconds)",
+                minimum=0.1,
+                maximum=3.0,
+                value=1.0,
+                step=0.1,
+                info="Wait time after each display update (higher=more time for image to load)"
+            )
+            
+            start_btn = gr.Button("🚀 Start Playing", variant="primary", size="lg")
+        
+        with gr.Column(scale=2):
+            image_output = gr.Image(
+                label="Game Stream",
+                height=600,
+                interactive=False
+            )
+            
+            with gr.Row():
+                with gr.Column():
+                    action_output = gr.Markdown(
+                        label="Actions",
+                        value="**Waiting to start...**"
+                    )
+                with gr.Column():
+                    stats_output = gr.Markdown(
+                        label="Statistics",
+                        value="**No data yet**"
+                    )
+            
+            gr.Markdown("""
+            ### 📝 Instructions
+            1. Adjust playback settings as needed
+            2. Click "Start Playing" to begin streaming
+            3. Game frames update in real-time with actions
+            
+            **Automatic Setup:**
+            - **Model**: nvidia/NitroGen checkpoint (ng.pt) from HuggingFace Hub
+            - **ROM**: Downloaded from configured URL
+            - **Save State**: Loaded from `./init.state` if available
+            - Model and ROM are cached in temp directory for faster subsequent runs
+            
+            **Tips:**
+            - **Display Every N Steps**: 1 = update every step, higher = faster but less frequent
+            - **Update Delay**: 1s default gives images time to load, reduce for faster updates
+            - **Frame Skip**: 16 = game runs 16 frames per inference (faster gameplay)
+            """)
+    
+    # Connect the button to the play function
+    start_btn.click(
+        fn=play_pokemon,
+        inputs=[
+            cfg_input,
+            ctx_input,
+            max_steps_input,
+            frame_skip_input,
+            button_threshold_input,
+            display_every_input,
+            update_delay_input
+        ],
+        outputs=[image_output, action_output, stats_output]
+    )
+
+if __name__ == "__main__":
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )
+

requirements.txt

@@ -0,0 +1,8 @@
+requests
+pillow
+opencv-python
+pyboy
+huggingface-hub
+gradio
+numpy
+nitrogen @ git+https://github.com/MineDojo/NitroGen.git@main