from __future__ import annotations import hashlib import json from dataclasses import dataclass from collections.abc import Callable from pathlib import Path import numpy as np import torch from diffusers.models import AutoencoderKLWan from diffusers.utils import export_to_video, load_image from transformers import AutoTokenizer, UMT5EncoderModel from .output import sha256, save_chunk_state, assemble_resumed_video from .diffusers import ( WorldCrafterPipeline, WorldCrafterScheduler, WorldCrafterTransformer3DModel, ) from .kernels import ( replace_all_norms_with_flash_norms, replace_rmsnorm_with_fp32, replace_rope_with_flash_rope, ) from .repencoder import ( RepEncoder, RepEncoderInferenceMemoryProvider, RepEncoderInferenceProviderConfig, ) from .ucpe.bridge import ( enable_ucpe_inference_sdpa_attention, load_ucpe_camera_adapter_weights, patch_worldcrafter_transformer_ucpe, ) CAMERA_CHUNK_FRAMES = 33 MODEL_HEIGHT = 384 MODEL_WIDTH = 640 @dataclass(frozen=True) class InferenceResult: video_path: Path summary_path: Path summary: dict[str, object] def load_camera(path: Path, num_chunks: int | None = None) -> np.ndarray: pose = np.load(path, allow_pickle=False) if pose.ndim == 3: pose = pose[None] if pose.ndim != 4 or pose.shape[-2:] not in ((3, 4), (4, 4)): raise ValueError(f"camera must be [B,T,3,4] or [B,T,4,4], got {pose.shape}") if pose.shape[0] != 1 or pose.shape[1] % CAMERA_CHUNK_FRAMES: raise ValueError( "WorldCrafter requires one camera trajectory containing complete 33-frame chunks" ) if not np.issubdtype(pose.dtype, np.floating) or not np.isfinite(pose).all(): raise ValueError("camera must contain finite floating-point c2w matrices") rotation = np.asarray(pose[..., :3, :3], dtype=np.float64) gram = np.swapaxes(rotation, -1, -2) @ rotation if not np.allclose(gram, np.eye(3), atol=5e-3, rtol=0.0): raise ValueError("camera rotations are not orthonormal") if not np.allclose(np.linalg.det(rotation), 1.0, atol=5e-3, rtol=0.0): raise ValueError("camera rotations must have determinant +1") if num_chunks is not None: frames = int(num_chunks) * CAMERA_CHUNK_FRAMES if pose.shape[1] < frames: raise ValueError( f"camera has {pose.shape[1] // CAMERA_CHUNK_FRAMES} chunks, " f"but {num_chunks} were requested" ) pose = pose[:, :frames] return np.ascontiguousarray(pose) def validate_weights(model_path: Path) -> dict[str, Path]: root = model_path.expanduser().resolve() config_path = root / "inference_config.json" config = json.loads(config_path.read_text()) if config_path.is_file() else {} shared = (root / config.get("shared_components", ".")).resolve() required = { "root": root, "transformer": root / "transformer", "adapter": root / "adapter", "repencoder": shared / "repencoder", "vae": shared / "vae", "scheduler": shared / "scheduler", "text_encoder": shared / "text_encoder", "tokenizer": shared / "tokenizer", } missing = [str(path) for path in required.values() if not path.exists()] for filename in ( required["adapter"] / "camera_adapter.pth", required["adapter"] / "pytorch_lora_weights.safetensors", required["repencoder"] / "model.safetensors", required["repencoder"] / "config.json", required["repencoder"] / "manifest.json", ): if not filename.is_file(): missing.append(str(filename)) if missing: raise FileNotFoundError( "WorldCrafter-Base is incomplete: " + ", ".join(missing) ) return required def configure_attention( transformer: WorldCrafterTransformer3DModel, backend: str ) -> str: if backend != "auto": transformer.set_attention_backend(backend) return backend for candidate in ("native", "_flash_3_hub", "flash_hub"): try: transformer.set_attention_backend(candidate) return candidate except (ImportError, RuntimeError, ValueError): continue raise RuntimeError("no supported attention backend is available") def load_model_adapter( pipe: WorldCrafterPipeline, adapter_path: Path ) -> dict[str, object]: from diffusers.loaders.peft import _SET_ADAPTER_SCALE_FN_MAPPING _SET_ADAPTER_SCALE_FN_MAPPING.setdefault( WorldCrafterTransformer3DModel.__name__, lambda _model_class, weights: weights ) state = WorldCrafterPipeline.lora_state_dict(str(adapter_path)) transformer_keys = [key for key in state if key.startswith("transformer.")] if not transformer_keys: raise RuntimeError("adapter does not contain transformer low-rank weights") name = "worldcrafter" pipe.load_lora_weights(str(adapter_path), adapter_name=name) pipe.set_adapters([name], adapter_weights=[1.0]) return {"name": name, "tensor_keys": len(transformer_keys)} class WorldCrafter: def __init__( self, *, pipeline: WorldCrafterPipeline, memory_provider: RepEncoderInferenceMemoryProvider, model_path: Path, device: torch.device, attention_backend: str, adapter_load: dict[str, object], height: int, width: int, ) -> None: self.model_type = "base" self.pipeline = pipeline self.memory_provider = memory_provider self.model_path = model_path self.device = device self.attention_backend = attention_backend self.adapter_load = adapter_load self.height = height self.width = width @classmethod def from_pretrained( cls, model_path: Path, *, model_type: str = "base", device: str = "cuda:0", height: int = MODEL_HEIGHT, width: int = MODEL_WIDTH, seed: int = 42, memory_fov_h_deg: float = 100.0, memory_fov_v_deg: float = 71.13349068444832, memory_fov_samples_per_axis: int = 10, attention_backend: str = "native", enable_compile: bool = False, ) -> "WorldCrafter": if model_type == "fast": from .model_loading import load_fast return load_fast( cls, model_path, device=device, height=height, width=width, seed=seed, memory_fov_h_deg=memory_fov_h_deg, memory_fov_v_deg=memory_fov_v_deg, memory_fov_samples_per_axis=memory_fov_samples_per_axis, attention_backend=attention_backend, enable_compile=enable_compile, ) if model_type != "base": raise ValueError(f"Unknown model type: {model_type}") torch_device = torch.device(device) if torch_device.type != "cuda" or not torch.cuda.is_available(): raise RuntimeError("WorldCrafter inference requires CUDA") if (height, width) != (MODEL_HEIGHT, MODEL_WIDTH): raise ValueError("WorldCrafter-Base is fixed to 384x640 inference") torch.cuda.set_device(torch_device) paths = validate_weights(model_path) enable_ucpe_inference_sdpa_attention() repencoder = RepEncoder.from_pretrained( paths["repencoder"], device=torch_device, compute_dtype="bf16", target_microbatch=4, ) memory_provider = RepEncoderInferenceMemoryProvider( repencoder, RepEncoderInferenceProviderConfig( seed=seed, trajectory_fov_horizontal_fov_degrees=memory_fov_h_deg, trajectory_fov_vertical_fov_degrees=memory_fov_v_deg, trajectory_fov_samples_per_axis=memory_fov_samples_per_axis, ), ) transformer = WorldCrafterTransformer3DModel.from_pretrained( paths["transformer"], torch_dtype=torch.bfloat16 ) patch_worldcrafter_transformer_ucpe( transformer=transformer, method="relray_absmap", height=height, width=width, attn_compress=8, adaptation_method="parallel", ) camera_adapter = load_ucpe_camera_adapter_weights(transformer, paths["adapter"]) if ( camera_adapter["loaded_tensor_keys"] != camera_adapter["expected_tensor_keys"] ): raise RuntimeError(f"camera adapter load is incomplete: {camera_adapter}") adapter_dtypes = { parameter.dtype for block in transformer.blocks for parameter in block.cam_self_attn.parameters() } if adapter_dtypes != {torch.float32}: raise RuntimeError(f"camera adapter dtypes are invalid: {adapter_dtypes}") if not enable_compile: transformer = replace_rmsnorm_with_fp32(transformer) transformer = replace_all_norms_with_flash_norms(transformer) replace_rope_with_flash_rope() resolved_backend = configure_attention(transformer, attention_backend) pipeline = WorldCrafterPipeline( tokenizer=AutoTokenizer.from_pretrained(paths["tokenizer"]), text_encoder=UMT5EncoderModel.from_pretrained( paths["text_encoder"], torch_dtype=torch.bfloat16 ), transformer=transformer, vae=AutoencoderKLWan.from_pretrained( paths["vae"], torch_dtype=torch.float32 ), scheduler=WorldCrafterScheduler.from_pretrained(paths["scheduler"]), ) adapter_load = load_model_adapter(pipeline, paths["adapter"]) pipeline = pipeline.to(torch_device) if enable_compile: torch.backends.cudnn.benchmark = True pipeline.text_encoder.compile( mode="max-autotune-no-cudagraphs", dynamic=False ) pipeline.vae.compile(mode="max-autotune-no-cudagraphs", dynamic=False) pipeline.transformer.compile( mode="max-autotune-no-cudagraphs", dynamic=False ) return cls( pipeline=pipeline, memory_provider=memory_provider, model_path=paths["root"], device=torch_device, attention_backend=resolved_backend, adapter_load=adapter_load, height=height, width=width, ) def generate( self, *, mode: str, camera_path: Path, output_path: Path, prompt: str, negative_prompt: str, image_path: Path | None = None, num_chunks: int | None = None, chunk_output_dir: Path | None = None, state_output_dir: Path | None = None, resume_from: Path | None = None, on_chunk_saved: Callable[[int, Path], None] | None = None, stop_after_chunk: int | None = None, num_inference_steps: int | None = None, guidance_scale: float | None = None, seed: int = 42, fps: int = 16, image_noise_sigma_min: float = 0.111, image_noise_sigma_max: float = 0.135, camera_x_fov: float = 100.0, camera_xi: float = 0.0, local_camera_path: Path | None = None, ) -> InferenceResult: if on_chunk_saved is not None and chunk_output_dir is None: raise ValueError("on_chunk_saved requires chunk_output_dir") is_fast = self.model_type == "fast" if num_inference_steps is None: num_inference_steps = 6 if is_fast else 50 if guidance_scale is None: guidance_scale = 1.0 if is_fast else 5.0 if is_fast: if guidance_scale != 1.0 or num_inference_steps != 6: raise ValueError( "Fast requires CFG=1 and six regular steps; the first T2V chunk uses twelve steps" ) if resume_from is not None or state_output_dir is not None: raise ValueError("Fast resume/state export is not yet validated") self.pipeline.resident_branches.switch("equal") self.pipeline.resident_branches.switches = 0 self.pipeline.stage_model_trace.clear() self.pipeline.fast_inference_mode = mode elif local_camera_path is not None: raise ValueError("--local-camera-path is only used by fast inference") if mode not in {"i2v", "t2v"}: raise ValueError(f"unsupported mode: {mode}") if not camera_path.is_file(): raise FileNotFoundError(camera_path) if mode == "i2v": if image_path is None or not image_path.is_file(): raise FileNotFoundError(image_path) image = load_image(str(image_path)).resize((self.width, self.height)) else: if image_path is not None: raise ValueError( "text-to-video inference does not accept an input image" ) image = None camera_c2w = load_camera(camera_path, num_chunks=num_chunks) num_frames = int(camera_c2w.shape[1]) total_chunks = num_frames // CAMERA_CHUNK_FRAMES camera = { "c2w": camera_c2w, "x_fov": torch.full( (1,), camera_x_fov, device=self.device, dtype=torch.float32 ), "xi": torch.full((1,), camera_xi, device=self.device, dtype=torch.float32), } if is_fast: if local_camera_path is not None: local_c2w = load_camera(local_camera_path, num_chunks=total_chunks) else: from .ucpe.bridge import _relative_pose_chunk local_c2w = torch.cat( [ _relative_pose_chunk( camera_c2w, chunk_index=k, window_num_frames=33, device=self.device, ) for k in range(total_chunks) ], dim=1, ) camera["pose"] = torch.as_tensor( local_c2w, device=self.device, dtype=torch.float32 ) camera["c2w"] = torch.as_tensor( camera_c2w, device=self.device, dtype=torch.float32 ) if chunk_output_dir is not None: chunk_output_dir.mkdir(parents=True, exist_ok=True) if resume_from is not None and chunk_output_dir is None: raise ValueError( "resuming requires --chunk-output-dir with completed prefix chunks" ) if state_output_dir is not None: state_output_dir.mkdir(parents=True, exist_ok=True) final_chunk_index = total_chunks - 1 if stop_after_chunk is not None: if stop_after_chunk < 0 or stop_after_chunk >= total_chunks: raise ValueError("stop_after_chunk must identify a generated chunk") final_chunk_index = int(stop_after_chunk) run_contract = { "mode": mode, "camera_sha256": sha256(camera_path), "image_sha256": sha256(image_path) if image_path is not None else None, "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(), "negative_prompt_sha256": hashlib.sha256( negative_prompt.encode("utf-8") ).hexdigest(), "num_chunks": total_chunks, "num_inference_steps": num_inference_steps, "guidance_scale": guidance_scale, "seed": seed, "image_noise_sigma_min": image_noise_sigma_min, "image_noise_sigma_max": image_noise_sigma_max, "camera_x_fov": camera_x_fov, "camera_xi": camera_xi, "repencoder_model_sha256": self.memory_provider.runtime.report[ "model_sha256" ], } resume_state: dict[str, object] | None = None prior_history_selection: list[dict[str, object]] = [] if resume_from is not None: if not resume_from.is_file(): raise FileNotFoundError(resume_from) resume_state = torch.load( resume_from, map_location="cpu", weights_only=False ) if resume_state.get("run_contract") != run_contract: raise ValueError( "resume checkpoint does not match this inference run contract" ) prior_history_selection = list(resume_state.get("history_selection", [])) if final_chunk_index < int(resume_state["next_chunk_index"]): raise ValueError("stop_after_chunk precedes the resume point") def save_chunk(chunk_index: int, current_video: torch.Tensor) -> None: if chunk_output_dir is None: return frames = self.pipeline.video_processor.postprocess_video( current_video, output_type="np" )[0] path = chunk_output_dir / f"chunk_{chunk_index:03d}_33f.mp4" export_to_video(frames, str(path), fps=fps) if on_chunk_saved is not None: on_chunk_saved(chunk_index, path) print(f"[worldcrafter] completed {path}", flush=True) def save_state(chunk_index: int, state: dict[str, object]) -> None: if state_output_dir is not None: save_chunk_state( chunk_index, state, state_output_dir=state_output_dir, run_contract=run_contract, history_selection=[ *prior_history_selection, *[ record.to_jsonable() for record in self.memory_provider.render_records ], ], ) self.memory_provider.reset_sequence() with torch.inference_mode(): frames = self.pipeline( prompt=prompt, negative_prompt=negative_prompt, height=self.height, width=self.width, num_frames=num_frames, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=torch.Generator(device=self.device).manual_seed(seed), memory_size=4, history_sizes=[2, 1], num_latent_frames_per_chunk=9, keep_first_frame=True, is_enable_stage2=is_fast, pyramid_num_inference_steps_list=None if is_fast else [2, 2, 2], is_skip_first_chunk=False, is_amplify_first_chunk=False, use_zero_init=False, zero_steps=1, image=image, image_noise_sigma_min=image_noise_sigma_min, image_noise_sigma_max=image_noise_sigma_max, video=None, video_noise_sigma_min=0.111, video_noise_sigma_max=0.135, camera_trajectory=camera, memory_provider=self.memory_provider, callback_on_chunk_end=( save_chunk if chunk_output_dir is not None else None ), callback_on_chunk_state=( save_state if state_output_dir is not None else None ), resume_state=resume_state, stop_after_chunk=final_chunk_index, ).frames[0] start_chunk = ( int(resume_state["next_chunk_index"]) if resume_state is not None else 0 ) expected_records = max(0, final_chunk_index - max(1, start_chunk) + 1) if len(self.memory_provider.render_records) != expected_records: raise RuntimeError( f"expected {expected_records} RepEncoder calls, " f"got {len(self.memory_provider.render_records)}" ) output_path.parent.mkdir(parents=True, exist_ok=True) if is_fast and len(frames) != (final_chunk_index + 1) * CAMERA_CHUNK_FRAMES: raise RuntimeError("Fast output must preserve every decoded RGB frame") if resume_state is None: export_to_video(frames, str(output_path), fps=fps) else: assemble_resumed_video(output_path, chunk_output_dir, final_chunk_index) history_selection = [ *prior_history_selection, *[record.to_jsonable() for record in self.memory_provider.render_records], ] summary = { "format": "worldcrafter_inference_v2", "mode": mode, "model_path": str(self.model_path), "image_path": str(image_path) if image_path is not None else None, "image_sha256": sha256(image_path) if image_path is not None else None, "camera_path": str(camera_path), "camera_sha256": sha256(camera_path), "camera_semantics": "global metric c2w; UCPE chunk-relative poses are derived internally", "output_path": str(output_path), "output_sha256": sha256(output_path), "prompt": prompt, "negative_prompt": negative_prompt, "num_frames": num_frames, "completed_through_chunk": final_chunk_index, "fps": fps, "seed": seed, "num_inference_steps": num_inference_steps, "guidance_scale": guidance_scale, "attention_backend": self.attention_backend, "adapter": self.adapter_load, "repencoder_model_sha256": self.memory_provider.runtime.report[ "model_sha256" ], "resumed_from": str(resume_from) if resume_from is not None else None, "history_selection": history_selection, } if is_fast: first_chunk_steps = 12 if mode == "t2v" else 6 expected_calls = first_chunk_steps + final_chunk_index * 6 if len(self.pipeline.stage_model_trace) != expected_calls: raise RuntimeError( "Fast forward count differs from the mode-specific DMD contract" ) summary.update( model_type="fast", local_camera_path=str(local_camera_path), local_camera_sha256=( sha256(local_camera_path) if local_camera_path else None ), first_chunk_inference_steps=first_chunk_steps, first_chunk_routing="4+8" if mode == "t2v" else "5+1", subsequent_chunk_routing="2+4" if mode == "t2v" else "5+1", fast=self.fast_report, stage_model_trace=self.pipeline.stage_model_trace, resident_branch_switches=self.pipeline.resident_branches.switches, ) summary_path = output_path.with_suffix(".json") summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") print(f"[worldcrafter] saved {output_path}", flush=True) return InferenceResult(output_path, summary_path, summary) __all__ = ["InferenceResult", "WorldCrafter", "load_camera", "sha256"]