From 73fd5772a7b3656d2d7524202f000ffc170b43e7 Mon Sep 17 00:00:00 2001 From: Sami Date: Wed, 4 Jun 2025 04:43:37 -0400 Subject: [PATCH 01/71] some thing idk --- requirements.txt | 9 +- webapp/action_builder.py | 319 +++++++++++++++++++++++++++++++++++++++ webapp/demo.py | 231 ++++++++++++++++++++++++++++ webapp/models.py | 131 ++++++++++++++++ webapp/render.py | 117 ++++++++++++++ webapp/samplers.py | 55 +++++++ 6 files changed, 860 insertions(+), 2 deletions(-) create mode 100644 webapp/action_builder.py create mode 100644 webapp/demo.py create mode 100644 webapp/models.py create mode 100644 webapp/render.py create mode 100644 webapp/samplers.py diff --git a/requirements.txt b/requirements.txt index 28f11ef2..31bfaa20 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,10 @@ opencv-python wandb einops -rotary-embeddings-pytorch -ema-pytorch \ No newline at end of file +rotary-embedding-torch +ema-pytorch +omegaconf +torch +toolz +termcolor +multimethod \ No newline at end of file diff --git a/webapp/action_builder.py b/webapp/action_builder.py new file mode 100644 index 00000000..65b3168f --- /dev/null +++ b/webapp/action_builder.py @@ -0,0 +1,319 @@ +import math +import torch +from enum import Enum +from dataclasses import dataclass +from typing import List, Dict, Optional, Union, Tuple, Callable + + +# Button mapping from the codebase +BUTTON_NAMES = ["W", "A", "S", "D", "LSHIFT", "SPACE", "R", "F", "E", "LMB", "RMB"] +BUTTON_INDICES = {name: idx for idx, name in enumerate(BUTTON_NAMES)} + + +class ActionPattern(Enum): + IDLE = "idle" + WALK_FORWARD = "walk_forward" + STRAFE_LEFT = "strafe_left" + STRAFE_RIGHT = "strafe_right" + WALK_BACKWARD = "walk_backward" + CIRCLE_STRAFE = "circle_strafe" + LOOK_AROUND = "look_around" + AIM_AND_SHOOT = "aim_and_shoot" + SPRINT_FORWARD = "sprint_forward" + RELOAD = "reload" + + +@dataclass(frozen=True) +class ActionConfig: + sequence_length: int + device: Union[str, torch.device] = 'cpu' + dtype: torch.dtype = torch.float32 + n_buttons: int = 11 + mouse_range: Tuple[float, float] = (-1.0, 1.0) + smooth_transitions: bool = True + random_seed: Optional[int] = None + + +class MouseGenerator: + @staticmethod + def _apply_smoothing(window_size: int, values: torch.Tensor) -> torch.Tensor: + """Apply smoothing to mouse movements using convolution.""" + if window_size <= 1: + return values + + kernel = torch.ones(window_size, device=values.device) / window_size + # Add padding to maintain sequence length + padding = window_size // 2 + padded = torch.nn.functional.pad(values, (padding, padding), mode='reflect') + return torch.nn.functional.conv1d( + padded.unsqueeze(0), + kernel.unsqueeze(0).unsqueeze(0), + padding=0 + ).squeeze(0) + + @staticmethod + def idle(config: ActionConfig) -> torch.Tensor: + """Generate idle mouse movement (minimal random drift).""" + return torch.randn(config.sequence_length, 2, + device=config.device, dtype=config.dtype) * 0.05 + + @staticmethod + def look_around(config: ActionConfig, + speed: float = 0.3, + amplitude: float = 0.7) -> torch.Tensor: + """Generate smooth looking around movement.""" + t = torch.linspace(0, 4 * math.pi, config.sequence_length, + device=config.device, dtype=config.dtype) + + # Create smooth sinusoidal movement + mouse_x = amplitude * torch.sin(t * speed) * torch.cos(t * speed * 0.3) + mouse_y = amplitude * torch.cos(t * speed * 0.7) * torch.sin(t * speed * 0.2) + + movement = torch.stack([mouse_x, mouse_y], dim=1) + + if config.smooth_transitions: + return MouseGenerator._apply_smoothing(5, movement.T).T + + return movement + + @staticmethod + def aim_tracking(config: ActionConfig, + target_speed: float = 0.1, + noise_level: float = 0.02) -> torch.Tensor: + """Generate aiming/tracking movement with micro-adjustments.""" + # Create base tracking movement + t = torch.linspace(0, 2 * math.pi, config.sequence_length, + device=config.device, dtype=config.dtype) + + # Smooth circular tracking + base_x = 0.3 * torch.sin(t * target_speed) + base_y = 0.2 * torch.cos(t * target_speed * 1.2) + + # Add realistic micro-movements + noise_x = torch.randn(config.sequence_length, device=config.device) * noise_level + noise_y = torch.randn(config.sequence_length, device=config.device) * noise_level + + return torch.stack([base_x + noise_x, base_y + noise_y], dim=1) + + @staticmethod + def custom_path(config: ActionConfig, + path_points: List[Tuple[float, float]], + interpolation: str = 'linear') -> torch.Tensor: + """Generate mouse movement following a custom path.""" + if len(path_points) < 2: + return MouseGenerator.idle(config) + + # Convert to tensors + points = torch.tensor(path_points, device=config.device, dtype=config.dtype) + + # Create interpolation indices + t = torch.linspace(0, len(points) - 1, config.sequence_length, + device=config.device, dtype=config.dtype) + + # Linear interpolation between points + indices = t.long() + weights = t - indices.float() + + # Handle edge case + indices = torch.clamp(indices, 0, len(points) - 2) + weights = weights.unsqueeze(1) + + interpolated = (1 - weights) * points[indices] + weights * points[indices + 1] + + return interpolated + + +class ButtonGenerator: + @staticmethod + def idle(config: ActionConfig) -> torch.Tensor: + """Generate idle button state (all buttons released).""" + return torch.zeros(config.sequence_length, config.n_buttons, + device=config.device, dtype=config.dtype) + + @staticmethod + def hold_buttons(config: ActionConfig, + button_names: List[str], + start_frame: int = 0, + duration: Optional[int] = None) -> torch.Tensor: + """Hold specific buttons for a duration.""" + buttons = torch.zeros(config.sequence_length, config.n_buttons, + device=config.device, dtype=config.dtype) + + end_frame = start_frame + (duration or config.sequence_length) + end_frame = min(end_frame, config.sequence_length) + + for button_name in button_names: + if button_name in BUTTON_INDICES: + idx = BUTTON_INDICES[button_name] + buttons[start_frame:end_frame, idx] = 1.0 + + return buttons + + @staticmethod + def tap_sequence(config: ActionConfig, + button_sequences: List[Tuple[str, int, int]]) -> torch.Tensor: + """Create button taps at specific times. + + Args: + button_sequences: List of (button_name, start_frame, duration) tuples + """ + buttons = torch.zeros(config.sequence_length, config.n_buttons, + device=config.device, dtype=config.dtype) + + for button_name, start_frame, duration in button_sequences: + if button_name in BUTTON_INDICES and start_frame < config.sequence_length: + idx = BUTTON_INDICES[button_name] + end_frame = min(start_frame + duration, config.sequence_length) + buttons[start_frame:end_frame, idx] = 1.0 + + return buttons + + @staticmethod + def pattern_from_name(config: ActionConfig, pattern: ActionPattern) -> torch.Tensor: + """Generate button pattern from predefined patterns.""" + if pattern == ActionPattern.WALK_FORWARD: + return ButtonGenerator.hold_buttons(config, ["W"]) + elif pattern == ActionPattern.STRAFE_LEFT: + return ButtonGenerator.hold_buttons(config, ["A"]) + elif pattern == ActionPattern.STRAFE_RIGHT: + return ButtonGenerator.hold_buttons(config, ["D"]) + elif pattern == ActionPattern.WALK_BACKWARD: + return ButtonGenerator.hold_buttons(config, ["S"]) + elif pattern == ActionPattern.SPRINT_FORWARD: + return ButtonGenerator.hold_buttons(config, ["W", "LSHIFT"]) + elif pattern == ActionPattern.RELOAD: + return ButtonGenerator.tap_sequence(config, [("R", 10, 20)]) + else: + return ButtonGenerator.idle(config) + + +class ActionSequenceBuilder: + def __init__(self, config: ActionConfig): + self.config = config + self.mouse_sequence = torch.zeros(config.sequence_length, 2, + device=config.device, dtype=config.dtype) + self.button_sequence = torch.zeros(config.sequence_length, config.n_buttons, + device=config.device, dtype=config.dtype) + + def add_mouse_segment(self, + start_frame: int, + end_frame: int, + generator_func: Callable[[ActionConfig], torch.Tensor], + **kwargs) -> 'ActionSequenceBuilder': + """Add a mouse movement segment.""" + segment_config = ActionConfig( + sequence_length=end_frame - start_frame, + device=self.config.device, + dtype=self.config.dtype, + n_buttons=self.config.n_buttons + ) + + segment = generator_func(segment_config, **kwargs) + self.mouse_sequence[start_frame:end_frame] = segment + return self + + def add_button_segment(self, + start_frame: int, + end_frame: int, + generator_func: Callable[[ActionConfig], torch.Tensor], + **kwargs) -> 'ActionSequenceBuilder': + """Add a button press segment.""" + segment_config = ActionConfig( + sequence_length=end_frame - start_frame, + device=self.config.device, + dtype=self.config.dtype, + n_buttons=self.config.n_buttons + ) + + segment = generator_func(segment_config, **kwargs) + self.button_sequence[start_frame:end_frame] = segment + return self + + def build(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Build and return the final action sequence.""" + return self.mouse_sequence, self.button_sequence + + +class ActionSequenceGenerator: + def __init__(self, config: ActionConfig): + self.config = config + if config.random_seed is not None: + torch.manual_seed(config.random_seed) + + def generate_pattern(self, + pattern: ActionPattern, + mouse_kwargs: Optional[Dict] = None, + button_kwargs: Optional[Dict] = None) -> Tuple[torch.Tensor, torch.Tensor]: + """Generate actions for a predefined pattern.""" + mouse_kwargs = mouse_kwargs or {} + button_kwargs = button_kwargs or {} + + if pattern == ActionPattern.IDLE: + mouse = MouseGenerator.idle(self.config) + buttons = ButtonGenerator.idle(self.config) + + elif pattern == ActionPattern.LOOK_AROUND: + mouse = MouseGenerator.look_around(self.config, **mouse_kwargs) + buttons = ButtonGenerator.idle(self.config) + + elif pattern == ActionPattern.AIM_AND_SHOOT: + mouse = MouseGenerator.aim_tracking(self.config, **mouse_kwargs) + # Add some shooting + shoot_times = [(i * 30, 5) for i in range(self.config.sequence_length // 30)] + buttons = ButtonGenerator.tap_sequence( + self.config, + [("LMB", start, dur) for start, dur in shoot_times] + ) + + elif pattern == ActionPattern.CIRCLE_STRAFE: + # Combine circular mouse movement with strafing + mouse = MouseGenerator.look_around(self.config, speed=0.2, amplitude=0.5) + buttons = ButtonGenerator.hold_buttons(self.config, ["A", "W"]) + + else: + mouse = MouseGenerator.idle(self.config) + buttons = ButtonGenerator.pattern_from_name(self.config, pattern) + + return mouse, buttons + + def generate_custom_sequence(self) -> ActionSequenceBuilder: + """Get a builder for creating custom action sequences.""" + return ActionSequenceBuilder(self.config) + + def generate_batch(self, + batch_size: int, + patterns: Optional[List[ActionPattern]] = None) -> Tuple[torch.Tensor, torch.Tensor]: + """Generate a batch of action sequences.""" + if patterns is None: + patterns = [ActionPattern.IDLE] * batch_size + + if len(patterns) != batch_size: + # Repeat or truncate patterns to match batch size + patterns = (patterns * (batch_size // len(patterns) + 1))[:batch_size] + + mouse_batch = [] + button_batch = [] + + for pattern in patterns: + mouse, buttons = self.generate_pattern(pattern) + mouse_batch.append(mouse) + button_batch.append(buttons) + + return torch.stack(mouse_batch), torch.stack(button_batch) + + +if __name__ == "__main__": + print("=== Action Generation Examples ===") + config = ActionConfig(sequence_length=100) + builder = ActionSequenceGenerator(config).generate_custom_sequence() + + mouse_custom, button_custom = (builder + .add_mouse_segment(0, 30, MouseGenerator.idle) + .add_mouse_segment(30, 70, MouseGenerator.look_around, speed=0.5) + .add_mouse_segment(70, 100, MouseGenerator.aim_tracking) + .add_button_segment(0, 50, ButtonGenerator.hold_buttons, button_names=["W"]) + .add_button_segment(50, 100, ButtonGenerator.hold_buttons, button_names=["W", "LSHIFT"]) + .build()) + + print(f"Custom sequence mouse shape: {mouse_custom.shape}, button shape: {button_custom.shape}") + print("Button names mapping:", BUTTON_INDICES) \ No newline at end of file diff --git a/webapp/demo.py b/webapp/demo.py new file mode 100644 index 00000000..bb516061 --- /dev/null +++ b/webapp/demo.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +""" +Demo script for OWL-WMS video generation pipeline. + +This script demonstrates how to use the complete pipeline to generate videos +with different action patterns and configurations. +""" + +import argparse +import sys +from pathlib import Path + +# Add webapp to path +sys.path.append(str(Path(__file__).parent)) + +from render import VideoRenderer, RenderConfig, quick_render, batch_render +from action_builder import ActionPattern, ActionSequenceGenerator, ActionConfig, MouseGenerator, ButtonGenerator + + +def demo_quick_renders(): + """Demonstrate quick rendering with different patterns.""" + print("🎮 Demo: Quick Renders") + print("-" * 30) + + patterns_to_try = [ + ActionPattern.LOOK_AROUND, + ActionPattern.AIM_AND_SHOOT, + ActionPattern.SPRINT_FORWARD, + ActionPattern.CIRCLE_STRAFE + ] + + for pattern in patterns_to_try: + print(f"Generating video with pattern: {pattern.value}") + result = quick_render( + pattern=pattern, + sequence_length=64, + output_dir=f"demo_output/{pattern.value}" + ) + print(f"✅ Generated: {result['video_path']}") + print(f" Render time: {result['render_time']:.2f}s") + print() + + +def demo_batch_render(): + """Demonstrate batch rendering with multiple patterns.""" + print("🎮 Demo: Batch Render") + print("-" * 30) + + patterns = [ + ActionPattern.WALK_FORWARD, + ActionPattern.STRAFE_LEFT, + ActionPattern.STRAFE_RIGHT, + ActionPattern.WALK_BACKWARD + ] + + print(f"Generating batch video with {len(patterns)} different patterns...") + result = batch_render( + patterns=patterns, + sequence_length=64, + output_dir="demo_output/batch" + ) + + print(f"✅ Generated batch: {result['video_path']}") + print(f" Video shape: {result['video_shape']}") + print(f" Render time: {result['render_time']:.2f}s") + print() + + +def demo_custom_actions(): + """Demonstrate custom action sequence creation.""" + print("🎮 Demo: Custom Action Sequences") + print("-" * 30) + + # Create a complex custom sequence + config = ActionConfig(sequence_length=120, random_seed=42) + generator = ActionSequenceGenerator(config) + + builder = generator.generate_custom_sequence() + + # Build a complex sequence: idle -> look around -> aim and move -> sprint + mouse_sequence, button_sequence = (builder + .add_mouse_segment(0, 30, MouseGenerator.idle) + .add_button_segment(0, 30, ButtonGenerator.idle) + + .add_mouse_segment(30, 60, MouseGenerator.look_around, speed=0.4, amplitude=0.6) + .add_button_segment(30, 60, ButtonGenerator.idle) + + .add_mouse_segment(60, 90, MouseGenerator.aim_tracking, target_speed=0.2) + .add_button_segment(60, 90, ButtonGenerator.hold_buttons, button_names=["W", "A"]) + + .add_mouse_segment(90, 120, MouseGenerator.look_around, speed=0.6) + .add_button_segment(90, 120, ButtonGenerator.hold_buttons, button_names=["W", "LSHIFT"]) + .build()) + + # Use these custom actions with the renderer + render_config = RenderConfig( + sequence_length=120, + action_pattern=ActionPattern.IDLE, # This will be overridden + random_seed=42, + output_dir="demo_output/custom" + ) + + renderer = VideoRenderer(render_config) + + # Manually set the actions and render + renderer.load_model() + renderer.create_sampler() + + print("Generating video with custom action sequence...") + generated_video = renderer.synthesize_video( + mouse_sequence.unsqueeze(0), # Add batch dimension + button_sequence.unsqueeze(0) + ) + + # Save the results + action_paths = renderer.save_actions( + mouse_sequence.unsqueeze(0), + button_sequence.unsqueeze(0), + "custom_demo" + ) + video_path = renderer.save_video(generated_video, "custom_demo") + + print(f"✅ Generated custom video: {video_path}") + print(f" Action files: {action_paths}") + print() + + +def demo_action_patterns(): + """Demonstrate all available action patterns.""" + print("🎮 Demo: All Action Patterns") + print("-" * 30) + + config = ActionConfig(sequence_length=60) + generator = ActionSequenceGenerator(config) + + print("Available action patterns:") + for pattern in ActionPattern: + print(f" - {pattern.value}") + + # Generate and show stats for each pattern + mouse, buttons = generator.generate_pattern(pattern) + mouse_range = f"[{mouse.min().item():.3f}, {mouse.max().item():.3f}]" + button_count = buttons.sum().item() + + print(f" Mouse range: {mouse_range}, Button presses: {button_count:.0f}") + + print() + + +def demo_with_args(): + """Demo with command line arguments.""" + parser = argparse.ArgumentParser(description="OWL-WMS Video Generation Demo") + parser.add_argument("--pattern", type=str, default="look_around", + choices=[p.value for p in ActionPattern], + help="Action pattern to use") + parser.add_argument("--length", type=int, default=64, + help="Sequence length") + parser.add_argument("--output", type=str, default="demo_output", + help="Output directory") + parser.add_argument("--device", type=str, default="auto", + help="Device to use (cuda/cpu/auto)") + parser.add_argument("--cfg-scale", type=float, default=1.3, + help="CFG scale for sampling") + parser.add_argument("--seed", type=int, default=None, + help="Random seed") + + args = parser.parse_args() + + # Find the pattern enum + pattern = None + for p in ActionPattern: + if p.value == args.pattern: + pattern = p + break + + if pattern is None: + print(f"❌ Unknown pattern: {args.pattern}") + return + + print(f"🎮 Custom Demo: {pattern.value}") + print("-" * 30) + + config = RenderConfig( + sequence_length=args.length, + action_pattern=pattern, + device=args.device, + output_dir=args.output, + cfg_scale=args.cfg_scale, + random_seed=args.seed + ) + + renderer = VideoRenderer(config) + result = renderer.render_video(filename_prefix=f"demo_{pattern.value}") + + print(f"✅ Generated: {result['video_path']}") + print(f" Shape: {result['video_shape']}") + print(f" Time: {result['render_time']:.2f}s") + + +def main(): + """Run all demos or handle command line arguments.""" + if len(sys.argv) > 1: + # Handle command line arguments + demo_with_args() + return + + print("🎬 OWL-WMS Video Generation Demo Suite") + print("=" * 50) + print() + + try: + # Run all demos + demo_action_patterns() + demo_quick_renders() + demo_batch_render() + demo_custom_actions() + + print("🎉 All demos completed successfully!") + print("📁 Check the demo_output/ directory for generated videos.") + print() + print("💡 Tip: Run with command line arguments for custom generation:") + print(" python demo.py --pattern aim_and_shoot --length 128 --seed 42") + + except Exception as e: + print(f"❌ Demo failed: {e}") + import traceback + traceback.print_exc() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/webapp/models.py b/webapp/models.py new file mode 100644 index 00000000..96ad9066 --- /dev/null +++ b/webapp/models.py @@ -0,0 +1,131 @@ +import torch +import torch.nn as nn +from pathlib import Path +from typing import Dict, Optional, Union +from dataclasses import dataclass + +from termcolor import colored + +from owl_wms.models import get_model_cls +from owl_wms.configs import Config + + +@dataclass(frozen=True) +class ModelPaths: + config: Path + checkpoint: Path + + @classmethod + def from_strings(cls, config_path: str, checkpoint_path: str) -> 'ModelPaths': + return cls( + config=Path(config_path), + checkpoint=Path(checkpoint_path) + ) + + def validate(self) -> None: + if not self.config.exists(): + raise FileNotFoundError(f"Config file not found: {self.config}") + if not self.checkpoint.exists(): + raise FileNotFoundError(f"Checkpoint file not found: {self.checkpoint}") + + +class ModelLoader: + DEFAULT_PATHS = ModelPaths.from_strings( + config_path='/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/basic.yml', + checkpoint_path='/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/ckpt_165k_ema.pt' + ) + + def __init__(self, paths: Optional[ModelPaths] = None): + self.paths = paths or self.DEFAULT_PATHS + self.paths.validate() + + @staticmethod + def _append_state_dict_prefix(state_dict: Dict, prefix: str = 'core.') -> Dict: + return {prefix+key: value for key, value in state_dict.items()} + + @staticmethod + def _count_parameters(model: nn.Module) -> int: + return sum(p.numel() for p in model.parameters()) + + def _load_config(self) -> Config: + return Config.from_yaml(str(self.paths.config)) + + def _load_checkpoint(self) -> Dict: + return torch.load(str(self.paths.checkpoint), map_location='cpu') + + def load_model(self, + device: Optional[Union[str, torch.device]] = None, + eval_mode: bool = True, + verbose: bool = True) -> nn.Module: + # Load configuration and create model + config = self._load_config() + model_cls = get_model_cls(config.model.model_id) + model = model_cls(config.model) + + # Load and filter state dict + checkpoint = self._load_checkpoint() + filtered_state_dict = self._append_state_dict_prefix(checkpoint) + model.load_state_dict(filtered_state_dict) + + # Configure model + if eval_mode: + model.eval() + + if device is not None: + model = model.to(device) + + # Print model information if requested + if verbose: + param_count = self._count_parameters(model) + print(f'{colored("Model loaded", "blue")}\t\t {colored("successfully", "green")}') + print(f'{colored("Parameters", "blue")} \t\t {colored(f"{param_count:,}", "green")}') + print(f'{colored("Config", "blue")} \t\t {colored(str(self.paths.config), "green", attrs=["bold"])}') + print(f'{colored("Checkpoint", "blue")} \t\t {colored(str(self.paths.checkpoint), "green", attrs=["bold"])}') + + return model + + +def load_model(config_path: Optional[str] = None, + checkpoint_path: Optional[str] = None, + device: Optional[Union[str, torch.device]] = None, + eval_mode: bool = True, + verbose: bool = True) -> nn.Module: + """ + Convenience function for loading models with custom paths. + + Args: + config_path: Path to model configuration file + checkpoint_path: Path to model checkpoint file + device: Target device for the model + eval_mode: Whether to set model to evaluation mode + verbose: Whether to print model information + + Returns: + Loaded PyTorch model + """ + if config_path or checkpoint_path: + # Use custom paths if provided + paths = ModelPaths.from_strings( + config_path or str(ModelLoader.DEFAULT_PATHS.config), + checkpoint_path or str(ModelLoader.DEFAULT_PATHS.checkpoint) + ) + loader = ModelLoader(paths) + else: + # Use default paths + loader = ModelLoader() + + return loader.load_model(device=device, eval_mode=eval_mode, verbose=verbose) + + +if __name__ == "__main__": + # Example usage with different approaches + + # Method 1: Using the class directly + print("=== Loading model using ModelLoader class ===") + loader = ModelLoader() + model = loader.load_model() + print(f"Model type: {type(model).__name__}") + + # Method 2: Using the convenience function + print("=== Loading model using convenience function ===") + model = load_model() \ No newline at end of file diff --git a/webapp/render.py b/webapp/render.py new file mode 100644 index 00000000..b586b056 --- /dev/null +++ b/webapp/render.py @@ -0,0 +1,117 @@ +import torch +from pathlib import Path + +from webapp.models import load_model +from webapp.samplers import create_sampler +from webapp.action_builder import ActionSequenceGenerator, ActionConfig, ActionPattern + +# Global configuration - easy to modify +SEQUENCE_LENGTH = 128 +DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' +OUTPUT_DIR = "generated_videos" +SAMPLER_TYPE = 'cfg' +DEFAULT_PATTERN = ActionPattern.LOOK_AROUND + + +def setup_output_dir(): + """Create output directory.""" + Path(OUTPUT_DIR).mkdir(exist_ok=True) + + +def generate_dummy_actions(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH): + """Generate dummy actions for video conditioning.""" + config = ActionConfig( + sequence_length=length, + device=DEVICE, + dtype=torch.float32 + ) + + generator = ActionSequenceGenerator(config) + mouse, buttons = generator.generate_pattern(pattern) + + # Add batch dimension + return mouse.unsqueeze(0), buttons.unsqueeze(0) + + +def synthesize_video(mouse_actions, button_actions, model, sampler): + """Generate video using model and sampler.""" + batch_size, sequence_length = mouse_actions.shape[:2] + + # Create dummy latent batch - hardcoded dimensions from config + dummy_batch = torch.randn( + batch_size, sequence_length, 128, 16, 16, + device=DEVICE, dtype=torch.float32 + ) + + # Ensure actions are on correct device + mouse_actions = mouse_actions.to(DEVICE) + button_actions = button_actions.to(DEVICE) + + # Generate video + with torch.no_grad(): + video = sampler( + dummy_batch=dummy_batch, + mouse=mouse_actions, + btn=button_actions + ) + + return video + + +def save_video(video_tensor, filename="generated_video"): + """Save video tensor to file.""" + setup_output_dir() + import time + timestamp = int(time.time()) + path = Path(OUTPUT_DIR) / f"{filename}_{timestamp}.pt" + torch.save(video_tensor.cpu(), path) + return str(path) + + +def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): + """Simple video rendering - just load, generate, save.""" + if verbose: + print(f"🎬 Rendering video with pattern: {pattern.value}") + + # Load model and create sampler + if verbose: + print("Loading model...") + model = load_model(device=DEVICE, verbose=verbose) + + if verbose: + print("Creating sampler...") + sampler = create_sampler(SAMPLER_TYPE, model) + + # Generate actions + if verbose: + print("Generating actions...") + mouse, buttons = generate_dummy_actions(pattern, length) + + # Synthesize video + if verbose: + print("Synthesizing video...") + video = synthesize_video(mouse, buttons, model, sampler) + + # Save video + if verbose: + print("Saving video...") + path = save_video(video, f"render_{pattern.value}") + + if verbose: + print(f"✅ Done! Video saved to: {path}") + print(f" Video shape: {video.shape}") + + return path + + +if __name__ == "__main__": + # Simple usage examples + print("🎮 Simple OWL-WMS Video Renderer") + + # Render with default settings + render_video() + + # Render with different pattern + render_video(ActionPattern.AIM_AND_SHOOT, length=64) + + diff --git a/webapp/samplers.py b/webapp/samplers.py new file mode 100644 index 00000000..7716421f --- /dev/null +++ b/webapp/samplers.py @@ -0,0 +1,55 @@ +import torch +from torch import nn, Tensor +from typing import Literal, Callable +from functools import partial, cache +from multimethod import multimethod +from owl_wms.sampling.cfg import CFGSampler, InpaintCFGSampler # TODO is there 'window' sampler? +from webapp.models import load_model + +SAMPLING_STEPS = 64 +BATCH_SHAPE = torch.empty(1, 128, 16, 256, 256) # TODO fix this shit +DECODE_FN = None +SCALE = 1 +CFG_SCALE = 1.3 + +MouseData = Tensor +ButtonData = Tensor +VideoData = Tensor + +@multimethod +def create_sampler(sampler_id: Literal['cfg'], model: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: + @cache # simple singleton + def _sampler(): return CFGSampler() + + return partial( + _sampler().__call__, + sampling_steps=SAMPLING_STEPS, + decode_fn=DECODE_FN, + scale=SCALE, + cfg_scale=CFG_SCALE, + model=model + ) + +@multimethod +def create_sampler(sampler_id: Literal['inpaint_cfg'], model: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: + @cache + def _sampler(): return InpaintCFGSampler() + + return partial( + _sampler().__call__, + sampling_steps=SAMPLING_STEPS, + decode_fn=DECODE_FN, + scale=SCALE, + cfg_scale=CFG_SCALE, + model=model + ) + +@multimethod +def create_sampler(sampler_id: Literal['window'], model: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: + raise NotImplementedError("Window sampler not implemented") + + +model = load_model() +cfg_sampler = create_sampler('cfg', model) + + From ba153a236594e9650dc1a98b14d93bccb69691d7 Mon Sep 17 00:00:00 2001 From: Sami Date: Wed, 4 Jun 2025 19:37:54 -0400 Subject: [PATCH 02/71] sampling and video generation --- owl_wms/models/__init__.py | 4 +- owl_wms/models/gamerft.py | 2 +- owl_wms/nn/embeddings.py | 2 +- requirements.txt | 9 +++- webapp/action_builder.py | 18 ++++--- webapp/demo.py | 10 ++-- webapp/models.py | 41 ++++++++++++--- webapp/render.py | 105 +++++++++++++++++++++++++++++-------- webapp/samplers.py | 27 ++++------ 9 files changed, 154 insertions(+), 64 deletions(-) diff --git a/owl_wms/models/__init__.py b/owl_wms/models/__init__.py index 34fcc572..53742586 100644 --- a/owl_wms/models/__init__.py +++ b/owl_wms/models/__init__.py @@ -1,8 +1,10 @@ -from .gamerft import GameRFT +from .gamerft import GameRFT, GameRFTCore def get_model_cls(model_id): if model_id == "game_rft": return GameRFT + if model_id == "game_rft_core": + return GameRFTCore diff --git a/owl_wms/models/gamerft.py b/owl_wms/models/gamerft.py index 3560984e..1e56c1d5 100644 --- a/owl_wms/models/gamerft.py +++ b/owl_wms/models/gamerft.py @@ -18,7 +18,7 @@ class GameRFTCore(nn.Module): def __init__(self, config): super().__init__() - + self.config = config self.transformer = UViT(config) self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) self.t_embed = TimestepEmbedding(config.d_model) diff --git a/owl_wms/nn/embeddings.py b/owl_wms/nn/embeddings.py index 01a6c1c2..465670a6 100644 --- a/owl_wms/nn/embeddings.py +++ b/owl_wms/nn/embeddings.py @@ -88,7 +88,7 @@ def forward(self, q, k): class LearnedPosEnc(nn.Module): def __init__(self, n_seq, dim): super().__init__() - + # n_seq is tokens_per_frame * n_frames, dim is d_model self.p = nn.Parameter(torch.randn(n_seq,dim)*0.02) def forward(self, x): diff --git a/requirements.txt b/requirements.txt index 31bfaa20..8c53b2cc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,11 @@ omegaconf torch toolz termcolor -multimethod \ No newline at end of file +multimethod +diffusers +vector-quantize-pytorch +torchtyping +imageio +einops +numpy +imageio[ffmpeg] diff --git a/webapp/action_builder.py b/webapp/action_builder.py index 65b3168f..d0e92288 100644 --- a/webapp/action_builder.py +++ b/webapp/action_builder.py @@ -41,15 +41,17 @@ def _apply_smoothing(window_size: int, values: torch.Tensor) -> torch.Tensor: if window_size <= 1: return values - kernel = torch.ones(window_size, device=values.device) / window_size - # Add padding to maintain sequence length + # values is [2, sequence_length] for mouse x,y coordinates + # Use groups=2 to smooth each channel independently + kernel = torch.ones(2, 1, window_size, device=values.device) / window_size padding = window_size // 2 padded = torch.nn.functional.pad(values, (padding, padding), mode='reflect') - return torch.nn.functional.conv1d( - padded.unsqueeze(0), - kernel.unsqueeze(0).unsqueeze(0), - padding=0 - ).squeeze(0) + smoothed = torch.nn.functional.conv1d( + padded.unsqueeze(0), # Add batch dim: [1, 2, sequence_length] + kernel, padding=0, + groups=2 # Each output channel only depends on corresponding input channel + ).squeeze(0) # Remove batch dim: [2, sequence_length] + return smoothed @staticmethod def idle(config: ActionConfig) -> torch.Tensor: @@ -314,6 +316,6 @@ def generate_batch(self, .add_button_segment(0, 50, ButtonGenerator.hold_buttons, button_names=["W"]) .add_button_segment(50, 100, ButtonGenerator.hold_buttons, button_names=["W", "LSHIFT"]) .build()) - + print(f"Custom sequence mouse shape: {mouse_custom.shape}, button shape: {button_custom.shape}") print("Button names mapping:", BUTTON_INDICES) \ No newline at end of file diff --git a/webapp/demo.py b/webapp/demo.py index bb516061..0054ae54 100644 --- a/webapp/demo.py +++ b/webapp/demo.py @@ -13,8 +13,8 @@ # Add webapp to path sys.path.append(str(Path(__file__).parent)) -from render import VideoRenderer, RenderConfig, quick_render, batch_render -from action_builder import ActionPattern, ActionSequenceGenerator, ActionConfig, MouseGenerator, ButtonGenerator +from webapp.render import render_video, save_video +from webapp.action_builder import ActionPattern, ActionSequenceGenerator, ActionConfig, MouseGenerator, ButtonGenerator def demo_quick_renders(): @@ -33,7 +33,7 @@ def demo_quick_renders(): print(f"Generating video with pattern: {pattern.value}") result = quick_render( pattern=pattern, - sequence_length=64, + sequence_length=60, output_dir=f"demo_output/{pattern.value}" ) print(f"✅ Generated: {result['video_path']}") @@ -56,7 +56,7 @@ def demo_batch_render(): print(f"Generating batch video with {len(patterns)} different patterns...") result = batch_render( patterns=patterns, - sequence_length=64, + sequence_length=60, output_dir="demo_output/batch" ) @@ -153,7 +153,7 @@ def demo_with_args(): parser.add_argument("--pattern", type=str, default="look_around", choices=[p.value for p in ActionPattern], help="Action pattern to use") - parser.add_argument("--length", type=int, default=64, + parser.add_argument("--length", type=int, default=60, help="Sequence length") parser.add_argument("--output", type=str, default="demo_output", help="Output directory") diff --git a/webapp/models.py b/webapp/models.py index 96ad9066..7c220d6e 100644 --- a/webapp/models.py +++ b/webapp/models.py @@ -8,7 +8,8 @@ from owl_wms.models import get_model_cls from owl_wms.configs import Config - +from owl_wms.utils.owl_vae_bridge import get_decoder_only +from owl_wms.utils import freeze @dataclass(frozen=True) class ModelPaths: @@ -64,8 +65,11 @@ def load_model(self, # Load and filter state dict checkpoint = self._load_checkpoint() - filtered_state_dict = self._append_state_dict_prefix(checkpoint) - model.load_state_dict(filtered_state_dict) + + if config.model.model_id == "game_rft": + checkpoint = self._append_state_dict_prefix(checkpoint) + + model.load_state_dict(checkpoint) # Configure model if eval_mode: @@ -83,13 +87,34 @@ def load_model(self, print(f'{colored("Checkpoint", "blue")} \t\t {colored(str(self.paths.checkpoint), "green", attrs=["bold"])}') return model + + def load_decoder(self, + device: Optional[Union[str, torch.device]] = None, + eval_mode: bool = True, + verbose: bool = True) -> nn.Module: + decoder = get_decoder_only() + freeze(decoder) + + if verbose: + print(f'{colored("Decoder loaded", "blue")}\t\t {colored("successfully", "green")}') + print(f'{colored("Parameters", "blue")} \t\t {colored(f"{self._count_parameters(decoder):,}", "green")}') + print(f'{colored("Config", "blue")} \t\t {colored(str(self.paths.config), "green", attrs=["bold"])}') + print(f'{colored("Checkpoint", "blue")} \t\t {colored(str(self.paths.checkpoint), "green", attrs=["bold"])}') + + if device is not None: + decoder = decoder.to(device) + + if eval_mode: + decoder.eval() + + return decoder -def load_model(config_path: Optional[str] = None, +def load_models(config_path: Optional[str] = None, checkpoint_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, eval_mode: bool = True, - verbose: bool = True) -> nn.Module: + verbose: bool = True) -> tuple[nn.Module, nn.Module]: """ Convenience function for loading models with custom paths. @@ -114,7 +139,9 @@ def load_model(config_path: Optional[str] = None, # Use default paths loader = ModelLoader() - return loader.load_model(device=device, eval_mode=eval_mode, verbose=verbose) + encoder = loader.load_model(device=device, eval_mode=eval_mode, verbose=verbose) + decoder = loader.load_decoder(device=device, eval_mode=eval_mode, verbose=verbose) + return encoder, decoder if __name__ == "__main__": @@ -128,4 +155,4 @@ def load_model(config_path: Optional[str] = None, # Method 2: Using the convenience function print("=== Loading model using convenience function ===") - model = load_model() \ No newline at end of file + encoder, decoder = load_models() \ No newline at end of file diff --git a/webapp/render.py b/webapp/render.py index b586b056..4ec63443 100644 --- a/webapp/render.py +++ b/webapp/render.py @@ -1,21 +1,30 @@ import torch +import math +import time +from datetime import datetime from pathlib import Path +import imageio +import numpy as np -from webapp.models import load_model +import einops as eo + +from webapp.models import load_models from webapp.samplers import create_sampler from webapp.action_builder import ActionSequenceGenerator, ActionConfig, ActionPattern -# Global configuration - easy to modify -SEQUENCE_LENGTH = 128 +HEIGHT = 256 +WIDTH = 256 +D_MODEL = 1024 +CHANNELS = 128 +SEQUENCE_LENGTH = 60 +TOKENS_PER_FRAME = 16 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' OUTPUT_DIR = "generated_videos" SAMPLER_TYPE = 'cfg' DEFAULT_PATTERN = ActionPattern.LOOK_AROUND -def setup_output_dir(): - """Create output directory.""" - Path(OUTPUT_DIR).mkdir(exist_ok=True) +def setup_output_dir(): Path(OUTPUT_DIR).mkdir(exist_ok=True) def generate_dummy_actions(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH): @@ -33,13 +42,15 @@ def generate_dummy_actions(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH): return mouse.unsqueeze(0), buttons.unsqueeze(0) -def synthesize_video(mouse_actions, button_actions, model, sampler): +def synthesize_video(mouse_actions, button_actions, encoder, decoder, sampler): """Generate video using model and sampler.""" batch_size, sequence_length = mouse_actions.shape[:2] # Create dummy latent batch - hardcoded dimensions from config dummy_batch = torch.randn( - batch_size, sequence_length, 128, 16, 16, + batch_size, sequence_length, CHANNELS, + int(math.sqrt(TOKENS_PER_FRAME)), # H + int(math.sqrt(TOKENS_PER_FRAME)), # W device=DEVICE, dtype=torch.float32 ) @@ -48,7 +59,7 @@ def synthesize_video(mouse_actions, button_actions, model, sampler): button_actions = button_actions.to(DEVICE) # Generate video - with torch.no_grad(): + with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): video = sampler( dummy_batch=dummy_batch, mouse=mouse_actions, @@ -58,14 +69,64 @@ def synthesize_video(mouse_actions, button_actions, model, sampler): return video -def save_video(video_tensor, filename="generated_video"): - """Save video tensor to file.""" +def save_video(video_tensor: torch.Tensor, filename="generated_video", fps=30): + """ + Save video tensor as MP4 file. + + Args: + video_tensor: Tensor with shape [batch_size, sequence_length, channels, height, width] + Expected range: [-1, 1] (VAE decoder output) + filename: Base filename (without extension) + fps: Frames per second for output video + + Returns: + str: Path to saved video file + """ setup_output_dir() - import time - timestamp = int(time.time()) - path = Path(OUTPUT_DIR) / f"{filename}_{timestamp}.pt" - torch.save(video_tensor.cpu(), path) - return str(path) + + # Generate timestamp for unique filename + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + output_path = Path(OUTPUT_DIR) / f"{filename}_{timestamp}.mp4" + + # Convert tensor to numpy and handle batch dimension + video_np: np.ndarray = video_tensor.float().cpu().detach().numpy() + + # Take first batch item if batch_size > 1 + if video_np.ndim == 5: # [batch, seq, channels, height, width] + video_np = video_np[0] # Take first batch item: [seq, channels, height, width] + + # Convert from [seq, channels, height, width] to [seq, height, width, channels] + video_np = eo.rearrange(video_np, 'seq c h w -> seq h w c') + + # Denormalize from actual range to [0,255] + print(f'Video range before denorm: [{video_np.min():.3f}, {video_np.max():.3f}]') + + # Normalize to [0, 1] using actual min/max + video_min, video_max = video_np.min(), video_np.max() + if video_max > video_min: # Avoid division by zero + video_np = (video_np - video_min) / (video_max - video_min) # -> [0, 1] + else: + video_np = np.zeros_like(video_np) # Handle edge case of constant values + + video_np = (video_np * 255.0).clip(0, 255).astype(np.uint8) # [0,1] -> [0,255] + + # Handle grayscale (single channel) by converting to RGB + if video_np.shape[-1] == 1: + video_np = np.repeat(video_np, 3, axis=-1) + elif video_np.shape[-1] > 3: + print(f'Warning: video has {video_np.shape[-1]} channels, taking only first 3') + video_np = video_np[:, :, :, :3] + + try: + imageio.mimsave(output_path, video_np, fps=fps, codec='libx264') + return str(output_path) + except Exception as e: + print(f"Warning: Could not save as MP4 ({e}), falling back to .pt format") + + # Fallback: save as PyTorch tensor + fallback_path = Path(OUTPUT_DIR) / f"{filename}_{timestamp}.pt" + torch.save(video_tensor.cpu(), fallback_path) + return str(fallback_path) def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): @@ -76,11 +137,11 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): # Load model and create sampler if verbose: print("Loading model...") - model = load_model(device=DEVICE, verbose=verbose) + encoder, decoder = load_models(device=DEVICE, verbose=verbose) if verbose: print("Creating sampler...") - sampler = create_sampler(SAMPLER_TYPE, model) + sampler = create_sampler(SAMPLER_TYPE, encoder, decoder) # Generate actions if verbose: @@ -90,7 +151,7 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): # Synthesize video if verbose: print("Synthesizing video...") - video = synthesize_video(mouse, buttons, model, sampler) + video = synthesize_video(mouse, buttons, encoder, decoder, sampler) # Save video if verbose: @@ -109,9 +170,7 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): print("🎮 Simple OWL-WMS Video Renderer") # Render with default settings - render_video() + render_video(verbose=True) # Render with different pattern - render_video(ActionPattern.AIM_AND_SHOOT, length=64) - - + render_video(ActionPattern.AIM_AND_SHOOT, verbose=True) diff --git a/webapp/samplers.py b/webapp/samplers.py index 7716421f..ca603a1f 100644 --- a/webapp/samplers.py +++ b/webapp/samplers.py @@ -4,12 +4,10 @@ from functools import partial, cache from multimethod import multimethod from owl_wms.sampling.cfg import CFGSampler, InpaintCFGSampler # TODO is there 'window' sampler? -from webapp.models import load_model +from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn -SAMPLING_STEPS = 64 -BATCH_SHAPE = torch.empty(1, 128, 16, 256, 256) # TODO fix this shit -DECODE_FN = None -SCALE = 1 +SAMPLING_STEPS = 60 +SCALE = 2.17 CFG_SCALE = 1.3 MouseData = Tensor @@ -17,39 +15,34 @@ VideoData = Tensor @multimethod -def create_sampler(sampler_id: Literal['cfg'], model: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['cfg'], encoder: nn.Module, decoder: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: @cache # simple singleton def _sampler(): return CFGSampler() return partial( _sampler().__call__, sampling_steps=SAMPLING_STEPS, - decode_fn=DECODE_FN, + decode_fn=make_batched_decode_fn(decoder), scale=SCALE, cfg_scale=CFG_SCALE, - model=model + model=encoder ) @multimethod -def create_sampler(sampler_id: Literal['inpaint_cfg'], model: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder: nn.Module, decoder: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: @cache def _sampler(): return InpaintCFGSampler() return partial( _sampler().__call__, sampling_steps=SAMPLING_STEPS, - decode_fn=DECODE_FN, + decode_fn=make_batched_decode_fn(decoder), scale=SCALE, cfg_scale=CFG_SCALE, - model=model + model=encoder ) @multimethod -def create_sampler(sampler_id: Literal['window'], model: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['window'], encoder: nn.Module, decoder: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: raise NotImplementedError("Window sampler not implemented") - -model = load_model() -cfg_sampler = create_sampler('cfg', model) - - From 74b0b3f0241ac02eef9e50400aefa131c1fac06f Mon Sep 17 00:00:00 2001 From: Sami Date: Thu, 5 Jun 2025 11:50:34 -0400 Subject: [PATCH 03/71] window sampler, utilities --- owl_wms/sampling/cfg.py | 61 +++++++++++++++++++++++++++++++++++++++++ webapp/render.py | 5 +++- webapp/samplers.py | 23 +++++++++++----- 3 files changed, 81 insertions(+), 8 deletions(-) diff --git a/owl_wms/sampling/cfg.py b/owl_wms/sampling/cfg.py index e40db404..9f97fbb7 100644 --- a/owl_wms/sampling/cfg.py +++ b/owl_wms/sampling/cfg.py @@ -58,6 +58,67 @@ def __call__(self, model, dummy_batch, mouse, btn, sampling_steps = 64, decode_f x = decode_fn(x) return x + +def zlerp(x, alpha): + z = torch.randn_like(x) + return x * (1. - alpha) + z * alpha + + +class WindowCFGSampler: + def __init__(self, n_steps = 20, cfg_scale = 1.3, window_length = 60, num_frames = 60, noise_prev = 0.2): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.noise_prev = noise_prev + + @torch.no_grad() + def __call__(self, model, dummy_batch, mouse, btn, decode_fn = None, scale = 1): + + x = torch.randn_like(dummy_batch) + ts = torch.ones(x.shape[0], x.shape[1], device=x.device,dtype=x.dtype) + dt = 1. / self.n_steps + + clean_history = dummy_batch.clone() + + def step_history(): + new_history = clean_history.clone()[:,-self.window_length:] # last 60 frames + b,n,c,h,w = new_history.shape + + new_history[:,:-1] = zlerp(new_history[:,1:],self.noise_prev) # pop off first frame and noise context + new_history[:,-1] = torch.randn(b,1,c,h,w) # Add noise to last + return new_history + + for _ in range(self.num_frames): + local_history = step_history() + ts_history = torch.ones(local_history.shape[0], local_history.shape[1], device=x.device,dtype=x.dtype) + ts_history[:,-1] = self.noise_prev + + for _ in range(self.n_steps): + # CFG Branches + x = local_history.clone() + ts = ts_history.clone() + cond_pred = model(x, ts, mouse, btn) + uncond_pred = model(x, ts, torch.zeros_like(mouse), torch.zeros_like(btn)) + pred = uncond_pred + self.cfg_scale * (cond_pred - uncond_pred) + + x = x - pred*dt + ts = ts - dt + + local_history[:,-1] = x[:,-1] + ts_history[:,-1] = ts[:,-1] + + # Frame is entirely cleaned now + new_frame = local_history[:,-1:] + clean_history = torch.cat([clean_history, new_frame], dim = 1) + + x = clean_history + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + return x + + if __name__ == "__main__": model = lambda x,t,m,b: x diff --git a/webapp/render.py b/webapp/render.py index 4ec63443..d2be2cf2 100644 --- a/webapp/render.py +++ b/webapp/render.py @@ -45,8 +45,9 @@ def generate_dummy_actions(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH): def synthesize_video(mouse_actions, button_actions, encoder, decoder, sampler): """Generate video using model and sampler.""" batch_size, sequence_length = mouse_actions.shape[:2] - + # Create dummy latent batch - hardcoded dimensions from config + # NOTE: Must match training scaling! Training divides latents by vae_scale=2.17 dummy_batch = torch.randn( batch_size, sequence_length, CHANNELS, int(math.sqrt(TOKENS_PER_FRAME)), # H @@ -174,3 +175,5 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): # Render with different pattern render_video(ActionPattern.AIM_AND_SHOOT, verbose=True) + render_video(ActionPattern.LOOK_AROUND, verbose=True) + render_video(ActionPattern.CIRCLE_STRAFE, verbose=True) diff --git a/webapp/samplers.py b/webapp/samplers.py index ca603a1f..43c05869 100644 --- a/webapp/samplers.py +++ b/webapp/samplers.py @@ -3,7 +3,7 @@ from typing import Literal, Callable from functools import partial, cache from multimethod import multimethod -from owl_wms.sampling.cfg import CFGSampler, InpaintCFGSampler # TODO is there 'window' sampler? +from owl_wms.sampling.cfg import CFGSampler, InpaintCFGSampler, WindowCFGSampler from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn SAMPLING_STEPS = 60 @@ -15,34 +15,43 @@ VideoData = Tensor @multimethod -def create_sampler(sampler_id: Literal['cfg'], encoder: nn.Module, decoder: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['cfg'], encoder: nn.Module, decoder: nn.Module, batch_size: int = 8) -> Callable[[MouseData, ButtonData], VideoData]: @cache # simple singleton def _sampler(): return CFGSampler() return partial( _sampler().__call__, sampling_steps=SAMPLING_STEPS, - decode_fn=make_batched_decode_fn(decoder), + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), scale=SCALE, cfg_scale=CFG_SCALE, model=encoder ) @multimethod -def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder: nn.Module, decoder: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder: nn.Module, decoder: nn.Module, batch_size: int = 8) -> Callable[[MouseData, ButtonData], VideoData]: @cache def _sampler(): return InpaintCFGSampler() return partial( _sampler().__call__, sampling_steps=SAMPLING_STEPS, - decode_fn=make_batched_decode_fn(decoder), + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), scale=SCALE, cfg_scale=CFG_SCALE, model=encoder ) @multimethod -def create_sampler(sampler_id: Literal['window'], encoder: nn.Module, decoder: nn.Module) -> Callable[[MouseData, ButtonData], VideoData]: - raise NotImplementedError("Window sampler not implemented") +def create_sampler(sampler_id: Literal['window'], encoder: nn.Module, decoder: nn.Module, batch_size: int = 8) -> Callable[[MouseData, ButtonData], VideoData]: + @cache + def _sampler(): return WindowCFGSampler() + return partial( + _sampler().__call__, + sampling_steps=SAMPLING_STEPS, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + scale=SCALE, + cfg_scale=CFG_SCALE, + model=encoder + ) From 69dd1de1476e1655841c4a33d7a390d9281e166a Mon Sep 17 00:00:00 2001 From: Sami Date: Thu, 5 Jun 2025 14:10:06 -0400 Subject: [PATCH 04/71] webapp backend --- owl_wms/configs.py | 11 +- requirements.txt | 1 + webapp/action_converter.py | 122 ++++++++++++++++++++ webapp/demo.py | 231 ------------------------------------- webapp/demo_streaming.py | 53 +++++++++ webapp/models.py | 13 ++- webapp/render.py | 2 +- webapp/streaming.py | 142 +++++++++++++++++++++++ webapp/user_session.py | 67 +++++++++++ 9 files changed, 403 insertions(+), 239 deletions(-) create mode 100644 webapp/action_converter.py delete mode 100644 webapp/demo.py create mode 100644 webapp/demo_streaming.py create mode 100644 webapp/streaming.py create mode 100644 webapp/user_session.py diff --git a/owl_wms/configs.py b/owl_wms/configs.py index e6b100ba..58e9653f 100644 --- a/owl_wms/configs.py +++ b/owl_wms/configs.py @@ -69,4 +69,13 @@ def from_yaml(cls, path): raw_cfg = yaml.safe_load(f) cfg = OmegaConf.create(raw_cfg) - return OmegaConf.structured(cls(**cfg)) \ No newline at end of file + return OmegaConf.structured(cls(**cfg)) + + +@dataclass +class WindowSamplingConfig: + n_steps : int = 20 + cfg_scale : float = 1.3 + window_length : int = 60 + num_frames : int = 60 + noise_prev : float = 0.2 diff --git a/requirements.txt b/requirements.txt index 8c53b2cc..42dd1bef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,3 +15,4 @@ imageio einops numpy imageio[ffmpeg] +accelerate \ No newline at end of file diff --git a/webapp/action_converter.py b/webapp/action_converter.py new file mode 100644 index 00000000..000e6571 --- /dev/null +++ b/webapp/action_converter.py @@ -0,0 +1,122 @@ +import torch +import asyncio +import time + +from webapp.streaming import StreamingConfig + +BUTTON_NAMES = ["W", "A", "S", "D", "LSHIFT", "SPACE", "R", "F", "E", "LMB", "RMB"] +BUTTON_INDICES = {name: idx for idx, name in enumerate(BUTTON_NAMES)} + +class ActionConverter: + """Converts WebSocket messages to model tensor format.""" + + def __init__(self, streaming_config: StreamingConfig): + self.streaming_config = streaming_config + self.device = streaming_config.device + + def websocket_to_action(self, ws_message: dict) -> tuple[torch.Tensor, torch.Tensor]: + """ + Convert WebSocket message to action tensors. + + Expected ws_message format: + { + "mouse_x": 0.1, # Mouse movement [-1, 1] + "mouse_y": -0.05, + "W": true, # Button states + "LMB": false, + # ... other buttons + } + + Returns: + mouse: [2] tensor + buttons: [n_buttons] tensor + """ + # Extract mouse movement + mouse_x = ws_message.get("mouse_x", 0.0) + mouse_y = ws_message.get("mouse_y", 0.0) + + # Clamp to valid range + mouse_x = max(min(mouse_x, self.streaming_config.mouse_range[1]), self.streaming_config.mouse_range[0]) + mouse_y = max(min(mouse_y, self.streaming_config.mouse_range[1]), self.streaming_config.mouse_range[0]) + + mouse = torch.tensor([mouse_x, mouse_y], device=self.device, dtype=torch.float32) + + # Extract button states + button_states = torch.zeros(self.streaming_config.n_buttons, device=self.device, dtype=torch.float32) + for button_name, idx in BUTTON_INDICES.items(): + if button_name in ws_message: + button_states[idx] = 1.0 if ws_message[button_name] else 0.0 + + return mouse, button_states + + def actions_to_batch(self, actions: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]: + """ + Convert list of individual actions to batch tensors. + + Args: + actions: List of (mouse, buttons) tuples + + Returns: + mouse_batch: [batch_size, sequence_length, 2] + button_batch: [batch_size, sequence_length, n_buttons] + """ + if not actions: + # Return empty batch + return ( + torch.zeros(1, 0, 2, device=self.device), + torch.zeros(1, 0, self.streaming_config.n_buttons, device=self.device) + ) + + mouse_list = [action[0] for action in actions] + button_list = [action[1] for action in actions] + + # Stack into sequences and add batch dimension + mouse_batch = torch.stack(mouse_list, dim=0).unsqueeze(0) # [1, seq_len, 2] + button_batch = torch.stack(button_list, dim=0).unsqueeze(0) # [1, seq_len, n_buttons] + + return mouse_batch, button_batch + + +class ActionCollector: + """Collects real-time actions into 8-frame batches.""" + + def __init__(self, streaming_config: StreamingConfig): + self.streaming_config = streaming_config + self.converter = ActionConverter(streaming_config) + self.action_queue = asyncio.Queue(maxsize=100) # Buffer incoming actions + self.current_batch = [] + + async def add_websocket_action(self, ws_message: dict): + """Add action from WebSocket message.""" + action = self.converter.websocket_to_action(ws_message) + await self.action_queue.put(action) + + async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: + """ + Collect actions for one batch (8 frames worth). + + Returns action tensors ready for model inference. + """ + actions = [] + batch_duration = self.streaming_config.batch_duration + start_time = time.time() + + # Collect actions for the batch duration (400ms for 8 frames) + while len(actions) < self.streaming_config.frames_per_batch: + try: + # Wait for next action, but don't wait forever + timeout = max(0.01, batch_duration - (time.time() - start_time)) + action = await asyncio.wait_for(self.action_queue.get(), timeout=timeout) + actions.append(action) + except asyncio.TimeoutError: + # If no action received, repeat the last action or use idle + if actions: + actions.append(actions[-1]) # Repeat last action + else: + # Generate idle action + idle_mouse = torch.zeros(2, device=self.streaming_config.device) + idle_buttons = torch.zeros(self.streaming_config.n_buttons, device=self.streaming_config.device) + actions.append((idle_mouse, idle_buttons)) + + # Convert to batch tensors + return self.converter.actions_to_batch(actions) diff --git a/webapp/demo.py b/webapp/demo.py deleted file mode 100644 index 0054ae54..00000000 --- a/webapp/demo.py +++ /dev/null @@ -1,231 +0,0 @@ -#!/usr/bin/env python3 -""" -Demo script for OWL-WMS video generation pipeline. - -This script demonstrates how to use the complete pipeline to generate videos -with different action patterns and configurations. -""" - -import argparse -import sys -from pathlib import Path - -# Add webapp to path -sys.path.append(str(Path(__file__).parent)) - -from webapp.render import render_video, save_video -from webapp.action_builder import ActionPattern, ActionSequenceGenerator, ActionConfig, MouseGenerator, ButtonGenerator - - -def demo_quick_renders(): - """Demonstrate quick rendering with different patterns.""" - print("🎮 Demo: Quick Renders") - print("-" * 30) - - patterns_to_try = [ - ActionPattern.LOOK_AROUND, - ActionPattern.AIM_AND_SHOOT, - ActionPattern.SPRINT_FORWARD, - ActionPattern.CIRCLE_STRAFE - ] - - for pattern in patterns_to_try: - print(f"Generating video with pattern: {pattern.value}") - result = quick_render( - pattern=pattern, - sequence_length=60, - output_dir=f"demo_output/{pattern.value}" - ) - print(f"✅ Generated: {result['video_path']}") - print(f" Render time: {result['render_time']:.2f}s") - print() - - -def demo_batch_render(): - """Demonstrate batch rendering with multiple patterns.""" - print("🎮 Demo: Batch Render") - print("-" * 30) - - patterns = [ - ActionPattern.WALK_FORWARD, - ActionPattern.STRAFE_LEFT, - ActionPattern.STRAFE_RIGHT, - ActionPattern.WALK_BACKWARD - ] - - print(f"Generating batch video with {len(patterns)} different patterns...") - result = batch_render( - patterns=patterns, - sequence_length=60, - output_dir="demo_output/batch" - ) - - print(f"✅ Generated batch: {result['video_path']}") - print(f" Video shape: {result['video_shape']}") - print(f" Render time: {result['render_time']:.2f}s") - print() - - -def demo_custom_actions(): - """Demonstrate custom action sequence creation.""" - print("🎮 Demo: Custom Action Sequences") - print("-" * 30) - - # Create a complex custom sequence - config = ActionConfig(sequence_length=120, random_seed=42) - generator = ActionSequenceGenerator(config) - - builder = generator.generate_custom_sequence() - - # Build a complex sequence: idle -> look around -> aim and move -> sprint - mouse_sequence, button_sequence = (builder - .add_mouse_segment(0, 30, MouseGenerator.idle) - .add_button_segment(0, 30, ButtonGenerator.idle) - - .add_mouse_segment(30, 60, MouseGenerator.look_around, speed=0.4, amplitude=0.6) - .add_button_segment(30, 60, ButtonGenerator.idle) - - .add_mouse_segment(60, 90, MouseGenerator.aim_tracking, target_speed=0.2) - .add_button_segment(60, 90, ButtonGenerator.hold_buttons, button_names=["W", "A"]) - - .add_mouse_segment(90, 120, MouseGenerator.look_around, speed=0.6) - .add_button_segment(90, 120, ButtonGenerator.hold_buttons, button_names=["W", "LSHIFT"]) - .build()) - - # Use these custom actions with the renderer - render_config = RenderConfig( - sequence_length=120, - action_pattern=ActionPattern.IDLE, # This will be overridden - random_seed=42, - output_dir="demo_output/custom" - ) - - renderer = VideoRenderer(render_config) - - # Manually set the actions and render - renderer.load_model() - renderer.create_sampler() - - print("Generating video with custom action sequence...") - generated_video = renderer.synthesize_video( - mouse_sequence.unsqueeze(0), # Add batch dimension - button_sequence.unsqueeze(0) - ) - - # Save the results - action_paths = renderer.save_actions( - mouse_sequence.unsqueeze(0), - button_sequence.unsqueeze(0), - "custom_demo" - ) - video_path = renderer.save_video(generated_video, "custom_demo") - - print(f"✅ Generated custom video: {video_path}") - print(f" Action files: {action_paths}") - print() - - -def demo_action_patterns(): - """Demonstrate all available action patterns.""" - print("🎮 Demo: All Action Patterns") - print("-" * 30) - - config = ActionConfig(sequence_length=60) - generator = ActionSequenceGenerator(config) - - print("Available action patterns:") - for pattern in ActionPattern: - print(f" - {pattern.value}") - - # Generate and show stats for each pattern - mouse, buttons = generator.generate_pattern(pattern) - mouse_range = f"[{mouse.min().item():.3f}, {mouse.max().item():.3f}]" - button_count = buttons.sum().item() - - print(f" Mouse range: {mouse_range}, Button presses: {button_count:.0f}") - - print() - - -def demo_with_args(): - """Demo with command line arguments.""" - parser = argparse.ArgumentParser(description="OWL-WMS Video Generation Demo") - parser.add_argument("--pattern", type=str, default="look_around", - choices=[p.value for p in ActionPattern], - help="Action pattern to use") - parser.add_argument("--length", type=int, default=60, - help="Sequence length") - parser.add_argument("--output", type=str, default="demo_output", - help="Output directory") - parser.add_argument("--device", type=str, default="auto", - help="Device to use (cuda/cpu/auto)") - parser.add_argument("--cfg-scale", type=float, default=1.3, - help="CFG scale for sampling") - parser.add_argument("--seed", type=int, default=None, - help="Random seed") - - args = parser.parse_args() - - # Find the pattern enum - pattern = None - for p in ActionPattern: - if p.value == args.pattern: - pattern = p - break - - if pattern is None: - print(f"❌ Unknown pattern: {args.pattern}") - return - - print(f"🎮 Custom Demo: {pattern.value}") - print("-" * 30) - - config = RenderConfig( - sequence_length=args.length, - action_pattern=pattern, - device=args.device, - output_dir=args.output, - cfg_scale=args.cfg_scale, - random_seed=args.seed - ) - - renderer = VideoRenderer(config) - result = renderer.render_video(filename_prefix=f"demo_{pattern.value}") - - print(f"✅ Generated: {result['video_path']}") - print(f" Shape: {result['video_shape']}") - print(f" Time: {result['render_time']:.2f}s") - - -def main(): - """Run all demos or handle command line arguments.""" - if len(sys.argv) > 1: - # Handle command line arguments - demo_with_args() - return - - print("🎬 OWL-WMS Video Generation Demo Suite") - print("=" * 50) - print() - - try: - # Run all demos - demo_action_patterns() - demo_quick_renders() - demo_batch_render() - demo_custom_actions() - - print("🎉 All demos completed successfully!") - print("📁 Check the demo_output/ directory for generated videos.") - print() - print("💡 Tip: Run with command line arguments for custom generation:") - print(" python demo.py --pattern aim_and_shoot --length 128 --seed 42") - - except Exception as e: - print(f"❌ Demo failed: {e}") - import traceback - traceback.print_exc() - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/webapp/demo_streaming.py b/webapp/demo_streaming.py new file mode 100644 index 00000000..d8715b64 --- /dev/null +++ b/webapp/demo_streaming.py @@ -0,0 +1,53 @@ +import asyncio +import termcolor +import torch +import json + +from webapp.streaming import StreamingConfig, StreamingFrameGenerator +from owl_wms.configs import TrainingConfig, WindowSamplingConfig, TransformerConfig as ModelConfig +from webapp.models import load_models +from webapp.user_session import UserGameSession + +CFG_PATH = 'checkpoints/wm/dcae_hf_cod/basic.yml' + +async def demo_streaming_session(): + """Demo the streaming system.""" + streaming_config = StreamingConfig( + fps=20, + frames_per_batch=8, + device='cuda' if torch.cuda.is_available() else 'cpu' + ) + sampling_config = WindowSamplingConfig() + encoder, decoder, train_run_config = load_models(device='cuda' if torch.cuda.is_available() else 'cpu', verbose=True) + model_config: ModelConfig = train_run_config.model + training_config: TrainingConfig = train_run_config.train + + frame_generator = StreamingFrameGenerator(encoder, decoder, + streaming_config, model_config, training_config, sampling_config) + session = UserGameSession(frame_generator) + + class MockWebSocket: + async def __aiter__(self): + # Generate dummy actions + while True: + action = { + "mouse_x": 0.1, + "mouse_y": 0.0, + "W": True, + "LMB": False + } + yield json.dumps(action) + await asyncio.sleep(0.05) # 20 FPS + + async def send(self, data): + print(f"📤 Sent: {data}") + + mock_ws = MockWebSocket() + await session.run_session(mock_ws) + + +if __name__ == "__main__": + print(termcolor.colored("OWL-WMS Real-time Streaming System", "green")) + print(termcolor.colored("=" * 50, "green")) + + asyncio.run(demo_streaming_session()) diff --git a/webapp/models.py b/webapp/models.py index 7c220d6e..0cd08652 100644 --- a/webapp/models.py +++ b/webapp/models.py @@ -111,10 +111,10 @@ def load_decoder(self, def load_models(config_path: Optional[str] = None, - checkpoint_path: Optional[str] = None, - device: Optional[Union[str, torch.device]] = None, - eval_mode: bool = True, - verbose: bool = True) -> tuple[nn.Module, nn.Module]: + checkpoint_path: Optional[str] = None, + device: Optional[Union[str, torch.device]] = None, + eval_mode: bool = True, + verbose: bool = True) -> tuple[nn.Module, nn.Module, Config]: """ Convenience function for loading models with custom paths. @@ -141,7 +141,7 @@ def load_models(config_path: Optional[str] = None, encoder = loader.load_model(device=device, eval_mode=eval_mode, verbose=verbose) decoder = loader.load_decoder(device=device, eval_mode=eval_mode, verbose=verbose) - return encoder, decoder + return encoder, decoder, loader._load_config() if __name__ == "__main__": @@ -155,4 +155,5 @@ def load_models(config_path: Optional[str] = None, # Method 2: Using the convenience function print("=== Loading model using convenience function ===") - encoder, decoder = load_models() \ No newline at end of file + encoder, decoder, config = load_models() + print(f"Config: {config}") diff --git a/webapp/render.py b/webapp/render.py index d2be2cf2..1705ed95 100644 --- a/webapp/render.py +++ b/webapp/render.py @@ -138,7 +138,7 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): # Load model and create sampler if verbose: print("Loading model...") - encoder, decoder = load_models(device=DEVICE, verbose=verbose) + encoder, decoder, model_config = load_models(device=DEVICE, verbose=verbose) if verbose: print("Creating sampler...") diff --git a/webapp/streaming.py b/webapp/streaming.py new file mode 100644 index 00000000..2a7c94fc --- /dev/null +++ b/webapp/streaming.py @@ -0,0 +1,142 @@ +import math +import time +import torch +import asyncio +from dataclasses import dataclass +from torch import nn + +from owl_wms.sampling.cfg import WindowCFGSampler +from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn +from owl_wms.configs import TrainingConfig, WindowSamplingConfig, TransformerConfig as ModelConfig + + +@dataclass +class StreamingConfig: + fps: int = 20 + frames_per_batch: int = 8 + window_length: int = 60 + device: str = 'cuda' + n_buttons: int = 11 + mouse_range: tuple[float, float] = (-1.0, 1.0) + + @property + def frame_interval(self) -> float: + return 1.0 / self.fps + + @property + def batch_duration(self) -> float: + return self.frames_per_batch / self.fps + +class FrameBuffer: + """Manages frame streaming at precise timing.""" + + def __init__(self, streaming_config: StreamingConfig): + self.streaming_config = streaming_config + self.frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches + self.last_frame_time = 0.0 + + async def add_frame_batch(self, frame_batch: torch.Tensor): + """Add a batch of frames to the streaming queue.""" + # frame_batch shape: [1, frames_per_batch, channels, height, width] + batch_size, num_frames = frame_batch.shape[:2] + + for i in range(num_frames): + frame = frame_batch[0, i] # take first batch cause 1 user only, then [channels, height, width] + await self.frame_queue.put(frame) + + async def get_next_frame(self) -> torch.Tensor: + """Get next frame for streaming at precise timing.""" + # Calculate when to release the next frame + now = time.time() + time_since_last = now - self.last_frame_time + time_to_wait = max(0, self.streaming_config.frame_interval - time_since_last) + + if time_to_wait > 0: + await asyncio.sleep(time_to_wait) + + frame = await self.frame_queue.get() + self.last_frame_time = time.time() + return frame + +class StreamingFrameGenerator: + """Wraps WindowCFGSampler for real-time 8-frame batch generation.""" + + def __init__(self, encoder: nn.Module, decoder: nn.Module, + streaming_config: StreamingConfig, + model_config: ModelConfig, + train_config: TrainingConfig, + sampling_config: WindowSamplingConfig): + self.streaming_config = streaming_config + self.model_config = model_config + self.train_config = train_config + self.sampling_config = sampling_config + + self.encoder = encoder + self.decoder = decoder + + # Create WindowCFGSampler for 8-frame generation + self.window_sampler = WindowCFGSampler( + window_length=self.streaming_config.window_length, + num_frames=self.streaming_config.frames_per_batch, # 8 frames per batch + noise_prev=self.sampling_config.noise_prev, + cfg_scale=self.sampling_config.cfg_scale + ) + + # Create batched decode function + self.decode_fn = make_batched_decode_fn(decoder, batch_size=8) + + # Initialize frame history (60 frames of dummy data) + self.frame_history = self._initialize_frame_history() + + def _initialize_frame_history(self) -> torch.Tensor: + """Initialize with dummy frame history for cold start.""" + # Generate random latent frames matching model's expected input + # Shape: [1, window_length, channels, height, width] + tokens_h = tokens_w = int(math.sqrt(self.model_config.tokens_per_frame)) + dummy_frames = torch.randn( + 1, self.streaming_config.window_length, + self.model_config.channels, tokens_h, tokens_w, + device=self.streaming_config.device, dtype=torch.float32 + ) + return dummy_frames + + async def generate_frame_batch(self, mouse_batch: torch.Tensor, button_batch: torch.Tensor) -> torch.Tensor: + """ + Generate 8 frames using WindowCFGSampler. + + Args: + mouse_batch: [1, 8, 2] + button_batch: [1, 8, n_buttons] + + Returns: + frame_batch: [1, 8, 3, 256, 256] - decoded RGB frames + """ + # Use current frame history as dummy_batch for the window sampler + dummy_batch = self.frame_history # [1, 60, channels, h, w] + + # Generate new frames + with torch.no_grad(): + new_frames = self.window_sampler( + model=self.encoder, + dummy_batch=dummy_batch, + mouse=mouse_batch, + btn=button_batch, + decode_fn=self.decode_fn, + scale=self.train_config.vae_scale + ) + + # Update frame history by sliding window + # Remove oldest 8 frames, add newest 8 frames + new_history = torch.cat([ + self.frame_history[:, self.streaming_config.frames_per_batch:], # Remove first 8 + new_frames # Add new 8 frames + ], dim=1) + self.frame_history = new_history + return new_frames + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + del self.encoder, self.decoder ; torch.cuda.empty_cache() + diff --git a/webapp/user_session.py b/webapp/user_session.py new file mode 100644 index 00000000..6c19861a --- /dev/null +++ b/webapp/user_session.py @@ -0,0 +1,67 @@ +import time +import json +import torch +import asyncio +import termcolor + +from webapp.action_converter import ActionCollector +from webapp.streaming import StreamingFrameGenerator, FrameBuffer + +class UserGameSession: + """Main orchestrator for single-user real-time gameplay.""" + def __init__(self, frame_generator: StreamingFrameGenerator): + self.frame_generator = frame_generator + self.action_collector = ActionCollector(frame_generator.streaming_config) + self.frame_buffer = FrameBuffer(frame_generator.streaming_config) + + async def run_session(self, websocket): + with self.frame_generator: + print(termcolor.colored(f"Starting streaming session at {self.frame_generator.streaming_config.fps} FPS", "green")) + print(termcolor.colored(f"Generating {self.frame_generator.streaming_config.frames_per_batch} frames per batch", "green")) + print(termcolor.colored(f"Batch duration: {self.frame_generator.streaming_config.batch_duration:.3f}s", "green")) + + async with asyncio.TaskGroup() as tg: + tg.create_task(self._action_input_loop(websocket)) + tg.create_task(self._frame_generation_loop()) + tg.create_task(self._frame_output_loop(websocket)) + + async def _action_input_loop(self, websocket): + async for message in websocket: + try: + action_data = json.loads(message) + await self.action_collector.add_websocket_action(action_data) + except Exception as e: + print(f"Error processing action: {e}") + + async def _frame_generation_loop(self): + """Generate frame batches continuously.""" + print(termcolor.colored("Frame generation loop started", "green")) + while True: + try: + # Collect 8 frames worth of actions + mouse_batch, button_batch = await self.action_collector.collect_batch() + # Generate 8 frames + frame_batch = await self.frame_generator.generate_frame_batch(mouse_batch, button_batch) + # Queue frames for streaming + await self.frame_buffer.add_frame_batch(frame_batch) + except Exception as e: + print(termcolor.colored(f"Error in frame generation: {e}", "red")) + await asyncio.sleep(0.1) # Brief pause before retry + + async def _frame_output_loop(self, websocket): + while True: + try: + frame = await self.frame_buffer.get_next_frame() + await self._send_frame_to_client(websocket, frame) + except Exception as e: + print(termcolor.colored(f"Error in frame streaming: {e}", "red")) + await asyncio.sleep(self.frame_generator.streaming_config.frame_interval) + + async def _send_frame_to_client(self, websocket, frame: torch.Tensor): + # For now, just send frame shape info + frame_info = { + "type": "frame", + "shape": list(frame.shape), + "timestamp": time.time() + } + await websocket.send(json.dumps(frame_info)) From d7b5a61d739f13655d71d135f955d9c0fe19108b Mon Sep 17 00:00:00 2001 From: Sami Date: Fri, 6 Jun 2025 03:28:06 -0400 Subject: [PATCH 05/71] webapp stuff --- .gitignore | 9 +++ configs/webapp/config.yaml | 23 ++++++ owl_wms/configs.py | 12 +-- owl_wms/sampling/cfg.py | 19 ++--- owl_wms/sampling/simple.py | 23 ++---- owl_wms/trainers/gamerft_trainer.py | 3 +- requirements.txt | 3 +- webapp/action_converter.py | 49 +++++++++---- webapp/demo_streaming.py | 53 -------------- webapp/samplers.py | 40 +++++++--- webapp/server.py | 54 ++++++++++++++ webapp/static/index.html | 30 ++++++++ webapp/static/todo | 1 + webapp/streaming.py | 106 ++++++++++----------------- webapp/user_session.py | 25 +++++-- webapp/utils/__init__.py | 0 webapp/{ => utils}/action_builder.py | 0 webapp/utils/configs.py | 56 ++++++++++++++ webapp/utils/demo_streaming.py | 71 ++++++++++++++++++ webapp/{ => utils}/models.py | 8 +- webapp/{ => utils}/render.py | 10 +-- 21 files changed, 395 insertions(+), 200 deletions(-) create mode 100644 .gitignore create mode 100644 configs/webapp/config.yaml delete mode 100644 webapp/demo_streaming.py create mode 100644 webapp/server.py create mode 100644 webapp/static/index.html create mode 100644 webapp/static/todo create mode 100644 webapp/utils/__init__.py rename webapp/{ => utils}/action_builder.py (100%) create mode 100644 webapp/utils/configs.py create mode 100644 webapp/utils/demo_streaming.py rename webapp/{ => utils}/models.py (97%) rename webapp/{ => utils}/render.py (94%) diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..86a9f629 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +checkpoints/* +venv/ +__pycache__/* +generated_videos/ +owl-vaes/ +*.pt +*.env +*.pyc +.vscode/* \ No newline at end of file diff --git a/configs/webapp/config.yaml b/configs/webapp/config.yaml new file mode 100644 index 00000000..5bd9bdbf --- /dev/null +++ b/configs/webapp/config.yaml @@ -0,0 +1,23 @@ +model_checkpoint_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/ckpt_165k_ema.pt" +model_config_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/basic.yml" +device: "cuda" + +stream_config: + fps: 20 + frames_per_batch: 8 + window_length: 60 + device: "cuda" + n_buttons: 11 + mouse_range: (-1.0, 1.0) + +sampling_config: + sampling_steps: 20 + cfg_scale: 1.3 + num_frames: 60 + noise_prev: 0.2 + window_length: 60 + +model_config: null # loaded at runtime from model_config_path +run_config: null # loaded at runtime from model_config_path + + diff --git a/owl_wms/configs.py b/owl_wms/configs.py index 58e9653f..5c5fba61 100644 --- a/owl_wms/configs.py +++ b/owl_wms/configs.py @@ -1,7 +1,7 @@ -from dataclasses import dataclass, field -from typing import List, Optional import yaml from omegaconf import OmegaConf +from dataclasses import dataclass + @dataclass class TransformerConfig: @@ -71,11 +71,3 @@ def from_yaml(cls, path): cfg = OmegaConf.create(raw_cfg) return OmegaConf.structured(cls(**cfg)) - -@dataclass -class WindowSamplingConfig: - n_steps : int = 20 - cfg_scale : float = 1.3 - window_length : int = 60 - num_frames : int = 60 - noise_prev : float = 0.2 diff --git a/owl_wms/sampling/cfg.py b/owl_wms/sampling/cfg.py index 9f97fbb7..f14096fb 100644 --- a/owl_wms/sampling/cfg.py +++ b/owl_wms/sampling/cfg.py @@ -65,8 +65,8 @@ def zlerp(x, alpha): class WindowCFGSampler: - def __init__(self, n_steps = 20, cfg_scale = 1.3, window_length = 60, num_frames = 60, noise_prev = 0.2): - self.n_steps = n_steps + def __init__(self, sampling_steps = 20, cfg_scale = 1.3, window_length = 60, num_frames = 60, noise_prev = 0.2): + self.sampling_steps = sampling_steps self.cfg_scale = cfg_scale self.window_length = window_length self.num_frames = num_frames @@ -77,7 +77,7 @@ def __call__(self, model, dummy_batch, mouse, btn, decode_fn = None, scale = 1): x = torch.randn_like(dummy_batch) ts = torch.ones(x.shape[0], x.shape[1], device=x.device,dtype=x.dtype) - dt = 1. / self.n_steps + dt = 1. / self.sampling_steps clean_history = dummy_batch.clone() @@ -94,7 +94,7 @@ def step_history(): ts_history = torch.ones(local_history.shape[0], local_history.shape[1], device=x.device,dtype=x.dtype) ts_history[:,-1] = self.noise_prev - for _ in range(self.n_steps): + for _ in range(self.sampling_steps): # CFG Branches x = local_history.clone() ts = ts_history.clone() @@ -113,16 +113,17 @@ def step_history(): clean_history = torch.cat([clean_history, new_frame], dim = 1) x = clean_history - if decode_fn is not None: - x = x * scale - x = decode_fn(x) - return x + pixels = None + + if decode_fn is not None: + pixels = decode_fn(x * scale) + return x, pixels if __name__ == "__main__": model = lambda x,t,m,b: x sampler = CFGSampler() - x = sampler(model, torch.randn(4, 128, 16, 128), + x, pixels = sampler(model, torch.randn(4, 128, 16, 128), torch.randn(4, 128, 2), torch.randn(4, 128, 11)) print(x.shape) \ No newline at end of file diff --git a/owl_wms/sampling/simple.py b/owl_wms/sampling/simple.py index 2dfb3907..3d5fb1d3 100644 --- a/owl_wms/sampling/simple.py +++ b/owl_wms/sampling/simple.py @@ -15,9 +15,10 @@ def __call__(self, model, dummy_batch, mouse, btn, sampling_steps = 64, decode_f ts = ts - dt if decode_fn is not None: - x = x * scale - x = decode_fn(x) - return x + pixels = decode_fn(x * scale) + else: + pixels = None + return x, pixels class InpaintSimpleSampler: @torch.no_grad() @@ -38,15 +39,7 @@ def __call__(self, model, dummy_batch, mouse, btn, sampling_steps = 64, decode_f ts[:, mid:] = ts[:, mid:] - dt if decode_fn is not None: - x = x * scale - x = decode_fn(x) - return x - - -if __name__ == "__main__": - model = lambda x,t,m,b: x - - sampler = Sampler() - x = sampler(model, torch.randn(4, 3, 64, 64), - torch.randn(4, 2), torch.randn(4, 8)) - print(x.shape) \ No newline at end of file + pixels = decode_fn(x * scale) + else: + pixels = None + return x, pixels diff --git a/owl_wms/trainers/gamerft_trainer.py b/owl_wms/trainers/gamerft_trainer.py index 2f685dd5..4918f133 100644 --- a/owl_wms/trainers/gamerft_trainer.py +++ b/owl_wms/trainers/gamerft_trainer.py @@ -178,7 +178,7 @@ def get_ema_core(): if self.total_step_counter % self.train_cfg.sample_interval == 0: with ctx, torch.no_grad(): n_samples = self.train_cfg.n_samples - samples = sampler( + latents, samples = sampler( get_ema_core(), batch_vid[:n_samples], batch_mouse[:n_samples], @@ -188,7 +188,6 @@ def get_ema_core(): scale=self.train_cfg.vae_scale ) # -> [b,n,c,h,w] wandb_dict['samples'] = to_wandb(samples, batch_mouse[:n_samples], batch_btn[:n_samples]) - if self.rank == 0: wandb.log(wandb_dict) diff --git a/requirements.txt b/requirements.txt index 42dd1bef..8b827cb3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +15,5 @@ imageio einops numpy imageio[ffmpeg] -accelerate \ No newline at end of file +accelerate +fastapi \ No newline at end of file diff --git a/webapp/action_converter.py b/webapp/action_converter.py index 000e6571..a1d4662d 100644 --- a/webapp/action_converter.py +++ b/webapp/action_converter.py @@ -93,30 +93,51 @@ async def add_websocket_action(self, ws_message: dict): async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: """ - Collect actions for one batch (8 frames worth). + Collect real actions, extend to model's expected window_length. - Returns action tensors ready for model inference. + Returns: + mouse_batch: [1, window_length, 2] + button_batch: [1, window_length, n_buttons] """ - actions = [] + # Collect real actions for frames_per_batch (8 frames) + real_actions = [] batch_duration = self.streaming_config.batch_duration start_time = time.time() - # Collect actions for the batch duration (400ms for 8 frames) - while len(actions) < self.streaming_config.frames_per_batch: + while len(real_actions) < self.streaming_config.frames_per_batch: try: - # Wait for next action, but don't wait forever timeout = max(0.01, batch_duration - (time.time() - start_time)) action = await asyncio.wait_for(self.action_queue.get(), timeout=timeout) - actions.append(action) + real_actions.append(action) except asyncio.TimeoutError: - # If no action received, repeat the last action or use idle - if actions: - actions.append(actions[-1]) # Repeat last action + # Fill with idle or repeated actions + if real_actions: + real_actions.append(real_actions[-1]) else: - # Generate idle action idle_mouse = torch.zeros(2, device=self.streaming_config.device) idle_buttons = torch.zeros(self.streaming_config.n_buttons, device=self.streaming_config.device) - actions.append((idle_mouse, idle_buttons)) + real_actions.append((idle_mouse, idle_buttons)) + + # Convert 8 real actions to batch tensors + mouse, button = self.converter.actions_to_batch(real_actions) + + # Extend to window_length for model compatibility + window_length = self.streaming_config.window_length + mouse_full = self._extend_to_window_length(mouse, window_length) + button_full = self._extend_to_window_length(button, window_length) + + return mouse_full, button_full + + def _extend_to_window_length(self, tensor_batch: torch.Tensor, target_length: int) -> torch.Tensor: + """Extend [1, frames_per_batch, features] to [1, window_length, features].""" + current_length = tensor_batch.shape[1] + + if current_length >= target_length: + return tensor_batch[:, :target_length, :] # Truncate if somehow longer + + # Repeat last action to fill remaining frames + last_action = tensor_batch[:, -1:, :] # [1, 1, features] + missing_frames = target_length - current_length + repeated = last_action.repeat(1, missing_frames, 1) # [1, missing_frames, features] - # Convert to batch tensors - return self.converter.actions_to_batch(actions) + return torch.cat([tensor_batch, repeated], dim=1) # [1, target_length, features] diff --git a/webapp/demo_streaming.py b/webapp/demo_streaming.py deleted file mode 100644 index d8715b64..00000000 --- a/webapp/demo_streaming.py +++ /dev/null @@ -1,53 +0,0 @@ -import asyncio -import termcolor -import torch -import json - -from webapp.streaming import StreamingConfig, StreamingFrameGenerator -from owl_wms.configs import TrainingConfig, WindowSamplingConfig, TransformerConfig as ModelConfig -from webapp.models import load_models -from webapp.user_session import UserGameSession - -CFG_PATH = 'checkpoints/wm/dcae_hf_cod/basic.yml' - -async def demo_streaming_session(): - """Demo the streaming system.""" - streaming_config = StreamingConfig( - fps=20, - frames_per_batch=8, - device='cuda' if torch.cuda.is_available() else 'cpu' - ) - sampling_config = WindowSamplingConfig() - encoder, decoder, train_run_config = load_models(device='cuda' if torch.cuda.is_available() else 'cpu', verbose=True) - model_config: ModelConfig = train_run_config.model - training_config: TrainingConfig = train_run_config.train - - frame_generator = StreamingFrameGenerator(encoder, decoder, - streaming_config, model_config, training_config, sampling_config) - session = UserGameSession(frame_generator) - - class MockWebSocket: - async def __aiter__(self): - # Generate dummy actions - while True: - action = { - "mouse_x": 0.1, - "mouse_y": 0.0, - "W": True, - "LMB": False - } - yield json.dumps(action) - await asyncio.sleep(0.05) # 20 FPS - - async def send(self, data): - print(f"📤 Sent: {data}") - - mock_ws = MockWebSocket() - await session.run_session(mock_ws) - - -if __name__ == "__main__": - print(termcolor.colored("OWL-WMS Real-time Streaming System", "green")) - print(termcolor.colored("=" * 50, "green")) - - asyncio.run(demo_streaming_session()) diff --git a/webapp/samplers.py b/webapp/samplers.py index 43c05869..e835f90a 100644 --- a/webapp/samplers.py +++ b/webapp/samplers.py @@ -13,45 +13,61 @@ MouseData = Tensor ButtonData = Tensor VideoData = Tensor +LatentData = Tensor @multimethod -def create_sampler(sampler_id: Literal['cfg'], encoder: nn.Module, decoder: nn.Module, batch_size: int = 8) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['cfg'], encoder: nn.Module, decoder: nn.Module, + batch_size: int = 8, + sampling_steps: int = SAMPLING_STEPS, + cfg_scale: float = CFG_SCALE, + scale: float = SCALE) -> Callable[[MouseData, ButtonData], + tuple[LatentData, VideoData]]: @cache # simple singleton def _sampler(): return CFGSampler() return partial( _sampler().__call__, - sampling_steps=SAMPLING_STEPS, + sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=SCALE, - cfg_scale=CFG_SCALE, + scale=scale, + cfg_scale=cfg_scale, model=encoder ) @multimethod -def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder: nn.Module, decoder: nn.Module, batch_size: int = 8) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder: nn.Module, decoder: nn.Module, + batch_size: int = 8, + sampling_steps: int = SAMPLING_STEPS, + cfg_scale: float = CFG_SCALE, + scale: float = SCALE) -> Callable[[MouseData, ButtonData], + tuple[LatentData, VideoData]]: @cache def _sampler(): return InpaintCFGSampler() return partial( _sampler().__call__, - sampling_steps=SAMPLING_STEPS, + sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=SCALE, - cfg_scale=CFG_SCALE, + scale=scale, + cfg_scale=cfg_scale, model=encoder ) @multimethod -def create_sampler(sampler_id: Literal['window'], encoder: nn.Module, decoder: nn.Module, batch_size: int = 8) -> Callable[[MouseData, ButtonData], VideoData]: +def create_sampler(sampler_id: Literal['window'], encoder: nn.Module, decoder: nn.Module, + batch_size: int = 8, + sampling_steps: int = SAMPLING_STEPS, + cfg_scale: float = CFG_SCALE, + scale: float = SCALE) -> Callable[[MouseData, ButtonData], + tuple[LatentData, VideoData]]: @cache def _sampler(): return WindowCFGSampler() return partial( _sampler().__call__, - sampling_steps=SAMPLING_STEPS, + sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=SCALE, - cfg_scale=CFG_SCALE, + scale=scale, + cfg_scale=cfg_scale, model=encoder ) diff --git a/webapp/server.py b/webapp/server.py new file mode 100644 index 00000000..66ba24fc --- /dev/null +++ b/webapp/server.py @@ -0,0 +1,54 @@ +import os +from torch import nn + +from contextlib import asynccontextmanager +from fastapi import FastAPI, WebSocket +from fastapi.staticfiles import StaticFiles + +from webapp.utils.models import load_models +from webapp.streaming import StreamingFrameGenerator +from webapp.user_session import UserGameSession +from webapp.utils.configs import WebappConfig + + +app = FastAPI() +app.mount("/static", StaticFiles(directory="static"), name="static") + + +# -- lifespan +encoder: nn.Module = None +decoder: nn.Module = None +config: WebappConfig = None +webapp_config_path = "../configs/webapp/config.yaml" ; assert os.path.exists(webapp_config_path) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + global encoder, decoder, config + with open(webapp_config_path, "r") as f: + config: WebappConfig = WebappConfig.from_yaml(f) + + encoder, decoder, _ = load_models( + checkpoint_path=config.model_checkpoint_path, + config_path=config.model_config_path, + device=config.device, verbose=True, + ) + + yield + encoder, decoder, config = None, None, None + + +@app.websocket("/ws/game") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + + # Create streaming session for this user + frame_generator = StreamingFrameGenerator(encoder, decoder, + streaming_config=config.stream_config, + model_config=config.model_config, + train_config=config.run_config, + sampling_config=config.sampling_config) + session = UserGameSession(frame_generator) + + # Run the session (your existing code!) + await session.run_session(websocket) diff --git a/webapp/static/index.html b/webapp/static/index.html new file mode 100644 index 00000000..c795566c --- /dev/null +++ b/webapp/static/index.html @@ -0,0 +1,30 @@ + + +Open-World Labs + + + + + \ No newline at end of file diff --git a/webapp/static/todo b/webapp/static/todo new file mode 100644 index 00000000..d142e548 --- /dev/null +++ b/webapp/static/todo @@ -0,0 +1 @@ +heavens help me its time for webdev \ No newline at end of file diff --git a/webapp/streaming.py b/webapp/streaming.py index 2a7c94fc..5127bc7d 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -2,30 +2,12 @@ import time import torch import asyncio -from dataclasses import dataclass from torch import nn +from dataclasses import dataclass -from owl_wms.sampling.cfg import WindowCFGSampler -from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn -from owl_wms.configs import TrainingConfig, WindowSamplingConfig, TransformerConfig as ModelConfig - - -@dataclass -class StreamingConfig: - fps: int = 20 - frames_per_batch: int = 8 - window_length: int = 60 - device: str = 'cuda' - n_buttons: int = 11 - mouse_range: tuple[float, float] = (-1.0, 1.0) - - @property - def frame_interval(self) -> float: - return 1.0 / self.fps - - @property - def batch_duration(self) -> float: - return self.frames_per_batch / self.fps +from webapp.samplers import create_sampler +from webapp.utils.configs import SamplingConfig, StreamingConfig +from owl_wms.configs import TrainingConfig, TransformerConfig as ModelConfig class FrameBuffer: """Manages frame streaming at precise timing.""" @@ -65,7 +47,7 @@ def __init__(self, encoder: nn.Module, decoder: nn.Module, streaming_config: StreamingConfig, model_config: ModelConfig, train_config: TrainingConfig, - sampling_config: WindowSamplingConfig): + sampling_config: SamplingConfig): self.streaming_config = streaming_config self.model_config = model_config self.train_config = train_config @@ -75,68 +57,58 @@ def __init__(self, encoder: nn.Module, decoder: nn.Module, self.decoder = decoder # Create WindowCFGSampler for 8-frame generation - self.window_sampler = WindowCFGSampler( - window_length=self.streaming_config.window_length, - num_frames=self.streaming_config.frames_per_batch, # 8 frames per batch - noise_prev=self.sampling_config.noise_prev, - cfg_scale=self.sampling_config.cfg_scale - ) - - # Create batched decode function - self.decode_fn = make_batched_decode_fn(decoder, batch_size=8) + self.sample_window_fn = create_sampler('window', encoder, decoder, + batch_size=1, + sampling_steps=self.sampling_config.sampling_steps, + cfg_scale=self.sampling_config.cfg_scale, + scale=self.train_config.vae_scale) + # Initialize frame history as empty tensor + self.frame_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) + + def add_to_history(self, frame_batch: torch.Tensor): + if self.frame_history.equal(torch.tensor([], device=self.streaming_config.device)): + self.frame_history = frame_batch + return + + self.frame_history = torch.cat([self.frame_history, frame_batch], dim=1) - # Initialize frame history (60 frames of dummy data) - self.frame_history = self._initialize_frame_history() + # cap this at around 60 frames + if self.frame_history.shape[1] > self.streaming_config.window_length: + self.frame_history = self.frame_history[:, -self.streaming_config.window_length:] + - def _initialize_frame_history(self) -> torch.Tensor: - """Initialize with dummy frame history for cold start.""" - # Generate random latent frames matching model's expected input - # Shape: [1, window_length, channels, height, width] + @property + def dummy_batch(self) -> torch.Tensor: + """Generate dummy autoencoder latents for cold start.""" tokens_h = tokens_w = int(math.sqrt(self.model_config.tokens_per_frame)) dummy_frames = torch.randn( 1, self.streaming_config.window_length, self.model_config.channels, tokens_h, tokens_w, - device=self.streaming_config.device, dtype=torch.float32 - ) + device=self.streaming_config.device, dtype=torch.bfloat16) return dummy_frames + async def generate_frame_batch(self, mouse_batch: torch.Tensor, button_batch: torch.Tensor) -> torch.Tensor: """ - Generate 8 frames using WindowCFGSampler. + Generate window_length frames, return first frames_per_batch for streaming. Args: - mouse_batch: [1, 8, 2] - button_batch: [1, 8, n_buttons] + mouse_batch: [1, window_length, 2] + button_batch: [1, window_length, n_buttons] Returns: - frame_batch: [1, 8, 3, 256, 256] - decoded RGB frames + frame_batch: [1, frames_per_batch, 3, 256, 256] - only streaming frames """ - # Use current frame history as dummy_batch for the window sampler - dummy_batch = self.frame_history # [1, 60, channels, h, w] - - # Generate new frames - with torch.no_grad(): - new_frames = self.window_sampler( - model=self.encoder, - dummy_batch=dummy_batch, - mouse=mouse_batch, - btn=button_batch, - decode_fn=self.decode_fn, - scale=self.train_config.vae_scale - ) - - # Update frame history by sliding window - # Remove oldest 8 frames, add newest 8 frames - new_history = torch.cat([ - self.frame_history[:, self.streaming_config.frames_per_batch:], # Remove first 8 - new_frames # Add new 8 frames - ], dim=1) - self.frame_history = new_history - return new_frames + with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): + latents, full_frames = self.sample_window_fn(dummy_batch=self.dummy_batch, mouse=mouse_batch, btn=button_batch) # [1, window_length, 3, 256, 256] + self.add_to_history(latents[0, :self.streaming_config.frames_per_batch, :, :, :]) + + # Take only first frames_per_batch for streaming + streaming_frames = full_frames[:, :self.streaming_config.frames_per_batch, :, :, :] + return streaming_frames # [1, frames_per_batch, 3, 256, 256] def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): del self.encoder, self.decoder ; torch.cuda.empty_cache() - diff --git a/webapp/user_session.py b/webapp/user_session.py index 6c19861a..f09a7b65 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -1,9 +1,13 @@ +import cv2 import time import json import torch +import base64 import asyncio import termcolor +import numpy as np +from fastapi import WebSocket from webapp.action_converter import ActionCollector from webapp.streaming import StreamingFrameGenerator, FrameBuffer @@ -45,7 +49,8 @@ async def _frame_generation_loop(self): # Queue frames for streaming await self.frame_buffer.add_frame_batch(frame_batch) except Exception as e: - print(termcolor.colored(f"Error in frame generation: {e}", "red")) + import traceback + print(termcolor.colored(f"Error in frame generation: {e} :\n {traceback.format_exc()}", "red")) await asyncio.sleep(0.1) # Brief pause before retry async def _frame_output_loop(self, websocket): @@ -57,11 +62,17 @@ async def _frame_output_loop(self, websocket): print(termcolor.colored(f"Error in frame streaming: {e}", "red")) await asyncio.sleep(self.frame_generator.streaming_config.frame_interval) - async def _send_frame_to_client(self, websocket, frame: torch.Tensor): - # For now, just send frame shape info - frame_info = { + async def _send_frame_to_client(self, websocket: WebSocket, frame: torch.Tensor): + # Convert frame to base64 JPEG + print(f"Frame shape: {frame.shape} - with stats: {frame.min()=}, {frame.max()=}, {frame.mean()=}, {frame.std()=}") + frame_np = frame.cpu().numpy().transpose(1, 2, 0) # CHW -> HWC + frame_np = ((frame_np + 1) * 127.5).clip(0, 255).astype(np.uint8) # Normalize + + _, buffer = cv2.imencode('.jpg', frame_np) + frame_base64 = base64.b64encode(buffer).decode('utf-8') + + await websocket.send(json.dumps({ "type": "frame", - "shape": list(frame.shape), + "data": frame_base64, "timestamp": time.time() - } - await websocket.send(json.dumps(frame_info)) + })) diff --git a/webapp/utils/__init__.py b/webapp/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/webapp/action_builder.py b/webapp/utils/action_builder.py similarity index 100% rename from webapp/action_builder.py rename to webapp/utils/action_builder.py diff --git a/webapp/utils/configs.py b/webapp/utils/configs.py new file mode 100644 index 00000000..0e8eb554 --- /dev/null +++ b/webapp/utils/configs.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import os +import yaml +from typing import Optional +from dataclasses import dataclass + +from owl_wms.configs import Config as RunConfig, TransformerConfig as ModelConfig + +@dataclass +class WebappConfig: + model_checkpoint_path : os.PathLike + model_config_path : os.PathLike + model_config : ModelConfig + run_config : RunConfig + stream_config : StreamingConfig + sampling_config : SamplingConfig + device : str = 'cuda' + + @classmethod + def from_yaml(cls, path: os.PathLike) -> WebappConfig: + + with (open(path, 'r') as wcp, + open(config['model_config_path'], 'r') as mcp, + open(config['run_config_path'], 'r') as rcp): + config = yaml.safe_load(wcp) + config['model_config'] = yaml.safe_load(mcp) + config['run_config'] = yaml.safe_load(rcp) + + return cls(**config) + + +@dataclass +class SamplingConfig: + sampling_steps : int = 20 + cfg_scale : float = 1.3 + num_frames : int = 60 + noise_prev : float = 0.2 + window_length : Optional[int] = 60 + +@dataclass +class StreamingConfig: + fps: int = 20 + frames_per_batch: int = 8 + window_length: int = 60 + device: str = 'cuda' + n_buttons: int = 11 + mouse_range: tuple[float, float] = (-1.0, 1.0) + + @property + def frame_interval(self) -> float: + return 1.0 / self.fps + + @property + def batch_duration(self) -> float: + return self.frames_per_batch / self.fps diff --git a/webapp/utils/demo_streaming.py b/webapp/utils/demo_streaming.py new file mode 100644 index 00000000..fdf263bf --- /dev/null +++ b/webapp/utils/demo_streaming.py @@ -0,0 +1,71 @@ +import asyncio +import torch +import termcolor + +from webapp.utils.render import generate_dummy_actions, save_video +from webapp.streaming import StreamingConfig, StreamingFrameGenerator +from webapp.utils.models import load_models +from webapp.utils.action_builder import ActionPattern +from webapp.utils.configs import SamplingConfig + + +async def demo_streaming_generation(pattern=ActionPattern.LOOK_AROUND): + """Generate one batch using StreamingFrameGenerator instead of regular sampler.""" + + print(termcolor.colored("🎮 OWL-WMS Streaming Demo", "green")) + print(termcolor.colored("=" * 50, "green")) + + # Configuration + device = 'cuda' if torch.cuda.is_available() else 'cpu' + + streaming_config = StreamingConfig( + fps=20, + frames_per_batch=8, + window_length=60, + device=device + ) + + sampling_config = SamplingConfig() + + # Load models (reuse render.py's load_models) + print("📦 Loading models...") + encoder, decoder, train_config = load_models(device=device, verbose=True) + model_config = train_config.model + training_config = train_config.train + + # Create streaming frame generator + print("🎬 Creating streaming frame generator...") + frame_generator = StreamingFrameGenerator( + encoder, decoder, + streaming_config, model_config, training_config, sampling_config + ) + + # Generate actions (reuse render.py's generate_dummy_actions) + print(f"🎯 Generating {pattern.value} actions...") + mouse_batch, button_batch = generate_dummy_actions(pattern, streaming_config.window_length) + + # Generate frames using streaming generator + print("🎨 Generating frames with streaming generator...") + with frame_generator: + frame_batch = await frame_generator.generate_frame_batch(mouse_batch, button_batch) + + print(f"Generated {frame_batch.shape[1]} frames with shape: {frame_batch.shape}") + + # Save video (reuse render.py's save_video) + print("💾 Saving video...") + output_path = save_video(frame_batch, f"streaming_demo_{pattern.value.lower()}", fps=streaming_config.fps) + + print(termcolor.colored(f"🎉 Demo complete! Video: {output_path}", "green")) + return output_path + + +if __name__ == "__main__": + # Try different patterns by changing this: + pattern = ActionPattern.LOOK_AROUND # or AIM_AND_SHOOT, CIRCLE_STRAFE, etc. + + print("Available patterns:") + for p in ActionPattern: + print(f" - {p.value}") + print() + + asyncio.run(demo_streaming_generation(pattern)) \ No newline at end of file diff --git a/webapp/models.py b/webapp/utils/models.py similarity index 97% rename from webapp/models.py rename to webapp/utils/models.py index 0cd08652..4766d9d4 100644 --- a/webapp/models.py +++ b/webapp/utils/models.py @@ -7,7 +7,7 @@ from termcolor import colored from owl_wms.models import get_model_cls -from owl_wms.configs import Config +from owl_wms.configs import Config as RunConfig from owl_wms.utils.owl_vae_bridge import get_decoder_only from owl_wms.utils import freeze @@ -48,8 +48,8 @@ def _append_state_dict_prefix(state_dict: Dict, prefix: str = 'core.') -> Dict: def _count_parameters(model: nn.Module) -> int: return sum(p.numel() for p in model.parameters()) - def _load_config(self) -> Config: - return Config.from_yaml(str(self.paths.config)) + def _load_config(self) -> RunConfig: + return RunConfig.from_yaml(str(self.paths.config)) def _load_checkpoint(self) -> Dict: return torch.load(str(self.paths.checkpoint), map_location='cpu') @@ -114,7 +114,7 @@ def load_models(config_path: Optional[str] = None, checkpoint_path: Optional[str] = None, device: Optional[Union[str, torch.device]] = None, eval_mode: bool = True, - verbose: bool = True) -> tuple[nn.Module, nn.Module, Config]: + verbose: bool = True) -> tuple[nn.Module, nn.Module, RunConfig]: """ Convenience function for loading models with custom paths. diff --git a/webapp/render.py b/webapp/utils/render.py similarity index 94% rename from webapp/render.py rename to webapp/utils/render.py index 1705ed95..418423f9 100644 --- a/webapp/render.py +++ b/webapp/utils/render.py @@ -8,9 +8,9 @@ import einops as eo -from webapp.models import load_models +from webapp.utils.models import load_models from webapp.samplers import create_sampler -from webapp.action_builder import ActionSequenceGenerator, ActionConfig, ActionPattern +from webapp.utils.action_builder import ActionSequenceGenerator, ActionConfig, ActionPattern HEIGHT = 256 WIDTH = 256 @@ -20,7 +20,7 @@ TOKENS_PER_FRAME = 16 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' OUTPUT_DIR = "generated_videos" -SAMPLER_TYPE = 'cfg' +SAMPLER_TYPE = 'window' DEFAULT_PATTERN = ActionPattern.LOOK_AROUND @@ -46,8 +46,6 @@ def synthesize_video(mouse_actions, button_actions, encoder, decoder, sampler): """Generate video using model and sampler.""" batch_size, sequence_length = mouse_actions.shape[:2] - # Create dummy latent batch - hardcoded dimensions from config - # NOTE: Must match training scaling! Training divides latents by vae_scale=2.17 dummy_batch = torch.randn( batch_size, sequence_length, CHANNELS, int(math.sqrt(TOKENS_PER_FRAME)), # H @@ -61,7 +59,7 @@ def synthesize_video(mouse_actions, button_actions, encoder, decoder, sampler): # Generate video with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): - video = sampler( + latents, video = sampler( dummy_batch=dummy_batch, mouse=mouse_actions, btn=button_actions From b4645ab99725127298be388681fc53fa146d92b0 Mon Sep 17 00:00:00 2001 From: Sami Date: Fri, 6 Jun 2025 03:29:31 -0400 Subject: [PATCH 06/71] slight cleanup --- webapp/action_converter.py | 2 +- webapp/streaming.py | 1 - webapp/{ => utils}/samplers.py | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) rename webapp/{ => utils}/samplers.py (99%) diff --git a/webapp/action_converter.py b/webapp/action_converter.py index a1d4662d..e32ecdd6 100644 --- a/webapp/action_converter.py +++ b/webapp/action_converter.py @@ -1,6 +1,6 @@ +import time import torch import asyncio -import time from webapp.streaming import StreamingConfig diff --git a/webapp/streaming.py b/webapp/streaming.py index 5127bc7d..936f4f18 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -3,7 +3,6 @@ import torch import asyncio from torch import nn -from dataclasses import dataclass from webapp.samplers import create_sampler from webapp.utils.configs import SamplingConfig, StreamingConfig diff --git a/webapp/samplers.py b/webapp/utils/samplers.py similarity index 99% rename from webapp/samplers.py rename to webapp/utils/samplers.py index e835f90a..dea3caf6 100644 --- a/webapp/samplers.py +++ b/webapp/utils/samplers.py @@ -1,4 +1,3 @@ -import torch from torch import nn, Tensor from typing import Literal, Callable from functools import partial, cache From 96a6b010eec2206b169413cd410e40eaf320bc40 Mon Sep 17 00:00:00 2001 From: Sami Date: Fri, 6 Jun 2025 03:38:42 -0400 Subject: [PATCH 07/71] stray refactor victim --- webapp/utils/render.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/webapp/utils/render.py b/webapp/utils/render.py index 418423f9..f5a7deb6 100644 --- a/webapp/utils/render.py +++ b/webapp/utils/render.py @@ -9,7 +9,7 @@ import einops as eo from webapp.utils.models import load_models -from webapp.samplers import create_sampler +from webapp.utils.samplers import create_sampler from webapp.utils.action_builder import ActionSequenceGenerator, ActionConfig, ActionPattern HEIGHT = 256 From 57e491f2f293b8477ce537753f0eba1b4ee51739 Mon Sep 17 00:00:00 2001 From: Sami Date: Fri, 6 Jun 2025 03:40:19 -0400 Subject: [PATCH 08/71] stray refactor victim --- webapp/streaming.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/webapp/streaming.py b/webapp/streaming.py index 936f4f18..cb108cfc 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -4,7 +4,7 @@ import asyncio from torch import nn -from webapp.samplers import create_sampler +from webapp.utils.samplers import create_sampler from webapp.utils.configs import SamplingConfig, StreamingConfig from owl_wms.configs import TrainingConfig, TransformerConfig as ModelConfig From 7a6cc4953081a02e65944b0219963c7f3e73410a Mon Sep 17 00:00:00 2001 From: Sami Date: Sun, 8 Jun 2025 03:11:42 -0400 Subject: [PATCH 09/71] Very basic webapp --- configs/webapp/config.yaml | 8 +- owl_wms/configs.py | 15 +- requirements.txt | 3 +- webapp/action_converter.py | 112 ++++---- webapp/server.py | 52 +++- webapp/static/index.html | 222 ++++++++++++++- webapp/streaming.py | 75 +++-- webapp/user_session.py | 48 ++-- webapp/utils/action_builder.py | 4 +- webapp/utils/configs.py | 19 +- webapp/utils/demo_streaming.py | 10 +- webapp/utils/render.py | 12 +- webapp/utils/visualize_overlay_actions.py | 333 ++++++++++++++++++++++ 13 files changed, 764 insertions(+), 149 deletions(-) create mode 100644 webapp/utils/visualize_overlay_actions.py diff --git a/configs/webapp/config.yaml b/configs/webapp/config.yaml index 5bd9bdbf..aa27f5a6 100644 --- a/configs/webapp/config.yaml +++ b/configs/webapp/config.yaml @@ -1,5 +1,5 @@ model_checkpoint_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/ckpt_165k_ema.pt" -model_config_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/basic.yml" +run_config_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/basic.yml" device: "cuda" stream_config: @@ -8,7 +8,8 @@ stream_config: window_length: 60 device: "cuda" n_buttons: 11 - mouse_range: (-1.0, 1.0) + n_mouse_axes: 2 + mouse_range: [-1.0, 1.0] sampling_config: sampling_steps: 20 @@ -17,7 +18,6 @@ sampling_config: noise_prev: 0.2 window_length: 60 -model_config: null # loaded at runtime from model_config_path -run_config: null # loaded at runtime from model_config_path +run_config: null # loaded at runtime from model_config_path, and used to access model and train config diff --git a/owl_wms/configs.py b/owl_wms/configs.py index 5c5fba61..85580792 100644 --- a/owl_wms/configs.py +++ b/owl_wms/configs.py @@ -6,18 +6,23 @@ @dataclass class TransformerConfig: model_id : str = None + channels : int = 128 + sample_size : int = 16 + patch_size : int = 1 n_layers : int = 12 n_heads : int = 12 d_model : int = 384 - patch_size : int = 1 - channels : int = 128 - sample_size : int = 16 + tokens_per_frame: int = 16 + n_buttons : int = 8 + n_mouse_axes : int = 2 cfg_prob : float = 0.1 - n_buttons : int = 8 - tokens_per_frame: int = 16 + n_frames : int = 60 + + causal : bool = False + @dataclass class TrainingConfig: diff --git a/requirements.txt b/requirements.txt index 8b827cb3..85210dd6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,4 +16,5 @@ einops numpy imageio[ffmpeg] accelerate -fastapi \ No newline at end of file +fastapi +uvicorn[standard] diff --git a/webapp/action_converter.py b/webapp/action_converter.py index e32ecdd6..6b7b1cba 100644 --- a/webapp/action_converter.py +++ b/webapp/action_converter.py @@ -3,17 +3,46 @@ import asyncio from webapp.streaming import StreamingConfig +from torch.nn import functional as F BUTTON_NAMES = ["W", "A", "S", "D", "LSHIFT", "SPACE", "R", "F", "E", "LMB", "RMB"] BUTTON_INDICES = {name: idx for idx, name in enumerate(BUTTON_NAMES)} + +def _interpolate(tensor_batch: torch.Tensor, + empty_action: torch.Tensor, + target_length: int) -> torch.Tensor: + + """ + Interpolate actions to target_length. + If tensor_batch is longer than target_length, subsample. + If tensor_batch is shorter than target_length, repeat with empty actions. + + Must provide empty_action, which is the action to repeat when the batch is shorter than target_length. + Must also provide target_length, which is the length to interpolate to. + Returns: + tensor_batch: [1, target_length, features] + """ + num_actions = tensor_batch.shape[1] + if num_actions >= target_length: + # subsample actions if somehow longer than frames_per_batch + downsampled = torch.arange(0, num_actions, num_actions // target_length) + return tensor_batch[:, downsampled, :] + + # Repeat with empty actions to fill remaining frames + num_missing_actions = target_length - num_actions + repeated = empty_action.repeat(1, num_missing_actions, 1) # [1, missing_frames, features] + + return torch.cat([tensor_batch, repeated], dim=1) # [1, target_length, features] + + class ActionConverter: """Converts WebSocket messages to model tensor format.""" def __init__(self, streaming_config: StreamingConfig): self.streaming_config = streaming_config self.device = streaming_config.device - + def websocket_to_action(self, ws_message: dict) -> tuple[torch.Tensor, torch.Tensor]: """ Convert WebSocket message to action tensors. @@ -61,19 +90,15 @@ def actions_to_batch(self, actions: list[tuple[torch.Tensor, torch.Tensor]]) -> button_batch: [batch_size, sequence_length, n_buttons] """ if not actions: - # Return empty batch + # Return empty batch - [B, 0, 2] and [B, 0, n_buttons] return ( torch.zeros(1, 0, 2, device=self.device), torch.zeros(1, 0, self.streaming_config.n_buttons, device=self.device) ) - - mouse_list = [action[0] for action in actions] - button_list = [action[1] for action in actions] - - # Stack into sequences and add batch dimension - mouse_batch = torch.stack(mouse_list, dim=0).unsqueeze(0) # [1, seq_len, 2] - button_batch = torch.stack(button_list, dim=0).unsqueeze(0) # [1, seq_len, n_buttons] - + + mouse_batch = torch.stack([action[0] for action in actions], dim=0).unsqueeze(0) # [1, seq_len, 2] + button_batch = torch.stack([action[1] for action in actions], dim=0).unsqueeze(0) # [1, seq_len, n_buttons] + return mouse_batch, button_batch @@ -85,7 +110,13 @@ def __init__(self, streaming_config: StreamingConfig): self.converter = ActionConverter(streaming_config) self.action_queue = asyncio.Queue(maxsize=100) # Buffer incoming actions self.current_batch = [] - + + # -- constants + self.empty_mouse = torch.zeros((streaming_config.frames_per_batch, streaming_config.n_mouse_axes), + device=streaming_config.device, dtype=torch.float32) + self.empty_buttons = torch.zeros((streaming_config.frames_per_batch, streaming_config.n_buttons), + device=streaming_config.device, dtype=torch.bool) + async def add_websocket_action(self, ws_message: dict): """Add action from WebSocket message.""" action = self.converter.websocket_to_action(ws_message) @@ -93,51 +124,34 @@ async def add_websocket_action(self, ws_message: dict): async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: """ - Collect real actions, extend to model's expected window_length. + Collect actions from the UI at whatever rate is sent over by the client. + This collects all the frames that have been supplied between the last frame generation and now. + + It takes in as many actions as the model generates frames at once. For example, in CausVid, if + the model generates 4 frames at a time, this function will return [1, 4, 2] and [1, 4, 11] + for mouse and button actions. + If, for one reason or another, we have <4 actions, we will fill the batch with idle actions. + Returns: - mouse_batch: [1, window_length, 2] - button_batch: [1, window_length, n_buttons] + mouse_batch: [1, X, 2] + button_batch: [1, X, n_buttons] """ - # Collect real actions for frames_per_batch (8 frames) real_actions = [] - batch_duration = self.streaming_config.batch_duration start_time = time.time() - while len(real_actions) < self.streaming_config.frames_per_batch: + while start_time + self.streaming_config.batch_duration > time.time(): try: - timeout = max(0.01, batch_duration - (time.time() - start_time)) - action = await asyncio.wait_for(self.action_queue.get(), timeout=timeout) + timeout = max(0.01, self.streaming_config.batch_duration - (time.time() - start_time)) + action = await asyncio.wait_for(self.action_queue.get(), timeout=timeout) real_actions.append(action) except asyncio.TimeoutError: - # Fill with idle or repeated actions - if real_actions: - real_actions.append(real_actions[-1]) - else: - idle_mouse = torch.zeros(2, device=self.streaming_config.device) - idle_buttons = torch.zeros(self.streaming_config.n_buttons, device=self.streaming_config.device) - real_actions.append((idle_mouse, idle_buttons)) - - # Convert 8 real actions to batch tensors - mouse, button = self.converter.actions_to_batch(real_actions) - - # Extend to window_length for model compatibility - window_length = self.streaming_config.window_length - mouse_full = self._extend_to_window_length(mouse, window_length) - button_full = self._extend_to_window_length(button, window_length) - - return mouse_full, button_full + pass - def _extend_to_window_length(self, tensor_batch: torch.Tensor, target_length: int) -> torch.Tensor: - """Extend [1, frames_per_batch, features] to [1, window_length, features].""" - current_length = tensor_batch.shape[1] - - if current_length >= target_length: - return tensor_batch[:, :target_length, :] # Truncate if somehow longer - - # Repeat last action to fill remaining frames - last_action = tensor_batch[:, -1:, :] # [1, 1, features] - missing_frames = target_length - current_length - repeated = last_action.repeat(1, missing_frames, 1) # [1, missing_frames, features] - - return torch.cat([tensor_batch, repeated], dim=1) # [1, target_length, features] + # Convert 8 real actions to batch tensors + mouse, button = self.converter.actions_to_batch(real_actions) + mouse = _interpolate(mouse, empty_action=self.empty_mouse, + target_length=self.streaming_config.frames_per_batch) + button = _interpolate(button, empty_action=self.empty_buttons, + target_length=self.streaming_config.frames_per_batch) + return mouse, button diff --git a/webapp/server.py b/webapp/server.py index 66ba24fc..5e5d0666 100644 --- a/webapp/server.py +++ b/webapp/server.py @@ -1,9 +1,10 @@ import os from torch import nn -from contextlib import asynccontextmanager -from fastapi import FastAPI, WebSocket -from fastapi.staticfiles import StaticFiles +from contextlib import asynccontextmanager +from fastapi import FastAPI, WebSocket +from fastapi.staticfiles import StaticFiles +from fastapi.responses import FileResponse from webapp.utils.models import load_models from webapp.streaming import StreamingFrameGenerator @@ -11,26 +12,23 @@ from webapp.utils.configs import WebappConfig -app = FastAPI() -app.mount("/static", StaticFiles(directory="static"), name="static") - +DEBUG = True # for funsies # -- lifespan encoder: nn.Module = None decoder: nn.Module = None config: WebappConfig = None -webapp_config_path = "../configs/webapp/config.yaml" ; assert os.path.exists(webapp_config_path) +webapp_config_path = "./configs/webapp/config.yaml" ; assert os.path.exists(webapp_config_path) @asynccontextmanager async def lifespan(app: FastAPI): global encoder, decoder, config - with open(webapp_config_path, "r") as f: - config: WebappConfig = WebappConfig.from_yaml(f) + config = WebappConfig.from_yaml(webapp_config_path) encoder, decoder, _ = load_models( checkpoint_path=config.model_checkpoint_path, - config_path=config.model_config_path, + config_path=config.run_config_path, device=config.device, verbose=True, ) @@ -38,6 +36,16 @@ async def lifespan(app: FastAPI): encoder, decoder, config = None, None, None +app = FastAPI(lifespan=lifespan) +app.mount("/assets", StaticFiles(directory="webapp/static"), name="assets") + + +@app.get("/") +async def read_root(): + """Serve the main game page.""" + return FileResponse("webapp/static/index.html") + + @app.websocket("/ws/game") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() @@ -45,10 +53,28 @@ async def websocket_endpoint(websocket: WebSocket): # Create streaming session for this user frame_generator = StreamingFrameGenerator(encoder, decoder, streaming_config=config.stream_config, - model_config=config.model_config, - train_config=config.run_config, - sampling_config=config.sampling_config) + model_config=config.run_config.model, + train_config=config.run_config.train, + sampling_config=config.sampling_config, + debug=DEBUG) session = UserGameSession(frame_generator) # Run the session (your existing code!) await session.run_session(websocket) + + +if __name__ == "__main__": + import uvicorn + + print("🚀 Starting OWL-WMS FastAPI Server...") + print("📡 WebSocket endpoint: ws://localhost:8000/ws/game") + print("🌐 Access via: http://localhost:8000") + print("🔄 Auto-reload enabled for development") + + uvicorn.run( + "webapp.server:app", + host="0.0.0.0", # Allow external connections + port=8000, + reload=True, # Auto-reload on file changes + log_level="info" + ) diff --git a/webapp/static/index.html b/webapp/static/index.html index c795566c..b15e0a9c 100644 --- a/webapp/static/index.html +++ b/webapp/static/index.html @@ -1,30 +1,234 @@ -Open-World Labs + + Open-World Labs - OWL-WMS + + +

🎮 OWL-WMS Game Stream

+
Connecting to server...
+

Move your mouse over the canvas to interact with the game world!

+ \ No newline at end of file diff --git a/webapp/streaming.py b/webapp/streaming.py index cb108cfc..16d9ec40 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -4,20 +4,24 @@ import asyncio from torch import nn -from webapp.utils.samplers import create_sampler -from webapp.utils.configs import SamplingConfig, StreamingConfig -from owl_wms.configs import TrainingConfig, TransformerConfig as ModelConfig +from webapp.utils.samplers import create_sampler +from webapp.utils.configs import SamplingConfig, StreamingConfig +from owl_wms.configs import TrainingConfig, TransformerConfig as ModelConfig +from webapp.utils.visualize_overlay_actions import _draw_video as _draw_action_overlays class FrameBuffer: - """Manages frame streaming at precise timing.""" - + """ + Manages frame streaming at precise timing, to adhere to a max FPS. + We need this because sometimes we could generate frames faster than the max FPS, + but, intuitively, the frame outputs will be the same 'set time' apart in the world state that we are approximating. + """ + def __init__(self, streaming_config: StreamingConfig): self.streaming_config = streaming_config self.frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches self.last_frame_time = 0.0 async def add_frame_batch(self, frame_batch: torch.Tensor): - """Add a batch of frames to the streaming queue.""" # frame_batch shape: [1, frames_per_batch, channels, height, width] batch_size, num_frames = frame_batch.shape[:2] @@ -26,8 +30,7 @@ async def add_frame_batch(self, frame_batch: torch.Tensor): await self.frame_queue.put(frame) async def get_next_frame(self) -> torch.Tensor: - """Get next frame for streaming at precise timing.""" - # Calculate when to release the next frame + """Get next frame for streaming at capped FPS.""" now = time.time() time_since_last = now - self.last_frame_time time_to_wait = max(0, self.streaming_config.frame_interval - time_since_last) @@ -40,21 +43,23 @@ async def get_next_frame(self) -> torch.Tensor: return frame class StreamingFrameGenerator: - """Wraps WindowCFGSampler for real-time 8-frame batch generation.""" + """Wraps WindowCFGSampler to generate frames.""" def __init__(self, encoder: nn.Module, decoder: nn.Module, streaming_config: StreamingConfig, model_config: ModelConfig, train_config: TrainingConfig, - sampling_config: SamplingConfig): + sampling_config: SamplingConfig, + debug: bool = False): self.streaming_config = streaming_config self.model_config = model_config self.train_config = train_config self.sampling_config = sampling_config + self.debug = debug self.encoder = encoder self.decoder = decoder - + # # Create WindowCFGSampler for 8-frame generation self.sample_window_fn = create_sampler('window', encoder, decoder, batch_size=1, @@ -63,22 +68,28 @@ def __init__(self, encoder: nn.Module, decoder: nn.Module, scale=self.train_config.vae_scale) # Initialize frame history as empty tensor self.frame_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) + self.mouse_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) + self.button_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) - def add_to_history(self, frame_batch: torch.Tensor): + def add_to_history(self, frame_batch: torch.Tensor, mouse_batch: torch.Tensor, button_batch: torch.Tensor): if self.frame_history.equal(torch.tensor([], device=self.streaming_config.device)): - self.frame_history = frame_batch + self.frame_history = frame_batch + self.mouse_history = mouse_batch + self.button_history = button_batch return - self.frame_history = torch.cat([self.frame_history, frame_batch], dim=1) - + self.frame_history = torch.cat([self.frame_history, frame_batch], dim=1) + self.mouse_history = torch.cat([self.mouse_history, mouse_batch], dim=1) + self.button_history = torch.cat([self.button_history, button_batch], dim=1) # cap this at around 60 frames if self.frame_history.shape[1] > self.streaming_config.window_length: - self.frame_history = self.frame_history[:, -self.streaming_config.window_length:] + self.frame_history = self.frame_history[:, -self.streaming_config.window_length:] + self.mouse_history = self.mouse_history[:, -self.streaming_config.window_length:] + self.button_history = self.button_history[:, -self.streaming_config.window_length:] - @property def dummy_batch(self) -> torch.Tensor: - """Generate dummy autoencoder latents for cold start.""" + """Dummy autoencoder latents for the sampler to initialize shapes.""" tokens_h = tokens_w = int(math.sqrt(self.model_config.tokens_per_frame)) dummy_frames = torch.randn( 1, self.streaming_config.window_length, @@ -86,25 +97,39 @@ def dummy_batch(self) -> torch.Tensor: device=self.streaming_config.device, dtype=torch.bfloat16) return dummy_frames - + def debug_generate_frame_batch(self, mouse_batch: torch.Tensor, button_batch: torch.Tensor) -> torch.Tensor: + # NOTE Debug mode, just draw overlays on a black video. + num_frames = mouse_batch.shape[1] + video = torch.zeros(num_frames, 256, 256, 3, + device=self.streaming_config.device, + dtype=torch.bfloat16) + frames = _draw_action_overlays(video, button_batch[0, ::], mouse_batch[0, ::]) + if frames == []: + return video + + frames = [torch.from_numpy(frame) for frame in frames] + return torch.stack(frames).permute(0, 3, 1, 2).unsqueeze(0) # [1, num_frames, c, h, w] + async def generate_frame_batch(self, mouse_batch: torch.Tensor, button_batch: torch.Tensor) -> torch.Tensor: """ Generate window_length frames, return first frames_per_batch for streaming. Args: - mouse_batch: [1, window_length, 2] + mouse_batch: [1, window_length, 2] button_batch: [1, window_length, n_buttons] Returns: frame_batch: [1, frames_per_batch, 3, 256, 256] - only streaming frames """ + if self.debug: + return self.debug_generate_frame_batch(mouse_batch, button_batch) + with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): - latents, full_frames = self.sample_window_fn(dummy_batch=self.dummy_batch, mouse=mouse_batch, btn=button_batch) # [1, window_length, 3, 256, 256] - self.add_to_history(latents[0, :self.streaming_config.frames_per_batch, :, :, :]) + latents, full_frames = self.sample_window_fn(dummy_batch=self.dummy_batch, + mouse=mouse_batch, btn=button_batch) # [1, window_length, 3, 256, 256] + self.add_to_history(latents[0, ::], mouse_batch[0, ::], button_batch[0, ::]) # ignore batch dimension (1 user per model) - # Take only first frames_per_batch for streaming - streaming_frames = full_frames[:, :self.streaming_config.frames_per_batch, :, :, :] - return streaming_frames # [1, frames_per_batch, 3, 256, 256] + return full_frames # [1, window_length, 3, 256, 256] def __enter__(self): return self diff --git a/webapp/user_session.py b/webapp/user_session.py index f09a7b65..9a344412 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -7,12 +7,15 @@ import termcolor import numpy as np -from fastapi import WebSocket -from webapp.action_converter import ActionCollector -from webapp.streaming import StreamingFrameGenerator, FrameBuffer +from fastapi import WebSocket +from webapp.action_converter import ActionCollector +from webapp.streaming import StreamingFrameGenerator, FrameBuffer + class UserGameSession: - """Main orchestrator for single-user real-time gameplay.""" + """ + Orchestrates receiving actions from the UI, generating frames, and displaying them. + """ def __init__(self, frame_generator: StreamingFrameGenerator): self.frame_generator = frame_generator self.action_collector = ActionCollector(frame_generator.streaming_config) @@ -25,54 +28,61 @@ async def run_session(self, websocket): print(termcolor.colored(f"Batch duration: {self.frame_generator.streaming_config.batch_duration:.3f}s", "green")) async with asyncio.TaskGroup() as tg: - tg.create_task(self._action_input_loop(websocket)) - tg.create_task(self._frame_generation_loop()) - tg.create_task(self._frame_output_loop(websocket)) + tg.create_task(self._action_input_loop (websocket)) + tg.create_task(self._frame_generation_loop ()) + tg.create_task(self._frame_display_loop (websocket)) async def _action_input_loop(self, websocket): - async for message in websocket: + while True: try: + message = await websocket.receive_text() action_data = json.loads(message) await self.action_collector.add_websocket_action(action_data) except Exception as e: print(f"Error processing action: {e}") + break # Exit loop if WebSocket disconnects async def _frame_generation_loop(self): """Generate frame batches continuously.""" print(termcolor.colored("Frame generation loop started", "green")) while True: try: - # Collect 8 frames worth of actions + # Collect multiple frames worth of actions, e.g. X frames, that have happened between the last frame generation and now. mouse_batch, button_batch = await self.action_collector.collect_batch() - # Generate 8 frames + # Generate Y frames from X actions by taking the X[-1]'th action. Typically, X >> Y, because they are sampled at uncapped FPS from the UI, + # whereas Y frames are sampled from the model one at a time. frame_batch = await self.frame_generator.generate_frame_batch(mouse_batch, button_batch) - # Queue frames for streaming + # Queue frames for streaming at a capped FPS. If model predictions speed up or slow down, it won't cause any dilation of frames being displayed. + # However, if the model predictions are too slow, the frames will be displayed at a lower FPS than the capped FPS. await self.frame_buffer.add_frame_batch(frame_batch) except Exception as e: import traceback print(termcolor.colored(f"Error in frame generation: {e} :\n {traceback.format_exc()}", "red")) - await asyncio.sleep(0.1) # Brief pause before retry + await asyncio.sleep(0.05) # Brief pause before retry - async def _frame_output_loop(self, websocket): + async def _frame_display_loop(self, websocket): while True: try: frame = await self.frame_buffer.get_next_frame() await self._send_frame_to_client(websocket, frame) except Exception as e: - print(termcolor.colored(f"Error in frame streaming: {e}", "red")) + import traceback + print(termcolor.colored(f"Error in frame streaming: {e} :\n {traceback.format_exc()}", "red")) await asyncio.sleep(self.frame_generator.streaming_config.frame_interval) async def _send_frame_to_client(self, websocket: WebSocket, frame: torch.Tensor): - # Convert frame to base64 JPEG - print(f"Frame shape: {frame.shape} - with stats: {frame.min()=}, {frame.max()=}, {frame.mean()=}, {frame.std()=}") + # TODO Do this more intelligently. I'm sure there's better tech to stream video to a UI. + # Convert frame to base64 JPEG\ frame_np = frame.cpu().numpy().transpose(1, 2, 0) # CHW -> HWC - frame_np = ((frame_np + 1) * 127.5).clip(0, 255).astype(np.uint8) # Normalize + # NOTE: Assumes -1 to 1 range for frames. + multiplier = 127.5 if frame.max() < 1 else 1 + frame_np = (frame_np * multiplier).clip(0, 255).astype(np.uint8) # Normalize _, buffer = cv2.imencode('.jpg', frame_np) frame_base64 = base64.b64encode(buffer).decode('utf-8') - await websocket.send(json.dumps({ + await websocket.send_json({ "type": "frame", "data": frame_base64, "timestamp": time.time() - })) + }) diff --git a/webapp/utils/action_builder.py b/webapp/utils/action_builder.py index d0e92288..47be882e 100644 --- a/webapp/utils/action_builder.py +++ b/webapp/utils/action_builder.py @@ -18,7 +18,7 @@ class ActionPattern(Enum): WALK_BACKWARD = "walk_backward" CIRCLE_STRAFE = "circle_strafe" LOOK_AROUND = "look_around" - AIM_AND_SHOOT = "aim_and_shoot" + SHOOT = "shoot" SPRINT_FORWARD = "sprint_forward" RELOAD = "reload" @@ -258,7 +258,7 @@ def generate_pattern(self, mouse = MouseGenerator.look_around(self.config, **mouse_kwargs) buttons = ButtonGenerator.idle(self.config) - elif pattern == ActionPattern.AIM_AND_SHOOT: + elif pattern == ActionPattern.SHOOT: mouse = MouseGenerator.aim_tracking(self.config, **mouse_kwargs) # Add some shooting shoot_times = [(i * 30, 5) for i in range(self.config.sequence_length // 30)] diff --git a/webapp/utils/configs.py b/webapp/utils/configs.py index 0e8eb554..7f5489ec 100644 --- a/webapp/utils/configs.py +++ b/webapp/utils/configs.py @@ -5,27 +5,25 @@ from typing import Optional from dataclasses import dataclass -from owl_wms.configs import Config as RunConfig, TransformerConfig as ModelConfig +from owl_wms.configs import Config as RunConfig, TransformerConfig as ModelConfig, TrainingConfig @dataclass class WebappConfig: model_checkpoint_path : os.PathLike - model_config_path : os.PathLike - model_config : ModelConfig run_config : RunConfig stream_config : StreamingConfig sampling_config : SamplingConfig + run_config_path : os.PathLike device : str = 'cuda' @classmethod def from_yaml(cls, path: os.PathLike) -> WebappConfig: - - with (open(path, 'r') as wcp, - open(config['model_config_path'], 'r') as mcp, - open(config['run_config_path'], 'r') as rcp): - config = yaml.safe_load(wcp) - config['model_config'] = yaml.safe_load(mcp) - config['run_config'] = yaml.safe_load(rcp) + # + with open(path, 'r') as wcp: + config = yaml.safe_load(wcp) + config['run_config'] = RunConfig.from_yaml(config['run_config_path']) + config['sampling_config'] = SamplingConfig(**config['sampling_config']) + config['stream_config'] = StreamingConfig(**config['stream_config']) return cls(**config) @@ -45,6 +43,7 @@ class StreamingConfig: window_length: int = 60 device: str = 'cuda' n_buttons: int = 11 + n_mouse_axes: int = 2 mouse_range: tuple[float, float] = (-1.0, 1.0) @property diff --git a/webapp/utils/demo_streaming.py b/webapp/utils/demo_streaming.py index fdf263bf..6ba54914 100644 --- a/webapp/utils/demo_streaming.py +++ b/webapp/utils/demo_streaming.py @@ -9,6 +9,8 @@ from webapp.utils.configs import SamplingConfig +DEBUG = True + async def demo_streaming_generation(pattern=ActionPattern.LOOK_AROUND): """Generate one batch using StreamingFrameGenerator instead of regular sampler.""" @@ -24,7 +26,6 @@ async def demo_streaming_generation(pattern=ActionPattern.LOOK_AROUND): window_length=60, device=device ) - sampling_config = SamplingConfig() # Load models (reuse render.py's load_models) @@ -37,7 +38,8 @@ async def demo_streaming_generation(pattern=ActionPattern.LOOK_AROUND): print("🎬 Creating streaming frame generator...") frame_generator = StreamingFrameGenerator( encoder, decoder, - streaming_config, model_config, training_config, sampling_config + streaming_config, model_config, training_config, sampling_config, + debug=DEBUG ) # Generate actions (reuse render.py's generate_dummy_actions) @@ -49,7 +51,7 @@ async def demo_streaming_generation(pattern=ActionPattern.LOOK_AROUND): with frame_generator: frame_batch = await frame_generator.generate_frame_batch(mouse_batch, button_batch) - print(f"Generated {frame_batch.shape[1]} frames with shape: {frame_batch.shape}") + print(f"Generated {frame_batch.shape[0]} frames with shape: {frame_batch.shape}") # Save video (reuse render.py's save_video) print("💾 Saving video...") @@ -61,7 +63,7 @@ async def demo_streaming_generation(pattern=ActionPattern.LOOK_AROUND): if __name__ == "__main__": # Try different patterns by changing this: - pattern = ActionPattern.LOOK_AROUND # or AIM_AND_SHOOT, CIRCLE_STRAFE, etc. + pattern = ActionPattern.CIRCLE_STRAFE # or AIM_AND_SHOOT, CIRCLE_STRAFE, etc. print("Available patterns:") for p in ActionPattern: diff --git a/webapp/utils/render.py b/webapp/utils/render.py index f5a7deb6..4cd8d7e3 100644 --- a/webapp/utils/render.py +++ b/webapp/utils/render.py @@ -82,12 +82,8 @@ def save_video(video_tensor: torch.Tensor, filename="generated_video", fps=30): str: Path to saved video file """ setup_output_dir() - - # Generate timestamp for unique filename - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - output_path = Path(OUTPUT_DIR) / f"{filename}_{timestamp}.mp4" - - # Convert tensor to numpy and handle batch dimension + + output_path = Path(OUTPUT_DIR) / f"{filename}.mp4" video_np: np.ndarray = video_tensor.float().cpu().detach().numpy() # Take first batch item if batch_size > 1 @@ -123,7 +119,7 @@ def save_video(video_tensor: torch.Tensor, filename="generated_video", fps=30): print(f"Warning: Could not save as MP4 ({e}), falling back to .pt format") # Fallback: save as PyTorch tensor - fallback_path = Path(OUTPUT_DIR) / f"{filename}_{timestamp}.pt" + fallback_path = Path(OUTPUT_DIR) / f"{filename}.pt" torch.save(video_tensor.cpu(), fallback_path) return str(fallback_path) @@ -172,6 +168,6 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): render_video(verbose=True) # Render with different pattern - render_video(ActionPattern.AIM_AND_SHOOT, verbose=True) + render_video(ActionPattern.SHOOT, verbose=True) render_video(ActionPattern.LOOK_AROUND, verbose=True) render_video(ActionPattern.CIRCLE_STRAFE, verbose=True) diff --git a/webapp/utils/visualize_overlay_actions.py b/webapp/utils/visualize_overlay_actions.py new file mode 100644 index 00000000..4ffb8a3e --- /dev/null +++ b/webapp/utils/visualize_overlay_actions.py @@ -0,0 +1,333 @@ +import imageio as imio +import cv2 +import numpy as np +import math +import torch +from typing import Optional +from contextlib import contextmanager + + +# Global configuration +KEYBINDS = ["W", "A", "S", "D", "LSHIFT", "SPACE", "R", "F", "E", "LMB", "RMB"] +MINIMUM_FRAME_SIZE = 512 +# Colors (BGR format for OpenCV) +COLOR_PRESSED = (50, 200, 50) # Green +COLOR_UNPRESSED = (100, 100, 100) # Gray +COLOR_TEXT = (255, 255, 255) # White +COLOR_BACKGROUND = (30, 30, 30) # Dark gray +COLOR_MOUSE_ARROW = COLOR_PRESSED # Green +COLOR_UNCERTAINTY = (100, 255, 255) # Yellow-ish +COLOR_LMB_SECTOR = COLOR_PRESSED # Green +COLOR_RMB_SECTOR = COLOR_PRESSED # Green + +# Key dimensions +KEY_SIZE = 30 +KEY_MARGIN = 5 +SHIFT_WIDTH = int(KEY_SIZE * 2 + KEY_MARGIN) # Two keys worth of width +SPACE_WIDTH = int(KEY_SIZE * 5) + +# Mouse compass dimensions +COMPASS_RADIUS = 48 # 80% of 80 +COMPASS_START_X_PERCENT = 0.80 + +# Mouse button arc dimensions +MOUSE_BUTTON_OFFSET = 6 # Pixels outside the main circle +MOUSE_BUTTON_THICKNESS = 6 # Increased thickness + +# Arrow scaling parameters +ARROW_SCALE_FACTOR = 1 # Base scaling factor for arrow length +ARROW_MIN_LENGTH = 55 # Minimum arrow length in pixels +ARROW_MAX_SCALE = 0.75 # Maximum scale relative to compass radius + +START_X_PERCENT = 0.12 +START_Y_PERCENT = 0.65 + + +@contextmanager +def _rescale_icons(ratio: float): + """ + Rescale all icon sizes to the ratio of the original video to the 512x512 video + """ + global KEY_SIZE, KEY_MARGIN, SHIFT_WIDTH, SPACE_WIDTH, COMPASS_RADIUS, MOUSE_BUTTON_OFFSET, MOUSE_BUTTON_THICKNESS + global ARROW_SCALE_FACTOR, ARROW_MIN_LENGTH, ARROW_MAX_SCALE + global START_X_PERCENT, START_Y_PERCENT + + try: + old_values = { + "KEY_SIZE": KEY_SIZE, + "KEY_MARGIN": KEY_MARGIN, + "SHIFT_WIDTH": SHIFT_WIDTH, + "SPACE_WIDTH": SPACE_WIDTH, + "COMPASS_RADIUS": COMPASS_RADIUS, + } + + KEY_SIZE *= ratio ; KEY_SIZE = int(KEY_SIZE) + KEY_MARGIN *= ratio ; KEY_MARGIN = int(KEY_MARGIN) + SHIFT_WIDTH *= ratio ; SHIFT_WIDTH = int(SHIFT_WIDTH) + SPACE_WIDTH *= ratio ; SPACE_WIDTH = int(SPACE_WIDTH) + COMPASS_RADIUS *= ratio ; COMPASS_RADIUS = int(COMPASS_RADIUS) + MOUSE_BUTTON_OFFSET *= ratio ; MOUSE_BUTTON_OFFSET = int(MOUSE_BUTTON_OFFSET) + MOUSE_BUTTON_THICKNESS *= ratio ; MOUSE_BUTTON_THICKNESS = int(MOUSE_BUTTON_THICKNESS) + ARROW_SCALE_FACTOR *= ratio ; ARROW_SCALE_FACTOR = float(ARROW_SCALE_FACTOR) + ARROW_MIN_LENGTH *= ratio ; ARROW_MIN_LENGTH = int(ARROW_MIN_LENGTH) + ARROW_MAX_SCALE *= ratio ; ARROW_MAX_SCALE = float(ARROW_MAX_SCALE) + yield + finally: + for key, value in old_values.items(): + globals()[key] = value + + +def _draw_buttons( + frame: np.ndarray, + button_sequence: list[bool], +) -> None: + """ + Draw keyboard buttons on the frame. + + Args: + frame: numpy array representing the image frame + button_sequence: list of bools corresponding to KEYBINDS (excluding LMB, RMB) + """ + # Starting position for keyboard layout (bottom left) + start_x = int(frame.shape[1] * START_X_PERCENT) + start_y = int(frame.shape[0] * START_Y_PERCENT) + + # Key positions - organized by rows + key_positions = { + # Top row: W E R (W above S) + "W": (start_x + (KEY_SIZE + KEY_MARGIN) * 1, start_y), + "E": (start_x + (KEY_SIZE + KEY_MARGIN) * 2, start_y), + "R": (start_x + (KEY_SIZE + KEY_MARGIN) * 3, start_y), + + # Middle row: A S D F + "A": (start_x + (KEY_SIZE + KEY_MARGIN) * 0, start_y + KEY_SIZE + KEY_MARGIN), + "S": (start_x + (KEY_SIZE + KEY_MARGIN) * 1, start_y + KEY_SIZE + KEY_MARGIN), + "D": (start_x + (KEY_SIZE + KEY_MARGIN) * 2, start_y + KEY_SIZE + KEY_MARGIN), + "F": (start_x + (KEY_SIZE + KEY_MARGIN) * 3, start_y + KEY_SIZE + KEY_MARGIN), + + # Bottom row: LSHIFT SPACE + "LSHIFT": (start_x - (KEY_SIZE + KEY_MARGIN), start_y + (KEY_SIZE + KEY_MARGIN) * 2), + "SPACE": (start_x + SHIFT_WIDTH + KEY_MARGIN, start_y + (KEY_SIZE + KEY_MARGIN) * 2), + } + + # Draw each key + for i, key in enumerate(KEYBINDS[:-2]): # Exclude LMB and RMB + if key in key_positions: + x, y = key_positions[key] + + # Determine key dimensions + if key == "LSHIFT": + width = SHIFT_WIDTH + height = KEY_SIZE + elif key == "SPACE": + width = SPACE_WIDTH + height = KEY_SIZE + else: + width = KEY_SIZE + height = KEY_SIZE + + # Determine color based on pressed state + color = COLOR_PRESSED if button_sequence[i] else COLOR_UNPRESSED + + # Draw key background + cv2.rectangle(frame, (x, y), (x + width, y + height), color, -1) + + # Draw key border + cv2.rectangle(frame, (x, y), (x + width, y + height), COLOR_TEXT, 1) + + # Draw key label + text_size = cv2.getTextSize(key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0] + text_x = x + (width - text_size[0]) // 2 + text_y = y + (height + text_size[1]) // 2 + cv2.putText(frame, key, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, + 0.5, COLOR_TEXT, 1, cv2.LINE_AA) + + +def _draw_mouse( + frame: np.ndarray, + LMB_on: bool, + RMB_on: bool, + mouse_delta: tuple[float, float], + center: tuple[int, int], +) -> None: + """ + Draw mouse compass with direction arrow and uncertainty cone. + Arrow length is proportional to mouse movement magnitude. + + Args: + frame: numpy array representing the image frame + LMB_on: bool indicating if left mouse button is pressed + RMB_on: bool indicating if right mouse button is pressed + mouse_vec: tuple of floats (x, y) representing mouse direction + mouse_std: tuple of floats (x_std, y_std) representing uncertainty + center: tuple (x, y) for compass center position + """ + # Draw compass circle + cv2.circle(frame, center, COMPASS_RADIUS, COLOR_TEXT, 2) + + # Calculate outer radius for mouse buttons (slightly outside the main compass) + button_radius = COMPASS_RADIUS + MOUSE_BUTTON_OFFSET + + # Draw LMB arc (top-left 45 degrees) - outside the main circle and thicker + color_lmb = COLOR_LMB_SECTOR if LMB_on else (60, 60, 60) # Gray when off + cv2.ellipse(frame, center, (button_radius, button_radius), + 0, 225, 270, color_lmb, MOUSE_BUTTON_THICKNESS) + + # Draw RMB arc (top-right 45 degrees) - outside the main circle and thicker + color_rmb = COLOR_RMB_SECTOR if RMB_on else (60, 60, 60) # Gray when off + cv2.ellipse(frame, center, (button_radius, button_radius), + 0, 270, 315, color_rmb, MOUSE_BUTTON_THICKNESS) + + text_loc_lmb = center[0] - COMPASS_RADIUS + 10, center[1] - COMPASS_RADIUS - 10 + cv2.putText(frame, 'LMB', text_loc_lmb, cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_TEXT, 1, cv2.LINE_AA) + text_loc_rmb = center[0], center[1] - COMPASS_RADIUS - 10 + cv2.putText(frame, 'RMB', text_loc_rmb, cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_TEXT, 1, cv2.LINE_AA) + + # Calculate mouse direction + mouse_x, mouse_y = mouse_delta + magnitude = math.sqrt(mouse_x**2 + mouse_y**2) + + if magnitude > 0: + # Normalize to unit vector + unit_x = mouse_x / magnitude + unit_y = mouse_y / magnitude + + # Calculate arrow length proportional to magnitude + # Scale by magnitude but clamp to reasonable bounds + arrow_scale = max( + ARROW_MIN_LENGTH, # Ensure minimum visible arrow + min( + COMPASS_RADIUS * magnitude * ARROW_SCALE_FACTOR, # Proportional to magnitude + COMPASS_RADIUS * ARROW_MAX_SCALE # Cap at maximum scale + ) + ) + + end_x = int(center[0] + unit_x * arrow_scale) + end_y = int(center[1] - unit_y * arrow_scale) # Negative because y-axis is inverted + + # Draw direction arrow with proportional length + cv2.arrowedLine(frame, center, (end_x, end_y), COLOR_MOUSE_ARROW, 3, tipLength=0.3) + + # Optional: Display magnitude as text for debugging + # magnitude_text = f"Mag: {magnitude:.2f}" + # cv2.putText(frame, magnitude_text, (center[0] - 40, center[1] + button_radius + 40), + # cv2.FONT_HERSHEY_SIMPLEX, 0.4, COLOR_TEXT, 1, cv2.LINE_AA) + + # Draw center dot + cv2.circle(frame, center, 3, COLOR_TEXT, -1) + + # Add labels + cv2.putText(frame, "Mouse", (center[0] - 25, center[1] + button_radius + 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, COLOR_TEXT, 1, cv2.LINE_AA) + +def _draw_frame( + frame: np.ndarray, + buttons: list[bool] | torch.Tensor, + mouse_delta: tuple[float, float], +) -> np.ndarray: + """ + Overlay keyboard & mouse on a single frame and return it. + """ + buttons, lmb, rmb = buttons[:-2], buttons[-2], buttons[-1] + _draw_buttons(frame, buttons) + + frame_height, frame_width = frame.shape[:2] + margin = 30 + compass_x = frame_width - margin - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET + compass_y = int(frame_height * START_Y_PERCENT) + compass_x = int(compass_x * COMPASS_START_X_PERCENT) + + # Ensure compass fits within frame boundaries + min_x = COMPASS_RADIUS + MOUSE_BUTTON_OFFSET + 20 + max_x = frame_width - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET - 20 + compass_x = max(min_x, min(compass_x, max_x)) + min_y = COMPASS_RADIUS + MOUSE_BUTTON_OFFSET + 30 + max_y = frame_height - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET - 30 + compass_y = max(min_y, min(compass_y, max_y)) + + _draw_mouse(frame, + LMB_on=lmb, RMB_on=rmb, + mouse_delta=mouse_delta, + center=(compass_x, compass_y)) + return frame + +def _draw_video( + video: torch.Tensor, + buttons: torch.Tensor, + mouse_delta: torch.Tensor, + save_path: Optional[str] = None, + fps: int = 30, + arrow_scale_factor: Optional[float] = None, + arrow_max_scale: Optional[float] = None, +) -> list[np.ndarray]: + """ + Draw video with input device monitoring overlays. + + Args: + video: torch tensor of video frames + buttons: torch tensor of button states + mouse_vec: torch tensor of mouse vectors + mouse_std: torch tensor of mouse standard deviations + save_path: optional path to save video + fps: frames per second for output video + arrow_scale_factor: optional override for arrow scaling + arrow_max_scale: optional override for maximum arrow scale + """ + # Update global arrow parameters if provided + global ARROW_SCALE_FACTOR, ARROW_MAX_SCALE + if arrow_scale_factor is not None: + ARROW_SCALE_FACTOR = arrow_scale_factor + if arrow_max_scale is not None: + ARROW_MAX_SCALE = arrow_max_scale + + video = video.float().cpu().numpy() + # upscale image to 512x512 only if < 512 on smallest dimension. + if video.shape[1] < MINIMUM_FRAME_SIZE or video.shape[2] < MINIMUM_FRAME_SIZE: + frames = [cv2.resize(frame, (512, 512), interpolation=cv2.INTER_CUBIC) for frame in video] + ratio = 1.0 + else: + frames = video + ratio = min(video.shape[1] / MINIMUM_FRAME_SIZE, video.shape[2] / MINIMUM_FRAME_SIZE) + + with _rescale_icons(ratio): + frames = [ + _draw_frame(frame, buttons[i], mouse_delta[i]) + for i, frame in enumerate(frames) + ] + if save_path is not None: + imio.mimsave(save_path, frames, fps=fps) + + return frames + + +# Example usage +if __name__ == "__main__": + # Example button states (9 buttons, excluding LMB and RMB) + button_states = [True, False, True, False, True, False, False, True, False] + + # Example mouse state + LMB_pressed = True + RMB_pressed = False + mouse_direction = (0.7, 0.7) # Diagonal up-right + mouse_uncertainty = (0.2, 0.1) + + vidpath, mousepath, buttonpath = '.pt' + + video = torch.load(vidpath, map_location='cpu', mmap=True) + mouse = torch.load(mousepath, map_location='cpu', mmap=True) + buttons = torch.load(buttonpath, map_location='cpu', mmap=True) # [1, window_length, n_buttons] + + min_len = min(len(video), len(mouse), len(buttons)) + video = video[:min_len] + mouse = mouse[:min_len] + buttons = buttons[:min_len] + + video = video.float() + buttons = buttons + mouse = mouse + + # You can adjust arrow scaling parameters here + _draw_video(video, buttons, mouse, + save_path="groundtruth.mp4", + arrow_scale_factor=0.3, # Adjust this to control arrow sensitivity + arrow_max_scale=0.9) # Maximum arrow length relative to compass From c53b6b8c2c950da77cfe49e67e28668b933de340 Mon Sep 17 00:00:00 2001 From: Sami Date: Sun, 8 Jun 2025 03:27:41 -0400 Subject: [PATCH 10/71] dont load models if debugging --- webapp/server.py | 12 ++++++------ webapp/utils/samplers.py | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/webapp/server.py b/webapp/server.py index 5e5d0666..21be481d 100644 --- a/webapp/server.py +++ b/webapp/server.py @@ -25,12 +25,12 @@ async def lifespan(app: FastAPI): global encoder, decoder, config config = WebappConfig.from_yaml(webapp_config_path) - - encoder, decoder, _ = load_models( - checkpoint_path=config.model_checkpoint_path, - config_path=config.run_config_path, - device=config.device, verbose=True, - ) + if not DEBUG: + encoder, decoder, _ = load_models( + checkpoint_path=config.model_checkpoint_path, + config_path=config.run_config_path, + device=config.device, verbose=True, + ) yield encoder, decoder, config = None, None, None diff --git a/webapp/utils/samplers.py b/webapp/utils/samplers.py index dea3caf6..63e1a437 100644 --- a/webapp/utils/samplers.py +++ b/webapp/utils/samplers.py @@ -15,7 +15,7 @@ LatentData = Tensor @multimethod -def create_sampler(sampler_id: Literal['cfg'], encoder: nn.Module, decoder: nn.Module, +def create_sampler(sampler_id: Literal['cfg'], encoder, decoder, batch_size: int = 8, sampling_steps: int = SAMPLING_STEPS, cfg_scale: float = CFG_SCALE, @@ -34,7 +34,7 @@ def _sampler(): return CFGSampler() ) @multimethod -def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder: nn.Module, decoder: nn.Module, +def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder, decoder, batch_size: int = 8, sampling_steps: int = SAMPLING_STEPS, cfg_scale: float = CFG_SCALE, @@ -53,7 +53,7 @@ def _sampler(): return InpaintCFGSampler() ) @multimethod -def create_sampler(sampler_id: Literal['window'], encoder: nn.Module, decoder: nn.Module, +def create_sampler(sampler_id: Literal['window'], encoder, decoder, batch_size: int = 8, sampling_steps: int = SAMPLING_STEPS, cfg_scale: float = CFG_SCALE, From 260114dffc19d53fcd1a9e9d0d2458974d98b94e Mon Sep 17 00:00:00 2001 From: Sami Date: Sun, 8 Jun 2025 17:21:46 -0400 Subject: [PATCH 11/71] fixed some issues with webapp and better mouse tracking --- configs/webapp/config.yaml | 2 +- webapp/action_converter.py | 44 ++++++++++++++-- webapp/static/index.html | 100 ++++++++++++++++++++++++++++--------- webapp/user_session.py | 54 ++++++++++++-------- 4 files changed, 152 insertions(+), 48 deletions(-) diff --git a/configs/webapp/config.yaml b/configs/webapp/config.yaml index aa27f5a6..dc1d87da 100644 --- a/configs/webapp/config.yaml +++ b/configs/webapp/config.yaml @@ -4,7 +4,7 @@ device: "cuda" stream_config: fps: 20 - frames_per_batch: 8 + frames_per_batch: 2 window_length: 60 device: "cuda" n_buttons: 11 diff --git a/webapp/action_converter.py b/webapp/action_converter.py index 6b7b1cba..cb801256 100644 --- a/webapp/action_converter.py +++ b/webapp/action_converter.py @@ -24,6 +24,7 @@ def _interpolate(tensor_batch: torch.Tensor, tensor_batch: [1, target_length, features] """ num_actions = tensor_batch.shape[1] + if num_actions >= target_length: # subsample actions if somehow longer than frames_per_batch downsampled = torch.arange(0, num_actions, num_actions // target_length) @@ -31,7 +32,12 @@ def _interpolate(tensor_batch: torch.Tensor, # Repeat with empty actions to fill remaining frames num_missing_actions = target_length - num_actions - repeated = empty_action.repeat(1, num_missing_actions, 1) # [1, missing_frames, features] + if num_missing_actions == target_length: + return empty_action.repeat(1, target_length, 1) + + # NOTE: Repeat last action for the remaining frames + last_action = tensor_batch[:, -1, :] + repeated = last_action.repeat(1, num_missing_actions, 1) # [1, missing_frames, features] return torch.cat([tensor_batch, repeated], dim=1) # [1, target_length, features] @@ -120,7 +126,10 @@ def __init__(self, streaming_config: StreamingConfig): async def add_websocket_action(self, ws_message: dict): """Add action from WebSocket message.""" action = self.converter.websocket_to_action(ws_message) - await self.action_queue.put(action) + # Add timestamp to track when action was received + timestamped_action = (action, time.time()) + + await self.action_queue.put(timestamped_action) async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: """ @@ -140,18 +149,45 @@ async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: real_actions = [] start_time = time.time() + # First, clear any stale actions from the queue (older than 1 second) + stale_threshold = start_time - 1.0 + temp_actions = [] + stale_count = 0 + + # Drain the queue and filter out stale actions + while not self.action_queue.empty(): + try: + timestamped_action = self.action_queue.get_nowait() + action, timestamp = timestamped_action + if timestamp >= stale_threshold: + temp_actions.append(action) + else: + stale_count += 1 + except asyncio.QueueEmpty: + break + + # Re-add fresh actions to the queue + for action in temp_actions: + try: + await self.action_queue.put((action, time.time())) + except asyncio.QueueFull: + break # Skip if queue is full + + # Now collect actions for the current batch while start_time + self.streaming_config.batch_duration > time.time(): try: timeout = max(0.01, self.streaming_config.batch_duration - (time.time() - start_time)) - action = await asyncio.wait_for(self.action_queue.get(), timeout=timeout) + timestamped_action = await asyncio.wait_for(self.action_queue.get(), timeout=timeout) + action, timestamp = timestamped_action real_actions.append(action) except asyncio.TimeoutError: pass - # Convert 8 real actions to batch tensors + # Convert real actions to batch tensors mouse, button = self.converter.actions_to_batch(real_actions) mouse = _interpolate(mouse, empty_action=self.empty_mouse, target_length=self.streaming_config.frames_per_batch) button = _interpolate(button, empty_action=self.empty_buttons, target_length=self.streaming_config.frames_per_batch) + return mouse, button diff --git a/webapp/static/index.html b/webapp/static/index.html index b15e0a9c..091b6929 100644 --- a/webapp/static/index.html +++ b/webapp/static/index.html @@ -75,19 +75,37 @@

🎮 OWL-WMS Game Stream

"RMB": false }; + // Mouse tracking for deltas with exponential smoothing + let lastMouseX = 0; + let lastMouseY = 0; + let isFirstMove = true; + let smoothedDeltaX = 0; + let smoothedDeltaY = 0; + const DELTA_SCALE = 0.8; // Scale down deltas for smoother movement + const DEAD_ZONE = 0.002; // Ignore tiny movements + const SMOOTHING_FACTOR = 0.3; // Lower = smoother, higher = more responsive + // WebSocket connection - const ws = new WebSocket('ws://localhost:8000/ws/game'); + const scheme = location.protocol === 'https:' ? 'wss' : 'ws'; + const ws = new WebSocket(`${scheme}://${location.hostname}:8000/ws/game`); + ws.onopen = () => { console.log('WebSocket connected'); statusDiv.textContent = '✅ Connected to game server'; statusDiv.className = 'connected'; + + // Start sending continuous keyboard updates at 60 FPS + startContinuousKeyboardUpdates(); }; ws.onclose = () => { console.log('WebSocket disconnected'); statusDiv.textContent = '❌ Disconnected from server'; statusDiv.className = 'disconnected'; + + // Stop continuous updates + stopContinuousKeyboardUpdates(); }; ws.onerror = (error) => { @@ -105,6 +123,26 @@

🎮 OWL-WMS Game Stream

} }; + let updateInterval = null; + + function startContinuousKeyboardUpdates() { + // Send keyboard state at 60 FPS, but only if keys are pressed + // NO MOUSE DATA - only keyboard + updateInterval = setInterval(() => { + const hasKeysPressed = Object.values(buttons).some(pressed => pressed); + if (hasKeysPressed) { + sendAction(0, 0, getCurrentButtons()); // Zero mouse deltas for keyboard-only updates + } + }, 16); + } + + function stopContinuousKeyboardUpdates() { + if (updateInterval) { + clearInterval(updateInterval); + updateInterval = null; + } + } + function displayFrame(frameData) { if (frameData.type === 'frame') { const img = new Image(); @@ -151,11 +189,7 @@

🎮 OWL-WMS Game Stream

if (buttonName && !buttons[buttonName]) { e.preventDefault(); buttons[buttonName] = true; - // Send current mouse position with new button state - const rect = canvas.getBoundingClientRect(); - const mouse_x = 0; // Default to center if no recent mouse event - const mouse_y = 0; - sendAction(mouse_x, mouse_y, getCurrentButtons()); + console.log(`🔽 Key pressed: ${buttonName}`); } }); @@ -179,20 +213,44 @@

🎮 OWL-WMS Game Stream

if (buttonName && buttons[buttonName]) { e.preventDefault(); buttons[buttonName] = false; - // Send current mouse position with new button state - const rect = canvas.getBoundingClientRect(); - const mouse_x = 0; // Default to center if no recent mouse event - const mouse_y = 0; - sendAction(mouse_x, mouse_y, getCurrentButtons()); + console.log(`🔼 Key released: ${buttonName}`); } }); - // Mouse event listeners + // Mouse event listeners - SMOOTH DELTAS canvas.addEventListener('mousemove', (e) => { const rect = canvas.getBoundingClientRect(); - const mouse_x = (e.clientX - rect.left) / rect.width * 2 - 1; - const mouse_y = -((e.clientY - rect.top) / rect.height * 2 - 1); - sendAction(mouse_x, mouse_y, getCurrentButtons()); + const currentMouseX = (e.clientX - rect.left) / rect.width * 2 - 1; + const currentMouseY = -((e.clientY - rect.top) / rect.height * 2 - 1); + + if (isFirstMove) { + // First move - just set position, don't send + lastMouseX = currentMouseX; + lastMouseY = currentMouseY; + isFirstMove = false; + return; + } + + // Calculate raw deltas + let deltaX = (currentMouseX - lastMouseX) * DELTA_SCALE; + let deltaY = (currentMouseY - lastMouseY) * DELTA_SCALE; + + // Apply dead zone - ignore tiny movements + if (Math.abs(deltaX) < DEAD_ZONE) deltaX = 0; + if (Math.abs(deltaY) < DEAD_ZONE) deltaY = 0; + + // Calculate smoothed deltas using exponential smoothing + smoothedDeltaX = deltaX * SMOOTHING_FACTOR + smoothedDeltaX * (1 - SMOOTHING_FACTOR); + smoothedDeltaY = deltaY * SMOOTHING_FACTOR + smoothedDeltaY * (1 - SMOOTHING_FACTOR); + + // Only send if there's actual movement after smoothing + if (smoothedDeltaX !== 0 || smoothedDeltaY !== 0) { + sendAction(smoothedDeltaX, smoothedDeltaY, getCurrentButtons()); + } + + // Update last position + lastMouseX = currentMouseX; + lastMouseY = currentMouseY; }); canvas.addEventListener('mousedown', (e) => { @@ -200,10 +258,8 @@

🎮 OWL-WMS Game Stream

if (e.button === 0) buttons["LMB"] = true; // Left mouse button if (e.button === 2) buttons["RMB"] = true; // Right mouse button - const rect = canvas.getBoundingClientRect(); - const mouse_x = (e.clientX - rect.left) / rect.width * 2 - 1; - const mouse_y = -((e.clientY - rect.top) / rect.height * 2 - 1); - sendAction(mouse_x, mouse_y, getCurrentButtons()); + // Send button state change with zero mouse delta + sendAction(0, 0, getCurrentButtons()); }); canvas.addEventListener('mouseup', (e) => { @@ -211,10 +267,8 @@

🎮 OWL-WMS Game Stream

if (e.button === 0) buttons["LMB"] = false; // Left mouse button if (e.button === 2) buttons["RMB"] = false; // Right mouse button - const rect = canvas.getBoundingClientRect(); - const mouse_x = (e.clientX - rect.left) / rect.width * 2 - 1; - const mouse_y = -((e.clientY - rect.top) / rect.height * 2 - 1); - sendAction(mouse_x, mouse_y, getCurrentButtons()); + // Send button state change with zero mouse delta + sendAction(0, 0, getCurrentButtons()); }); // Prevent context menu on right click diff --git a/webapp/user_session.py b/webapp/user_session.py index 9a344412..04bc629a 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -39,8 +39,13 @@ async def _action_input_loop(self, websocket): action_data = json.loads(message) await self.action_collector.add_websocket_action(action_data) except Exception as e: - print(f"Error processing action: {e}") - break # Exit loop if WebSocket disconnects + # Check if this is a WebSocket disconnect + if "websocket.close" in str(e) or "response already completed" in str(e) or "WebSocket" in str(e): + print(termcolor.colored("🔌 WebSocket disconnected - stopping action input", "yellow")) + break + else: + print(f"Error processing action: {e}") + break # Exit loop on any other error async def _frame_generation_loop(self): """Generate frame batches continuously.""" @@ -66,23 +71,32 @@ async def _frame_display_loop(self, websocket): frame = await self.frame_buffer.get_next_frame() await self._send_frame_to_client(websocket, frame) except Exception as e: - import traceback - print(termcolor.colored(f"Error in frame streaming: {e} :\n {traceback.format_exc()}", "red")) - await asyncio.sleep(self.frame_generator.streaming_config.frame_interval) + # Check if this is a WebSocket disconnect + if "websocket.close" in str(e) or "response already completed" in str(e): + print(termcolor.colored("🔌 WebSocket disconnected - stopping frame stream", "yellow")) + break + else: + import traceback + print(termcolor.colored(f"Error in frame streaming: {e} :\n {traceback.format_exc()}", "red")) + await asyncio.sleep(self.frame_generator.streaming_config.frame_interval) async def _send_frame_to_client(self, websocket: WebSocket, frame: torch.Tensor): - # TODO Do this more intelligently. I'm sure there's better tech to stream video to a UI. - # Convert frame to base64 JPEG\ - frame_np = frame.cpu().numpy().transpose(1, 2, 0) # CHW -> HWC - # NOTE: Assumes -1 to 1 range for frames. - multiplier = 127.5 if frame.max() < 1 else 1 - frame_np = (frame_np * multiplier).clip(0, 255).astype(np.uint8) # Normalize - - _, buffer = cv2.imencode('.jpg', frame_np) - frame_base64 = base64.b64encode(buffer).decode('utf-8') - - await websocket.send_json({ - "type": "frame", - "data": frame_base64, - "timestamp": time.time() - }) + try: + # TODO Do this more intelligently. I'm sure there's better tech to stream video to a UI. + # Convert frame to base64 JPEG\ + frame_np = frame.cpu().numpy().transpose(1, 2, 0) # CHW -> HWC + # NOTE: Assumes -1 to 1 range for frames. + multiplier = 127.5 if frame.max() < 1 else 1 + frame_np = (frame_np * multiplier).clip(0, 255).astype(np.uint8) # Normalize + + _, buffer = cv2.imencode('.jpg', frame_np) + frame_base64 = base64.b64encode(buffer).decode('utf-8') + + await websocket.send_json({ + "type": "frame", + "data": frame_base64, + "timestamp": time.time() + }) + except Exception as e: + # Re-raise to be caught by the display loop + raise e From f92377856af0c0fb1b00fba5ca990137d626fedc Mon Sep 17 00:00:00 2001 From: Sami Date: Sun, 8 Jun 2025 19:55:48 -0400 Subject: [PATCH 12/71] cleanup code, add action display margin --- configs/webapp/config.yaml | 1 + webapp/action_converter.py | 51 ++++++++++---------- webapp/static/index.html | 17 ++++++- webapp/streaming.py | 57 ++++++++++++++++------- webapp/user_session.py | 13 +++--- webapp/utils/configs.py | 1 + webapp/utils/visualize_overlay_actions.py | 5 +- 7 files changed, 92 insertions(+), 53 deletions(-) diff --git a/configs/webapp/config.yaml b/configs/webapp/config.yaml index dc1d87da..9f464a0a 100644 --- a/configs/webapp/config.yaml +++ b/configs/webapp/config.yaml @@ -10,6 +10,7 @@ stream_config: n_buttons: 11 n_mouse_axes: 2 mouse_range: [-1.0, 1.0] + action_margin_px_height: 150 sampling_config: sampling_steps: 20 diff --git a/webapp/action_converter.py b/webapp/action_converter.py index cb801256..135c0bb5 100644 --- a/webapp/action_converter.py +++ b/webapp/action_converter.py @@ -9,7 +9,7 @@ BUTTON_INDICES = {name: idx for idx, name in enumerate(BUTTON_NAMES)} -def _interpolate(tensor_batch: torch.Tensor, +def _interpolate(actions: torch.Tensor, empty_action: torch.Tensor, target_length: int) -> torch.Tensor: @@ -21,25 +21,25 @@ def _interpolate(tensor_batch: torch.Tensor, Must provide empty_action, which is the action to repeat when the batch is shorter than target_length. Must also provide target_length, which is the length to interpolate to. Returns: - tensor_batch: [1, target_length, features] + actions: [target_length, features] """ - num_actions = tensor_batch.shape[1] + num_actions = actions.shape[0] if num_actions >= target_length: # subsample actions if somehow longer than frames_per_batch - downsampled = torch.arange(0, num_actions, num_actions // target_length) - return tensor_batch[:, downsampled, :] + downsampled = torch.arange(0, num_actions, step=(num_actions // target_length)) + return actions[downsampled, :] # Repeat with empty actions to fill remaining frames num_missing_actions = target_length - num_actions if num_missing_actions == target_length: - return empty_action.repeat(1, target_length, 1) + return empty_action.repeat(target_length, 1) # NOTE: Repeat last action for the remaining frames - last_action = tensor_batch[:, -1, :] - repeated = last_action.repeat(1, num_missing_actions, 1) # [1, missing_frames, features] + last_action = actions[-1, :] + repeated = last_action.repeat(num_missing_actions, 1) # [missing_frames, features] - return torch.cat([tensor_batch, repeated], dim=1) # [1, target_length, features] + return torch.cat([actions, repeated], dim=0) # [target_length, features] class ActionConverter: @@ -92,20 +92,20 @@ def actions_to_batch(self, actions: list[tuple[torch.Tensor, torch.Tensor]]) -> actions: List of (mouse, buttons) tuples Returns: - mouse_batch: [batch_size, sequence_length, 2] - button_batch: [batch_size, sequence_length, n_buttons] + mouse: [sequence_length, 2] + button: [sequence_length, n_buttons] """ if not actions: - # Return empty batch - [B, 0, 2] and [B, 0, n_buttons] + # Return empty batch - [0, 2] and [0, n_buttons] return ( - torch.zeros(1, 0, 2, device=self.device), - torch.zeros(1, 0, self.streaming_config.n_buttons, device=self.device) + torch.zeros(0, 2, device=self.device), + torch.zeros(0, self.streaming_config.n_buttons, device=self.device) ) - mouse_batch = torch.stack([action[0] for action in actions], dim=0).unsqueeze(0) # [1, seq_len, 2] - button_batch = torch.stack([action[1] for action in actions], dim=0).unsqueeze(0) # [1, seq_len, n_buttons] + mouse = torch.stack([action[0] for action in actions], dim=0) # [seq_len, 2] + button = torch.stack([action[1] for action in actions], dim=0) # [seq_len, n_buttons] - return mouse_batch, button_batch + return mouse, button class ActionCollector: @@ -115,7 +115,6 @@ def __init__(self, streaming_config: StreamingConfig): self.streaming_config = streaming_config self.converter = ActionConverter(streaming_config) self.action_queue = asyncio.Queue(maxsize=100) # Buffer incoming actions - self.current_batch = [] # -- constants self.empty_mouse = torch.zeros((streaming_config.frames_per_batch, streaming_config.n_mouse_axes), @@ -131,20 +130,20 @@ async def add_websocket_action(self, ws_message: dict): await self.action_queue.put(timestamped_action) - async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: + async def collect_actions(self) -> tuple[torch.Tensor, torch.Tensor]: """ Collect actions from the UI at whatever rate is sent over by the client. This collects all the frames that have been supplied between the last frame generation and now. It takes in as many actions as the model generates frames at once. For example, in CausVid, if - the model generates 4 frames at a time, this function will return [1, 4, 2] and [1, 4, 11] + the model generates 4 frames at a time, this function will return [4, 2] and [4, 11] for mouse and button actions. If, for one reason or another, we have <4 actions, we will fill the batch with idle actions. Returns: - mouse_batch: [1, X, 2] - button_batch: [1, X, n_buttons] + mouse: [X, 2] + button: [X, n_buttons] """ real_actions = [] start_time = time.time() @@ -185,9 +184,9 @@ async def collect_batch(self) -> tuple[torch.Tensor, torch.Tensor]: # Convert real actions to batch tensors mouse, button = self.converter.actions_to_batch(real_actions) - mouse = _interpolate(mouse, empty_action=self.empty_mouse, - target_length=self.streaming_config.frames_per_batch) - button = _interpolate(button, empty_action=self.empty_buttons, - target_length=self.streaming_config.frames_per_batch) + mouse = _interpolate(mouse, empty_action=self.empty_mouse, + target_length=self.streaming_config.frames_per_batch) + button = _interpolate(button,empty_action=self.empty_buttons, + target_length=self.streaming_config.frames_per_batch) return mouse, button diff --git a/webapp/static/index.html b/webapp/static/index.html index 091b6929..8db8e8bf 100644 --- a/webapp/static/index.html +++ b/webapp/static/index.html @@ -52,8 +52,23 @@

🎮 OWL-WMS Game Stream

Connecting to server...
- +

Move your mouse over the canvas to interact with the game world!

+ - - \ No newline at end of file diff --git a/webapp/static/todo b/webapp/static/todo deleted file mode 100644 index d142e548..00000000 --- a/webapp/static/todo +++ /dev/null @@ -1 +0,0 @@ -heavens help me its time for webdev \ No newline at end of file diff --git a/webapp/streaming.py b/webapp/streaming.py index e19b08f7..7d3642a8 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -2,8 +2,9 @@ import time import torch import asyncio -from torch import nn import numpy as np +from torch import nn +from functools import cache from webapp.utils.samplers import create_sampler from webapp.utils.configs import SamplingConfig, StreamingConfig @@ -18,29 +19,36 @@ class FrameBuffer: """ def __init__(self, streaming_config: StreamingConfig): - self.streaming_config = streaming_config - self.frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches - self.last_frame_time = 0.0 + self.streaming_config = streaming_config + self.video_frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches + self.buttons_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches + self.mouse_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches + self.last_frame_time = 0.0 - async def queue_frames(self, frames: torch.Tensor): - # frames shape: [frames_per_batch, channels, height, width] - num_frames = frames.shape[0] + async def queue_frames(self, video_frames: torch.Tensor, mouse: torch.Tensor, button: torch.Tensor): + # video_frames shape: [frames_per_batch, channels, height, width] + # overlay_frames shape: [frames_per_batch, channels, overlay_height, width] + num_frames = video_frames.shape[0] for i in range(num_frames): - await self.frame_queue.put(frames[i]) + await self.video_frame_queue.put(video_frames[i]) + await self.buttons_queue.put(button[i]) + await self.mouse_queue.put(mouse[i]) - async def get_next_frame(self) -> torch.Tensor: - """Get next frame for streaming at capped FPS.""" + async def get_next_frames(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Get next video and overlay frames for streaming at capped FPS.""" now = time.time() time_since_last = now - self.last_frame_time - time_to_wait = max(0, self.streaming_config.frame_interval - time_since_last) - + time_to_wait = max(0, self.streaming_config.frame_interval - time_since_last) + if time_to_wait > 0: await asyncio.sleep(time_to_wait) - frame = await self.frame_queue.get() + video_frame = await self.video_frame_queue.get() + button = await self.buttons_queue.get() + mouse = await self.mouse_queue.get() self.last_frame_time = time.time() - return frame + return video_frame, button, mouse class StreamingFrameGenerator: """Wraps WindowCFGSampler to generate frames.""" @@ -60,34 +68,44 @@ def __init__(self, encoder: nn.Module, decoder: nn.Module, self.encoder = encoder self.decoder = decoder # - # Create WindowCFGSampler for 8-frame generation + # Create WindowCFGSampler for frame generation self.sample_window_fn = create_sampler('window', encoder, decoder, batch_size=1, sampling_steps=self.sampling_config.sampling_steps, + vae_scale=self.sampling_config.vae_scale, cfg_scale=self.sampling_config.cfg_scale, - scale=self.train_config.vae_scale) + window_length=self.sampling_config.window_length, + num_frames=self.sampling_config.num_frames, + noise_prev=self.sampling_config.noise_prev) # Initialize frame history as empty tensor - self.frame_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) - self.mouse_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) + self.latent_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) + self.mouse_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) self.button_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) def add_to_history(self, frame_batch: torch.Tensor, mouse_batch: torch.Tensor, button_batch: torch.Tensor): - if self.frame_history.equal(torch.tensor([], device=self.streaming_config.device)): - self.frame_history = frame_batch + if self.latent_history.equal(torch.tensor([], device=self.streaming_config.device)): + self.latent_history = frame_batch self.mouse_history = mouse_batch self.button_history = button_batch return - self.frame_history = torch.cat([self.frame_history, frame_batch], dim=1) - self.mouse_history = torch.cat([self.mouse_history, mouse_batch], dim=1) - self.button_history = torch.cat([self.button_history, button_batch], dim=1) + self.latent_history = torch.cat([self.latent_history, frame_batch], dim=0) + self.mouse_history = torch.cat([self.mouse_history, mouse_batch], dim=0) + self.button_history = torch.cat([self.button_history, button_batch], dim=0) # cap this at around 60 frames - if self.frame_history.shape[1] > self.streaming_config.window_length: - self.frame_history = self.frame_history[:, -self.streaming_config.window_length:] - self.mouse_history = self.mouse_history[:, -self.streaming_config.window_length:] - self.button_history = self.button_history[:, -self.streaming_config.window_length:] + if self.latent_history.shape[0] > self.streaming_config.window_length: + self.latent_history = self.latent_history[-self.streaming_config.window_length:] + self.mouse_history = self.mouse_history [-self.streaming_config.window_length:] + self.button_history = self.button_history[-self.streaming_config.window_length:] + + def get_latent_history_batch(self) -> torch.Tensor: + if self.latent_history.equal(torch.tensor([], device=self.streaming_config.device)): + return self.dummy_batch + + return self.latent_history.unsqueeze(0) @property + @cache def dummy_batch(self) -> torch.Tensor: """Dummy autoencoder latents for the sampler to initialize shapes.""" tokens_h = tokens_w = int(math.sqrt(self.model_config.tokens_per_frame)) @@ -96,19 +114,7 @@ def dummy_batch(self) -> torch.Tensor: self.model_config.channels, tokens_h, tokens_w, device=self.streaming_config.device, dtype=torch.bfloat16) return dummy_frames - - def debug_generate_frame_batch(self, mouse_batch: torch.Tensor, button_batch: torch.Tensor) -> torch.Tensor: - # NOTE Debug mode, just draw overlays on a black video. - num_frames = mouse_batch.shape[1] - video = torch.zeros(num_frames, 256, 256, 3, - device=self.streaming_config.device, - dtype=torch.bfloat16) - frames = _draw_action_overlays(video, button_batch[0, ::], mouse_batch[0, ::]) - if frames == []: - return video - - frames = [torch.from_numpy(frame) for frame in frames] - return torch.stack(frames).permute(0, 3, 1, 2).unsqueeze(0) # [1, num_frames, c, h, w] + def overlay_actions(self, video: torch.Tensor, @@ -126,33 +132,65 @@ def overlay_actions(self, action_video = torch.stack(action_video).permute(0, 3, 1, 2) # [n h w c] -> [n c h w] return action_video + def create_overlay_only(self, + video: torch.Tensor, + mouse: torch.Tensor, button: torch.Tensor, + action_margin_px_height: int = 150) -> torch.Tensor: + num_frames, channels, height, width = video.shape + # Create overlay-only frames with just the action margin height + action_video = torch.zeros((num_frames, action_margin_px_height, width, channels), # [n h w c] + device=self.streaming_config.device, dtype=torch.bfloat16) + action_video_np: list[np.ndarray] = _draw_action_overlays(action_video, button, mouse) + action_video = [torch.from_numpy(frame) for frame in action_video_np] + action_video = torch.stack(action_video).permute(0, 3, 1, 2) # [n h w c] -> [n c h w] + return action_video + async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor, - action_margin_px_height: int = 150) -> torch.Tensor: + action_margin_px_height: int = 150) -> tuple[torch.Tensor, torch.Tensor]: """ - Generate window_length frames, return first frames_per_batch for streaming. + Generate window_length frames, return separate video and overlay frames for streaming. Args: mouse: [window_length, 2] button: [window_length, n_buttons] Returns: - frame_batch: [frames_per_batch, 3, 256, 256] - only streaming frames + tuple: (video_frames, overlay_frames) + video_frames: [frames_per_batch, 3, 256, 256] - pure video frames + overlay_frames: [frames_per_batch, 3, action_margin_px_height, 256] - action overlay frames """ if self.debug: num_frames = mouse.shape[0] - full_frames = torch.randn(num_frames, 3, 256, 256, - device=self.streaming_config.device, dtype=torch.bfloat16) + # Create gradient from white to black to white across columns + col_indices = torch.arange(256, device=self.streaming_config.device) + # Create gradient that goes from 1 to 0 to 1 + gradient = torch.where( + col_indices < 128, + 1.0 - (col_indices / 127.0), # First half: 1 to 0 + (col_indices - 128) / 127.0 # Second half: 0 to 1 + ).view(1, 1, 1, -1) # [1, 1, 1, 256] + full_frames = gradient.expand(num_frames, 3, 256, 256).to(torch.bfloat16) + # to between 0 and 1 + full_frames = (full_frames - full_frames.min()) / (full_frames.max() - full_frames.min()) + # between 0 and 255 + full_frames = (full_frames * 255).to(torch.uint8) else: with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): - latents, full_frames = self.sample_window_fn(dummy_batch=self.dummy_batch, - mouse=mouse.unsqueeze(0), btn=button.unsqueeze(0)) # [1, window_length, 3, 256, 256] - # remove batch dimension - latents = latents [0, ::] - full_frames = full_frames [0, ::] - self.add_to_history(latents, mouse, button) # ignore batch dimension (1 user per model) - - action_overlayed_frames = self.overlay_actions(full_frames, mouse, button, action_margin_px_height) - return action_overlayed_frames # [window_length, 3, 256+action_margin_px_height, 256] + latents, full_frames = self.sample_window_fn(dummy_batch=self.get_latent_history_batch(), + mouse=mouse.float().unsqueeze(0), + btn=button.float().unsqueeze(0)) # [1, window_length, 3, 256, 256] + # remove batch dimension, then take only the frames we generated, since the WindowCFGSampler appends the history (which is of window_length=60) + latents = latents [0, -self.sampling_config.num_frames:] + full_frames = full_frames [0, -self.sampling_config.num_frames:] + # then, convert the frames to a pixel-range of [0-255] from [-1,1] + full_frames = (full_frames + 1) / 2 + full_frames = (full_frames * 255).to(torch.uint8) + self.add_to_history(latents, mouse, button) + + # Create overlay frames separately + overlay_frames = self.create_overlay_only(full_frames, mouse, button, action_margin_px_height) + + return full_frames, overlay_frames # [window_length, 3, 256, 256], [window_length, 3, action_margin_px_height, 256] def __enter__(self): return self diff --git a/webapp/user_session.py b/webapp/user_session.py index d99cbc25..59fd9515 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -56,10 +56,10 @@ async def _frame_generation_loop(self): mouse, button = await self.action_collector.collect_actions() # Generate Y frames from X actions by taking the X[-1]'th action. Typically, X >> Y, because they are sampled at uncapped FPS from the UI, # whereas Y frames are sampled from the model one at a time. - frames = await self.frame_generator.generate_frames(mouse, button) + video_frames, overlay_frames = await self.frame_generator.generate_frames(mouse, button) # Queue frames for streaming at a capped FPS. If model predictions speed up or slow down, it won't cause any dilation of frames being displayed. # However, if the model predictions are too slow, the frames will be displayed at a lower FPS than the capped FPS. - await self.frame_buffer.queue_frames(frames) + await self.frame_buffer.queue_frames(video_frames, mouse, button) except Exception as e: import traceback print(termcolor.colored(f"Error in frame generation: {e} :\n {traceback.format_exc()}", "red")) @@ -68,8 +68,8 @@ async def _frame_generation_loop(self): async def _frame_display_loop(self, websocket: WebSocket): while True: try: - frame = await self.frame_buffer.get_next_frame() - await self._send_frame_to_client(websocket, frame) + video_frame, button, mouse = await self.frame_buffer.get_next_frames() + await self._send_frames_to_client(websocket, video_frame, button, mouse) except Exception as e: # Check if this is a WebSocket disconnect if "websocket.close" in str(e) or "response already completed" in str(e): @@ -80,22 +80,31 @@ async def _frame_display_loop(self, websocket: WebSocket): print(termcolor.colored(f"Error in frame streaming: {e} :\n {traceback.format_exc()}", "red")) await asyncio.sleep(self.frame_generator.streaming_config.frame_interval) - async def _send_frame_to_client(self, websocket: WebSocket, frame: torch.Tensor): + async def _send_frames_to_client(self, + websocket: WebSocket, + video_frame: torch.Tensor, + button: torch.Tensor, + mouse: torch.Tensor): try: # TODO Do this more intelligently. I'm sure there's better tech to stream video to a UI. - # Convert frame to base64 JPEG\ - frame_np = frame.cpu().numpy().transpose(1, 2, 0) # CHW -> HWC - # NOTE: Assumes -1 to 1 range for frames. - multiplier = 127.5 if frame.max() < 1 else 1 - frame_np = (frame_np * multiplier).clip(0, 255).astype(np.uint8) # Normalize + # Convert video frame to base64 JPEG + video_frame_np = video_frame.float().cpu().numpy().transpose(1, 2, 0) # CHW -> HWC + # Normalize frame data to 0-255 range + if video_frame_np.max() <= 1.0: + video_frame_np = (video_frame_np * 255).clip(0, 255).astype(np.uint8) + else: + video_frame_np = video_frame_np.clip(0, 255).astype(np.uint8) + + _, video_buffer = cv2.imencode('.jpg', video_frame_np) + video_base64 = base64.b64encode(video_buffer).decode('utf-8') - _, buffer = cv2.imencode('.jpg', frame_np) - frame_base64 = base64.b64encode(buffer).decode('utf-8') await websocket.send_json({ - "type": "frame", - "data": frame_base64, - "timestamp": time.time() + "type": "frame", + "video_data": video_base64, + "button_data": self.action_collector.converter.buttons_to_dict(button), + "mouse_data": self.action_collector.converter.mouse_to_dict(mouse), + "timestamp": time.time() }) except Exception as e: # Re-raise to be caught by the display loop diff --git a/webapp/utils/configs.py b/webapp/utils/configs.py index aeb8cd29..4cda5c58 100644 --- a/webapp/utils/configs.py +++ b/webapp/utils/configs.py @@ -31,10 +31,12 @@ def from_yaml(cls, path: os.PathLike) -> WebappConfig: @dataclass class SamplingConfig: sampling_steps : int = 20 + vae_scale: float = 1.0 cfg_scale : float = 1.3 + window_length : int = 60 num_frames : int = 60 noise_prev : float = 0.2 - window_length : Optional[int] = 60 + @dataclass class StreamingConfig: diff --git a/webapp/utils/samplers.py b/webapp/utils/samplers.py index 63e1a437..7dad32c8 100644 --- a/webapp/utils/samplers.py +++ b/webapp/utils/samplers.py @@ -1,4 +1,4 @@ -from torch import nn, Tensor +from torch import Tensor from typing import Literal, Callable from functools import partial, cache from multimethod import multimethod @@ -19,7 +19,8 @@ def create_sampler(sampler_id: Literal['cfg'], encoder, decoder, batch_size: int = 8, sampling_steps: int = SAMPLING_STEPS, cfg_scale: float = CFG_SCALE, - scale: float = SCALE) -> Callable[[MouseData, ButtonData], + vae_scale: float = SCALE, + **kwargs) -> Callable[[MouseData, ButtonData], tuple[LatentData, VideoData]]: @cache # simple singleton def _sampler(): return CFGSampler() @@ -28,7 +29,7 @@ def _sampler(): return CFGSampler() _sampler().__call__, sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=scale, + scale=vae_scale, cfg_scale=cfg_scale, model=encoder ) @@ -38,7 +39,8 @@ def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder, decoder, batch_size: int = 8, sampling_steps: int = SAMPLING_STEPS, cfg_scale: float = CFG_SCALE, - scale: float = SCALE) -> Callable[[MouseData, ButtonData], + vae_scale: float = SCALE, + **kwargs) -> Callable[[MouseData, ButtonData], tuple[LatentData, VideoData]]: @cache def _sampler(): return InpaintCFGSampler() @@ -47,7 +49,7 @@ def _sampler(): return InpaintCFGSampler() _sampler().__call__, sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=scale, + scale=vae_scale, cfg_scale=cfg_scale, model=encoder ) @@ -57,7 +59,8 @@ def create_sampler(sampler_id: Literal['window'], encoder, decoder, batch_size: int = 8, sampling_steps: int = SAMPLING_STEPS, cfg_scale: float = CFG_SCALE, - scale: float = SCALE) -> Callable[[MouseData, ButtonData], + vae_scale: float = SCALE, + **kwargs) -> Callable[[MouseData, ButtonData], tuple[LatentData, VideoData]]: @cache def _sampler(): return WindowCFGSampler() @@ -66,7 +69,8 @@ def _sampler(): return WindowCFGSampler() _sampler().__call__, sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=scale, + scale=vae_scale, cfg_scale=cfg_scale, - model=encoder + model=encoder, + **kwargs ) diff --git a/webapp/utils/visualize_overlay_actions.py b/webapp/utils/visualize_overlay_actions.py index 3005d670..03903b3f 100644 --- a/webapp/utils/visualize_overlay_actions.py +++ b/webapp/utils/visualize_overlay_actions.py @@ -9,7 +9,7 @@ # Global configuration KEYBINDS = ["W", "A", "S", "D", "LSHIFT", "SPACE", "R", "F", "E", "LMB", "RMB"] -MINIMUM_FRAME_SIZE = 512 +MINIMUM_FRAME_SIZE = 150 # Colors (BGR format for OpenCV) COLOR_PRESSED = (50, 200, 50) # Green COLOR_UNPRESSED = (100, 100, 100) # Gray @@ -21,13 +21,13 @@ COLOR_RMB_SECTOR = COLOR_PRESSED # Green # Key dimensions -KEY_SIZE = 30 +KEY_SIZE = 20 KEY_MARGIN = 5 SHIFT_WIDTH = int(KEY_SIZE * 2 + KEY_MARGIN) # Two keys worth of width SPACE_WIDTH = int(KEY_SIZE * 5) # Mouse compass dimensions -COMPASS_RADIUS = 48 # 80% of 80 +COMPASS_RADIUS = 24 # 80% of 80 COMPASS_START_X_PERCENT = 0.80 # Mouse button arc dimensions @@ -39,8 +39,8 @@ ARROW_MIN_LENGTH = 55 # Minimum arrow length in pixels ARROW_MAX_SCALE = 0.75 # Maximum scale relative to compass radius -START_X_PERCENT = 0.12 -START_Y_PERCENT = 0.65 +START_X_PERCENT = 0.18 +START_Y_PERCENT = 0.85 @contextmanager @@ -59,6 +59,9 @@ def _rescale_icons(ratio: float): "SHIFT_WIDTH": SHIFT_WIDTH, "SPACE_WIDTH": SPACE_WIDTH, "COMPASS_RADIUS": COMPASS_RADIUS, + "MOUSE_BUTTON_OFFSET": MOUSE_BUTTON_OFFSET, + "MOUSE_BUTTON_THICKNESS": MOUSE_BUTTON_THICKNESS, + "ARROW_SCALE_FACTOR": ARROW_SCALE_FACTOR, } KEY_SIZE *= ratio ; KEY_SIZE = int(KEY_SIZE) @@ -69,30 +72,64 @@ def _rescale_icons(ratio: float): MOUSE_BUTTON_OFFSET *= ratio ; MOUSE_BUTTON_OFFSET = int(MOUSE_BUTTON_OFFSET) MOUSE_BUTTON_THICKNESS *= ratio ; MOUSE_BUTTON_THICKNESS = int(MOUSE_BUTTON_THICKNESS) ARROW_SCALE_FACTOR *= ratio ; ARROW_SCALE_FACTOR = float(ARROW_SCALE_FACTOR) - ARROW_MIN_LENGTH *= ratio ; ARROW_MIN_LENGTH = int(ARROW_MIN_LENGTH) - ARROW_MAX_SCALE *= ratio ; ARROW_MAX_SCALE = float(ARROW_MAX_SCALE) + # Note: ARROW_MIN_LENGTH and ARROW_MAX_SCALE are now calculated relative to COMPASS_RADIUS + # so they don't need separate scaling yield finally: for key, value in old_values.items(): globals()[key] = value +def _get_adaptive_positioning(frame_width: int, frame_height: int) -> tuple[float, float, float]: + """ + Calculate adaptive positioning percentages based on frame dimensions. + Returns (keyboard_start_y_percent, mouse_start_y_percent, start_x_percent) + """ + aspect_ratio = frame_width / frame_height + + # Calculate how much vertical space the keyboard needs + keyboard_height_needed = KEY_SIZE * 3 + KEY_MARGIN * 2 + 20 # 3 rows + margins + buffer + + # Adaptive Y positioning - ensure keyboard fits + if frame_height <= keyboard_height_needed + 40: # Very short frame + keyboard_start_y_percent = 0.1 # Start near top + mouse_start_y_percent = 0.6 # Place mouse in middle-bottom + elif frame_height < 300: # Short frame + keyboard_start_y_percent = 0.4 + mouse_start_y_percent = 0.75 + else: # Normal/tall frame + keyboard_start_y_percent = 0.75 + mouse_start_y_percent = 0.85 + + # Adaptive X positioning based on aspect ratio + if aspect_ratio > 1.5: # Wide frame + start_x_percent = 0.08 + elif aspect_ratio < 0.75: # Tall frame + start_x_percent = 0.15 + else: # Near square frame + start_x_percent = 0.12 + + return keyboard_start_y_percent, mouse_start_y_percent, start_x_percent + + + def _draw_buttons( frame: np.ndarray, button_sequence: list[bool], ) -> None: """ - Draw keyboard buttons on the frame. - - Args: - frame: numpy array representing the image frame - button_sequence: list of bools corresponding to KEYBINDS (excluding LMB, RMB) + Draw keyboard buttons on the frame with adaptive positioning. """ - # Starting position for keyboard layout (bottom left) - start_x = int(frame.shape[1] * START_X_PERCENT) - start_y = int(frame.shape[0] * START_Y_PERCENT) + frame_height, frame_width = frame.shape[:2] + + # Get adaptive positioning + keyboard_start_y_percent, _, start_x_percent = _get_adaptive_positioning(frame_width, frame_height) - # Key positions - organized by rows + # Starting position for keyboard layout + start_x = int(frame_width * start_x_percent) + start_y = int(frame_height * keyboard_start_y_percent) + + # Rest of the function remains the same... key_positions = { # Top row: W E R (W above S) "W": (start_x + (KEY_SIZE + KEY_MARGIN) * 1, start_y), @@ -110,11 +147,15 @@ def _draw_buttons( "SPACE": (start_x + SHIFT_WIDTH + KEY_MARGIN, start_y + (KEY_SIZE + KEY_MARGIN) * 2), } - # Draw each key + # Draw each key (rest remains the same) for i, key in enumerate(KEYBINDS[:-2]): # Exclude LMB and RMB if key in key_positions: x, y = key_positions[key] + # Ensure keys stay within frame bounds + if x < 0 or y < 0 or x + KEY_SIZE > frame_width or y + KEY_SIZE > frame_height: + continue + # Determine key dimensions if key == "LSHIFT": width = SHIFT_WIDTH @@ -126,6 +167,10 @@ def _draw_buttons( width = KEY_SIZE height = KEY_SIZE + # Final bounds check with actual key dimensions + if x + width > frame_width or y + height > frame_height: + continue + # Determine color based on pressed state color = COLOR_PRESSED if button_sequence[i] else COLOR_UNPRESSED @@ -142,7 +187,6 @@ def _draw_buttons( cv2.putText(frame, key, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_TEXT, 1, cv2.LINE_AA) - def _draw_mouse( frame: np.ndarray, LMB_on: bool, @@ -151,15 +195,14 @@ def _draw_mouse( center: tuple[int, int], ) -> None: """ - Draw mouse compass with direction arrow and uncertainty cone. + Draw mouse compass with direction arrow and properly positioned labels. Arrow length is proportional to mouse movement magnitude. Args: frame: numpy array representing the image frame LMB_on: bool indicating if left mouse button is pressed RMB_on: bool indicating if right mouse button is pressed - mouse_vec: tuple of floats (x, y) representing mouse direction - mouse_std: tuple of floats (x_std, y_std) representing uncertainty + mouse_delta: tuple of floats (x, y) representing mouse direction center: tuple (x, y) for compass center position """ # Draw compass circle @@ -178,10 +221,14 @@ def _draw_mouse( cv2.ellipse(frame, center, (button_radius, button_radius), 0, 270, 315, color_rmb, MOUSE_BUTTON_THICKNESS) - text_loc_lmb = center[0] - COMPASS_RADIUS + 10, center[1] - COMPASS_RADIUS - 10 - cv2.putText(frame, 'LMB', text_loc_lmb, cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_TEXT, 1, cv2.LINE_AA) - text_loc_rmb = center[0], center[1] - COMPASS_RADIUS - 10 - cv2.putText(frame, 'RMB', text_loc_rmb, cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_TEXT, 1, cv2.LINE_AA) + # Fixed label positioning - spread them out more and position them better + # LMB label - position it to the left of the compass + text_loc_lmb = (center[0] - COMPASS_RADIUS - 15, center[1] - COMPASS_RADIUS + 5) + cv2.putText(frame, 'LMB', text_loc_lmb, cv2.FONT_HERSHEY_SIMPLEX, 0.4, COLOR_TEXT, 1, cv2.LINE_AA) + + # RMB label - position it to the right of the compass + text_loc_rmb = (center[0] + COMPASS_RADIUS - 10, center[1] - COMPASS_RADIUS + 5) + cv2.putText(frame, 'RMB', text_loc_rmb, cv2.FONT_HERSHEY_SIMPLEX, 0.4, COLOR_TEXT, 1, cv2.LINE_AA) # Calculate mouse direction mouse_x, mouse_y = mouse_delta @@ -192,13 +239,17 @@ def _draw_mouse( unit_x = mouse_x / magnitude unit_y = mouse_y / magnitude - # Calculate arrow length proportional to magnitude + # Fixed arrow scaling - ensure arrow stays within compass bounds + # Calculate minimum and maximum arrow lengths relative to compass size + min_arrow_length = COMPASS_RADIUS * 0.3 # 30% of compass radius + max_arrow_length = COMPASS_RADIUS * 0.8 # 80% of compass radius (stays within circle) + # Scale by magnitude but clamp to reasonable bounds arrow_scale = max( - ARROW_MIN_LENGTH, # Ensure minimum visible arrow + min_arrow_length, # Minimum visible arrow min( COMPASS_RADIUS * magnitude * ARROW_SCALE_FACTOR, # Proportional to magnitude - COMPASS_RADIUS * ARROW_MAX_SCALE # Cap at maximum scale + max_arrow_length # Cap at 80% of compass radius ) ) @@ -206,7 +257,7 @@ def _draw_mouse( end_y = int(center[1] - unit_y * arrow_scale) # Negative because y-axis is inverted # Draw direction arrow with proportional length - cv2.arrowedLine(frame, center, (end_x, end_y), COLOR_MOUSE_ARROW, 3, tipLength=0.3) + cv2.arrowedLine(frame, center, (end_x, end_y), COLOR_MOUSE_ARROW, 2, tipLength=0.4) # Optional: Display magnitude as text for debugging # magnitude_text = f"Mag: {magnitude:.2f}" @@ -214,11 +265,12 @@ def _draw_mouse( # cv2.FONT_HERSHEY_SIMPLEX, 0.4, COLOR_TEXT, 1, cv2.LINE_AA) # Draw center dot - cv2.circle(frame, center, 3, COLOR_TEXT, -1) + cv2.circle(frame, center, 2, COLOR_TEXT, -1) - # Add labels - cv2.putText(frame, "Mouse", (center[0] - 25, center[1] + button_radius + 20), - cv2.FONT_HERSHEY_SIMPLEX, 0.6, COLOR_TEXT, 1, cv2.LINE_AA) + # Add "Mouse" label below the compass + cv2.putText(frame, "Mouse", (center[0] - 20, center[1] + button_radius + 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLOR_TEXT, 1, cv2.LINE_AA) + def _draw_frame( frame: np.ndarray, @@ -226,23 +278,28 @@ def _draw_frame( mouse_delta: tuple[float, float], ) -> np.ndarray: """ - Overlay keyboard & mouse on a single frame and return it. + Overlay keyboard & mouse on a single frame with adaptive positioning. """ buttons, lmb, rmb = buttons[:-2], buttons[-2], buttons[-1] _draw_buttons(frame, buttons) frame_height, frame_width = frame.shape[:2] - margin = 30 - compass_x = frame_width - margin - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET - compass_y = int(frame_height * START_Y_PERCENT) - compass_x = int(compass_x * COMPASS_START_X_PERCENT) + + # Get adaptive positioning + _, mouse_start_y_percent, _ = _get_adaptive_positioning(frame_width, frame_height) + + # Calculate compass position + compass_x_percent = COMPASS_START_X_PERCENT + compass_x = int(frame_width * compass_x_percent) - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET + compass_y = int(frame_height * mouse_start_y_percent) # Ensure compass fits within frame boundaries min_x = COMPASS_RADIUS + MOUSE_BUTTON_OFFSET + 20 max_x = frame_width - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET - 20 compass_x = max(min_x, min(compass_x, max_x)) + min_y = COMPASS_RADIUS + MOUSE_BUTTON_OFFSET + 30 - max_y = frame_height - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET - 30 + max_y = frame_height - COMPASS_RADIUS - MOUSE_BUTTON_OFFSET - 50 compass_y = max(min_y, min(compass_y, max_y)) _draw_mouse(frame, @@ -280,13 +337,25 @@ def _draw_video( ARROW_MAX_SCALE = arrow_max_scale video = video.float().cpu().numpy() - # upscale image to 512x512 only if < 512 on smallest dimension. - if video.shape[1] < MINIMUM_FRAME_SIZE or video.shape[2] < MINIMUM_FRAME_SIZE: - frames = [cv2.resize(frame, (512, 512), interpolation=cv2.INTER_CUBIC) for frame in video] + + # Get original dimensions + original_height, original_width = video.shape[1], video.shape[2] + + # Calculate scaling to meet minimum size while preserving aspect ratio + if original_height < MINIMUM_FRAME_SIZE or original_width < MINIMUM_FRAME_SIZE: + # Calculate scale factor to make smallest dimension equal to MINIMUM_FRAME_SIZE + scale_factor = MINIMUM_FRAME_SIZE / min(original_height, original_width) + new_height = int(original_height * scale_factor) + new_width = int(original_width * scale_factor) + + frames = [cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_CUBIC) for frame in video] + # UI scaling ratio - keep icons at base size since we're scaling up small frames ratio = 1.0 else: frames = video - ratio = min(video.shape[1] / MINIMUM_FRAME_SIZE, video.shape[2] / MINIMUM_FRAME_SIZE) + # UI scaling ratio - scale icons proportional to how much larger the frame is than minimum + # This ensures icons don't become too large on big frames + ratio = min(original_height, original_width) / MINIMUM_FRAME_SIZE with _rescale_icons(ratio): frames = [ @@ -294,7 +363,7 @@ def _draw_video( for i, frame in enumerate(frames) ] if save_path is not None: - imio.mimsave(save_path, frames, fps=fps) + imio.mimsave(save_path, [f.astype(np.uint8) for f in frames], fps=fps) return frames From 32f4358bee99c9005d2d0d9393a4ae35b77b4ca8 Mon Sep 17 00:00:00 2001 From: Sami Date: Mon, 9 Jun 2025 05:10:26 -0400 Subject: [PATCH 14/71] fixes to parseargs --- Dockerfile | 18 ++++++----- webapp/server.py | 73 +++++++++++++++++++++++++----------------- webapp/user_session.py | 14 +++++++- 3 files changed, 67 insertions(+), 38 deletions(-) diff --git a/Dockerfile b/Dockerfile index 4be55f44..e3280b20 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,20 +6,22 @@ ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV PYTHONDONTWRITEBYTECODE=1 -# Install system dependencies +# Install system dependencies (without python3.12 first) RUN apt-get update && apt-get install -y --no-install-recommends \ wget \ curl \ python3-pip \ - python3.12 \ - python3.12-dev \ - python3.12-distutils \ - python3.12-venv \ git \ software-properties-common \ - && add-apt-repository ppa:deadsnakes/ppa \ + && rm -rf /var/lib/apt/lists/* + +# Add deadsnakes PPA and install Python 3.12 +RUN add-apt-repository ppa:deadsnakes/ppa \ && apt-get update \ - && apt-get install -y python3.12 python3.12-dev python3.12-distutils \ + && apt-get install -y --no-install-recommends \ + python3.12 \ + python3.12-dev \ + python3.12-venv \ && rm -rf /var/lib/apt/lists/* # Set Python 3.12 as default @@ -49,4 +51,4 @@ COPY . /app EXPOSE 8000 # Set the default command to run the web server -CMD ["python3", "webapp/server.py", "--port", "8000", "--no-debug"] \ No newline at end of file +CMD ["python3", "webapp/server.py", "--port", "8000", "--no-debug"] diff --git a/webapp/server.py b/webapp/server.py index cda23317..6cd04212 100644 --- a/webapp/server.py +++ b/webapp/server.py @@ -12,7 +12,7 @@ from webapp.utils.configs import WebappConfig -DEBUG = True # for funsies +DEBUG = True # -- lifespan encoder: nn.Module = None @@ -23,7 +23,7 @@ @asynccontextmanager async def lifespan(app: FastAPI): - global encoder, decoder, config + global encoder, decoder, config, DEBUG config = WebappConfig.from_yaml(webapp_config_path) if not DEBUG: encoder, decoder, _ = load_models( @@ -36,29 +36,32 @@ async def lifespan(app: FastAPI): encoder, decoder, config = None, None, None -app = FastAPI(lifespan=lifespan) -app.mount("/assets", StaticFiles(directory="webapp/static"), name="assets") - - -@app.get("/") -async def read_root(): - """Serve the main game page.""" - return FileResponse("webapp/static/index.html") - - -@app.websocket("/ws/game") -async def websocket_endpoint(websocket: WebSocket): - await websocket.accept() +def run(): + """Create and configure the FastAPI app with routes.""" + app = FastAPI(lifespan=lifespan) + + @app.get("/") + async def read_root(): + """Serve the main game page.""" + return FileResponse("webapp/static/index.html") + + @app.websocket("/ws/game") + async def websocket_endpoint(websocket: WebSocket): + global DEBUG + await websocket.accept() + + # Create streaming session for this user + frame_generator = StreamingFrameGenerator(encoder, decoder, + streaming_config=config.stream_config, + model_config=config.run_config.model, + train_config=config.run_config.train, + sampling_config=config.sampling_config, + debug=DEBUG) + session = UserGameSession(frame_generator) + await session.run_session(websocket) - # Create streaming session for this user - frame_generator = StreamingFrameGenerator(encoder, decoder, - streaming_config=config.stream_config, - model_config=config.run_config.model, - train_config=config.run_config.train, - sampling_config=config.sampling_config, - debug=DEBUG) - session = UserGameSession(frame_generator) - await session.run_session(websocket) + app.mount("/assets", StaticFiles(directory="webapp/static"), name="assets") + return app def main(): @@ -72,19 +75,31 @@ def main(): parser.add_argument("--port", type=int, default=8000, help="Port to run the server on") args = parser.parse_args() - assert not (args.debug and args.no_debug), "Cannot have both debug and no-debug flags" - DEBUG = args.debug or not args.no_debug + # Fix the DEBUG logic + if args.debug and args.no_debug: + raise ValueError("Cannot have both --debug and --no-debug flags") + + if args.debug: + DEBUG = True + elif args.no_debug: + DEBUG = False + # Otherwise keep the default value (True) + + # Create app AFTER setting DEBUG + app = run() print("🚀 Starting OWL-WMS FastAPI Server...") print("📡 WebSocket endpoint: ws://localhost:8000/ws/game") print("🌐 Access via: http://localhost:8000") print("🔄 Auto-reload enabled for development") - + print("🔄 DEBUG is set to:", DEBUG) + print("🔄 PORT is set to:", args.port) + uvicorn.run( - "webapp.server:app", + app, # Pass the app object directly instead of module string host="0.0.0.0", # Allow external connections port=args.port, - reload=True, # Auto-reload on file changes + reload=False, # Can't use reload with app object log_level="info" ) diff --git a/webapp/user_session.py b/webapp/user_session.py index 59fd9515..c96f709a 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -68,11 +68,19 @@ async def _frame_generation_loop(self): async def _frame_display_loop(self, websocket: WebSocket): while True: try: + # Check if WebSocket is still connected before processing + if websocket.client_state.name != 'CONNECTED': + print(termcolor.colored("🔌 WebSocket no longer connected - stopping frame stream", "yellow")) + break + video_frame, button, mouse = await self.frame_buffer.get_next_frames() await self._send_frames_to_client(websocket, video_frame, button, mouse) except Exception as e: # Check if this is a WebSocket disconnect - if "websocket.close" in str(e) or "response already completed" in str(e): + if ("websocket.close" in str(e) or + "response already completed" in str(e) or + "Cannot call \"send\" once a close message has been sent" in str(e) or + "RuntimeError" in str(e)): print(termcolor.colored("🔌 WebSocket disconnected - stopping frame stream", "yellow")) break else: @@ -86,6 +94,10 @@ async def _send_frames_to_client(self, button: torch.Tensor, mouse: torch.Tensor): try: + # Check WebSocket state before sending + if websocket.client_state.name != 'CONNECTED': + raise RuntimeError("WebSocket is not connected") + # TODO Do this more intelligently. I'm sure there's better tech to stream video to a UI. # Convert video frame to base64 JPEG video_frame_np = video_frame.float().cpu().numpy().transpose(1, 2, 0) # CHW -> HWC From 2cc32188cba8ee8f3539a313576a591c1717409b Mon Sep 17 00:00:00 2001 From: shahbuland Date: Mon, 9 Jun 2025 21:00:02 +0000 Subject: [PATCH 15/71] add causvid trainer --- configs/causvid.yml | 69 +++++++++ owl_wms/models/gamerft.py | 19 ++- owl_wms/trainers/__init__.py | 5 +- owl_wms/trainers/causvid.py | 270 +++++++++++++++++++++++++++++++++++ 4 files changed, 358 insertions(+), 5 deletions(-) create mode 100644 configs/causvid.yml create mode 100644 owl_wms/trainers/causvid.py diff --git a/configs/causvid.yml b/configs/causvid.yml new file mode 100644 index 00000000..065a7f9b --- /dev/null +++ b/configs/causvid.yml @@ -0,0 +1,69 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: game_rft + sample_size: 4 + channels: 128 + + n_layers: 17 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 16 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.0 + n_frames: 30 + + causal: false + +train: + trainer_id: causvid + data_id: cod_latent + data_kwargs: + window_length: 30 + root: ../cod_data/BlackOpsColdWar + add_optical_flow: false + + target_batch_size: 256 + batch_size: 32 + + epochs: 200 + + opt: AdamW + opt_kwargs: + lr: 2.0e-6 + weight_decay: 1.0e-4 + eps: 1.0e-15 + betas: [0.9, 0.95] + + scheduler: null + + checkpoint_dir: checkpoints/360p + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: window + sampler_kwargs: + n_steps: 20 + cfg_scale: 1.3 + window_length: 30 + num_frames: 60 + noise_prev: 0.2 + only_return_generated: true + + n_samples: 8 + + vae_id: 720pr3dc + vae_batch_size: 4 + vae_scale: 0.35 + vae_cfg_path: configs/owl_vaes/128x_cod_stage2.yml + vae_ckpt_path: 720p_cod_vae_30m_35k_steps.pt + + teacher_ckpt: null # Set later TODO + +wandb: + name: shahbuland + project: video_models + run_name: v2 \ No newline at end of file diff --git a/owl_wms/models/gamerft.py b/owl_wms/models/gamerft.py index ade60112..f5a047fa 100644 --- a/owl_wms/models/gamerft.py +++ b/owl_wms/models/gamerft.py @@ -57,15 +57,17 @@ def __init__(self, config): self.core = GameRFTCore(config) self.cfg_prob = config.cfg_prob - def forward(self, x, mouse, btn): + def forward(self, x, mouse, btn, return_dict = False, cfg_prob = None): # x is [b,n,c,h,w] # mouse is [b,n,2] # btn is [b,n,n_buttons] b,n,c,h,w = x.shape # Apply classifier-free guidance dropout - if self.cfg_prob > 0.: - mask = torch.rand(b, device=x.device) < self.cfg_prob + if cfg_prob is None: + cfg_prob = self.cfg_prob + if cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob null_mouse = torch.zeros_like(mouse) null_btn = torch.zeros_like(btn) @@ -85,7 +87,16 @@ def forward(self, x, mouse, btn): pred = self.core(lerpd, ts, mouse, btn) diff_loss = F.mse_loss(pred, target) - return diff_loss + if not return_dict: + return diff_loss + else: + return { + 'diffusion_loss' : diff_loss, + 'lerpd' : lerpd, + 'pred' : pred, + 'ts': ts, + 'z': z + } if __name__ == "__main__": from ..configs import Config diff --git a/owl_wms/trainers/__init__.py b/owl_wms/trainers/__init__.py index 86d17cba..20d7ac4c 100644 --- a/owl_wms/trainers/__init__.py +++ b/owl_wms/trainers/__init__.py @@ -2,4 +2,7 @@ def get_trainer_cls(trainer_id): if trainer_id == "rft": - return RFTTrainer \ No newline at end of file + return RFTTrainer + if trainer_id == "causvid": + from .causvid import CausVidTrainer + return CausVidTrainer \ No newline at end of file diff --git a/owl_wms/trainers/causvid.py b/owl_wms/trainers/causvid.py new file mode 100644 index 00000000..4c586d1f --- /dev/null +++ b/owl_wms/trainers/causvid.py @@ -0,0 +1,270 @@ +import torch +from ema_pytorch import EMA +import wandb +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist +import einops as eo +from copy import deepcopy + +from .base import BaseTrainer + +from ..utils import freeze, unfreeze, Timer, find_unused_params, versatile_load +from ..schedulers import get_scheduler_cls +from ..models import get_model_cls +from ..sampling import get_sampler_cls +from ..data import get_loader +from ..utils.logging import LogHelper, to_wandb +from ..muon import init_muon +from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn + +class CausVidTrainer(BaseTrainer): + """ + CausVid Trainer + + :param train_cfg: Configuration for training + :param logging_cfg: Configuration for logging + :param model_cfg: Configuration for model + :param global_rank: Rank across all devices. + :param local_rank: Rank for current device on this process. + :param world_size: Overall number of devices + """ + def __init__(self,*args,**kwargs): + super().__init__(*args,**kwargs) + + model_id = self.model_cfg.model_id + + student_cfg = deepcopy(self.model_cfg) + teacher_cfg = deepcopy(self.model_cfg) + + student_cfg.causal = True + teacher_cfg.causal = False + + self.model = get_model_cls(model_id)(student_cfg) + self.score_real = get_model_cls(model_id)(teacher_cfg) + + self.score_real.load_state_dict(versatile_load(self.train_cfg.teacher_ckpt)) + self.score_fake = deepcopy(self.score_real) + + freeze(self.score_real) + + # Print model size + if self.rank == 0: + n_params = sum(p.numel() for p in self.model.parameters()) + print(f"Model has {n_params:,} parameters") + + self.ema = None + self.opt = None + self.s_fake_opt = None + self.scheduler = None + self.s_fake_scaler = None + self.scaler = None + + self.total_step_counter = 0 + self.decoder = get_decoder_only() + freeze(self.decoder) + + def save(self): + save_dict = { + 'model' : self.model.state_dict(), + 'ema' : self.ema.state_dict(), + 'opt' : self.opt.state_dict(), + 'scaler' : self.scaler.state_dict(), + 'score_fake': self.score_fake.state_dict(), + 's_fake_opt': self.s_fake_opt.state_dict(), + 's_fake_scaler': self.s_fake_scaler.state_dict(), + 'steps': self.total_step_counter + } + if self.scheduler is not None: + save_dict['scheduler'] = self.scheduler.state_dict() + super().save(save_dict) + + def load(self): + has_ckpt = False + try: + if self.train_cfg.resume_ckpt is not None: + save_dict = super().load(self.train_cfg.resume_ckpt) + has_ckpt = True + except: + print("Error loading checkpoint") + + if not has_ckpt: + return + + + self.model.load_state_dict(save_dict['model']) + self.ema.load_state_dict(save_dict['ema']) + self.opt.load_state_dict(save_dict['opt']) + if self.scheduler is not None and 'scheduler' in save_dict: + self.scheduler.load_state_dict(save_dict['scheduler']) + self.scaler.load_state_dict(save_dict['scaler']) + self.score_fake.load_state_dict(save_dict['score_fake']) + self.s_fake_opt.load_state_dict(save_dict['s_fake_opt']) + self.s_fake_scaler.load_state_dict(save_dict['s_fake_scaler']) + self.total_step_counter = save_dict['steps'] + + def train(self): + torch.cuda.set_device(self.local_rank) + + # Prepare model and ema + self.model = self.model.cuda().train() + self.decoder = self.decoder.cuda().eval().bfloat16() + self.score_real = self.score_real.cuda().eval().bfloat16() + self.score_fake = self.score_fake.cuda().train() + + if self.world_size > 1: + self.model = DDP(self.model) + self.score_fake = DDP(self.score_fake) + + freeze(self.decoder) + freeze(self.score_real) + + #torch.compile(self.score_real, dynamic = False) + + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + + self.ema = EMA( + self.model, + beta = 0.999, + update_after_step = 0, + update_every = 1 + ) + # Hard coded stuff, probably #TODO figure out where to put this? + self.update_ratio = 5 + self.cfg_scale = 1.3 + + def get_ema_core(): + if self.world_size > 1: + return self.ema.ema_model.module.core + else: + return self.ema.ema_model.core + + # Don't use MUON pls + self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) + self.s_fake_opt = getattr(torch.optim, self.train_cfg.opt)(self.score_fake.parameters(), **self.train_cfg.opt_kwargs) + + if self.train_cfg.scheduler is not None: + self.scheduler = get_scheduler_cls(self.train_cfg.scheduler)(self.opt, **self.train_cfg.scheduler_kwargs) + + # Scaler + self.s_fake_scaler = torch.amp.GradScaler() + self.scaler = torch.amp.GradScaler() + ctx = torch.amp.autocast('cuda',torch.bfloat16) + + self.load() + + # Timer reset + timer = Timer() + timer.reset() + metrics = LogHelper() + if self.rank == 0: + wandb.watch(self.get_module(), log = 'all') + + # Dataset setup + loader = get_loader(self.train_cfg.data_id, self.train_cfg.batch_size, **self.train_cfg.data_kwargs) + sampler = get_sampler_cls(self.train_cfg.sampler_id)() + + # Simplifiying assumptions: data will never stop iter, no grad accum + + def sample_from_gen(vid, mouse, btn): + model_out = self.model(vid, mouse, btn, return_dict = True) + ts = model_out['ts'][:,None,None,None] # [b,n,c,h,w] + lerpd = model_out['lerpd'] # [b,n,c,h,w] + pred = model_out['pred'] # [b,n,c,h,w] + + samples = lerpd - pred*ts + return samples + + def get_dmd_loss(vid, mouse, btn): + s_real_fn = self.score_real.core + s_fake_fn = self.score_fake.module.core + + with torch.no_grad(): + b,n,c,h,w = vid.shape + ts = torch.randn(b,n,device=vid.device,dtype=vid.dtype).sigmoid() + z = torch.randn_like(vid) + ts_exp = ts[:,:,None,None,None] + lerpd = vid * (1. - ts_exp) + z * ts_exp + + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + s_real_uncond = s_real_fn(lerpd, ts, null_mouse, null_btn) + s_real_cond = s_real_fn(lerpd, ts, mouse, btn) + s_real = s_real_uncond + self.cfg_scale * (s_real_cond - s_real_uncond) + + s_fake = s_fake_fn(lerpd, ts, mouse, btn) + + grad = (s_fake - s_real) + + # Normalizer? + p_real = (vid - s_real) + normalizer = torch.abs(p_real).mean(dim=[1,2,3,4],keepdim=True) + grad = grad / (normalizer + 1.0e-6) + + grad = torch.nan_to_num(grad) + dmd_loss = 0.5 * F.mse_loss(vid.double(), vid.double() - grad.double()) + # ^ simplify to 0.5 * 2 * (vid - vid + grad) = grad, neat! + return dmd_loss + + def optimizer_step(loss, model, scaler, optimizer): + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + scaler.step(optimizer) + optimizer.zero_grad(set_to_none=True) + scaler.update() + + loader = iter(loader) + while True: + freeze(self.model) + unfreeze(self.score_fake) + for _ in range(self.update_ratio): + batch_vid, batch_mouse, batch_btn = next(loader) + with ctx: + with torch.no_grad(): + samples = sample_from_gen(batch_vid, batch_mouse, batch_btn) + s_fake_loss = self.score_fake(samples, batch_mouse, batch_btn) + + optimizer_step(s_fake_loss, self.score_fake, self.s_fake_scaler, self.s_fake_opt) + + metrics.log('s_fake_loss', s_fake_loss) + unfreeze(self.model) + freeze(self.score_fake) + + batch_vid, batch_mouse, batch_btn = next(loader) + with ctx: + samples = sample_from_gen(batch_vid, batch_mouse, batch_btn) + dmd_loss = get_dmd_loss(samples, batch_mouse, batch_btn) + metrics.log('dmd_loss', dmd_loss) + + optimizer_step(dmd_loss, self.model, self.scaler, self.opt) + self.ema.update() + + with torch.no_grad(): + wandb_dict = metrics.pop() + wandb_dict['time'] = timer.hit() + timer.reset() + + if self.total_step_counter % self.train_cfg.sample_interval == 0: + with ctx, torch.no_grad(): + n_samples = self.train_cfg.n_samples + samples, sample_mouse, sample_button = sampler( + get_ema_core(), + batch_vid[:n_samples], + batch_mouse[:n_samples], + batch_btn[:n_samples], + decode_fn = decode_fn, + scale=self.train_cfg.vae_scale + ) # -> [b,n,c,h,w] + if self.rank == 0: wandb_dict['samples'] = to_wandb(samples, sample_mouse, sample_button) + + if self.rank == 0: + wandb.log(wandb_dict) + + self.total_step_counter += 1 + if self.total_step_counter % self.train_cfg.save_interval == 0: + if self.rank == 0: + self.save() + + self.barrier() \ No newline at end of file From 25e74f348320f4d8f46902f20406576e3dadf512 Mon Sep 17 00:00:00 2001 From: Sami Date: Mon, 9 Jun 2025 18:17:22 -0400 Subject: [PATCH 16/71] slightly more changes after new samplers --- owl_wms/sampling/cfg.py | 65 ++----------------- owl_wms/sampling/simple.py | 2 +- owl_wms/sampling/window.py | 8 +-- webapp/streaming.py | 33 ++++++---- webapp/utils/samplers.py | 128 +++++++++++++++++++++++++++---------- 5 files changed, 124 insertions(+), 112 deletions(-) diff --git a/owl_wms/sampling/cfg.py b/owl_wms/sampling/cfg.py index 75b8c0a8..2e2dc976 100644 --- a/owl_wms/sampling/cfg.py +++ b/owl_wms/sampling/cfg.py @@ -29,10 +29,10 @@ def __call__(self, model, dummy_batch, mouse, btn, decode_fn = None, scale = 1): x = x - pred*dt ts = ts - dt + pixels = None if decode_fn is not None: - x = x * scale - x = decode_fn(x) - return x, mouse, btn + pixels = decode_fn(x * scale) + return x, pixels, mouse, btn class InpaintCFGSampler(CFGSampler): @torch.no_grad() @@ -63,10 +63,10 @@ def __call__(self, model, dummy_batch, mouse, btn, decode_fn = None, scale = 1): x[:, mid:] = x[:, mid:] - pred[:, mid:]*dt ts[:, mid:] = ts[:, mid:] - dt + pixels = None if decode_fn is not None: - x = x * scale - x = decode_fn(x) - return x, mouse, btn + pixels = decode_fn(x * scale) + return x, pixels, mouse, btn def zlerp(x, alpha): @@ -74,59 +74,6 @@ def zlerp(x, alpha): return x * (1. - alpha) + z * alpha -class WindowCFGSampler: - @torch.no_grad() - def __call__(self, model, dummy_batch, mouse, btn, - sampling_steps = 64, decode_fn = None, scale = 1, cfg_scale = 1.3, - window_length = 60, num_frames = 60, noise_prev = 0.2): - - x = torch.randn_like(dummy_batch) - ts = torch.ones(x.shape[0], x.shape[1], device=x.device,dtype=x.dtype) - dt = 1. / sampling_steps - - clean_history = dummy_batch.clone() - - def step_history(): - new_history = clean_history.clone()[:,-window_length:] # last 60 frames - b,n,c,h,w = new_history.shape - # Shift over by 1: - # [ 0 1 2 3 4 ] - # reassigned to -> zlerp([ 1 2 3 4 ]) + [ randn ] - new_history[:,:-1] = zlerp(new_history[:,1:],noise_prev) # pop off first frame and noise context - new_history[:,-1] = torch.randn(b,1,c,h,w) # Add noise to last - return new_history - - for _ in range(num_frames): - local_history = step_history() - ts_history = torch.ones(local_history.shape[0], local_history.shape[1], device=x.device,dtype=x.dtype) - ts_history[:,-1] = noise_prev - - for _ in range(sampling_steps): - # CFG Branches - x = local_history.clone() - ts = ts_history.clone() - cond_pred = model(x, ts, mouse, btn) - uncond_pred = model(x, ts, torch.zeros_like(mouse), torch.zeros_like(btn)) - pred = uncond_pred + cfg_scale * (cond_pred - uncond_pred) - - x = x - pred*dt - ts = ts - dt - - local_history[:,-1] = x[:,-1] - ts_history[:,-1] = ts[:,-1] - - # Frame is entirely cleaned now - new_frame = local_history[:,-1:] - clean_history = torch.cat([clean_history, new_frame], dim = 1) - - x = clean_history - pixels = None - - if decode_fn is not None: - pixels = decode_fn(x * scale) - return x, pixels - - if __name__ == "__main__": model = lambda x,t,m,b: x diff --git a/owl_wms/sampling/simple.py b/owl_wms/sampling/simple.py index 4d20a7e6..602cd7b6 100644 --- a/owl_wms/sampling/simple.py +++ b/owl_wms/sampling/simple.py @@ -53,7 +53,7 @@ def __call__(self, model, dummy_batch, mouse, btn, decode_fn = None, scale = 1): if __name__ == "__main__": model = lambda x,t,m,b: x - sampler = Sampler() + sampler = SimpleSampler() x = sampler(model, torch.randn(4, 3, 64, 64), torch.randn(4, 2), torch.randn(4, 8)) print(x.shape) diff --git a/owl_wms/sampling/window.py b/owl_wms/sampling/window.py index 5fdf837e..ff220054 100644 --- a/owl_wms/sampling/window.py +++ b/owl_wms/sampling/window.py @@ -95,11 +95,11 @@ def step_history(): extended_mouse = extended_mouse[:,-num_frames:] extended_btn = extended_btn[:,-num_frames:] + pixels = None if decode_fn is not None: - x = x * scale - x = decode_fn(x) - - return x, extended_mouse, extended_btn + pixels = decode_fn(x * scale) + + return x, pixels, extended_mouse, extended_btn def test_window_cfg_sampler(): diff --git a/webapp/streaming.py b/webapp/streaming.py index 7d3642a8..b25dc1fa 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -4,7 +4,8 @@ import asyncio import numpy as np from torch import nn -from functools import cache +from typing import Optional +from functools import cached_property from webapp.utils.samplers import create_sampler from webapp.utils.configs import SamplingConfig, StreamingConfig @@ -71,19 +72,20 @@ def __init__(self, encoder: nn.Module, decoder: nn.Module, # Create WindowCFGSampler for frame generation self.sample_window_fn = create_sampler('window', encoder, decoder, batch_size=1, - sampling_steps=self.sampling_config.sampling_steps, + n_steps=self.sampling_config.sampling_steps, vae_scale=self.sampling_config.vae_scale, cfg_scale=self.sampling_config.cfg_scale, window_length=self.sampling_config.window_length, num_frames=self.sampling_config.num_frames, - noise_prev=self.sampling_config.noise_prev) + noise_prev=self.sampling_config.noise_prev, + only_return_generated=True) # Initialize frame history as empty tensor - self.latent_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) - self.mouse_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) - self.button_history: torch.Tensor = torch.tensor([], device=self.streaming_config.device) + self.latent_history: Optional[torch.Tensor] = None + self.mouse_history: Optional[torch.Tensor] = None + self.button_history: Optional[torch.Tensor] = None def add_to_history(self, frame_batch: torch.Tensor, mouse_batch: torch.Tensor, button_batch: torch.Tensor): - if self.latent_history.equal(torch.tensor([], device=self.streaming_config.device)): + if self.latent_history is None: self.latent_history = frame_batch self.mouse_history = mouse_batch self.button_history = button_batch @@ -99,13 +101,12 @@ def add_to_history(self, frame_batch: torch.Tensor, mouse_batch: torch.Tensor, b self.button_history = self.button_history[-self.streaming_config.window_length:] def get_latent_history_batch(self) -> torch.Tensor: - if self.latent_history.equal(torch.tensor([], device=self.streaming_config.device)): + if self.latent_history.numel() == 0: return self.dummy_batch return self.latent_history.unsqueeze(0) - @property - @cache + @cached_property def dummy_batch(self) -> torch.Tensor: """Dummy autoencoder latents for the sampler to initialize shapes.""" tokens_h = tokens_w = int(math.sqrt(self.model_config.tokens_per_frame)) @@ -159,6 +160,9 @@ async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor, video_frames: [frames_per_batch, 3, 256, 256] - pure video frames overlay_frames: [frames_per_batch, 3, action_margin_px_height, 256] - action overlay frames """ + mouse = mouse.to(self.streaming_config.device) + button = button.to(self.streaming_config.device) + if self.debug: num_frames = mouse.shape[0] # Create gradient from white to black to white across columns @@ -175,13 +179,14 @@ async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor, # between 0 and 255 full_frames = (full_frames * 255).to(torch.uint8) else: - with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): + device_type = 'cuda' if self.streaming_config.device.type == 'cuda' else 'cpu' + with torch.no_grad(), torch.autocast(device_type=device_type, dtype=torch.bfloat16): latents, full_frames = self.sample_window_fn(dummy_batch=self.get_latent_history_batch(), mouse=mouse.float().unsqueeze(0), btn=button.float().unsqueeze(0)) # [1, window_length, 3, 256, 256] - # remove batch dimension, then take only the frames we generated, since the WindowCFGSampler appends the history (which is of window_length=60) - latents = latents [0, -self.sampling_config.num_frames:] - full_frames = full_frames [0, -self.sampling_config.num_frames:] + # remove batch dimension, it returns the frames we generated, since we pass in only_return_generated=True + latents = latents [0, ::] + full_frames = full_frames [0, ::] # then, convert the frames to a pixel-range of [0-255] from [-1,1] full_frames = (full_frames + 1) / 2 full_frames = (full_frames * 255).to(torch.uint8) diff --git a/webapp/utils/samplers.py b/webapp/utils/samplers.py index 7dad32c8..2f1bf688 100644 --- a/webapp/utils/samplers.py +++ b/webapp/utils/samplers.py @@ -2,7 +2,9 @@ from typing import Literal, Callable from functools import partial, cache from multimethod import multimethod -from owl_wms.sampling.cfg import CFGSampler, InpaintCFGSampler, WindowCFGSampler +from owl_wms.sampling.cfg import CFGSampler +from owl_wms.sampling.simple import SimpleSampler, InpaintSimpleSampler +from owl_wms.sampling.window import WindowCFGSampler from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn SAMPLING_STEPS = 60 @@ -16,61 +18,119 @@ @multimethod def create_sampler(sampler_id: Literal['cfg'], encoder, decoder, - batch_size: int = 8, - sampling_steps: int = SAMPLING_STEPS, - cfg_scale: float = CFG_SCALE, - vae_scale: float = SCALE, - **kwargs) -> Callable[[MouseData, ButtonData], - tuple[LatentData, VideoData]]: - @cache # simple singleton - def _sampler(): return CFGSampler() + batch_size: int = 8, + n_steps: int = 20, + cfg_scale: float = 1.3, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create CFG sampler with its specific parameters.""" + + @cache + def _sampler(): + return CFGSampler(n_steps=n_steps, cfg_scale=cfg_scale) return partial( _sampler().__call__, - sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), scale=vae_scale, - cfg_scale=cfg_scale, model=encoder ) -@multimethod -def create_sampler(sampler_id: Literal['inpaint_cfg'], encoder, decoder, - batch_size: int = 8, - sampling_steps: int = SAMPLING_STEPS, - cfg_scale: float = CFG_SCALE, - vae_scale: float = SCALE, - **kwargs) -> Callable[[MouseData, ButtonData], - tuple[LatentData, VideoData]]: + +@multimethod +def create_sampler(sampler_id: Literal['simple'], encoder, decoder, + batch_size: int = 8, + n_steps: int = 64, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create Simple sampler with its specific parameters.""" + @cache - def _sampler(): return InpaintCFGSampler() + def _sampler(): + return SimpleSampler(n_steps=n_steps) + + return partial( + _sampler().__call__, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + scale=vae_scale, + model=encoder + ) + +@multimethod +def create_sampler(sampler_id: Literal['inpaint_simple'], encoder, decoder, + batch_size: int = 8, + n_steps: int = 64, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create Inpaint Simple sampler with its specific parameters.""" + + @cache + def _sampler(): + return InpaintSimpleSampler(n_steps=n_steps) + return partial( _sampler().__call__, - sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), scale=vae_scale, - cfg_scale=cfg_scale, model=encoder ) + @multimethod def create_sampler(sampler_id: Literal['window'], encoder, decoder, - batch_size: int = 8, - sampling_steps: int = SAMPLING_STEPS, - cfg_scale: float = CFG_SCALE, - vae_scale: float = SCALE, - **kwargs) -> Callable[[MouseData, ButtonData], - tuple[LatentData, VideoData]]: + batch_size: int = 8, + n_steps: int = 20, + cfg_scale: float = 1.3, + window_length: int = 60, + num_frames: int = 60, + noise_prev: float = 0.2, + only_return_generated: bool = False, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create Window CFG sampler with its specific parameters.""" + @cache - def _sampler(): return WindowCFGSampler() - + def _sampler(): + return WindowCFGSampler( + n_steps=n_steps, + cfg_scale=cfg_scale, + window_length=window_length, + num_frames=num_frames, + noise_prev=noise_prev, + only_return_generated=only_return_generated + ) + return partial( _sampler().__call__, - sampling_steps=sampling_steps, decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), scale=vae_scale, - cfg_scale=cfg_scale, - model=encoder, - **kwargs + model=encoder + ) + + +# Example usage: +if __name__ == "__main__": + # Each sampler type can be created with its specific parameters + import webapp.utils.models + encoder, decoder, model_config = webapp.utils.models.load_models() + + # CFG sampler + cfg_sampler = create_sampler( + 'cfg', encoder, decoder, + n_steps=25, cfg_scale=1.5 + ) + + # Simple sampler + simple_sampler = create_sampler( + 'simple', encoder, decoder, + n_steps=50 + ) + + # Window sampler with all its specific params + window_sampler = create_sampler( + 'window', encoder, decoder, + n_steps=30, cfg_scale=1.4, + window_length=80, num_frames=120, + noise_prev=0.3, only_return_generated=True ) From e856fe4fc372c71548f0f37d2fd38e2c5a44a5a9 Mon Sep 17 00:00:00 2001 From: Sami Date: Mon, 9 Jun 2025 18:22:07 -0400 Subject: [PATCH 17/71] renaming for clarity --- webapp/utils/samplers.py | 136 --------------------------------------- 1 file changed, 136 deletions(-) delete mode 100644 webapp/utils/samplers.py diff --git a/webapp/utils/samplers.py b/webapp/utils/samplers.py deleted file mode 100644 index 2f1bf688..00000000 --- a/webapp/utils/samplers.py +++ /dev/null @@ -1,136 +0,0 @@ -from torch import Tensor -from typing import Literal, Callable -from functools import partial, cache -from multimethod import multimethod -from owl_wms.sampling.cfg import CFGSampler -from owl_wms.sampling.simple import SimpleSampler, InpaintSimpleSampler -from owl_wms.sampling.window import WindowCFGSampler -from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn - -SAMPLING_STEPS = 60 -SCALE = 2.17 -CFG_SCALE = 1.3 - -MouseData = Tensor -ButtonData = Tensor -VideoData = Tensor -LatentData = Tensor - -@multimethod -def create_sampler(sampler_id: Literal['cfg'], encoder, decoder, - batch_size: int = 8, - n_steps: int = 20, - cfg_scale: float = 1.3, - vae_scale: float = 1.0, - **kwargs) -> Callable: - """Create CFG sampler with its specific parameters.""" - - @cache - def _sampler(): - return CFGSampler(n_steps=n_steps, cfg_scale=cfg_scale) - - return partial( - _sampler().__call__, - decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=vae_scale, - model=encoder - ) - - -@multimethod -def create_sampler(sampler_id: Literal['simple'], encoder, decoder, - batch_size: int = 8, - n_steps: int = 64, - vae_scale: float = 1.0, - **kwargs) -> Callable: - """Create Simple sampler with its specific parameters.""" - - @cache - def _sampler(): - return SimpleSampler(n_steps=n_steps) - - return partial( - _sampler().__call__, - decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=vae_scale, - model=encoder - ) - - -@multimethod -def create_sampler(sampler_id: Literal['inpaint_simple'], encoder, decoder, - batch_size: int = 8, - n_steps: int = 64, - vae_scale: float = 1.0, - **kwargs) -> Callable: - """Create Inpaint Simple sampler with its specific parameters.""" - - @cache - def _sampler(): - return InpaintSimpleSampler(n_steps=n_steps) - - return partial( - _sampler().__call__, - decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=vae_scale, - model=encoder - ) - - -@multimethod -def create_sampler(sampler_id: Literal['window'], encoder, decoder, - batch_size: int = 8, - n_steps: int = 20, - cfg_scale: float = 1.3, - window_length: int = 60, - num_frames: int = 60, - noise_prev: float = 0.2, - only_return_generated: bool = False, - vae_scale: float = 1.0, - **kwargs) -> Callable: - """Create Window CFG sampler with its specific parameters.""" - - @cache - def _sampler(): - return WindowCFGSampler( - n_steps=n_steps, - cfg_scale=cfg_scale, - window_length=window_length, - num_frames=num_frames, - noise_prev=noise_prev, - only_return_generated=only_return_generated - ) - - return partial( - _sampler().__call__, - decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), - scale=vae_scale, - model=encoder - ) - - -# Example usage: -if __name__ == "__main__": - # Each sampler type can be created with its specific parameters - import webapp.utils.models - encoder, decoder, model_config = webapp.utils.models.load_models() - - # CFG sampler - cfg_sampler = create_sampler( - 'cfg', encoder, decoder, - n_steps=25, cfg_scale=1.5 - ) - - # Simple sampler - simple_sampler = create_sampler( - 'simple', encoder, decoder, - n_steps=50 - ) - - # Window sampler with all its specific params - window_sampler = create_sampler( - 'window', encoder, decoder, - n_steps=30, cfg_scale=1.4, - window_length=80, num_frames=120, - noise_prev=0.3, only_return_generated=True - ) From 12327c60ddc1448f9e0c486c8fc8485a0d0e03ce Mon Sep 17 00:00:00 2001 From: Sami Date: Mon, 9 Jun 2025 18:22:19 -0400 Subject: [PATCH 18/71] renaming for clarity --- webapp/streaming.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/webapp/streaming.py b/webapp/streaming.py index b25dc1fa..9e1cbb7e 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -7,7 +7,7 @@ from typing import Optional from functools import cached_property -from webapp.utils.samplers import create_sampler +from webapp.utils.create_samplers import create_sampler from webapp.utils.configs import SamplingConfig, StreamingConfig from owl_wms.configs import TrainingConfig, TransformerConfig as ModelConfig from webapp.utils.visualize_overlay_actions import _draw_video as _draw_action_overlays From d0b73da42d817edcd3337e013d88cb33cb71ca1c Mon Sep 17 00:00:00 2001 From: Sami Date: Mon, 9 Jun 2025 23:42:29 -0400 Subject: [PATCH 19/71] forgot file --- webapp/utils/create_samplers.py | 136 ++++++++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 webapp/utils/create_samplers.py diff --git a/webapp/utils/create_samplers.py b/webapp/utils/create_samplers.py new file mode 100644 index 00000000..2f1bf688 --- /dev/null +++ b/webapp/utils/create_samplers.py @@ -0,0 +1,136 @@ +from torch import Tensor +from typing import Literal, Callable +from functools import partial, cache +from multimethod import multimethod +from owl_wms.sampling.cfg import CFGSampler +from owl_wms.sampling.simple import SimpleSampler, InpaintSimpleSampler +from owl_wms.sampling.window import WindowCFGSampler +from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn + +SAMPLING_STEPS = 60 +SCALE = 2.17 +CFG_SCALE = 1.3 + +MouseData = Tensor +ButtonData = Tensor +VideoData = Tensor +LatentData = Tensor + +@multimethod +def create_sampler(sampler_id: Literal['cfg'], encoder, decoder, + batch_size: int = 8, + n_steps: int = 20, + cfg_scale: float = 1.3, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create CFG sampler with its specific parameters.""" + + @cache + def _sampler(): + return CFGSampler(n_steps=n_steps, cfg_scale=cfg_scale) + + return partial( + _sampler().__call__, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + scale=vae_scale, + model=encoder + ) + + +@multimethod +def create_sampler(sampler_id: Literal['simple'], encoder, decoder, + batch_size: int = 8, + n_steps: int = 64, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create Simple sampler with its specific parameters.""" + + @cache + def _sampler(): + return SimpleSampler(n_steps=n_steps) + + return partial( + _sampler().__call__, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + scale=vae_scale, + model=encoder + ) + + +@multimethod +def create_sampler(sampler_id: Literal['inpaint_simple'], encoder, decoder, + batch_size: int = 8, + n_steps: int = 64, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create Inpaint Simple sampler with its specific parameters.""" + + @cache + def _sampler(): + return InpaintSimpleSampler(n_steps=n_steps) + + return partial( + _sampler().__call__, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + scale=vae_scale, + model=encoder + ) + + +@multimethod +def create_sampler(sampler_id: Literal['window'], encoder, decoder, + batch_size: int = 8, + n_steps: int = 20, + cfg_scale: float = 1.3, + window_length: int = 60, + num_frames: int = 60, + noise_prev: float = 0.2, + only_return_generated: bool = False, + vae_scale: float = 1.0, + **kwargs) -> Callable: + """Create Window CFG sampler with its specific parameters.""" + + @cache + def _sampler(): + return WindowCFGSampler( + n_steps=n_steps, + cfg_scale=cfg_scale, + window_length=window_length, + num_frames=num_frames, + noise_prev=noise_prev, + only_return_generated=only_return_generated + ) + + return partial( + _sampler().__call__, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + scale=vae_scale, + model=encoder + ) + + +# Example usage: +if __name__ == "__main__": + # Each sampler type can be created with its specific parameters + import webapp.utils.models + encoder, decoder, model_config = webapp.utils.models.load_models() + + # CFG sampler + cfg_sampler = create_sampler( + 'cfg', encoder, decoder, + n_steps=25, cfg_scale=1.5 + ) + + # Simple sampler + simple_sampler = create_sampler( + 'simple', encoder, decoder, + n_steps=50 + ) + + # Window sampler with all its specific params + window_sampler = create_sampler( + 'window', encoder, decoder, + n_steps=30, cfg_scale=1.4, + window_length=80, num_frames=120, + noise_prev=0.3, only_return_generated=True + ) From cd9f1b5aa1384eb0e6ac224d8daf9898c2e68587 Mon Sep 17 00:00:00 2001 From: Sami Date: Tue, 10 Jun 2025 01:20:46 -0400 Subject: [PATCH 20/71] idk --- webapp/utils/models.py | 2 +- webapp/utils/render.py | 35 ++++++++++++++++++++--------------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/webapp/utils/models.py b/webapp/utils/models.py index 4766d9d4..5734e165 100644 --- a/webapp/utils/models.py +++ b/webapp/utils/models.py @@ -92,7 +92,7 @@ def load_decoder(self, device: Optional[Union[str, torch.device]] = None, eval_mode: bool = True, verbose: bool = True) -> nn.Module: - decoder = get_decoder_only() + decoder = get_decoder_only(vae_id='dcae', cfg_path=str(self.paths.config), ckpt_path=str(self.paths.checkpoint)) freeze(decoder) if verbose: diff --git a/webapp/utils/render.py b/webapp/utils/render.py index 4cd8d7e3..069901cc 100644 --- a/webapp/utils/render.py +++ b/webapp/utils/render.py @@ -9,7 +9,7 @@ import einops as eo from webapp.utils.models import load_models -from webapp.utils.samplers import create_sampler +from webapp.utils.create_samplers import create_sampler, CFG_SCALE from webapp.utils.action_builder import ActionSequenceGenerator, ActionConfig, ActionPattern HEIGHT = 256 @@ -22,7 +22,7 @@ OUTPUT_DIR = "generated_videos" SAMPLER_TYPE = 'window' DEFAULT_PATTERN = ActionPattern.LOOK_AROUND - +VAE_SCALE = 2.17 def setup_output_dir(): Path(OUTPUT_DIR).mkdir(exist_ok=True) @@ -59,7 +59,7 @@ def synthesize_video(mouse_actions, button_actions, encoder, decoder, sampler): # Generate video with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): - latents, video = sampler( + latents, video, mouse, button = sampler( dummy_batch=dummy_batch, mouse=mouse_actions, btn=button_actions @@ -124,19 +124,27 @@ def save_video(video_tensor: torch.Tensor, filename="generated_video", fps=30): return str(fallback_path) -def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): +def render_video( + pattern=DEFAULT_PATTERN, + length=SEQUENCE_LENGTH, + verbose=True, + sampler_type=SAMPLER_TYPE, + vae_scale=VAE_SCALE, + cfg_scale=CFG_SCALE, + encoder=None, decoder=None): """Simple video rendering - just load, generate, save.""" if verbose: print(f"🎬 Rendering video with pattern: {pattern.value}") # Load model and create sampler - if verbose: - print("Loading model...") - encoder, decoder, model_config = load_models(device=DEVICE, verbose=verbose) + if encoder is None or decoder is None: + if verbose: + print("Loading model...") + encoder, decoder, model_config = load_models(device=DEVICE, verbose=verbose) if verbose: print("Creating sampler...") - sampler = create_sampler(SAMPLER_TYPE, encoder, decoder) + sampler = create_sampler(sampler_type, encoder, decoder, vae_scale=vae_scale, cfg_scale=cfg_scale) # Generate actions if verbose: @@ -151,7 +159,7 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): # Save video if verbose: print("Saving video...") - path = save_video(video, f"render_{pattern.value}") + path = save_video(video, f"{sampler_type}_render_{pattern.value}_{vae_scale=:.2f}_{cfg_scale=:.2f}") if verbose: print(f"✅ Done! Video saved to: {path}") @@ -165,9 +173,6 @@ def render_video(pattern=DEFAULT_PATTERN, length=SEQUENCE_LENGTH, verbose=True): print("🎮 Simple OWL-WMS Video Renderer") # Render with default settings - render_video(verbose=True) - - # Render with different pattern - render_video(ActionPattern.SHOOT, verbose=True) - render_video(ActionPattern.LOOK_AROUND, verbose=True) - render_video(ActionPattern.CIRCLE_STRAFE, verbose=True) + encoder, decoder, model_config = load_models(device=DEVICE, verbose=True) + render_video(verbose=True, sampler_type='window', encoder=encoder, decoder=decoder) + render_video(verbose=True, sampler_type='cfg', encoder=encoder, decoder=decoder) From 5a7af71f45a38ab8ad50aa8173b8e58e71e0f3f4 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 19:24:07 +0000 Subject: [PATCH 21/71] Allow data to have keyframes --- owl_wms/configs.py | 1 + owl_wms/data/s3_cod_latent.py | 47 ++++++++++++++++++++++++++++------- 2 files changed, 39 insertions(+), 9 deletions(-) diff --git a/owl_wms/configs.py b/owl_wms/configs.py index d572a3c9..5648ca52 100644 --- a/owl_wms/configs.py +++ b/owl_wms/configs.py @@ -18,6 +18,7 @@ class TransformerConfig: cfg_prob : float = 0.1 n_buttons : int = 8 tokens_per_frame : int = 16 + n_frames : int = 120 causal : bool = False diff --git a/owl_wms/data/s3_cod_latent.py b/owl_wms/data/s3_cod_latent.py index d223d315..8fdcdddf 100644 --- a/owl_wms/data/s3_cod_latent.py +++ b/owl_wms/data/s3_cod_latent.py @@ -33,13 +33,14 @@ def pop(self): BUCKET_NAME="cod-data-latent-360x640to5x8" class S3CoDLatentDataset(IterableDataset): - def __init__(self, window_length=120, file_share_max=20, rank=0, world_size=1): + def __init__(self, window_length=120, file_share_max=20, rank=0, world_size=1, include_keyframe = False): super().__init__() self.window = window_length self.file_share_max = file_share_max self.rank = rank self.world_size = world_size + self.include_keyframe = include_keyframe # Queue parameters self.max_tars = 2 @@ -140,8 +141,23 @@ def background_load_data(self): latent_slice = latent[window_start:window_start+self.window].float() mouse_slice = mouse[window_start:window_start+self.window] button_slice = button[window_start:window_start+self.window] - - self.data_queue.add((latent_slice, mouse_slice, button_slice)) + + if self.include_keyframe: + # Sample keyframe from nearby in video but not in window + buffer = 400 + valid_range_start = max(0, window_start - buffer) + valid_range_end = min(len(latent), window_start + self.window + buffer) + + # Exclude the actual window frames + valid_frames = list(range(valid_range_start, window_start)) + \ + list(range(window_start + self.window, valid_range_end)) + + if valid_frames: + keyframe_idx = random.choice(valid_frames) + latent_keyframe = latent[keyframe_idx].float().unsqueeze(0) + self.data_queue.add((latent_slice, latent_keyframe, mouse_slice, button_slice)) + else: + self.data_queue.add((latent_slice, mouse_slice, button_slice)) except Exception as e: print(f"Error processing tar: {e}") @@ -157,12 +173,25 @@ def __iter__(self): time.sleep(0.1) def collate_fn(batch): - # batch is list of triples - latents, mouses, buttons = zip(*batch) - latents = torch.stack(latents) # [b,n,c,h,w] - mouses = torch.stack(mouses) # [b,n,2] - buttons = torch.stack(buttons) # [b,n,n_buttons] - return latents, mouses, buttons + # batch is list of triples or quads + items = zip(*batch) + items = list(items) + + if len(items) == 3: + # No keyframe case + latents, mouses, buttons = items + latents = torch.stack(latents) # [b,n,c,h,w] + mouses = torch.stack(mouses) # [b,n,2] + buttons = torch.stack(buttons) # [b,n,n_buttons] + return latents, mouses, buttons + else: + # With keyframe case + latents, keyframes, mouses, buttons = items + latents = torch.stack(latents) # [b,n,c,h,w] + keyframes = torch.stack(keyframes) # [b,1,c,h,w] + mouses = torch.stack(mouses) # [b,n,2] + buttons = torch.stack(buttons) # [b,n,n_buttons] + return latents, keyframes, mouses, buttons def get_loader(batch_size, **data_kwargs): if dist.is_initialized(): From aacc01b6cc0b78fee49b349ad18549972d711294 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 19:24:14 +0000 Subject: [PATCH 22/71] fix kv caching --- owl_wms/nn/attn.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/owl_wms/nn/attn.py b/owl_wms/nn/attn.py index 38cee69e..ee54ccc6 100644 --- a/owl_wms/nn/attn.py +++ b/owl_wms/nn/attn.py @@ -53,6 +53,7 @@ def forward(self, x, kv_cache = None): mask = None else: mask = create_block_causal_mask(x.shape[1], self.tokens_per_frame).to(x.device) + mask = mask.to(device=x.device,dtype=x.dtype) mask = mask.unsqueeze(0).repeat(x.shape[0], 1, 1) if kv_cache is not None: @@ -82,9 +83,6 @@ def __init__(self, config): dim = config.d_model - self.norm1 = LayerNorm(dim) - self.norm2 = LayerNorm(dim) - self.attn = Attn(config) self.mlp = MLP(config) From e86756e969153aa84f438ef563c620f6694b7314 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 19:24:54 +0000 Subject: [PATCH 23/71] shortcut model, mmdit, kv cache for mmdit --- owl_wms/models/gamerft_shortcut.py | 287 +++++++++++++++++++++++++++++ owl_wms/nn/embeddings.py | 29 ++- owl_wms/nn/kv_cache.py | 13 +- owl_wms/nn/mmattn.py | 286 ++++++++++++++++++++++++++++ 4 files changed, 612 insertions(+), 3 deletions(-) create mode 100644 owl_wms/models/gamerft_shortcut.py create mode 100644 owl_wms/nn/mmattn.py diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py new file mode 100644 index 00000000..959b2e9e --- /dev/null +++ b/owl_wms/models/gamerft_shortcut.py @@ -0,0 +1,287 @@ +""" +Causal-First RFT With Shortcut objective +""" + +import torch +from torch import nn +import torch.nn.functional as F + +import einops as eo + +from ..nn.embeddings import ( + TimestepEmbedding, + StepEmbedding, + ControlEmbedding, + LearnedPosEnc +) +from ..nn.attn import UViT, FinalLayer +from ..nn.mmattn import MMUViT + +class ShortcutGameRFTCore(nn.Module): + def __init__(self, config): + super().__init__() + + self.transformer = MMUViT(config) + self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) + + self.step_embed = StepEmbedding(config.d_model) + self.t_embed = TimestepEmbedding(config.d_model) + + self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) + self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) + + self.pos_enc = LearnedPosEnc(config.tokens_per_frame * config.n_frames, config.d_model) + + self.proj_y_in = nn.Linear(config.channels, config.d_model, bias = False) + self.pos_enc_y = LearnedPosEnc(config.tokens_per_frame, config.d_model) + + def sample(self, x, y, mouse, btn, kv_cache = None, d = None, t = None): + """ + This is a function that largely abstracts + away most things for the specific case where + you are only generating the one next token + """ + if x is None: + x = torch.randn_like(y) + + b,n,c,h,w = x.shape + if t is None: + t = torch.ones_like(x[:,:,0,0,0]) + if d is None: + d = torch.ones_like(x[:,:,0,0,0]) + + return self.forward(x, t, mouse, btn, d, kv_cache) + + def forward(self, x, y, t, mouse, btn, d, kv_cache = None): + # x is [b,n,c,h,w] + # y is [b,1,c,h,w] + # t is [b,n] + # d is [b,n] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + + print(mouse.shape) + print(btn.shape) + print(x.shape) + + ctrl_cond = self.control_embed(mouse, btn) + t_cond = self.t_embed(t) + d_cond = self.step_embed(d) + + cond = ctrl_cond + t_cond + d_cond # [b,n,d] + + b,n,c,h,w = x.shape + x = eo.rearrange(x, 'b n c h w -> b (n h w) c') + y = eo.rearrange(y, 'b n c h w -> b (n h w) c') + + x = self.proj_in(x) + x = self.pos_enc(x) + + y = self.proj_y_in(y) + y = self.pos_enc_y(y) + + x = self.transformer(x, y, cond, kv_cache) + x = self.proj_out(x, cond) # -> [b,n*hw,c] + x = eo.rearrange(x, 'b (n h w) c -> b n c h w', n=n,h=h,w=w) + + return x + +def sample_discrete_timesteps(steps, eps = 1.0e-6): + # steps is Tensor([1,4,2,64,16]) as an example + b = len(steps) + n = len(steps[0]) + + ts_list = [] + ts = torch.rand(b, n, device=steps.device, dtype=steps.dtype) * (steps - eps) + ts = ts.clamp(eps).ceil() / steps + """ + Example, if d was all 2, ts would be [0,2] + so do clamp, then ceil will be 1 or 2 (0, 2] + then do t / 2 and get 0.5 or 1.0, our desired timesteps + """ + return ts + +def sample_steps(b, n, device, dtype, min_val = 0): + valid = torch.tensor([2**i for i in range(min_val, 8)]) # [1,2,...,128] + inds = torch.randint(low=0,high=len(valid), size = (b,n)) + steps = valid[inds].to(device=device,dtype=dtype) + return steps + +class ShortcutGameRFT(nn.Module): + def __init__(self, config): + super().__init__() + + self.core = GameRFTCore(config) + self.cfg_prob = config.cfg_prob + + self.ema = None + self.sc_frac = 0.25 + self.cfg_strength = 1.3 + + def set_ema_core(self, ema): + if hasattr(ema.ema_model, 'module'): + self.ema = ema.ema_model.module.core + else: + self.ema = ema.ema_model.core + + @torch.no_grad() + @torch.compile() + def get_sc_targets(self, x, y, mouse, btn): + steps_slow = sample_steps(len(x), x.device, x.dtype, min_val = 1) + steps_fast = steps_slow / 2 + + dt_slow = 1./steps_slow + dt_fast = 1./steps_fast + + def expand(t): + #b,c,h,w = x.shape + #t = eo.repeat(t,'b -> b c h w',c=c,h=h,w=w) + #return t + return t[:,:,None,None,None] + + ts = sample_discrete_timesteps(steps_fast) + cfg_mask = torch.isclose(steps_slow, 128) + cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] + + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + pred_1_uncond = self.ema(x, y, ts, null_mouse, null_btn, steps_slow) + pred_1_cond = self.ema(x, y, ts, mouse, btn) + pred_1_cfg = pred_1_uncond + self.cfg_scale * (pred_1_cond - pred_1_uncond) + pred_1 = torch.where(cfg_mask, pred_1_cfg, pred_1_cond) + + x_new = x - pred_1 * expand(dt_slow) + ts_new = ts - dt_slow + + pred_2_uncond = self.ema(x_new, y, ts_new, null_mouse, null_btn, steps_slow) + pred_2_cond = self.ema(x_new, y, ts_new, mouse, btn) + pred_2_cfg = pred_2_uncond + self.cfg_scale * (pred_2_cond - pred_2_uncond) + pred_2 = torch.where(cfg_mask, pred_2_cfg, pred_2_cond) + + pred = 0.5 * (pred_1 + pred_2) + return pred, steps_fast, ts + + def get_sc_loss(self, x, y, mouse, btn): + target, steps, ts = self.get_sc_targets(x, y, mouse, btn) + pred = self.core(x, y, ts, mouse, btn, steps) + sc_loss = F.mse_loss(pred, target) + return sc_loss + + def forward(self, x, y, mouse, btn, return_dict = False, cfg_prob = None): + # x is [b,n,c,h,w] + # y (seed frame) is [b,1,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + _,n,c,h,w = x.shape + + # Split batches between consistency/rf + b = int(len(x) * (1 - self.sc_frac)) + x,x_sc = x[:b], x[b:] + y,y_sc = y[:b], y[b:] + mouse,mouse_sc = mouse[:b], mouse[b:] + btn,btn_sc = btn[:b], btn[b:] + + # Consistency inputs are not autoregressive + n_tok_sc = self.config.n_frames * self.config.tokens_per_frame + x_sc = x_sc[:,:self.config.n_frames] + mouse_sc = mouse_sc[:,:self.config.n_frames] + btn_sc = btn_sc[:,:self.config.n_frames] + + # Apply classifier-free guidance dropout + if cfg_prob is None: + cfg_prob = self.cfg_prob + if cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + # Where mask is True, replace with zeros + mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) + btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) + + with torch.no_grad(): + d = torch.ones_like(x[:,:,0,0,0]) + ts = sample_discrete_timesteps(d) + ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() + + ts_exp = eo.repeat(ts, 'b n -> b n 1 1 1') + z = torch.randn_like(x) + + lerpd = x * (1. - ts_exp) + z * ts_exp + target = z - x + + pred = self.core(lerpd, y, ts, mouse, btn, d) + diff_loss = F.mse_loss(pred, target) + sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc) + + if not return_dict: + return diff_loss + else: + return { + 'diffusion_loss' : diff_loss, + 'lerpd' : lerpd, + 'pred' : pred, + 'ts': ts, + 'z': z + } + +if __name__ == "__main__": + from ..configs import TransformerConfig + from ..nn.kv_cache import KVCache + + cfg = TransformerConfig( + None, # model_id + 6, # n_layers + 6, # n_heads + 384, # d_model + 1, # patch_size + 128, # channels + 16, # sample_size + 0.1, # cfg_prob + 11, # n_buttons + 16, # tokens_per_frame + 10, # n_frames + True # causal + ) + + model = ShortcutGameRFTCore(cfg).bfloat16().cuda() + + NUM_FRAMES = 10 + x = torch.randn(1, NUM_FRAMES, 128, 4, 4).bfloat16().cuda() + y = torch.randn(1, 1, 128, 4, 4).bfloat16().cuda() + mouse = torch.randn(1, NUM_FRAMES, 2).bfloat16().cuda() + btn = torch.randn(1, NUM_FRAMES, 11).bfloat16().cuda() + t = torch.full((1, NUM_FRAMES), 0.25, device='cuda', dtype=torch.bfloat16) + d = torch.full((1, NUM_FRAMES), 4, device='cuda', dtype=torch.bfloat16) + + cache = KVCache(cfg).to(device='cuda', dtype=torch.bfloat16) + cache.reset(1) + + with torch.no_grad(): + # First pass - generate cache for all frames + cache.enable_cache_updates() + out = model(x, y, t, mouse, btn, d, cache) + print(f"Initial cache length: {len(cache)}") + print(f"Initial cache shape: {cache.cache[0][0].shape}") + + # Generate single new frame with t=1, d=1 + new_x = torch.randn(1, 1, 128, 4, 4).bfloat16().cuda() + new_mouse = torch.randn(1, 1, 2).bfloat16().cuda() + new_btn = torch.randn(1, 1, 11).bfloat16().cuda() + new_t = torch.ones(1, 1, device='cuda', dtype=torch.bfloat16) + new_d = torch.ones(1, 1, device='cuda', dtype=torch.bfloat16) + + # Disable cache updates for inference + cache.disable_cache_updates() + new_out = model(new_x, y, new_t, new_mouse, new_btn, new_d, cache) + print(f"After inference cache length: {len(cache)}") + print(f"After inference cache shape: {cache.cache[0][0].shape}") + + # Re-enable cache updates and update cache with t=0.25, d=4 + cache.enable_cache_updates() + new_t = torch.full((1, 1), 0.25, device='cuda', dtype=torch.bfloat16) + new_d = torch.full((1, 1), 4, device='cuda', dtype=torch.bfloat16) + new_out = model(new_x, y, new_t, new_mouse, new_btn, new_d, cache) + print(f"Final cache length: {len(cache)}") + print(f"Final cache shape: {cache.cache[0][0].shape}") \ No newline at end of file diff --git a/owl_wms/nn/embeddings.py b/owl_wms/nn/embeddings.py index 01a6c1c2..adbfd32a 100644 --- a/owl_wms/nn/embeddings.py +++ b/owl_wms/nn/embeddings.py @@ -2,6 +2,8 @@ from torch import nn import torch.nn.functional as F +import math + import einops as eo from .mlp import MLPCustom @@ -89,11 +91,16 @@ class LearnedPosEnc(nn.Module): def __init__(self, n_seq, dim): super().__init__() + self.n_seq = n_seq self.p = nn.Parameter(torch.randn(n_seq,dim)*0.02) def forward(self, x): b,n,d = x.shape - p = eo.repeat(self.p, 'n d -> b n d', b = b) + if n < self.n_seq: + # Only add positional embeddings for the last n tokens + p = eo.repeat(self.p[-n:], 'n d -> b n d', b=b) + else: + p = eo.repeat(self.p, 'n d -> b n d', b=b) return x + p class SinCosEmbed(nn.Module): @@ -152,6 +159,26 @@ def forward(self, x): x = self.mlp(x) return x +class StepEmbedding(nn.Module): + def __init__(self, d_out, d_in=512, max_steps=128): + super().__init__() + + self.mlp = MLPCustom(d_in, dim_middle = 4 * d_out, dim_out=d_out) + self.max_steps = max_steps + mult = 1000 / math.log2(max_steps) + self.sincos = SinCosEmbed(d_in, theta=300, mult=mult) + + def forward(self, steps): + if not isinstance(steps, torch.Tensor): + steps = torch.tensor(steps, device=self.mlp.fc_uv.weight.device, dtype=self.mlp.fc_uv.weight.dtype) + if steps.ndim == 0: + steps = steps.unsqueeze(0) + + # Map steps to [0, log2(max_steps)] + t = (math.log2(self.max_steps) - torch.log2(steps.float())).to(steps.dtype) + embs = self.sincos(t) + return self.mlp(embs) + class ConditionEmbedding(nn.Module): def __init__(self, n_classes, dim): super().__init__() diff --git a/owl_wms/nn/kv_cache.py b/owl_wms/nn/kv_cache.py index 99d33eb9..ca153f87 100644 --- a/owl_wms/nn/kv_cache.py +++ b/owl_wms/nn/kv_cache.py @@ -14,6 +14,8 @@ def __init__(self, config : TransformerConfig): self.should_update = False + self.max_length = config.tokens_per_frame * config.n_frames + def enable_cache_updates(self): self.should_update = True @@ -23,6 +25,7 @@ def disable_cache_updates(self): def to(self, device = 'cuda', dtype = torch.bfloat16): self.device = device self.dtype = dtype + return self def reset(self, batch_size = 1): self.shape = (batch_size, self.config.n_heads, 0, self.config.d_model//self.config.n_heads) @@ -46,11 +49,17 @@ def push(self, new_k, new_v, layer_ind): @torch.no_grad() def update(self, new_k, new_v, layer_ind): assert self.cache is not None, "Must reset cache before using" - self.cache[layer_ind] = (new_k,new_v) + + def tuple_truncate(k, v): + k = k[:,:,-self.max_length:] + v = v[:,:,-self.max_length:] + return k, v + + self.cache[layer_ind] = tuple_truncate(new_k,new_v) def __len__(self): assert self.cache is not None, "Must reset cache before using" - return self.cache[0].shape[2] + return self.cache[0][0].shape[2] def shape(self): return self.shape \ No newline at end of file diff --git a/owl_wms/nn/mmattn.py b/owl_wms/nn/mmattn.py new file mode 100644 index 00000000..ca3b003d --- /dev/null +++ b/owl_wms/nn/mmattn.py @@ -0,0 +1,286 @@ +import torch +from torch import nn +import torch.nn.functional as F + +from .normalization import LayerNorm, RMSNorm, QKNorm +from .mlp import MLP + +import einops as eo + +from .modulation import AdaLN, Gate +#from .embeddings import FlatVideoRoPE +from rotary_embedding_torch import RotaryEmbedding + +torch.backends.cuda.enable_flash_sdp(enabled = True) + +from einops._torch_specific import allow_ops_in_compiled_graph +allow_ops_in_compiled_graph() + +""" +This code makes the assumption that there are some +tokens from another modality that must always be attended to +""" + +def create_block_causal_mask_with_mm(tokens, context_tokens, tokens_per_frame): + frames = tokens // tokens_per_frame + + # Create base causal mask, nothing is masked + total_tokens = tokens + context_tokens + mask = torch.zeros(total_tokens, total_tokens) + + # Allow attention within each frame + for i in range(frames): + start = i * tokens_per_frame + end = (i + 1) * tokens_per_frame + mask[start:end, end:tokens] = True # Can't see future frames + + # Context tokens can attend to everything (no masking needed) + # Regular tokens can attend to all context tokens (no masking needed) + # The zeros in mask[tokens:, :] allow context to attend to everything + # The zeros in mask[:, tokens:] allow tokens to attend to all context + + return mask + +class MMAttn(nn.Module): + """ + MMDiT style attention + """ + def __init__(self, config : 'TransformerConfig'): + super().__init__() + + self.n_heads = config.n_heads + + self.qkv_1 = nn.Linear(config.d_model, 3 * config.d_model) + self.qkv_2 = nn.Linear(config.d_model, 3 * config.d_model) + + self.out_1 = nn.Linear(config.d_model, config.d_model) + self.out_2 = nn.Linear(config.d_model, config.d_model) + + self.qk_norm_1 = QKNorm(config.d_model // config.n_heads) + self.qk_norm_2 = QKNorm(config.d_model // config.n_heads) + + self.config = config + self.causal = config.causal + + def split(self, qkv): + return eo.rearrange(qkv, 'b n (three h d) -> three b h n d', three = 3, h = self.n_heads) + + def merge(self, x): + return eo.rearrange(x, 'b h n d -> b n (h d)') + + def forward(self, x_1, x_2, kv_cache=None): + n1 = x_1.shape[1] + + q1,k1,v1 = self.split(self.qkv_1(x_1)) + q2,k2,v2 = self.split(self.qkv_2(x_2)) + + q1,k1 = self.qk_norm_1(q1,k1) + q2,k2 = self.qk_norm_2(q2,k2) + + if not self.causal or (kv_cache is not None and len(kv_cache) > 0): + mask = None + else: + mask = create_block_causal_mask_with_mm(x_1.shape[1], x_2.shape[1], self.config.tokens_per_frame) + mask = mask.to(device=x_1.device,dtype=x_1.dtype) + mask = mask.unsqueeze(0).repeat(x_1.shape[0],1,1) + + if kv_cache is not None: + if len(kv_cache) > 0: + old_k, old_v = kv_cache.get(self.layer_ind) + + new_k = torch.cat([old_k, k1], dim=2).contiguous() + new_v = torch.cat([old_v, v1], dim=2).contiguous() + else: + new_k = k1.contiguous() + new_v = v1.contiguous() + + if kv_cache.should_update: + kv_cache.update(new_k, new_v, self.layer_ind) + + k = torch.cat([new_k, k2], dim=-2) + v = torch.cat([new_v, v2], dim=-2) + q = torch.cat([q1, q2], dim=-2) + + x = F.scaled_dot_product_attention(q, k, v, attn_mask = mask) + x = x[:,:,-q.shape[2]:] # Only keep latest outputs + x = self.merge(x) + else: + q = torch.cat([q1,q2],dim=-2) + k = torch.cat([k1,k2],dim=-2) + v = torch.cat([v1,v2],dim=-2) + + x = F.scaled_dot_product_attention(q,k,v, attn_mask = mask) + x = self.merge(x) + + x_1, x_2 = x[:,:n1], x[:,n1:] + x_1 = self.out_1(x_1) + x_2 = self.out_2(x_2) + + return x_1, x_2 + +class MMDiTBlock(nn.Module): + def __init__(self, config): + super().__init__() + + dim = config.d_model + + self.attn = MMAttn(config) + + self.mlp_1 = MLP(config) + self.mlp_2 = MLP(config) + + # Stream 1 - AdaLN and gating + self.adaln1_1 = AdaLN(dim) + self.gate1_1 = Gate(dim) + self.adaln2_1 = AdaLN(dim) + self.gate2_1 = Gate(dim) + + # Stream 2 - Standard LayerNorm + self.ln1_2 = nn.LayerNorm(dim) + self.ln2_2 = nn.LayerNorm(dim) + + def forward(self, x, y, cond, kv_cache = None): + res1_x = x.clone() + res1_y = y.clone() + + # First attention block + x = self.adaln1_1(x, cond) + y = self.ln1_2(y) + + x, y = self.attn(x, y, kv_cache) + + x = self.gate1_1(x, cond) + + x = res1_x + x + y = res1_y + y + + # Second MLP block + res2_x = x.clone() + res2_y = y.clone() + + x = self.adaln2_1(x, cond) + y = self.ln2_2(y) + + x = self.mlp_1(x) + y = self.mlp_2(y) + + x = self.gate2_1(x, cond) + + x = res2_x + x + y = res2_y + y + + return x, y + +class MMUViT(nn.Module): + def __init__(self, config): + super().__init__() + + blocks = [] + for i in range(config.n_layers): + blocks.append(MMDiTBlock(config)) + blocks[-1].attn.layer_ind = i + + self.blocks = nn.ModuleList(blocks) + + # For odd number of layers, need linear projections for skip connections + n_skip_connections = config.n_layers // 2 + skip_projs = [] + for _ in range(n_skip_connections): + skip_projs.append(nn.Linear(config.d_model * 2, config.d_model)) + self.skip_projs = nn.ModuleList(skip_projs) + + def forward(self, x, y, cond, kv_cache = None): + # Cache early block outputs for skip connections + early_features = [] + n_blocks = len(self.blocks) + mid_idx = n_blocks // 2 + + # Early blocks + for i in range(mid_idx): + x,y = self.blocks[i](x, y, cond, kv_cache) + early_features.append(x) + + # Middle block (if odd number of layers) + x,y = self.blocks[mid_idx](x, y, cond, kv_cache) + + # Late blocks with skip connections + for i in range(mid_idx + 1, n_blocks): + # Get corresponding early block output + early_idx = n_blocks - 1 - i + early_feat = early_features[early_idx] + + # Concatenate early and current features + skip_idx = i - (mid_idx + 1) + x = torch.cat([x, early_feat], dim=-1) + x = self.skip_projs[skip_idx](x) + + x,y = self.blocks[i](x, y, cond, kv_cache) + + return x + + +def test_fwd_with_cache(): + from ..configs import TransformerConfig + from .kv_cache import KVCache + + import matplotlib.pyplot as plt + + cfg = TransformerConfig( + None, + 6, + 6, + 384, + 1, + 128, + 4, + 0.1, + 8, + 16, + True + ) + + model = MMUViT(cfg).bfloat16().cuda() + + NUM_FRAMES = 10 + x = torch.randn(1,16*NUM_FRAMES,384).bfloat16().cuda() + y = torch.randn(1,16,384).bfloat16().cuda() + cond=torch.randn(1,16,384).bfloat16().cuda() + + cache = KVCache(cfg).to(device='cuda',dtype=torch.bfloat16) + cache.reset(1) + + with torch.no_grad(): + cache.enable_cache_updates() + out = model(x,y,cond,cache) + + new_x = torch.randn(1,16,384).bfloat16().cuda() + cond = torch.randn(1,1,384).bfloat16().cuda() + + print(len(cache)) + print(cache.cache[0][0].shape) + new_out = model(new_x, y, cond, cache) + + print(len(cache)) + print(cache.cache[0][0].shape) + +def test_mask(): + import matplotlib.pyplot as plt + + n_frames = 10 + n_tok_per_frame = 16 + n_context = 16 + + mask = create_block_causal_mask_with_mm(n_frames*n_tok_per_frame, n_context, n_tok_per_frame) + + plt.figure(figsize=(10,10)) + plt.imshow(mask.float().cpu().numpy(), cmap='gray') + plt.colorbar() + plt.title(f'Block Causal Mask with MM ({n_frames*n_tok_per_frame} tokens, {n_context} context, {n_tok_per_frame} per frame)') + plt.xlabel('Key Position') + plt.ylabel('Query Position') + plt.savefig('test_mm_mask.png') + plt.close() + + +if __name__ == "__main__": + test_mask() \ No newline at end of file From 0fc78eaad77570b4461e61cb3e2f0dea02e01dc4 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 20:35:06 +0000 Subject: [PATCH 24/71] add shortcut sampler and debug kv cache --- owl_wms/nn/kv_cache.py | 18 ++++++ owl_wms/sampling/__init__.py | 5 +- owl_wms/sampling/shortcut_sampler.py | 95 ++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 owl_wms/sampling/shortcut_sampler.py diff --git a/owl_wms/nn/kv_cache.py b/owl_wms/nn/kv_cache.py index ca153f87..fd2de00d 100644 --- a/owl_wms/nn/kv_cache.py +++ b/owl_wms/nn/kv_cache.py @@ -15,6 +15,7 @@ def __init__(self, config : TransformerConfig): self.should_update = False self.max_length = config.tokens_per_frame * config.n_frames + self.noise_caches = 0.0 def enable_cache_updates(self): self.should_update = True @@ -36,6 +37,9 @@ def reset(self, batch_size = 1): def get(self, layer_ind): assert self.cache is not None, "Must reset cache before using" k,v = self.cache[layer_ind] + if self.noise_caches > 0.0: + k = k + torch.randn_like(k) * self.noise_caches + v = v + torch.randn_like(v) * self.noise_caches return k,v @torch.no_grad() @@ -56,6 +60,20 @@ def tuple_truncate(k, v): return k, v self.cache[layer_ind] = tuple_truncate(new_k,new_v) + + @torch.no_grad() + def truncate(self, truncate_amt): + """ + Truncate frames from the KV cache + """ + truncate_amt = truncate_amt * self.config.tokens_per_frame + def tuple_truncate(k, v): + k = k[:,:,truncate_amt:] + v = v[:,:,truncate_amt:] + return k, v + + for i in range(self.config.n_layers): + self.cache[layer_ind] = tuple_truncate(self.cache[layer_ind]) def __len__(self): assert self.cache is not None, "Must reset cache before using" diff --git a/owl_wms/sampling/__init__.py b/owl_wms/sampling/__init__.py index c35f799e..9488daf6 100644 --- a/owl_wms/sampling/__init__.py +++ b/owl_wms/sampling/__init__.py @@ -8,4 +8,7 @@ def get_sampler_cls(sampler_id): elif sampler_id == "cfg": return CFGSampler elif sampler_id == "window": - return WindowCFGSampler \ No newline at end of file + return WindowCFGSampler + elif sampler_id == "shortcut": + from .shortcut_sampler import CacheShortcutSampler + return CacheShortcutSampler \ No newline at end of file diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py new file mode 100644 index 00000000..13b07c57 --- /dev/null +++ b/owl_wms/sampling/shortcut_sampler.py @@ -0,0 +1,95 @@ +""" +Shortcut sampler, with cache! +""" + +import torch +from torch import nn +import torch.nn.functional as F + +from tqdm import tqdm + +from ..utils import batch_permute_to_length +from ..nn.kv_cache import KVCache + +def zlerp(x, alpha): + z = torch.randn_like(x) + return x * (1. - alpha) + z * alpha + +class CacheShortcutSampler: + """ + Window CFG Sampler samples new frames one by one, by inpainting the final frame. + This is basically diffusion forcing. + + :param n_steps: Number of diffusion steps for each frame (diffusoin steps) + :param cfg_scale: CFG scale for each frame + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param noise_prev: Noise previous frame + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, window_length = 60, num_frames = 60, noise_prev = 0.25, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.noise_prev = noise_prev + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale = 1): + # dummy_batch is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + history = history[:,:self.window_length] + new_frames = [] + alpha = 0.25 # This number is special for our sampler + + # Extended fake controls to use during sampling + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, num_frames + self.window_length) + + # Generate cache over history + noisy_history = zlerp(history.clone(), alpha) + ts = torch.ones_like(noisy_history[:,:,0,0,0]) * alpha + d = torch.ones_like(noisy_history[:,:,0,0,0]) * round(1./alpha) + ts_single = ts[:,0].unsqueeze(1) + d_single = d[:,0].unsqueeze(1) + + cache = KVCache(model.config) + cache.reset(history.shape[0]) + + cache.enable_cache_updates() + _ = model.sample(noisy_history, keyframe, mouse, btn, cache, ts, d) + cache.disable_cache_updates() + + # Cache is now built! + + for frame_idx in tqdm(range(num_frames)): + cache.truncate(1) # Drop first frame + + # Generate new frame + cache.disable_cache_updates() + mouse = extended_mouse[:,self.window_length+frame_idx].unsqueeze(1) + btn = extended_btn[:,self.window_length+frame_idx].unsqueeze(1) + new_frame = model.sample(None, keyframe, mouse, btn, cache) # [b,1,c,h,w] + new_frames.append(new_frame) + + # Add that frame to the cache + cache.enable_cache_updates() + new_frame_noisy = zlerp(new_frame, alpha) + _ = model.sample(new_frame_noisy, keyframe, mouse, btn, cache, ts_single, d_single) + + new_frames = torch.cat(new_frames, dim = 1) + x = torch.cat([history,new_frames], dim = 1) + + if self.only_return_generated: + x = x[:,-num_frames:] + extended_mouse = extended_mouse[:,-num_frames:] + extended_btn = extended_btn[:,-num_frames:] + + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + + return x, extended_mouse, extended_btn \ No newline at end of file From 11aa5f1509b47fbac98208e0006e5f7363eee6aa Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 20:35:23 +0000 Subject: [PATCH 25/71] fix mask for batched input --- owl_wms/nn/attn.py | 1 + owl_wms/nn/mmattn.py | 1 + 2 files changed, 2 insertions(+) diff --git a/owl_wms/nn/attn.py b/owl_wms/nn/attn.py index ee54ccc6..9bbb381d 100644 --- a/owl_wms/nn/attn.py +++ b/owl_wms/nn/attn.py @@ -55,6 +55,7 @@ def forward(self, x, kv_cache = None): mask = create_block_causal_mask(x.shape[1], self.tokens_per_frame).to(x.device) mask = mask.to(device=x.device,dtype=x.dtype) mask = mask.unsqueeze(0).repeat(x.shape[0], 1, 1) + mask = mask.unsqueeze(1) if kv_cache is not None: old_k, old_v = kv_cache.get(self.layer_ind) diff --git a/owl_wms/nn/mmattn.py b/owl_wms/nn/mmattn.py index ca3b003d..9127e14e 100644 --- a/owl_wms/nn/mmattn.py +++ b/owl_wms/nn/mmattn.py @@ -83,6 +83,7 @@ def forward(self, x_1, x_2, kv_cache=None): mask = create_block_causal_mask_with_mm(x_1.shape[1], x_2.shape[1], self.config.tokens_per_frame) mask = mask.to(device=x_1.device,dtype=x_1.dtype) mask = mask.unsqueeze(0).repeat(x_1.shape[0],1,1) + mask = mask.unsqueeze(1) # head dim if kv_cache is not None: if len(kv_cache) > 0: From 11dfba7322931be6bc8fb6a51eb5d85f2588fa42 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 20:35:55 +0000 Subject: [PATCH 26/71] debug shortcut model --- owl_wms/models/__init__.py | 5 +- owl_wms/models/gamerft_shortcut.py | 98 ++++++++++++++++++------------ 2 files changed, 63 insertions(+), 40 deletions(-) diff --git a/owl_wms/models/__init__.py b/owl_wms/models/__init__.py index 34fcc572..817978ef 100644 --- a/owl_wms/models/__init__.py +++ b/owl_wms/models/__init__.py @@ -1,8 +1,9 @@ -from .gamerft import GameRFT - def get_model_cls(model_id): if model_id == "game_rft": + from .gamerft import GameRFT return GameRFT + if model_id == "game_rft_shortcut": + from .gamerft_shortcut import ShortcutGameRFT diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index 959b2e9e..d028b5bb 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -35,11 +35,15 @@ def __init__(self, config): self.proj_y_in = nn.Linear(config.channels, config.d_model, bias = False) self.pos_enc_y = LearnedPosEnc(config.tokens_per_frame, config.d_model) - def sample(self, x, y, mouse, btn, kv_cache = None, d = None, t = None): + self.config = config + + def sample(self, x, y, mouse, btn, kv_cache = None, t = None, d = None): """ This is a function that largely abstracts away most things for the specific case where you are only generating the one next token + + The return is one step sample always """ if x is None: x = torch.randn_like(y) @@ -50,7 +54,7 @@ def sample(self, x, y, mouse, btn, kv_cache = None, d = None, t = None): if d is None: d = torch.ones_like(x[:,:,0,0,0]) - return self.forward(x, t, mouse, btn, d, kv_cache) + return x - self.forward(x, t, mouse, btn, d, kv_cache) def forward(self, x, y, t, mouse, btn, d, kv_cache = None): # x is [b,n,c,h,w] @@ -60,10 +64,6 @@ def forward(self, x, y, t, mouse, btn, d, kv_cache = None): # mouse is [b,n,2] # btn is [b,n,n_buttons] - print(mouse.shape) - print(btn.shape) - print(x.shape) - ctrl_cond = self.control_embed(mouse, btn) t_cond = self.t_embed(t) d_cond = self.step_embed(d) @@ -88,8 +88,7 @@ def forward(self, x, y, t, mouse, btn, d, kv_cache = None): def sample_discrete_timesteps(steps, eps = 1.0e-6): # steps is Tensor([1,4,2,64,16]) as an example - b = len(steps) - n = len(steps[0]) + b,n = steps.shape ts_list = [] ts = torch.rand(b, n, device=steps.device, dtype=steps.dtype) * (steps - eps) @@ -111,14 +110,17 @@ class ShortcutGameRFT(nn.Module): def __init__(self, config): super().__init__() - self.core = GameRFTCore(config) + self.core = ShortcutGameRFTCore(config) self.cfg_prob = config.cfg_prob self.ema = None self.sc_frac = 0.25 - self.cfg_strength = 1.3 + self.sc_max_steps = 128 + self.cfg_scale = 1.3 + + self.config = config - def set_ema_core(self, ema): + def set_ema(self, ema): if hasattr(ema.ema_model, 'module'): self.ema = ema.ema_model.module.core else: @@ -127,7 +129,7 @@ def set_ema_core(self, ema): @torch.no_grad() @torch.compile() def get_sc_targets(self, x, y, mouse, btn): - steps_slow = sample_steps(len(x), x.device, x.dtype, min_val = 1) + steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) steps_fast = steps_slow / 2 dt_slow = 1./steps_slow @@ -140,14 +142,14 @@ def expand(t): return t[:,:,None,None,None] ts = sample_discrete_timesteps(steps_fast) - cfg_mask = torch.isclose(steps_slow, 128) + cfg_mask = torch.isclose(steps_slow, torch.ones_like(steps_slow)*128) cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] null_mouse = torch.zeros_like(mouse) null_btn = torch.zeros_like(btn) pred_1_uncond = self.ema(x, y, ts, null_mouse, null_btn, steps_slow) - pred_1_cond = self.ema(x, y, ts, mouse, btn) + pred_1_cond = self.ema(x, y, ts, mouse, btn, steps_slow) pred_1_cfg = pred_1_uncond + self.cfg_scale * (pred_1_cond - pred_1_uncond) pred_1 = torch.where(cfg_mask, pred_1_cfg, pred_1_cond) @@ -155,7 +157,7 @@ def expand(t): ts_new = ts - dt_slow pred_2_uncond = self.ema(x_new, y, ts_new, null_mouse, null_btn, steps_slow) - pred_2_cond = self.ema(x_new, y, ts_new, mouse, btn) + pred_2_cond = self.ema(x_new, y, ts_new, mouse, btn, steps_slow) pred_2_cfg = pred_2_uncond + self.cfg_scale * (pred_2_cond - pred_2_uncond) pred_2 = torch.where(cfg_mask, pred_2_cfg, pred_2_cond) @@ -168,7 +170,7 @@ def get_sc_loss(self, x, y, mouse, btn): sc_loss = F.mse_loss(pred, target) return sc_loss - def forward(self, x, y, mouse, btn, return_dict = False, cfg_prob = None): + def forward(self, x, y, mouse, btn): # x is [b,n,c,h,w] # y (seed frame) is [b,1,c,h,w] # mouse is [b,n,2] @@ -182,16 +184,8 @@ def forward(self, x, y, mouse, btn, return_dict = False, cfg_prob = None): mouse,mouse_sc = mouse[:b], mouse[b:] btn,btn_sc = btn[:b], btn[b:] - # Consistency inputs are not autoregressive - n_tok_sc = self.config.n_frames * self.config.tokens_per_frame - x_sc = x_sc[:,:self.config.n_frames] - mouse_sc = mouse_sc[:,:self.config.n_frames] - btn_sc = btn_sc[:,:self.config.n_frames] - # Apply classifier-free guidance dropout - if cfg_prob is None: - cfg_prob = self.cfg_prob - if cfg_prob > 0.0: + if self.cfg_prob > 0.0: mask = torch.rand(b, device=x.device) <= self.cfg_prob null_mouse = torch.zeros_like(mouse) null_btn = torch.zeros_like(btn) @@ -201,7 +195,7 @@ def forward(self, x, y, mouse, btn, return_dict = False, cfg_prob = None): btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) with torch.no_grad(): - d = torch.ones_like(x[:,:,0,0,0]) + d = torch.ones_like(x[:,:,0,0,0])*self.sc_max_steps ts = sample_discrete_timesteps(d) ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() @@ -215,18 +209,9 @@ def forward(self, x, y, mouse, btn, return_dict = False, cfg_prob = None): diff_loss = F.mse_loss(pred, target) sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc) - if not return_dict: - return diff_loss - else: - return { - 'diffusion_loss' : diff_loss, - 'lerpd' : lerpd, - 'pred' : pred, - 'ts': ts, - 'z': z - } + return diff_loss, sc_loss -if __name__ == "__main__": +def test_inference_cache(): from ..configs import TransformerConfig from ..nn.kv_cache import KVCache @@ -284,4 +269,41 @@ def forward(self, x, y, mouse, btn, return_dict = False, cfg_prob = None): new_d = torch.full((1, 1), 4, device='cuda', dtype=torch.bfloat16) new_out = model(new_x, y, new_t, new_mouse, new_btn, new_d, cache) print(f"Final cache length: {len(cache)}") - print(f"Final cache shape: {cache.cache[0][0].shape}") \ No newline at end of file + print(f"Final cache shape: {cache.cache[0][0].shape}") + +def test_wrapper(): + from ..configs import TransformerConfig + from ema_pytorch import EMA + from copy import deepcopy + + cfg = TransformerConfig( + None, # model_id + 6, # n_layers + 6, # n_heads + 384, # d_model + 1, # patch_size + 128, # channels + 16, # sample_size + 0.1, # cfg_prob + 11, # n_buttons + 16, # tokens_per_frame + 10, # n_frames + True # causal + ) + + model = ShortcutGameRFT(cfg).bfloat16().cuda() + ema = EMA(model, beta=0.999,update_after_step=0,update_every=1) + model.set_ema(ema) + + NUM_FRAMES = 10 + x = torch.randn(4, NUM_FRAMES, 128, 4, 4).bfloat16().cuda() + y = torch.randn(4, 1, 128, 4, 4).bfloat16().cuda() + mouse = torch.randn(4, NUM_FRAMES, 2).bfloat16().cuda() + btn = torch.randn(4, NUM_FRAMES, 11).bfloat16().cuda() + + with torch.no_grad(): + loss_1, loss_2 = model(x, y, mouse, btn) + print(loss_1, loss_2) + +if __name__ == "__main__": + test_wrapper() \ No newline at end of file From 666e20eeca821cc0ced261114a4366aa926cbad4 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 20:40:31 +0000 Subject: [PATCH 27/71] add shortcut sampler without cache --- owl_wms/sampling/shortcut_sampler.py | 84 +++++++++++++++++++++++++--- 1 file changed, 77 insertions(+), 7 deletions(-) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index 13b07c57..0d1be95a 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -17,22 +17,18 @@ def zlerp(x, alpha): class CacheShortcutSampler: """ - Window CFG Sampler samples new frames one by one, by inpainting the final frame. - This is basically diffusion forcing. + Shortcut CFG sampler builds cache with 4 step diffusion. + Samples new frames in 1 step. - :param n_steps: Number of diffusion steps for each frame (diffusoin steps) - :param cfg_scale: CFG scale for each frame :param window_length: Number of frames to use for each frame generation step :param num_frames: Number of new frames to sample - :param noise_prev: Noise previous frame :param only_return_generated: Whether to only return the generated frames """ - def __init__(self, window_length = 60, num_frames = 60, noise_prev = 0.25, only_return_generated = False): + def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): self.n_steps = n_steps self.cfg_scale = cfg_scale self.window_length = window_length self.num_frames = num_frames - self.noise_prev = noise_prev self.only_return_generated = only_return_generated @torch.no_grad() @@ -82,7 +78,81 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale new_frames = torch.cat(new_frames, dim = 1) x = torch.cat([history,new_frames], dim = 1) + + if self.only_return_generated: + x = x[:,-num_frames:] + extended_mouse = extended_mouse[:,-num_frames:] + extended_btn = extended_btn[:,-num_frames:] + + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + + return x, extended_mouse, extended_btn + +class WindowShortcutSampler: + """ + Same as above but with no cache + + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale = 1): + # dummy_batch is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + history = history[:,:self.window_length] + new_frames = [] + alpha = 0.25 # This number is special for our sampler + + # Extended fake controls to use during sampling + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, num_frames + self.window_length) + + # Generate cache over history + noisy_history = zlerp(history.clone(), alpha) + ts = torch.ones_like(noisy_history[:,:,0,0,0]) * alpha + d = torch.ones_like(noisy_history[:,:,0,0,0]) * round(1./alpha) + ts_single = ts[:,0].unsqueeze(1) + d_single = d[:,0].unsqueeze(1) + + cache = KVCache(model.config) + cache.reset(history.shape[0]) + + cache.enable_cache_updates() + _ = model.sample(noisy_history, keyframe, mouse, btn, cache, ts, d) + cache.disable_cache_updates() + + # Cache is now built! + for frame_idx in tqdm(range(num_frames)): + cache.truncate(1) # Drop first frame + + # Generate new frame + cache.disable_cache_updates() + mouse = extended_mouse[:,self.window_length+frame_idx].unsqueeze(1) + btn = extended_btn[:,self.window_length+frame_idx].unsqueeze(1) + new_frame = model.sample(None, keyframe, mouse, btn, cache) # [b,1,c,h,w] + new_frames.append(new_frame) + + # Add that frame to the cache + cache.enable_cache_updates() + new_frame_noisy = zlerp(new_frame, alpha) + _ = model.sample(new_frame_noisy, keyframe, mouse, btn, cache, ts_single, d_single) + + new_frames = torch.cat(new_frames, dim = 1) + x = torch.cat([history,new_frames], dim = 1) + if self.only_return_generated: x = x[:,-num_frames:] extended_mouse = extended_mouse[:,-num_frames:] From 511305e3eed2b72e9e678d31a54d0a524b81b0ff Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 20:40:54 +0000 Subject: [PATCH 28/71] add shortcut sampler without cache --- owl_wms/sampling/shortcut_sampler.py | 67 ++++++++++++++-------------- 1 file changed, 34 insertions(+), 33 deletions(-) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index 0d1be95a..2edc310d 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -107,7 +107,7 @@ def __init__(self, window_length = 60, num_frames = 60, only_return_generated = @torch.no_grad() def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale = 1): - # dummy_batch is [b,n,c,h,w] + # history is [b,n,c,h,w] # mouse is [b,n,2] # btn is [b,n,n_button] @@ -117,49 +117,50 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale alpha = 0.25 # This number is special for our sampler # Extended fake controls to use during sampling - extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, num_frames + self.window_length) - - # Generate cache over history - noisy_history = zlerp(history.clone(), alpha) - ts = torch.ones_like(noisy_history[:,:,0,0,0]) * alpha - d = torch.ones_like(noisy_history[:,:,0,0,0]) * round(1./alpha) - ts_single = ts[:,0].unsqueeze(1) - d_single = d[:,0].unsqueeze(1) + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, self.num_frames + self.window_length) - cache = KVCache(model.config) - cache.reset(history.shape[0]) + # Initialize window history + window_history = history.clone() - cache.enable_cache_updates() - _ = model.sample(noisy_history, keyframe, mouse, btn, cache, ts, d) - cache.disable_cache_updates() + for frame_idx in tqdm(range(self.num_frames)): + # Setup window history + x = window_history[:,-self.window_length:].clone() + + # Noise all but last frame to alpha + x[:,:-1] = zlerp(x[:,:-1], alpha) + # Last frame starts as random noise + x[:,-1] = torch.randn_like(x[:,-1]) + + # Setup timesteps - alpha for context, 1.0 for generated + ts = torch.ones_like(x[:,:,0,0,0]) + ts[:,:-1] = alpha + + # Setup diffusion steps - 4 for context, 1 for generated + d = torch.ones_like(x[:,:,0,0,0]) + d[:,:-1] = 4 - # Cache is now built! - - for frame_idx in tqdm(range(num_frames)): - cache.truncate(1) # Drop first frame + # Get current controls + curr_mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] + curr_btn = extended_btn[:,frame_idx:frame_idx+self.window_length] # Generate new frame - cache.disable_cache_updates() - mouse = extended_mouse[:,self.window_length+frame_idx].unsqueeze(1) - btn = extended_btn[:,self.window_length+frame_idx].unsqueeze(1) - new_frame = model.sample(None, keyframe, mouse, btn, cache) # [b,1,c,h,w] + pred = model.sample(x, keyframe, curr_mouse, curr_btn, None, ts, d) + new_frame = pred[:,-1:] # Take only the last frame new_frames.append(new_frame) - # Add that frame to the cache - cache.enable_cache_updates() - new_frame_noisy = zlerp(new_frame, alpha) - _ = model.sample(new_frame_noisy, keyframe, mouse, btn, cache, ts_single, d_single) + # Add new frame to window history + window_history = torch.cat([window_history, new_frame], dim=1) - new_frames = torch.cat(new_frames, dim = 1) - x = torch.cat([history,new_frames], dim = 1) + new_frames = torch.cat(new_frames, dim=1) + x = torch.cat([history, new_frames], dim=1) if self.only_return_generated: - x = x[:,-num_frames:] - extended_mouse = extended_mouse[:,-num_frames:] - extended_btn = extended_btn[:,-num_frames:] + x = x[:,-self.num_frames:] + extended_mouse = extended_mouse[:,-self.num_frames:] + extended_btn = extended_btn[:,-self.num_frames:] if decode_fn is not None: - x = x * scale + x = x * scale x = decode_fn(x) - + return x, extended_mouse, extended_btn \ No newline at end of file From 1126c7ff5c5276bc6b03d1dd175ba1305d477fe5 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 22:37:31 +0000 Subject: [PATCH 29/71] switch default to rope, remove other pos encs --- owl_wms/models/gamerft_shortcut.py | 3 - owl_wms/nn/attn.py | 2 - owl_wms/nn/embeddings.py | 74 ----------------------- owl_wms/nn/mmattn.py | 10 +++- owl_wms/nn/rope.py | 96 ++++++++++++++++++++++++++++++ 5 files changed, 104 insertions(+), 81 deletions(-) create mode 100644 owl_wms/nn/rope.py diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index d028b5bb..38f11b7b 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -30,8 +30,6 @@ def __init__(self, config): self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) - self.pos_enc = LearnedPosEnc(config.tokens_per_frame * config.n_frames, config.d_model) - self.proj_y_in = nn.Linear(config.channels, config.d_model, bias = False) self.pos_enc_y = LearnedPosEnc(config.tokens_per_frame, config.d_model) @@ -75,7 +73,6 @@ def forward(self, x, y, t, mouse, btn, d, kv_cache = None): y = eo.rearrange(y, 'b n c h w -> b (n h w) c') x = self.proj_in(x) - x = self.pos_enc(x) y = self.proj_y_in(y) y = self.pos_enc_y(y) diff --git a/owl_wms/nn/attn.py b/owl_wms/nn/attn.py index 9bbb381d..6174b72a 100644 --- a/owl_wms/nn/attn.py +++ b/owl_wms/nn/attn.py @@ -8,8 +8,6 @@ import einops as eo from .modulation import AdaLN, Gate -#from .embeddings import FlatVideoRoPE -from rotary_embedding_torch import RotaryEmbedding torch.backends.cuda.enable_flash_sdp(enabled = True) diff --git a/owl_wms/nn/embeddings.py b/owl_wms/nn/embeddings.py index adbfd32a..c10ebddd 100644 --- a/owl_wms/nn/embeddings.py +++ b/owl_wms/nn/embeddings.py @@ -13,80 +13,6 @@ ) import einops as eo -class VideoRoPE(nn.Module): - """ - Video RoPE embedding for when latents are 3D [n,h,w] - """ - def __init__(self, config : 'TransformerConfig'): - super().__init__() - - dim_head = config.d_model // config.n_heads - self.pos_emb = RotaryEmbedding( - dim = dim_head//8, - freqs_for = 'pixel', - max_freq = 256 - ) - n_patches = config.sample_size // config.patch_size - self.tokens_per_frame = n_patches**2 - - self.rearrange_in = lambda x: eo.rearrange(x, 'b h (n_t n_y n_x) d -> b h n_t n_y n_x d', n_y = n_patches) - self.rearrange_out = lambda x: eo.rearrange(x, 'b h n_t n_y n_x d -> b h (n_t n_y n_x) d') - self.get_freqs = lambda n_t: self.pos_emb.get_axial_freqs(n_t, n_patches, n_patches) - - def forward(self, q, k): - # q k both [b,h,n,d] - q = self.rearrange_in(q) - k = self.rearrange_in(k) - - n_t = q.shape[2] - freqs = self.get_freqs(n_t) - - q = apply_rotary_emb(freqs.float(), q.float()).to(q.dtype) - k = apply_rotary_emb(freqs.float(), k.float()).to(k.dtype) - - q = self.rearrange_out(q) - k = self.rearrange_out(k) - - return q, k - -class FlatVideoRoPE(nn.Module): - """ - Video RoPE embedding for when latents are 2d [n,m] (1D Frame Tokenization) - """ - def __init__(self, config : 'TransformerConfig'): - super().__init__() - - dim_head = config.d_model // config.n_heads - self.pos_emb = RotaryEmbedding( - dim = dim_head//4, - freqs_for = 'pixel', - max_freq = 256 - ) - self.pos_emb.freqs.requires_grad = False - self.tokens_per_frame = config.sample_size - - self.rearrange_in = lambda x: eo.rearrange(x, 'b h (n_t m) d -> b h n_t m d', m = self.tokens_per_frame) - self.rearrange_out = lambda x: eo.rearrange(x, 'b h n_t m d -> b h (n_t m) d') - self.get_freqs = lambda n_t: self.pos_emb.get_axial_freqs(n_t, self.tokens_per_frame) - - def forward(self, q, k): - # q k both [b,h,n,d] - q = self.rearrange_in(q) - k = self.rearrange_in(k) - - n_t = q.shape[2] - with torch.no_grad(): - freqs = self.get_freqs(n_t) - - q = apply_rotary_emb(freqs.float(), q.float()).to(q.dtype) - k = apply_rotary_emb(freqs.float(), k.float()).to(k.dtype) - - q = self.rearrange_out(q) - k = self.rearrange_out(k) - - return q, k - - class LearnedPosEnc(nn.Module): def __init__(self, n_seq, dim): super().__init__() diff --git a/owl_wms/nn/mmattn.py b/owl_wms/nn/mmattn.py index 9127e14e..7ed4948e 100644 --- a/owl_wms/nn/mmattn.py +++ b/owl_wms/nn/mmattn.py @@ -8,8 +8,7 @@ import einops as eo from .modulation import AdaLN, Gate -#from .embeddings import FlatVideoRoPE -from rotary_embedding_torch import RotaryEmbedding +from .rope import FlatVideoRoPE torch.backends.cuda.enable_flash_sdp(enabled = True) @@ -62,6 +61,8 @@ def __init__(self, config : 'TransformerConfig'): self.config = config self.causal = config.causal + self.rope = FlatVideoRoPE(config) + def split(self, qkv): return eo.rearrange(qkv, 'b n (three h d) -> three b h n d', three = 3, h = self.n_heads) @@ -98,6 +99,8 @@ def forward(self, x_1, x_2, kv_cache=None): if kv_cache.should_update: kv_cache.update(new_k, new_v, self.layer_ind) + q1, new_k = self.rope(q1, new_k) + k = torch.cat([new_k, k2], dim=-2) v = torch.cat([new_v, v2], dim=-2) q = torch.cat([q1, q2], dim=-2) @@ -106,10 +109,13 @@ def forward(self, x_1, x_2, kv_cache=None): x = x[:,:,-q.shape[2]:] # Only keep latest outputs x = self.merge(x) else: + q1, k1 = self.rope(q1,k1) + q = torch.cat([q1,q2],dim=-2) k = torch.cat([k1,k2],dim=-2) v = torch.cat([v1,v2],dim=-2) + print("====") x = F.scaled_dot_product_attention(q,k,v, attn_mask = mask) x = self.merge(x) diff --git a/owl_wms/nn/rope.py b/owl_wms/nn/rope.py new file mode 100644 index 00000000..a2b9ec76 --- /dev/null +++ b/owl_wms/nn/rope.py @@ -0,0 +1,96 @@ +""" +Variants of RoPE were becoming heavy for embeddings so +I made a unique script for all of them here +""" + +from rotary_embedding_torch import ( + RotaryEmbedding, + apply_rotary_emb +) +import einops as eo +import torch +from torch import nn + +class VideoRoPE(nn.Module): + """ + Video RoPE embedding for when latents are 3D [n,h,w] + """ + def __init__(self, config : 'TransformerConfig'): + super().__init__() + + dim_head = config.d_model // config.n_heads + self.pos_emb = RotaryEmbedding( + dim = dim_head//8, + freqs_for = 'pixel', + max_freq = 256 + ) + n_patches = config.sample_size // config.patch_size + self.tokens_per_frame = n_patches**2 + + self.rearrange_in = lambda x: eo.rearrange(x, 'b h (n_t n_y n_x) d -> b h n_t n_y n_x d', n_y = n_patches) + self.rearrange_out = lambda x: eo.rearrange(x, 'b h n_t n_y n_x d -> b h (n_t n_y n_x) d') + self.get_freqs = lambda n_t: self.pos_emb.get_axial_freqs(n_t, n_patches, n_patches) + + def forward(self, q, k): + # q k both [b,h,n,d] + q = self.rearrange_in(q) + k = self.rearrange_in(k) + + n_t = q.shape[2] + freqs = self.get_freqs(n_t) + + q = apply_rotary_emb(freqs.float(), q.float()).to(q.dtype) + k = apply_rotary_emb(freqs.float(), k.float()).to(k.dtype) + + q = self.rearrange_out(q) + k = self.rearrange_out(k) + + return q, k + +class FlatVideoRoPE(nn.Module): + """ + Half-flat of RoPE that treats [n_frames, tokens_per_frame] as [n_frames, tokens_per_frame] image + """ + def __init__(self, config): + super().__init__() + + dim_head = config.d_model // config.n_heads + self.pos_emb = RotaryEmbedding( + dim = dim_head//4, + freqs_for='pixel', + max_freq=256 + ) + + self.m = config.tokens_per_frame + + def pad_q(self, q, k): + # Pad Q when it's needed for kv caching + q_len = q.shape[2] + k_len = k.shape[2] + + def forward(self, q, k): + # q|k is [b,h,n_frames*tokens_per_frame,d] + n = k.shape[2]//self.m + m = self.m + + truncate = n + if q.shape[2] < n * m: + truncate = q.shape[2]//m # How many frames is q? + + q = eo.rearrange(q, 'b h (n m) d -> b h n m d', n=n,m=m) + k = eo.rearrange(k, 'b h (n m) d -> b h n m d', n=n,m=m) + + freqs = self.pos_emb.get_axial_freqs(n,m) + q = apply_rotary_emb(freqs[-truncate:], q) + k = apply_rotary_emb(freqs, k) + + q = eo.rearrange(q, 'b h n m d -> b h (n m) d') + k = eo.rearrange(k, 'b h n m d -> b h (n m) d') + + if truncate is not None: + q = q[:,:,-truncate*m:] + + return q,k + + + From 98e12a6d2d30ae872f29b1b3c5c00e3ccb2ef1ff Mon Sep 17 00:00:00 2001 From: shahbuland Date: Tue, 10 Jun 2025 22:56:36 +0000 Subject: [PATCH 30/71] fix mask, add new trainer to init --- owl_wms/nn/mmattn.py | 9 +++++---- owl_wms/trainers/__init__.py | 5 ++++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/owl_wms/nn/mmattn.py b/owl_wms/nn/mmattn.py index 7ed4948e..66b0ecc2 100644 --- a/owl_wms/nn/mmattn.py +++ b/owl_wms/nn/mmattn.py @@ -33,11 +33,12 @@ def create_block_causal_mask_with_mm(tokens, context_tokens, tokens_per_frame): end = (i + 1) * tokens_per_frame mask[start:end, end:tokens] = True # Can't see future frames - # Context tokens can attend to everything (no masking needed) - # Regular tokens can attend to all context tokens (no masking needed) - # The zeros in mask[tokens:, :] allow context to attend to everything + # Context tokens can only attend to themselves + mask[tokens:, :tokens] = True # Mask out attention to regular tokens + + # Regular tokens can still attend to all context tokens (no masking needed) # The zeros in mask[:, tokens:] allow tokens to attend to all context - + return mask class MMAttn(nn.Module): diff --git a/owl_wms/trainers/__init__.py b/owl_wms/trainers/__init__.py index 20d7ac4c..0a78fb14 100644 --- a/owl_wms/trainers/__init__.py +++ b/owl_wms/trainers/__init__.py @@ -5,4 +5,7 @@ def get_trainer_cls(trainer_id): return RFTTrainer if trainer_id == "causvid": from .causvid import CausVidTrainer - return CausVidTrainer \ No newline at end of file + return CausVidTrainer + if trainer_id == "shortcut": + from .shortcut_trainer import ShortcutTrainer + return ShortcutTrainer \ No newline at end of file From cc2be00469927487376e2f340a628ce6cce28e8b Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:11:04 +0000 Subject: [PATCH 31/71] remove wrong comment --- owl_wms/trainers/gamerft_trainer.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/owl_wms/trainers/gamerft_trainer.py b/owl_wms/trainers/gamerft_trainer.py index 0b6cfe93..6a3dd120 100644 --- a/owl_wms/trainers/gamerft_trainer.py +++ b/owl_wms/trainers/gamerft_trainer.py @@ -1,7 +1,3 @@ -""" -Trainer for reconstruction only -""" - import torch from ema_pytorch import EMA import wandb From f0f688227e0746f38489720905200262fde9dd6f Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:11:29 +0000 Subject: [PATCH 32/71] add shortcut trainer --- owl_wms/trainers/shortcut_trainer.py | 202 +++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 owl_wms/trainers/shortcut_trainer.py diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py new file mode 100644 index 00000000..37d7401b --- /dev/null +++ b/owl_wms/trainers/shortcut_trainer.py @@ -0,0 +1,202 @@ +import torch +from ema_pytorch import EMA +import wandb +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist +import einops as eo + +from .base import BaseTrainer + +from ..utils import freeze, Timer, find_unused_params +from ..schedulers import get_scheduler_cls +from ..models import get_model_cls +from ..sampling import get_sampler_cls +from ..data import get_loader +from ..utils.logging import LogHelper, to_wandb +from ..muon import init_muon +from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn + +class ShortcutTrainer(BaseTrainer): + """ + Trainer for rectified flow transformer with shortcut + + :param train_cfg: Configuration for training + :param logging_cfg: Configuration for logging + :param model_cfg: Configuration for model + :param global_rank: Rank across all devices. + :param local_rank: Rank for current device on this process. + :param world_size: Overall number of devices + """ + def __init__(self,*args,**kwargs): + super().__init__(*args,**kwargs) + + model_id = self.model_cfg.model_id + self.model = get_model_cls(model_id)(self.model_cfg) + + # Print model size + if self.rank == 0: + n_params = sum(p.numel() for p in self.model.parameters()) + print(f"Model has {n_params:,} parameters") + + self.ema = None + self.opt = None + self.scheduler = None + self.scaler = None + + self.total_step_counter = 0 + self.decoder = get_decoder_only( + self.train_cfg.vae_id, + self.train_cfg.vae_cfg_path, + self.train_cfg.vae_ckpt_path + ) + + freeze(self.decoder) + + def save(self): + save_dict = { + 'model' : self.model.state_dict(), + 'ema' : self.ema.state_dict(), + 'opt' : self.opt.state_dict(), + 'scaler' : self.scaler.state_dict(), + 'steps': self.total_step_counter + } + if self.scheduler is not None: + save_dict['scheduler'] = self.scheduler.state_dict() + super().save(save_dict) + + def load(self): + has_ckpt = False + try: + if self.train_cfg.resume_ckpt is not None: + save_dict = super().load(self.train_cfg.resume_ckpt) + has_ckpt = True + except: + print("Error loading checkpoint") + + if not has_ckpt: + return + + + self.model.load_state_dict(save_dict['model']) + self.ema.load_state_dict(save_dict['ema']) + self.opt.load_state_dict(save_dict['opt']) + if self.scheduler is not None and 'scheduler' in save_dict: + self.scheduler.load_state_dict(save_dict['scheduler']) + self.scaler.load_state_dict(save_dict['scaler']) + self.total_step_counter = save_dict['steps'] + + def train(self): + torch.cuda.set_device(self.local_rank) + + # Prepare model and ema + self.model = self.model.cuda().train() + if self.world_size > 1: + self.model = DDP(self.model, device_ids=[self.local_rank]) + self.decoder = self.decoder.cuda().eval().bfloat16() + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + + self.ema = EMA( + self.model, + beta = 0.999, + update_after_step = 0, + update_every = 1 + ) + #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) + + def get_ema_core(): + if self.world_size > 1: + return self.ema.ema_model.module.core + else: + return self.ema.ema_model.core + + # No muon pls + self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) + + if self.train_cfg.scheduler is not None: + self.scheduler = get_scheduler_cls(self.train_cfg.scheduler)(self.opt, **self.train_cfg.scheduler_kwargs) + + # Grad accum setup and scaler + accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size + accum_steps = max(1, accum_steps) + self.scaler = torch.amp.GradScaler() + ctx = torch.amp.autocast('cuda',torch.bfloat16) + + self.load() + + # Timer reset + timer = Timer() + timer.reset() + metrics = LogHelper() + if self.rank == 0: + wandb.watch(self.get_module(), log = 'all') + + # Dataset setup + loader = get_loader(self.train_cfg.data_id, self.train_cfg.batch_size, **self.train_cfg.data_kwargs) + sampler = get_sampler_cls(self.train_cfg.sampler_id)(**self.train_cfg.sampler_kwargs) + + local_step = 0 + for _ in range(self.train_cfg.epochs): + for batch_vid, batch_keyframe, batch_mouse, batch_btn in loader: + batch_vid = batch_vid.cuda().bfloat16() / self.train_cfg.vae_scale + batch_keyframe = batch_keyframe.cuda().bfloat16() + batch_mouse = batch_mouse.cuda().bfloat16() + batch_btn = batch_btn.cuda().bfloat16() + + with ctx: + diff_loss, sc_loss = self.model(batch_vid,batch_mouse,batch_btn) / accum_steps + loss = diff_loss + sc_loss + + self.scaler.scale(loss).backward() + #find_unused_params(self.model) + + metrics.log('diffusion_loss', diff_loss) + metrics.log('shortcut_loss', sc_loss) + + local_step += 1 + if local_step % accum_steps == 0: + # Updates + self.scaler.unscale_(self.opt) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + + self.scaler.step(self.opt) + self.opt.zero_grad(set_to_none=True) + + self.scaler.update() + + if self.scheduler is not None: + self.scheduler.step() + self.ema.update() + + # Do logging + with torch.no_grad(): + wandb_dict = metrics.pop() + wandb_dict['time'] = timer.hit() + wandb_dict['lr'] = self.opt.param_groups[0]['lr'] + timer.reset() + + # Sampling commented out for now + if self.total_step_counter % self.train_cfg.sample_interval == 0: + with ctx, torch.no_grad(): + n_samples = self.train_cfg.n_samples + samples, sample_mouse, sample_button = sampler( + get_ema_core(), + batch_vid[:n_samples], + batch_keyframe[:n_samples], + batch_mouse[:n_samples], + batch_btn[:n_samples], + decode_fn = decode_fn, + scale=self.train_cfg.vae_scale + ) # -> [b,n,c,h,w] + if self.rank == 0: wandb_dict['samples'] = to_wandb(samples, sample_mouse, sample_button) + + + if self.rank == 0: + wandb.log(wandb_dict) + + self.total_step_counter += 1 + if self.total_step_counter % self.train_cfg.save_interval == 0: + if self.rank == 0: + self.save() + + self.barrier() \ No newline at end of file From 1286ec091f228c77e825760542bfa61dcc0ea315 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:13:39 +0000 Subject: [PATCH 33/71] add shortcut config --- configs/shortcut.yml | 64 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 configs/shortcut.yml diff --git a/configs/shortcut.yml b/configs/shortcut.yml new file mode 100644 index 00000000..4e42b119 --- /dev/null +++ b/configs/shortcut.yml @@ -0,0 +1,64 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: game_rft_shortcut + sample_size: 4 + channels: 128 + + n_layers: 17 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 16 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 120 + + causal: false + +train: + trainer_id: shortcut + data_id: cod_s3 + data_kwargs: + window_length: 120 + root: ../cod_data/BlackOpsColdWar + include_keyframe: true + + target_batch_size: 256 + batch_size: 16 + + epochs: 200 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + weight_decay: 0.1 + eps: 1.0e-15 + betas: [0.9, 0.95] + + scheduler: null + + checkpoint_dir: checkpoints/360p + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: shortcut_sampler + sampler_kwargs: + window_length: 120 + num_frames: 120 + only_return_generated: true + + n_samples: 8 + + vae_id: 720pr3dc + vae_batch_size: 4 + vae_scale: 0.35 + vae_cfg_path: configs/owl_vaes/128x_cod_stage2.yml + vae_ckpt_path: 720p_cod_vae_30m_35k_steps.pt + +wandb: + name: shahbuland + project: video_models + run_name: v2 \ No newline at end of file From c81b012b5df3f7c9d49ed9f32801a1f9e3752cdd Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:16:42 +0000 Subject: [PATCH 34/71] add rft shortcut --- owl_wms/models/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/owl_wms/models/__init__.py b/owl_wms/models/__init__.py index 817978ef..5f70facd 100644 --- a/owl_wms/models/__init__.py +++ b/owl_wms/models/__init__.py @@ -4,6 +4,7 @@ def get_model_cls(model_id): return GameRFT if model_id == "game_rft_shortcut": from .gamerft_shortcut import ShortcutGameRFT + return ShortcutGameRFT From 080061482029f98850110a36bf080e2c8034246e Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:21:20 +0000 Subject: [PATCH 35/71] add bucket name to config --- configs/shortcut.yml | 4 ++-- owl_wms/data/s3_cod_latent.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/configs/shortcut.yml b/configs/shortcut.yml index 4e42b119..acd4441b 100644 --- a/configs/shortcut.yml +++ b/configs/shortcut.yml @@ -22,7 +22,7 @@ train: data_id: cod_s3 data_kwargs: window_length: 120 - root: ../cod_data/BlackOpsColdWar + bucket_name: cod-data-latent-360x640to5x8 include_keyframe: true target_batch_size: 256 @@ -44,7 +44,7 @@ train: sample_interval: 1000 save_interval: 5000 - sampler_id: shortcut_sampler + sampler_id: shortcut sampler_kwargs: window_length: 120 num_frames: 120 diff --git a/owl_wms/data/s3_cod_latent.py b/owl_wms/data/s3_cod_latent.py index 8fdcdddf..f7eb59bb 100644 --- a/owl_wms/data/s3_cod_latent.py +++ b/owl_wms/data/s3_cod_latent.py @@ -33,7 +33,7 @@ def pop(self): BUCKET_NAME="cod-data-latent-360x640to5x8" class S3CoDLatentDataset(IterableDataset): - def __init__(self, window_length=120, file_share_max=20, rank=0, world_size=1, include_keyframe = False): + def __init__(self, window_length=120, file_share_max=20, rank=0, world_size=1, bucket_name = BUCKET_NAME, include_keyframe = False): super().__init__() self.window = window_length @@ -41,6 +41,7 @@ def __init__(self, window_length=120, file_share_max=20, rank=0, world_size=1, i self.rank = rank self.world_size = world_size self.include_keyframe = include_keyframe + self.bucket_name = bucket_name # Queue parameters self.max_tars = 2 @@ -80,7 +81,7 @@ def background_download_tars(self): tar_path = self.random_sample_prefix() try: # Download tar directly to memory - response = self.s3_client.get_object(Bucket=BUCKET_NAME, Key=tar_path) + response = self.s3_client.get_object(Bucket=self.bucket_name, Key=tar_path) tar_data = response['Body'].read() self.tar_queue.add(tar_data) except Exception as e: From 91b7b8a144b67f9ddec2852c1bb391b75d931ee7 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:22:56 +0000 Subject: [PATCH 36/71] remove redundant args --- owl_wms/sampling/shortcut_sampler.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index 2edc310d..b41a0bc4 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -25,8 +25,6 @@ class CacheShortcutSampler: :param only_return_generated: Whether to only return the generated frames """ def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): - self.n_steps = n_steps - self.cfg_scale = cfg_scale self.window_length = window_length self.num_frames = num_frames self.only_return_generated = only_return_generated From e3edba3c12467266e532d2eea6ed28765a6cfbed Mon Sep 17 00:00:00 2001 From: Sami Date: Tue, 10 Jun 2025 20:47:44 -0400 Subject: [PATCH 37/71] merge --- owl_wms/models/gamerft_shortcut.py | 309 +++++++++++++++++++++++++ owl_wms/nn/kv_cache.py | 22 +- owl_wms/sampling/shortcut_sampler.py | 326 +++++++++++++++++++++++++++ owl_wms/utils/ddp.py | 2 +- 4 files changed, 657 insertions(+), 2 deletions(-) create mode 100644 owl_wms/models/gamerft_shortcut.py create mode 100644 owl_wms/sampling/shortcut_sampler.py diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py new file mode 100644 index 00000000..866bf1d0 --- /dev/null +++ b/owl_wms/models/gamerft_shortcut.py @@ -0,0 +1,309 @@ +""" +Causal-First RFT With Shortcut objective +""" + +import torch +from torch import nn +import torch.nn.functional as F + +import einops as eo + +from ..nn.embeddings import ( + TimestepEmbedding, + StepEmbedding, + ControlEmbedding, + LearnedPosEnc +) +from ..nn.attn import UViT, FinalLayer +from ..nn.mmattn import MMUViT + +class ShortcutGameRFTCore(nn.Module): + def __init__(self, config): + super().__init__() + + self.transformer = MMUViT(config) + self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) + + self.step_embed = StepEmbedding(config.d_model) + self.t_embed = TimestepEmbedding(config.d_model) + + self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) + self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) + + self.pos_enc = LearnedPosEnc(config.tokens_per_frame * config.n_frames, config.d_model) + + self.proj_y_in = nn.Linear(config.channels, config.d_model, bias = False) + self.pos_enc_y = LearnedPosEnc(config.tokens_per_frame, config.d_model) + + self.config = config + + def sample(self, x, y, mouse, btn, kv_cache = None, t = None, d = None): + """ + This is a function that largely abstracts + away most things for the specific case where + you are only generating the one next token + + The return is one step sample always + """ + if x is None: + x = torch.randn_like(y) + + b,n,c,h,w = x.shape + if t is None: + t = torch.ones_like(x[:,:,0,0,0]) + if d is None: + d = torch.ones_like(x[:,:,0,0,0]) + + return x - self.forward(x, y, t, mouse, btn, d, kv_cache) + + def forward(self, x, y, t, mouse, btn, d, kv_cache = None): + # x is [b,n,c,h,w] + # y is [b,1,c,h,w] + # t is [b,n] + # d is [b,n] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + + ctrl_cond = self.control_embed(mouse, btn) + t_cond = self.t_embed(t) + d_cond = self.step_embed(d) + + cond = ctrl_cond + t_cond + d_cond # [b,n,d] + + b,n,c,h,w = x.shape + x = eo.rearrange(x, 'b n c h w -> b (n h w) c') + y = eo.rearrange(y, 'b n c h w -> b (n h w) c') + + x = self.proj_in(x) + x = self.pos_enc(x) + + y = self.proj_y_in(y) + y = self.pos_enc_y(y) + + x = self.transformer(x, y, cond, kv_cache) + x = self.proj_out(x, cond) # -> [b,n*hw,c] + x = eo.rearrange(x, 'b (n h w) c -> b n c h w', n=n,h=h,w=w) + + return x + +def sample_discrete_timesteps(steps, eps = 1.0e-6): + # steps is Tensor([1,4,2,64,16]) as an example + b,n = steps.shape + + ts_list = [] + ts = torch.rand(b, n, device=steps.device, dtype=steps.dtype) * (steps - eps) + ts = ts.clamp(eps).ceil() / steps + """ + Example, if d was all 2, ts would be [0,2] + so do clamp, then ceil will be 1 or 2 (0, 2] + then do t / 2 and get 0.5 or 1.0, our desired timesteps + """ + return ts + +def sample_steps(b, n, device, dtype, min_val = 0): + valid = torch.tensor([2**i for i in range(min_val, 8)]) # [1,2,...,128] + inds = torch.randint(low=0,high=len(valid), size = (b,n)) + steps = valid[inds].to(device=device,dtype=dtype) + return steps + +class ShortcutGameRFT(nn.Module): + def __init__(self, config): + super().__init__() + + self.core = ShortcutGameRFTCore(config) + self.cfg_prob = config.cfg_prob + + self.ema = None + self.sc_frac = 0.25 + self.sc_max_steps = 128 + self.cfg_scale = 1.3 + + self.config = config + + def set_ema(self, ema): + if hasattr(ema.ema_model, 'module'): + self.ema = ema.ema_model.module.core + else: + self.ema = ema.ema_model.core + + @torch.no_grad() + @torch.compile() + def get_sc_targets(self, x, y, mouse, btn): + steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) + steps_fast = steps_slow / 2 + + dt_slow = 1./steps_slow + dt_fast = 1./steps_fast + + def expand(t): + #b,c,h,w = x.shape + #t = eo.repeat(t,'b -> b c h w',c=c,h=h,w=w) + #return t + return t[:,:,None,None,None] + + ts = sample_discrete_timesteps(steps_fast) + cfg_mask = torch.isclose(steps_slow, torch.ones_like(steps_slow)*128) + cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] + + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + pred_1_uncond = self.ema(x, y, ts, null_mouse, null_btn, steps_slow) + pred_1_cond = self.ema(x, y, ts, mouse, btn, steps_slow) + pred_1_cfg = pred_1_uncond + self.cfg_scale * (pred_1_cond - pred_1_uncond) + pred_1 = torch.where(cfg_mask, pred_1_cfg, pred_1_cond) + + x_new = x - pred_1 * expand(dt_slow) + ts_new = ts - dt_slow + + pred_2_uncond = self.ema(x_new, y, ts_new, null_mouse, null_btn, steps_slow) + pred_2_cond = self.ema(x_new, y, ts_new, mouse, btn, steps_slow) + pred_2_cfg = pred_2_uncond + self.cfg_scale * (pred_2_cond - pred_2_uncond) + pred_2 = torch.where(cfg_mask, pred_2_cfg, pred_2_cond) + + pred = 0.5 * (pred_1 + pred_2) + return pred, steps_fast, ts + + def get_sc_loss(self, x, y, mouse, btn): + target, steps, ts = self.get_sc_targets(x, y, mouse, btn) + pred = self.core(x, y, ts, mouse, btn, steps) + sc_loss = F.mse_loss(pred, target) + return sc_loss + + def forward(self, x, y, mouse, btn): + # x is [b,n,c,h,w] + # y (seed frame) is [b,1,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + _,n,c,h,w = x.shape + + # Split batches between consistency/rf + b = int(len(x) * (1 - self.sc_frac)) + x,x_sc = x[:b], x[b:] + y,y_sc = y[:b], y[b:] + mouse,mouse_sc = mouse[:b], mouse[b:] + btn,btn_sc = btn[:b], btn[b:] + + # Apply classifier-free guidance dropout + if self.cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + # Where mask is True, replace with zeros + mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) + btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) + + with torch.no_grad(): + d = torch.ones_like(x[:,:,0,0,0])*self.sc_max_steps + ts = sample_discrete_timesteps(d) + ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() + + ts_exp = eo.repeat(ts, 'b n -> b n 1 1 1') + z = torch.randn_like(x) + + lerpd = x * (1. - ts_exp) + z * ts_exp + target = z - x + + pred = self.core(lerpd, y, ts, mouse, btn, d) + diff_loss = F.mse_loss(pred, target) + sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc) + + return diff_loss, sc_loss + +def test_inference_cache(): + from ..configs import TransformerConfig + from ..nn.kv_cache import KVCache + + cfg = TransformerConfig( + None, # model_id + 6, # n_layers + 6, # n_heads + 384, # d_model + 1, # patch_size + 128, # channels + 16, # sample_size + 0.1, # cfg_prob + 11, # n_buttons + 16, # tokens_per_frame + 10, # n_frames + True # causal + ) + + model = ShortcutGameRFTCore(cfg).bfloat16().cuda() + + NUM_FRAMES = 10 + x = torch.randn(1, NUM_FRAMES, 128, 4, 4).bfloat16().cuda() + y = torch.randn(1, 1, 128, 4, 4).bfloat16().cuda() + mouse = torch.randn(1, NUM_FRAMES, 2).bfloat16().cuda() + btn = torch.randn(1, NUM_FRAMES, 11).bfloat16().cuda() + t = torch.full((1, NUM_FRAMES), 0.25, device='cuda', dtype=torch.bfloat16) + d = torch.full((1, NUM_FRAMES), 4, device='cuda', dtype=torch.bfloat16) + + cache = KVCache(cfg).to(device='cuda', dtype=torch.bfloat16) + cache.reset(1) + + with torch.no_grad(): + # First pass - generate cache for all frames + cache.enable_cache_updates() + out = model(x, y, t, mouse, btn, d, cache) + print(f"Initial cache length: {len(cache)}") + print(f"Initial cache shape: {cache.cache[0][0].shape}") + + # Generate single new frame with t=1, d=1 + new_x = torch.randn(1, 1, 128, 4, 4).bfloat16().cuda() + new_mouse = torch.randn(1, 1, 2).bfloat16().cuda() + new_btn = torch.randn(1, 1, 11).bfloat16().cuda() + new_t = torch.ones(1, 1, device='cuda', dtype=torch.bfloat16) + new_d = torch.ones(1, 1, device='cuda', dtype=torch.bfloat16) + + # Disable cache updates for inference + cache.disable_cache_updates() + new_out = model(new_x, y, new_t, new_mouse, new_btn, new_d, cache) + print(f"After inference cache length: {len(cache)}") + print(f"After inference cache shape: {cache.cache[0][0].shape}") + + # Re-enable cache updates and update cache with t=0.25, d=4 + cache.enable_cache_updates() + new_t = torch.full((1, 1), 0.25, device='cuda', dtype=torch.bfloat16) + new_d = torch.full((1, 1), 4, device='cuda', dtype=torch.bfloat16) + new_out = model(new_x, y, new_t, new_mouse, new_btn, new_d, cache) + print(f"Final cache length: {len(cache)}") + print(f"Final cache shape: {cache.cache[0][0].shape}") + +def test_wrapper(): + from ..configs import TransformerConfig + from ema_pytorch import EMA + from copy import deepcopy + + cfg = TransformerConfig( + None, # model_id + 6, # n_layers + 6, # n_heads + 384, # d_model + 1, # patch_size + 128, # channels + 16, # sample_size + 0.1, # cfg_prob + 11, # n_buttons + 16, # tokens_per_frame + 10, # n_frames + True # causal + ) + + model = ShortcutGameRFT(cfg).bfloat16().cuda() + ema = EMA(model, beta=0.999,update_after_step=0,update_every=1) + model.set_ema(ema) + + NUM_FRAMES = 10 + x = torch.randn(4, NUM_FRAMES, 128, 4, 4).bfloat16().cuda() + y = torch.randn(4, 1, 128, 4, 4).bfloat16().cuda() + mouse = torch.randn(4, NUM_FRAMES, 2).bfloat16().cuda() + btn = torch.randn(4, NUM_FRAMES, 11).bfloat16().cuda() + + with torch.no_grad(): + loss_1, loss_2 = model(x, y, mouse, btn) + print(loss_1, loss_2) + +if __name__ == "__main__": + test_wrapper() \ No newline at end of file diff --git a/owl_wms/nn/kv_cache.py b/owl_wms/nn/kv_cache.py index 99d33eb9..82399d8f 100644 --- a/owl_wms/nn/kv_cache.py +++ b/owl_wms/nn/kv_cache.py @@ -46,7 +46,27 @@ def push(self, new_k, new_v, layer_ind): @torch.no_grad() def update(self, new_k, new_v, layer_ind): assert self.cache is not None, "Must reset cache before using" - self.cache[layer_ind] = (new_k,new_v) + + def tuple_truncate(k, v): + k = k[:,:,-self.max_length:] + v = v[:,:,-self.max_length:] + return k, v + + self.cache[layer_ind] = tuple_truncate(new_k,new_v) + + @torch.no_grad() + def truncate(self, truncate_amt): + """ + Truncate frames from the KV cache + """ + truncate_amt = truncate_amt * self.config.tokens_per_frame + def tuple_truncate(k, v): + k = k[:,:,truncate_amt:] + v = v[:,:,truncate_amt:] + return k, v + + for i in range(self.config.n_layers): + self.cache[i] = tuple_truncate(self.cache[i]) def __len__(self): assert self.cache is not None, "Must reset cache before using" diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py new file mode 100644 index 00000000..103c130a --- /dev/null +++ b/owl_wms/sampling/shortcut_sampler.py @@ -0,0 +1,326 @@ +import math +import pathlib +import torch +from torch import Module +import torch.nn.functional + +from tqdm import tqdm +from typing import Optional + +from ..utils import batch_permute_to_length +from ..nn.kv_cache import KVCache +from ..models.gamerft_shortcut import ShortcutGameRFT + +def zlerp(x, alpha): + z = torch.randn_like(x) + return x * (1. - alpha) + z * alpha + + +class InferenceCachedShortcutSampler: + + ALPHA = 0.25 + + def __init__(self, + model: ShortcutGameRFT, + window_length = 60, # TODO What's the difference between this and num_frames? + num_frames = 1, # TODO I am assuming num_frames is how many we generate at once, window_length is history temporal dim. + only_return_generated = False, # When true, don't return the history alongside the frame latents. + vae_scale=2.17, # TODO Double-check + decode_fn: Optional[Module] = None, + initial_history_path: Optional[pathlib.Path] = None): + # -- + self.model: ShortcutGameRFT = model + self.window_length = window_length + self.num_frames = num_frames + + self.vae_scale = vae_scale + self.only_return_generated = only_return_generated + + # -- + self._cache_built = False + self.cache = KVCache(model.config) + self.decode_fn = decode_fn + self.initial_history_path = initial_history_path + self.initial_history_bWchw = self.load_cached_history(self.initial_history_path) + + N = self.window_length + C = self.model.config.channels + H = W = int(math.sqrt(self.model.config.tokens_per_frame)) + + if self.initial_history_bWchw.ndim == 4: + self.initial_history_bWchw = self.initial_history_bWchw.unsqueeze(1) + + # NOTE Should batch-size be checked to be 1? + assert tuple(self.initial_history_bWchw.shape) == (1, N, C, H, W), \ + f'Initial history must have shape (B=1, {N=}, {C=}, {H=}, {W=}), ' \ + f'but got {tuple(self.initial_history_bWchw.shape)}' + + def init_cache(self, + frames_bWchw, # [B, W, c, h, w] - NOTE history of frames + keyframe_b1chw, # [B, 1, c, h, w] - NOTE + mouse_bW2, # [B, W, 2] + button_bW11, # [B, W, 11] + ts_bW, # [B, W] + d_bW): # [B, W] + if self._cache_built: + print(f'WARNING: Cache already built but called `init_cache` again - ignoring.') + return + + B, N, *_ = frames_bWchw.shape + + self.cache.reset(B) ; self.cache.enable_cache_updates() + + # -- noise the history and fwd to kv cache + self.model.core.sample(x=zlerp(frames_bWchw, self.ALPHA), + y=keyframe_b1chw, + mouse=mouse_bW2, + btn=button_bW11, + cache=self.cache, + ts=ts_bW, d=d_bW) + + self.cache.disable_cache_updates() ; self._cache_built = True + print(f'Cache initialized for {B} x {N} frames - {[[i.shape for i in elt] + for elt in self.cache.cache]}') + return self.cache + + + def load_cached_history(self, initial_history_path: pathlib.Path | None): + if initial_history_path is not None: + return torch.load(initial_history_path) + + # NOTE Generate history, then save it & return it. + history = ... + # NOTE This is so we avoid generating the history with a compiled model. + torch.save(history, initial_history_path) + return history + + + def __call__(self, + ctxt_frame_b1chw, # [B, 1, c, h, w] - NOTE Keyframe + mouse_b1_2, # [B, 1, 2] + button_b1_11, # [B, 1, 11] + ts_alpha_b1, # [B, 1] - overall denoising timestamp (e.g. 128) + d_alpha_b1, # [B, 1] - denoising step budget (e.g. 4) + ) -> torch.Tensor: # [B, 1, c, h, w] + # 1. ---- generate next frame ---- + self.cache.disable_cache_updates() + # 1.A) -- use the full context, including entire action history, to generate the next frame given + # -- cache. + frame = self.model.core.sample(None, ctxt_frame_b1chw, + mouse_b1_2, button_b1_11, + self.cache, ts=None, d=None) # NOTE simulating one-step sampling + # 1.B) -- add frame latent to display + # 2. ---- repopulate cache ---- + self.cache.enable_cache_updates() ; self.cache.truncate(1) + self.model.core.sample( x=zlerp(frame, self.ALPHA), # diffuse with noised frame to repopulate cache + y=ctxt_frame_b1chw, + mouse=mouse_b1_2, + btn=button_b1_11, + cache=self.cache, + ts=ts_alpha_b1, d=d_alpha_b1) + self.cache.disable_cache_updates() + return frame + + # NOTE NOTE NOTE NOTE NOTE NOTE NOTE + # NOTE W = window_size N = num_frames + # NOTE T = window_size + num_frames + # NOTE NOTE NOTE NOTE NOTE NOTE NOTE + @torch.no_grad() + def generate_frames(self, + history_bWchw, # [B, window_size, c, h, w] - TODO: MP4 from CoD initially, and after that it's just KV cache. + keyframe_b1chw, # [B, 1, c, h, w] - NOTE: Keyframe is just a png from CoD that is encoded. + mouse_bT2, # [B, T, 2] - Actions taken by the user. + button_bT11, # [B, T, 11] - Actions taken by the user. + ) -> torch.Tensor: # [B, T, c, h, w] - either latent or rgb. + + if not self._cache_built: + print(f'WARNING: Cache not built, but called `generate_frames` - initializing cache.') + self.init_cache(history_bWchw, keyframe_b1chw, mouse_bT2, button_bT11) + + # If does not have batch-size, add it. This sampler is going to be used for single-user inference so batch-size is always 1. + # The caller might not specify the batch-size, so we have this here. + if history_bWchw.ndim == 4: + history_bWchw = history_bWchw.unsqueeze(1) + + history_bWchw = history_bWchw[:, -self.window_length:, ::] + + assert history_bWchw.shape[1] == self.window_length, \ + f'Window history must be at least {self.window_length} frames long, but got {history_bWchw.shape}' + + ts_alpha_bW = torch.ones_like(history_bWchw[:,:,0,0,0]) * self.ALPHA + d_alpha_bW = torch.ones_like(history_bWchw[:,:,0,0,0]) * round(1./self.ALPHA) + + ts_alpha_b1 = ts_alpha_bW[:,0].unsqueeze(1) + d_alpha_b1 = d_alpha_bW [:,0].unsqueeze(1) + + frames_latent = [] + for frame_idx in range(self.num_frames): + btn_atom = button_bT11[:, self.window_length+frame_idx].unsqueeze(1) + mouse_atom = mouse_bT2 [:, self.window_length+frame_idx].unsqueeze(1) + frame = self.__call__(ctxt_frame_b1chw=keyframe_b1chw, + mouse_b1_2=mouse_atom, button_b1_11=btn_atom, + ts_alpha_b1=ts_alpha_b1, d_alpha_b1=d_alpha_b1) + frames_latent += [frame] + + frames_latent = torch.cat(frames_latent, dim=1) + + if self.only_return_generated: frames_latent = frames_latent[:,-self.num_frames:] + + if self.decode_fn is not None: + frames_rgb = self.decode_fn(frames_latent * self.vae_scale) + return frames_rgb, mouse_bT2, button_bT11 + + return frames_latent, mouse_bT2, button_bT11 + + +class CacheShortcutSampler: + """ + Shortcut CFG sampler builds cache with 4 step diffusion. + Samples new frames in 1 step. + + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale = 1): + # dummy_batch is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + history = history[:,:self.window_length] + new_frames = [] + alpha = 0.25 # This number is special for our sampler + + # Extended fake controls to use during sampling + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, num_frames + self.window_length) + + # Generate cache over history + noisy_history = zlerp(history.clone(), alpha) + ts = torch.ones_like(noisy_history[:,:,0,0,0]) * alpha + d = torch.ones_like(noisy_history[:,:,0,0,0]) * round(1./alpha) + ts_single = ts[:,0].unsqueeze(1) + d_single = d[:,0].unsqueeze(1) + + cache = KVCache(model.config) + cache.reset(history.shape[0]) + + cache.enable_cache_updates() + _ = model.sample(noisy_history, keyframe, mouse, btn, cache, ts, d) + cache.disable_cache_updates() + + # Cache is now built! + + for frame_idx in tqdm(range(num_frames)): + cache.truncate(1) # Drop first frame + + # Generate new frame + cache.disable_cache_updates() + mouse = extended_mouse[:,self.window_length+frame_idx].unsqueeze(1) + btn = extended_btn[:,self.window_length+frame_idx].unsqueeze(1) + # N+1 + new_frame = model.sample(None, keyframe, mouse, btn, cache) # [b,1,c,h,w] + new_frames.append(new_frame) + + # Add that frame to the cache + cache.enable_cache_updates() + new_frame_noisy = zlerp(new_frame, alpha) + # N+2, noisy(N+1) gets cached + _ = model.sample(new_frame_noisy, keyframe, mouse, btn, cache, ts_single, d_single) + + new_frames = torch.cat(new_frames, dim = 1) + x = torch.cat([history,new_frames], dim = 1) + + if self.only_return_generated: + x = x[:,-num_frames:] + extended_mouse = extended_mouse[:,-num_frames:] + extended_btn = extended_btn[:,-num_frames:] + + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + + return x, extended_mouse, extended_btn + +class WindowShortcutSampler: + """ + Same as above but with no cache + + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale = 1): + # history is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + history = history[:,:self.window_length] + new_frames = [] + alpha = 0.25 # This number is special for our sampler + + # Extended fake controls to use during sampling + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, self.num_frames + self.window_length) + + # Initialize window history + window_history = history.clone() + + for frame_idx in tqdm(range(self.num_frames)): + # Setup window history + x = window_history[:,-self.window_length:].clone() + + # Noise all but last frame to alpha + x[:,:-1] = zlerp(x[:,:-1], alpha) + # Last frame starts as random noise + x[:,-1] = torch.randn_like(x[:,-1]) + + # Setup timesteps - alpha for context, 1.0 for generated + ts = torch.ones_like(x[:,:,0,0,0]) + ts[:,:-1] = alpha + + # Setup diffusion steps - 4 for context, 1 for generated + d = torch.ones_like(x[:,:,0,0,0]) + d[:,:-1] = 4 + + # Get current controls + curr_mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] + curr_btn = extended_btn[:,frame_idx:frame_idx+self.window_length] + + # Generate new frame + pred = model.sample(x, keyframe, curr_mouse, curr_btn, None, ts, d) + new_frame = pred[:,-1:] # Take only the last frame + new_frames.append(new_frame) + + # Add new frame to window history + window_history = torch.cat([window_history, new_frame], dim=1) + + new_frames = torch.cat(new_frames, dim=1) + x = torch.cat([history, new_frames], dim=1) + + if self.only_return_generated: + x = x[:,-self.num_frames:] + extended_mouse = extended_mouse[:,-self.num_frames:] + extended_btn = extended_btn[:,-self.num_frames:] + + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + + return x, extended_mouse, extended_btn \ No newline at end of file diff --git a/owl_wms/utils/ddp.py b/owl_wms/utils/ddp.py index 8bc23e26..7296a3a2 100644 --- a/owl_wms/utils/ddp.py +++ b/owl_wms/utils/ddp.py @@ -23,4 +23,4 @@ def setup(force=False): def cleanup(): if dist.is_available() and dist.is_initialized(): - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() From f0fcf42f077984a63b1ec09d52134221cdca0993 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 00:58:19 +0000 Subject: [PATCH 38/71] debugs --- owl_wms/models/gamerft_shortcut.py | 6 +++--- owl_wms/nn/kv_cache.py | 4 ++-- owl_wms/nn/mmattn.py | 3 +-- owl_wms/nn/rope.py | 2 +- owl_wms/sampling/shortcut_sampler.py | 6 ++++-- owl_wms/trainers/shortcut_trainer.py | 9 +++++++-- 6 files changed, 18 insertions(+), 12 deletions(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index 38f11b7b..c35c5e3b 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -52,7 +52,7 @@ def sample(self, x, y, mouse, btn, kv_cache = None, t = None, d = None): if d is None: d = torch.ones_like(x[:,:,0,0,0]) - return x - self.forward(x, t, mouse, btn, d, kv_cache) + return x - self.forward(x, y, t, mouse, btn, d, kv_cache) def forward(self, x, y, t, mouse, btn, d, kv_cache = None): # x is [b,n,c,h,w] @@ -123,8 +123,8 @@ def set_ema(self, ema): else: self.ema = ema.ema_model.core + #@torch.compile() @torch.no_grad() - @torch.compile() def get_sc_targets(self, x, y, mouse, btn): steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) steps_fast = steps_slow / 2 @@ -303,4 +303,4 @@ def test_wrapper(): print(loss_1, loss_2) if __name__ == "__main__": - test_wrapper() \ No newline at end of file + test_wrapper() diff --git a/owl_wms/nn/kv_cache.py b/owl_wms/nn/kv_cache.py index fd2de00d..fea43f0c 100644 --- a/owl_wms/nn/kv_cache.py +++ b/owl_wms/nn/kv_cache.py @@ -73,11 +73,11 @@ def tuple_truncate(k, v): return k, v for i in range(self.config.n_layers): - self.cache[layer_ind] = tuple_truncate(self.cache[layer_ind]) + self.cache[i] = tuple_truncate(*self.cache[i]) def __len__(self): assert self.cache is not None, "Must reset cache before using" return self.cache[0][0].shape[2] def shape(self): - return self.shape \ No newline at end of file + return self.shape diff --git a/owl_wms/nn/mmattn.py b/owl_wms/nn/mmattn.py index 66b0ecc2..00855779 100644 --- a/owl_wms/nn/mmattn.py +++ b/owl_wms/nn/mmattn.py @@ -116,7 +116,6 @@ def forward(self, x_1, x_2, kv_cache=None): k = torch.cat([k1,k2],dim=-2) v = torch.cat([v1,v2],dim=-2) - print("====") x = F.scaled_dot_product_attention(q,k,v, attn_mask = mask) x = self.merge(x) @@ -291,4 +290,4 @@ def test_mask(): if __name__ == "__main__": - test_mask() \ No newline at end of file + test_mask() diff --git a/owl_wms/nn/rope.py b/owl_wms/nn/rope.py index a2b9ec76..2a5402da 100644 --- a/owl_wms/nn/rope.py +++ b/owl_wms/nn/rope.py @@ -77,7 +77,7 @@ def forward(self, q, k): if q.shape[2] < n * m: truncate = q.shape[2]//m # How many frames is q? - q = eo.rearrange(q, 'b h (n m) d -> b h n m d', n=n,m=m) + q = eo.rearrange(q, 'b h (n m) d -> b h n m d', n=q.shape[2]//m,m=m) k = eo.rearrange(k, 'b h (n m) d -> b h n m d', n=n,m=m) freqs = self.pos_emb.get_axial_freqs(n,m) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index b41a0bc4..67be7315 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -34,6 +34,7 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale # dummy_batch is [b,n,c,h,w] # mouse is [b,n,2] # btn is [b,n,n_button] + num_frames = self.num_frames # output will be [b,n+self.num_frames,c,h,w] history = history[:,:self.window_length] @@ -52,8 +53,9 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale cache = KVCache(model.config) cache.reset(history.shape[0]) - + cache.enable_cache_updates() + _ = model.sample(noisy_history, keyframe, mouse, btn, cache, ts, d) cache.disable_cache_updates() @@ -161,4 +163,4 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale x = x * scale x = decode_fn(x) - return x, extended_mouse, extended_btn \ No newline at end of file + return x, extended_mouse, extended_btn diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index 37d7401b..4c0b0c56 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -102,6 +102,10 @@ def train(self): update_after_step = 0, update_every = 1 ) + if self.world_size > 1: + self.model.module.set_ema(self.ema) + else: + self.model.set_ema(self.ema) #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) def get_ema_core(): @@ -144,8 +148,9 @@ def get_ema_core(): batch_btn = batch_btn.cuda().bfloat16() with ctx: - diff_loss, sc_loss = self.model(batch_vid,batch_mouse,batch_btn) / accum_steps + diff_loss, sc_loss = self.model(batch_vid,batch_keyframe,batch_mouse,batch_btn) loss = diff_loss + sc_loss + loss = loss / accum_steps self.scaler.scale(loss).backward() #find_unused_params(self.model) @@ -199,4 +204,4 @@ def get_ema_core(): if self.rank == 0: self.save() - self.barrier() \ No newline at end of file + self.barrier() From 60496b5aa090bc41a9445e5fdecab337b6664663 Mon Sep 17 00:00:00 2001 From: Sami Date: Tue, 10 Jun 2025 21:19:38 -0400 Subject: [PATCH 39/71] inference shortcut diffusion sampler but with kv cache --- owl_wms/sampling/shortcut_sampler.py | 133 +++++++++++++++++---------- owl_wms/utils/owl_vae_bridge.py | 36 +++++++- 2 files changed, 120 insertions(+), 49 deletions(-) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index 103c130a..edcf4b86 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -1,20 +1,52 @@ +import cv2 import math import pathlib import torch from torch import Module -import torch.nn.functional - from tqdm import tqdm from typing import Optional -from ..utils import batch_permute_to_length from ..nn.kv_cache import KVCache +from ..utils import batch_permute_to_length from ..models.gamerft_shortcut import ShortcutGameRFT + def zlerp(x, alpha): z = torch.randn_like(x) return x * (1. - alpha) + z * alpha +def load_mp4_as_tensor(mp4_path: pathlib.Path) -> torch.Tensor: + """Load MP4 as tensor in format [N, C=3, H, W] with values in [-1, 1]""" + video = cv2.VideoCapture(str(mp4_path)) + + if not video.isOpened(): + raise ValueError(f"Could not open video file: {mp4_path}") + + frames = [] + while True: + ret, frame = video.read() + if not ret: + break + + # Convert BGR to RGB + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # Convert to torch tensor and normalize to [-1, 1] + frame = torch.from_numpy(frame).float() / 127.5 - 1.0 + + # Rearrange from [H, W, C] to [C, H, W] + frame = frame.permute(2, 0, 1) + + frames.append(frame) + + video.release() + + if not frames: + raise ValueError(f"No frames found in video: {mp4_path}") + + # Stack to [N, C, H, W] + return torch.stack(frames) + class InferenceCachedShortcutSampler: @@ -22,12 +54,14 @@ class InferenceCachedShortcutSampler: def __init__(self, model: ShortcutGameRFT, - window_length = 60, # TODO What's the difference between this and num_frames? - num_frames = 1, # TODO I am assuming num_frames is how many we generate at once, window_length is history temporal dim. - only_return_generated = False, # When true, don't return the history alongside the frame latents. - vae_scale=2.17, # TODO Double-check + window_length = 60, + num_frames = 1, + only_return_generated = False, + vae_scale = 2.17, decode_fn: Optional[Module] = None, - initial_history_path: Optional[pathlib.Path] = None): + initial_history_pt_path: Optional[pathlib.Path] = None, + initial_history_mp4_path: Optional[pathlib.Path] = None, + encoder: Optional[Module] = None): # -- self.model: ShortcutGameRFT = model self.window_length = window_length @@ -40,24 +74,46 @@ def __init__(self, self._cache_built = False self.cache = KVCache(model.config) self.decode_fn = decode_fn - self.initial_history_path = initial_history_path - self.initial_history_bWchw = self.load_cached_history(self.initial_history_path) + self.initial_history_pt_path = initial_history_pt_path + self.initial_history_mp4_path = initial_history_mp4_path + self.encoder = encoder + + assert initial_history_pt_path is not None or initial_history_mp4_path is not None, \ + 'Either initial_history_pt_path or initial_history_mp4_path must be provided' + + if initial_history_mp4_path is not None: + assert encoder is not None, \ + 'Encoder must be provided if initial_history_mp4_path is provided' + + self.initial_history_bWchw = self.init_history(self.initial_history_pt_path, self.initial_history_mp4_path) + self.keyframe_b1chw = self.initial_history_bWchw[:,0] + + def init_history(self, + initial_history_pt_path: pathlib.Path | None, + initial_history_mp4_path: pathlib.Path | None) -> torch.Tensor: + + if initial_history_pt_path is not None: + history_wchw = torch.load(initial_history_pt_path) + else: + history_wrgb = load_mp4_as_tensor(initial_history_mp4_path).unsqueeze(0) # add batch dim + history_wchw = self.encoder(history_wrgb) + # NOTE This is so we avoid generating the history with a compiled model. + torch.save(history_wchw, initial_history_mp4_path.absolute().replace('.mp4', '.pt')) N = self.window_length C = self.model.config.channels H = W = int(math.sqrt(self.model.config.tokens_per_frame)) - if self.initial_history_bWchw.ndim == 4: - self.initial_history_bWchw = self.initial_history_bWchw.unsqueeze(1) - - # NOTE Should batch-size be checked to be 1? - assert tuple(self.initial_history_bWchw.shape) == (1, N, C, H, W), \ + assert tuple(history_wchw.shape) == (1, N, C, H, W), \ f'Initial history must have shape (B=1, {N=}, {C=}, {H=}, {W=}), ' \ - f'but got {tuple(self.initial_history_bWchw.shape)}' + f'but got {tuple(history_wchw.shape)}' + + return history_wchw + def init_cache(self, frames_bWchw, # [B, W, c, h, w] - NOTE history of frames - keyframe_b1chw, # [B, 1, c, h, w] - NOTE + keyframe_b1chw, # [B, 1, c, h, w] - NOTE keyframe conditioning mouse_bW2, # [B, W, 2] button_bW11, # [B, W, 11] ts_bW, # [B, W] @@ -83,33 +139,19 @@ def init_cache(self, for elt in self.cache.cache]}') return self.cache - - def load_cached_history(self, initial_history_path: pathlib.Path | None): - if initial_history_path is not None: - return torch.load(initial_history_path) - - # NOTE Generate history, then save it & return it. - history = ... - # NOTE This is so we avoid generating the history with a compiled model. - torch.save(history, initial_history_path) - return history - - def __call__(self, - ctxt_frame_b1chw, # [B, 1, c, h, w] - NOTE Keyframe - mouse_b1_2, # [B, 1, 2] - button_b1_11, # [B, 1, 11] - ts_alpha_b1, # [B, 1] - overall denoising timestamp (e.g. 128) - d_alpha_b1, # [B, 1] - denoising step budget (e.g. 4) + ctxt_frame_b1chw, # [B, 1, c, h, w] - NOTE Keyframe conditioning + mouse_b1_2, # [B, 1, 2] - NOTE mouse actions + button_b1_11, # [B, 1, 11] - NOTE button actions + ts_alpha_b1, # [B, 1] - NOTE overall denoising timestamp (e.g. 128) + d_alpha_b1, # [B, 1] - NOTE denoising step budget (e.g. 4) ) -> torch.Tensor: # [B, 1, c, h, w] # 1. ---- generate next frame ---- self.cache.disable_cache_updates() - # 1.A) -- use the full context, including entire action history, to generate the next frame given - # -- cache. + # 1.A) -- use the full context, including entire action history, to generate the next frame given cache. frame = self.model.core.sample(None, ctxt_frame_b1chw, mouse_b1_2, button_b1_11, self.cache, ts=None, d=None) # NOTE simulating one-step sampling - # 1.B) -- add frame latent to display # 2. ---- repopulate cache ---- self.cache.enable_cache_updates() ; self.cache.truncate(1) self.model.core.sample( x=zlerp(frame, self.ALPHA), # diffuse with noised frame to repopulate cache @@ -121,21 +163,16 @@ def __call__(self, self.cache.disable_cache_updates() return frame - # NOTE NOTE NOTE NOTE NOTE NOTE NOTE - # NOTE W = window_size N = num_frames - # NOTE T = window_size + num_frames - # NOTE NOTE NOTE NOTE NOTE NOTE NOTE @torch.no_grad() def generate_frames(self, - history_bWchw, # [B, window_size, c, h, w] - TODO: MP4 from CoD initially, and after that it's just KV cache. - keyframe_b1chw, # [B, 1, c, h, w] - NOTE: Keyframe is just a png from CoD that is encoded. - mouse_bT2, # [B, T, 2] - Actions taken by the user. - button_bT11, # [B, T, 11] - Actions taken by the user. - ) -> torch.Tensor: # [B, T, c, h, w] - either latent or rgb. + history_bWchw, # [B, W, c, h, w] - NOTE: MP4 from CoD initially, and after that it's just KV cache. + mouse_bT2, # [B, W+N, 2] - Actions taken by the user. + button_bT11, # [B, W+N, 11] - Actions taken by the user. + ) -> torch.Tensor: # [B, W+N, c, h, w] - either latent or rgb. if not self._cache_built: print(f'WARNING: Cache not built, but called `generate_frames` - initializing cache.') - self.init_cache(history_bWchw, keyframe_b1chw, mouse_bT2, button_bT11) + self.init_cache(history_bWchw, self.keyframe_b1chw, mouse_bT2, button_bT11) # If does not have batch-size, add it. This sampler is going to be used for single-user inference so batch-size is always 1. # The caller might not specify the batch-size, so we have this here. @@ -157,7 +194,7 @@ def generate_frames(self, for frame_idx in range(self.num_frames): btn_atom = button_bT11[:, self.window_length+frame_idx].unsqueeze(1) mouse_atom = mouse_bT2 [:, self.window_length+frame_idx].unsqueeze(1) - frame = self.__call__(ctxt_frame_b1chw=keyframe_b1chw, + frame = self.__call__(ctxt_frame_b1chw=self.keyframe_b1chw, mouse_b1_2=mouse_atom, button_b1_11=btn_atom, ts_alpha_b1=ts_alpha_b1, d_alpha_b1=d_alpha_b1) frames_latent += [frame] diff --git a/owl_wms/utils/owl_vae_bridge.py b/owl_wms/utils/owl_vae_bridge.py index 331a8582..50b4f524 100644 --- a/owl_wms/utils/owl_vae_bridge.py +++ b/owl_wms/utils/owl_vae_bridge.py @@ -34,6 +34,21 @@ def get_decoder_only(vae_id, cfg_path, ckpt_path): model = model.bfloat16().cuda().eval() return model +def get_encoder_only(vae_id, cfg_path, ckpt_path): + if vae_id == "dcae": + model_id = "mit-han-lab/dc-ae-f64c128-mix-1.0-diffusers" + model = AutoencoderDC.from_pretrained(model_id).bfloat16().cuda().eval() + del model.decoder # Keep encoder only + return model.encoder + elif vae_id == "720pr3dc": + cfg = Config.from_yaml(cfg_path).model + model = get_model_cls(cfg.model_id)(cfg) + model.load_state_dict(torch.load(ckpt_path, map_location='cpu',weights_only=False)) + del model.decoder # Keep encoder only + model = model.encoder + model = model.bfloat16().cuda().eval() + return model + @torch.no_grad() def _make_batched_decode_fn(decoder, batch_size = 8): def decode(x): @@ -70,4 +85,23 @@ def decode(x): x = x.view(b,n,c,h,w).contiguous() return x - return decode \ No newline at end of file + return decode + +@torch.no_grad() +def make_batched_encode_fn(encoder, batch_size=8): + def encode(x): + # x is [b,n,c,h,w] RGB frames + b,n,c,h,w = x.shape + x = x.view(b*n,c,h,w).contiguous() + + batches = x.split(batch_size) + batch_out = [] + for batch in batches: + batch_out.append(encoder(batch).bfloat16()) + + x = torch.cat(batch_out) # [b*n,latent_c,latent_h,latent_w] + _,latent_c,latent_h,latent_w = x.shape + x = x.view(b,n,latent_c,latent_h,latent_w).contiguous() + + return x + return encode \ No newline at end of file From 7f51b544ab7b118ed87f84202e70ff4772b97761 Mon Sep 17 00:00:00 2001 From: Sami Date: Tue, 10 Jun 2025 21:24:07 -0400 Subject: [PATCH 40/71] test --- owl_wms/sampling/shortcut_sampler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index edcf4b86..4347adb4 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -15,6 +15,7 @@ def zlerp(x, alpha): z = torch.randn_like(x) return x * (1. - alpha) + z * alpha +# TODO Test def load_mp4_as_tensor(mp4_path: pathlib.Path) -> torch.Tensor: """Load MP4 as tensor in format [N, C=3, H, W] with values in [-1, 1]""" video = cv2.VideoCapture(str(mp4_path)) From 3c9994a7050e84f88fcca1ca6236d6489c849a7d Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 01:24:08 +0000 Subject: [PATCH 41/71] debug --- owl_wms/models/gamerft_shortcut.py | 11 +++++++---- owl_wms/trainers/shortcut_trainer.py | 12 ++++++------ 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index c35c5e3b..d0578e83 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -16,6 +16,7 @@ ) from ..nn.attn import UViT, FinalLayer from ..nn.mmattn import MMUViT +from ..utils import freeze class ShortcutGameRFTCore(nn.Module): def __init__(self, config): @@ -110,7 +111,6 @@ def __init__(self, config): self.core = ShortcutGameRFTCore(config) self.cfg_prob = config.cfg_prob - self.ema = None self.sc_frac = 0.25 self.sc_max_steps = 128 self.cfg_scale = 1.3 @@ -119,10 +119,13 @@ def __init__(self, config): def set_ema(self, ema): if hasattr(ema.ema_model, 'module'): - self.ema = ema.ema_model.module.core + ema = ema.ema_model.module.core else: - self.ema = ema.ema_model.core - + ema = ema.ema_model.core + + self.ema = ema + freeze(self.ema) + #@torch.compile() @torch.no_grad() def get_sc_targets(self, x, y, mouse, btn): diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index 4c0b0c56..2c6f565e 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -102,10 +102,6 @@ def train(self): update_after_step = 0, update_every = 1 ) - if self.world_size > 1: - self.model.module.set_ema(self.ema) - else: - self.model.set_ema(self.ema) #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) def get_ema_core(): @@ -120,6 +116,11 @@ def get_ema_core(): if self.train_cfg.scheduler is not None: self.scheduler = get_scheduler_cls(self.train_cfg.scheduler)(self.opt, **self.train_cfg.scheduler_kwargs) + if self.world_size > 1: + self.model.module.set_ema(self.ema) + else: + self.model.set_ema(self.ema) + # Grad accum setup and scaler accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size accum_steps = max(1, accum_steps) @@ -153,8 +154,7 @@ def get_ema_core(): loss = loss / accum_steps self.scaler.scale(loss).backward() - #find_unused_params(self.model) - + metrics.log('diffusion_loss', diff_loss) metrics.log('shortcut_loss', sc_loss) From cde4a595b639c6ff8e408e2e63dfe15a3fd59ef9 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 01:31:33 +0000 Subject: [PATCH 42/71] add model --- owl_wms/models/gamerft_shortcut.py | 93 +++++++++++++--------------- owl_wms/trainers/shortcut_trainer.py | 11 +--- 2 files changed, 44 insertions(+), 60 deletions(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index d0578e83..33ebb164 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -104,6 +104,44 @@ def sample_steps(b, n, device, dtype, min_val = 0): steps = valid[inds].to(device=device,dtype=dtype) return steps +#@torch.compile() +@torch.no_grad() +def get_sc_targets(ema, x, y, mouse, btn, cfg_scale): + steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) + steps_fast = steps_slow / 2 + + dt_slow = 1./steps_slow + dt_fast = 1./steps_fast + + def expand(t): + #b,c,h,w = x.shape + #t = eo.repeat(t,'b -> b c h w',c=c,h=h,w=w) + #return t + return t[:,:,None,None,None] + + ts = sample_discrete_timesteps(steps_fast) + cfg_mask = torch.isclose(steps_slow, torch.ones_like(steps_slow)*128) + cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] + + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + pred_1_uncond = ema(x, y, ts, null_mouse, null_btn, steps_slow) + pred_1_cond = ema(x, y, ts, mouse, btn, steps_slow) + pred_1_cfg = pred_1_uncond + cfg_scale * (pred_1_cond - pred_1_uncond) + pred_1 = torch.where(cfg_mask, pred_1_cfg, pred_1_cond) + + x_new = x - pred_1 * expand(dt_slow) + ts_new = ts - dt_slow + + pred_2_uncond = ema(x_new, y, ts_new, null_mouse, null_btn, steps_slow) + pred_2_cond = ema(x_new, y, ts_new, mouse, btn, steps_slow) + pred_2_cfg = pred_2_uncond + cfg_scale * (pred_2_cond - pred_2_uncond) + pred_2 = torch.where(cfg_mask, pred_2_cfg, pred_2_cond) + + pred = 0.5 * (pred_1 + pred_2) + return pred, steps_fast, ts + class ShortcutGameRFT(nn.Module): def __init__(self, config): super().__init__() @@ -116,61 +154,14 @@ def __init__(self, config): self.cfg_scale = 1.3 self.config = config - - def set_ema(self, ema): - if hasattr(ema.ema_model, 'module'): - ema = ema.ema_model.module.core - else: - ema = ema.ema_model.core - - self.ema = ema - freeze(self.ema) - #@torch.compile() - @torch.no_grad() - def get_sc_targets(self, x, y, mouse, btn): - steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) - steps_fast = steps_slow / 2 - - dt_slow = 1./steps_slow - dt_fast = 1./steps_fast - - def expand(t): - #b,c,h,w = x.shape - #t = eo.repeat(t,'b -> b c h w',c=c,h=h,w=w) - #return t - return t[:,:,None,None,None] - - ts = sample_discrete_timesteps(steps_fast) - cfg_mask = torch.isclose(steps_slow, torch.ones_like(steps_slow)*128) - cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] - - null_mouse = torch.zeros_like(mouse) - null_btn = torch.zeros_like(btn) - - pred_1_uncond = self.ema(x, y, ts, null_mouse, null_btn, steps_slow) - pred_1_cond = self.ema(x, y, ts, mouse, btn, steps_slow) - pred_1_cfg = pred_1_uncond + self.cfg_scale * (pred_1_cond - pred_1_uncond) - pred_1 = torch.where(cfg_mask, pred_1_cfg, pred_1_cond) - - x_new = x - pred_1 * expand(dt_slow) - ts_new = ts - dt_slow - - pred_2_uncond = self.ema(x_new, y, ts_new, null_mouse, null_btn, steps_slow) - pred_2_cond = self.ema(x_new, y, ts_new, mouse, btn, steps_slow) - pred_2_cfg = pred_2_uncond + self.cfg_scale * (pred_2_cond - pred_2_uncond) - pred_2 = torch.where(cfg_mask, pred_2_cfg, pred_2_cond) - - pred = 0.5 * (pred_1 + pred_2) - return pred, steps_fast, ts - - def get_sc_loss(self, x, y, mouse, btn): - target, steps, ts = self.get_sc_targets(x, y, mouse, btn) + def get_sc_loss(self, x, y, mouse, btn, ema): + target, steps, ts = self.get_sc_targets(ema, x, y, mouse, btn) pred = self.core(x, y, ts, mouse, btn, steps) sc_loss = F.mse_loss(pred, target) return sc_loss - def forward(self, x, y, mouse, btn): + def forward(self, x, y, mouse, btn, ema): # x is [b,n,c,h,w] # y (seed frame) is [b,1,c,h,w] # mouse is [b,n,2] @@ -207,7 +198,7 @@ def forward(self, x, y, mouse, btn): pred = self.core(lerpd, y, ts, mouse, btn, d) diff_loss = F.mse_loss(pred, target) - sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc) + sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) return diff_loss, sc_loss diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index 2c6f565e..da648ae0 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -102,6 +102,7 @@ def train(self): update_after_step = 0, update_every = 1 ) + freeze(self.ema) #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) def get_ema_core(): @@ -113,14 +114,6 @@ def get_ema_core(): # No muon pls self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) - if self.train_cfg.scheduler is not None: - self.scheduler = get_scheduler_cls(self.train_cfg.scheduler)(self.opt, **self.train_cfg.scheduler_kwargs) - - if self.world_size > 1: - self.model.module.set_ema(self.ema) - else: - self.model.set_ema(self.ema) - # Grad accum setup and scaler accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size accum_steps = max(1, accum_steps) @@ -149,7 +142,7 @@ def get_ema_core(): batch_btn = batch_btn.cuda().bfloat16() with ctx: - diff_loss, sc_loss = self.model(batch_vid,batch_keyframe,batch_mouse,batch_btn) + diff_loss, sc_loss = self.model(batch_vid,batch_keyframe,batch_mouse,batch_btn, get_ema_core()) loss = diff_loss + sc_loss loss = loss / accum_steps From 2fa9d5bc668d41e5c0c50deab0cfd6ec9ba508a4 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 01:42:14 +0000 Subject: [PATCH 43/71] debug --- owl_wms/models/gamerft_shortcut.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index 33ebb164..a57a4756 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -156,7 +156,8 @@ def __init__(self, config): self.config = config def get_sc_loss(self, x, y, mouse, btn, ema): - target, steps, ts = self.get_sc_targets(ema, x, y, mouse, btn) + with torch.no_grad(): + target, steps, ts = get_sc_targets(ema, x, y, mouse, btn, self.cfg_scale) pred = self.core(x, y, ts, mouse, btn, steps) sc_loss = F.mse_loss(pred, target) return sc_loss From 67b4375b7c3514ecc43af3fd4d4d325be7bc7ac9 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 01:48:46 +0000 Subject: [PATCH 44/71] debug: --- owl_wms/models/gamerft_shortcut.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index a57a4756..bc877a94 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -158,8 +158,8 @@ def __init__(self, config): def get_sc_loss(self, x, y, mouse, btn, ema): with torch.no_grad(): target, steps, ts = get_sc_targets(ema, x, y, mouse, btn, self.cfg_scale) - pred = self.core(x, y, ts, mouse, btn, steps) - sc_loss = F.mse_loss(pred, target) + pred = self.core(x.detach(), y.detach(), ts.detach(), mouse.detach(), btn.detach(), steps.detach()) + sc_loss = F.mse_loss(pred, target.detach()) return sc_loss def forward(self, x, y, mouse, btn, ema): @@ -197,9 +197,9 @@ def forward(self, x, y, mouse, btn, ema): lerpd = x * (1. - ts_exp) + z * ts_exp target = z - x - pred = self.core(lerpd, y, ts, mouse, btn, d) - diff_loss = F.mse_loss(pred, target) - sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) + pred = self.core(lerpd.detach(), y.detach(), ts.detach(), mouse.detach(), btn.detach(), d.detach()) + diff_loss = F.mse_loss(pred, target.detach()) + sc_loss = self.get_sc_loss(x_sc.detach(), y_sc.detach(), mouse_sc.detach(), btn_sc.detach(), ema) return diff_loss, sc_loss From 0dd7e9121ee754e378d92f363bf82b2ef18d7846 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 01:57:46 +0000 Subject: [PATCH 45/71] debug --- owl_wms/models/gamerft_shortcut.py | 41 +++++++++++++++--------------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index bc877a94..c71bb8cc 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -156,10 +156,9 @@ def __init__(self, config): self.config = config def get_sc_loss(self, x, y, mouse, btn, ema): - with torch.no_grad(): - target, steps, ts = get_sc_targets(ema, x, y, mouse, btn, self.cfg_scale) - pred = self.core(x.detach(), y.detach(), ts.detach(), mouse.detach(), btn.detach(), steps.detach()) - sc_loss = F.mse_loss(pred, target.detach()) + target, steps, ts = get_sc_targets(ema, x, y, mouse, btn, self.cfg_scale) + pred = self.core(x, y, ts, mouse, btn, steps) + sc_loss = F.mse_loss(pred, target) return sc_loss def forward(self, x, y, mouse, btn, ema): @@ -167,18 +166,19 @@ def forward(self, x, y, mouse, btn, ema): # y (seed frame) is [b,1,c,h,w] # mouse is [b,n,2] # btn is [b,n,n_buttons] - _,n,c,h,w = x.shape - - # Split batches between consistency/rf - b = int(len(x) * (1 - self.sc_frac)) - x,x_sc = x[:b], x[b:] - y,y_sc = y[:b], y[b:] - mouse,mouse_sc = mouse[:b], mouse[b:] - btn,btn_sc = btn[:b], btn[b:] - - # Apply classifier-free guidance dropout - if self.cfg_prob > 0.0: - mask = torch.rand(b, device=x.device) <= self.cfg_prob + with torch.no_grad(): + _,n,c,h,w = x.shape + + # Split batches between consistency/rf + b = int(len(x) * (1 - self.sc_frac)) + x,x_sc = x[:b], x[b:] + y,y_sc = y[:b], y[b:] + mouse,mouse_sc = mouse[:b], mouse[b:] + btn,btn_sc = btn[:b], btn[b:] + + # Apply classifier-free guidance dropout + if self.cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob null_mouse = torch.zeros_like(mouse) null_btn = torch.zeros_like(btn) @@ -186,7 +186,6 @@ def forward(self, x, y, mouse, btn, ema): mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) - with torch.no_grad(): d = torch.ones_like(x[:,:,0,0,0])*self.sc_max_steps ts = sample_discrete_timesteps(d) ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() @@ -197,10 +196,10 @@ def forward(self, x, y, mouse, btn, ema): lerpd = x * (1. - ts_exp) + z * ts_exp target = z - x - pred = self.core(lerpd.detach(), y.detach(), ts.detach(), mouse.detach(), btn.detach(), d.detach()) - diff_loss = F.mse_loss(pred, target.detach()) - sc_loss = self.get_sc_loss(x_sc.detach(), y_sc.detach(), mouse_sc.detach(), btn_sc.detach(), ema) - + pred = self.core(lerpd, y, ts, mouse, btn, d) + diff_loss = F.mse_loss(pred, target) + sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) + return diff_loss, sc_loss def test_inference_cache(): From ac6119d2efcfe14330395027473fc6b3d5a5aa5e Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 02:04:17 +0000 Subject: [PATCH 46/71] debug --- owl_wms/models/gamerft_shortcut.py | 5 +++-- owl_wms/trainers/shortcut_trainer.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index c71bb8cc..6b0419af 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -198,8 +198,9 @@ def forward(self, x, y, mouse, btn, ema): pred = self.core(lerpd, y, ts, mouse, btn, d) diff_loss = F.mse_loss(pred, target) - sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) - + #sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) + sc_loss = diff_loss + return diff_loss, sc_loss def test_inference_cache(): diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index da648ae0..1d7687f2 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -102,7 +102,6 @@ def train(self): update_after_step = 0, update_every = 1 ) - freeze(self.ema) #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) def get_ema_core(): @@ -145,7 +144,9 @@ def get_ema_core(): diff_loss, sc_loss = self.model(batch_vid,batch_keyframe,batch_mouse,batch_btn, get_ema_core()) loss = diff_loss + sc_loss loss = loss / accum_steps - + + find_unused_params(self.model) + exit() self.scaler.scale(loss).backward() metrics.log('diffusion_loss', diff_loss) From d4257ef88b7ec94256dd238ac1b2a4a7bb9b14c8 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 02:28:09 +0000 Subject: [PATCH 47/71] debug stuff --- configs/shortcut.yml | 24 ++++++++++++------------ owl_wms/models/gamerft_shortcut.py | 3 +-- owl_wms/nn/rope.py | 7 ++++--- owl_wms/trainers/shortcut_trainer.py | 7 ++++--- 4 files changed, 21 insertions(+), 20 deletions(-) diff --git a/configs/shortcut.yml b/configs/shortcut.yml index acd4441b..891a9707 100644 --- a/configs/shortcut.yml +++ b/configs/shortcut.yml @@ -1,19 +1,19 @@ # Config for a simple 256 -> 16 autoencoder model: model_id: game_rft_shortcut - sample_size: 4 - channels: 128 + sample_size: 5 + channels: 64 - n_layers: 17 - n_heads: 16 - d_model: 1024 + n_layers: 5 + n_heads: 6 + d_model: 384 - tokens_per_frame: 16 + tokens_per_frame: 40 n_buttons: 11 n_mouse_axes: 2 cfg_prob: 0.1 - n_frames: 120 + n_frames: 60 causal: false @@ -21,11 +21,11 @@ train: trainer_id: shortcut data_id: cod_s3 data_kwargs: - window_length: 120 + window_length: 60 bucket_name: cod-data-latent-360x640to5x8 include_keyframe: true - target_batch_size: 256 + target_batch_size: 16 batch_size: 16 epochs: 200 @@ -46,8 +46,8 @@ train: sampler_id: shortcut sampler_kwargs: - window_length: 120 - num_frames: 120 + window_length: 60 + num_frames: 60 only_return_generated: true n_samples: 8 @@ -61,4 +61,4 @@ train: wandb: name: shahbuland project: video_models - run_name: v2 \ No newline at end of file + run_name: v2 diff --git a/owl_wms/models/gamerft_shortcut.py b/owl_wms/models/gamerft_shortcut.py index 6b0419af..d3f56c7d 100644 --- a/owl_wms/models/gamerft_shortcut.py +++ b/owl_wms/models/gamerft_shortcut.py @@ -198,8 +198,7 @@ def forward(self, x, y, mouse, btn, ema): pred = self.core(lerpd, y, ts, mouse, btn, d) diff_loss = F.mse_loss(pred, target) - #sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) - sc_loss = diff_loss + sc_loss = self.get_sc_loss(x_sc, y_sc, mouse_sc, btn_sc, ema) return diff_loss, sc_loss diff --git a/owl_wms/nn/rope.py b/owl_wms/nn/rope.py index 2a5402da..8c5b8cfd 100644 --- a/owl_wms/nn/rope.py +++ b/owl_wms/nn/rope.py @@ -80,9 +80,10 @@ def forward(self, q, k): q = eo.rearrange(q, 'b h (n m) d -> b h n m d', n=q.shape[2]//m,m=m) k = eo.rearrange(k, 'b h (n m) d -> b h n m d', n=n,m=m) - freqs = self.pos_emb.get_axial_freqs(n,m) - q = apply_rotary_emb(freqs[-truncate:], q) - k = apply_rotary_emb(freqs, k) + with torch.no_grad(): + freqs = self.pos_emb.get_axial_freqs(n,m) + q = apply_rotary_emb(freqs[-truncate:].detach(), q) + k = apply_rotary_emb(freqs.detach(), k) q = eo.rearrange(q, 'b h n m d -> b h (n m) d') k = eo.rearrange(k, 'b h n m d -> b h (n m) d') diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index 1d7687f2..c048b338 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -16,6 +16,7 @@ from ..utils.logging import LogHelper, to_wandb from ..muon import init_muon from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn +from ..models.gamerft_shortcut import get_sc_targets class ShortcutTrainer(BaseTrainer): """ @@ -93,6 +94,7 @@ def train(self): self.model = self.model.cuda().train() if self.world_size > 1: self.model = DDP(self.model, device_ids=[self.local_rank]) + self.decoder = self.decoder.cuda().eval().bfloat16() decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) @@ -145,10 +147,9 @@ def get_ema_core(): loss = diff_loss + sc_loss loss = loss / accum_steps - find_unused_params(self.model) - exit() self.scaler.scale(loss).backward() - + find_unused_params(self.model) + exit() metrics.log('diffusion_loss', diff_loss) metrics.log('shortcut_loss', sc_loss) From f4fde6d065d2ec1895ebb9fee31f26449de85df9 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 02:32:26 +0000 Subject: [PATCH 48/71] debug --- owl_wms/trainers/shortcut_trainer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index c048b338..3a545fd3 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -148,8 +148,7 @@ def get_ema_core(): loss = loss / accum_steps self.scaler.scale(loss).backward() - find_unused_params(self.model) - exit() + metrics.log('diffusion_loss', diff_loss) metrics.log('shortcut_loss', sc_loss) From 2b7142ccc09eae2fd0ba29fb34048060bc38a144 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 03:39:41 +0000 Subject: [PATCH 49/71] unused params thing --- owl_wms/trainers/shortcut_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/owl_wms/trainers/shortcut_trainer.py b/owl_wms/trainers/shortcut_trainer.py index 3a545fd3..1c96a9a4 100644 --- a/owl_wms/trainers/shortcut_trainer.py +++ b/owl_wms/trainers/shortcut_trainer.py @@ -93,7 +93,7 @@ def train(self): # Prepare model and ema self.model = self.model.cuda().train() if self.world_size > 1: - self.model = DDP(self.model, device_ids=[self.local_rank]) + self.model = DDP(self.model, find_unused_parameters=True) self.decoder = self.decoder.cuda().eval().bfloat16() decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) From 867851c03c79587622b7be168bf5dbcf74069170 Mon Sep 17 00:00:00 2001 From: Sami Date: Wed, 11 Jun 2025 00:33:14 -0400 Subject: [PATCH 50/71] slight refactors, setting up for latents and history --- Dockerfile | 3 ++ checkpoints/wm/dcae_hf_cod/basic.yml | 67 ++++++++++++++++++++++++++++ configs/webapp/config.yaml | 4 +- webapp/checkpoints/shortcut.yaml | 67 ++++++++++++++++++++++++++++ webapp/config.yaml | 25 +++++++++++ webapp/utils/models.py | 4 +- 6 files changed, 166 insertions(+), 4 deletions(-) create mode 100644 checkpoints/wm/dcae_hf_cod/basic.yml create mode 100644 webapp/checkpoints/shortcut.yaml create mode 100644 webapp/config.yaml diff --git a/Dockerfile b/Dockerfile index e3280b20..702b781d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,6 +5,7 @@ FROM nvidia/cuda:12.8.1-runtime-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV PYTHONDONTWRITEBYTECODE=1 +ENV PYTHONPATH=/app # Install system dependencies (without python3.12 first) RUN apt-get update && apt-get install -y --no-install-recommends \ @@ -44,6 +45,8 @@ RUN uv pip install --system torch torchvision --index-url https://download.pytor # Install other requirements from requirements.txt RUN uv pip install --system -r requirements.txt +RUN git submodule update --init --recursive + # Copy the entire application COPY . /app diff --git a/checkpoints/wm/dcae_hf_cod/basic.yml b/checkpoints/wm/dcae_hf_cod/basic.yml new file mode 100644 index 00000000..541a3991 --- /dev/null +++ b/checkpoints/wm/dcae_hf_cod/basic.yml @@ -0,0 +1,67 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: game_rft_core + sample_size: 4 + channels: 128 + + n_layers: 17 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 16 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 60 + + causal: false + +train: + trainer_id: rft + data_id: cod_latent + data_kwargs: + window_length: 60 + root: ../cod_data/BlackOpsColdWar + add_optical_flow: false + + target_batch_size: 320 + batch_size: 40 + + epochs: 200 + + opt: Muon + opt_kwargs: + lr: 1.0e-3 + momentum: 0.95 + adamw_lr: 1.0e-4 + adamw_wd: 1.0e-4 + adamw_eps: 1.0e-15 + adamw_betas: [0.9, 0.95] + adamw_keys: [core.proj_in, core.proj_out.proj] + + scheduler: null + + checkpoint_dir: checkpoints/v2 + resume_ckpt: checkpoints/v2/step_165000.pt + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: window + sampler_kwargs: + n_steps: 32 + cfg_scale: 1.3 + window_length: 60 + num_frames: 120 + noise_prev: 0.2 + only_return_generated: true + + vae_batch_size: 16 + vae_scale: 2.17 + n_samples: 8 + +wandb: + name: shahbuland + project: video_models + run_name: v2 \ No newline at end of file diff --git a/configs/webapp/config.yaml b/configs/webapp/config.yaml index ed6943f7..1de97cc9 100644 --- a/configs/webapp/config.yaml +++ b/configs/webapp/config.yaml @@ -1,5 +1,5 @@ -model_checkpoint_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/ckpt_165k_ema.pt" -run_config_path: "/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/basic.yml" +model_checkpoint_path: "checkpoints/wm/dcae_hf_cod/ckpt_165k_ema.pt" +run_config_path: "checkpoints/wm/dcae_hf_cod/basic.yml" device: "cuda" stream_config: diff --git a/webapp/checkpoints/shortcut.yaml b/webapp/checkpoints/shortcut.yaml new file mode 100644 index 00000000..fd6a1bfb --- /dev/null +++ b/webapp/checkpoints/shortcut.yaml @@ -0,0 +1,67 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: game_rft_shortcut + sample_size: 4 + channels: 128 + + n_layers: 17 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 16 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 120 + + causal: true + +train: + trainer_id: rft + data_id: cod_latent + data_kwargs: + window_length: 60 + root: ../cod_data/BlackOpsColdWar + add_optical_flow: false + + target_batch_size: 320 + batch_size: 40 + + epochs: 200 + + opt: Muon + opt_kwargs: + lr: 1.0e-3 + momentum: 0.95 + adamw_lr: 1.0e-4 + adamw_wd: 1.0e-4 + adamw_eps: 1.0e-15 + adamw_betas: [0.9, 0.95] + adamw_keys: [core.proj_in, core.proj_out.proj] + + scheduler: null + + checkpoint_dir: webapp/checkpoints/shortcut + resume_ckpt: webapp/checkpoints/shortcut/step_165000.pt + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: window + sampler_kwargs: + n_steps: 32 + cfg_scale: 1.3 + window_length: 60 + num_frames: 120 + noise_prev: 0.2 + only_return_generated: true + + vae_batch_size: 16 + vae_scale: 2.17 + n_samples: 8 + +wandb: + name: shahbuland + project: video_models + run_name: v2 \ No newline at end of file diff --git a/webapp/config.yaml b/webapp/config.yaml new file mode 100644 index 00000000..a5dd154a --- /dev/null +++ b/webapp/config.yaml @@ -0,0 +1,25 @@ +model_checkpoint_path: "webapp/checkpoints/shortcut/step_165000.pt" +run_config_path: "webapp/checkpoints/shortcut.yaml" +device: "cuda" + +stream_config: + fps: 20 + frames_per_batch: 60 + window_length: 60 + device: "cuda" + n_buttons: 11 + n_mouse_axes: 2 + mouse_range: [-1.0, 1.0] + action_margin_px_height: 150 + +sampling_config: + sampling_steps: 20 + vae_scale: 1.0 + cfg_scale: 1.3 + window_length: 60 + num_frames: 1 + noise_prev: 0.25 + +run_config: null # loaded at runtime from model_config_path, and used to access model and train config + + diff --git a/webapp/utils/models.py b/webapp/utils/models.py index 5734e165..e43d1d07 100644 --- a/webapp/utils/models.py +++ b/webapp/utils/models.py @@ -32,8 +32,8 @@ def validate(self) -> None: class ModelLoader: DEFAULT_PATHS = ModelPaths.from_strings( - config_path='/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/basic.yml', - checkpoint_path='/home/sami/owl-wms/checkpoints/wm/dcae_hf_cod/ckpt_165k_ema.pt' + config_path='webapp/checkpoints/shortcut.yaml', + checkpoint_path='webapp/checkpoints/shortcut/step_165000.pt' ) def __init__(self, paths: Optional[ModelPaths] = None): From a91e6986305338b0cb0ce1277f4491f4f8f38d64 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 05:01:24 +0000 Subject: [PATCH 51/71] add alt shortcut sampler --- owl_wms/models/gamerft_shortcut_simple.py | 190 +++++++++++++++++++++ owl_wms/nn/attn.py | 2 +- owl_wms/sampling/__init__.py | 5 +- owl_wms/sampling/shortcut_sampler.py | 75 +++++++- owl_wms/trainers/__init__.py | 3 + owl_wms/trainers/shortcut_trainer_2.py | 199 ++++++++++++++++++++++ 6 files changed, 470 insertions(+), 4 deletions(-) create mode 100644 owl_wms/models/gamerft_shortcut_simple.py create mode 100644 owl_wms/trainers/shortcut_trainer_2.py diff --git a/owl_wms/models/gamerft_shortcut_simple.py b/owl_wms/models/gamerft_shortcut_simple.py new file mode 100644 index 00000000..6cdc8147 --- /dev/null +++ b/owl_wms/models/gamerft_shortcut_simple.py @@ -0,0 +1,190 @@ +""" +Causal-First RFT With Shortcut objective +""" + +import torch +from torch import nn +import torch.nn.functional as F + +import einops as eo + +from ..nn.embeddings import ( + TimestepEmbedding, + StepEmbedding, + ControlEmbedding, + LearnedPosEnc +) +from ..nn.attn import UViT, FinalLayer +from ..nn.mmattn import MMUViT +from ..utils import freeze + +class ShortcutGameRFTCore(nn.Module): + def __init__(self, config): + super().__init__() + + self.transformer = UViT(config) + self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) + + self.step_embed = StepEmbedding(config.d_model) + self.t_embed = TimestepEmbedding(config.d_model) + + self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) + self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) + + self.config = config + + def sample(self, x, mouse, btn, kv_cache = None, t = None, d = None): + """ + This is a function that largely abstracts + away most things for the specific case where + you are only generating the one next token + + The return is one step sample always + """ + + b,n,c,h,w = x.shape + if t is None: + t = torch.ones_like(x[:,:,0,0,0]) + if d is None: + d = torch.ones_like(x[:,:,0,0,0]) + + return x - self.forward(x, t, mouse, btn, d, kv_cache) + + def forward(self, x, t, mouse, btn, d, kv_cache = None): + # x is [b,n,c,h,w] + # y is [b,1,c,h,w] + # t is [b,n] + # d is [b,n] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + + ctrl_cond = self.control_embed(mouse, btn) + t_cond = self.t_embed(t) + d_cond = self.step_embed(d) + + cond = ctrl_cond + t_cond + d_cond # [b,n,d] + + b,n,c,h,w = x.shape + x = eo.rearrange(x, 'b n c h w -> b (n h w) c') + x = self.proj_in(x) + x = self.transformer(x, cond, kv_cache) + x = self.proj_out(x, cond) # -> [b,n*hw,c] + x = eo.rearrange(x, 'b (n h w) c -> b n c h w', n=n,h=h,w=w) + + return x + +def sample_discrete_timesteps(steps, eps = 1.0e-6): + # steps is Tensor([1,4,2,64,16]) as an example + b,n = steps.shape + + ts_list = [] + ts = torch.rand(b, n, device=steps.device, dtype=steps.dtype) * (steps - eps) + ts = ts.clamp(eps).ceil() / steps + """ + Example, if d was all 2, ts would be [0,2] + so do clamp, then ceil will be 1 or 2 (0, 2] + then do t / 2 and get 0.5 or 1.0, our desired timesteps + """ + return ts + +def sample_steps(b, n, device, dtype, min_val = 0): + valid = torch.tensor([2**i for i in range(min_val, 8)]) # [1,2,...,128] + inds = torch.randint(low=0,high=len(valid), size = (b,n)) + steps = valid[inds].to(device=device,dtype=dtype) + return steps + +#@torch.compile() +@torch.no_grad() +def get_sc_targets(ema, x, mouse, btn, cfg_scale): + steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) + steps_fast = steps_slow / 2 + + dt_slow = 1./steps_slow + dt_fast = 1./steps_fast + + def expand(t): + #b,c,h,w = x.shape + #t = eo.repeat(t,'b -> b c h w',c=c,h=h,w=w) + #return t + return t[:,:,None,None,None] + + ts = sample_discrete_timesteps(steps_fast) + cfg_mask = torch.isclose(steps_slow, torch.ones_like(steps_slow)*128) + cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] + + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + pred_1_uncond = ema(x, ts, null_mouse, null_btn, steps_slow) + pred_1_cond = ema(x, ts, mouse, btn, steps_slow) + pred_1_cfg = pred_1_uncond + cfg_scale * (pred_1_cond - pred_1_uncond) + pred_1 = torch.where(cfg_mask, pred_1_cfg, pred_1_cond) + + x_new = x - pred_1 * expand(dt_slow) + ts_new = ts - dt_slow + + pred_2_uncond = ema(x_new, ts_new, null_mouse, null_btn, steps_slow) + pred_2_cond = ema(x_new, ts_new, mouse, btn, steps_slow) + pred_2_cfg = pred_2_uncond + cfg_scale * (pred_2_cond - pred_2_uncond) + pred_2 = torch.where(cfg_mask, pred_2_cfg, pred_2_cond) + + pred = 0.5 * (pred_1 + pred_2) + return pred, steps_fast, ts + +class ShortcutGameRFT(nn.Module): + def __init__(self, config): + super().__init__() + + self.core = ShortcutGameRFTCore(config) + self.cfg_prob = config.cfg_prob + + self.sc_frac = 0.25 + self.sc_max_steps = 128 + self.cfg_scale = 1.3 + + self.config = config + + def get_sc_loss(self, x, mouse, btn, ema): + target, steps, ts = get_sc_targets(ema, x, mouse, btn, self.cfg_scale) + pred = self.core(x, ts, mouse, btn, steps) + sc_loss = F.mse_loss(pred, target) + return sc_loss + + def forward(self, x, mouse, btn, ema): + # x is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + with torch.no_grad(): + _,n,c,h,w = x.shape + + # Split batches between consistency/rf + b = int(len(x) * (1 - self.sc_frac)) + x,x_sc = x[:b], x[b:] + mouse,mouse_sc = mouse[:b], mouse[b:] + btn,btn_sc = btn[:b], btn[b:] + + # Apply classifier-free guidance dropout + if self.cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + # Where mask is True, replace with zeros + mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) + btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) + + d = torch.ones_like(x[:,:,0,0,0])*self.sc_max_steps + ts = sample_discrete_timesteps(d) + ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() + + ts_exp = eo.repeat(ts, 'b n -> b n 1 1 1') + z = torch.randn_like(x) + + lerpd = x * (1. - ts_exp) + z * ts_exp + target = z - x + + pred = self.core(lerpd, ts, mouse, btn, d) + diff_loss = F.mse_loss(pred, target) + sc_loss = self.get_sc_loss(x_sc, mouse_sc, btn_sc, ema) + + return diff_loss, sc_loss \ No newline at end of file diff --git a/owl_wms/nn/attn.py b/owl_wms/nn/attn.py index 6174b72a..c1077a84 100644 --- a/owl_wms/nn/attn.py +++ b/owl_wms/nn/attn.py @@ -47,7 +47,7 @@ def forward(self, x, kv_cache = None): q,k,v = eo.rearrange(self.qkv(x), 'b n (three h d) -> three b h n d', three = 3, h = self.n_heads) q,k = self.qk_norm(q,k) - if not self.causal or len(kv_cache) > 0: + if not self.causal or (kv_cache is not None and len(kv_cache) > 0): mask = None else: mask = create_block_causal_mask(x.shape[1], self.tokens_per_frame).to(x.device) diff --git a/owl_wms/sampling/__init__.py b/owl_wms/sampling/__init__.py index 9488daf6..3de909d5 100644 --- a/owl_wms/sampling/__init__.py +++ b/owl_wms/sampling/__init__.py @@ -11,4 +11,7 @@ def get_sampler_cls(sampler_id): return WindowCFGSampler elif sampler_id == "shortcut": from .shortcut_sampler import CacheShortcutSampler - return CacheShortcutSampler \ No newline at end of file + return CacheShortcutSampler + elif sampler_id == "shortcut_2": + from .shortcut_sampler import WindowShortcutSamplerNoKeyframe + return WindowShortcutSamplerNoKeyframe \ No newline at end of file diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index 67be7315..dd61de9e 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -99,8 +99,6 @@ class WindowShortcutSampler: :param only_return_generated: Whether to only return the generated frames """ def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): - self.n_steps = n_steps - self.cfg_scale = cfg_scale self.window_length = window_length self.num_frames = num_frames self.only_return_generated = only_return_generated @@ -164,3 +162,76 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale x = decode_fn(x) return x, extended_mouse, extended_btn + +class WindowShortcutSamplerNoKeyframe: + """ + Same as above but with no cache + + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + self.window_length = window_length + self.num_frames = num_frames + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, history, mouse, btn, decode_fn = None, scale = 1): + # history is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + history = history[:,:self.window_length] + new_frames = [] + alpha = 0.25 # This number is special for our sampler + + # Extended fake controls to use during sampling + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, self.num_frames + self.window_length) + + # Initialize window history + window_history = history.clone() + + for frame_idx in tqdm(range(self.num_frames)): + # Setup window history + x = window_history[:,-self.window_length:].clone() + + # Noise all but last frame to alpha + x[:,:-1] = zlerp(x[:,:-1], alpha) + # Last frame starts as random noise + x[:,-1] = torch.randn_like(x[:,-1]) + + # Setup timesteps - alpha for context, 1.0 for generated + ts = torch.ones_like(x[:,:,0,0,0]) + ts[:,:-1] = alpha + + # Setup diffusion steps - 4 for context, 1 for generated + d = torch.ones_like(x[:,:,0,0,0]) + d[:,:-1] = 4 + + # Get current controls + curr_mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] + curr_btn = extended_btn[:,frame_idx:frame_idx+self.window_length] + + # Generate new frame + pred = model.sample(x, keyframe, curr_mouse, curr_btn, None, ts, d) + new_frame = pred[:,-1:] # Take only the last frame + new_frames.append(new_frame) + + # Add new frame to window history + window_history = torch.cat([window_history, new_frame], dim=1) + + new_frames = torch.cat(new_frames, dim=1) + x = torch.cat([history, new_frames], dim=1) + + if self.only_return_generated: + x = x[:,-self.num_frames:] + extended_mouse = extended_mouse[:,-self.num_frames:] + extended_btn = extended_btn[:,-self.num_frames:] + + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + + return x, extended_mouse, extended_btn \ No newline at end of file diff --git a/owl_wms/trainers/__init__.py b/owl_wms/trainers/__init__.py index 0a78fb14..09a4bebc 100644 --- a/owl_wms/trainers/__init__.py +++ b/owl_wms/trainers/__init__.py @@ -8,4 +8,7 @@ def get_trainer_cls(trainer_id): return CausVidTrainer if trainer_id == "shortcut": from .shortcut_trainer import ShortcutTrainer + return ShortcutTrainer + if trainer_id == "shortcut_2": + from .shortcut_trainer_2 import ShortcutTrainer return ShortcutTrainer \ No newline at end of file diff --git a/owl_wms/trainers/shortcut_trainer_2.py b/owl_wms/trainers/shortcut_trainer_2.py new file mode 100644 index 00000000..7a449fc3 --- /dev/null +++ b/owl_wms/trainers/shortcut_trainer_2.py @@ -0,0 +1,199 @@ +import torch +from ema_pytorch import EMA +import wandb +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist +import einops as eo + +from .base import BaseTrainer + +from ..utils import freeze, Timer, find_unused_params +from ..schedulers import get_scheduler_cls +from ..models import get_model_cls +from ..sampling import get_sampler_cls +from ..data import get_loader +from ..utils.logging import LogHelper, to_wandb +from ..muon import init_muon +from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn +from ..models.gamerft_shortcut import get_sc_targets + +class ShortcutTrainer(BaseTrainer): + """ + Trainer for rectified flow transformer with shortcut + + :param train_cfg: Configuration for training + :param logging_cfg: Configuration for logging + :param model_cfg: Configuration for model + :param global_rank: Rank across all devices. + :param local_rank: Rank for current device on this process. + :param world_size: Overall number of devices + """ + def __init__(self,*args,**kwargs): + super().__init__(*args,**kwargs) + + model_id = self.model_cfg.model_id + self.model = get_model_cls(model_id)(self.model_cfg) + + # Print model size + if self.rank == 0: + n_params = sum(p.numel() for p in self.model.parameters()) + print(f"Model has {n_params:,} parameters") + + self.ema = None + self.opt = None + self.scheduler = None + self.scaler = None + + self.total_step_counter = 0 + self.decoder = get_decoder_only( + self.train_cfg.vae_id, + self.train_cfg.vae_cfg_path, + self.train_cfg.vae_ckpt_path + ) + + freeze(self.decoder) + + def save(self): + save_dict = { + 'model' : self.model.state_dict(), + 'ema' : self.ema.state_dict(), + 'opt' : self.opt.state_dict(), + 'scaler' : self.scaler.state_dict(), + 'steps': self.total_step_counter + } + if self.scheduler is not None: + save_dict['scheduler'] = self.scheduler.state_dict() + super().save(save_dict) + + def load(self): + has_ckpt = False + try: + if self.train_cfg.resume_ckpt is not None: + save_dict = super().load(self.train_cfg.resume_ckpt) + has_ckpt = True + except: + print("Error loading checkpoint") + + if not has_ckpt: + return + + + self.model.load_state_dict(save_dict['model']) + self.ema.load_state_dict(save_dict['ema']) + self.opt.load_state_dict(save_dict['opt']) + if self.scheduler is not None and 'scheduler' in save_dict: + self.scheduler.load_state_dict(save_dict['scheduler']) + self.scaler.load_state_dict(save_dict['scaler']) + self.total_step_counter = save_dict['steps'] + + def train(self): + torch.cuda.set_device(self.local_rank) + + # Prepare model and ema + self.model = self.model.cuda().train() + if self.world_size > 1: + self.model = DDP(self.model, device_ids=[self.local_rank]) + + self.decoder = self.decoder.cuda().eval().bfloat16() + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + + self.ema = EMA( + self.model, + beta = 0.999, + update_after_step = 0, + update_every = 1 + ) + #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) + + def get_ema_core(): + if self.world_size > 1: + return self.ema.ema_model.module.core + else: + return self.ema.ema_model.core + + # No muon pls + self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) + + # Grad accum setup and scaler + accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size + accum_steps = max(1, accum_steps) + self.scaler = torch.amp.GradScaler() + ctx = torch.amp.autocast('cuda',torch.bfloat16) + + self.load() + + # Timer reset + timer = Timer() + timer.reset() + metrics = LogHelper() + if self.rank == 0: + wandb.watch(self.get_module(), log = 'all') + + # Dataset setup + loader = get_loader(self.train_cfg.data_id, self.train_cfg.batch_size, **self.train_cfg.data_kwargs) + sampler = get_sampler_cls(self.train_cfg.sampler_id)(**self.train_cfg.sampler_kwargs) + + local_step = 0 + for _ in range(self.train_cfg.epochs): + for batch_vid, batch_mouse, batch_btn in loader: + batch_vid = batch_vid.cuda().bfloat16() / self.train_cfg.vae_scale + batch_mouse = batch_mouse.cuda().bfloat16() + batch_btn = batch_btn.cuda().bfloat16() + + with ctx: + diff_loss, sc_loss = self.model(batch_vid,batch_mouse,batch_btn, get_ema_core()) + loss = diff_loss + sc_loss + loss = loss / accum_steps + + self.scaler.scale(loss).backward() + + metrics.log('diffusion_loss', diff_loss) + metrics.log('shortcut_loss', sc_loss) + + local_step += 1 + if local_step % accum_steps == 0: + # Updates + self.scaler.unscale_(self.opt) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + + self.scaler.step(self.opt) + self.opt.zero_grad(set_to_none=True) + + self.scaler.update() + + if self.scheduler is not None: + self.scheduler.step() + self.ema.update() + + # Do logging + with torch.no_grad(): + wandb_dict = metrics.pop() + wandb_dict['time'] = timer.hit() + wandb_dict['lr'] = self.opt.param_groups[0]['lr'] + timer.reset() + + # Sampling commented out for now + if self.total_step_counter % self.train_cfg.sample_interval == 0: + with ctx, torch.no_grad(): + n_samples = self.train_cfg.n_samples + samples, sample_mouse, sample_button = sampler( + get_ema_core(), + batch_vid[:n_samples], + batch_mouse[:n_samples], + batch_btn[:n_samples], + decode_fn = decode_fn, + scale=self.train_cfg.vae_scale + ) # -> [b,n,c,h,w] + if self.rank == 0: wandb_dict['samples'] = to_wandb(samples, sample_mouse, sample_button) + + + if self.rank == 0: + wandb.log(wandb_dict) + + self.total_step_counter += 1 + if self.total_step_counter % self.train_cfg.save_interval == 0: + if self.rank == 0: + self.save() + + self.barrier() From 5b5a04d01456298a483ff04550e091433b837c9d Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 05:14:45 +0000 Subject: [PATCH 52/71] push it --- owl_wms/models/__init__.py | 3 +++ owl_wms/sampling/shortcut_sampler.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/owl_wms/models/__init__.py b/owl_wms/models/__init__.py index 5f70facd..91c28793 100644 --- a/owl_wms/models/__init__.py +++ b/owl_wms/models/__init__.py @@ -5,6 +5,9 @@ def get_model_cls(model_id): if model_id == "game_rft_shortcut": from .gamerft_shortcut import ShortcutGameRFT return ShortcutGameRFT + if model_id == "shortcut_2": + from .gamerft_shortcut_simple import ShortcutGameRFT + return ShortcutGameRFT diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index dd61de9e..ed23f173 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -215,7 +215,7 @@ def __call__(self, model, history, mouse, btn, decode_fn = None, scale = 1): curr_btn = extended_btn[:,frame_idx:frame_idx+self.window_length] # Generate new frame - pred = model.sample(x, keyframe, curr_mouse, curr_btn, None, ts, d) + pred = model.sample(x, curr_mouse, curr_btn, None, ts, d) new_frame = pred[:,-1:] # Take only the last frame new_frames.append(new_frame) From f02ef0b4e0a7688fa56498699995ca81801ff9b8 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 05:17:17 +0000 Subject: [PATCH 53/71] a --- configs/shortcut_2.yml | 64 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 configs/shortcut_2.yml diff --git a/configs/shortcut_2.yml b/configs/shortcut_2.yml new file mode 100644 index 00000000..9942f3c5 --- /dev/null +++ b/configs/shortcut_2.yml @@ -0,0 +1,64 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: shortcut_2 + sample_size: 5 + channels: 64 + + n_layers: 5 + n_heads: 6 + d_model: 384 + + tokens_per_frame: 40 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 60 + + causal: true + +train: + trainer_id: shortcut_2 + data_id: cod_s3 + data_kwargs: + window_length: 60 + bucket_name: cod-data-latent-360x640to5x8 + include_keyframe: false + + target_batch_size: 16 + batch_size: 16 + + epochs: 200 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + weight_decay: 0.1 + eps: 1.0e-15 + betas: [0.9, 0.95] + + scheduler: null + + checkpoint_dir: checkpoints/360p + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: shortcut_2 + sampler_kwargs: + window_length: 60 + num_frames: 60 + only_return_generated: true + + n_samples: 8 + + vae_id: 720pr3dc + vae_batch_size: 4 + vae_scale: 0.35 + vae_cfg_path: configs/owl_vaes/128x_cod_stage2.yml + vae_ckpt_path: 720p_cod_vae_30m_35k_steps.pt + +wandb: + name: shahbuland + project: video_models + run_name: v2 From 1ca3e6f1041e017f16549c1e5447902e5b59a384 Mon Sep 17 00:00:00 2001 From: Sami Date: Wed, 11 Jun 2025 01:22:20 -0400 Subject: [PATCH 54/71] self-forcing --- configs/self_forcing.yaml | 76 ++++++ owl_wms/models/causal_gamerft.py | 298 +++++++++++++++++++++ owl_wms/trainers/self_forcing.py | 439 +++++++++++++++++++++++++++++++ 3 files changed, 813 insertions(+) create mode 100644 configs/self_forcing.yaml create mode 100644 owl_wms/models/causal_gamerft.py create mode 100644 owl_wms/trainers/self_forcing.py diff --git a/configs/self_forcing.yaml b/configs/self_forcing.yaml new file mode 100644 index 00000000..606fc6ef --- /dev/null +++ b/configs/self_forcing.yaml @@ -0,0 +1,76 @@ +# Config for Self-Forcing training with autoregressive rollout +model: + model_id: game_rft + sample_size: [5,8] + channels: 64 + + n_layers: 17 + n_heads: 20 + d_model: 1280 + + tokens_per_frame: 40 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 30 + + causal: true # Enable causal attention + +train: + trainer_id: self_forcing + data_id: cod_latent + data_kwargs: + window_length: 30 + root: ../cod_data/BlackOpsColdWar + add_optical_flow: false + + target_batch_size: 256 + batch_size: 32 + + epochs: 200 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + weight_decay: 1.0e-4 + eps: 1.0e-15 + betas: [0.9, 0.95] + + scheduler: null + + checkpoint_dir: checkpoints/self_forcing + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: window + sampler_kwargs: + n_steps: 20 + cfg_scale: 1.3 + window_length: 30 + num_frames: 60 + noise_prev: 0.2 + only_return_generated: true + + n_samples: 8 + + vae_id: 720pr3dc + vae_batch_size: 4 + vae_scale: 0.35 + vae_cfg_path: configs/owl_vaes/128x_cod_stage2.yml + vae_ckpt_path: 720p_cod_vae_30m_35k_steps.pt + + # Self-forcing specific parameters + teacher_ckpt: checkpoints/bidirectional_teacher/best.pt # Pretrained bidirectional model + loss_type: dmd # Options: dmd, sid, gan + gradient_steps: 2 # Number of steps to backprop through + rollout_steps: 5 # Total autoregressive rollout length + stochastic_steps: true # Random gradient truncation + update_ratio: 5 # Critic updates per generator update + cfg_scale: 1.3 # Classifier-free guidance scale + +wandb: + name: shahbuland + project: video_models + run_name: self_forcing_dmd \ No newline at end of file diff --git a/owl_wms/models/causal_gamerft.py b/owl_wms/models/causal_gamerft.py new file mode 100644 index 00000000..8d1de154 --- /dev/null +++ b/owl_wms/models/causal_gamerft.py @@ -0,0 +1,298 @@ +""" +Causal GameRFT model with KV cache support for efficient autoregressive generation +""" + +import torch +from torch import nn +import torch.nn.functional as F +import einops as eo + +from ..nn.embeddings import ( + TimestepEmbedding, + ControlEmbedding, + LearnedPosEnc +) +from ..nn.attn import UViT, FinalLayer +from ..nn.kv_cache import KVCache + +class CausalGameRFTCore(nn.Module): + """ + Core model with causal attention support and KV caching + """ + def __init__(self, config): + super().__init__() + + # Modify config for causal attention + self.causal = config.causal + self.tokens_per_frame = config.tokens_per_frame + self.n_frames = config.n_frames + + # Initialize transformer with causal flag + config_copy = deepcopy(config) + config_copy.causal = self.causal + self.transformer = UViT(config_copy) + + self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) + self.t_embed = TimestepEmbedding(config.d_model) + + self.proj_in = nn.Linear(config.channels, config.d_model, bias=False) + self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) + + self.pos_enc = LearnedPosEnc(config.tokens_per_frame * config.n_frames, config.d_model) + + # For caching context frames + self.cached_frames = None + self.cached_frame_count = 0 + + def forward(self, x, t, mouse, btn, kv_cache=None, use_cache=False): + """ + Forward pass with optional KV caching + + Args: + x: Input frames [b,n,c,h,w] + t: Timesteps [b,n] + mouse: Mouse inputs [b,n,2] + btn: Button inputs [b,n,n_buttons] + kv_cache: Optional KVCache object for caching + use_cache: Whether to use cached context + """ + # Handle caching for autoregressive generation + if use_cache and self.cached_frames is not None: + # Concatenate cached frames with new input + x = torch.cat([self.cached_frames, x], dim=1) + t = torch.cat([self.cached_timesteps, t], dim=1) + mouse = torch.cat([self.cached_mouse, mouse], dim=1) + btn = torch.cat([self.cached_btn, btn], dim=1) + + # Update cache with new frames + self.cached_frames = x + self.cached_timesteps = t + self.cached_mouse = mouse + self.cached_btn = btn + elif use_cache: + # Initialize cache + self.cached_frames = x + self.cached_timesteps = t + self.cached_mouse = mouse + self.cached_btn = btn + + # Standard forward pass + ctrl_cond = self.control_embed(mouse, btn) + t_cond = self.t_embed(t) + + cond = ctrl_cond + t_cond # [b,n,d] + + b, n, c, h, w = x.shape + x = eo.rearrange(x, 'b n c h w -> b (n h w) c') + + x = self.proj_in(x) + x = self.pos_enc(x) + + # Pass KV cache to transformer if provided + x = self.transformer(x, cond, kv_cache=kv_cache) + + x = self.proj_out(x, cond) # -> [b,n*hw,c] + x = eo.rearrange(x, 'b (n h w) c -> b n c h w', n=n, h=h, w=w) + + # If using cache, only return the newly generated frames + if use_cache and self.cached_frame_count > 0: + x = x[:, -1:] # Return only last frame + + return x + + def reset_cache(self): + """Reset the cached context frames""" + self.cached_frames = None + self.cached_timesteps = None + self.cached_mouse = None + self.cached_btn = None + self.cached_frame_count = 0 + +class CausalGameRFT(nn.Module): + """ + Causal GameRFT model for autoregressive video generation + """ + def __init__(self, config): + super().__init__() + + self.core = CausalGameRFTCore(config) + self.cfg_prob = config.cfg_prob + self.causal = config.causal + + def forward(self, x, mouse, btn, return_dict=False, cfg_prob=None, kv_cache=None, use_cache=False): + """ + Forward pass with diffusion loss computation + + For training: Standard diffusion loss + For generation: Can use KV cache for efficiency + """ + b, n, c, h, w = x.shape + + # Apply classifier-free guidance dropout + if cfg_prob is None: + cfg_prob = self.cfg_prob + if cfg_prob > 0.0 and self.training: + mask = torch.rand(b, device=x.device) <= cfg_prob + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + # Where mask is True, replace with zeros + mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) + btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) + + if self.training: + # Standard diffusion training + with torch.no_grad(): + ts = torch.rand(b, n, device=x.device, dtype=x.dtype).sigmoid() + + ts_exp = eo.repeat(ts, 'b n -> b n 1 1 1') + z = torch.randn_like(x) + + lerpd = x * (1. - ts_exp) + z * ts_exp + target = z - x + + pred = self.core(lerpd, ts, mouse, btn) + diff_loss = F.mse_loss(pred, target) + + if not return_dict: + return diff_loss + else: + return { + 'diffusion_loss': diff_loss, + 'lerpd': lerpd, + 'pred': pred, + 'ts': ts, + 'z': z + } + else: + # Generation mode - can use KV cache + return self.core(x, torch.zeros(b, n, device=x.device), mouse, btn, + kv_cache=kv_cache, use_cache=use_cache) + + def generate_next_frame(self, context, mouse, btn, num_steps=50, cfg_scale=1.3, kv_cache=None): + """ + Generate the next frame given context frames + + Args: + context: Context frames [b, n_context, c, h, w] + mouse: Mouse inputs for context + 1 new frame [b, n_context+1, 2] + btn: Button inputs for context + 1 new frame [b, n_context+1, n_buttons] + num_steps: Number of denoising steps + cfg_scale: Classifier-free guidance scale + kv_cache: Optional KV cache for efficiency + + Returns: + next_frame: Generated next frame [b, 1, c, h, w] + """ + b = context.shape[0] + device = context.device + + # Initialize next frame with noise + next_frame = torch.randn(b, 1, *context.shape[2:], device=device) + + # Combine context and noisy next frame + full_sequence = torch.cat([context, next_frame], dim=1) + + # Denoising loop + for step in range(num_steps): + t = (1.0 - step / num_steps) * torch.ones(b, full_sequence.shape[1], device=device) + + # Zero timestep for context frames (they're clean) + t[:, :-1] = 0.0 + + with torch.no_grad(): + # Conditional prediction + pred_cond = self.core(full_sequence, t, mouse, btn, kv_cache=kv_cache) + + if cfg_scale > 1.0: + # Unconditional prediction for CFG + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + pred_uncond = self.core(full_sequence, t, null_mouse, null_btn, kv_cache=kv_cache) + + # Apply CFG + pred = pred_uncond + cfg_scale * (pred_cond - pred_uncond) + else: + pred = pred_cond + + # Update only the last frame + noise_level = t[:, -1:, None, None, None] + full_sequence[:, -1:] = full_sequence[:, -1:] - pred[:, -1:] * noise_level / num_steps + + return full_sequence[:, -1:] + + def generate_sequence(self, initial_frame, mouse, btn, num_frames, + num_steps=50, cfg_scale=1.3, use_kv_cache=True): + """ + Generate a full sequence autoregressively + + Args: + initial_frame: Starting frame [b, 1, c, h, w] + mouse: Mouse inputs for full sequence [b, num_frames, 2] + btn: Button inputs for full sequence [b, num_frames, n_buttons] + num_frames: Number of frames to generate + num_steps: Denoising steps per frame + cfg_scale: CFG scale + use_kv_cache: Whether to use KV caching for efficiency + + Returns: + sequence: Generated sequence [b, num_frames, c, h, w] + """ + generated = [initial_frame] + + # Initialize KV cache if requested + kv_cache = KVCache(self.core.transformer.config) if use_kv_cache else None + if kv_cache: + kv_cache.reset(initial_frame.shape[0]) + + # Reset any internal caching + self.core.reset_cache() + + for frame_idx in range(1, num_frames): + # Get context and actions up to current frame + context = torch.cat(generated, dim=1) + frame_mouse = mouse[:, :frame_idx+1] + frame_btn = btn[:, :frame_idx+1] + + # Generate next frame + next_frame = self.generate_next_frame( + context, frame_mouse, frame_btn, + num_steps=num_steps, cfg_scale=cfg_scale, + kv_cache=kv_cache + ) + + generated.append(next_frame) + + return torch.cat(generated, dim=1) + + +if __name__ == "__main__": + from ..configs import Config + from copy import deepcopy + + # Test causal model + cfg = Config.from_yaml("configs/basic.yml").model + cfg.causal = True + model = CausalGameRFT(cfg).cuda().bfloat16() + + # Test training forward pass + with torch.no_grad(): + x = torch.randn(2, 30, 128, 4, 4, device='cuda', dtype=torch.bfloat16) + mouse = torch.randn(2, 30, 2, device='cuda', dtype=torch.bfloat16) + btn = torch.randn(2, 30, 11, device='cuda', dtype=torch.bfloat16) + + loss = model(x, mouse, btn) + print(f"Training loss: {loss.item()}") + + # Test generation + model.eval() + with torch.no_grad(): + initial = torch.randn(2, 1, 128, 4, 4, device='cuda', dtype=torch.bfloat16) + mouse_seq = torch.randn(2, 10, 2, device='cuda', dtype=torch.bfloat16) + btn_seq = torch.randn(2, 10, 11, device='cuda', dtype=torch.bfloat16) + + generated = model.generate_sequence( + initial, mouse_seq, btn_seq, + num_frames=10, num_steps=20, cfg_scale=1.3 + ) + print(f"Generated sequence shape: {generated.shape}") \ No newline at end of file diff --git a/owl_wms/trainers/self_forcing.py b/owl_wms/trainers/self_forcing.py new file mode 100644 index 00000000..e3eb4c0f --- /dev/null +++ b/owl_wms/trainers/self_forcing.py @@ -0,0 +1,439 @@ +""" +Self-Forcing Trainer for Game World Model +Implements autoregressive self-rollout training with proper gradient truncation +""" + +import torch +from ema_pytorch import EMA +import wandb +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist +import einops as eo +from copy import deepcopy + +from .base import BaseTrainer +from ..utils import freeze, unfreeze, Timer, find_unused_params, versatile_load +from ..schedulers import get_scheduler_cls +from ..models import get_model_cls +from ..sampling import get_sampler_cls +from ..data import get_loader +from ..utils.logging import LogHelper, to_wandb +from ..muon import init_muon +from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn +from copy import deepcopy + +class SelfForcingTrainer(BaseTrainer): + """ + Self-Forcing Trainer implementing autoregressive self-rollout training + + Key differences from CausVid: + 1. Uses self-generated context during training (not ground truth) + 2. Implements gradient truncation with stochastic steps + 3. Supports DMD, SiD, and GAN losses on correct distribution + """ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.model_cfg = self.config.model + + model_id = self.model_cfg.model_id + + # Create student (causal) and teacher (non-causal) configs + student_cfg = deepcopy(self.model_cfg) + teacher_cfg = deepcopy(self.model_cfg) + + student_cfg.causal = True + teacher_cfg.causal = False + + # Initialize models + self.model = get_model_cls(model_id)(student_cfg) + self.score_real = get_model_cls(model_id)(teacher_cfg) + + # Load pretrained teacher + if self.train_cfg.teacher_ckpt: + self.score_real.load_state_dict(versatile_load(self.train_cfg.teacher_ckpt)) + freeze(self.score_real) + + # Initialize fake score for DMD/SiD losses + self.score_fake = deepcopy(self.score_real) + + # Print model size + if self.rank == 0: + n_params = sum(p.numel() for p in self.model.parameters()) + print(f"Model has {n_params:,} parameters") + + self.ema = None + self.opt = None + self.s_fake_opt = None + self.scheduler = None + self.s_fake_scaler = None + self.scaler = None + + self.total_step_counter = 0 + + # Initialize VAE decoder + self.decoder = get_decoder_only( + self.train_cfg.vae_id, + self.train_cfg.vae_cfg_path, + self.train_cfg.vae_ckpt_path + ) + freeze(self.decoder) + + # Self-forcing specific parameters + self.loss_type = self.train_cfg.get('loss_type', 'dmd') # dmd, sid, or gan + self.gradient_steps = self.train_cfg.get('gradient_steps', 1) # Number of steps to backprop + self.rollout_steps = self.train_cfg.get('rollout_steps', 5) # Total rollout length + self.stochastic_steps = self.train_cfg.get('stochastic_steps', True) # Random gradient truncation + self.update_ratio = self.train_cfg.get('update_ratio', 5) # Critic updates per generator update + self.cfg_scale = self.train_cfg.get('cfg_scale', 1.3) + + def save(self): + save_dict = { + 'model': self.model.state_dict(), + 'ema': self.ema.state_dict(), + 'opt': self.opt.state_dict(), + 'scaler': self.scaler.state_dict(), + 'score_fake': self.score_fake.state_dict(), + 's_fake_opt': self.s_fake_opt.state_dict(), + 's_fake_scaler': self.s_fake_scaler.state_dict(), + 'steps': self.total_step_counter + } + if self.scheduler is not None: + save_dict['scheduler'] = self.scheduler.state_dict() + super().save(save_dict) + + def load(self): + has_ckpt = False + try: + if self.train_cfg.resume_ckpt is not None: + save_dict = super().load(self.train_cfg.resume_ckpt) + has_ckpt = True + except: + print("Error loading checkpoint") + + if not has_ckpt: + return + + self.model.load_state_dict(save_dict['model']) + self.ema.load_state_dict(save_dict['ema']) + self.opt.load_state_dict(save_dict['opt']) + if self.scheduler is not None and 'scheduler' in save_dict: + self.scheduler.load_state_dict(save_dict['scheduler']) + self.scaler.load_state_dict(save_dict['scaler']) + self.score_fake.load_state_dict(save_dict['score_fake']) + self.s_fake_opt.load_state_dict(save_dict['s_fake_opt']) + self.s_fake_scaler.load_state_dict(save_dict['s_fake_scaler']) + self.total_step_counter = save_dict['steps'] + + def autoregressive_rollout(self, initial_latent, mouse, btn, decode_fn=None): + """ + Perform autoregressive rollout with gradient truncation + + Args: + initial_latent: Initial frame(s) [b, init_frames, c, h, w] + mouse: Mouse inputs for entire sequence [b, n, 2] + btn: Button inputs for entire sequence [b, n, n_buttons] + decode_fn: Optional decode function for visualization + + Returns: + generated_latents: Full generated sequence + gradient_mask: Boolean mask indicating which frames get gradients + """ + b = initial_latent.shape[0] + device = initial_latent.device + + # Initialize output with initial frames + generated_latents = [initial_latent] + + # Determine gradient truncation point + if self.stochastic_steps: + # Randomly select which steps to backprop through + grad_start = torch.randint( + max(0, self.rollout_steps - self.gradient_steps), + self.rollout_steps, + (1,) + ).item() + else: + # Always backprop through last gradient_steps + grad_start = self.rollout_steps - self.gradient_steps + + # Generate frames autoregressively + for step in range(self.rollout_steps): + # Get context from previously generated frames + context = torch.cat(generated_latents, dim=1) + context_frames = context.shape[1] + + # Get corresponding actions + step_mouse = mouse[:, :context_frames] + step_btn = btn[:, :context_frames] + + # Determine if this step needs gradients + needs_grad = step >= grad_start + + with torch.set_grad_enabled(needs_grad): + # Add noise to last frame for next prediction + noisy_next = torch.randn(b, 1, *initial_latent.shape[2:], device=device) + + # Prepare input + model_input = torch.cat([context, noisy_next], dim=1) + model_mouse = mouse[:, :context_frames + 1] + model_btn = btn[:, :context_frames + 1] + + # Generate next frame + with torch.amp.autocast('cuda', torch.bfloat16): + # Run diffusion model to denoise + ts = torch.ones(b, context_frames + 1, device=device) + ts[:, -1] = 0.99 # High noise for last frame + ts[:, :-1] = 0.0 # Clean context + + pred = self.model.core(model_input, ts, model_mouse, model_btn) + next_frame = model_input - pred * ts[:, :, None, None, None] + next_frame = next_frame[:, -1:] # Take only the newly generated frame + + generated_latents.append(next_frame) + + # Concatenate all generated frames + full_sequence = torch.cat(generated_latents, dim=1) + + # Create gradient mask + gradient_mask = torch.zeros(b, full_sequence.shape[1], dtype=torch.bool, device=device) + gradient_mask[:, initial_latent.shape[1] + grad_start:] = True + + return full_sequence, gradient_mask + + def compute_dmd_loss(self, generated, mouse, btn, gradient_mask): + """Compute DMD loss on generated sequence""" + s_real_fn = self.score_real.core + s_fake_fn = self.score_fake.module.core if self.world_size > 1 else self.score_fake.core + + with torch.no_grad(): + b, n, c, h, w = generated.shape + ts = torch.rand(b, n, device=generated.device).sigmoid() + z = torch.randn_like(generated) + ts_exp = ts[:, :, None, None, None] + lerpd = generated * (1. - ts_exp) + z * ts_exp + + # Compute real score with CFG + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + s_real_uncond = s_real_fn(lerpd, ts, null_mouse, null_btn) + s_real_cond = s_real_fn(lerpd, ts, mouse, btn) + s_real = s_real_uncond + self.cfg_scale * (s_real_cond - s_real_uncond) + + # Compute fake score + s_fake = s_fake_fn(lerpd, ts, mouse, btn) + + # DMD gradient + grad = (s_fake - s_real) + + # Normalize + p_real = (generated - s_real) + normalizer = torch.abs(p_real).mean(dim=[2, 3, 4], keepdim=True) + grad = grad / (normalizer + 1e-6) + grad = torch.nan_to_num(grad) + + # Apply gradient mask + if gradient_mask is not None: + grad = grad * gradient_mask[:, :, None, None, None] + + dmd_loss = 0.5 * F.mse_loss(generated.double(), (generated - grad).double()) + + return dmd_loss + + def compute_sid_loss(self, generated, mouse, btn, gradient_mask): + """Compute SiD loss on generated sequence""" + # Similar to DMD but with SiD formulation + # Implementation details from SiD paper + raise NotImplementedError("SiD loss not yet implemented") + + def compute_gan_loss(self, generated, mouse, btn, gradient_mask, train_generator=True): + """Compute GAN loss on generated sequence""" + # Implementation for GAN-based loss + raise NotImplementedError("GAN loss not yet implemented") + + def train(self): + torch.cuda.set_device(self.local_rank) + + # Prepare models + self.model = self.model.cuda().train() + self.decoder = self.decoder.cuda().eval().bfloat16() + self.score_real = self.score_real.cuda().eval().bfloat16() + self.score_fake = self.score_fake.cuda().train() + + if self.world_size > 1: + self.model = DDP(self.model) + self.score_fake = DDP(self.score_fake) + + freeze(self.decoder) + freeze(self.score_real) + + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + + # Initialize EMA + self.ema = EMA( + self.model, + beta=0.999, + update_after_step=0, + update_every=1 + ) + + # Initialize optimizers + self.opt = getattr(torch.optim, self.train_cfg.opt)( + self.model.parameters(), + **self.train_cfg.opt_kwargs + ) + self.s_fake_opt = getattr(torch.optim, self.train_cfg.opt)( + self.score_fake.parameters(), + **self.train_cfg.opt_kwargs + ) + + if self.train_cfg.scheduler is not None: + self.scheduler = get_scheduler_cls(self.train_cfg.scheduler)( + self.opt, + **self.train_cfg.scheduler_kwargs + ) + + # Scalers for mixed precision + self.s_fake_scaler = torch.amp.GradScaler() + self.scaler = torch.amp.GradScaler() + ctx = torch.amp.autocast('cuda', torch.bfloat16) + + self.load() + + # Setup logging + timer = Timer() + timer.reset() + metrics = LogHelper() + if self.rank == 0: + wandb.watch(self.get_module(), log='all') + + # Dataset setup + loader = get_loader( + self.train_cfg.data_id, + self.train_cfg.batch_size, + **self.train_cfg.data_kwargs + ) + sampler = get_sampler_cls(self.train_cfg.sampler_id)(**self.train_cfg.sampler_kwargs) + + def optimizer_step(loss, model, scaler, optimizer): + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + scaler.step(optimizer) + optimizer.zero_grad(set_to_none=True) + scaler.update() + + # Training loop + loader = iter(loader) + while True: + # Update critic/fake score + if self.loss_type in ['dmd', 'sid']: + freeze(self.model) + unfreeze(self.score_fake) + + for _ in range(self.update_ratio): + batch_vid, batch_mouse, batch_btn = next(loader) + + # Get initial frames + initial_frames = batch_vid[:, :1] # Use first frame as initial + + with torch.no_grad(): + # Generate sequence autoregressively + generated, _ = self.autoregressive_rollout( + initial_frames, + batch_mouse, + batch_btn + ) + + # Train fake score on generated data + with ctx: + s_fake_loss = self.score_fake( + generated.detach(), + batch_mouse[:, :generated.shape[1]], + batch_btn[:, :generated.shape[1]] + ) + + optimizer_step(s_fake_loss, self.score_fake, self.s_fake_scaler, self.s_fake_opt) + metrics.log('s_fake_loss', s_fake_loss) + + # Update generator + unfreeze(self.model) + freeze(self.score_fake) + + batch_vid, batch_mouse, batch_btn = next(loader) + initial_frames = batch_vid[:, :1] + + # Generate with gradients + with ctx: + generated, gradient_mask = self.autoregressive_rollout( + initial_frames, + batch_mouse, + batch_btn + ) + + # Compute loss based on selected type + if self.loss_type == 'dmd': + loss = self.compute_dmd_loss( + generated, + batch_mouse[:, :generated.shape[1]], + batch_btn[:, :generated.shape[1]], + gradient_mask + ) + elif self.loss_type == 'sid': + loss = self.compute_sid_loss( + generated, + batch_mouse[:, :generated.shape[1]], + batch_btn[:, :generated.shape[1]], + gradient_mask + ) + elif self.loss_type == 'gan': + loss = self.compute_gan_loss( + generated, + batch_mouse[:, :generated.shape[1]], + batch_btn[:, :generated.shape[1]], + gradient_mask, + train_generator=True + ) + + metrics.log(f'{self.loss_type}_loss', loss) + + optimizer_step(loss, self.model, self.scaler, self.opt) + self.ema.update() + + # Logging and visualization + with torch.no_grad(): + wandb_dict = metrics.pop() + wandb_dict['time'] = timer.hit() + timer.reset() + + if self.total_step_counter % self.train_cfg.sample_interval == 0: + with ctx, torch.no_grad(): + n_samples = self.train_cfg.n_samples + + # Get EMA model for sampling + ema_core = self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core + + # Sample using the trained model + samples, sample_mouse, sample_button = sampler( + ema_core, + batch_vid[:n_samples, :1], # Initial frame + batch_mouse[:n_samples], + batch_btn[:n_samples], + decode_fn=decode_fn, + scale=self.train_cfg.vae_scale + ) + + if self.rank == 0: + wandb_dict['samples'] = to_wandb(samples, sample_mouse, sample_button) + + if self.rank == 0: + wandb.log(wandb_dict) + + self.total_step_counter += 1 + if self.total_step_counter % self.train_cfg.save_interval == 0: + if self.rank == 0: + self.save() + + self.barrier() \ No newline at end of file From eb58fc8089b2d5414d567c30354d8bae2a0e7dee Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 05:22:47 +0000 Subject: [PATCH 55/71] add back muon --- configs/shortcut_2.yml | 16 +++++++++++----- owl_wms/trainers/shortcut_trainer_2.py | 13 +++++++++---- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/configs/shortcut_2.yml b/configs/shortcut_2.yml index 9942f3c5..15e34bc6 100644 --- a/configs/shortcut_2.yml +++ b/configs/shortcut_2.yml @@ -30,12 +30,18 @@ train: epochs: 200 - opt: AdamW + opt: Muon opt_kwargs: - lr: 1.0e-4 - weight_decay: 0.1 - eps: 1.0e-15 - betas: [0.9, 0.95] + lr: 1.0e-3 + momentum: 0.95 + adamw_lr: 1.0e-4 + adamw_wd: 1.0e-4 + adamw_eps: 1.0e-15 + adamw_betas: [0.9, 0.95] + adamw_keys: [ + core.proj_in, + core.proj_out.proj + ] scheduler: null diff --git a/owl_wms/trainers/shortcut_trainer_2.py b/owl_wms/trainers/shortcut_trainer_2.py index 7a449fc3..d0e9309a 100644 --- a/owl_wms/trainers/shortcut_trainer_2.py +++ b/owl_wms/trainers/shortcut_trainer_2.py @@ -112,8 +112,11 @@ def get_ema_core(): else: return self.ema.ema_model.core - # No muon pls - self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) + # Set up optimizer and scheduler + if self.train_cfg.opt.lower() == "muon": + self.opt = init_muon(self.model, rank=self.rank,world_size=self.world_size,**self.train_cfg.opt_kwargs) + else: + self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) # Grad accum setup and scaler accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size @@ -154,8 +157,10 @@ def get_ema_core(): local_step += 1 if local_step % accum_steps == 0: # Updates - self.scaler.unscale_(self.opt) - torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + if self.train_cfg.opt.lower() != "muon": + self.scaler.unscale_(self.opt) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + self.scaler.step(self.opt) self.opt.zero_grad(set_to_none=True) From c9d5c0816091cddb1191fa628faf854a08e23491 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 05:24:24 +0000 Subject: [PATCH 56/71] fix it --- configs/shortcut_2.yml | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/configs/shortcut_2.yml b/configs/shortcut_2.yml index 15e34bc6..7bdc93e3 100644 --- a/configs/shortcut_2.yml +++ b/configs/shortcut_2.yml @@ -4,16 +4,16 @@ model: sample_size: 5 channels: 64 - n_layers: 5 - n_heads: 6 - d_model: 384 + n_layers: 13 + n_heads: 12 + d_model: 768 tokens_per_frame: 40 n_buttons: 11 n_mouse_axes: 2 cfg_prob: 0.1 - n_frames: 60 + n_frames: 30 causal: true @@ -21,7 +21,7 @@ train: trainer_id: shortcut_2 data_id: cod_s3 data_kwargs: - window_length: 60 + window_length: 30 bucket_name: cod-data-latent-360x640to5x8 include_keyframe: false @@ -52,8 +52,8 @@ train: sampler_id: shortcut_2 sampler_kwargs: - window_length: 60 - num_frames: 60 + window_length: 30 + num_frames: 120 only_return_generated: true n_samples: 8 From 36aa9e460ef1de661d89f1fd22fd1fcb617d45d2 Mon Sep 17 00:00:00 2001 From: Sami Date: Wed, 11 Jun 2025 01:50:25 -0400 Subject: [PATCH 57/71] self forcing cfg --- configs/self_forcing.yaml | 7 +- owl_wms/models/__init__.py | 4 +- owl_wms/sampling/shortcut_sampler.py | 292 ++++++++++++++++++++++++++- owl_wms/trainers/__init__.py | 8 +- owl_wms/trainers/self_forcing.py | 90 ++++++++- train.py | 3 + 6 files changed, 378 insertions(+), 26 deletions(-) diff --git a/configs/self_forcing.yaml b/configs/self_forcing.yaml index 606fc6ef..9b951031 100644 --- a/configs/self_forcing.yaml +++ b/configs/self_forcing.yaml @@ -55,14 +55,15 @@ train: n_samples: 8 - vae_id: 720pr3dc + vae_id: dcae vae_batch_size: 4 vae_scale: 0.35 vae_cfg_path: configs/owl_vaes/128x_cod_stage2.yml vae_ckpt_path: 720p_cod_vae_30m_35k_steps.pt # Self-forcing specific parameters - teacher_ckpt: checkpoints/bidirectional_teacher/best.pt # Pretrained bidirectional model + # teacher_ckpt: checkpoints/bidirectional_teacher/best.pt # Pretrained bidirectional model + teacher_ckpt: null # Pretrained bidirectional model loss_type: dmd # Options: dmd, sid, gan gradient_steps: 2 # Number of steps to backprop through rollout_steps: 5 # Total autoregressive rollout length @@ -71,6 +72,6 @@ train: cfg_scale: 1.3 # Classifier-free guidance scale wandb: - name: shahbuland + name: samibg project: video_models run_name: self_forcing_dmd \ No newline at end of file diff --git a/owl_wms/models/__init__.py b/owl_wms/models/__init__.py index 5f70facd..9c8ea80a 100644 --- a/owl_wms/models/__init__.py +++ b/owl_wms/models/__init__.py @@ -5,7 +5,9 @@ def get_model_cls(model_id): if model_id == "game_rft_shortcut": from .gamerft_shortcut import ShortcutGameRFT return ShortcutGameRFT - + if model_id == "causal_game_rft": + from .causal_gamerft import CausalGameRFT + return CausalGameRFT diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index 67be7315..a236f8be 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -1,20 +1,215 @@ -""" -Shortcut sampler, with cache! -""" - +import cv2 +import math +import pathlib import torch -from torch import nn -import torch.nn.functional as F - +from torch import Module from tqdm import tqdm +from typing import Optional -from ..utils import batch_permute_to_length from ..nn.kv_cache import KVCache +from ..utils import batch_permute_to_length +from ..models.gamerft_shortcut import ShortcutGameRFT + def zlerp(x, alpha): z = torch.randn_like(x) return x * (1. - alpha) + z * alpha +def load_mp4_as_tensor(mp4_path: pathlib.Path) -> torch.Tensor: + """Load MP4 as tensor in format [N, C=3, H, W] with values in [-1, 1]""" + video = cv2.VideoCapture(str(mp4_path)) + + if not video.isOpened(): + raise ValueError(f"Could not open video file: {mp4_path}") + + frames = [] + while True: + ret, frame = video.read() + if not ret: + break + + # Convert BGR to RGB + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # Convert to torch tensor and normalize to [-1, 1] + frame = torch.from_numpy(frame).float() / 127.5 - 1.0 + + # Rearrange from [H, W, C] to [C, H, W] + frame = frame.permute(2, 0, 1) + + frames.append(frame) + + video.release() + + if not frames: + raise ValueError(f"No frames found in video: {mp4_path}") + + # Stack to [N, C, H, W] + return torch.stack(frames) + + +class InferenceCachedShortcutSampler: + + ALPHA = 0.25 + + def __init__(self, + model: ShortcutGameRFT, + window_length = 60, + num_frames = 1, + only_return_generated = False, + vae_scale = 2.17, + decode_fn: Optional[Module] = None, + initial_history_pt_path: Optional[pathlib.Path] = None, + initial_history_mp4_path: Optional[pathlib.Path] = None, + encoder: Optional[Module] = None): + # -- + self.model: ShortcutGameRFT = model + self.window_length = window_length + self.num_frames = num_frames + + self.vae_scale = vae_scale + self.only_return_generated = only_return_generated + + # -- + self._cache_built = False + self.cache = KVCache(model.config) + self.decode_fn = decode_fn + self.initial_history_pt_path = initial_history_pt_path + self.initial_history_mp4_path = initial_history_mp4_path + self.encoder = encoder + + assert initial_history_pt_path is not None or initial_history_mp4_path is not None, \ + 'Either initial_history_pt_path or initial_history_mp4_path must be provided' + + if initial_history_mp4_path is not None: + assert encoder is not None, \ + 'Encoder must be provided if initial_history_mp4_path is provided' + + self.initial_history_bWchw = self.init_history(self.initial_history_pt_path, self.initial_history_mp4_path) + self.keyframe_b1chw = self.initial_history_bWchw[:,0] + + def init_history(self, + initial_history_pt_path: pathlib.Path | None, + initial_history_mp4_path: pathlib.Path | None) -> torch.Tensor: + + if initial_history_pt_path is not None: + history_wchw = torch.load(initial_history_pt_path) + else: + history_wrgb = load_mp4_as_tensor(initial_history_mp4_path).unsqueeze(0) # add batch dim + history_wchw = self.encoder(history_wrgb) + # NOTE This is so we avoid generating the history with a compiled model. + torch.save(history_wchw, initial_history_mp4_path.absolute().replace('.mp4', '.pt')) + + N = self.window_length + C = self.model.config.channels + H = W = int(math.sqrt(self.model.config.tokens_per_frame)) + + assert tuple(history_wchw.shape) == (1, N, C, H, W), \ + f'Initial history must have shape (B=1, {N=}, {C=}, {H=}, {W=}), ' \ + f'but got {tuple(history_wchw.shape)}' + + return history_wchw + + + def init_cache(self, + frames_bWchw, # [B, W, c, h, w] - NOTE history of frames + keyframe_b1chw, # [B, 1, c, h, w] - NOTE keyframe conditioning + mouse_bW2, # [B, W, 2] + button_bW11, # [B, W, 11] + ts_bW, # [B, W] + d_bW): # [B, W] + if self._cache_built: + print(f'WARNING: Cache already built but called `init_cache` again - ignoring.') + return + + B, N, *_ = frames_bWchw.shape + + self.cache.reset(B) ; self.cache.enable_cache_updates() + + # -- noise the history and fwd to kv cache + self.model.core.sample(x=zlerp(frames_bWchw, self.ALPHA), + y=keyframe_b1chw, + mouse=mouse_bW2, + btn=button_bW11, + cache=self.cache, + ts=ts_bW, d=d_bW) + + self.cache.disable_cache_updates() ; self._cache_built = True + print(f'Cache initialized for {B} x {N} frames - {[[i.shape for i in elt] + for elt in self.cache.cache]}') + return self.cache + + def __call__(self, + ctxt_frame_b1chw, # [B, 1, c, h, w] - NOTE Keyframe conditioning + mouse_b1_2, # [B, 1, 2] - NOTE mouse actions + button_b1_11, # [B, 1, 11] - NOTE button actions + ts_alpha_b1, # [B, 1] - NOTE overall denoising timestamp (e.g. 128) + d_alpha_b1, # [B, 1] - NOTE denoising step budget (e.g. 4) + ) -> torch.Tensor: # [B, 1, c, h, w] + # 1. ---- generate next frame ---- + self.cache.disable_cache_updates() + # 1.A) -- use the full context, including entire action history, to generate the next frame given cache. + frame = self.model.core.sample(None, ctxt_frame_b1chw, + mouse_b1_2, button_b1_11, + self.cache, ts=None, d=None) # NOTE simulating one-step sampling + # 2. ---- repopulate cache ---- + self.cache.enable_cache_updates() ; self.cache.truncate(1) + self.model.core.sample( x=zlerp(frame, self.ALPHA), # diffuse with noised frame to repopulate cache + y=ctxt_frame_b1chw, + mouse=mouse_b1_2, + btn=button_b1_11, + cache=self.cache, + ts=ts_alpha_b1, d=d_alpha_b1) + self.cache.disable_cache_updates() + return frame + + @torch.no_grad() + def generate_frames(self, + history_bWchw, # [B, W, c, h, w] - NOTE: MP4 from CoD initially, and after that it's just KV cache. + mouse_bT2, # [B, W+N, 2] - Actions taken by the user. + button_bT11, # [B, W+N, 11] - Actions taken by the user. + ) -> torch.Tensor: # [B, W+N, c, h, w] - either latent or rgb. + + if not self._cache_built: + print(f'WARNING: Cache not built, but called `generate_frames` - initializing cache.') + self.init_cache(history_bWchw, self.keyframe_b1chw, mouse_bT2, button_bT11) + + # If does not have batch-size, add it. This sampler is going to be used for single-user inference so batch-size is always 1. + # The caller might not specify the batch-size, so we have this here. + if history_bWchw.ndim == 4: + history_bWchw = history_bWchw.unsqueeze(1) + + history_bWchw = history_bWchw[:, -self.window_length:, ::] + + assert history_bWchw.shape[1] == self.window_length, \ + f'Window history must be at least {self.window_length} frames long, but got {history_bWchw.shape}' + + ts_alpha_bW = torch.ones_like(history_bWchw[:,:,0,0,0]) * self.ALPHA + d_alpha_bW = torch.ones_like(history_bWchw[:,:,0,0,0]) * round(1./self.ALPHA) + + ts_alpha_b1 = ts_alpha_bW[:,0].unsqueeze(1) + d_alpha_b1 = d_alpha_bW [:,0].unsqueeze(1) + + frames_latent = [] + for frame_idx in range(self.num_frames): + btn_atom = button_bT11[:, self.window_length+frame_idx].unsqueeze(1) + mouse_atom = mouse_bT2 [:, self.window_length+frame_idx].unsqueeze(1) + frame = self.__call__(ctxt_frame_b1chw=self.keyframe_b1chw, + mouse_b1_2=mouse_atom, button_b1_11=btn_atom, + ts_alpha_b1=ts_alpha_b1, d_alpha_b1=d_alpha_b1) + frames_latent += [frame] + + frames_latent = torch.cat(frames_latent, dim=1) + + if self.only_return_generated: frames_latent = frames_latent[:,-self.num_frames:] + + if self.decode_fn is not None: + frames_rgb = self.decode_fn(frames_latent * self.vae_scale) + return frames_rgb, mouse_bT2, button_bT11 + + return frames_latent, mouse_bT2, button_bT11 + + class CacheShortcutSampler: """ Shortcut CFG sampler builds cache with 4 step diffusion. @@ -25,6 +220,8 @@ class CacheShortcutSampler: :param only_return_generated: Whether to only return the generated frames """ def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale self.window_length = window_length self.num_frames = num_frames self.only_return_generated = only_return_generated @@ -34,7 +231,6 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale # dummy_batch is [b,n,c,h,w] # mouse is [b,n,2] # btn is [b,n,n_button] - num_frames = self.num_frames # output will be [b,n+self.num_frames,c,h,w] history = history[:,:self.window_length] @@ -53,9 +249,8 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale cache = KVCache(model.config) cache.reset(history.shape[0]) - - cache.enable_cache_updates() + cache.enable_cache_updates() _ = model.sample(noisy_history, keyframe, mouse, btn, cache, ts, d) cache.disable_cache_updates() @@ -68,12 +263,14 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale cache.disable_cache_updates() mouse = extended_mouse[:,self.window_length+frame_idx].unsqueeze(1) btn = extended_btn[:,self.window_length+frame_idx].unsqueeze(1) + # N+1 new_frame = model.sample(None, keyframe, mouse, btn, cache) # [b,1,c,h,w] new_frames.append(new_frame) # Add that frame to the cache cache.enable_cache_updates() new_frame_noisy = zlerp(new_frame, alpha) + # N+2, noisy(N+1) gets cached _ = model.sample(new_frame_noisy, keyframe, mouse, btn, cache, ts_single, d_single) new_frames = torch.cat(new_frames, dim = 1) @@ -164,3 +361,76 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale x = decode_fn(x) return x, extended_mouse, extended_btn + +class WindowShortcutSamplerNoKeyframe: + """ + Same as above but with no cache + + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + self.window_length = window_length + self.num_frames = num_frames + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, history, mouse, btn, decode_fn = None, scale = 1): + # history is [b,n,c,h,w] + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + history = history[:,:self.window_length] + new_frames = [] + alpha = 0.25 # This number is special for our sampler + + # Extended fake controls to use during sampling + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, self.num_frames + self.window_length) + + # Initialize window history + window_history = history.clone() + + for frame_idx in tqdm(range(self.num_frames)): + # Setup window history + x = window_history[:,-self.window_length:].clone() + + # Noise all but last frame to alpha + x[:,:-1] = zlerp(x[:,:-1], alpha) + # Last frame starts as random noise + x[:,-1] = torch.randn_like(x[:,-1]) + + # Setup timesteps - alpha for context, 1.0 for generated + ts = torch.ones_like(x[:,:,0,0,0]) + ts[:,:-1] = alpha + + # Setup diffusion steps - 4 for context, 1 for generated + d = torch.ones_like(x[:,:,0,0,0]) + d[:,:-1] = 4 + + # Get current controls + curr_mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] + curr_btn = extended_btn[:,frame_idx:frame_idx+self.window_length] + + # Generate new frame + pred = model.sample(x, keyframe, curr_mouse, curr_btn, None, ts, d) + new_frame = pred[:,-1:] # Take only the last frame + new_frames.append(new_frame) + + # Add new frame to window history + window_history = torch.cat([window_history, new_frame], dim=1) + + new_frames = torch.cat(new_frames, dim=1) + x = torch.cat([history, new_frames], dim=1) + + if self.only_return_generated: + x = x[:,-self.num_frames:] + extended_mouse = extended_mouse[:,-self.num_frames:] + extended_btn = extended_btn[:,-self.num_frames:] + + if decode_fn is not None: + x = x * scale + x = decode_fn(x) + + return x, extended_mouse, extended_btn \ No newline at end of file diff --git a/owl_wms/trainers/__init__.py b/owl_wms/trainers/__init__.py index 0a78fb14..983d9c43 100644 --- a/owl_wms/trainers/__init__.py +++ b/owl_wms/trainers/__init__.py @@ -8,4 +8,10 @@ def get_trainer_cls(trainer_id): return CausVidTrainer if trainer_id == "shortcut": from .shortcut_trainer import ShortcutTrainer - return ShortcutTrainer \ No newline at end of file + return ShortcutTrainer + if trainer_id == "self_forcing": # Add this case + from .self_forcing import SelfForcingTrainer # Add this import + return SelfForcingTrainer + + + diff --git a/owl_wms/trainers/self_forcing.py b/owl_wms/trainers/self_forcing.py index e3eb4c0f..0d9a17da 100644 --- a/owl_wms/trainers/self_forcing.py +++ b/owl_wms/trainers/self_forcing.py @@ -35,8 +35,6 @@ class SelfForcingTrainer(BaseTrainer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.model_cfg = self.config.model - model_id = self.model_cfg.model_id # Create student (causal) and teacher (non-causal) configs @@ -243,15 +241,87 @@ def compute_dmd_loss(self, generated, mouse, btn, gradient_mask): return dmd_loss def compute_sid_loss(self, generated, mouse, btn, gradient_mask): - """Compute SiD loss on generated sequence""" - # Similar to DMD but with SiD formulation - # Implementation details from SiD paper - raise NotImplementedError("SiD loss not yet implemented") + """ + Compute Score identity Distillation (SiD) loss on generated sequence + Based on "Score identity Distillation" paper + """ + s_real_fn = self.score_real.core + s_fake_fn = self.score_fake.module.core if self.world_size > 1 else self.score_fake.core + + with torch.no_grad(): + b, n, c, h, w = generated.shape + ts = torch.rand(b, n, device=generated.device).sigmoid() + z = torch.randn_like(generated) + ts_exp = ts[:, :, None, None, None] + lerpd = generated * (1. - ts_exp) + z * ts_exp + + # Compute real score with CFG + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + s_real_uncond = s_real_fn(lerpd, ts, null_mouse, null_btn) + s_real_cond = s_real_fn(lerpd, ts, mouse, btn) + s_real = s_real_uncond + self.cfg_scale * (s_real_cond - s_real_uncond) + + # Compute fake score + s_fake = s_fake_fn(lerpd, ts, mouse, btn) + + # SiD loss formulation + # L = (s_real - s_fake) * ((s_real - x) - α(s_real - s_fake)) + alpha = self.train_cfg.get('sid_alpha', 1.0) + + diff_score = s_real - s_fake + diff_real = s_real - generated + + sid_loss = diff_score * (diff_real - alpha * diff_score) + + # Normalize + with torch.no_grad(): + normalizer = torch.abs(diff_real).mean(dim=[2, 3, 4], keepdim=True) + sid_loss = sid_loss / (normalizer + 1e-6) + + # Apply gradient mask + if gradient_mask is not None: + sid_loss = sid_loss * gradient_mask[:, :, None, None, None] + + sid_loss = torch.nan_to_num(sid_loss).mean() + + return sid_loss def compute_gan_loss(self, generated, mouse, btn, gradient_mask, train_generator=True): - """Compute GAN loss on generated sequence""" - # Implementation for GAN-based loss - raise NotImplementedError("GAN loss not yet implemented") + """ + Simplified GAN loss using the score networks as discriminators + This avoids needing a separate discriminator implementation + """ + with torch.no_grad(): + b, n, c, h, w = generated.shape + # Use a fixed small noise level for discrimination + ts = torch.ones(b, n, device=generated.device) * 0.01 + noise = torch.randn_like(generated) * 0.01 + noisy_generated = generated + noise + + if train_generator: + # Generator loss: make fake score match real score + s_fake = self.score_fake.module.core(noisy_generated, ts, mouse, btn) if self.world_size > 1 else self.score_fake.core(noisy_generated, ts, mouse, btn) + s_real = self.score_real.core(noisy_generated, ts, mouse, btn) + + # L2 loss between scores + score_diff = (s_fake - s_real) ** 2 + + # Apply gradient mask + if gradient_mask is not None: + score_diff = score_diff * gradient_mask[:, :, None, None, None] + + gan_loss = score_diff.mean() + return gan_loss + else: + # Train fake score to distinguish real from fake + # This is similar to standard DMD critic training + with torch.no_grad(): + # For real data, we'd need ground truth + # For now, return a placeholder + return torch.tensor(0.0, device=generated.device) + def train(self): torch.cuda.set_device(self.local_rank) @@ -436,4 +506,4 @@ def optimizer_step(loss, model, scaler, optimizer): if self.rank == 0: self.save() - self.barrier() \ No newline at end of file + self.barrier() diff --git a/train.py b/train.py index 64576921..9126bbd5 100644 --- a/train.py +++ b/train.py @@ -6,6 +6,9 @@ from owl_wms.utils.ddp import cleanup, setup if __name__ == "__main__": + import sys + sys.argv[1:] = ["--config_path", "configs/self_forcing.yaml"] + parser = argparse.ArgumentParser() parser.add_argument("--config_path", type=str, help="Path to config YAML file") From 2a6410de61082080cde755dd42360fff339ef993 Mon Sep 17 00:00:00 2001 From: Sami Date: Wed, 11 Jun 2025 02:08:26 -0400 Subject: [PATCH 58/71] stupid inference window shortcut sampler without keyframe --- owl_wms/sampling/shortcut_sampler.py | 189 ++++++++++++++++++++------- 1 file changed, 140 insertions(+), 49 deletions(-) diff --git a/owl_wms/sampling/shortcut_sampler.py b/owl_wms/sampling/shortcut_sampler.py index a236f8be..43690cca 100644 --- a/owl_wms/sampling/shortcut_sampler.py +++ b/owl_wms/sampling/shortcut_sampler.py @@ -362,75 +362,166 @@ def __call__(self, model, history, keyframe, mouse, btn, decode_fn = None, scale return x, extended_mouse, extended_btn -class WindowShortcutSamplerNoKeyframe: - """ - Same as above but with no cache +class InferenceWindowShortcutSamplerNoKeyframe: + """ + Window-based shortcut sampler without keyframe conditioning or KV cache. + Generates frames using sliding window approach with diffusion forcing. + + :param model: The shortcut diffusion model :param window_length: Number of frames to use for each frame generation step - :param num_frames: Number of new frames to sample + :param num_frames: Number of new frames to sample per generate_frames call :param only_return_generated: Whether to only return the generated frames + :param vae_scale: Scale factor for VAE decoding + :param decode_fn: Optional decoder function + :param initial_history_pt_path: Path to pre-encoded initial history tensor + :param initial_history_mp4_path: Path to MP4 file for initial history + :param encoder: Encoder module (required if using MP4 path) """ - def __init__(self, window_length = 60, num_frames = 60, only_return_generated = False): + + ALPHA = 0.25 # Noise level for context frames (step 3 of 4-step diffusion) + + def __init__(self, + model, + window_length = 60, + num_frames = 1, + only_return_generated = False, + vae_scale = 2.17, # TODO Shab trained a new VAE so this needs to be updated. + decode_fn: Optional[Module] = None, + initial_history_pt_path: Optional[pathlib.Path] = None, + initial_history_mp4_path: Optional[pathlib.Path] = None, + encoder: Optional[Module] = None): + + self.model = model self.window_length = window_length self.num_frames = num_frames + self.vae_scale = vae_scale self.only_return_generated = only_return_generated + self.decode_fn = decode_fn + + self.initial_history_pt_path = initial_history_pt_path + self.initial_history_mp4_path = initial_history_mp4_path + self.encoder = encoder + + assert initial_history_pt_path is not None or initial_history_mp4_path is not None, \ + 'Either initial_history_pt_path or initial_history_mp4_path must be provided' + + if initial_history_mp4_path is not None: + assert encoder is not None, \ + 'Encoder must be provided if initial_history_mp4_path is provided' + + self.initial_history_bWchw = self.init_history(self.initial_history_pt_path, self.initial_history_mp4_path) + + def init_history(self, + initial_history_pt_path: pathlib.Path | None, + initial_history_mp4_path: pathlib.Path | None) -> torch.Tensor: + """Initialize history from either .pt file or MP4 file""" + + if initial_history_pt_path is not None: + history_wchw = torch.load(initial_history_pt_path) + else: + history_wrgb = load_mp4_as_tensor(initial_history_mp4_path).unsqueeze(0) # add batch dim + history_wchw = self.encoder(history_wrgb) + # Save encoded version to avoid re-encoding + torch.save(history_wchw, initial_history_mp4_path.absolute().replace('.mp4', '.pt')) + + N = self.window_length + C = self.model.config.channels if hasattr(self.model, 'config') else history_wchw.shape[2] + H = W = int(math.sqrt(self.model.config.tokens_per_frame)) if hasattr(self.model, 'config') else history_wchw.shape[3] + + assert tuple(history_wchw.shape) == (1, N, C, H, W), \ + f'Initial history must have shape (B=1, {N=}, {C=}, {H=}, {W=}), ' \ + f'but got {tuple(history_wchw.shape)}' + + return history_wchw + + def __call__(self, + window_history_bWchw, # [B, W, c, h, w] - Current window of frames + mouse_bW2, # [B, W, 2] - Mouse actions for window + button_bW11, # [B, W, 11] - Button actions for window + ) -> torch.Tensor: # [B, 1, c, h, w] - Generated frame + """Generate a single frame given current window history""" + + # Setup window for generation + x = window_history_bWchw[:, -self.window_length:].clone() + + # Noise all but last frame to alpha level (diffusion forcing) + x[:, :-1] = zlerp(x[:, :-1], self.ALPHA) + # Last frame starts as random noise + x[:, -1] = torch.randn_like(x[:, -1]) + + # Setup timesteps - ALPHA for context frames, 1.0 for generated frame + ts = torch.ones_like(x[:, :, 0, 0, 0]) + ts[:, :-1] = self.ALPHA + + # Setup diffusion steps - 4 for context frames, 1 for generated frame + d = torch.ones_like(x[:, :, 0, 0, 0]) + d[:, :-1] = 4 # Context frames use 4-step budget + + # Generate new frame using window + pred = self.model.sample(x, mouse_bW2, button_bW11, None, ts, d) + new_frame = pred[:, -1:] # Take only the last (generated) frame + + return new_frame @torch.no_grad() - def __call__(self, model, history, mouse, btn, decode_fn = None, scale = 1): - # history is [b,n,c,h,w] - # mouse is [b,n,2] - # btn is [b,n,n_button] + def generate_frames(self, + history_bWchw, # [B, W, c, h, w] - Initial history + mouse_bT2, # [B, W+N, 2] - Mouse actions for entire sequence + button_bT11, # [B, W+N, 11] - Button actions for entire sequence + ) -> torch.Tensor: # [B, W+N, c, h, w] - Generated sequence + """Generate multiple frames using sliding window approach""" + + # Handle batch dimension + if history_bWchw.ndim == 4: + history_bWchw = history_bWchw.unsqueeze(0) - # output will be [b,n+self.num_frames,c,h,w] - history = history[:,:self.window_length] - new_frames = [] - alpha = 0.25 # This number is special for our sampler + history_bWchw = history_bWchw[:, -self.window_length:] + + assert history_bWchw.shape[1] == self.window_length, \ + f'History must be exactly {self.window_length} frames long, but got {history_bWchw.shape[1]}' - # Extended fake controls to use during sampling - extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, self.num_frames + self.window_length) + # Extended controls for generation + extended_mouse, extended_btn = batch_permute_to_length( + mouse_bT2[:, :self.window_length], + button_bT11[:, :self.window_length], + self.num_frames + self.window_length + ) # Initialize window history - window_history = history.clone() + window_history = history_bWchw.clone() + frames_latent = [] - for frame_idx in tqdm(range(self.num_frames)): - # Setup window history - x = window_history[:,-self.window_length:].clone() + for frame_idx in range(self.num_frames): + # Get current window controls + curr_mouse = extended_mouse[:, frame_idx:frame_idx + self.window_length] + curr_btn = extended_btn[:, frame_idx:frame_idx + self.window_length] - # Noise all but last frame to alpha - x[:,:-1] = zlerp(x[:,:-1], alpha) - # Last frame starts as random noise - x[:,-1] = torch.randn_like(x[:,-1]) - - # Setup timesteps - alpha for context, 1.0 for generated - ts = torch.ones_like(x[:,:,0,0,0]) - ts[:,:-1] = alpha + # Generate single frame + new_frame = self.__call__( + window_history_bWchw=window_history, + mouse_bW2=curr_mouse, + button_bW11=curr_btn + ) - # Setup diffusion steps - 4 for context, 1 for generated - d = torch.ones_like(x[:,:,0,0,0]) - d[:,:-1] = 4 - - # Get current controls - curr_mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] - curr_btn = extended_btn[:,frame_idx:frame_idx+self.window_length] - - # Generate new frame - pred = model.sample(x, keyframe, curr_mouse, curr_btn, None, ts, d) - new_frame = pred[:,-1:] # Take only the last frame - new_frames.append(new_frame) + frames_latent.append(new_frame) - # Add new frame to window history + # Add new frame to window history for next iteration window_history = torch.cat([window_history, new_frame], dim=1) - new_frames = torch.cat(new_frames, dim=1) - x = torch.cat([history, new_frames], dim=1) + # Combine all generated frames + frames_latent = torch.cat(frames_latent, dim=1) + + # Combine with original history + full_sequence = torch.cat([history_bWchw, frames_latent], dim=1) if self.only_return_generated: - x = x[:,-self.num_frames:] - extended_mouse = extended_mouse[:,-self.num_frames:] - extended_btn = extended_btn[:,-self.num_frames:] + full_sequence = full_sequence[:, -self.num_frames:] + extended_mouse = extended_mouse[:, -self.num_frames:] + extended_btn = extended_btn[:, -self.num_frames:] - if decode_fn is not None: - x = x * scale - x = decode_fn(x) + if self.decode_fn is not None: + frames_rgb = self.decode_fn(full_sequence * self.vae_scale) + return frames_rgb, extended_mouse, extended_btn - return x, extended_mouse, extended_btn \ No newline at end of file + return full_sequence, extended_mouse, extended_btn \ No newline at end of file From c368f5544ea254a0522e4063a898d7acf830069a Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 14:42:10 +0000 Subject: [PATCH 59/71] update shortcut config, fix rope, add reqs --- configs/shortcut_2.yml | 28 ++++++++++++++-------------- owl_wms/nn/attn.py | 9 +++++++-- requirements.txt | 1 + 3 files changed, 22 insertions(+), 16 deletions(-) diff --git a/configs/shortcut_2.yml b/configs/shortcut_2.yml index 7bdc93e3..b6a79a58 100644 --- a/configs/shortcut_2.yml +++ b/configs/shortcut_2.yml @@ -1,32 +1,32 @@ # Config for a simple 256 -> 16 autoencoder model: model_id: shortcut_2 - sample_size: 5 - channels: 64 + sample_size: 4 + channels: 128 n_layers: 13 - n_heads: 12 - d_model: 768 + n_heads: 16 + d_model: 1024 - tokens_per_frame: 40 + tokens_per_frame: 16 n_buttons: 11 n_mouse_axes: 2 cfg_prob: 0.1 n_frames: 30 - causal: true + causal: false train: trainer_id: shortcut_2 data_id: cod_s3 data_kwargs: window_length: 30 - bucket_name: cod-data-latent-360x640to5x8 + bucket_name: cod-data-latent-360x640to4x4 include_keyframe: false - target_batch_size: 16 - batch_size: 16 + target_batch_size: 256 + batch_size: 32 epochs: 200 @@ -35,7 +35,7 @@ train: lr: 1.0e-3 momentum: 0.95 adamw_lr: 1.0e-4 - adamw_wd: 1.0e-4 + adamw_wd: 0.1 adamw_eps: 1.0e-15 adamw_betas: [0.9, 0.95] adamw_keys: [ @@ -60,11 +60,11 @@ train: vae_id: 720pr3dc vae_batch_size: 4 - vae_scale: 0.35 - vae_cfg_path: configs/owl_vaes/128x_cod_stage2.yml - vae_ckpt_path: 720p_cod_vae_30m_35k_steps.pt + vae_scale: 0.13 + vae_cfg_path: configs/owl_vaes/cod_128x.yml + vae_ckpt_path: checkpoints/owl_vaes/cod_128x_30k_ema.pt wandb: name: shahbuland project: video_models - run_name: v2 + run_name: bidir_shortcut diff --git a/owl_wms/nn/attn.py b/owl_wms/nn/attn.py index c1077a84..a57106ff 100644 --- a/owl_wms/nn/attn.py +++ b/owl_wms/nn/attn.py @@ -8,6 +8,7 @@ import einops as eo from .modulation import AdaLN, Gate +from .rope import FlatVideoRoPE torch.backends.cuda.enable_flash_sdp(enabled = True) @@ -40,6 +41,8 @@ def __init__(self, config : 'TransformerConfig'): self.qk_norm = QKNorm(config.d_model // config.n_heads) self.layer_ind = None + self.rope = FlatVideoRoPE(config) + self.tokens_per_frame = config.tokens_per_frame self.causal = config.causal @@ -62,7 +65,8 @@ def forward(self, x, kv_cache = None): new_k = torch.cat([old_k, k], dim = 2).contiguous() new_v = torch.cat([old_v, v], dim = 2).contiguous() - + + q,new_k = self.rope(q,new_k) if kv_cache.should_update: kv_cache.update(new_k, new_v, self.layer_ind) @@ -70,6 +74,7 @@ def forward(self, x, kv_cache = None): x = F.scaled_dot_product_attention(q, new_k, new_v, attn_mask = mask) x = x[:,:,-q.shape[2]:] # Skip cached outputs (not relevant now) else: + q,k = self.rope(q,k) x = F.scaled_dot_product_attention(q,k,v, attn_mask = mask) x = eo.rearrange(x, 'b h n d -> b n (h d)') @@ -243,4 +248,4 @@ def test_kv_cache(): print("Cache test complete") if __name__ == "__main__": - test_attn_mask() \ No newline at end of file + test_attn_mask() diff --git a/requirements.txt b/requirements.txt index 909273a1..e259d76c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ diffusers accelerate boto3 python-dotenv +wandb[media] From 671556476ee79358a48f7b8dfba2a6cd41bc60d5 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 16:10:55 +0000 Subject: [PATCH 60/71] audio stuffs --- owl_wms/configs.py | 2 + owl_wms/models/gamerft_shortcut_audio.py | 220 +++++++++++++++++++++ owl_wms/trainers/shortcut_trainer_audio.py | 217 ++++++++++++++++++++ owl_wms/utils/logging.py | 32 ++- owl_wms/utils/owl_vae_bridge.py | 26 ++- 5 files changed, 482 insertions(+), 15 deletions(-) create mode 100644 owl_wms/models/gamerft_shortcut_audio.py create mode 100644 owl_wms/trainers/shortcut_trainer_audio.py diff --git a/owl_wms/configs.py b/owl_wms/configs.py index 5648ca52..eebc1366 100644 --- a/owl_wms/configs.py +++ b/owl_wms/configs.py @@ -13,11 +13,13 @@ class TransformerConfig: patch_size : int = 1 channels : int = 128 + audio_channels : int = 64 sample_size : int = 16 cfg_prob : float = 0.1 n_buttons : int = 8 tokens_per_frame : int = 16 + audio_tokens : int = 0 n_frames : int = 120 causal : bool = False diff --git a/owl_wms/models/gamerft_shortcut_audio.py b/owl_wms/models/gamerft_shortcut_audio.py new file mode 100644 index 00000000..33b47a4f --- /dev/null +++ b/owl_wms/models/gamerft_shortcut_audio.py @@ -0,0 +1,220 @@ +""" +Shortcut simple with audio +""" + +import torch +from torch import nn +import torch.nn.functional as F + +import einops as eo + +from ..nn.embeddings import ( + TimestepEmbedding, + StepEmbedding, + ControlEmbedding, + LearnedPosEnc +) +from ..nn.attn import UViT, FinalLayer +from ..nn.mmattn import MMUViT +from ..utils import freeze + +class ShortcutGameRFTCore(nn.Module): + def __init__(self, config): + super().__init__() + + self.transformer = UViT(config) + self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) + + self.step_embed = StepEmbedding(config.d_model) + self.t_embed = TimestepEmbedding(config.d_model) + + self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) + self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) + + self.audio_proj = nn.Linear(config.audio_channels, config.d_model, bias = False) + self.audio_proj_out = nn.Linear(config.d_model, config.audio_channels, bias = False) + + self.config = config + + def sample(self, x, audio, mouse, btn, kv_cache = None, t = None, d = None): + """ + This is a function that largely abstracts + away most things for the specific case where + you are only generating the one next token + + The return is one step sample always + """ + + b,n,c,h,w = x.shape + if t is None: + t = torch.ones_like(x[:,:,0,0,0]) + if d is None: + d = torch.ones_like(x[:,:,0,0,0]) + + pred_x, pred_audio = self.forward(x, audio, t, mouse, btn, d, kv_cache) + return x - pred_x, audio - pred_audio + + def forward(self, x, audio, t, mouse, btn, d, kv_cache = None): + # x is [b,n,c,h,w] + # a is [b,c,n] + # t is [b,n] + # d is [b,n] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + + ctrl_cond = self.control_embed(mouse, btn) + t_cond = self.t_embed(t) + d_cond = self.step_embed(d) + + cond = ctrl_cond + t_cond + d_cond # [b,n,d] + + audio = self.audio_proj(audio.transpose(-1,-2))[:,:,None] # -> [b,n,1,d] + + b,n,c,h,w = x.shape + x = eo.rearrange(x, 'b n c h w -> b (n h w) c') + x = self.proj_in(x) + x = eo.rearrange(x, 'b (n f) d -> b n f d', n = n) + x = torch.cat([x, audio], dim=-2) # [b,n,f,d] + x = eo.rearrange(x, 'b n f d -> b (n f) d') + + x = self.transformer(x, cond, kv_cache) + + x = eo.rearrange(x, 'b (n f) c -> b n f w', n=n) + audio = x[:,:,-1] # [b,n,d] + x = x[:,:,:-1] # [b,n,f,d] + + x = eo.rearrange(x, 'b n f d -> b (n f) d') + x = self.proj_out(x, cond) # -> [b,n*hw,c] + x = eo.rearrange(x, 'b (n h w) c -> b n c h w', n=n,h=h,w=w) + audio = self.audio_proj_out(audio).transpose(-1,-2) # [b,d,n] + + return x, audio + +def sample_discrete_timesteps(steps, eps = 1.0e-6): + # steps is Tensor([1,4,2,64,16]) as an example + b,n = steps.shape + + ts_list = [] + ts = torch.rand(b, n, device=steps.device, dtype=steps.dtype) * (steps - eps) + ts = ts.clamp(eps).ceil() / steps + """ + Example, if d was all 2, ts would be [0,2] + so do clamp, then ceil will be 1 or 2 (0, 2] + then do t / 2 and get 0.5 or 1.0, our desired timesteps + """ + return ts + +def sample_steps(b, n, device, dtype, min_val = 0): + valid = torch.tensor([2**i for i in range(min_val, 8)]) # [1,2,...,128] + inds = torch.randint(low=0,high=len(valid), size = (b,n)) + steps = valid[inds].to(device=device,dtype=dtype) + return steps + +#@torch.compile() +@torch.no_grad() +def get_sc_targets(ema, x, audio, mouse, btn, cfg_scale): + steps_slow = sample_steps(x.shape[0], x.shape[1], x.device, x.dtype, min_val = 1) + steps_fast = steps_slow / 2 + + dt_slow = 1./steps_slow + dt_fast = 1./steps_fast + + def expand(t): + #b,c,h,w = x.shape + #t = eo.repeat(t,'b -> b c h w',c=c,h=h,w=w) + #return t + return t[:,:,None,None,None] + + ts = sample_discrete_timesteps(steps_fast) + cfg_mask = torch.isclose(steps_slow, torch.ones_like(steps_slow)*128) + cfg_mask = expand(cfg_mask) # -> [b,n,1,1,1] + + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + pred_1_x_uncond, pred_1_a_uncond = ema(x, audio, ts, null_mouse, null_btn, steps_slow) + pred_1_x_cond, pred_1_a_cond = ema(x, audio, ts, mouse, btn, steps_slow) + pred_1_x_cfg = pred_1_x_uncond + cfg_scale * (pred_1_x_cond - pred_1_x_uncond) + pred_1_a_cfg = pred_1_a_uncond + cfg_scale * (pred_1_a_cond - pred_1_a_uncond) + pred_1_x = torch.where(cfg_mask, pred_1_x_cfg, pred_1_x_cond) + pred_1_a = pred_1_a_cfg if cfg_mask.any() else pred_1_a_cond + + x_new = x - pred_1_x * expand(dt_slow) + audio_new = audio - pred_1_a * dt_slow[:,None,:] + ts_new = ts - dt_slow + + pred_2_x_uncond, pred_2_a_uncond = ema(x_new, audio_new, ts_new, null_mouse, null_btn, steps_slow) + pred_2_x_cond, pred_2_a_cond = ema(x_new, audio_new, ts_new, mouse, btn, steps_slow) + pred_2_x_cfg = pred_2_x_uncond + cfg_scale * (pred_2_x_cond - pred_2_x_uncond) + pred_2_a_cfg = pred_2_a_uncond + cfg_scale * (pred_2_a_cond - pred_2_a_uncond) + pred_2_x = torch.where(cfg_mask, pred_2_x_cfg, pred_2_x_cond) + pred_2_a = pred_2_a_cfg if cfg_mask.any() else pred_2_a_cond + + pred_x = 0.5 * (pred_1_x + pred_2_x) + pred_a = 0.5 * (pred_1_a + pred_2_a) + return (pred_x, pred_a), steps_fast, ts + +class ShortcutGameRFT(nn.Module): + def __init__(self, config): + super().__init__() + + self.core = ShortcutGameRFTCore(config) + self.cfg_prob = config.cfg_prob + + self.sc_frac = 0.25 + self.sc_max_steps = 128 + self.cfg_scale = 1.3 + + self.config = config + + def get_sc_loss(self, x, audio, mouse, btn, ema): + (target_x, target_a), steps, ts = get_sc_targets(ema, x, audio, mouse, btn, self.cfg_scale) + pred_x, pred_a = self.core(x, audio, ts, mouse, btn, steps) + sc_loss_x = F.mse_loss(pred_x, target_x) + sc_loss_a = F.mse_loss(pred_a, target_a) + return sc_loss_x + sc_loss_a + + def forward(self, x, audio, mouse, btn, ema): + # x is [b,n,c,h,w] + # audio is [b,c,n] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + with torch.no_grad(): + _,n,c,h,w = x.shape + + # Split batches between consistency/rf + b = int(len(x) * (1 - self.sc_frac)) + x,x_sc = x[:b], x[b:] + audio,audio_sc = audio[:b], audio[b:] + mouse,mouse_sc = mouse[:b], mouse[b:] + btn,btn_sc = btn[:b], btn[b:] + + # Apply classifier-free guidance dropout + if self.cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + # Where mask is True, replace with zeros + mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) + btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) + + d = torch.ones_like(x[:,:,0,0,0])*self.sc_max_steps + ts = sample_discrete_timesteps(d) + ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() + + ts_exp = eo.repeat(ts, 'b n -> b n 1 1 1') + z_x = torch.randn_like(x) + z_a = torch.randn_like(audio) + + lerpd_x = x * (1. - ts_exp) + z_x * ts_exp + lerpd_a = audio * (1. - ts[:,None,:]) + z_a * ts[:,None,:] + target_x = z_x - x + target_a = z_a - audio + + pred_x, pred_a = self.core(lerpd_x, lerpd_a, ts, mouse, btn, d) + diff_loss_x = F.mse_loss(pred_x, target_x) + diff_loss_a = F.mse_loss(pred_a, target_a) + sc_loss = self.get_sc_loss(x_sc, audio_sc, mouse_sc, btn_sc, ema) + + return diff_loss_x + diff_loss_a, sc_loss \ No newline at end of file diff --git a/owl_wms/trainers/shortcut_trainer_audio.py b/owl_wms/trainers/shortcut_trainer_audio.py new file mode 100644 index 00000000..5f1d7263 --- /dev/null +++ b/owl_wms/trainers/shortcut_trainer_audio.py @@ -0,0 +1,217 @@ +import torch +from ema_pytorch import EMA +import wandb +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist +import einops as eo + +from .base import BaseTrainer + +from ..utils import freeze, Timer, find_unused_params +from ..schedulers import get_scheduler_cls +from ..models import get_model_cls +from ..sampling import get_sampler_cls +from ..data import get_loader +from ..utils.logging import LogHelper, to_wandb +from ..muon import init_muon +from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn, make_batched_audio_decode_fn +from ..models.gamerft_shortcut import get_sc_targets + +class ShortcutTrainer(BaseTrainer): + """ + Trainer for rectified flow transformer with shortcut + + :param train_cfg: Configuration for training + :param logging_cfg: Configuration for logging + :param model_cfg: Configuration for model + :param global_rank: Rank across all devices. + :param local_rank: Rank for current device on this process. + :param world_size: Overall number of devices + """ + def __init__(self,*args,**kwargs): + super().__init__(*args,**kwargs) + + model_id = self.model_cfg.model_id + self.model = get_model_cls(model_id)(self.model_cfg) + + # Print model size + if self.rank == 0: + n_params = sum(p.numel() for p in self.model.parameters()) + print(f"Model has {n_params:,} parameters") + + self.ema = None + self.opt = None + self.scheduler = None + self.scaler = None + + self.total_step_counter = 0 + self.decoder = get_decoder_only( + self.train_cfg.image_vae_id, + self.train_cfg.image_vae_cfg_path, + self.train_cfg.image_vae_ckpt_path + ) + + self.audio_decoder = get_decoder_only( + self.train_cfg.audio_vae_id, + self.train_cfg.audio_vae_cfg_path, + self.train_cfg.audio_vae_ckpt_path + ) + + freeze(self.decoder) + freeze(self.audio_decoder) + + def save(self): + save_dict = { + 'model' : self.model.state_dict(), + 'ema' : self.ema.state_dict(), + 'opt' : self.opt.state_dict(), + 'scaler' : self.scaler.state_dict(), + 'steps': self.total_step_counter + } + if self.scheduler is not None: + save_dict['scheduler'] = self.scheduler.state_dict() + super().save(save_dict) + + def load(self): + has_ckpt = False + try: + if self.train_cfg.resume_ckpt is not None: + save_dict = super().load(self.train_cfg.resume_ckpt) + has_ckpt = True + except: + print("Error loading checkpoint") + + if not has_ckpt: + return + + + self.model.load_state_dict(save_dict['model']) + self.ema.load_state_dict(save_dict['ema']) + self.opt.load_state_dict(save_dict['opt']) + if self.scheduler is not None and 'scheduler' in save_dict: + self.scheduler.load_state_dict(save_dict['scheduler']) + self.scaler.load_state_dict(save_dict['scaler']) + self.total_step_counter = save_dict['steps'] + + def train(self): + torch.cuda.set_device(self.local_rank) + + # Prepare model and ema + self.model = self.model.cuda().train() + if self.world_size > 1: + self.model = DDP(self.model, device_ids=[self.local_rank]) + + self.decoder = self.decoder.cuda().eval().bfloat16() + self.audio_decoder = self.audio_decoder.cuda().eval().bfloat16() + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + audio_decode_fn = make_batched_audio_decode_fn(self.audio_decoder, self.train_cfg.vae_batch_size) + + self.ema = EMA( + self.model, + beta = 0.999, + update_after_step = 0, + update_every = 1 + ) + #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) + + def get_ema_core(): + if self.world_size > 1: + return self.ema.ema_model.module.core + else: + return self.ema.ema_model.core + + # Set up optimizer and scheduler + if self.train_cfg.opt.lower() == "muon": + self.opt = init_muon(self.model, rank=self.rank,world_size=self.world_size,**self.train_cfg.opt_kwargs) + else: + self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) + + # Grad accum setup and scaler + accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size + accum_steps = max(1, accum_steps) + self.scaler = torch.amp.GradScaler() + ctx = torch.amp.autocast('cuda',torch.bfloat16) + + self.load() + + # Timer reset + timer = Timer() + timer.reset() + metrics = LogHelper() + if self.rank == 0: + wandb.watch(self.get_module(), log = 'all') + + # Dataset setup + loader = get_loader(self.train_cfg.data_id, self.train_cfg.batch_size, **self.train_cfg.data_kwargs) + sampler = get_sampler_cls(self.train_cfg.sampler_id)(**self.train_cfg.sampler_kwargs) + + local_step = 0 + for _ in range(self.train_cfg.epochs): + for batch_vid, batch_audio, batch_mouse, batch_btn in loader: + batch_vid = batch_vid.cuda().bfloat16() / self.train_cfg.vae_scale + batch_audio = batch_audio.cuda().bfloat16() + batch_mouse = batch_mouse.cuda().bfloat16() + batch_btn = batch_btn.cuda().bfloat16() + + with ctx: + diff_loss, sc_loss = self.model(batch_vid, batch_audio, batch_mouse, batch_btn, get_ema_core()) + loss = diff_loss + sc_loss + loss = loss / accum_steps + + self.scaler.scale(loss).backward() + + metrics.log('diffusion_loss', diff_loss) + metrics.log('shortcut_loss', sc_loss) + + local_step += 1 + if local_step % accum_steps == 0: + # Updates + if self.train_cfg.opt.lower() != "muon": + self.scaler.unscale_(self.opt) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + + + self.scaler.step(self.opt) + self.opt.zero_grad(set_to_none=True) + + self.scaler.update() + + if self.scheduler is not None: + self.scheduler.step() + self.ema.update() + + # Do logging + with torch.no_grad(): + wandb_dict = metrics.pop() + wandb_dict['time'] = timer.hit() + wandb_dict['lr'] = self.opt.param_groups[0]['lr'] + timer.reset() + + # Sampling commented out for now + if self.total_step_counter % self.train_cfg.sample_interval == 0: + with ctx, torch.no_grad(): + n_samples = self.train_cfg.n_samples + samples, sample_audio, sample_mouse, sample_button = sampler( + get_ema_core(), + batch_vid[:n_samples], + batch_audio[:n_samples], + batch_mouse[:n_samples], + batch_btn[:n_samples], + image_decode_fn = decode_fn, + audio_decode_fn = audio_decode_fn, + image_vae_scale=self.train_cfg.image_vae_scale, + audio_vae_scale=self.train_cfg.audio_vae_scale + ) # -> [b,n,c,h,w] + if self.rank == 0: wandb_dict['samples'] = to_wandb_av(samples, sample_audio, sample_mouse, sample_button) + + + if self.rank == 0: + wandb.log(wandb_dict) + + self.total_step_counter += 1 + if self.total_step_counter % self.train_cfg.save_interval == 0: + if self.rank == 0: + self.save() + + self.barrier() diff --git a/owl_wms/utils/logging.py b/owl_wms/utils/logging.py index e2ef70c1..9f2bca7e 100644 --- a/owl_wms/utils/logging.py +++ b/owl_wms/utils/logging.py @@ -74,4 +74,34 @@ def to_wandb(x, batch_mouse, batch_btn, gather = False, max_samples = 8): x = eo.rearrange(x, '(r c) n d h w -> n d (r h) (c w)', r = 2, c = 4) return wandb.Video(x, format='gif',fps=60) - \ No newline at end of file + +@torch.no_grad() +def to_wandb_av(x, audio, batch_mouse, batch_btn, gather = False, max_samples = 8): + # x is [b,n,c,h,w] + x = x.clamp(-1, 1) + x = x[:max_samples] + audio = audio[:max_samples] + + if dist.is_initialized() and gather: + gathered_x = [None for _ in range(dist.get_world_size())] + gathered_audio = [None for _ in range(dist.get_world_size())] + dist.all_gather(gathered_x, x) + dist.all_gather(gathered_audio, audio) + x = torch.cat(gathered_x, dim=0) + audio = torch.cat(gathered_audio, dim=0) + + # Get labels on frames + x = draw_frames(x, batch_mouse, batch_btn) # -> [b,n,c,h,w] [0,255] uint8 np + + # Convert audio to numpy float32 [-1,1] + audio = audio.cpu().float().numpy() + + # Create list of video/audio pairs + samples = [] + for i in range(len(x)): + video = wandb.Video(x[i], format='gif', fps=60) + # Sample rate assumed to be 48kHz based on common audio standards + audio_sample = wandb.Audio(audio[i].T, sample_rate=48000) + samples.append((video, audio_sample)) + + return samples \ No newline at end of file diff --git a/owl_wms/utils/owl_vae_bridge.py b/owl_wms/utils/owl_vae_bridge.py index 331a8582..c306f18f 100644 --- a/owl_wms/utils/owl_vae_bridge.py +++ b/owl_wms/utils/owl_vae_bridge.py @@ -25,7 +25,7 @@ def get_decoder_only(vae_id, cfg_path, ckpt_path): model = AutoencoderDC.from_pretrained(model_id).bfloat16().cuda().eval() del model.encoder return model.decoder - elif vae_id == "720pr3dc": + else: cfg = Config.from_yaml(cfg_path).model model = get_model_cls(cfg.model_id)(cfg) model.load_state_dict(torch.load(ckpt_path, map_location='cpu',weights_only=False)) @@ -35,18 +35,18 @@ def get_decoder_only(vae_id, cfg_path, ckpt_path): return model @torch.no_grad() -def _make_batched_decode_fn(decoder, batch_size = 8): +def make_batched_decode_fn(decoder, batch_size = 8): def decode(x): - # x is [b,n,m,d] - b,n,m,d = x.shape - x = x.view(b*n,m,d).contiguous() + # x is [b,n,c,h,w] + b,n,c,h,w = x.shape + x = x.view(b*n,c,h,w).contiguous() batches = x.split(batch_size) batch_out = [] for batch in batches: batch_out.append(decoder(batch).bfloat16()) - x = torch.cat(batch_out) # [b*n,3,256,256] + x = torch.cat(batch_out) # [b*n,c,h,w] _,c,h,w = x.shape x = x.view(b,n,c,h,w).contiguous() @@ -54,20 +54,18 @@ def decode(x): return decode @torch.no_grad() -def make_batched_decode_fn(decoder, batch_size = 8): +def make_batched_audio_decode_fn(decoder, batch_size = 8): def decode(x): - # x is [b,n,c,h,w] - b,n,c,h,w = x.shape - x = x.view(b*n,c,h,w).contiguous() + # x is [b,c,n] audio samples + b,c,n = x.shape - batches = x.split(batch_size) + batches = x.contiguous().split(batch_size) batch_out = [] for batch in batches: batch_out.append(decoder(batch).bfloat16()) - x = torch.cat(batch_out) # [b*n,c,h,w] - _,c,h,w = x.shape - x = x.view(b,n,c,h,w).contiguous() + x = torch.cat(batch_out) # [b,c,n] + x = x.transpose(-1,-2).contiguous() # [b,n,2] return x return decode \ No newline at end of file From 28c6e82285aef615bc392de63d0f20d18d6f7c31 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 16:52:26 +0000 Subject: [PATCH 61/71] add new config --- configs/360p_v2.yml | 70 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 configs/360p_v2.yml diff --git a/configs/360p_v2.yml b/configs/360p_v2.yml new file mode 100644 index 00000000..bf71727a --- /dev/null +++ b/configs/360p_v2.yml @@ -0,0 +1,70 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: game_rft + sample_size: 4 + channels: 128 + + n_layers: 13 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 16 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 30 + + causal: false + +train: + trainer_id: rft + data_id: cod_s3 + data_kwargs: + window_length: 30 + bucket_name: cod-data-latent-360x640to4x4 + include_keyframe: false + + target_batch_size: 256 + batch_size: 32 + + epochs: 200 + + opt: Muon + opt_kwargs: + lr: 1.0e-3 + momentum: 0.95 + adamw_lr: 1.0e-4 + adamw_wd: 1.0e-4 + adamw_eps: 1.0e-15 + adamw_betas: [0.9, 0.95] + adamw_keys: [core.proj_in, core.proj_out.proj] + + scheduler: null + + checkpoint_dir: checkpoints/360p + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: window + sampler_kwargs: + n_steps: 10 + cfg_scale: 1.3 + window_length: 30 + num_frames: 60 + noise_prev: 0.2 + only_return_generated: true + + n_samples: 8 + + vae_id: 720pr3dc + vae_batch_size: 4 + vae_scale: 0.13 + vae_cfg_path: configs/owl_vaes/cod_128x.yml + vae_ckpt_path: checkpoints/owl_vaes/cod_128x_30k_ema.pt + +wandb: + name: shahbuland + project: video_models + run_name: v3 \ No newline at end of file From 93bb0d7eb9eaa205313a2ede16c6c993bee5efdd Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 16:56:40 +0000 Subject: [PATCH 62/71] fix config --- configs/360p_v2.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/360p_v2.yml b/configs/360p_v2.yml index bf71727a..9a3ced68 100644 --- a/configs/360p_v2.yml +++ b/configs/360p_v2.yml @@ -54,7 +54,7 @@ train: window_length: 30 num_frames: 60 noise_prev: 0.2 - only_return_generated: true + only_return_generated: false n_samples: 8 From a0f498bdbd8f76ae9f743bc1d43626aaff7ca362 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 17:06:15 +0000 Subject: [PATCH 63/71] audio dataset added --- owl_wms/data/s3_cod_latent_audio.py | 204 ++++++++++++++++++++++++++++ 1 file changed, 204 insertions(+) create mode 100644 owl_wms/data/s3_cod_latent_audio.py diff --git a/owl_wms/data/s3_cod_latent_audio.py b/owl_wms/data/s3_cod_latent_audio.py new file mode 100644 index 00000000..e5f08d6f --- /dev/null +++ b/owl_wms/data/s3_cod_latent_audio.py @@ -0,0 +1,204 @@ +import boto3 +import threading +from dotenv import load_dotenv +import os + +load_dotenv() + +import torch +import random +from torch.utils.data import IterableDataset, DataLoader +import torch.distributed as dist +import tarfile +import io +import time + +class RandomizedQueue: + def __init__(self): + self.items = [] + + def add(self, item): + idx = random.randint(0, len(self.items)) + self.items.insert(idx, item) + + def pop(self): + if not self.items: + return None + idx = random.randint(0, len(self.items) - 1) + return self.items.pop(idx) + +TOTAL_SHARDS = 1 +NUM_SUBDIRS=1 +NUM_TARS=9 +BUCKET_NAME="cod-data-latent-360x640to4x4" + +class S3CoDLatentAudioDataset(IterableDataset): + def __init__(self, window_length=120, file_share_max=20, rank=0, world_size=1, bucket_name = BUCKET_NAME): + super().__init__() + + self.window = window_length + self.file_share_max = file_share_max + self.rank = rank + self.world_size = world_size + self.bucket_name = bucket_name + + # Queue parameters + self.max_tars = 2 + self.max_data = 1000 + + # Initialize queues + self.tar_queue = RandomizedQueue() + self.data_queue = RandomizedQueue() + + # Setup S3 client + self.s3_client = boto3.client( + 's3', + endpoint_url=os.environ['AWS_ENDPOINT_URL_S3'], + aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'], + aws_secret_access_key=os.environ['AWS_SECRET_ACCESS_KEY'], + region_name=os.environ['AWS_REGION'], + ) + + # Start background threads + self.tar_thread = threading.Thread(target=self.background_download_tars, daemon=True) + self.data_thread = threading.Thread(target=self.background_load_data, daemon=True) + self.tar_thread.start() + self.data_thread.start() + + def random_sample_prefix(self): + # For now just 2 shards (00, 01) + shard = random.randint(0, TOTAL_SHARDS-1) + # Each shard has 1000 subdirs + subdir = random.randint(0, NUM_SUBDIRS-1) + # Each subdir has multiple tars + tar_num = random.randint(0, NUM_TARS-1) + return f"{shard:02d}/{subdir:04d}/{tar_num:04d}.tar" + + def background_download_tars(self): + while True: + if len(self.tar_queue.items) < self.max_tars: + tar_path = self.random_sample_prefix() + try: + # Download tar directly to memory + response = self.s3_client.get_object(Bucket=self.bucket_name, Key=tar_path) + tar_data = response['Body'].read() + self.tar_queue.add(tar_data) + except Exception as e: + print(f"Error downloading tar {tar_path}: {e}") + else: + time.sleep(1) + + def process_tensor_file(self, tar, base_name, suffix): + try: + f = tar.extractfile(f"{base_name}.{suffix}.pt") + if f is not None: + tensor_data = f.read() + tensor = torch.load(io.BytesIO(tensor_data)) + return tensor + except: + return None + return None + + def background_load_data(self): + while True: + if len(self.data_queue.items) < self.max_data: + tar_data = self.tar_queue.pop() + if tar_data is None: + time.sleep(1) + continue + + try: + tar_file = io.BytesIO(tar_data) + with tarfile.open(fileobj=tar_file) as tar: + members = tar.getmembers() + base_names = set() + + # Get unique base names + for member in members: + if member.name.endswith('.latent.pt'): + base_names.add(member.name.split('.')[0]) + + for base_name in base_names: + # Load all tensors for this base name + latent = self.process_tensor_file(tar, base_name, "latent") + mouse = self.process_tensor_file(tar, base_name, "mouse") + button = self.process_tensor_file(tar, base_name, "buttons") + audio = self.process_tensor_file(tar, base_name, "audiolatent") + + if all(t is not None for t in [latent, mouse, button, audio]): + min_len = min(len(latent), len(mouse), len(button), len(audio)) + + # Sample multiple windows if requested + for _ in range(self.file_share_max): + if len(self.data_queue.items) >= self.max_data: + break + + max_start = min_len - self.window + if max_start <= 0: + continue + + window_start = random.randint(0, max_start) + + latent_slice = latent[window_start:window_start+self.window].float() + mouse_slice = mouse[window_start:window_start+self.window] + button_slice = button[window_start:window_start+self.window] + audio_slice = audio[window_start:window_start+self.window] + + self.data_queue.add((latent_slice, mouse_slice, button_slice, audio_slice)) + + except Exception as e: + print(f"Error processing tar: {e}") + else: + time.sleep(1) + + def __iter__(self): + while True: + item = self.data_queue.pop() + if item is not None: + yield item + else: + time.sleep(0.1) + +def collate_fn(batch): + # batch is list of quadruples + latents, mouses, buttons, audios = zip(*batch) + + latents = torch.stack(latents) # [b,n,c,h,w] + mouses = torch.stack(mouses) # [b,n,2] + buttons = torch.stack(buttons) # [b,n,n_buttons] + audios = torch.stack(audios) # [b,n,d] + + return latents, mouses, buttons, audios + +def get_loader(batch_size, **data_kwargs): + if dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + + ds = S3CoDLatentDataset(rank=rank, world_size=world_size, **data_kwargs) + return DataLoader(ds, batch_size=batch_size, collate_fn=collate_fn) + +if __name__ == "__main__": + import time + loader = get_loader(16, window_length = 120, file_share_max = 20, include_audio = True) + + start = time.time() + batch = next(iter(loader)) + end = time.time() + first_time = end - start + + start = time.time() + batch = next(iter(loader)) + end = time.time() + second_time = end - start + + x,y,z = batch + print(f"Time to load first batch: {first_time:.2f}s") + print(f"Time to load second batch: {second_time:.2f}s") + print(f"Video shape: {x.shape}") + print(x.std()) + print(f"Mouse shape: {y.shape}") + print(f"Button shape: {z.shape}") From d84c8ffd9db15291c71638a8512de58707dd38d1 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 18:01:39 +0000 Subject: [PATCH 64/71] av --- configs/av.yml | 69 +++++++ owl_wms/models/__init__.py | 3 + owl_wms/models/gamerft_audio.py | 148 +++++++++++++++ owl_wms/models/gamerft_shortcut_simple.py | 2 + owl_wms/sampling/__init__.py | 5 +- owl_wms/sampling/av_window.py | 138 ++++++++++++++ owl_wms/trainers/__init__.py | 5 +- owl_wms/trainers/av_trainer.py | 209 ++++++++++++++++++++++ owl_wms/utils/owl_vae_bridge.py | 3 +- 9 files changed, 579 insertions(+), 3 deletions(-) create mode 100644 configs/av.yml create mode 100644 owl_wms/models/gamerft_audio.py create mode 100644 owl_wms/sampling/av_window.py create mode 100644 owl_wms/trainers/av_trainer.py diff --git a/configs/av.yml b/configs/av.yml new file mode 100644 index 00000000..1066c406 --- /dev/null +++ b/configs/av.yml @@ -0,0 +1,69 @@ +model: + model_id: game_rft_audio + sample_size: 4 + channels: 128 + + n_layers: 13 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 17 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 30 + + causal: false + +train: + trainer_id: av + data_id: cod_s3_audio + data_kwargs: + window_length: 30 + bucket_name: cod-data-latent-360x640to4x4 + include_keyframe: false + + target_batch_size: 256 + batch_size: 32 + + epochs: 200 + + opt: Muon + opt_kwargs: + lr: 1.0e-3 + momentum: 0.95 + adamw_lr: 1.0e-4 + adamw_wd: 1.0e-4 + adamw_eps: 1.0e-15 + adamw_betas: [0.9, 0.95] + adamw_keys: [core.proj_in, core.proj_out.proj] + + scheduler: null + + checkpoint_dir: checkpoints/360p + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: av_window + sampler_kwargs: + n_steps: 10 + cfg_scale: 1.3 + window_length: 30 + num_frames: 60 + noise_prev: 0.2 + only_return_generated: false + + n_samples: 8 + + vae_id: 720pr3dc + vae_batch_size: 4 + vae_scale: 0.13 + vae_cfg_path: configs/owl_vaes/cod_128x.yml + vae_ckpt_path: checkpoints/owl_vaes/cod_128x_30k_ema.pt + +wandb: + name: shahbuland + project: video_models + run_name: v3 \ No newline at end of file diff --git a/owl_wms/models/__init__.py b/owl_wms/models/__init__.py index 91c28793..0ad456b7 100644 --- a/owl_wms/models/__init__.py +++ b/owl_wms/models/__init__.py @@ -8,6 +8,9 @@ def get_model_cls(model_id): if model_id == "shortcut_2": from .gamerft_shortcut_simple import ShortcutGameRFT return ShortcutGameRFT + if model_id == "game_rft_audio": + from .gamerft_audio import GameRFT + return GameRFT diff --git a/owl_wms/models/gamerft_audio.py b/owl_wms/models/gamerft_audio.py new file mode 100644 index 00000000..22df09f2 --- /dev/null +++ b/owl_wms/models/gamerft_audio.py @@ -0,0 +1,148 @@ +""" +GameRFT with Audio +""" + +import torch +from torch import nn +import torch.nn.functional as F + +import einops as eo + +from ..nn.embeddings import ( + TimestepEmbedding, + ControlEmbedding, + LearnedPosEnc +) +from ..nn.attn import UViT, FinalLayer + +class GameRFTCore(nn.Module): + def __init__(self, config): + super().__init__() + + self.transformer = UViT(config) + self.control_embed = ControlEmbedding(config.n_buttons, config.d_model) + self.t_embed = TimestepEmbedding(config.d_model) + + self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) + self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) + + self.audio_proj_in = nn.Layer(config.audio_channels, config.d_model, bias=False) + self.audio_proj_out = FinalLayer(None, config.d_model, config.audio_channels) + + self.pos_enc = LearnedPosEnc(config.tokens_per_frame * config.n_frames, config.d_model) + + def forward(self, x, audio, t, mouse, btn): + # x is [b,n,c,h,w] + # audio is [b,n,c] + # t is [b,n] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + + ctrl_cond = self.control_embed(mouse, btn) + t_cond = self.t_embed(t) + + cond = ctrl_cond + t_cond # [b,n,d] + + b,n,c,h,w = x.shape + x = eo.rearrange(x, 'b n c h w -> b (n h w) c') + + x = self.proj_in(x) + audio = self.proj_in(audio).unsqueeze(-2) # [b,n,1,d] + + x = eo.rearrange(x, 'b (n f) d -> b n f d', n = n) + x = torch.cat([x, audio], dim = -2) + x = eo.rearrange(x, 'b n f d -> b (n f) d') + + x = self.pos_enc(x) + x = self.transformer(x, cond) + + # Split into video and audio tokens + x = eo.rearrange(x, 'b (n f) d -> b n f d', n=n) + video, audio = x[...,:-1,:], x[...,-1:,:] + + # Project video tokens + video = eo.rearrange(video, 'b n f d -> b (n f) d') + video = self.proj_out(video, cond) + video = eo.rearrange(video, 'b (n h w) c -> b n c h w', n=n, h=h, w=w) + + # Project audio tokens + audio = eo.rearrange(audio, 'b n 1 d -> b n d') + audio = self.audio_proj_out(audio, cond) + + return video, audio + +class GameRFT(nn.Module): + def __init__(self, config): + super().__init__() + + self.core = GameRFTCore(config) + self.cfg_prob = config.cfg_prob + + def forward(self, x, audio, mouse, btn, return_dict = False, cfg_prob = None): + # x is [b,n,c,h,w] + # audio is [b,n,c] + # mouse is [b,n,2] + # btn is [b,n,n_buttons] + b,n,c,h,w = x.shape + + # Apply classifier-free guidance dropout + if cfg_prob is None: + cfg_prob = self.cfg_prob + if cfg_prob > 0.0: + mask = torch.rand(b, device=x.device) <= self.cfg_prob + null_mouse = torch.zeros_like(mouse) + null_btn = torch.zeros_like(btn) + + # Where mask is True, replace with zeros + mouse = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_mouse, mouse) + btn = torch.where(mask.unsqueeze(-1).unsqueeze(-1), null_btn, btn) + + with torch.no_grad(): + ts = torch.randn(b,n,device=x.device,dtype=x.dtype).sigmoid() + + # Video noise + ts_exp = eo.repeat(ts, 'b n -> b n 1 1 1') + z_video = torch.randn_like(x) + lerpd_video = x * (1. - ts_exp) + z_video * ts_exp + target_video = z_video - x + + # Audio noise + ts_exp_audio = ts.unsqueeze(-1) + z_audio = torch.randn_like(audio) + lerpd_audio = audio * (1. - ts_exp_audio) + z_audio * ts_exp_audio + target_audio = z_audio - audio + + pred_video, pred_audio = self.core(lerpd_video, lerpd_audio, ts, mouse, btn) + video_loss = F.mse_loss(pred_video, target_video) + audio_loss = F.mse_loss(pred_audio, target_audio) + diff_loss = video_loss + audio_loss + + if not return_dict: + return diff_loss + else: + return { + 'diffusion_loss': diff_loss, + 'video_loss': video_loss, + 'audio_loss': audio_loss, + 'lerpd_video': lerpd_video, + 'lerpd_audio': lerpd_audio, + 'pred_video': pred_video, + 'pred_audio': pred_audio, + 'ts': ts, + 'z_video': z_video, + 'z_audio': z_audio + } + +if __name__ == "__main__": + from ..configs import Config + + cfg = Config.from_yaml("configs/basic.yml").model + model = GameRFT(cfg).cuda().bfloat16() + + with torch.no_grad(): + x = torch.randn(1, 128, 16, 256, device='cuda', dtype=torch.bfloat16) + mouse = torch.randn(1, 128, 2, device='cuda', dtype=torch.bfloat16) + btn = torch.randn(1, 128, 11, device='cuda', dtype=torch.bfloat16) + + loss = model(x, mouse, btn) + print(f"Loss: {loss.item()}") \ No newline at end of file diff --git a/owl_wms/models/gamerft_shortcut_simple.py b/owl_wms/models/gamerft_shortcut_simple.py index 6cdc8147..54d031ce 100644 --- a/owl_wms/models/gamerft_shortcut_simple.py +++ b/owl_wms/models/gamerft_shortcut_simple.py @@ -32,6 +32,8 @@ def __init__(self, config): self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) self.config = config + if config.audio_tokens > 0: + self.audio_proj = nn.Linear(config.audio_channels, config.d_model, bias = False) def sample(self, x, mouse, btn, kv_cache = None, t = None, d = None): """ diff --git a/owl_wms/sampling/__init__.py b/owl_wms/sampling/__init__.py index 3de909d5..b2bc45cc 100644 --- a/owl_wms/sampling/__init__.py +++ b/owl_wms/sampling/__init__.py @@ -14,4 +14,7 @@ def get_sampler_cls(sampler_id): return CacheShortcutSampler elif sampler_id == "shortcut_2": from .shortcut_sampler import WindowShortcutSamplerNoKeyframe - return WindowShortcutSamplerNoKeyframe \ No newline at end of file + return WindowShortcutSamplerNoKeyframe + elif sampler_id == "av_window": + from .av_window import AVWindowSampler + return AVWindowSampler \ No newline at end of file diff --git a/owl_wms/sampling/av_window.py b/owl_wms/sampling/av_window.py new file mode 100644 index 00000000..78359184 --- /dev/null +++ b/owl_wms/sampling/av_window.py @@ -0,0 +1,138 @@ +import torch +from torch import nn +import torch.nn.functional as F + +from tqdm import tqdm + +from ..utils import batch_permute_to_length + +def zlerp(x, alpha): + z = torch.randn_like(x) + return x * (1. - alpha) + z * alpha + +class AVWindowSampler: + """ + Window CFG Sampler samples new frames one by one, by inpainting the final frame. + This is basically diffusion forcing. + + :param n_steps: Number of diffusion steps for each frame (diffusoin steps) + :param cfg_scale: CFG scale for each frame + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param noise_prev: Noise previous frame + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, n_steps = 20, cfg_scale = 1.3, window_length = 60, num_frames = 60, noise_prev = 0.2, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.noise_prev = noise_prev + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, dummy_batch, audio, mouse, btn, decode_fn = None, audio_decode_fn = None, image_scale = 1, audio_scale = 1): + # dummy_batch is [b,n,c,h,w] + # audio is [b,n,c] and should be treated same as video (it'e being generated) + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + + sampling_steps = self.n_steps + num_frames = self.num_frames + + dt = 1. / sampling_steps + + clean_history = dummy_batch.clone() + clean_audio_history = audio.clone() + + extended_mouse, extended_btn = batch_permute_to_length(mouse, btn, num_frames + self.window_length) + + def step_history(): + # Video history + new_history = clean_history.clone()[:,-self.window_length:] + b,n,c,h,w = new_history.shape + new_history[:,:-1] = zlerp(new_history[:,1:],self.noise_prev) + new_history[:,-1] = torch.randn_like(new_history[:,0]) + + # Audio history + new_audio = clean_audio_history.clone()[:,-self.window_length:] + new_audio[:,:-1] = zlerp(new_audio[:,1:],self.noise_prev) + new_audio[:,-1] = torch.randn_like(new_audio[:,0]) + + return new_history, new_audio + + for frame_idx in tqdm(range(num_frames)): + local_history, local_audio = step_history() + ts_history = torch.ones(local_history.shape[0], local_history.shape[1], device=local_history.device,dtype=local_history.dtype) + ts_history[:,:-1] = self.noise_prev + + mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] + btn = extended_btn[:,frame_idx:frame_idx+self.window_length] + + mouse_batch = torch.cat([mouse, torch.zeros_like(mouse)], dim=0) + btn_batch = torch.cat([btn, torch.zeros_like(btn)], dim=0) + for _ in range(sampling_steps): + # CFG Branches + x = local_history.clone() + a = local_audio.clone() + ts = ts_history.clone() + + x_batch = torch.cat([x, x], dim=0) + a_batch = torch.cat([a, a], dim=0) + ts_batch = torch.cat([ts, ts], dim=0) + + pred_video_batch, pred_audio_batch = model(x_batch, a_batch, ts_batch, mouse_batch, btn_batch) + + # Split predictions back into conditional and unconditional + cond_pred_video, uncond_pred_video = pred_video_batch.chunk(2) + cond_pred_audio, uncond_pred_audio = pred_audio_batch.chunk(2) + + pred_video = uncond_pred_video + self.cfg_scale * (cond_pred_video - uncond_pred_video) + pred_audio = uncond_pred_audio + self.cfg_scale * (cond_pred_audio - uncond_pred_audio) + + x = x - pred_video*dt + a = a - pred_audio*dt + ts = ts - dt + + local_history[:,-1] = x[:,-1] + local_audio[:,-1] = a[:,-1] + ts_history[:,-1] = ts[:,-1] + + # Frame is entirely cleaned now + new_frame = local_history[:,-1:] + new_audio = local_audio[:,-1:] + clean_history = torch.cat([clean_history, new_frame], dim=1) + clean_audio_history = torch.cat([clean_audio_history, new_audio], dim=1) + + x = clean_history + audio = clean_audio_history + if self.only_return_generated: + x = x[:,-num_frames:] + audio = audio[:,-num_frames:] + extended_mouse = extended_mouse[:,-num_frames:] + extended_btn = extended_btn[:,-num_frames:] + + if decode_fn is not None: + x = x * image_scale + x = decode_fn(x) + + if audio_decode_fn is not None: + audio = audio * audio_scale + audio = audio_decode_fn(audio) + + return x, audio, extended_mouse, extended_btn + +def test_window_cfg_sampler(): + sampler = WindowCFGSampler() + model = lambda x, ts, mouse, btn: x + dummy_batch = torch.randn(1, 32, 128, 4, 4) + + mouse = torch.zeros(1, 32, 2) + btn = torch.zeros(1, 32, 11) + x = sampler(model, dummy_batch, mouse, btn) + print(x.shape) + +if __name__ == "__main__": + test_window_cfg_sampler() diff --git a/owl_wms/trainers/__init__.py b/owl_wms/trainers/__init__.py index 09a4bebc..4a6826fc 100644 --- a/owl_wms/trainers/__init__.py +++ b/owl_wms/trainers/__init__.py @@ -11,4 +11,7 @@ def get_trainer_cls(trainer_id): return ShortcutTrainer if trainer_id == "shortcut_2": from .shortcut_trainer_2 import ShortcutTrainer - return ShortcutTrainer \ No newline at end of file + return ShortcutTrainer + if trainer_id == "av": + from .av_trainer import AVRFTTrainer + return AVRFTTrainer \ No newline at end of file diff --git a/owl_wms/trainers/av_trainer.py b/owl_wms/trainers/av_trainer.py new file mode 100644 index 00000000..55bb3488 --- /dev/null +++ b/owl_wms/trainers/av_trainer.py @@ -0,0 +1,209 @@ +import torch +from ema_pytorch import EMA +import wandb +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP +import torch.distributed as dist +import einops as eo + +from .base import BaseTrainer + +from ..utils import freeze, Timer, find_unused_params +from ..schedulers import get_scheduler_cls +from ..models import get_model_cls +from ..sampling import get_sampler_cls +from ..data import get_loader +from ..utils.logging import LogHelper, to_wandb_av +from ..muon import init_muon +from ..utils.owl_vae_bridge import get_decoder_only, make_batched_decode_fn, make_batched_audio_decode_fn + +class AVRFTTrainer(BaseTrainer): + """ + Trainer for rectified flow transformer + + :param train_cfg: Configuration for training + :param logging_cfg: Configuration for logging + :param model_cfg: Configuration for model + :param global_rank: Rank across all devices. + :param local_rank: Rank for current device on this process. + :param world_size: Overall number of devices + """ + def __init__(self,*args,**kwargs): + super().__init__(*args,**kwargs) + + model_id = self.model_cfg.model_id + self.model = get_model_cls(model_id)(self.model_cfg) + + # Print model size + if self.rank == 0: + n_params = sum(p.numel() for p in self.model.parameters()) + print(f"Model has {n_params:,} parameters") + + self.ema = None + self.opt = None + self.scheduler = None + self.scaler = None + + self.total_step_counter = 0 + self.decoder = get_decoder_only( + self.train_cfg.vae_id, + self.train_cfg.vae_cfg_path, + self.train_cfg.vae_ckpt_path + ) + + self.audio_decoder = get_decoder_only( + self.train_cfg.audio_vae_id, + self.train_cfg.audio_vae_cfg_path, + self.train_cfg.audio_vae_ckpt_path + ) + + freeze(self.decoder) + + def save(self): + save_dict = { + 'model' : self.model.state_dict(), + 'ema' : self.ema.state_dict(), + 'opt' : self.opt.state_dict(), + 'scaler' : self.scaler.state_dict(), + 'steps': self.total_step_counter + } + if self.scheduler is not None: + save_dict['scheduler'] = self.scheduler.state_dict() + super().save(save_dict) + + def load(self): + has_ckpt = False + try: + if self.train_cfg.resume_ckpt is not None: + save_dict = super().load(self.train_cfg.resume_ckpt) + has_ckpt = True + except: + print("Error loading checkpoint") + + if not has_ckpt: + return + + + self.model.load_state_dict(save_dict['model']) + self.ema.load_state_dict(save_dict['ema']) + self.opt.load_state_dict(save_dict['opt']) + if self.scheduler is not None and 'scheduler' in save_dict: + self.scheduler.load_state_dict(save_dict['scheduler']) + self.scaler.load_state_dict(save_dict['scaler']) + self.total_step_counter = save_dict['steps'] + + def train(self): + torch.cuda.set_device(self.local_rank) + + # Prepare model and ema + self.model = self.model.cuda().train() + if self.world_size > 1: + self.model = DDP(self.model, device_ids=[self.local_rank]) + self.decoder = self.decoder.cuda().eval().bfloat16() + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + + self.ema = EMA( + self.model, + beta = 0.999, + update_after_step = 0, + update_every = 1 + ) + #torch.compile(self.ema.ema_model.module.core if self.world_size > 1 else self.ema.ema_model.core, dynamic=False, fullgraph=True) + + def get_ema_core(): + if self.world_size > 1: + return self.ema.ema_model.module.core + else: + return self.ema.ema_model.core + + # Set up optimizer and scheduler + if self.train_cfg.opt.lower() == "muon": + self.opt = init_muon(self.model, rank=self.rank,world_size=self.world_size,**self.train_cfg.opt_kwargs) + else: + self.opt = getattr(torch.optim, self.train_cfg.opt)(self.model.parameters(), **self.train_cfg.opt_kwargs) + + if self.train_cfg.scheduler is not None: + self.scheduler = get_scheduler_cls(self.train_cfg.scheduler)(self.opt, **self.train_cfg.scheduler_kwargs) + + # Grad accum setup and scaler + accum_steps = self.train_cfg.target_batch_size // self.train_cfg.batch_size // self.world_size + accum_steps = max(1, accum_steps) + self.scaler = torch.amp.GradScaler() + ctx = torch.amp.autocast('cuda',torch.bfloat16) + + self.load() + + # Timer reset + timer = Timer() + timer.reset() + metrics = LogHelper() + if self.rank == 0: + wandb.watch(self.get_module(), log = 'all') + + # Dataset setup + loader = get_loader(self.train_cfg.data_id, self.train_cfg.batch_size, **self.train_cfg.data_kwargs) + sampler = get_sampler_cls(self.train_cfg.sampler_id)(**self.train_cfg.sampler_kwargs) + + local_step = 0 + for _ in range(self.train_cfg.epochs): + for batch_vid, batch_audio, batch_mouse, batch_btn in loader: + batch_vid = batch_vid.cuda().bfloat16() / self.train_cfg.vae_scale + batch_audio = batch_audio.cuda().bfloat16() / self.train_cfg.audio_vae_scale + batch_mouse = batch_mouse.cuda().bfloat16() + batch_btn = batch_btn.cuda().bfloat16() + + with ctx: + loss = self.model(batch_vid,batch_audio,batch_mouse,batch_btn) / accum_steps + + self.scaler.scale(loss).backward() + #find_unused_params(self.model) + + metrics.log('diffusion_loss', loss) + + local_step += 1 + if local_step % accum_steps == 0: + # Updates + if self.train_cfg.opt.lower() != "muon": + self.scaler.unscale_(self.opt) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + + self.scaler.step(self.opt) + self.opt.zero_grad(set_to_none=True) + + self.scaler.update() + + if self.scheduler is not None: + self.scheduler.step() + self.ema.update() + + # Do logging + with torch.no_grad(): + wandb_dict = metrics.pop() + wandb_dict['time'] = timer.hit() + timer.reset() + + # Sampling commented out for now + if self.total_step_counter % self.train_cfg.sample_interval == 0: + with ctx, torch.no_grad(): + n_samples = self.train_cfg.n_samples + samples, audio, sample_mouse, sample_button = sampler( + get_ema_core(), + batch_vid[:n_samples], + batch_audio[:n_samples], + batch_mouse[:n_samples], + batch_btn[:n_samples], + decode_fn = decode_fn, + scale=self.train_cfg.vae_scale + ) # -> [b,n,c,h,w] + if self.rank == 0: wandb_dict['samples'] = to_wandb_av(samples, audio, sample_mouse, sample_button) + + + if self.rank == 0: + wandb.log(wandb_dict) + + self.total_step_counter += 1 + if self.total_step_counter % self.train_cfg.save_interval == 0: + if self.rank == 0: + self.save() + + self.barrier() \ No newline at end of file diff --git a/owl_wms/utils/owl_vae_bridge.py b/owl_wms/utils/owl_vae_bridge.py index c306f18f..e4cfd74a 100644 --- a/owl_wms/utils/owl_vae_bridge.py +++ b/owl_wms/utils/owl_vae_bridge.py @@ -56,7 +56,8 @@ def decode(x): @torch.no_grad() def make_batched_audio_decode_fn(decoder, batch_size = 8): def decode(x): - # x is [b,c,n] audio samples + # x is [b,n,c] audio samples + x = x.transpose(1,2) b,c,n = x.shape batches = x.contiguous().split(batch_size) From 7022ce2b42706926443df79e4e12288c91160562 Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 18:30:20 +0000 Subject: [PATCH 65/71] debug --- configs/av.yml | 16 +++++++++++----- owl_wms/data/__init__.py | 5 ++++- owl_wms/data/s3_cod_latent_audio.py | 4 ++-- owl_wms/models/gamerft_audio.py | 6 +++--- owl_wms/trainers/av_trainer.py | 9 +++++++-- 5 files changed, 27 insertions(+), 13 deletions(-) diff --git a/configs/av.yml b/configs/av.yml index 1066c406..f8755bdd 100644 --- a/configs/av.yml +++ b/configs/av.yml @@ -2,6 +2,7 @@ model: model_id: game_rft_audio sample_size: 4 channels: 128 + audio_channels: 64 n_layers: 13 n_heads: 16 @@ -22,10 +23,9 @@ train: data_kwargs: window_length: 30 bucket_name: cod-data-latent-360x640to4x4 - include_keyframe: false - target_batch_size: 256 - batch_size: 32 + target_batch_size: 16 + batch_size: 16 epochs: 200 @@ -57,13 +57,19 @@ train: n_samples: 8 - vae_id: 720pr3dc + vae_id: null vae_batch_size: 4 vae_scale: 0.13 + audio_vae_scale: 0.17 + vae_cfg_path: configs/owl_vaes/cod_128x.yml vae_ckpt_path: checkpoints/owl_vaes/cod_128x_30k_ema.pt + audio_vae_id: null + audio_vae_cfg_path: configs/owl_vaes/cod_audio.yml + audio_vae_ckpt_path: checkpoints/owl_vaes/cod_audio_20k_ema.pt + wandb: name: shahbuland project: video_models - run_name: v3 \ No newline at end of file + run_name: av \ No newline at end of file diff --git a/owl_wms/data/__init__.py b/owl_wms/data/__init__.py index 564abd25..ce809017 100644 --- a/owl_wms/data/__init__.py +++ b/owl_wms/data/__init__.py @@ -8,4 +8,7 @@ def get_loader(data_id, batch_size, **data_kwargs): return local_cod_latent.get_loader(batch_size, **data_kwargs) elif data_id == "cod_s3": from . import s3_cod_latent - return s3_cod_latent.get_loader(batch_size, **data_kwargs) \ No newline at end of file + return s3_cod_latent.get_loader(batch_size, **data_kwargs) + elif data_id == "cod_s3_audio": + from . import s3_cod_latent_audio + return s3_cod_latent_audio.get_loader(batch_size, **data_kwargs) \ No newline at end of file diff --git a/owl_wms/data/s3_cod_latent_audio.py b/owl_wms/data/s3_cod_latent_audio.py index e5f08d6f..1a8b7873 100644 --- a/owl_wms/data/s3_cod_latent_audio.py +++ b/owl_wms/data/s3_cod_latent_audio.py @@ -168,7 +168,7 @@ def collate_fn(batch): buttons = torch.stack(buttons) # [b,n,n_buttons] audios = torch.stack(audios) # [b,n,d] - return latents, mouses, buttons, audios + return latents, audios, mouses, buttons def get_loader(batch_size, **data_kwargs): if dist.is_initialized(): @@ -178,7 +178,7 @@ def get_loader(batch_size, **data_kwargs): rank = 0 world_size = 1 - ds = S3CoDLatentDataset(rank=rank, world_size=world_size, **data_kwargs) + ds = S3CoDLatentAudioDataset(rank=rank, world_size=world_size, **data_kwargs) return DataLoader(ds, batch_size=batch_size, collate_fn=collate_fn) if __name__ == "__main__": diff --git a/owl_wms/models/gamerft_audio.py b/owl_wms/models/gamerft_audio.py index 22df09f2..d28026b3 100644 --- a/owl_wms/models/gamerft_audio.py +++ b/owl_wms/models/gamerft_audio.py @@ -26,7 +26,7 @@ def __init__(self, config): self.proj_in = nn.Linear(config.channels, config.d_model, bias = False) self.proj_out = FinalLayer(config.sample_size, config.d_model, config.channels) - self.audio_proj_in = nn.Layer(config.audio_channels, config.d_model, bias=False) + self.audio_proj_in = nn.Linear(config.audio_channels, config.d_model, bias=False) self.audio_proj_out = FinalLayer(None, config.d_model, config.audio_channels) self.pos_enc = LearnedPosEnc(config.tokens_per_frame * config.n_frames, config.d_model) @@ -47,7 +47,7 @@ def forward(self, x, audio, t, mouse, btn): x = eo.rearrange(x, 'b n c h w -> b (n h w) c') x = self.proj_in(x) - audio = self.proj_in(audio).unsqueeze(-2) # [b,n,1,d] + audio = self.audio_proj_in(audio).unsqueeze(-2) # [b,n,1,d] x = eo.rearrange(x, 'b (n f) d -> b n f d', n = n) x = torch.cat([x, audio], dim = -2) @@ -111,7 +111,7 @@ def forward(self, x, audio, mouse, btn, return_dict = False, cfg_prob = None): z_audio = torch.randn_like(audio) lerpd_audio = audio * (1. - ts_exp_audio) + z_audio * ts_exp_audio target_audio = z_audio - audio - + pred_video, pred_audio = self.core(lerpd_video, lerpd_audio, ts, mouse, btn) video_loss = F.mse_loss(pred_video, target_video) audio_loss = F.mse_loss(pred_audio, target_audio) diff --git a/owl_wms/trainers/av_trainer.py b/owl_wms/trainers/av_trainer.py index 55bb3488..e482dd42 100644 --- a/owl_wms/trainers/av_trainer.py +++ b/owl_wms/trainers/av_trainer.py @@ -100,7 +100,10 @@ def train(self): if self.world_size > 1: self.model = DDP(self.model, device_ids=[self.local_rank]) self.decoder = self.decoder.cuda().eval().bfloat16() + self.audio_decoder = self.audio_decoder.cuda().eval().bfloat16() + decode_fn = make_batched_decode_fn(self.decoder, self.train_cfg.vae_batch_size) + audio_decode_fn = make_batched_audio_decode_fn(self.audio_decoder, self.train_cfg.vae_batch_size) self.ema = EMA( self.model, @@ -192,8 +195,10 @@ def get_ema_core(): batch_audio[:n_samples], batch_mouse[:n_samples], batch_btn[:n_samples], - decode_fn = decode_fn, - scale=self.train_cfg.vae_scale + decode_fn, + audio_decode_fn, + self.train_cfg.vae_scale, + self.train_cfg.audio_vae_scale ) # -> [b,n,c,h,w] if self.rank == 0: wandb_dict['samples'] = to_wandb_av(samples, audio, sample_mouse, sample_button) From 3ee6d10d059b064ccb5f4951a4dfb911632abc6d Mon Sep 17 00:00:00 2001 From: shahbuland Date: Wed, 11 Jun 2025 18:57:07 +0000 Subject: [PATCH 66/71] configs for av model --- configs/av.yml | 4 +-- configs/owl_vaes/cod_128x.yml | 53 +++++++++++++++++++++++++++++++ configs/owl_vaes/cod_audio.yml | 57 ++++++++++++++++++++++++++++++++++ owl_wms/trainers/av_trainer.py | 5 ++- owl_wms/utils/logging.py | 17 +++++----- 5 files changed, 125 insertions(+), 11 deletions(-) create mode 100644 configs/owl_vaes/cod_128x.yml create mode 100644 configs/owl_vaes/cod_audio.yml diff --git a/configs/av.yml b/configs/av.yml index f8755bdd..72811dfb 100644 --- a/configs/av.yml +++ b/configs/av.yml @@ -24,8 +24,8 @@ train: window_length: 30 bucket_name: cod-data-latent-360x640to4x4 - target_batch_size: 16 - batch_size: 16 + target_batch_size: 256 + batch_size: 32 epochs: 200 diff --git a/configs/owl_vaes/cod_128x.yml b/configs/owl_vaes/cod_128x.yml new file mode 100644 index 00000000..76502ca9 --- /dev/null +++ b/configs/owl_vaes/cod_128x.yml @@ -0,0 +1,53 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: dcae + sample_size: [360,640] + channels: 3 + latent_size: 4 + latent_channels: 128 + + noise_decoder_inputs: 0.0 + ch_0: 128 + ch_max: 1024 + + encoder_blocks_per_stage: [3, 3, 3, 3, 3, 3, 3, 3] + decoder_blocks_per_stage: [3, 3, 3, 3, 3, 3, 3, 3] + + checkpoint_grads: true + +train: + trainer_id: rec + data_id: s3_cod + target_batch_size: 128 + batch_size: 16 + + epochs: 200 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + weight_decay: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-15 + + lpips_type: convnext + loss_weights: + latent_reg: 1.0e-6 + lpips: 10.0 + se_reg: 0.0 + + scheduler: LinearWarmup + scheduler_kwargs: + warmup_steps: 3000 + min_lr: 5.0e-6 + + checkpoint_dir: checkpoints/cod_128x + resume_ckpt: null #checkpoints/2d_64x/step_10000.pt + + sample_interval: 1000 + save_interval: 5000 + +wandb: + name: ${env:WANDB_USER_NAME} + project: new_vaes + run_name: 128x_cod \ No newline at end of file diff --git a/configs/owl_vaes/cod_audio.yml b/configs/owl_vaes/cod_audio.yml new file mode 100644 index 00000000..ad198e7c --- /dev/null +++ b/configs/owl_vaes/cod_audio.yml @@ -0,0 +1,57 @@ +model: + model_id: audio_ae + + channels: 2 + latent_channels: 64 + ch_0: 128 + ch_max: 512 + + strides: [3, 5, 7, 7, 1] + + eq: true + checkpoint_grads: true + +train: + trainer_id: audio_rec + data_id: local_cod_audio + data_kwargs: + window_length: 88200 + root: "../cod_download/raw" + + target_batch_size: 128 + batch_size: 16 + epochs: 100 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + eps: 1.0e-15 + betas: [0.9, 0.95] + weight_decay: 1.0e-4 + + loss_weights: + recon: 2.5 + stft: 1.5 + kl: 1.0e-5 + lr_ms_ratio: 0.5 + hubert: 0.0 + crt: 4.0 + + sample_rate: 44100 + n_fft_list: [1024, 2048, 512] + + scheduler: LinearWarmup + scheduler_kwargs: + warmup_steps: 1500 + min_lr: 1.0e-6 + + checkpoint_dir: checkpoints/audio_ae + sample_interval: 500 + save_interval: 5000 + + resume_ckpt: null + +wandb: + name: ${env:WANDB_USER_NAME} + project: owl_audio_vaes + run_name: audio_ae_baseline \ No newline at end of file diff --git a/owl_wms/trainers/av_trainer.py b/owl_wms/trainers/av_trainer.py index e482dd42..4aa78301 100644 --- a/owl_wms/trainers/av_trainer.py +++ b/owl_wms/trainers/av_trainer.py @@ -200,7 +200,10 @@ def get_ema_core(): self.train_cfg.vae_scale, self.train_cfg.audio_vae_scale ) # -> [b,n,c,h,w] - if self.rank == 0: wandb_dict['samples'] = to_wandb_av(samples, audio, sample_mouse, sample_button) + if self.rank == 0: + video, audio = to_wandb_av(samples, audio, sample_mouse, sample_button) + wandb_dict['samples'] = video + wandb_dict['audio_samples'] = audio if self.rank == 0: diff --git a/owl_wms/utils/logging.py b/owl_wms/utils/logging.py index 9f2bca7e..fc6912c1 100644 --- a/owl_wms/utils/logging.py +++ b/owl_wms/utils/logging.py @@ -78,6 +78,7 @@ def to_wandb(x, batch_mouse, batch_btn, gather = False, max_samples = 8): @torch.no_grad() def to_wandb_av(x, audio, batch_mouse, batch_btn, gather = False, max_samples = 8): # x is [b,n,c,h,w] + # audio is [b,n,2] x = x.clamp(-1, 1) x = x[:max_samples] audio = audio[:max_samples] @@ -96,12 +97,12 @@ def to_wandb_av(x, audio, batch_mouse, batch_btn, gather = False, max_samples = # Convert audio to numpy float32 [-1,1] audio = audio.cpu().float().numpy() - # Create list of video/audio pairs - samples = [] - for i in range(len(x)): - video = wandb.Video(x[i], format='gif', fps=60) - # Sample rate assumed to be 48kHz based on common audio standards - audio_sample = wandb.Audio(audio[i].T, sample_rate=48000) - samples.append((video, audio_sample)) + # Create grid of videos like in to_wandb + if max_samples == 8: + x = eo.rearrange(x, '(r c) n d h w -> n d (r h) (c w)', r = 2, c = 4) + + # Create video and audio objects + video = wandb.Video(x, format='gif', fps=60) + audio_samples = [wandb.Audio(audio[i], sample_rate=44100) for i in range(len(audio))] - return samples \ No newline at end of file + return video, audio_samples \ No newline at end of file From 8f7bd7dbb82f9c0b0d59f9944fbf844ba8c65a97 Mon Sep 17 00:00:00 2001 From: Sami Date: Thu, 12 Jun 2025 01:03:48 -0500 Subject: [PATCH 67/71] cfg --- owl_wms/configs.py | 5 + owl_wms/data/s3_cod_latent_audio.py | 2 +- owl_wms/sampling/av_window.py | 138 ++++++++++++++++- owl_wms/utils/__init__.py | 2 +- owl_wms/utils/owl_vae_bridge.py | 20 --- webapp/checkpoints/configs/ae.yml | 53 +++++++ webapp/checkpoints/configs/audio.yml | 57 +++++++ webapp/checkpoints/configs/av.yml | 75 +++++++++ webapp/checkpoints/shortcut.yaml | 67 -------- webapp/config.yaml | 25 --- webapp/server.py | 26 +--- webapp/streaming.py | 155 +++++-------------- webapp/user_session.py | 2 +- webapp/utils/av_window_inference_pipeline.py | 127 +++++++++++++++ webapp/utils/configs.py | 38 +++-- webapp/utils/create_samplers.py | 46 +++++- webapp/webapp_config.yaml | 27 ++++ 17 files changed, 598 insertions(+), 267 deletions(-) create mode 100644 webapp/checkpoints/configs/ae.yml create mode 100644 webapp/checkpoints/configs/audio.yml create mode 100644 webapp/checkpoints/configs/av.yml delete mode 100644 webapp/checkpoints/shortcut.yaml delete mode 100644 webapp/config.yaml create mode 100644 webapp/utils/av_window_inference_pipeline.py create mode 100644 webapp/webapp_config.yaml diff --git a/owl_wms/configs.py b/owl_wms/configs.py index 5283e5e8..7b463e00 100644 --- a/owl_wms/configs.py +++ b/owl_wms/configs.py @@ -64,6 +64,11 @@ class TrainingConfig: vae_scale : float = 0.34 vae_batch_size: int = 4 + audio_vae_id : str = None + audio_vae_cfg_path : str = None + audio_vae_ckpt_path : str = None + audio_vae_scale : float = 0.17 + @dataclass class WANDBConfig: name : str = None diff --git a/owl_wms/data/s3_cod_latent_audio.py b/owl_wms/data/s3_cod_latent_audio.py index 1a8b7873..1aeed32c 100644 --- a/owl_wms/data/s3_cod_latent_audio.py +++ b/owl_wms/data/s3_cod_latent_audio.py @@ -183,7 +183,7 @@ def get_loader(batch_size, **data_kwargs): if __name__ == "__main__": import time - loader = get_loader(16, window_length = 120, file_share_max = 20, include_audio = True) + loader = get_loader(1, window_length = 30, file_share_max = 20) start = time.time() batch = next(iter(loader)) diff --git a/owl_wms/sampling/av_window.py b/owl_wms/sampling/av_window.py index 78359184..ee4715ad 100644 --- a/owl_wms/sampling/av_window.py +++ b/owl_wms/sampling/av_window.py @@ -124,15 +124,141 @@ def step_history(): return x, audio, extended_mouse, extended_btn -def test_window_cfg_sampler(): - sampler = WindowCFGSampler() + +class Inference_AV_WindowSampler: + """ + Window CFG Sampler samples new frames one by one, by inpainting the final frame. + This is basically diffusion forcing. + + :param n_steps: Number of diffusion steps for each frame (diffusoin steps) + :param cfg_scale: CFG scale for each frame + :param window_length: Number of frames to use for each frame generation step + :param num_frames: Number of new frames to sample + :param noise_prev: Noise previous frame + :param only_return_generated: Whether to only return the generated frames + """ + def __init__(self, n_steps = 20, cfg_scale = 1.3, window_length = 60, num_frames = 60, noise_prev = 0.2, only_return_generated = False): + self.n_steps = n_steps + self.cfg_scale = cfg_scale + self.window_length = window_length + self.num_frames = num_frames + self.noise_prev = noise_prev + self.only_return_generated = only_return_generated + + @torch.no_grad() + def __call__(self, model, dummy_batch, audio, mouse, btn, decode_fn = None, audio_decode_fn = None, image_scale = 1, audio_scale = 1): + # dummy_batch is [b,n,c,h,w] + # audio is [b,n,c] and should be treated same as video (it'e being generated) + # mouse is [b,n,2] + # btn is [b,n,n_button] + + # output will be [b,n+self.num_frames,c,h,w] + + sampling_steps = self.n_steps + num_frames = self.num_frames + + dt = 1. / sampling_steps + + clean_history = dummy_batch.clone() + clean_audio_history = audio.clone() + + assert mouse.shape[1] == num_frames + self.window_length + assert btn.shape[1] == num_frames + self.window_length + + extended_mouse, extended_btn = mouse, btn + + def step_history(): + # Video history + new_history = clean_history.clone()[:,-self.window_length:] + b,n,c,h,w = new_history.shape + new_history[:,:-1] = zlerp(new_history[:,1:],self.noise_prev) + new_history[:,-1] = torch.randn_like(new_history[:,0]) + + # Audio history + new_audio = clean_audio_history.clone()[:,-self.window_length:] + new_audio[:,:-1] = zlerp(new_audio[:,1:],self.noise_prev) + new_audio[:,-1] = torch.randn_like(new_audio[:,0]) + + return new_history, new_audio + + for frame_idx in tqdm(range(num_frames)): + local_history, local_audio = step_history() + ts_history = torch.ones(local_history.shape[0], local_history.shape[1], device=local_history.device,dtype=local_history.dtype) + ts_history[:,:-1] = self.noise_prev + + mouse = extended_mouse[:,frame_idx:frame_idx+self.window_length] + btn = extended_btn[:,frame_idx:frame_idx+self.window_length] + + mouse_batch = torch.cat([mouse, torch.zeros_like(mouse)], dim=0) + btn_batch = torch.cat([btn, torch.zeros_like(btn)], dim=0) + for _ in range(sampling_steps): + # CFG Branches + x = local_history.clone() + a = local_audio.clone() + ts = ts_history.clone() + + x_batch = torch.cat([x, x], dim=0) + a_batch = torch.cat([a, a], dim=0) + ts_batch = torch.cat([ts, ts], dim=0) + + pred_video_batch, pred_audio_batch = model(x_batch, a_batch, ts_batch, mouse_batch, btn_batch) + + # Split predictions back into conditional and unconditional + cond_pred_video, uncond_pred_video = pred_video_batch.chunk(2) + cond_pred_audio, uncond_pred_audio = pred_audio_batch.chunk(2) + + pred_video = uncond_pred_video + self.cfg_scale * (cond_pred_video - uncond_pred_video) + pred_audio = uncond_pred_audio + self.cfg_scale * (cond_pred_audio - uncond_pred_audio) + + x = x - pred_video*dt + a = a - pred_audio*dt + ts = ts - dt + + local_history[:,-1] = x[:,-1] + local_audio[:,-1] = a[:,-1] + ts_history[:,-1] = ts[:,-1] + + # Frame is entirely cleaned now + new_frame = local_history[:,-1:] + new_audio = local_audio[:,-1:] + clean_history = torch.cat([clean_history, new_frame], dim=1) + clean_audio_history = torch.cat([clean_audio_history, new_audio], dim=1) + + x = clean_history + audio = clean_audio_history + if self.only_return_generated: + x = x[:,-num_frames:] + audio = audio[:,-num_frames:] + extended_mouse = extended_mouse[:,-num_frames:] + extended_btn = extended_btn[:,-num_frames:] + + pixel_latents = x + pixels = None + + audio_latents = audio + audio_wav = None + + if decode_fn is not None: + pixels = decode_fn(pixel_latents * image_scale) + + if audio_decode_fn is not None: + audio_wav = audio_decode_fn(audio_latents * audio_scale) + + return ( + pixel_latents, audio_latents, # NOTE Need this for history + pixels, audio_wav, # NOTE Need this for rendering + extended_mouse, extended_btn, + clean_history, clean_audio_history + ) + + +def test_av_window_sampler(): + sampler = AVWindowSampler() model = lambda x, ts, mouse, btn: x dummy_batch = torch.randn(1, 32, 128, 4, 4) - + audio = torch.randn(1, 32, 128) mouse = torch.zeros(1, 32, 2) btn = torch.zeros(1, 32, 11) - x = sampler(model, dummy_batch, mouse, btn) - print(x.shape) if __name__ == "__main__": - test_window_cfg_sampler() + test_av_window_sampler() diff --git a/owl_wms/utils/__init__.py b/owl_wms/utils/__init__.py index 312b4a21..a88c23ed 100644 --- a/owl_wms/utils/__init__.py +++ b/owl_wms/utils/__init__.py @@ -96,7 +96,7 @@ def batch_permute_to_length(mouse, button, length): # Calculate how many times we need to double n to exceed length n = mouse.shape[1] factor = 0 - doubled_length = n + doubled_length = mouse.shape[1] while doubled_length < length: factor += 1 doubled_length *= 2 diff --git a/owl_wms/utils/owl_vae_bridge.py b/owl_wms/utils/owl_vae_bridge.py index 79a2b4e2..4c341a6c 100644 --- a/owl_wms/utils/owl_vae_bridge.py +++ b/owl_wms/utils/owl_vae_bridge.py @@ -69,25 +69,6 @@ def decode(x): return decode @torch.no_grad() -<<<<<<< HEAD -def make_batched_encode_fn(encoder, batch_size=8): - def encode(x): - # x is [b,n,c,h,w] RGB frames - b,n,c,h,w = x.shape - x = x.view(b*n,c,h,w).contiguous() - - batches = x.split(batch_size) - batch_out = [] - for batch in batches: - batch_out.append(encoder(batch).bfloat16()) - - x = torch.cat(batch_out) # [b*n,latent_c,latent_h,latent_w] - _,latent_c,latent_h,latent_w = x.shape - x = x.view(b,n,latent_c,latent_h,latent_w).contiguous() - - return x - return encode -======= def make_batched_audio_decode_fn(decoder, batch_size = 8): def decode(x): # x is [b,n,c] audio samples @@ -104,4 +85,3 @@ def decode(x): return x return decode ->>>>>>> causvid diff --git a/webapp/checkpoints/configs/ae.yml b/webapp/checkpoints/configs/ae.yml new file mode 100644 index 00000000..570f334b --- /dev/null +++ b/webapp/checkpoints/configs/ae.yml @@ -0,0 +1,53 @@ +# Config for a simple 256 -> 16 autoencoder +model: + model_id: dcae + sample_size: [360,640] + channels: 3 + latent_size: 4 + latent_channels: 128 + + noise_decoder_inputs: 0.0 + ch_0: 128 + ch_max: 1024 + + encoder_blocks_per_stage: [3, 3, 3, 3, 3, 3, 3, 3] + decoder_blocks_per_stage: [3, 3, 3, 3, 3, 3, 3, 3] + + checkpoint_grads: true + +train: + trainer_id: rec + data_id: s3_cod + target_batch_size: 128 + batch_size: 16 + + epochs: 200 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + weight_decay: 1.0e-4 + betas: [0.9, 0.95] + eps: 1.0e-15 + + lpips_type: convnext + loss_weights: + latent_reg: 1.0e-6 + lpips: 10.0 + se_reg: 0.0 + + scheduler: LinearWarmup + scheduler_kwargs: + warmup_steps: 3000 + min_lr: 5.0e-6 + + checkpoint_dir: webapp/checkpoints/models/cod_128x_30k_ema + resume_ckpt: null + + sample_interval: 1000 + save_interval: 5000 + +wandb: + name: ${env:WANDB_USER_NAME} + project: new_vaes + run_name: 128x_cod \ No newline at end of file diff --git a/webapp/checkpoints/configs/audio.yml b/webapp/checkpoints/configs/audio.yml new file mode 100644 index 00000000..1fde568d --- /dev/null +++ b/webapp/checkpoints/configs/audio.yml @@ -0,0 +1,57 @@ +model: + model_id: audio_ae + + channels: 2 + latent_channels: 64 + ch_0: 128 + ch_max: 512 + + strides: [3, 5, 7, 7, 1] + + eq: true + checkpoint_grads: true + +train: + trainer_id: audio_rec + data_id: local_cod_audio + data_kwargs: + window_length: 88200 + root: "../cod_download/raw" + + target_batch_size: 128 + batch_size: 16 + epochs: 100 + + opt: AdamW + opt_kwargs: + lr: 1.0e-4 + eps: 1.0e-15 + betas: [0.9, 0.95] + weight_decay: 1.0e-4 + + loss_weights: + recon: 2.5 + stft: 1.5 + kl: 1.0e-5 + lr_ms_ratio: 0.5 + hubert: 0.0 + crt: 4.0 + + sample_rate: 44100 + n_fft_list: [1024, 2048, 512] + + scheduler: LinearWarmup + scheduler_kwargs: + warmup_steps: 1500 + min_lr: 1.0e-6 + + checkpoint_dir: webapp/checkpoints/models/audio_20k_ema.pt + sample_interval: 500 + save_interval: 5000 + + resume_ckpt: null + +wandb: + name: ${env:WANDB_USER_NAME} + project: owl_audio_vaes + run_name: audio_ae_baseline \ No newline at end of file diff --git a/webapp/checkpoints/configs/av.yml b/webapp/checkpoints/configs/av.yml new file mode 100644 index 00000000..65085dcc --- /dev/null +++ b/webapp/checkpoints/configs/av.yml @@ -0,0 +1,75 @@ +model: + model_id: game_rft_audio + sample_size: 4 + channels: 128 + audio_channels: 64 + + n_layers: 13 + n_heads: 16 + d_model: 1024 + + tokens_per_frame: 17 + n_buttons: 11 + n_mouse_axes: 2 + + cfg_prob: 0.1 + n_frames: 30 + + causal: false + +train: + trainer_id: av + data_id: cod_s3_audio + data_kwargs: + window_length: 30 + bucket_name: cod-data-latent-360x640to4x4 + + target_batch_size: 256 + batch_size: 32 + + epochs: 200 + + opt: Muon + opt_kwargs: + lr: 1.0e-3 + momentum: 0.95 + adamw_lr: 1.0e-4 + adamw_wd: 1.0e-4 + adamw_eps: 1.0e-15 + adamw_betas: [0.9, 0.95] + adamw_keys: [core.proj_in, core.proj_out.proj] + + scheduler: null + + checkpoint_dir: checkpoints/360p + + sample_interval: 1000 + save_interval: 5000 + + sampler_id: av_window + sampler_kwargs: + n_steps: 10 + cfg_scale: 1.3 + window_length: 30 + num_frames: 60 + noise_prev: 0.2 + only_return_generated: false + + n_samples: 8 + + vae_id: null + vae_batch_size: 4 + vae_scale: 0.13 + audio_vae_scale: 0.17 + + vae_cfg_path: webapp/checkpoints/configs/ae.yml + vae_ckpt_path: webapp/checkpoints/models/cod_128x_30k_ema.pt + + audio_vae_id: null + audio_vae_cfg_path: webapp/checkpoints/configs/audio.yml + audio_vae_ckpt_path: webapp/checkpoints/models/cod_audio_20k_ema.pt + +wandb: + name: shahbuland + project: video_models + run_name: av \ No newline at end of file diff --git a/webapp/checkpoints/shortcut.yaml b/webapp/checkpoints/shortcut.yaml deleted file mode 100644 index fd6a1bfb..00000000 --- a/webapp/checkpoints/shortcut.yaml +++ /dev/null @@ -1,67 +0,0 @@ -# Config for a simple 256 -> 16 autoencoder -model: - model_id: game_rft_shortcut - sample_size: 4 - channels: 128 - - n_layers: 17 - n_heads: 16 - d_model: 1024 - - tokens_per_frame: 16 - n_buttons: 11 - n_mouse_axes: 2 - - cfg_prob: 0.1 - n_frames: 120 - - causal: true - -train: - trainer_id: rft - data_id: cod_latent - data_kwargs: - window_length: 60 - root: ../cod_data/BlackOpsColdWar - add_optical_flow: false - - target_batch_size: 320 - batch_size: 40 - - epochs: 200 - - opt: Muon - opt_kwargs: - lr: 1.0e-3 - momentum: 0.95 - adamw_lr: 1.0e-4 - adamw_wd: 1.0e-4 - adamw_eps: 1.0e-15 - adamw_betas: [0.9, 0.95] - adamw_keys: [core.proj_in, core.proj_out.proj] - - scheduler: null - - checkpoint_dir: webapp/checkpoints/shortcut - resume_ckpt: webapp/checkpoints/shortcut/step_165000.pt - - sample_interval: 1000 - save_interval: 5000 - - sampler_id: window - sampler_kwargs: - n_steps: 32 - cfg_scale: 1.3 - window_length: 60 - num_frames: 120 - noise_prev: 0.2 - only_return_generated: true - - vae_batch_size: 16 - vae_scale: 2.17 - n_samples: 8 - -wandb: - name: shahbuland - project: video_models - run_name: v2 \ No newline at end of file diff --git a/webapp/config.yaml b/webapp/config.yaml deleted file mode 100644 index a5dd154a..00000000 --- a/webapp/config.yaml +++ /dev/null @@ -1,25 +0,0 @@ -model_checkpoint_path: "webapp/checkpoints/shortcut/step_165000.pt" -run_config_path: "webapp/checkpoints/shortcut.yaml" -device: "cuda" - -stream_config: - fps: 20 - frames_per_batch: 60 - window_length: 60 - device: "cuda" - n_buttons: 11 - n_mouse_axes: 2 - mouse_range: [-1.0, 1.0] - action_margin_px_height: 150 - -sampling_config: - sampling_steps: 20 - vae_scale: 1.0 - cfg_scale: 1.3 - window_length: 60 - num_frames: 1 - noise_prev: 0.25 - -run_config: null # loaded at runtime from model_config_path, and used to access model and train config - - diff --git a/webapp/server.py b/webapp/server.py index 6cd04212..cf0c6ba1 100644 --- a/webapp/server.py +++ b/webapp/server.py @@ -1,39 +1,30 @@ import os -from torch import nn +from dotenv import load_dotenv from contextlib import asynccontextmanager from fastapi import FastAPI, WebSocket from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse -from webapp.utils.models import load_models from webapp.streaming import StreamingFrameGenerator from webapp.user_session import UserGameSession from webapp.utils.configs import WebappConfig +load_dotenv() DEBUG = True # -- lifespan -encoder: nn.Module = None -decoder: nn.Module = None config: WebappConfig = None -webapp_config_path = "./configs/webapp/config.yaml" ; assert os.path.exists(webapp_config_path) +webapp_config_path = "./webapp/webapp_config.yaml" ; assert os.path.exists(webapp_config_path) @asynccontextmanager async def lifespan(app: FastAPI): - global encoder, decoder, config, DEBUG + global config, DEBUG config = WebappConfig.from_yaml(webapp_config_path) - if not DEBUG: - encoder, decoder, _ = load_models( - checkpoint_path=config.model_checkpoint_path, - config_path=config.run_config_path, - device=config.device, verbose=True, - ) - yield - encoder, decoder, config = None, None, None + config = None def run(): @@ -51,11 +42,8 @@ async def websocket_endpoint(websocket: WebSocket): await websocket.accept() # Create streaming session for this user - frame_generator = StreamingFrameGenerator(encoder, decoder, - streaming_config=config.stream_config, - model_config=config.run_config.model, - train_config=config.run_config.train, - sampling_config=config.sampling_config, + frame_generator = StreamingFrameGenerator(streaming_config=config.stream_config, + run_config=config.run_config, debug=DEBUG) session = UserGameSession(frame_generator) await session.run_session(websocket) diff --git a/webapp/streaming.py b/webapp/streaming.py index 9e1cbb7e..91cb362b 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -1,16 +1,12 @@ -import math import time import torch import asyncio -import numpy as np from torch import nn -from typing import Optional -from functools import cached_property -from webapp.utils.create_samplers import create_sampler -from webapp.utils.configs import SamplingConfig, StreamingConfig -from owl_wms.configs import TrainingConfig, TransformerConfig as ModelConfig -from webapp.utils.visualize_overlay_actions import _draw_video as _draw_action_overlays +from webapp.utils.configs import StreamingConfig +from webapp.utils.av_window_inference_pipeline import AV_WindowInferencePipeline +from owl_wms.configs import Config as RunConfig + class FrameBuffer: """ @@ -54,111 +50,39 @@ async def get_next_frames(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tenso class StreamingFrameGenerator: """Wraps WindowCFGSampler to generate frames.""" - def __init__(self, encoder: nn.Module, decoder: nn.Module, - streaming_config: StreamingConfig, - model_config: ModelConfig, - train_config: TrainingConfig, - sampling_config: SamplingConfig, - debug: bool = False): + def __init__(self, + streaming_config: StreamingConfig, + run_config: RunConfig, + debug: bool = False): + + self.run_config = run_config self.streaming_config = streaming_config - self.model_config = model_config - self.train_config = train_config - self.sampling_config = sampling_config self.debug = debug - self.encoder = encoder - self.decoder = decoder - # - # Create WindowCFGSampler for frame generation - self.sample_window_fn = create_sampler('window', encoder, decoder, - batch_size=1, - n_steps=self.sampling_config.sampling_steps, - vae_scale=self.sampling_config.vae_scale, - cfg_scale=self.sampling_config.cfg_scale, - window_length=self.sampling_config.window_length, - num_frames=self.sampling_config.num_frames, - noise_prev=self.sampling_config.noise_prev, - only_return_generated=True) - # Initialize frame history as empty tensor - self.latent_history: Optional[torch.Tensor] = None - self.mouse_history: Optional[torch.Tensor] = None - self.button_history: Optional[torch.Tensor] = None - - def add_to_history(self, frame_batch: torch.Tensor, mouse_batch: torch.Tensor, button_batch: torch.Tensor): - if self.latent_history is None: - self.latent_history = frame_batch - self.mouse_history = mouse_batch - self.button_history = button_batch - return - - self.latent_history = torch.cat([self.latent_history, frame_batch], dim=0) - self.mouse_history = torch.cat([self.mouse_history, mouse_batch], dim=0) - self.button_history = torch.cat([self.button_history, button_batch], dim=0) - # cap this at around 60 frames - if self.latent_history.shape[0] > self.streaming_config.window_length: - self.latent_history = self.latent_history[-self.streaming_config.window_length:] - self.mouse_history = self.mouse_history [-self.streaming_config.window_length:] - self.button_history = self.button_history[-self.streaming_config.window_length:] - - def get_latent_history_batch(self) -> torch.Tensor: - if self.latent_history.numel() == 0: - return self.dummy_batch - - return self.latent_history.unsqueeze(0) - - @cached_property - def dummy_batch(self) -> torch.Tensor: - """Dummy autoencoder latents for the sampler to initialize shapes.""" - tokens_h = tokens_w = int(math.sqrt(self.model_config.tokens_per_frame)) - dummy_frames = torch.randn( - 1, self.streaming_config.window_length, - self.model_config.channels, tokens_h, tokens_w, - device=self.streaming_config.device, dtype=torch.bfloat16) - return dummy_frames - + self.av_window_inference_pipeline = AV_WindowInferencePipeline( + config = self.run_config, + video_latent_history = self.streaming_config.video_latent_history, + audio_latent_history = self.streaming_config.audio_latent_history, + mouse_history = self.streaming_config.mouse_history, + button_history = self.streaming_config.button_history, + return_only_generated = True, + compile = True + ) + - def overlay_actions(self, - video: torch.Tensor, - mouse: torch.Tensor, button: torch.Tensor, - action_margin_px_height: int = 150) -> torch.Tensor: - num_frames, channels, height, width = video.shape - action_height = height + action_margin_px_height - action_width = width - action_video = torch.zeros((num_frames, action_height, action_width, channels), # [n h w c] - device=self.streaming_config.device, dtype=torch.bfloat16) - # Copy video into top portion of action video - action_video[:, :height, :, :] = video.permute(0, 2, 3, 1) # [n c h w] -> [n h w c] - action_video_np: list[np.ndarray] = _draw_action_overlays(action_video, button, mouse) - action_video = [torch.from_numpy(frame) for frame in action_video_np] - action_video = torch.stack(action_video).permute(0, 3, 1, 2) # [n h w c] -> [n c h w] - return action_video - - def create_overlay_only(self, - video: torch.Tensor, - mouse: torch.Tensor, button: torch.Tensor, - action_margin_px_height: int = 150) -> torch.Tensor: - num_frames, channels, height, width = video.shape - # Create overlay-only frames with just the action margin height - action_video = torch.zeros((num_frames, action_margin_px_height, width, channels), # [n h w c] - device=self.streaming_config.device, dtype=torch.bfloat16) - action_video_np: list[np.ndarray] = _draw_action_overlays(action_video, button, mouse) - action_video = [torch.from_numpy(frame) for frame in action_video_np] - action_video = torch.stack(action_video).permute(0, 3, 1, 2) # [n h w c] -> [n c h w] - return action_video - - async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor, - action_margin_px_height: int = 150) -> tuple[torch.Tensor, torch.Tensor]: + + async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Generate window_length frames, return separate video and overlay frames for streaming. Args: - mouse: [window_length, 2] - button: [window_length, n_buttons] + mouse: [window_length, 2] , user input mouse + button: [window_length, n_buttons] , user input button Returns: - tuple: (video_frames, overlay_frames) + tuple: (video_frames, audio_frames) video_frames: [frames_per_batch, 3, 256, 256] - pure video frames - overlay_frames: [frames_per_batch, 3, action_margin_px_height, 256] - action overlay frames + audio_frames: [frames_per_batch, 2] - audio frames """ mouse = mouse.to(self.streaming_config.device) button = button.to(self.streaming_config.device) @@ -174,6 +98,7 @@ async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor, (col_indices - 128) / 127.0 # Second half: 0 to 1 ).view(1, 1, 1, -1) # [1, 1, 1, 256] full_frames = gradient.expand(num_frames, 3, 256, 256).to(torch.bfloat16) + audio_frames = torch.randn(num_frames, 2, device=self.streaming_config.device, dtype=torch.bfloat16) # to between 0 and 1 full_frames = (full_frames - full_frames.min()) / (full_frames.max() - full_frames.min()) # between 0 and 255 @@ -181,24 +106,22 @@ async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor, else: device_type = 'cuda' if self.streaming_config.device.type == 'cuda' else 'cpu' with torch.no_grad(), torch.autocast(device_type=device_type, dtype=torch.bfloat16): - latents, full_frames = self.sample_window_fn(dummy_batch=self.get_latent_history_batch(), - mouse=mouse.float().unsqueeze(0), - btn=button.float().unsqueeze(0)) # [1, window_length, 3, 256, 256] - # remove batch dimension, it returns the frames we generated, since we pass in only_return_generated=True - latents = latents [0, ::] - full_frames = full_frames [0, ::] + full_frames, audio_frames = self.av_window_inference_pipeline( + mouse=mouse.float().unsqueeze(0), # NOTE Need batch dimension + button=button.float().unsqueeze(0) # NOTE Need batch dimension + ) # [1, window_length, 3, 256, 256], [1, window_length, 2] + + full_frames = full_frames [0, ::] # remove batch dimension + audio_frames = audio_frames [0, ::] # remove batch dimension # then, convert the frames to a pixel-range of [0-255] from [-1,1] - full_frames = (full_frames + 1) / 2 - full_frames = (full_frames * 255).to(torch.uint8) - self.add_to_history(latents, mouse, button) + full_frames = (full_frames + 1) / 2 + full_frames = (full_frames * 255).to(torch.uint8) # [window_length, 3, 256, 256] + + return full_frames, audio_frames # [window_length, 3, 256, 256], [window_length, 2] - # Create overlay frames separately - overlay_frames = self.create_overlay_only(full_frames, mouse, button, action_margin_px_height) - - return full_frames, overlay_frames # [window_length, 3, 256, 256], [window_length, 3, action_margin_px_height, 256] def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): - del self.encoder, self.decoder ; torch.cuda.empty_cache() + del torch.cuda.empty_cache() diff --git a/webapp/user_session.py b/webapp/user_session.py index c96f709a..ed44d4a1 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -56,7 +56,7 @@ async def _frame_generation_loop(self): mouse, button = await self.action_collector.collect_actions() # Generate Y frames from X actions by taking the X[-1]'th action. Typically, X >> Y, because they are sampled at uncapped FPS from the UI, # whereas Y frames are sampled from the model one at a time. - video_frames, overlay_frames = await self.frame_generator.generate_frames(mouse, button) + video_frames, audio_frames = await self.frame_generator.generate_frames(mouse, button) # Queue frames for streaming at a capped FPS. If model predictions speed up or slow down, it won't cause any dilation of frames being displayed. # However, if the model predictions are too slow, the frames will be displayed at a lower FPS than the capped FPS. await self.frame_buffer.queue_frames(video_frames, mouse, button) diff --git a/webapp/utils/av_window_inference_pipeline.py b/webapp/utils/av_window_inference_pipeline.py new file mode 100644 index 00000000..f106ca22 --- /dev/null +++ b/webapp/utils/av_window_inference_pipeline.py @@ -0,0 +1,127 @@ +import os +import torch +from torch.nn import Module +from owl_wms.models import get_model_cls +from owl_wms.utils.owl_vae_bridge import get_decoder_only +from owl_wms.configs import Config as RunConfig +from owl_wms.models.gamerft_audio import GameRFTCore + +def zlerp(x, alpha): + return x * (1. - alpha) + alpha * torch.randn_like(x) + +class AV_WindowInferencePipeline: + def __init__(self, + config: RunConfig, + video_latent_history: torch.Tensor, + audio_latent_history: torch.Tensor, + mouse_history: torch.Tensor, + button_history: torch.Tensor, + ckpt_path: str = "av_dfot_35k_ema_200m.pt", + alpha: float = 0.2, + cfg_scale: float = 1.3, + sampling_steps: int = 10, + audio_f: int = 735, + return_only_generated: bool = True, + compile: bool = True): + + self.return_only_generated = return_only_generated + self.config = config + + + self.model: GameRFTCore = get_model_cls(self.config.model.model_id)(self.config.model).core + state_dict = torch .load(ckpt_path, map_location="cpu") + self.model .load_state_dict(state_dict) + self.model .eval() + + self.frame_decoder: Module = get_decoder_only( + None, + self.config.train.vae_cfg_path, + self.config.train.vae_ckpt_path + ) + self.frame_decoder.eval() + + self.audio_decoder: Module = get_decoder_only( + None, + self.config.train.audio_vae_cfg_path, + self.config.train.audio_vae_ckpt_path + ) + self.audio_decoder.eval() + + self.frame_scale = self.config.train.vae_scale + self.audio_scale = self.config.train.audio_vae_scale + + self.history_buffer = video_latent_history / self.frame_scale + self.audio_buffer = audio_latent_history / self.audio_scale + self.mouse_buffer = mouse_history + self.button_buffer = button_history + + self.alpha = alpha + self.cfg_scale = cfg_scale + self.sampling_steps = sampling_steps + self.audio_f = audio_f + + if compile: + print(f'Compiling models...') + torch.compile(self.model) + torch.compile(self.frame_decoder) + torch.compile(self.audio_decoder) + + + @torch.no_grad() + def __call__(self, + user_input_mouse: torch.Tensor, + user_input_button: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + + # [2,] float and [11,] bool + print(f'user_input_mouse.shape: {user_input_mouse.shape}') + print(f'user_input_button.shape: {user_input_button.shape}') + + noised_history = zlerp(self.history_buffer[:,1:], self.alpha) + noised_audio = zlerp(self.audio_buffer[:,1:], self.alpha) + + noised_history = torch.cat([noised_history, torch.randn_like(noised_history[:,0:1])], dim = 1) + noised_audio = torch.cat([noised_audio, torch.randn_like(noised_audio[:,0:1])], dim = 1) + + user_input_mouse = user_input_mouse[None,None,:] + user_input_button = user_input_button [None,None,:] + + self.mouse_buffer = torch.cat([self.mouse_buffer[:,1:],user_input_mouse],dim=1) + self.button_buffer = torch.cat([self.button_buffer[:,1:],user_input_button],dim=1) + + dt = 1. / self.sampling_steps + + x = noised_history + a = noised_audio + ts = torch.ones_like(noised_history[:,:,0,0,0]) + ts[:,:-1] = self.alpha + + mouse_batch = torch.cat([self.mouse_buffer, torch.zeros_like(user_input_mouse)], dim=0) + btn_batch = torch.cat([self.button_buffer, torch.zeros_like(user_input_button)], dim=0) + + for _ in range(self.sampling_steps): + x_batch = torch.cat([x, x], dim=0) + a_batch = torch.cat([a, a], dim=0) + ts_batch = torch.cat([ts, ts], dim=0) + + video_rollout, audio_rollout = self.model(x_batch,a_batch,ts_batch,mouse_batch,btn_batch) + + cond_pred_video, uncond_pred_video = video_rollout.chunk(2) + cond_pred_audio, uncond_pred_audio = audio_rollout.chunk(2) + + pred_video = uncond_pred_video + self.cfg_scale * (cond_pred_video - uncond_pred_video) + pred_audio = uncond_pred_audio + self.cfg_scale * (cond_pred_audio - uncond_pred_audio) + + x[:,-1] = x[:,-1] - dt * pred_video[:,-1] + a[:,-1] = a[:,-1] - dt * pred_audio[:,-1] + ts[:,-1] = ts[:,-1] - dt + + new_frame = x [:,-1:] # [1,1,c,h,w] + new_audio = audio[:,-1:] # [1,1,c] + + self.history_buffer = torch.cat([self.history_buffer[:,1:], new_frame], dim=1) + self.audio_buffer = torch.cat([self.audio_buffer[:,1:], new_audio], dim=1) + + frame = self.frame_decoder(new_frame[0] * self.frame_scale).squeeze() # [c,h,w] + audio = self.audio_decoder(self.audio_buffer * self.audio_scale).squeeze()[-self.audio_f:] # [735,2] + + return frame, audio diff --git a/webapp/utils/configs.py b/webapp/utils/configs.py index 4cda5c58..40877303 100644 --- a/webapp/utils/configs.py +++ b/webapp/utils/configs.py @@ -2,10 +2,10 @@ import os import yaml -from typing import Optional +import torch from dataclasses import dataclass -from owl_wms.configs import Config as RunConfig, TransformerConfig as ModelConfig, TrainingConfig +from owl_wms.configs import Config as RunConfig @dataclass class WebappConfig: @@ -41,18 +41,38 @@ class SamplingConfig: @dataclass class StreamingConfig: fps: int = 20 - frames_per_batch: int = 8 - window_length: int = 60 + frames_per_batch: int = 1 device: str = 'cuda' n_buttons: int = 11 n_mouse_axes: int = 2 mouse_range: tuple[float, float] = (-1.0, 1.0) - action_margin_px_height: int = 150 + video_latent_history_path: os.PathLike = None + audio_latent_history_path: os.PathLike = None + mouse_history_path: os.PathLike = None + button_history_path: os.PathLike = None @property - def frame_interval(self) -> float: - return 1.0 / self.fps + def frame_interval(self) -> float: return 1.0 / self.fps @property - def batch_duration(self) -> float: - return self.frames_per_batch / self.fps + def batch_duration(self) -> float: return self.frames_per_batch / self.fps + + @property + def video_latent_history(self) -> torch.Tensor: + if self.video_latent_history_path is None: raise ValueError("video_latent_history_path is not set") + return torch.load(self.video_latent_history_path) + + @property + def audio_latent_history(self) -> torch.Tensor: + if self.audio_latent_history_path is None: raise ValueError("audio_latent_history_path is not set") + return torch.load(self.audio_latent_history_path) + + @property + def mouse_history(self) -> torch.Tensor: + if self.mouse_history_path is None: raise ValueError("mouse_history_path is not set") + return torch.load(self.mouse_history_path) + + @property + def button_history(self) -> torch.Tensor: + if self.button_history_path is None: raise ValueError("button_history_path is not set") + return torch.load(self.button_history_path) \ No newline at end of file diff --git a/webapp/utils/create_samplers.py b/webapp/utils/create_samplers.py index 2f1bf688..87fbbd8a 100644 --- a/webapp/utils/create_samplers.py +++ b/webapp/utils/create_samplers.py @@ -5,7 +5,8 @@ from owl_wms.sampling.cfg import CFGSampler from owl_wms.sampling.simple import SimpleSampler, InpaintSimpleSampler from owl_wms.sampling.window import WindowCFGSampler -from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn +from owl_wms.sampling.av_window import Inference_AV_WindowSampler +from owl_wms.utils.owl_vae_bridge import make_batched_decode_fn, make_batched_audio_decode_fn SAMPLING_STEPS = 60 SCALE = 2.17 @@ -109,7 +110,40 @@ def _sampler(): ) -# Example usage: +@multimethod +def create_sampler(sampler_id: Literal['av_window'], + encoder, decoder, audio_decoder, + batch_size: int = 8, + n_steps: int = 20, + cfg_scale: float = 1.3, + window_length: int = 60, + num_frames: int = 60, + noise_prev: float = 0.2, + only_return_generated: bool = True, + vae_scale: float = 1.0, + **kwargs) -> Callable: + + @cache + def _sampler(): + return Inference_AV_WindowSampler( + n_steps=n_steps, + cfg_scale=cfg_scale, + window_length=window_length, + num_frames=num_frames, + noise_prev=noise_prev, + only_return_generated=only_return_generated + ) + + # TODO `dummy_batch` is history, `audio` is audio history, mouse is mouse history, btn is button history + return partial( + _sampler().__call__, + decode_fn=make_batched_decode_fn(decoder, batch_size=batch_size), + audio_decode_fn=make_batched_audio_decode_fn(audio_decoder, batch_size=batch_size), + scale=vae_scale, + model=encoder + ) + + if __name__ == "__main__": # Each sampler type can be created with its specific parameters import webapp.utils.models @@ -134,3 +168,11 @@ def _sampler(): window_length=80, num_frames=120, noise_prev=0.3, only_return_generated=True ) + + # AV Window sampler with all its specific params + av_window_sampler = create_sampler( + 'av_window', encoder, decoder, + n_steps=10, cfg_scale=1.3, + window_length=30, num_frames=60, + noise_prev=0.2, only_return_generated=False + ) \ No newline at end of file diff --git a/webapp/webapp_config.yaml b/webapp/webapp_config.yaml new file mode 100644 index 00000000..6502ac80 --- /dev/null +++ b/webapp/webapp_config.yaml @@ -0,0 +1,27 @@ +model_checkpoint_path: "webapp/checkpoints/av_dfot_35k_ema_200m.pt" +run_config_path: "webapp/checkpoints/av.yml" +device: "cuda" + +stream_config: + fps: 20 + frames_per_batch: 1 + device: "cuda" + n_buttons: 11 + n_mouse_axes: 2 + mouse_range: [-1.0, 1.0] + video_latent_history_path: "webapp/static/histories/base/video_latent.pt" + audio_latent_history_path: "webapp/static/histories/base/audio_latent.pt" + mouse_history_path: "webapp/static/histories/base/mouse.pt" + button_history_path: "webapp/static/histories/base/button.pt" + +sampling_config: + sampling_steps: 20 + vae_scale: 1.0 + cfg_scale: 1.3 + window_length: 60 + num_frames: 1 + noise_prev: 0.25 + +run_config: null # loaded at runtime from model_config_path, and used to access model and train config + + From e5ae56d39b68dceadad214bd0d3c3f664b90107f Mon Sep 17 00:00:00 2001 From: Sami Date: Thu, 12 Jun 2025 01:05:25 -0500 Subject: [PATCH 68/71] smapler --- webapp/server.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/webapp/server.py b/webapp/server.py index cf0c6ba1..2f599162 100644 --- a/webapp/server.py +++ b/webapp/server.py @@ -56,6 +56,7 @@ def main(): global DEBUG import argparse import uvicorn + parser = argparse.ArgumentParser() parser.add_argument("--debug", action="store_true", help="Enable debug mode") @@ -63,10 +64,7 @@ def main(): parser.add_argument("--port", type=int, default=8000, help="Port to run the server on") args = parser.parse_args() - # Fix the DEBUG logic - if args.debug and args.no_debug: - raise ValueError("Cannot have both --debug and --no-debug flags") - + if args.debug: DEBUG = True elif args.no_debug: From 3f6b22a68cc431711128969aa05672fb8293d972 Mon Sep 17 00:00:00 2001 From: Sami Date: Thu, 12 Jun 2025 03:24:12 -0500 Subject: [PATCH 69/71] demogit status! --- owl-vaes | 2 +- owl_wms/data/s3_cod_latent_audio.py | 6 + requirements.txt | 6 +- webapp/server.py | 3 +- webapp/static/index.html | 289 +++++++++---------- webapp/streaming.py | 49 ++-- webapp/user_session.py | 56 +++- webapp/utils/av_window_inference_pipeline.py | 62 ++-- webapp/utils/configs.py | 2 +- webapp/webapp_config.yaml | 14 +- 10 files changed, 276 insertions(+), 213 deletions(-) diff --git a/owl-vaes b/owl-vaes index cdde9f3e..f8b08b22 160000 --- a/owl-vaes +++ b/owl-vaes @@ -1 +1 @@ -Subproject commit cdde9f3e93cbae5ed77c99e6d7f29926ba89af42 +Subproject commit f8b08b227fccc4a391b71187eeb530d7c39d4740 diff --git a/owl_wms/data/s3_cod_latent_audio.py b/owl_wms/data/s3_cod_latent_audio.py index 1aeed32c..998df356 100644 --- a/owl_wms/data/s3_cod_latent_audio.py +++ b/owl_wms/data/s3_cod_latent_audio.py @@ -187,6 +187,12 @@ def get_loader(batch_size, **data_kwargs): start = time.time() batch = next(iter(loader)) + video_latent, audio_latent, mouse, button = batch + BASE_IDX=1 + torch.save(video_latent, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/video_latent.pt") + torch.save(audio_latent, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/audio_latent.pt") + torch.save(mouse, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/mouse.pt") + torch.save(button, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/buttons.pt") end = time.time() first_time = end - start diff --git a/requirements.txt b/requirements.txt index 80aae6c3..541bc77d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +--extra-index-url https://download.pytorch.org/whl/cu128 opencv-python wandb einops @@ -5,8 +6,7 @@ rotary-embedding-torch ema-pytorch omegaconf torch -toolz -termcolor +torchvision multimethod diffusers vector-quantize-pytorch @@ -22,3 +22,5 @@ accelerate boto3 python-dotenv wandb[media] +alias-free-torch +taskgroup # only here to backport to 3.10 cause hotel wifi is too slow to remake venv diff --git a/webapp/server.py b/webapp/server.py index 2f599162..d66ed616 100644 --- a/webapp/server.py +++ b/webapp/server.py @@ -1,5 +1,5 @@ import os -from dotenv import load_dotenv +from dotenv import load_dotenv ; load_dotenv() from contextlib import asynccontextmanager from fastapi import FastAPI, WebSocket @@ -10,7 +10,6 @@ from webapp.user_session import UserGameSession from webapp.utils.configs import WebappConfig -load_dotenv() DEBUG = True diff --git a/webapp/static/index.html b/webapp/static/index.html index d718f3fc..4142d931 100644 --- a/webapp/static/index.html +++ b/webapp/static/index.html @@ -42,6 +42,34 @@ .disconnected { background: #f44336; } .connecting { background: #ff9800; } + /* ---------------- Audio Controls ---------------- */ + #audioControls { + margin-top: 10px; + text-align: center; + } + + .audio-btn { + background: var(--key-bg); + border: 2px solid var(--key-border); + color: var(--text); + padding: 8px 16px; + margin: 0 5px; + cursor: pointer; + font-family: monospace; + font-size: 12px; + } + + .audio-btn:hover { + background: var(--brand); + color: #000; + } + + .audio-btn.active { + background: var(--brand); + color: #000; + } + /* ------------------------------------------------- */ + #gameContainer { position: relative; border: 2px solid var(--key-border); @@ -80,24 +108,10 @@ min-width: 400px; } - .keyboard-section { - flex: 0 0 auto; - } + .keyboard-section { flex: 0 0 auto; } + .mouse-section { flex: 0 0 auto; display: flex; flex-direction: column; align-items: center; } - .mouse-section { - flex: 0 0 auto; - display: flex; - flex-direction: column; - align-items: center; - } - - .key-row { - display: flex; - align-items: center; - justify-content: center; - margin-bottom: 10px; - gap: 8px; - } + .key-row { display: flex; align-items: center; justify-content: center; margin-bottom: 10px; gap: 8px; } .key { background: var(--key-bg); @@ -112,42 +126,13 @@ transition: all 0.1s ease; } - .key.active { - background: var(--key-active); - border-color: var(--key-active); - color: #000; - } - - .key.pressed { - transform: scale(0.95); - } - - .key.spacebar { - min-width: 120px; - } + .key.active { background: var(--key-active); border-color: var(--key-active); color: #000; } + .key.pressed { transform: scale(0.95); } + .key.spacebar { min-width: 120px; } + .key.shift { min-width: 60px; } - .key.shift { - min-width: 60px; - } - - .mouse-area { - display: flex; - align-items: center; - gap: 20px; - margin-top: 10px; - } - - .mouse-delta { - position: relative; - width: 80px; - height: 80px; - border: 2px solid var(--key-border); - border-radius: 50%; - background: var(--key-bg); - display: flex; - align-items: center; - justify-content: center; - } + .mouse-area { display: flex; align-items: center; gap: 20px; margin-top: 10px; } + .mouse-delta { position: relative; width: 80px; height: 80px; border: 2px solid var(--key-border); border-radius: 50%; background: var(--key-bg); display: flex; align-items: center; justify-content: center; } .mouse-compass { position: absolute; @@ -159,10 +144,7 @@ transform: translate(-50%, -50%); } - .mouse-buttons { - display: flex; - gap: 8px; - } + .mouse-buttons { display: flex; gap: 8px; } .mouse-btn { background: var(--key-bg); @@ -175,26 +157,11 @@ transition: all 0.1s ease; } - .mouse-btn:first-child { - border-top-left-radius: 4px; - border-bottom-left-radius: 4px; - } - - .mouse-btn:last-child { - border-top-right-radius: 4px; - border-bottom-right-radius: 4px; - margin-left: -2px; /* Overlap borders */ - } - - .mouse-btn.active { - background: var(--key-active); - border-color: var(--key-active); - color: #000; - } + .mouse-btn:first-child { border-top-left-radius: 4px; border-bottom-left-radius: 4px; } + .mouse-btn:last-child { border-top-right-radius: 4px; border-bottom-right-radius: 4px; margin-left: -2px; } - .mouse-btn.pressed { - transform: scale(0.95); - } + .mouse-btn.active { background: var(--key-active); border-color: var(--key-active); color: #000; } + .mouse-btn.pressed { transform: scale(0.95); } p { margin-top: 24px; text-align: center; } @@ -202,6 +169,14 @@

🎮 OWL‑WMS Game Stream

Connecting to server…
+ + +
+ + Volume: + +
+
@@ -224,7 +199,6 @@

🎮 OWL‑WMS Game Stream

SPACE
-
@@ -241,38 +215,91 @@

🎮 OWL‑WMS Game Stream

Move your mouse over the canvas to interact with the game world!

- \ No newline at end of file + diff --git a/webapp/streaming.py b/webapp/streaming.py index 91cb362b..004a897f 100644 --- a/webapp/streaming.py +++ b/webapp/streaming.py @@ -11,29 +11,33 @@ class FrameBuffer: """ Manages frame streaming at precise timing, to adhere to a max FPS. - We need this because sometimes we could generate frames faster than the max FPS, - but, intuitively, the frame outputs will be the same 'set time' apart in the world state that we are approximating. """ def __init__(self, streaming_config: StreamingConfig): self.streaming_config = streaming_config - self.video_frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches - self.buttons_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches - self.mouse_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Buffer 2 batches + self.video_frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) + self.audio_frame_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) # Add audio queue + self.buttons_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) + self.mouse_queue = asyncio.Queue(maxsize=streaming_config.frames_per_batch * 2) self.last_frame_time = 0.0 - async def queue_frames(self, video_frames: torch.Tensor, mouse: torch.Tensor, button: torch.Tensor): + async def queue_frames(self, + video_frames: torch.Tensor, # [t,c,h,w] + audio_frames: torch.Tensor, # [t,?,2] + mouse: torch.Tensor, + button: torch.Tensor): # video_frames shape: [frames_per_batch, channels, height, width] - # overlay_frames shape: [frames_per_batch, channels, overlay_height, width] + # audio_frames shape: [frames_per_batch, 2] num_frames = video_frames.shape[0] for i in range(num_frames): await self.video_frame_queue.put(video_frames[i]) + await self.audio_frame_queue.put(audio_frames[i]) # Queue audio frames await self.buttons_queue.put(button[i]) await self.mouse_queue.put(mouse[i]) - async def get_next_frames(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Get next video and overlay frames for streaming at capped FPS.""" + async def get_next_frames(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Get next video, audio, and input frames for streaming at capped FPS.""" now = time.time() time_since_last = now - self.last_frame_time time_to_wait = max(0, self.streaming_config.frame_interval - time_since_last) @@ -42,10 +46,12 @@ async def get_next_frames(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tenso await asyncio.sleep(time_to_wait) video_frame = await self.video_frame_queue.get() + audio_frame = await self.audio_frame_queue.get() button = await self.buttons_queue.get() mouse = await self.mouse_queue.get() self.last_frame_time = time.time() - return video_frame, button, mouse + return video_frame, audio_frame, button, mouse + class StreamingFrameGenerator: """Wraps WindowCFGSampler to generate frames.""" @@ -61,6 +67,7 @@ def __init__(self, self.av_window_inference_pipeline = AV_WindowInferencePipeline( config = self.run_config, + ckpt_path = self.streaming_config.model_checkpoint_path, video_latent_history = self.streaming_config.video_latent_history, audio_latent_history = self.streaming_config.audio_latent_history, mouse_history = self.streaming_config.mouse_history, @@ -70,7 +77,6 @@ def __init__(self, ) - async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Generate window_length frames, return separate video and overlay frames for streaming. @@ -104,24 +110,23 @@ async def generate_frames(self, mouse: torch.Tensor, button: torch.Tensor) -> tu # between 0 and 255 full_frames = (full_frames * 255).to(torch.uint8) else: - device_type = 'cuda' if self.streaming_config.device.type == 'cuda' else 'cpu' + device_type = 'cuda' if self.streaming_config.device == 'cuda' else 'cpu' with torch.no_grad(), torch.autocast(device_type=device_type, dtype=torch.bfloat16): full_frames, audio_frames = self.av_window_inference_pipeline( - mouse=mouse.float().unsqueeze(0), # NOTE Need batch dimension - button=button.float().unsqueeze(0) # NOTE Need batch dimension - ) # [1, window_length, 3, 256, 256], [1, window_length, 2] + user_input_mouse=mouse.float().unsqueeze(0), # NOTE Need batch dimension + user_input_button=button.float().unsqueeze(0) # NOTE Need batch dimension + ) # [3, 256, 256], [f, 2] - full_frames = full_frames [0, ::] # remove batch dimension - audio_frames = audio_frames [0, ::] # remove batch dimension - # then, convert the frames to a pixel-range of [0-255] from [-1,1] - full_frames = (full_frames + 1) / 2 - full_frames = (full_frames * 255).to(torch.uint8) # [window_length, 3, 256, 256] + # convert the frames to a pixel-range of [0-255] from [-1,1] + full_frames = (full_frames + 1) / 2 # NOTE for some reason this is slightly off of [-1, 1] + full_frames = (full_frames * 255).to(torch.uint8) # [3, 256, 256] + full_frames = torch.clip(full_frames, 0, 255) # bandaid for the [-1,1] - return full_frames, audio_frames # [window_length, 3, 256, 256], [window_length, 2] + return full_frames.unsqueeze(0), audio_frames.unsqueeze(0) # [t, 3, 256, 256], [t, f, 2] where t = 1 for frame-by-frame rollouts def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): - del torch.cuda.empty_cache() + torch.cuda.empty_cache() diff --git a/webapp/user_session.py b/webapp/user_session.py index ed44d4a1..ef4a1f22 100644 --- a/webapp/user_session.py +++ b/webapp/user_session.py @@ -1,3 +1,5 @@ +import io +import wave import cv2 import time import json @@ -10,7 +12,7 @@ from fastapi import WebSocket from webapp.action_converter import ActionCollector from webapp.streaming import StreamingFrameGenerator, FrameBuffer - +from taskgroup import TaskGroup class UserGameSession: """ @@ -27,7 +29,7 @@ async def run_session(self, websocket: WebSocket): print(termcolor.colored(f"Generating {self.frame_generator.streaming_config.frames_per_batch} frames per batch", "green")) print(termcolor.colored(f"Batch duration: {self.frame_generator.streaming_config.batch_duration:.3f}s", "green")) - async with asyncio.TaskGroup() as tg: + async with TaskGroup() as tg: tg.create_task(self._action_input_loop (websocket)) tg.create_task(self._frame_generation_loop ()) tg.create_task(self._frame_display_loop (websocket)) @@ -56,10 +58,10 @@ async def _frame_generation_loop(self): mouse, button = await self.action_collector.collect_actions() # Generate Y frames from X actions by taking the X[-1]'th action. Typically, X >> Y, because they are sampled at uncapped FPS from the UI, # whereas Y frames are sampled from the model one at a time. - video_frames, audio_frames = await self.frame_generator.generate_frames(mouse, button) + video_frames, audio_frames = await self.frame_generator.generate_frames(mouse, button) # TODO What are dims of mouse, button here? # Queue frames for streaming at a capped FPS. If model predictions speed up or slow down, it won't cause any dilation of frames being displayed. # However, if the model predictions are too slow, the frames will be displayed at a lower FPS than the capped FPS. - await self.frame_buffer.queue_frames(video_frames, mouse, button) + await self.frame_buffer.queue_frames(video_frames, audio_frames, mouse, button) # TODO What should we pass into here? mouse should have 1 frame, button should have 1 frame? except Exception as e: import traceback print(termcolor.colored(f"Error in frame generation: {e} :\n {traceback.format_exc()}", "red")) @@ -73,8 +75,8 @@ async def _frame_display_loop(self, websocket: WebSocket): print(termcolor.colored("🔌 WebSocket no longer connected - stopping frame stream", "yellow")) break - video_frame, button, mouse = await self.frame_buffer.get_next_frames() - await self._send_frames_to_client(websocket, video_frame, button, mouse) + video_frame, audio_frames, button, mouse = await self.frame_buffer.get_next_frames() + await self._send_frames_to_client(websocket, video_frame, audio_frames, button, mouse) except Exception as e: # Check if this is a WebSocket disconnect if ("websocket.close" in str(e) or @@ -91,6 +93,7 @@ async def _frame_display_loop(self, websocket: WebSocket): async def _send_frames_to_client(self, websocket: WebSocket, video_frame: torch.Tensor, + audio_frames: torch.Tensor, # Add audio_frame parameter button: torch.Tensor, mouse: torch.Tensor): try: @@ -98,10 +101,8 @@ async def _send_frames_to_client(self, if websocket.client_state.name != 'CONNECTED': raise RuntimeError("WebSocket is not connected") - # TODO Do this more intelligently. I'm sure there's better tech to stream video to a UI. - # Convert video frame to base64 JPEG - video_frame_np = video_frame.float().cpu().numpy().transpose(1, 2, 0) # CHW -> HWC - # Normalize frame data to 0-255 range + # Convert video frame to base64 JPEG (existing code) + video_frame_np = video_frame.float().cpu().numpy().transpose(1, 2, 0) if video_frame_np.max() <= 1.0: video_frame_np = (video_frame_np * 255).clip(0, 255).astype(np.uint8) else: @@ -110,15 +111,48 @@ async def _send_frames_to_client(self, _, video_buffer = cv2.imencode('.jpg', video_frame_np) video_base64 = base64.b64encode(video_buffer).decode('utf-8') + # Convert audio frame to base64 WAV + audio_base64 = self._encode_audio_to_wav(audio_frames) await websocket.send_json({ "type": "frame", "video_data": video_base64, + "audio_data": audio_base64, # Add audio data "button_data": self.action_collector.converter.buttons_to_dict(button), "mouse_data": self.action_collector.converter.mouse_to_dict(mouse), "timestamp": time.time() }) except Exception as e: - # Re-raise to be caught by the display loop raise e + + def _encode_audio_to_wav(self, audio_frames: torch.Tensor, sample_rate: int = 44100) -> str: + """ + Convert audio tensor to base64 encoded WAV data. + + Args: + audio_frames: [window_length, 2] tensor representing stereo audio for one frame + sample_rate: Audio sample rate (default 44100 Hz) + """ + # Convert to numpy and ensure it's in the right format + audio_np = audio_frames.float().cpu().numpy() + + # Normalize audio to [-1, 1] range if needed + if audio_np.dtype == torch.bfloat16 or audio_np.max() > 1.0: + audio_np = np.clip(audio_np, -1.0, 1.0) + + # Convert to 16-bit PCM (standard for WAV) + audio_int16 = (audio_np * 32767).astype(np.int16) + + # Create WAV data in memory + wav_buffer = io.BytesIO() + with wave.open(wav_buffer, 'wb') as wav_file: + wav_file.setnchannels(2) # Stereo + wav_file.setsampwidth(2) # 2 bytes per sample (16-bit) + wav_file.setframerate(sample_rate) + wav_file.writeframes(audio_int16.tobytes()) + + # Get WAV data and encode as base64 + wav_data = wav_buffer.getvalue() + return base64.b64encode(wav_data).decode('utf-8') + # TODO Audio diff --git a/webapp/utils/av_window_inference_pipeline.py b/webapp/utils/av_window_inference_pipeline.py index f106ca22..cb3fd3fe 100644 --- a/webapp/utils/av_window_inference_pipeline.py +++ b/webapp/utils/av_window_inference_pipeline.py @@ -1,4 +1,5 @@ import os +import time import torch from torch.nn import Module from owl_wms.models import get_model_cls @@ -9,6 +10,19 @@ def zlerp(x, alpha): return x * (1. - alpha) + alpha * torch.randn_like(x) + +def print_duration(func): + """Decorator that logs the input and output of a function.""" + def wrapper(*args, **kwargs): + start_time = time.time() + result = func(*args, **kwargs) + end_time = time.time() + execution_time = end_time - start_time + print(f"{func.__name__} took {execution_time:.3f} seconds to execute, which would yield FPS of {1/execution_time:.3f}") + return result + return wrapper + + class AV_WindowInferencePipeline: def __init__(self, config: RunConfig, @@ -22,16 +36,18 @@ def __init__(self, sampling_steps: int = 10, audio_f: int = 735, return_only_generated: bool = True, - compile: bool = True): + compile: bool = True, + device: str = 'cuda'): self.return_only_generated = return_only_generated self.config = config - + self.device = device self.model: GameRFTCore = get_model_cls(self.config.model.model_id)(self.config.model).core state_dict = torch .load(ckpt_path, map_location="cpu") self.model .load_state_dict(state_dict) self.model .eval() + self.model .to(self.device) self.frame_decoder: Module = get_decoder_only( None, @@ -39,6 +55,7 @@ def __init__(self, self.config.train.vae_ckpt_path ) self.frame_decoder.eval() + self.frame_decoder.to(self.device) self.audio_decoder: Module = get_decoder_only( None, @@ -46,19 +63,20 @@ def __init__(self, self.config.train.audio_vae_ckpt_path ) self.audio_decoder.eval() + self.audio_decoder.to(self.device) self.frame_scale = self.config.train.vae_scale self.audio_scale = self.config.train.audio_vae_scale - self.history_buffer = video_latent_history / self.frame_scale - self.audio_buffer = audio_latent_history / self.audio_scale - self.mouse_buffer = mouse_history - self.button_buffer = button_history + self.history_buffer = (video_latent_history / self.frame_scale).to(self.device) + self.audio_buffer = (audio_latent_history / self.audio_scale).to(self.device) + self.mouse_buffer = mouse_history.to(self.device) + self.button_buffer = button_history.to(self.device) self.alpha = alpha self.cfg_scale = cfg_scale self.sampling_steps = sampling_steps - self.audio_f = audio_f + self.audio_f = audio_f if compile: print(f'Compiling models...') @@ -66,15 +84,12 @@ def __init__(self, torch.compile(self.frame_decoder) torch.compile(self.audio_decoder) - + @print_duration @torch.no_grad() def __call__(self, - user_input_mouse: torch.Tensor, - user_input_button: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - - # [2,] float and [11,] bool - print(f'user_input_mouse.shape: {user_input_mouse.shape}') - print(f'user_input_button.shape: {user_input_button.shape}') + user_input_mouse: torch.Tensor, # b,1,2 + user_input_button: torch.Tensor # b,1,11 + ) -> tuple[torch.Tensor, torch.Tensor]: noised_history = zlerp(self.history_buffer[:,1:], self.alpha) noised_audio = zlerp(self.audio_buffer[:,1:], self.alpha) @@ -82,9 +97,6 @@ def __call__(self, noised_history = torch.cat([noised_history, torch.randn_like(noised_history[:,0:1])], dim = 1) noised_audio = torch.cat([noised_audio, torch.randn_like(noised_audio[:,0:1])], dim = 1) - user_input_mouse = user_input_mouse[None,None,:] - user_input_button = user_input_button [None,None,:] - self.mouse_buffer = torch.cat([self.mouse_buffer[:,1:],user_input_mouse],dim=1) self.button_buffer = torch.cat([self.button_buffer[:,1:],user_input_button],dim=1) @@ -95,8 +107,11 @@ def __call__(self, ts = torch.ones_like(noised_history[:,:,0,0,0]) ts[:,:-1] = self.alpha - mouse_batch = torch.cat([self.mouse_buffer, torch.zeros_like(user_input_mouse)], dim=0) - btn_batch = torch.cat([self.button_buffer, torch.zeros_like(user_input_button)], dim=0) + # mouse_batch = torch.cat([self.mouse_buffer, torch.zeros_like(user_input_mouse)], dim=0) + # btn_batch = torch.cat([self.button_buffer, torch.zeros_like(user_input_button)], dim=0) + # TODO Who knows bruh idk. I think this is to get cfg - uncond is just no actions (zeros) + mouse_batch = torch.cat([self.mouse_buffer, torch.zeros_like(self.mouse_buffer)], dim=0) + btn_batch = torch.cat([self.button_buffer, torch.zeros_like(self.button_buffer)], dim=0) for _ in range(self.sampling_steps): x_batch = torch.cat([x, x], dim=0) @@ -115,13 +130,16 @@ def __call__(self, a[:,-1] = a[:,-1] - dt * pred_audio[:,-1] ts[:,-1] = ts[:,-1] - dt - new_frame = x [:,-1:] # [1,1,c,h,w] - new_audio = audio[:,-1:] # [1,1,c] + new_frame = x[:,-1:] # [1,1,c,h,w] + new_audio = a[:,-1:] # [1,1,c] self.history_buffer = torch.cat([self.history_buffer[:,1:], new_frame], dim=1) self.audio_buffer = torch.cat([self.audio_buffer[:,1:], new_audio], dim=1) frame = self.frame_decoder(new_frame[0] * self.frame_scale).squeeze() # [c,h,w] - audio = self.audio_decoder(self.audio_buffer * self.audio_scale).squeeze()[-self.audio_f:] # [735,2] + audio = self.audio_decoder( + self.audio_buffer.permute(0,2,1) # need this as [b,c,t] for some reason + * self.audio_scale + ).squeeze()[-self.audio_f:].T # [735,2] return frame, audio diff --git a/webapp/utils/configs.py b/webapp/utils/configs.py index 40877303..9a3a222e 100644 --- a/webapp/utils/configs.py +++ b/webapp/utils/configs.py @@ -9,7 +9,6 @@ @dataclass class WebappConfig: - model_checkpoint_path : os.PathLike run_config : RunConfig stream_config : StreamingConfig sampling_config : SamplingConfig @@ -40,6 +39,7 @@ class SamplingConfig: @dataclass class StreamingConfig: + model_checkpoint_path : os.PathLike fps: int = 20 frames_per_batch: int = 1 device: str = 'cuda' diff --git a/webapp/webapp_config.yaml b/webapp/webapp_config.yaml index 6502ac80..bc4b4b73 100644 --- a/webapp/webapp_config.yaml +++ b/webapp/webapp_config.yaml @@ -1,18 +1,18 @@ -model_checkpoint_path: "webapp/checkpoints/av_dfot_35k_ema_200m.pt" -run_config_path: "webapp/checkpoints/av.yml" +run_config_path: "webapp/checkpoints/configs/av.yml" device: "cuda" stream_config: - fps: 20 + fps: 10 frames_per_batch: 1 device: "cuda" n_buttons: 11 n_mouse_axes: 2 mouse_range: [-1.0, 1.0] - video_latent_history_path: "webapp/static/histories/base/video_latent.pt" - audio_latent_history_path: "webapp/static/histories/base/audio_latent.pt" - mouse_history_path: "webapp/static/histories/base/mouse.pt" - button_history_path: "webapp/static/histories/base/button.pt" + video_latent_history_path: "webapp/static/histories/base1/video_latent.pt" + audio_latent_history_path: "webapp/static/histories/base1/audio_latent.pt" + mouse_history_path: "webapp/static/histories/base1/mouse.pt" + button_history_path: "webapp/static/histories/base1/buttons.pt" + model_checkpoint_path: "webapp/checkpoints/models/av_dfot_35k_ema_200m.pt" sampling_config: sampling_steps: 20 From 28def2f716b3384e03a5bbc8bff8b48a001e85cc Mon Sep 17 00:00:00 2001 From: Sami Date: Thu, 12 Jun 2025 03:34:49 -0500 Subject: [PATCH 70/71] hodgepodge to get new history contexts --- owl_wms/data/s3_cod_latent_audio.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/owl_wms/data/s3_cod_latent_audio.py b/owl_wms/data/s3_cod_latent_audio.py index 998df356..4c401b92 100644 --- a/owl_wms/data/s3_cod_latent_audio.py +++ b/owl_wms/data/s3_cod_latent_audio.py @@ -188,11 +188,14 @@ def get_loader(batch_size, **data_kwargs): start = time.time() batch = next(iter(loader)) video_latent, audio_latent, mouse, button = batch - BASE_IDX=1 - torch.save(video_latent, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/video_latent.pt") - torch.save(audio_latent, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/audio_latent.pt") - torch.save(mouse, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/mouse.pt") - torch.save(button, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/buttons.pt") + for i in range(2, 6): + BASE_IDX=i + basedir = f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}" + os.makedirs(basedir, exist_ok=True) + torch.save(video_latent, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/video_latent.pt") + torch.save(audio_latent, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/audio_latent.pt") + torch.save(mouse, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/mouse.pt") + torch.save(button, f"/home/louis/owl-wms/webapp/static/histories/base{BASE_IDX}/buttons.pt") end = time.time() first_time = end - start From 5f3ccf591034f7a4c6b96c1d2c914579517f18e9 Mon Sep 17 00:00:00 2001 From: Sami Date: Sat, 14 Jun 2025 16:59:02 -0400 Subject: [PATCH 71/71] merge with cvpr demo --- owl-vaes | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/owl-vaes b/owl-vaes index f8b08b22..cdde9f3e 160000 --- a/owl-vaes +++ b/owl-vaes @@ -1 +1 @@ -Subproject commit f8b08b227fccc4a391b71187eeb530d7c39d4740 +Subproject commit cdde9f3e93cbae5ed77c99e6d7f29926ba89af42