File size: 9,511 Bytes
59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba c3781c5 59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """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()
|