#!/usr/bin/env python3 """ Unified single-chunk inference for all Echo-Memory memory families. Supports: no_memory, context_k1/k5/k20, framepack_weight, framepack_len_r2/r4, framepack_hybrid_r2/r4, spatial_mem, spatial_concat_text, spatial_inject_none, spatial_cross_attn_readout, videossm_hybrid, block_wise_ssm. Memory type can be specified explicitly via --memory_type or auto-detected from the checkpoint path (--memory_type auto). Examples: # Auto-detect memory type from checkpoint path python inference/unified_inference.py \ --ckpt ./ckpts/context_k1/epoch-0.safetensors \ --prompt "A toy bear on a table, the camera rotates around it" \ --output_path output.mp4 # Explicit memory type python inference/unified_inference.py \ --ckpt ./ckpts/my_checkpoint.safetensors \ --memory_type context_k1 \ --prompt "A scene" \ --output_path output.mp4 # With context image (first frame conditioning) python inference/unified_inference.py \ --ckpt ./ckpts/context_k1/epoch-0.safetensors \ --context_image assets/opendomain_revisit/1774363417.png \ --action_path env/action_rotation_left_45.json \ --prompt "A toy bear on a table" \ --output_path output.mp4 """ from __future__ import annotations import argparse import os import sys # Ensure repo root is in sys.path _script_dir = os.path.dirname(os.path.abspath(__file__)) _repo_root = os.path.abspath(os.path.join(_script_dir, "..")) if _repo_root not in sys.path: sys.path.insert(0, _repo_root) # memory_baseline_runtime has no heavy deps — safe to import at module level from env.memory_baseline_runtime import ( MemoryProfile, MEMORY_PROFILE_REGISTRY, infer_memory_profile_spec, ) # --------------------------------------------------------------------------- # Memory type → profile mapping # --------------------------------------------------------------------------- # Friendly name → registry profile_id _REGISTRY_ALIAS = { "no_memory": "no_memory_extra_two_chunk", "framepack_weight": "framepack_weight_only", "framepack_len_r2": "framepack_lencompress_r2", "framepack_len_r4": "framepack_lencompress_r4", "framepack_hybrid_r2": "framepack_hybrid_r2_weight_two_chunk", "framepack_hybrid_r4": "framepack_hybrid_r4_weight_two_chunk", "spatial_mem": "spatial_mem", "spatial_concat_text": "spatial_concat_text_two_chunk", "spatial_inject_none": "spatial_inject_none_two_chunk", "spatial_cross_attn_readout": "spatial_cross_attn_readout_two_chunk", "geometry_spatial_mem": "geometry_spatial_mem", "videossm_hybrid": "videossm_hybrid_legacy", "block_wise_ssm": "block_wise_ssm_two_chunk", "cgla_memory": "cgla_memory_two_chunk", "prope_memory": "prope_memory_two_chunk", "ucpe_memory": "ucpe_memory_two_chunk", } # context_k* are not in the registry; they use default pipe flags # with only context_override differing _CONTEXT_K_PROFILES = { "context_k1": MemoryProfile(context_override=1), "context_k5": MemoryProfile(context_override=5), "context_k20": MemoryProfile(context_override=20), } # Build profile_id → MemoryProfile lookup from the registry _REGISTRY_PROFILES = {spec.profile_id: spec.profile for spec in MEMORY_PROFILE_REGISTRY} ALL_MEMORY_TYPES = ["auto"] + sorted( set(_REGISTRY_ALIAS.keys()) | set(_CONTEXT_K_PROFILES.keys()) ) def resolve_memory_profile(memory_type: str, ckpt_path: str) -> MemoryProfile: """Resolve --memory_type to a MemoryProfile.""" if memory_type == "auto": spec = infer_memory_profile_spec(ckpt_path) if spec is None: print( f"[unified_inference] WARNING: --memory_type=auto but checkpoint path " f"does not match any known memory profile. Running with no memory flags.\n" f" ckpt: {ckpt_path}\n" f" Hint: use --memory_type to specify explicitly.", file=sys.stderr, flush=True, ) return MemoryProfile() print(f"[unified_inference] Auto-detected memory profile: {spec.profile_id}") return spec.profile if memory_type in _CONTEXT_K_PROFILES: print(f"[unified_inference] Using context learning profile: {memory_type}") return _CONTEXT_K_PROFILES[memory_type] if memory_type in _REGISTRY_ALIAS: profile_id = _REGISTRY_ALIAS[memory_type] profile = _REGISTRY_PROFILES[profile_id] print(f"[unified_inference] Using memory profile: {memory_type} ({profile_id})") return profile print( f"[unified_inference] ERROR: unknown --memory_type '{memory_type}'. " f"Available: {', '.join(ALL_MEMORY_TYPES)}", file=sys.stderr, ) sys.exit(1) def apply_profile_to_pipe(pipe, profile: MemoryProfile) -> None: """Apply a MemoryProfile directly to the pipeline object.""" pipe.use_framepack_memory = bool(profile.use_framepack_memory) pipe.context_temporal_decay = float(profile.context_temporal_decay or 1.0) pipe.context_attention_weight = float(profile.context_attention_weight or 1.0) pipe.use_framepack_length_compress = bool(profile.use_framepack_length_compress) pipe.framepack_ratio = int(profile.framepack_ratio or 2) pipe.use_spatial_memory = bool(profile.use_spatial_memory) pipe.spatial_memory_tokens = int(profile.spatial_memory_tokens or 64) if profile.spatial_memory_inject_mode: pipe.spatial_memory_inject_mode = str(profile.spatial_memory_inject_mode) pipe.use_spatial_memory_legacy = bool(profile.use_spatial_memory_legacy) pipe.use_geometry_spatial_memory = bool(profile.use_geometry_spatial_memory) if profile.geometry_spatial_memory_inject_mode: pipe.geometry_spatial_memory_inject_mode = str( profile.geometry_spatial_memory_inject_mode ) pipe.use_block_wise_ssm = bool(getattr(profile, "use_block_wise_ssm", False)) pipe.use_videossm_hybrid = bool(getattr(profile, "use_videossm_hybrid", False)) pipe.use_cgla_memory = bool(getattr(profile, "use_cgla_memory", False)) # Warn if spatial memory requested but module not loaded from checkpoint if ( pipe.use_spatial_memory and not pipe.use_spatial_memory_legacy and getattr(pipe, "spatial_memory_module", None) is None ): raise RuntimeError( "Spatial token-grid profile requested, but checkpoint has no " "spatial_memory_module weights. Refusing to silently substitute the " "legacy adaptive pool; select an explicit legacy profile instead." ) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Unified single-chunk inference for all Echo-Memory memory families.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Memory types: auto Auto-detect from checkpoint path no_memory No memory (I2V floor baseline) context_k1/k5/k20 Raw context with 1/5/20 frames framepack_weight FramePack temporal decay reweighting framepack_len_r2/r4 FramePack length compression ratio 2/4 framepack_hybrid_r2/r4 FramePack hybrid (length + weight) spatial_mem Spatial grid memory (64 tokens) spatial_concat_text Spatial memory via text KV concatenation spatial_inject_none Spatial memory with withheld read-out spatial_cross_attn_readout Spatial memory via cross-attention videossm_hybrid Legacy VideoSSM hybrid (temporal-conv baseline) block_wise_ssm Block-wise recurrent SSM (paper-aligned) cgla_memory Camera-guided linear attention (pose-conditioned GLA) prope_memory PRoPE: CGLA + camera-pose rotary PE (use_pose_rope) ucpe_memory UCPE: PRoPE + absolute-orientation camera encoding """, ) # Required parser.add_argument("--ckpt", type=str, required=True, help="Path to fine-tuned .safetensors checkpoint") parser.add_argument("--prompt", type=str, required=True, help="Text prompt describing the scene") parser.add_argument("--output_path", type=str, required=True, help="Output video path (.mp4)") # Memory selection parser.add_argument("--memory_type", type=str, default="auto", choices=ALL_MEMORY_TYPES, help="Memory type (default: auto-detect from checkpoint path)") # Model paths parser.add_argument("--base_model", type=str, default=os.environ.get("WAN_BASE_MODEL", ""), help="Wan2.1 base model directory (default: $WAN_BASE_MODEL)") parser.add_argument("--tokenizer_path", type=str, default=None, help="Local tokenizer path (default: /google/umt5-xxl when present)") # Context image parser.add_argument("--context_image", type=str, default=None, help="Path to first-frame context image (enables context memory)") parser.add_argument( "--geometry_memory_video", type=str, default=None, help="TSDF/point-cloud-rendered static condition video for geometry-grounded Spatial Memory", ) # Action control parser.add_argument("--action_path", type=str, default=None, help="Path to action JSON file (81-frame camera trajectory)") # Generation parameters parser.add_argument("--height", type=int, default=352) parser.add_argument("--width", type=int, default=640) parser.add_argument("--num_frames", type=int, default=81, help="Number of frames per chunk (default: 81)") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--num_inference_steps", type=int, default=50) parser.add_argument("--sigma_shift", type=float, default=15.0) parser.add_argument("--cfg_scale", type=float, default=5.0) parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt (default: standard quality filter)") # Output parser.add_argument("--fps", type=int, default=15, help="Output video FPS") return parser def main(): args = build_parser().parse_args() # ── Validate paths ────────────────────────────────────────────────── # Heavy imports deferred so --help works without GPU/conda environment import torch from PIL import Image from env.loop_utils import load_pipeline_and_ckpt, DEFAULT_NEGATIVE_PROMPT from env.run_replay_loop_two_chunk import run_one_chunk, encode_context_frames_per_frame from diffsynth import save_video neg_prompt = args.negative_prompt if args.negative_prompt else DEFAULT_NEGATIVE_PROMPT if not args.base_model: print("ERROR: --base_model or $WAN_BASE_MODEL must be set.", file=sys.stderr) sys.exit(1) dit_path = os.path.join(args.base_model, "diffusion_pytorch_model.safetensors") text_encoder_path = os.path.join(args.base_model, "models_t5_umt5-xxl-enc-bf16.pth") vae_path = os.path.join(args.base_model, "Wan2.1_VAE.pth") tokenizer_path = args.tokenizer_path or os.path.join(args.base_model, "google", "umt5-xxl") if not os.path.isdir(tokenizer_path): tokenizer_path = None for p in [dit_path, text_encoder_path, vae_path]: if not os.path.isfile(p): print(f"ERROR: base model file not found: {p}", file=sys.stderr) sys.exit(1) if not os.path.isfile(args.ckpt): print(f"ERROR: checkpoint not found: {args.ckpt}", file=sys.stderr) sys.exit(1) # ── Resolve memory profile ────────────────────────────────────────── profile = resolve_memory_profile(args.memory_type, args.ckpt) # ── Load pipeline + checkpoint ────────────────────────────────────── print(f"[unified_inference] Loading pipeline from {args.base_model}") print(f"[unified_inference] Loading checkpoint from {args.ckpt}") pipe = load_pipeline_and_ckpt( ckpt_path=args.ckpt, dit_path=dit_path, text_encoder_path=text_encoder_path, vae_path=vae_path, device="cuda", add_action_attn=False, action_use_temporal_attention=True, tokenizer_path=tokenizer_path, ) # ── Apply memory flags ────────────────────────────────────────────── apply_profile_to_pipe(pipe, profile) if args.memory_type == "geometry_spatial_mem" or args.geometry_memory_video: if getattr(pipe, "geometry_spatial_memory_module", None) is None: print( "ERROR: geometry Spatial Memory checkpoint does not contain " "geometry_spatial_memory_module weights.", file=sys.stderr, ) sys.exit(1) pipe.use_geometry_spatial_memory = True # ── Encode context image (if provided) ────────────────────────────── context_latents = None context_actions_t = None num_context_frames = 0 if args.context_image: if not os.path.isfile(args.context_image): print(f"ERROR: context image not found: {args.context_image}", file=sys.stderr) sys.exit(1) print(f"[unified_inference] Encoding context image: {args.context_image}") ctx_pil = Image.open(args.context_image).convert("RGB").resize( (args.width, args.height), Image.LANCZOS ) pipe.load_models_to_device(["vae"]) with torch.no_grad(): context_latents = encode_context_frames_per_frame( pipe, [ctx_pil], pipe.device ) num_context_frames = 1 # Identity RT for context frame (no relative pose change) identity_rt = [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0] context_actions_t = torch.tensor([identity_rt], dtype=torch.float32) geometry_memory_latents = None if args.geometry_memory_video: if not os.path.isfile(args.geometry_memory_video): print( f"ERROR: geometry memory video not found: {args.geometry_memory_video}", file=sys.stderr, ) sys.exit(1) import imageio.v3 as iio geometry_frames = [ Image.fromarray(frame).convert("RGB").resize( (args.width, args.height), Image.Resampling.LANCZOS, ) for frame in iio.imiter(args.geometry_memory_video) ] if not geometry_frames: print("ERROR: geometry memory video contains no frames.", file=sys.stderr) sys.exit(1) pipe.load_models_to_device(["vae"]) with torch.no_grad(): geometry_video = pipe.preprocess_video(geometry_frames) if geometry_video.dim() == 4: geometry_video = geometry_video.unsqueeze(0) geometry_memory_latents = pipe.vae.encode( [geometry_video[i] for i in range(geometry_video.shape[0])], device=pipe.device, tiled=False, tile_size=None, tile_stride=None, ).to(dtype=pipe.torch_dtype, device=pipe.device) print( "[unified_inference] Encoded geometry memory video: " f"{tuple(geometry_memory_latents.shape)}" ) # ── Generate ──────────────────────────────────────────────────────── print(f"[unified_inference] Generating {args.num_frames} frames @ {args.width}x{args.height}") frames = run_one_chunk( pipe=pipe, prompt=args.prompt, use_negative_prompt=neg_prompt, action_path=args.action_path, context_latents=context_latents, num_context_frames=num_context_frames, context_actions_t=context_actions_t, geometry_memory_latents=geometry_memory_latents, chunk_frames=args.num_frames, h=args.height, w=args.width, seed=args.seed, sigma_shift=args.sigma_shift, num_inference_steps=args.num_inference_steps, cfg_scale=args.cfg_scale, log_prefix="[unified_inference]", ) # ── Save video ────────────────────────────────────────────────────── os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) save_video(frames, args.output_path, fps=args.fps, quality=5) print(f"[unified_inference] Video saved to {args.output_path}") if __name__ == "__main__": main()