#!/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 )