"""Released Geometry Forcing backbone under the fixed WorldMem RE10K protocol.""" import json from pathlib import Path import sys import time sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import torch from hydra import compose, initialize_config_dir from omegaconf import OmegaConf, open_dict from algorithms.dfot import DFoTGeometryForcing from utils.hydra_utils import unwrap_shortcuts class NativeGeometryForcing(DFoTGeometryForcing): def _sample_sequence(self, *args, conditions=None, **kwargs): # The final partial keyframe batch still needs 16 camera entries. return super()._sample_sequence( *args, conditions=self._pad_to_max_tokens(conditions), **kwargs ) class GeometryForcingAdapter: def __init__(self, checkpoint, device="cuda"): repo = Path(__file__).resolve().parents[1] overrides = unwrap_shortcuts( ["dataset=realestate10k", "algorithm=dfot_geometry_forcing", "experiment=video_generation", "@diffusion/continuous", "algorithm.compile=false", "algorithm.alignment.alignment_coeff=0", "algorithm.logging.metrics=[]", "dataset.context_length=1", "dataset.frame_skip=1", "dataset.n_frames=301", "algorithm.diffusion.sampling_timesteps=50", "algorithm.diffusion.ddim_sampling_eta=0.0", "algorithm.backbone.use_checkpointing=[false,false,false,false]", "algorithm.tasks.prediction.keyframe_density=0.0625", "algorithm.tasks.prediction.sliding_context_len=8", "algorithm.tasks.prediction.history_guidance.name=stabilized_vanilla", "+algorithm.tasks.prediction.history_guidance.guidance_scale=4.0", "+algorithm.tasks.prediction.history_guidance.stabilization_level=0.02", "+algorithm.tasks.prediction.history_guidance.visualize=false", "algorithm.tasks.interpolation.history_guidance.name=vanilla", "+algorithm.tasks.interpolation.history_guidance.guidance_scale=1.5", "+algorithm.tasks.interpolation.history_guidance.visualize=false", "algorithm.tasks.interpolation.max_batch_size=4"], str(repo / "configurations"), "config", ) with initialize_config_dir(config_dir=str(repo / "configurations"), version_base=None): cfg = compose(config_name="config", overrides=overrides) with open_dict(cfg.algorithm): cfg.algorithm._name = "dfot_geometry_forcing" self.model = NativeGeometryForcing(cfg.algorithm) checkpoint = Path(checkpoint) payload = torch.load(checkpoint, map_location="cpu", weights_only=False, mmap=True) state = payload.get("state_dict", payload) prefix = "diffusion_model.model." state = {key.replace("diffusion_model._orig_mod.", "diffusion_model."): value for key, value in state.items()} backbone = {key[len(prefix):]: value for key, value in state.items() if key.startswith(prefix)} self.model.diffusion_model.model.load_state_dict(backbone, strict=True) self.model.requires_grad_(False).eval().to(device) parameters = list(self.model.parameters()) self.provenance = { "method": "Geometry Forcing", "repository": "https://github.com/CIntellifusion/GeometryForcing", "commit": "f5dfc6c5c3cbd2c0dad0121c0f7b161d2886e683", "checkpoint": str(checkpoint), "checkpoint_source": "https://huggingface.co/Haoyuwu/GeometryForcing/blob/main/geometry_forcing_state_dict.ckpt", "checkpoint_bytes": checkpoint.stat().st_size, "weight_policy": "released state_dict; no EMA substitution", "checkpoint_pretrained_ema": payload.get("pretrained_ema"), "loaded_backbone_tensors": len(backbone), "discarded_checkpoint_keys": [key for key in state if not key.startswith(prefix)], "parameter_tensors": len(parameters), "parameters": sum(p.numel() for p in parameters), "trainable_tensors": 0, "trainable_parameters": 0, "frozen_parameters": sum(p.numel() for p in parameters), "optimizer_groups": [], "lr_stages": [], "adapters_lora": [], "activation_checkpointing": False, "batch_size_per_gpu": 1, "sampler": "official ContinuousDiffusion DDIM", "sampling_steps": 50, "ddim_eta": 0.0, "history_guidance": {"prediction": {"scale": 4.0, "stabilization": 0.02}, "interpolation": {"scale": 1.5}}, "generation_workflow": "released _predict_videos: sparse keyframes followed by interpolation", "keyframe_density": 0.0625, "keyframe_sliding_context": 8, "interpolation_max_batch_size": 4, "selector": "last observed RGB only; subsequent context consists of generated keyframes", "model_window": 16, "active_window": "native plan, up to 16; not a 12-slot comparison", "padding": "native short-batch noise padding; repeat final camera to 16 entries", "observed_rgb_consumed": 1, "observed_anchor_timeline": 99, "observed_frames": 100, "generated_frames_full": 300, "resolution": [256, 256], "trajectory": "G0..G199,G199..G0", "stride": 1, "text_conditioning": None, "vae": None, "training_only_alignment": "disabled; not involved in generation and not present in released inference backbone", "deviation_from_official_demo": "common scenes and cameras; anchor99 + 300 future frames instead of first-image + 255; common metrics retained", "inference_precision": "float16 autocast; fp32 camera processing and model weights", "torch": torch.__version__, } self.resolved_config = OmegaConf.to_container(cfg, resolve=True) @torch.inference_mode() def generate(self, observed_rgb, cameras, num_generated=300, seed=0, on_frame=None): """RGB is 100 x 3 x 256 x 256 in [0,1]; cameras are timeline-aligned K+w2c.""" device = self.model.device torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) self.model.generator = torch.Generator(device=device).manual_seed(seed) length = num_generated + 1 self.model.cfg.n_frames = length xs = torch.zeros((1, length, *self.model.x_shape), device=device) # Future RGB is never provided to the native prediction/interpolation workflow. xs[:, :1] = self.model._normalize_x(observed_rgb[-1:].to(device)[None]) conditions = cameras[99:100 + num_generated].to(device)[None] start = time.monotonic() with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=device.type == "cuda"): sample = self.model._predict_videos(xs, conditions=conditions) prediction = self.model._unnormalize_x(sample[:, 1:])[0].float() if prediction.shape != (num_generated, 3, 256, 256) or not torch.isfinite(prediction).all(): raise ValueError("Geometry Forcing did not produce the complete finite RGB suffix") prediction = prediction.clamp(0, 1).cpu() self.provenance["keyframe_timeline_indices"] = (torch.linspace(0, length - 1, round(length / 16)).round().long() + 99).tolist() print(json.dumps({"generated": num_generated, "workflow": "native_keyframes_interpolation", "elapsed_seconds": round(time.monotonic() - start, 2)}), flush=True) if on_frame is not None: for offset, rgb in enumerate(prediction): on_frame(offset, rgb) return prediction def main(): import argparse root = Path(__file__).resolve().parents[2] sys.path.insert(0, str(root / "shared")) from protocol import case_seed, load_re10k_case, select_cases, save_case parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", type=Path, default=root / "checkpoints/geometry-forcing/geometry_forcing_state_dict.ckpt") parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--data-root", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--scope", choices=["monitor8", "final100"], default="monitor8") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--limit", type=int, default=1) parser.add_argument("--max-generated", type=int, default=300) parser.add_argument("--rank", type=int, default=0) parser.add_argument("--world-size", type=int, default=1) args = parser.parse_args() rows = select_cases(args.manifest, args.scope, rank=args.rank, world_size=args.world_size, limit=args.limit) adapter = GeometryForcingAdapter(args.checkpoint) args.output.mkdir(parents=True, exist_ok=True) (args.output / f"provenance_rank{args.rank}.json").write_text(json.dumps(adapter.provenance, indent=2)) (args.output / f"config_rank{args.rank}.json").write_text(json.dumps(adapter.resolved_config, indent=2)) for row in rows: case = load_re10k_case(row, args.data_root) seed = case_seed(row, args.seed) prediction = adapter.generate(case["rgb"][:100], case["cameras"], args.max_generated, seed) save_case(args.output, row, prediction, {**adapter.provenance, "seed": seed, "generated_frames": args.max_generated}) if __name__ == "__main__": main()