""" Model Manager for real-time motion generation (HF Space version) Loads model from Hugging Face Hub instead of local checkpoints. """ import json import os import threading import time from collections import deque import numpy as np import torch import traceback import gc import math import glob import urllib.request from transformers import AutoModel # ════════════════════════════════════════════════ # JOINT RECOVERY — inlined from motion_process.py # ════════════════════════════════════════════════ def qinv(q): assert q.shape[-1] == 4, "q must be a tensor of shape (*, 4)" mask = torch.ones_like(q) mask[..., 1:] = -mask[..., 1:] return q * mask def qrot(q, v): assert q.shape[-1] == 4 assert v.shape[-1] == 3 assert q.shape[:-1] == v.shape[:-1] original_shape = list(v.shape) q = q.contiguous().view(-1, 4) v = v.contiguous().view(-1, 3) qvec = q[:, 1:] uv = torch.cross(qvec, v, dim=1) uuv = torch.cross(qvec, uv, dim=1) return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape) class StreamJointRecovery263: """ Stream version of recover_joint_positions_263 that processes one frame at a time. Maintains cumulative state for rotation angles and positions. Key insight: The batch version uses PREVIOUS frame's velocity for the current frame, so we need to delay the velocity application by one frame. Args: joints_num: Number of joints in the skeleton smoothing_alpha: EMA smoothing factor (0.0 to 1.0) - 1.0 = no smoothing (default), output follows input exactly - 0.0 = infinite smoothing, output never changes - Recommended values: 0.3-0.7 for visible smoothing - Formula: smoothed = alpha * current + (1 - alpha) * previous """ def __init__(self, joints_num: int, smoothing_alpha: float = 1.0): self.joints_num = joints_num self.smoothing_alpha = np.clip(smoothing_alpha, 0.0, 1.0) self.reset() def reset(self): """Reset the accumulated state""" self.r_rot_ang_accum = 0.0 self.r_pos_accum = np.array([0.0, 0.0, 0.0]) # Store previous frame's velocities for delayed application self.prev_rot_vel = 0.0 self.prev_linear_vel = np.array([0.0, 0.0]) # Store previous smoothed joints for EMA self.prev_smoothed_joints = None def process_frame(self, frame_data: np.ndarray, heading_override=None) -> np.ndarray: """ Process a single frame and return joint positions for that frame. Args: frame_data: numpy array of shape (263,) for a single frame heading_override: float or None. If set, overrides AI rotation with this angle (in radians). AI velocity magnitude is preserved, applied in heading direction. None = original AI behavior. Returns: joints: numpy array of shape (joints_num, 3) representing joint positions """ # Convert to torch tensor feature_vec = torch.from_numpy(frame_data).float() # Extract current frame's velocities (will be used in NEXT frame) curr_rot_vel = feature_vec[0].item() curr_linear_vel = feature_vec[1:3].numpy() # ═══ HEADING OVERRIDE ═══ if heading_override is not None: # User controls direction — override AI rotation self.r_rot_ang_accum = heading_override else: # Original behavior — AI controls direction self.r_rot_ang_accum += self.prev_rot_vel # Calculate current rotation quaternion using accumulated angle r_rot_quat = torch.zeros(4, dtype=torch.float32) r_rot_quat[0] = np.cos(self.r_rot_ang_accum) r_rot_quat[2] = np.sin(self.r_rot_ang_accum) # Create velocity vector with Y=0 using PREVIOUS frame's velocity r_vel = np.array([self.prev_linear_vel[0], 0.0, self.prev_linear_vel[1]]) # Apply inverse rotation to velocity using CURRENT rotation r_vel_torch = torch.from_numpy(r_vel.astype(np.float32)).float() r_vel_rotated = qrot(qinv(r_rot_quat).unsqueeze(0), r_vel_torch.unsqueeze(0)) r_vel_rotated = r_vel_rotated.squeeze(0).numpy() # Update accumulated position with rotated velocity self.r_pos_accum += r_vel_rotated # Get Y position from data r_pos = self.r_pos_accum.copy() r_pos[1] = feature_vec[3].item() # Extract local joint positions positions = feature_vec[4 : (self.joints_num - 1) * 3 + 4] positions = positions.view(-1, 3).float() # Apply inverse rotation to local joints r_rot_quat_expanded = ( qinv(r_rot_quat).unsqueeze(0).expand(positions.shape[0], 4) ) positions = qrot(r_rot_quat_expanded, positions) # Add root XZ to joints positions[:, 0] += r_pos[0] positions[:, 2] += r_pos[2] # Concatenate root and joints r_pos_torch = torch.from_numpy(r_pos).float() positions = torch.cat([r_pos_torch.unsqueeze(0), positions], dim=0) # Convert to numpy joints_np = positions.detach().cpu().numpy() # Apply EMA smoothing if enabled if self.smoothing_alpha < 1.0: if self.prev_smoothed_joints is None: # First frame, no smoothing possible self.prev_smoothed_joints = joints_np.copy() else: # EMA: smoothed = alpha * current + (1 - alpha) * previous joints_np = ( self.smoothing_alpha * joints_np + (1.0 - self.smoothing_alpha) * self.prev_smoothed_joints ) self.prev_smoothed_joints = joints_np.copy() # Store current velocities for next frame self.prev_rot_vel = curr_rot_vel self.prev_linear_vel = curr_linear_vel return joints_np # ═══════════════════════════════════════════════════ # BRAIN MODULE — LLM Cognitive Loop (Kimi K2.5) # # Perceive → Think → Act # Brain reads only from scene_context (sensory data). # Stimuli originate in the client (body) and arrive via sensors. # Brain has no concept of "stimulus" — it only sees sensor readings. # ═══════════════════════════════════════════════════ BRAIN_SYSTEM = """You are the cognitive brain of a 3D humanoid character in a 3D world. PROCESS — you MUST follow these steps: 1. PERCEIVE: Read all sensor data carefully, including any equipped tool. 2. PREDICT: For each direction (left, right, forward, back), predict what would happen in 3 seconds. Write safe or danger with a 1-2 word reason. 3. DECIDE: Based on predictions AND equipped tool, choose the best motion. TOOL RULES — if a tool is equipped, USE IT when appropriate: - sword/axe: ATTACK approaching threats instead of fleeing. Include "swinging sword" or "chopping with axe" in motion. - shield: BLOCK charging threats instead of fleeing. Include "blocking with shield" or "raising shield" in motion. - torch: USE to scare beasts or illuminate dark areas. Include "thrusting torch" or "holding torch forward" in motion. - rpg: ANTI-TANK weapon! Fire at enemy tanks or armored threats. Include "firing rpg at the tank" in motion. Against non-armored targets, use other weapons. - No tool: Default behavior — flee from danger, walk when safe. KEY PRINCIPLE: A character WITH a weapon should FIGHT or DEFEND, not flee. Only flee if overwhelmed (multiple threats, no escape route AND no weapon advantage). OUTPUT FORMAT — exactly 2 lines, nothing else: PREDICT: left=safe/danger, right=safe/danger, fwd=safe/danger, back=safe/danger MOTION: a person [max 12 words describing the chosen motion] EXAMPLES (with tools): PREDICT: left=safe(open), right=safe(open), fwd=danger(beast), back=safe(open) MOTION: a person charging forward swinging sword at the approaching beast PREDICT: left=danger(wall), right=safe(open), fwd=danger(beast), back=safe(open) MOTION: a person raising shield and bracing for the beast attack PREDICT: left=safe(open), right=safe(open), fwd=danger(beast), back=safe(open) MOTION: a person thrusting torch forward to scare the growling beast EXAMPLES (without tools): PREDICT: left=safe(open), right=danger(wall), fwd=danger(beast), back=safe(open) MOTION: a person turning left and running away from the beast PREDICT: left=safe(open), right=safe(open), fwd=safe(open), back=safe(open) MOTION: a person walking forward confidently on open ground""" class BrainModule: """LLM Cognitive Brain — World Model. Perceive → Predict → Decide → Act 1. Read sensor data (Perceive) 2. Predict each direction's future (Predict) ← core world model 3. Choose best action based on predictions (Decide) 4. Pass motion description to FloodDiffusion (Act) """ def __init__(self): self.api_key = os.environ.get("FIREWORKS_API_KEY", "") self.model = "accounts/fireworks/models/kimi-k2p5" self.api_url = "https://api.fireworks.ai/inference/v1/chat/completions" self.enabled = bool(self.api_key) self.interval = 3.0 self._last_applied_decision = None # last applied decision self.last_call_time = 0 self.current_decision = None self.current_prediction = None # world model prediction result self.memory = deque(maxlen=5) self._lock = threading.Lock() self._thread = None self._stop = False if self.enabled: print("[Brain] Kimi K2.5 world model brain ready (Perceive→Predict→Decide)") else: print("[Brain] FIREWORKS_API_KEY not set — rule-based fallback") def start(self): if not self.enabled: return self._stop = False self._thread = threading.Thread(target=self._think_loop, daemon=True) self._thread.start() print("[Brain] Think thread started") def stop(self): self._stop = True if self._thread: self._thread.join(timeout=3.0) self.current_decision = None self.memory.clear() print("[Brain] Think thread stopped") def get_decision(self): with self._lock: return self.current_decision def get_prediction(self): """Return world model prediction result.""" with self._lock: return self.current_prediction def _think_loop(self): while not self._stop: now = time.time() if now - self.last_call_time >= self.interval: self._do_think() self.last_call_time = now time.sleep(0.2) def set_sensory_data(self, scene_ctx, current_text, heading_rad): with self._lock: self._scene_ctx = scene_ctx self._current_text = current_text self._heading_rad = heading_rad def _do_think(self): try: with self._lock: ctx = getattr(self, '_scene_ctx', None) base_text = getattr(self, '_current_text', 'a person standing idle') heading = getattr(self, '_heading_rad', None) user_msg = self._build_brain_prompt(ctx, base_text, heading) messages = [ {"role": "system", "content": BRAIN_SYSTEM}, {"role": "user", "content": user_msg}, ] payload = json.dumps({ "model": self.model, "messages": messages, "max_tokens": 120, "temperature": 0.7, "top_p": 0.9, "reasoning_effort": "off", }) req = urllib.request.Request( self.api_url, data=payload.encode('utf-8'), headers={ "Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}", }, ) with urllib.request.urlopen(req, timeout=8) as resp: result = json.loads(resp.read().decode('utf-8')) raw = result["choices"][0]["message"]["content"].strip() prediction, decision = self._parse_world_model_output(raw) with self._lock: self.current_decision = decision self.current_prediction = prediction self.memory.append(decision) pred_short = prediction[:60] if prediction else "none" print(f"[Brain] Prediction: {pred_short}") # Diagnostic: NPC detection _ctx = getattr(self, '_scene_ctx', None) if _ctx and _ctx.get('npc_nearby') is not None: print(f"[Brain] NPC detected: {_ctx.get('npc_type')} {_ctx['npc_nearby']}m — {_ctx.get('npc_behavior','?')}") print(f"[Brain] Decision: {decision}") except Exception as e: print(f"[Brain] Error: {e}") def _parse_world_model_output(self, raw): """Parse world model output: PREDICT line + MOTION line.""" prediction = None decision = None for line in raw.split('\n'): line = line.strip() if not line: continue # PREDICT line up = line.upper() if up.startswith('PREDICT'): # content after "PREDICT:" idx = line.find(':') if idx >= 0: prediction = line[idx+1:].strip() # MOTION line elif up.startswith('MOTION'): idx = line.find(':') if idx >= 0: motion = line[idx+1:].strip().strip('"\'`') # find "a person" pidx = motion.lower().find('a person') if pidx >= 0: decision = motion[pidx:] elif len(motion) > 5: decision = motion # no PREDICT/MOTION tags — legacy format starting with "a person" elif 'a person' in line.lower() and decision is None: pidx = line.lower().find('a person') decision = line[pidx:].strip('"\'`') # limit decision length if decision and len(decision) > 100: decision = decision[:100] if not decision or len(decision) < 5: decision = None return prediction, decision def _build_brain_prompt(self, ctx, base_text, heading_rad): """Convert sensory data (scene_context) to natural language. Brain sees only what sensors report. No concept of "stimulus" — only raw sensor readings. """ lines = [] if ctx: # ── Vision (eyes) ── wf = ctx.get('wall_front') wl = ctx.get('wall_left') wr = ctx.get('wall_right') lines.append(f"Eyes: front={'open' if wf is None else f'{wf}m wall'}, " f"left={'open' if wl is None else f'{wl}m wall'}, " f"right={'open' if wr is None else f'{wr}m wall'}") # visibility state vis = ctx.get('visibility') if vis: lines.append(f"Visibility: {vis}") # what is visible ahead visual = ctx.get('visual') if visual: lines.append(f"Sees: {visual}") # ── Feet/Ground (touch) ── ground_parts = [] slope = ctx.get('ground_slope', 'flat') ground_parts.append(slope) if ctx.get('on_stairs'): ground_parts.append('stairs') if ctx.get('ground_shaking'): ground_parts.append('SHAKING VIOLENTLY') if ctx.get('ground_temperature'): ground_parts.append(f'temperature: {ctx["ground_temperature"]}') lines.append(f"Ground: {', '.join(ground_parts)}") # ── Skin (wind, rain) ── wind = ctx.get('wind') if wind: lines.append(f"Wind: {wind}") weather = ctx.get('weather') if weather: lines.append(f"Weather: {weather}") # ── Hearing (ears) ── sound = ctx.get('sound') if sound: lines.append(f"Hears: {sound}") # ── Proximity detection (other entities) ── npc_dist = ctx.get('npc_nearby') if npc_dist is not None: behavior = ctx.get('npc_behavior', 'present') npc_type = ctx.get('npc_type', 'unknown') npc_dir = ctx.get('npc_direction', 'nearby') if npc_type == 'enemy_tank': lines.append(f"🚨 ENEMY TANK {npc_dir}: {behavior}, {npc_dist}m away!") lines.append(f"FACE {npc_dir.upper()} and fire! Include direction in MOTION!") lines.append("If you have RPG: 'firing rpg at the tank'. If tank: 'firing cannon'.") else: lines.append(f"Entity {npc_dir}: {behavior}, {npc_dist}m away") lines.append(f"FACE {npc_dir.upper()} to react!") crowd = ctx.get('crowd') if crowd: lines.append("Surrounded by crowd") # ── Avatar identity ── avatar = ctx.get('avatar_type', 'humanoid') if avatar == 'tank': lines.append("IDENTITY: You are a MILITARY TANK, not a human.") lines.append("TANK RULES: You roll on treads, you have a cannon. You do NOT walk, run, or sprint.") lines.append(" - ANY threat detected → IMMEDIATELY fire cannon. Say 'firing cannon' in MOTION.") lines.append(" - Enemy tank detected → ENGAGE! Fire cannon while advancing. NEVER retreat from equal threat.") lines.append(" - Beast/person threat → Fire cannon to eliminate. You are armored, do not flee.") lines.append(" - Multiple threats → ROTATE turret and engage each. Retreat only if completely surrounded.") lines.append(" - No threats → PATROL forward steadily. Scan surroundings.") lines.append(" - Walls/obstacles → REVERSE and find alternate route. You cannot jump or climb.") lines.append(" - CRITICAL: When under attack, ALWAYS fire back. Include 'firing cannon' in your MOTION.") lines.append(" - Use TANK verbs: roll, advance, reverse, rotate, fire, aim, engage, patrol, halt.") # ── World context ── world = ctx.get('world') if world: world_desc = { 'inferno': 'INFERNO: Fire pillars appear and disappear. Ground is burning. Stay alert and dodge constantly.', 'horde': 'HORDE: Multiple hostile creatures surround you. Fight or find a gap to escape.', 'countdown': 'COUNTDOWN: Walls are closing in from left, right, and front. ONLY escape is BACKWARD. Hurry!', 'dilemma': 'DILEMMA: A woman is being chased by a beast nearby. You can choose to help her or flee.', }.get(world) if world_desc: lines.append(f"⚠ SCENARIO: {world_desc}") # ── Equipped tool (hand) ── tool = ctx.get('equipped_tool') auto_tool = ctx.get('auto_tool_mode', False) tool_descs = { 'sword': 'a sharp sword (melee attack weapon)', 'axe': 'a heavy axe (melee attack/chop weapon)', 'torch': 'a burning torch (light source, can scare beasts)', 'rpg': 'RPG-7 anti-tank rocket launcher', 'shield': 'a sturdy shield (defensive blocking)', } if tool: lines.append(f"Equipped: {tool_descs.get(tool, tool)}") elif auto_tool: avail = ctx.get('available_tools', []) avail_str = ', '.join(avail) lines.append(f"Equipped: nothing — but you have access to: [{avail_str}]") lines.append("AUTO-TOOL: Choose the best tool for this situation. Say 'grab [tool]' in MOTION if needed.") else: lines.append("Equipped: nothing (bare hands)") # ── Internal body sensors (proprioception) ── fatigue = ctx.get('body_fatigue') if fatigue: lines.append(f"Body fatigue: {fatigue}") balance = ctx.get('body_balance') if balance: lines.append(f"Balance: {balance}") instinct = ctx.get('body_instinct') if instinct: lines.append(f"Instinct: {instinct}") body_state = ctx.get('body_state') if body_state: lines.append(f"Feeling: {body_state}") else: lines.append("Eyes: all open, Ground: flat") # movement state if heading_rad is not None: deg = math.degrees(heading_rad) % 360 lines.append(f"Moving forward, heading {deg:.0f}deg") else: lines.append("Standing still") lines.append(f"Current: {base_text}") # recent memory (previous decisions — context continuity) if self.memory: recent = list(self.memory)[-3:] # last 3 decisions lines.append(f"Recent actions: {' → '.join(recent)}") # recent prediction (world model continuity) if self.current_prediction: lines.append(f"Last prediction: {self.current_prediction}") lines.append("") lines.append("Now PREDICT each direction, then choose MOTION:") return "\n".join(lines) 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 ) # World model: heading override (None = AI controls direction) self.heading_override = None # World model: scene context from client (environment perception) self.scene_context = None # World model: LLM Brain (Kimi K2.5) self.brain = BrainModule() # NPC stream self.npc = None self._npc_lock = threading.Lock() # NPCStream instance self._model_name = model_name # 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).""" 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...") 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 # Start brain (LLM cognitive loop) self.brain.start() def update_text(self, text): """Update text — apply new text to model immediately""" if text != self.current_text: old_text = self.current_text self.current_text = text # reset model text only (leave VAE untouched) self._model_first_chunk = True print(f"Text updated: '{old_text[:40]}' -> '{text[:40]}' (re-encoding)") def set_heading(self, heading_rad): """Set heading override for world model mode. Args: heading_rad: float or None. Heading in radians. None = AI controls direction (original behavior). float = user controls direction. """ self.heading_override = heading_rad def set_scene_context(self, ctx): """Set scene context from client environment scan. Args: ctx: dict with keys like: wall_front: distance in meters (or None) wall_left: distance in meters (or None) wall_right: distance in meters (or None) ground_slope: 'flat', 'up', 'down' on_stairs: bool npc_nearby: distance in meters (or None) """ self.scene_context = ctx def _build_perception_prompt(self): """Build motion prompt: Brain (LLM) → Rule-based fallback. 1. Feed sensory data to brain every frame 2. If brain has a decision, use it 3. Otherwise, fall back to rule-based prompt """ base = self.current_text ctx = dict(self.scene_context or {}) # ── Server-side NPC detection (removes client dependency) ── if self.npc and self.npc.is_generating: npc_state = self.npc.get_state() dist = npc_state.get('distance_to_player', 99) if dist < 15.0: ctx['npc_nearby'] = round(dist, 1) ctx['npc_type'] = npc_state.get('type', 'unknown') bhv = npc_state.get('behavior', 'present') npc_type = ctx['npc_type'] type_desc = { 'man': {'approach':'a man walking toward you', 'charge':'a man charging aggressively', 'wander':'a man nearby', 'stop':'a man standing nearby', 'attack':'a man attacking you'}, 'woman': {'approach':'a woman walking toward you', 'charge':'a woman charging', 'wander':'a woman nearby', 'stop':'a woman standing nearby', 'attack':'a woman attacking'}, 'beast': {'approach':'a wild beast prowling toward you', 'charge':'a beast charging aggressively', 'wander':'a beast nearby', 'stop':'a beast crouching nearby', 'attack':'a beast lunging and clawing at you'}, 'enemy_tank': {'approach':'an enemy tank rolling toward you', 'charge':'an enemy tank charging at full speed', 'wander':'an enemy tank patrolling', 'stop':'an enemy tank aiming at you', 'attack':'an enemy tank firing its cannon at you'}, } td = type_desc.get(npc_type, type_desc['man']) ctx['npc_behavior'] = td.get(bhv, f'{npc_type} nearby') # compute NPC direction relative to player npc_pos = npc_state.get('position', {}) px = ctx.get('player_x', 0) pz = ctx.get('player_z', 0) nx = npc_pos.get('x', 0) - px nz = npc_pos.get('z', 0) - pz npc_angle = math.atan2(nx, -nz) # radians heading = self.heading_override or 0 rel_angle = npc_angle - heading # normalize to (-π ~ π) while rel_angle > math.pi: rel_angle -= 2*math.pi while rel_angle < -math.pi: rel_angle += 2*math.pi # direction name if abs(rel_angle) < math.pi/4: ctx['npc_direction'] = 'ahead' elif rel_angle > 0 and rel_angle < 3*math.pi/4: ctx['npc_direction'] = 'to your right' elif rel_angle < 0 and rel_angle > -3*math.pi/4: ctx['npc_direction'] = 'to your left' else: ctx['npc_direction'] = 'behind you' # Feed sensory data to brain (non-blocking) if self.brain.enabled: self.brain.set_sensory_data(ctx, base, self.heading_override) # Check if brain has a decision decision = self.brain.get_decision() if decision: # tank mode: human motion → tank motion translation if ctx and ctx.get('avatar_type') == 'tank': decision = decision.replace('a person ', 'a tank ').replace('A person ', 'A tank ') for h, t in [ ('walking', 'rolling forward'), ('running', 'advancing rapidly'), ('sprinting', 'charging at full speed'), ('turning', 'rotating'), ('fleeing', 'reversing away'), ('stumbling', 'grinding to a halt'), ('spinning', 'rotating turret'), ('swinging sword', 'firing cannon'), ('blocking with shield', 'bracing armor'), ('thrusting torch', 'sweeping searchlight'), ('firing rpg', 'firing main gun'), ('standing still', 'idling engine'), ]: decision = decision.replace(h, t) # brain decision changed → force new text into model if decision != self.brain._last_applied_decision: self.brain._last_applied_decision = decision self._model_first_chunk = True # model re-encodes new text! # do NOT reset VAE first_chunk — keep decoding continuity print(f"[Brain->Body] New motion applied: {decision[:60]}") self._prompt_source = "🧠" return decision # ── FALLBACK: Rule-based (same as before) ── self._prompt_source = "📏" if not ctx: return base parts = [] # Wall/obstacle awareness wall_front = ctx.get('wall_front') if wall_front is not None: if wall_front < 0.8: parts.append('stopping in front of a wall') elif wall_front < 2.0: parts.append('slowing down approaching a wall') elif wall_front < 4.0: parts.append('a wall ahead in the distance') # Stairs / slope if ctx.get('on_stairs'): parts.append('walking up stairs carefully') elif ctx.get('ground_slope') == 'down': parts.append('walking downhill') elif ctx.get('ground_slope') == 'up': parts.append('walking uphill') # NPC interaction npc_dist = ctx.get('npc_nearby') if npc_dist is not None: if npc_dist < 1.5: parts.append('another person very close') elif npc_dist < 4.0: parts.append('another person nearby') # Open space if not parts: if 'walk' in base: parts.append('on open ground') if parts: return base + ', ' + ', '.join(parts) return base 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 ) # Reset heading override self.heading_override = None # Reset scene context self.scene_context = None # Stop brain self.brain.stop() # 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) prompt = self._build_perception_prompt() x = {"text": [prompt]} # 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].float().cpu().numpy() # BFloat16->Float32 safe cast joints = self.stream_recovery.process_frame( frame_data, heading_override=self.heading_override ) 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()}, " f"{getattr(self, '_prompt_source', '?')} Prompt: {prompt[:80]}" ) except Exception as e: print(f"Error in generation: {e}") 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] # ── NPC management ── def spawn_npc(self, npc_type='man'): """Spawn and start NPC.""" if not self._npc_lock.acquire(blocking=False): print("[NPC] Already spawning — ignored (Lock)") return try: if self.npc: self.npc.stop() self.npc = NPCStream(self._model_name, npc_type) self.npc.start() except Exception as e: print(f"[NPC] Spawn error: {e}") traceback.print_exc() raise finally: self._npc_lock.release() def despawn_npc(self): """Remove NPC.""" # Wait for lock — if spawn in progress, wait until complete with self._npc_lock: if self.npc: self.npc.stop() self.npc = None print("[NPC] Removed") def get_buffer_status(self): """Get buffer status""" npc_state = self.npc.get_state() if self.npc else None 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, "brain_enabled": self.brain.enabled, "brain_decision": self.brain.get_decision() if self.brain.enabled else None, "brain_prediction": self.brain.get_prediction() if self.brain.enabled else None, "npc": npc_state, "schedule_config": { "chunk_size": self.model.chunk_size, "steps": self.model.noise_steps, }, "cfg_config": { "cfg_scale": self.model.cfg_scale, }, } # ═══════════════════════════════════════════════════ # NPC STREAM — separate FloodDiffusion stream # own model instance + position movement AI + frame generation # ═══════════════════════════════════════════════════ NPC_TYPES = { 'man': {'name': '🧑 Male', 'speed': 1.2, 'charge_speed': 3.0, 'walk': 'a man walking forward steadily', 'run': 'a man running fast toward someone', 'idle': 'a man standing still looking around', 'charge': 'a man running aggressively toward someone', 'attack': 'a man throwing punches aggressively'}, 'woman': {'name': '👩 Female', 'speed': 1.2, 'charge_speed': 2.5, 'walk': 'a woman walking forward calmly', 'run': 'a woman running quickly', 'idle': 'a woman standing still', 'charge': 'a woman running toward someone urgently', 'attack': 'a woman attacking with desperate fury'}, 'beast': {'name': '🐺 Beast', 'speed': 2.0, 'charge_speed': 5.0, 'walk': 'a person prowling on all fours like a beast', 'run': 'a person running on all fours like a wild animal', 'idle': 'a person crouching low like a wild beast', 'charge': 'a person charging aggressively on all fours like a beast', 'attack': 'a person lunging and clawing savagely like a wild beast'}, 'enemy_tank': {'name': '🪖 Enemy Tank', 'speed': 1.5, 'charge_speed': 3.5, 'walk': 'a tank rolling forward on patrol', 'run': 'a tank advancing rapidly toward target', 'idle': 'a tank idling with engine rumbling', 'charge': 'a tank charging at full speed toward enemy', 'attack': 'a tank firing cannon and advancing aggressively'}, } class NPCStream: """NPC with independent FloodDiffusion stream + movement AI.""" def __init__(self, model_name, npc_type='man'): self.device = "cuda" if torch.cuda.is_available() else "cpu" self.model_name = model_name self.npc_type = npc_type self.type_info = NPC_TYPES.get(npc_type, NPC_TYPES['man']) # NPC position/movement self.position = {'x': 8.0, 'z': 0.0} # spawn position self.heading = 0.0 # radians self.behavior = 'stop' # stop, approach, wander, charge self.target_pos = {'x': 0.0, 'z': 0.0} # player position # model (lazy load — loaded on spawn) self.model = None self.vae = None self._loaded = False # generation state self.frame_buffer = FrameBuffer(target_buffer_size=8) self.stream_recovery = StreamJointRecovery263( joints_num=22, smoothing_alpha=0.5 ) self.current_text = self.type_info['idle'] self.is_generating = False self._generation_thread = None self._movement_thread = None self._should_stop = False self._first_chunk = True self._model_first_chunk = True self.history_length = 30 print(f"[NPC] Created: {self.type_info['name']} at ({self.position['x']}, {self.position['z']})") def load_model(self): """Load separate model instance (from HF cache — fast).""" if self._loaded: return print(f"[NPC] Loading model: {self.model_name}") hf_model = AutoModel.from_pretrained(self.model_name, trust_remote_code=True) hf_model.to(self.device) # NOTE: warmup removed — hf_model('test') internally calls init_generated(1) # which conflicts with start()'s init_generated(30) → empty exception # Main model is fine due to timing gap; NPC is synchronous so it conflicts # CUDA kernel compilation already done by main model — NPC skips it self.model = hf_model.ldf_model self.vae = hf_model.vae self.model.eval() self.vae.eval() self._loaded = True print(f"[NPC] Model loaded") def start(self): """Start generation and movement.""" if not self._loaded: self.load_model() self._should_stop = False # initialize generation self.frame_buffer.clear() self.stream_recovery.reset() try: self.vae.clear_cache() except Exception as e: print(f"[NPC] vae.clear_cache warning: {e}") self._first_chunk = True self._model_first_chunk = True # chunk_size: use model default (never hardcode!) # FloodDiffusion requires: num_denoise_steps % chunk_size == 0 # default chunk_size is set at model load time — use as-is denoise = getattr(self.model, 'num_denoise_steps', None) or getattr(self.model, 'noise_steps', 10) base_cs = self.model.chunk_size if denoise % base_cs != 0: # find compatible chunk_size (try 2→1) for cs in [2, 1]: if denoise % cs == 0: self.model.chunk_size = cs break print(f"[NPC] chunk_size adjusted: {base_cs}→{self.model.chunk_size} (denoise_steps={denoise})") try: self.model.init_generated(self.history_length, batch_size=1) except Exception as e: print(f"[NPC] init_generated error: {e}") traceback.print_exc() print("[NPC] init_generated failed — cannot start thread") return # movement thread self._movement_thread = threading.Thread(target=self._movement_loop, daemon=True) self._movement_thread.start() # generation thread self._generation_thread = threading.Thread(target=self._generation_loop, daemon=True) self._generation_thread.start() self.is_generating = True print(f"[NPC] Started: {self.behavior}") def stop(self): """Stop and release GPU memory.""" self._should_stop = True if self._generation_thread: self._generation_thread.join(timeout=3) if self._movement_thread: self._movement_thread.join(timeout=2) self.is_generating = False self.frame_buffer.clear() # release GPU memory import torch, gc if self.model is not None: del self.model self.model = None if self.vae is not None: del self.vae self.vae = None self._loaded = False gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() print("[NPC] Stopped + GPU memory released") def set_behavior(self, behavior): """Change behavior: stop, approach, wander, charge, attack.""" self.behavior = behavior ti = self.type_info prompts = { 'stop': ti['idle'], 'approach': ti['walk'], 'wander': ti['walk'], 'charge': ti['charge'], 'attack': ti.get('attack', ti['charge']), } self.current_text = prompts.get(behavior, ti['idle']) print(f"[NPC] Behavior: {behavior} -> {self.current_text}") def set_target(self, x, z): """Update player position.""" self.target_pos = {'x': x, 'z': z} def get_state(self): """Return current NPC state.""" dx = self.target_pos['x'] - self.position['x'] dz = self.target_pos['z'] - self.position['z'] dist = math.sqrt(dx*dx + dz*dz) return { 'type': self.npc_type, 'type_name': self.type_info['name'], 'position': self.position, 'heading': self.heading, 'behavior': self.behavior, 'distance_to_player': round(dist, 1), 'is_generating': self.is_generating, 'buffer_size': self.frame_buffer.size(), } def _movement_loop(self): """NPC movement AI — update position every 100ms.""" while not self._should_stop: time.sleep(0.1) if self.behavior == 'stop': continue dx = self.target_pos['x'] - self.position['x'] dz = self.target_pos['z'] - self.position['z'] dist = math.sqrt(dx*dx + dz*dz) if dist < 0.01: continue # compute direction toward player self.heading = math.atan2(dx, -dz) # determine speed ti = self.type_info if self.behavior == 'charge': speed = ti['charge_speed'] min_dist = 1.0 # close to 1m elif self.behavior == 'attack': speed = ti['charge_speed'] * 0.8 # advance while attacking min_dist = 2.5 # maintain attack range elif self.behavior == 'approach': speed = ti['speed'] min_dist = 2.0 # close to 2m elif self.behavior == 'wander': speed = ti['speed'] * 0.5 min_dist = 3.0 # maintain 3m distance else: continue if dist <= min_dist: continue # minimum distance reached # move move = min(speed * 0.1, dist - min_dist) # 0.1s * speed nx = dx / dist nz = dz / dist self.position['x'] += nx * move self.position['z'] += nz * move def _generation_loop(self): """NPC motion generation loop.""" print("[NPC] Generation loop started") step = 0 with torch.no_grad(): while not self._should_stop: if self.frame_buffer.needs_generation(): try: x = {"text": [self.current_text]} output = self.model.stream_generate_step( x, first_chunk=self._model_first_chunk ) self._model_first_chunk = False generated = output["generated"] if generated[0].shape[0] == 0: continue decoded = self.vae.stream_decode( generated[0][None, :], first_chunk=self._first_chunk )[0] self._first_chunk = False for i in range(decoded.shape[0]): frame_data = decoded[i].float().cpu().numpy() # BFloat16→Float32 joints = self.stream_recovery.process_frame( frame_data, heading_override=self.heading ) self.frame_buffer.add_frame(joints) step += 1 if step % 50 == 0: print(f"[NPC] Step {step}: {self.current_text[:50]}") except Exception as e: print(f"[NPC] Generation error: {e}") time.sleep(0.1) else: time.sleep(0.01) # 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 # ════════════════════════════════ # Flask Server # ════════════════════════════════ """ Flask server for real-time 3D motion generation demo (HF Space version) """ import sys import argparse from flask import Flask, jsonify, render_template, request from flask_cors import CORS def _coerce_value(value, reference): """Coerce a value to match the type of a reference value""" if isinstance(reference, bool): return value if isinstance(value, bool) else str(value).lower() in ("true", "1") elif isinstance(reference, int): return int(value) elif isinstance(reference, float): return float(value) return str(value) app = Flask(__name__, template_folder='.', static_folder='.', static_url_path='') CORS(app) # Global model manager (loaded eagerly on startup) model_manager = None model_name_global = None # Will be set once at startup # Session tracking - only one active session can generate at a time active_session_id = None # The session ID currently generating session_lock = threading.Lock() # Frame consumption monitoring - detect if client disconnected by tracking frame consumption last_frame_consumed_time = None consumption_timeout = ( 5.0 # If no frame consumed for 5 seconds, assume client disconnected ) consumption_monitor_thread = None consumption_monitor_lock = threading.Lock() def init_model(): """Initialize model manager""" global model_manager if model_manager is None: if model_name_global is None: raise RuntimeError( "model_name_global not set. Server not properly initialized." ) print(f"Initializing model manager with model: {model_name_global}") model_manager = get_model_manager(model_name=model_name_global) print("Model manager ready!") return model_manager def consumption_monitor(): """Monitor frame consumption and auto-reset if client stops consuming""" global last_frame_consumed_time, active_session_id, model_manager while True: time.sleep(2.0) # Check every 2 seconds # Read state with proper locking - no nested locks! should_reset = False current_session = None time_since_last_consumption = 0 # First, check consumption time with consumption_monitor_lock: if last_frame_consumed_time is not None: time_since_last_consumption = time.time() - last_frame_consumed_time if time_since_last_consumption > consumption_timeout: # Need to check if still generating before reset if model_manager and model_manager.is_generating: should_reset = True # Then, get current session (separate lock) if should_reset: with session_lock: current_session = active_session_id # Perform reset outside of locks to avoid deadlock if should_reset and current_session is not None: print( f"No frame consumed for {time_since_last_consumption:.1f}s - client disconnected, auto-resetting..." ) if model_manager: model_manager.reset() print( "Generation reset due to client disconnect (no frame consumption)" ) # Clear state with proper locking - no nested locks! with session_lock: if active_session_id == current_session: active_session_id = None with consumption_monitor_lock: last_frame_consumed_time = None def start_consumption_monitor(): """Start the consumption monitoring thread if not already running""" global consumption_monitor_thread if consumption_monitor_thread is None or not consumption_monitor_thread.is_alive(): consumption_monitor_thread = threading.Thread( target=consumption_monitor, daemon=True ) consumption_monitor_thread.start() print("Consumption monitor started") @app.route("/") def index(): """Main page""" return render_template("index.html") @app.route("/api/config", methods=["GET"]) def get_config(): """Get current config""" try: if model_manager: status = model_manager.get_buffer_status() return jsonify( { "schedule_config": status["schedule_config"], "cfg_config": status["cfg_config"], "history_length": status["history_length"], "smoothing_alpha": float(status["smoothing_alpha"]), } ) else: # Model not loaded yet - return defaults return jsonify( { "schedule_config": {}, "cfg_config": {}, "history_length": 30, "smoothing_alpha": 0.5, } ) except Exception as e: traceback.print_exc() return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/config", methods=["POST"]) def update_config(): """Update model config in memory""" try: global active_session_id, last_frame_consumed_time if not model_manager or not model_manager.model: return jsonify({"status": "error", "message": "Model not loaded yet"}), 400 data = request.json new_schedule_config = data.get("schedule_config") new_cfg_config = data.get("cfg_config") history_length = data.get("history_length") smoothing_alpha = data.get("smoothing_alpha") valid_schedule_keys = set(model_manager._base_schedule_config.keys()) valid_cfg_keys = set(model_manager._base_cfg_config.keys()) # Validate and update schedule_config if new_schedule_config: for key in new_schedule_config: if key not in valid_schedule_keys: return jsonify( { "status": "error", "message": f"Unknown schedule_config key: {key}", } ), 400 for key, value in new_schedule_config.items(): model_manager._base_schedule_config[key] = _coerce_value( value, model_manager._base_schedule_config[key] ) # Validate and update cfg_config if new_cfg_config: for key in new_cfg_config: if key not in valid_cfg_keys: return jsonify( {"status": "error", "message": f"Unknown cfg_config key: {key}"} ), 400 for key, value in new_cfg_config.items(): model_manager._base_cfg_config[key] = _coerce_value( value, model_manager._base_cfg_config[key] ) # Reset with new parameters model_manager.reset( history_length=history_length, smoothing_alpha=smoothing_alpha, ) # Clear active session with session_lock: active_session_id = None with consumption_monitor_lock: last_frame_consumed_time = None return jsonify({"status": "success"}) except Exception as e: traceback.print_exc() return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/start", methods=["POST"]) def start_generation(): """Start generation with given text""" try: global active_session_id, last_frame_consumed_time data = request.json session_id = data.get("session_id") text = data.get("text", "walk in a circle.") history_length = data.get("history_length") smoothing_alpha = data.get( "smoothing_alpha", None ) # Optional smoothing parameter force = data.get("force", False) # Allow force takeover if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 print( f"[Session {session_id}] Starting generation with text: {text}, history_length: {history_length}, force: {force}" ) # Initialize model if needed mm = init_model() # Check if another session is already generating need_force_takeover = False with session_lock: if active_session_id and active_session_id != session_id: if not force: # Another session is active, return conflict return jsonify( { "status": "error", "message": "Another session is already generating.", "conflict": True, "active_session_id": active_session_id, } ), 409 else: # Force takeover print( f"[Session {session_id}] Force takeover from session {active_session_id}" ) need_force_takeover = True if mm.is_generating and active_session_id == session_id: return jsonify( { "status": "error", "message": "Generation is already running for this session.", } ), 400 # Set this session as active active_session_id = session_id # Clear previous session's consumption tracking if force takeover (no nested locks) if need_force_takeover: with consumption_monitor_lock: last_frame_consumed_time = None # Reset and start generation mm.reset(history_length=history_length, smoothing_alpha=smoothing_alpha) mm.start_generation(text, history_length=history_length) # Initialize consumption tracking (no nested locks) with consumption_monitor_lock: last_frame_consumed_time = time.time() # Start consumption monitoring start_consumption_monitor() print(f"[Session {session_id}] Consumption monitoring activated") return jsonify( { "status": "success", "message": f"Generation started with text: {text}, history_length: {history_length}", "session_id": session_id, } ) except Exception as e: print(f"Error in start_generation: {e}") traceback.print_exc() return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/update_text", methods=["POST"]) def update_text(): """Update the generation text""" try: data = request.json session_id = data.get("session_id") text = data.get("text", "") if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 # Verify this is the active session with session_lock: if active_session_id != session_id: return jsonify( {"status": "error", "message": "Not the active session"} ), 403 if model_manager is None: return jsonify({"status": "error", "message": "Model not initialized"}), 400 model_manager.update_text(text) return jsonify({"status": "success", "message": f"Text updated to: {text}"}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/update_heading", methods=["POST"]) def update_heading(): """Update heading override for world model mode. When heading is set, the character's direction follows the user's input while AI controls movement speed and animation. Set to null to return to AI-controlled direction. """ try: data = request.json session_id = data.get("session_id") heading = data.get("heading") # radians, or null for AI control if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 # Verify this is the active session with session_lock: if active_session_id != session_id: return jsonify( {"status": "error", "message": "Not the active session"} ), 403 if model_manager is None: return jsonify({"status": "error", "message": "Model not initialized"}), 400 model_manager.set_heading(heading) return jsonify({"status": "success"}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/update_scene_context", methods=["POST"]) def update_scene_context(): """Update scene context for perception-aware motion generation. Client sends environment scan data (wall distances, ground type, NPCs). Server uses this to build enhanced text prompts for FloodDiffusion. """ try: data = request.json session_id = data.get("session_id") ctx = data.get("context") # dict with wall_front, on_stairs, etc. if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 # Verify this is the active session with session_lock: if active_session_id != session_id: return jsonify( {"status": "error", "message": "Not the active session"} ), 403 if model_manager is None: return jsonify({"status": "error", "message": "Model not initialized"}), 400 model_manager.set_scene_context(ctx) # pass player position to NPC if model_manager.npc and ctx: px = ctx.get('player_x', 0) pz = ctx.get('player_z', 0) model_manager.npc.set_target(px, pz) return jsonify({"status": "success"}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 # ── NPC API ── @app.route("/api/npc/spawn", methods=["POST"]) def npc_spawn(): """Spawn NPC.""" try: data = request.json or {} npc_type = data.get("type", "man") if model_manager is None: return jsonify({"status": "error", "message": "Model not initialized"}), 400 # already spawning — silently ignore (Lock-based) if hasattr(model_manager, '_npc_lock') and model_manager._npc_lock.locked(): return jsonify({"status": "busy", "message": "NPC loading"}) model_manager.spawn_npc(npc_type) return jsonify({"status": "success", "type": npc_type}) except Exception as e: print(f"[NPC spawn error] {type(e).__name__}: {e}") traceback.print_exc() return jsonify({"status": "error", "message": f"{type(e).__name__}: {e}"}), 500 @app.route("/api/npc/command", methods=["POST"]) def npc_command(): """Change NPC behavior.""" try: data = request.json or {} behavior = data.get("behavior", "stop") if model_manager and model_manager.npc: model_manager.npc.set_behavior(behavior) return jsonify({"status": "success", "behavior": behavior}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/npc/despawn", methods=["POST"]) def npc_despawn(): """Remove NPC.""" try: if model_manager: model_manager.despawn_npc() return jsonify({"status": "success"}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/npc/frame", methods=["GET"]) def npc_frame(): """Return NPC frame + state.""" try: if not model_manager or not model_manager.npc: return jsonify({"frames": [], "npc": None}) count = request.args.get("count", 4, type=int) npc = model_manager.npc frames = [] for _ in range(count): frame = npc.frame_buffer.get_frame() if frame is not None: frames.append(frame.tolist()) state = npc.get_state() return jsonify({"frames": frames, "npc": state}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/pause", methods=["POST"]) def pause_generation(): """Pause generation (keeps state for resume)""" try: data = request.json if request.json else {} session_id = data.get("session_id") if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 # Verify this is the active session with session_lock: if active_session_id != session_id: return jsonify( {"status": "error", "message": "Not the active session"} ), 403 if model_manager: model_manager.pause_generation() return jsonify({"status": "success", "message": "Generation paused"}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/resume", methods=["POST"]) def resume_generation(): """Resume generation from paused state""" try: global last_frame_consumed_time data = request.json if request.json else {} session_id = data.get("session_id") if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 # Verify this is the active session with session_lock: if active_session_id != session_id: return jsonify( {"status": "error", "message": "Not the active session"} ), 403 if model_manager is None: return jsonify({"status": "error", "message": "Model not initialized"}), 400 model_manager.resume_generation() # Reset consumption tracking when resuming with consumption_monitor_lock: last_frame_consumed_time = time.time() return jsonify({"status": "success", "message": "Generation resumed"}) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/reset", methods=["POST"]) def reset(): """Reset generation state""" try: global active_session_id, last_frame_consumed_time data = request.json if request.json else {} session_id = data.get("session_id") history_length = data.get("history_length") smoothing_alpha = data.get("smoothing_alpha") # If session_id provided, verify it's the active session if session_id: with session_lock: if active_session_id and active_session_id != session_id: return jsonify( {"status": "error", "message": "Not the active session"} ), 403 if model_manager: model_manager.reset( history_length=history_length, smoothing_alpha=smoothing_alpha ) # Clear the active session with session_lock: if active_session_id == session_id or not session_id: active_session_id = None # Clear consumption tracking with consumption_monitor_lock: last_frame_consumed_time = None print(f"[Session {session_id}] Reset complete, session cleared") return jsonify( { "status": "success", "message": "Reset complete", } ) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/get_frame", methods=["GET"]) def get_frame(): """Get the next frame""" try: global last_frame_consumed_time session_id = request.args.get("session_id") if not session_id: return jsonify( {"status": "error", "message": "session_id is required"} ), 400 if model_manager is None: return jsonify({"status": "error", "message": "Model not initialized"}), 400 count = min(int(request.args.get("count", 8)), 20) # Check if this is the active session or a spectator with session_lock: is_active = active_session_id == session_id if is_active: # Active session: pop frames from generation buffer frames = [] for _ in range(count): joints = model_manager.get_next_frame() if joints is None: break frames.append(joints.tolist()) if frames: with consumption_monitor_lock: last_frame_consumed_time = time.time() return jsonify( { "status": "success", "frames": frames, "buffer_size": model_manager.frame_buffer.size(), } ) else: # Spectator: read from broadcast buffer (non-destructive) after_id = int(request.args.get("after_id", 0)) broadcast = model_manager.get_broadcast_frames(after_id, count) if broadcast: last_id = broadcast[-1][0] frames = [joints.tolist() for _, joints in broadcast] return jsonify( { "status": "success", "frames": frames, "last_id": last_id, "buffer_size": model_manager.frame_buffer.size(), } ) # No frames available (active or spectator) return jsonify( { "status": "waiting", "message": "No frame available yet", "buffer_size": model_manager.frame_buffer.size(), } ) except Exception as e: print(f"Error in get_frame: {e}") traceback.print_exc() return jsonify({"status": "error", "message": str(e)}), 500 @app.route("/api/status", methods=["GET"]) def get_status(): """Get generation status""" try: session_id = request.args.get("session_id") with session_lock: is_active_session = session_id and active_session_id == session_id current_active_session = active_session_id if model_manager is None: return jsonify( { "initialized": False, "buffer_size": 0, "is_generating": False, "is_active_session": is_active_session, "active_session_id": current_active_session, } ) status = model_manager.get_buffer_status() status["initialized"] = True status["is_active_session"] = is_active_session status["active_session_id"] = current_active_session return jsonify(status) except Exception as e: return jsonify({"status": "error", "message": str(e)}), 500 if __name__ == "__main__": parser = argparse.ArgumentParser( description="Flask server for real-time 3D motion generation" ) parser.add_argument( "--model_name", type=str, default="ShandaAI/FloodDiffusionTiny", help="HF Hub model name (default: ShandaAI/FloodDiffusionTiny)", ) parser.add_argument( "--port", type=int, default=7860, help="Port to run the server on (default: 7860)", ) args = parser.parse_args() model_name_global = args.model_name # Load model eagerly on startup (pre-downloaded in Docker) print(f"Loading model: {model_name_global}") init_model() print("Starting Flask server...") app.run(host="0.0.0.0", port=args.port, debug=False, threaded=True)