""" Model Manager for real-time motion generation (HF Space version) Loads model from Hugging Face Hub instead of local checkpoints. """ import threading import time from collections import deque import numpy as np import torch from motion_process import StreamJointRecovery263 class FrameBuffer: """ Thread-safe frame buffer that maintains a queue of generated frames """ def __init__(self, target_buffer_size=4): self.buffer = deque(maxlen=100) # Max 100 frames in buffer self.target_size = target_buffer_size self.lock = threading.Lock() def add_frame(self, joints): """Add a frame to the buffer""" with self.lock: self.buffer.append(joints) def get_frame(self): """Get the next frame from buffer""" with self.lock: if len(self.buffer) > 0: return self.buffer.popleft() return None def size(self): """Get current buffer size""" with self.lock: return len(self.buffer) def clear(self): """Clear the buffer""" with self.lock: self.buffer.clear() def needs_generation(self): """Check if buffer needs more frames""" return self.size() < self.target_size class ModelManager: """ Manages model loading from HF Hub and real-time frame generation """ def __init__(self, model_name): self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {self.device}") # Load models from HF Hub self.vae, self.model = self._load_models(model_name) # Build config dicts from model's individual attributes (HF model API) self._base_schedule_config = { "chunk_size": self.model.chunk_size, "steps": self.model.noise_steps, } self._base_cfg_config = { "cfg_scale": self.model.cfg_scale, } # Frame buffer (for active session) self.frame_buffer = FrameBuffer(target_buffer_size=16) # Broadcast buffer (for spectators) - append-only with frame IDs self.broadcast_frames = deque(maxlen=200) self.broadcast_id = 0 self.broadcast_lock = threading.Lock() # Stream joint recovery with smoothing self.smoothing_alpha = 0.5 # Default: medium smoothing self.stream_recovery = StreamJointRecovery263( joints_num=22, smoothing_alpha=self.smoothing_alpha ) # Generation state self.current_text = "" self.is_generating = False self.generation_thread = None self.should_stop = False # Model generation state self.first_chunk = True # For VAE stream_decode self._model_first_chunk = True # For model stream_generate_step self.history_length = 30 print("ModelManager initialized successfully") def _patch_attention_sdpa(self, model_name): """Patch flash_attention() to include SDPA fallback for GPUs without flash-attn (e.g., T4).""" import glob import os hf_cache = os.path.join(os.path.expanduser("~"), ".cache", "huggingface") patterns = [ os.path.join( hf_cache, "hub", "models--" + model_name.replace("/", "--"), "snapshots", "*", "ldf_models", "tools", "attention.py", ), os.path.join( hf_cache, "modules", "transformers_modules", model_name, "*", "ldf_models", "tools", "attention.py", ), ] # Use the assert + next line as target to ensure idempotent patching target = ( ' assert q.device.type == "cuda" and q.size(-1) <= 256\n' "\n" " # params\n" ) replacement = ( ' assert q.device.type == "cuda" and q.size(-1) <= 256\n' "\n" " # SDPA fallback when flash-attn is not available (e.g., T4 GPU)\n" " if not FLASH_ATTN_2_AVAILABLE and not FLASH_ATTN_3_AVAILABLE:\n" " out_dtype = q.dtype\n" " b, lq, nq, c = q.shape\n" " lk = k.size(1)\n" " q = q.transpose(1, 2).to(dtype)\n" " k = k.transpose(1, 2).to(dtype)\n" " v = v.transpose(1, 2).to(dtype)\n" " attn_mask = None\n" " is_causal_flag = causal\n" " if k_lens is not None:\n" " k_lens = k_lens.to(q.device)\n" " valid = torch.arange(lk, device=q.device).unsqueeze(0) < k_lens.unsqueeze(1)\n" " attn_mask = torch.where(valid[:, None, None, :], 0.0, float('-inf')).to(dtype=dtype)\n" " is_causal_flag = False\n" " if causal:\n" " cm = torch.triu(torch.ones(lq, lk, device=q.device, dtype=torch.bool), diagonal=1)\n" " attn_mask = attn_mask.masked_fill(cm[None, None, :, :], float('-inf'))\n" " out = torch.nn.functional.scaled_dot_product_attention(\n" " q, k, v, attn_mask=attn_mask, is_causal=is_causal_flag, dropout_p=dropout_p\n" " )\n" " return out.transpose(1, 2).contiguous().to(out_dtype)\n" "\n" " # params\n" ) for pattern in patterns: for filepath in glob.glob(pattern): with open(filepath, "r") as f: content = f.read() if "SDPA fallback" in content: print(f"Already patched: {filepath}") continue if target in content: content = content.replace(target, replacement, 1) with open(filepath, "w") as f: f.write(content) print(f"Patched with SDPA fallback: {filepath}") def _load_models(self, model_name): """Load VAE and diffusion models from HF Hub""" torch.set_float32_matmul_precision("high") # Pre-download model files to hub cache print(f"Downloading model from HF Hub: {model_name}") from huggingface_hub import snapshot_download snapshot_download(model_name) # Patch flash_attention with SDPA fallback for T4 (no flash-attn) self._patch_attention_sdpa(model_name) print("Loading model...") from transformers import AutoModel hf_model = AutoModel.from_pretrained(model_name, trust_remote_code=True) hf_model.to(self.device) # Trigger lazy loading / warmup print("Warming up model...") _ = hf_model("test", length=1) # Access underlying streaming components model = hf_model.ldf_model vae = hf_model.vae model.eval() vae.eval() print("Models loaded successfully") return vae, model def start_generation(self, text, history_length=None): """Start or update generation with new text""" self.current_text = text if history_length is not None: self.history_length = history_length if not self.is_generating: # Reset state before starting (only once at the beginning) self.frame_buffer.clear() self.stream_recovery.reset() self.vae.clear_cache() self.first_chunk = True self._model_first_chunk = True # Restore model params from base config self.model.chunk_size = self._base_schedule_config["chunk_size"] self.model.noise_steps = self._base_schedule_config["steps"] self.model.cfg_scale = self._base_cfg_config["cfg_scale"] self.model.init_generated(self.history_length, batch_size=1) print( f"Model initialized with history length: {self.history_length}" ) # Start generation thread self.should_stop = False self.generation_thread = threading.Thread(target=self._generation_loop) self.generation_thread.daemon = True self.generation_thread.start() self.is_generating = True def update_text(self, text): """Update text without resetting state (continuous generation with new text)""" if text != self.current_text: old_text = self.current_text self.current_text = text # Don't reset first_chunk, stream_recovery, or vae cache # This allows continuous generation with text changes print(f"Text updated: '{old_text}' -> '{text}' (continuous generation)") def pause_generation(self): """Pause generation (keeps all state)""" self.should_stop = True if self.generation_thread: self.generation_thread.join(timeout=2.0) self.is_generating = False print("Generation paused (state preserved)") def resume_generation(self): """Resume generation from paused state""" if self.is_generating: print("Already generating, ignoring resume") return # Restart generation thread with existing state self.should_stop = False self.generation_thread = threading.Thread(target=self._generation_loop) self.generation_thread.daemon = True self.generation_thread.start() self.is_generating = True print("Generation resumed") def reset(self, history_length=None, smoothing_alpha=None): """Reset generation state completely Args: history_length: History window length for the model smoothing_alpha: EMA smoothing factor (0.0 to 1.0) - 1.0 = no smoothing (default) - 0.0 = infinite smoothing - Recommended: 0.3-0.7 for visible smoothing """ # Stop if running if self.is_generating: self.pause_generation() # Clear everything self.frame_buffer.clear() self.vae.clear_cache() self.first_chunk = True if history_length is not None: self.history_length = history_length # Update smoothing alpha if provided and recreate stream recovery if smoothing_alpha is not None: self.smoothing_alpha = np.clip(smoothing_alpha, 0.0, 1.0) print(f"Smoothing alpha updated to: {self.smoothing_alpha}") # Recreate stream recovery with new smoothing alpha self.stream_recovery = StreamJointRecovery263( joints_num=22, smoothing_alpha=self.smoothing_alpha ) # Restore model params from base config self.model.chunk_size = self._base_schedule_config["chunk_size"] self.model.noise_steps = self._base_schedule_config["steps"] self.model.cfg_scale = self._base_cfg_config["cfg_scale"] self._model_first_chunk = True # Initialize model self.model.init_generated(self.history_length, batch_size=1) print( f"Model reset - history: {self.history_length}, smoothing: {self.smoothing_alpha}" ) def _generation_loop(self): """Main generation loop that runs in background thread""" print("Generation loop started") step_count = 0 total_gen_time = 0 with torch.no_grad(): while not self.should_stop: # Check if buffer needs more frames if self.frame_buffer.needs_generation(): try: step_start = time.time() # Generate one token (produces frames from VAE) x = {"text": [self.current_text]} # Generate from model (1 token) output = self.model.stream_generate_step( x, first_chunk=self._model_first_chunk ) self._model_first_chunk = False generated = output["generated"] # Skip if no frames committed yet if generated[0].shape[0] == 0: continue # Decode with VAE (1 token -> 4 frames) decoded = self.vae.stream_decode( generated[0][None, :], first_chunk=self.first_chunk )[0] self.first_chunk = False # Convert each frame to joints for i in range(decoded.shape[0]): frame_data = decoded[i].cpu().numpy() joints = self.stream_recovery.process_frame(frame_data) self.frame_buffer.add_frame(joints) # Also add to broadcast buffer for spectators with self.broadcast_lock: self.broadcast_id += 1 self.broadcast_frames.append( (self.broadcast_id, joints) ) step_time = time.time() - step_start total_gen_time += step_time step_count += 1 # Print performance stats every 10 steps if step_count % 10 == 0: avg_time = total_gen_time / step_count fps = decoded.shape[0] / avg_time print( f"[Generation] Step {step_count}: {step_time * 1000:.1f}ms, " f"Avg: {avg_time * 1000:.1f}ms, " f"FPS: {fps:.1f}, " f"Buffer: {self.frame_buffer.size()}" ) except Exception as e: print(f"Error in generation: {e}") import traceback traceback.print_exc() time.sleep(0.1) else: # Buffer is full, wait a bit time.sleep(0.01) print("Generation loop stopped") def get_next_frame(self): """Get the next frame from buffer""" return self.frame_buffer.get_frame() def get_broadcast_frames(self, after_id, count=8): """Get frames from broadcast buffer after the given ID (for spectators).""" with self.broadcast_lock: frames = [ (fid, joints) for fid, joints in self.broadcast_frames if fid > after_id ] return frames[:count] def get_buffer_status(self): """Get buffer status""" return { "buffer_size": self.frame_buffer.size(), "target_size": self.frame_buffer.target_size, "is_generating": self.is_generating, "current_text": self.current_text, "smoothing_alpha": self.smoothing_alpha, "history_length": self.history_length, "schedule_config": { "chunk_size": self.model.chunk_size, "steps": self.model.noise_steps, }, "cfg_config": { "cfg_scale": self.model.cfg_scale, }, } # Global model manager instance _model_manager = None def get_model_manager(model_name=None): """Get or create the global model manager instance""" global _model_manager if _model_manager is None: _model_manager = ModelManager(model_name) return _model_manager