| """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): |
| |
| 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) |
| |
| 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() |
|
|
