"""Streaming multi-view VIDEO weather/style transfer = keyframe transfer (pluggable keyframe model) + temporal interpolation (Wan Fun-InP+canny), producing a full-length multi-camera video. STREAMING / CAUSAL DESIGN ------------------------- Videos are treated as a *source* that emits frames in a streaming fashion (see `Mp4CameraSource`; swap it for a CARLA sim source later by reimplementing `read_step`). The re-renderer (`StreamingRerenderer`) is the combination of the per-frame keyframe model and the interpolation model. From the outside it looks CAUSAL: it consumes the stream chunk-by-chunk (8 new frames per chunk) and emits 8 new output frames per chunk. Inside a chunk it is non-causal — the LAST frame of the chunk (the next keyframe) is generated first, then the interpolation model fills the frames between the previous keyframe and this new keyframe. This front-facing causal shape is what lets the same code drive a live CARLA stream later. START: read frame 0 -> keyframe model -> output frame 0 (fresh). CHUNK j (frames 8j..8j+8, 9 frames, overlap 1 with prev chunk): 1. keyframe model on frame 8j+8 (the chunk's LAST frame), SDEdit-anchored to keyframe 0 for cross-time consistency; 2. interpolation model fills frames (8j, 8j+8] from prev keyframe + new keyframe + clip canny; 3. emit the 8 new output frames (the overlap frame == prev keyframe, already emitted). KEYFRAME MODELS (--kf-model) — the only part that differs between the two variants; everything else (streaming loop, interpolation, tone-match, grid, IO) is shared: anchor : 2-view ANCHOR Fun-Control. front_wide (slot 0) is transferred standalone, then each other camera (slot 1) is transferred using the transferred front_wide as its reference view (cross-view consistency). A SINGLE input video makes the per-other-view phase a no-op. syncfz : pure Wan-T2V syncfz-K4, image-phase, per-camera and INDEPENDENT (no cross-view anchor); consistency comes only from the shared base Gaussian + SDEdit-to-keyframe-0. Both keyframe models expose the same `transfer(...)` interface, so the streaming re-renderer drives either one identically. Tone-match (default on) locks each keyframe's exposure to keyframe 0 to kill inter-keyframe flicker (--no-tone-match to disable; turn off for long-horizon lighting changes). ARGUMENTS --------- --kf-model {anchor,syncfz} which keyframe backend to use (default anchor). --videos v0.mp4 v1.mp4 ... list of per-camera videos. The FIRST is front_wide. A single video runs single-view transfer. --num-frames N N==1 -> save a single output IMAGE; N==-1 -> process the whole video; otherwise process N frames. If N-1 is not a multiple of the 8-frame chunk, extra frames are generated (the source duplicates its last frame past EOS) and the output is cropped back to N. OUTPUT ------ Per-camera prediction mp4s (or pngs) + a grid: column 0 is [input front_wide (top), output front_wide (bottom)] at full size; every other camera is tiled as thumbnails to the right, 5 per column, adding columns as needed and zero-padding to match the front column's height. Run inside the cosmos-predict2-local docker. """ import argparse import copy import math import os import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F ANY4D_ROOT = os.path.dirname(os.path.abspath(__file__)) if ANY4D_ROOT not in sys.path: sys.path.insert(0, ANY4D_ROOT) os.chdir(ANY4D_ROOT) MODEL_VIEWS = 2 # anchor model: slot 0 = anchor, slot 1 = target ANCHOR_SLOT, TARGET_SLOT = 0, 1 NUM_FRAMES = 9 # InP clip: 9 raw -> 3 latent; clips overlap by 1 frame CHUNK = NUM_FRAMES - 1 # 8 new frames emitted per streaming chunk ROWS_PER_COL = 5 # grid: non-front cameras tiled 5 per column CROP_HW = (506, 900) TARGET_H, TARGET_W = 704, 1280 SNOW_PROMPT = ( 'Driving on a snow-covered road during heavy snowfall. Thick snowflakes fall through ' 'the air, the road and surroundings are blanketed in deep white snow, trees are dusted ' 'with snow, and the sky is overcast and gray. A high quality video captured by a camera ' 'mounted on a car.' ) NEG_PROMPT = ( 'sunny, clear sky, blue sky, sunshine, dry road, ' '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，' '彩色路面，黄色路面，红色路面，' '最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，' '画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，' '静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走' ) # The InP model was trained with a fixed dummy prompt (text_encoder_path=''); the weather is # already baked into the transferred keyframes, so interpolation only needs neutral text. DUMMY_PROMPT = 'a high quality video of driving, realistic, cinematic' def _init_dist_single_gpu(): if torch.distributed.is_initialized(): return import random torch.distributed.init_process_group( backend='gloo', init_method=f'tcp://127.0.0.1:{random.randint(29000, 30000)}', world_size=1, rank=0) # --------------------------------------------------------------------------- # # Frame / canny helpers # --------------------------------------------------------------------------- # def _crop_resize(img_3hw): _, H0, W0 = img_3hw.shape ch, cw = min(CROP_HW[0], H0), min(CROP_HW[1], W0) top, left = (H0 - ch) // 2, (W0 - cw) // 2 img = img_3hw[:, top:top + ch, left:left + cw] img = F.interpolate(img.unsqueeze(0), size=(TARGET_H, TARGET_W), mode='bilinear', align_corners=False) return img.squeeze(0).clamp(0.0, 1.0) def _canny_pixels(rgb_3hw, lo, hi): """(3,H,W) rgb[0,1] -> (3,H,W) canny pseudo-RGB in {0,1} (matches training).""" import cv2 img = (rgb_3hw.permute(1, 2, 0).clamp(0, 1).cpu().numpy() * 255.0).astype(np.uint8) gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) edge = cv2.Canny(gray, lo, hi).astype(np.float32) / 255.0 return torch.from_numpy(np.repeat(edge[None], 3, axis=0)) def _match_tone(img_3hw, ref_3hw): """Per-channel mean/std transfer so img adopts ref's global exposure/tone (keeps structure). Applied to every keyframe after the first (locking exposure to keyframe 0) to remove the inter-keyframe brightness flicker; optional via --no-tone-match. NOTE: for long-horizon streaming with genuine lighting changes, locking to frame 0 is undesirable -> turn it off.""" out = img_3hw.clone() for c in range(3): im, isd = img_3hw[c].mean(), img_3hw[c].std() rm, rsd = ref_3hw[c].mean(), ref_3hw[c].std() out[c] = (img_3hw[c] - im) / (isd + 1e-5) * rsd + rm return out.clamp(0, 1) def _enc_img(pipe, img_3hw): """(3,H,W)[0,1] -> VAE latent (1,16,1,Hl,Wl) for the T=1 keyframe pipe.""" x = (img_3hw.to('cuda', torch.bfloat16) * 2 - 1).unsqueeze(0).unsqueeze(2) with torch.no_grad(): return pipe.tokenizer.encode(x) # --------------------------------------------------------------------------- # # Streaming frame source (mp4-backed; swap for CARLA) # --------------------------------------------------------------------------- # class Mp4CameraSource: """Emits preprocessed multi-camera frames one timestep at a time. This is the *source* the re-renderer streams from. For real streaming re-rendering (e.g. CARLA) reimplement `read_step` to block on the simulator and return the current frame per camera; the re-renderer above does not care where frames come from. Past end-of-stream (finite mp4) the last real frame is duplicated, which is exactly what the num-frames padding relies on. """ def __init__(self, video_paths): import imageio self.paths = list(video_paths) self.num_cams = len(self.paths) self.readers = [imageio.get_reader(p, 'ffmpeg') for p in self.paths] self._cache = {} # (cam, t) -> (3,H,W) preprocessed self._last_real = {} # cam -> last real preprocessed frame (for EOS duplication) self._total = None def total_frames(self): """Number of real frames available (min across cameras); None if unknown (live stream).""" if self._total is None: lens = [] for r in self.readers: n = None try: n = int(r.count_frames()) except Exception: try: n = int(r.get_length()) except Exception: n = None if n and n > 0 and n != float('inf'): lens.append(n) self._total = min(lens) if len(lens) == self.num_cams else None return self._total def _read_raw(self, cam, t): try: fr = self.readers[cam].get_data(t) except Exception: return None return torch.from_numpy(np.asarray(fr)).permute(2, 0, 1).float() / 255.0 def read_step(self, t): """Return list[(3,H,W)] preprocessed frames for all cameras at timestep t.""" frames = [] for cam in range(self.num_cams): key = (cam, t) if key in self._cache: frames.append(self._cache[key]); continue raw = self._read_raw(cam, t) if raw is None: if cam not in self._last_real: raise SystemExit(f'camera {cam} ({self.paths[cam]}) has no readable frame at t=0') img = self._last_real[cam].clone() # EOS: duplicate last real frame else: img = _crop_resize(raw) self._last_real[cam] = img self._cache[key] = img frames.append(img) return frames # --------------------------------------------------------------------------- # # Keyframe pipe (2-view anchor Fun-Control) — see run_multiview_style_transfer.py # --------------------------------------------------------------------------- # def _kf_video_entries(nv): e = [dict(rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), rgb0_ref=(18, 34, 'load', None), rgb0_edge_control=(34, 50, 'load', None))] for v in range(1, nv): e.append({f'rgb{v}': (0, 16, 'copy:rgb0/load', 'copy:rgb0/load'), f'rgb{v}_input_mask': (16, 17, 'copy:rgb0_input_mask/load', None), f'rgb{v}_output_mask': (17, 18, 'copy:rgb0_output_mask/load', None), f'rgb{v}_ref': (18, 34, 'copy:rgb0_ref/load', None), f'rgb{v}_edge_control': (34, 50, 'copy:rgb0_edge_control/load', None)}) return e def _build_kf_cfg(sn_cutoff, sn_channels, sn_seed): from custom.any4d.a4d_config import Any4DConfig from custom.any4d import a4d_surgery cfg = Any4DConfig( vidar_active=True, any4d_active=True, use_wan_backbone=True, wan_variant='control', num_views=MODEL_VIEWS, video_entries=_kf_video_entries(MODEL_VIEWS), video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', view_emb_std=0.06, block_gate_fix=True, harmonize_streams=True, harmonize_frames=True, load_modals=['rgb'], use_structured_noise=True, structured_noise_channels=sn_channels, structured_noise_seed=sn_seed, val_structured_noise=True, val_structured_noise_cutoff=sn_cutoff) return a4d_surgery.validate_config(cfg) def build_kf_pipe(ckpt, vae_path, sn_cutoff, sn_channels, sn_seed, use_ema=True): from cosmos_predict2.configs.base.config_video2world import PREDICT2_VIDEO2WORLD_PIPELINE_2B from custom.wan.wan_vae_adapter import WanVAEAdapter from custom.wan.a4d_network_wan_control import Any4DWanControlDiT from custom.wan.a4d_pipe_wan_control import build_wan_control_pipeline from imaginaire.lazy_config import LazyCall as L pc = copy.deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B) pc.net = L(Any4DWanControlDiT)(any4d_config=None, view_emb_std=0.06) pc.any4d_active = True pc.tokenizer = L(WanVAEAdapter)(vae_pth=vae_path, chunk_duration=1, name='wan_vae', load_mean_std=False, dtype='bfloat16') pc.state_ch = 16; pc.state_t = 1; pc.conditioner.text.dropout_rate = 0.0 try: pc.ema.enabled = False except Exception: pass if getattr(pc, 'extra_nets', None) is None: pc.extra_nets = nn.ModuleDict() cfg = _build_kf_cfg(sn_cutoff, sn_channels, sn_seed) pipe = build_wan_control_pipeline(pc, dit_path='', text_encoder_path='', vae_path=vae_path, any4d_config=cfg) pipe.any4d_config = cfg sd = torch.load(ckpt, map_location='cpu', weights_only=True, mmap=True) prefix = 'net_ema.' if use_ema else 'net.' sub = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)} assert sub, f'no {prefix}* in {ckpt}' pipe.load_checkpoint(sub, any4d_config=cfg); del sd, sub # NOTE: on this branch (anchor-frontwide) Any4DWanControlDiT._to_wan_48ch natively routes the # ref slot [18:34] -> Wan image-cond channel [32:48], so no monkeypatch is needed here. pipe.dit = pipe.dit.to('cuda', torch.bfloat16).eval() print('[kf] anchor Fun-Control pipe ready (native ref->y routing)') return pipe, cfg def _decode_pixels(pipe, samples, slot): lat = samples['y0_pred_streams'][f'v{slot}'][:, 0:16].to(torch.bfloat16) with torch.no_grad(): raw = pipe.tokenizer.decode(lat) return ((raw.clamp(-1, 1)[:, :, 0] + 1) / 2).squeeze(0).clamp(0, 1).cpu().float() # (3,H,W) def _kf_batch(cfg, views, text, device): from custom.any4d.a4d_logistics import impute_data_batch_masks pos, pos_m, neg, neg_m = text a4d = {} for v, sp in views.items(): rgb, edge, known = sp['rgb'], sp['edge'], sp['known'] ref = sp.get('ref', None) B, C, Tl, Hl, Wl = rgb.shape a4d[f'rgb{v}'] = rgb a4d[f'rgb{v}_ref'] = ref if ref is not None else torch.zeros_like(rgb) a4d[f'rgb{v}_edge_control'] = edge im = torch.full((B, 1, Tl, Hl, Wl), 1.0 if known else 0.0, device=device, dtype=torch.bfloat16) a4d[f'rgb{v}_input_mask'] = im a4d[f'rgb{v}_output_mask'] = 1.0 - im a4d[f'rgb{v}_supervise_mask'] = 1.0 - im db = {'fps': torch.tensor([10.0], dtype=torch.bfloat16, device=device), 'padding_mask': torch.zeros((1, 1, TARGET_H, TARGET_W), dtype=torch.bfloat16, device=device), 'prompt': [''], 'a4d_raw': {}, 'a4d_latent': a4d, 'anydata': {}, 't5_text_embeddings': pos, 't5_text_mask': pos_m, 'neg_t5_text_embeddings': neg, 'neg_t5_text_mask': neg_m} return impute_data_batch_masks(db, cfg) class AnchorKeyframer: """Keyframe backend: 2-view ANCHOR Fun-Control model (cross-view). front_wide (anchor slot) is transferred standalone; each other camera (target slot) is transferred using the transferred front_wide as its reference view. Single input video -> the per-other-view phase is a no-op.""" name = 'anchor' def __init__(self, pipe, cfg): self.pipe, self.cfg = pipe, cfg def transfer(self, cam_rgb, cam_canny, text, steps, guidance, seed, num_cams, prev_lat=None, sdedit_strength=1.0): """cam_rgb/cam_canny: {cam_idx -> (3,H,W)}. SAME seed for every view/keyframe -> shared base Gaussian. If prev_lat is given, each camera's keyframe is SDEdited from it (front_wide via v0 anchor slot, others via v1 target slot). Returns ({cam->pixels}, {cam->full-stream latent}).""" device = torch.device('cuda') pix, lat = {}, {} # Phase 1: front_wide standalone (anchor slot). SDEdit from prev front_wide latent. rgb0_lat = _enc_img(self.pipe, cam_rgb[0]); edge0_lat = _enc_img(self.pipe, cam_canny[0]) db = _kf_batch(self.cfg, {ANCHOR_SLOT: dict(rgb=rgb0_lat, edge=edge0_lat, known=False)}, text, device) sd0 = {f'v{ANCHOR_SLOT}': prev_lat[0]} if prev_lat is not None else None with torch.no_grad(): s = self.pipe.generate(db, seed=seed, num_sampling_steps=steps, guidance=guidance, is_negative_prompt=True, phase='val', sdedit_init_streams=sd0, sdedit_strength=sdedit_strength) anchor_full = s['y0_pred_streams'][f'v{ANCHOR_SLOT}'].detach().clone() # full stream (for next-time SDEdit) anchor_lat = anchor_full[:, 0:16] lat[0] = anchor_full pix[0] = _decode_pixels(self.pipe, s, ANCHOR_SLOT) # Phase 2 (multi-view only): each other camera, transferred front_wide as reference view (v0 # known); target = v1. Empty loop for single-view -> the multi-view code does nothing. for i in range(1, num_cams): rgbi = _enc_img(self.pipe, cam_rgb[i]); edgei = _enc_img(self.pipe, cam_canny[i]) views = {ANCHOR_SLOT: dict(rgb=anchor_lat, edge=edge0_lat, known=True, ref=anchor_lat), TARGET_SLOT: dict(rgb=rgbi, edge=edgei, known=False)} db = _kf_batch(self.cfg, views, text, device) sdi = {f'v{TARGET_SLOT}': prev_lat[i]} if prev_lat is not None else None with torch.no_grad(): s = self.pipe.generate(db, seed=seed, num_sampling_steps=steps, guidance=guidance, is_negative_prompt=True, phase='val', sdedit_init_streams=sdi, sdedit_strength=sdedit_strength) lat[i] = s['y0_pred_streams'][f'v{TARGET_SLOT}'].detach().clone() pix[i] = _decode_pixels(self.pipe, s, TARGET_SLOT) return pix, lat # --------------------------------------------------------------------------- # # Keyframe pipe (pure Wan-T2V, syncfz K=4, image phase) — run_t2i_syncfz_k4.py # --------------------------------------------------------------------------- # def _build_t2v_cfg(sn_channels, sn_seed): from custom.any4d import a4d_surgery from custom.any4d.a4d_config import Any4DConfig # Full-frequency structured noise, phase from the (image-phase) rgb latent. At inference # _t2v_batch fills both the rgb span and the rgb0_edge_control slot with the same rgb latent, # so the phase read is identical either way (val frequency follows full_frequency=True). cfg = Any4DConfig( vidar_active=True, any4d_active=True, use_wan_backbone=True, wan_variant='t2v', num_views=1, video_entries=[dict(rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), rgb0_edge_control=(18, 34, 'load', None))], lowdim_adaln_entries=dict(), video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', block_gate_fix=True, harmonize_streams=True, harmonize_frames=True, load_modals=['rgb'], use_structured_noise=True, structured_noise_channels=sn_channels, structured_noise_seed=sn_seed, structured_noise_full_frequency=True, val_structured_noise=True) return a4d_surgery.validate_config(cfg) def build_t2v_pipe(ckpt, vae_path, sn_channels=4, sn_seed=0, use_ema=True): from cosmos_predict2.configs.base.config_video2world import PREDICT2_VIDEO2WORLD_PIPELINE_2B from custom.wan.wan_vae_adapter import WanVAEAdapter from custom.wan.a4d_network_wan_t2v import Any4DWanT2VDiT from custom.wan.a4d_pipe_wan import build_wan_pipeline from imaginaire.lazy_config import LazyCall as L pc = copy.deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B) pc.any4d_active = True pc.tokenizer = L(WanVAEAdapter)(vae_pth=vae_path, chunk_duration=1, name='wan_vae', load_mean_std=False, dtype='bfloat16') pc.state_ch = 16; pc.state_t = 1; pc.conditioner.text.dropout_rate = 0.0 try: pc.ema.enabled = False except Exception: pass if getattr(pc, 'extra_nets', None) is None: pc.extra_nets = nn.ModuleDict() pc.net = L(Any4DWanT2VDiT)(any4d_config=None) cfg = _build_t2v_cfg(sn_channels, sn_seed) pipe = build_wan_pipeline(pc, dit_path='', text_encoder_path='', vae_path=vae_path, any4d_config=cfg) pipe.any4d_config = cfg prefix = 'net_ema.' if use_ema else 'net.' raw = torch.load(ckpt, map_location='cpu', weights_only=True, mmap=True) sub = {k[len(prefix):]: v for k, v in raw.items() if k.startswith(prefix)} assert sub, f'no {prefix}* in {ckpt}' m, u = pipe.dit.load_state_dict(sub, strict=False) print(f'[t2v] syncfz pipe loaded {len(sub)} (sn_P_raw={any("sn_P_raw" in k for k in sub)}) ' f'missing={len(m)} unexpected={len(u)}') del raw, sub pipe.dit = pipe.dit.to('cuda', torch.bfloat16).eval() return pipe, cfg def _enc1(pipe, img_3hw): x = (img_3hw.to('cuda', torch.bfloat16) * 2 - 1).unsqueeze(1).unsqueeze(0) # (1,3,1,H,W) with torch.no_grad(): return pipe.tokenizer.encode(x) def _t2v_batch(pipe, rgb_3hw, text): device, dtype = torch.device('cuda'), torch.bfloat16 Hl, Wl = TARGET_H // 8, TARGET_W // 8 rgb_lat = _enc1(pipe, rgb_3hw).to(dtype) # target region + (image-phase source) im = torch.zeros((1, 1, 1, Hl, Wl), device=device, dtype=dtype) a4d = {'rgb0': rgb_lat, 'rgb0_input_mask': im, 'rgb0_output_mask': 1 - im, 'rgb0_supervise_mask': 1 - im, 'rgb0_edge_control': rgb_lat, # image phase: fill [18:34] with rgb latent 'rgb0_edge_control_input_mask': torch.ones_like(im), 'rgb0_edge_control_output_mask': torch.zeros_like(im), 'rgb0_edge_control_supervise_mask': torch.zeros_like(im), 'video_dims': [(1, Hl, Wl)]} pos, pos_m, neg, neg_m = text return {'fps': torch.tensor([10.0], dtype=dtype, device=device), 'padding_mask': torch.zeros((1, 1, TARGET_H, TARGET_W), dtype=dtype, device=device), 'prompt': [''], 'a4d_raw': {}, 'a4d_latent': a4d, 'anydata': {}, 't5_text_embeddings': pos, 't5_text_mask': pos_m, 'neg_t5_text_embeddings': neg, 'neg_t5_text_mask': neg_m} def _decode1(pipe, samples): lat = samples['y0_pred_streams']['v0'][:, 0:16].to(torch.bfloat16) with torch.no_grad(): raw = pipe.tokenizer.decode(lat) return ((raw + 1) / 2).clamp(0, 1).squeeze(0)[:, 0].cpu().float() # (3,H,W) class SyncfzKeyframer: """Keyframe backend: pure Wan-T2V syncfz-K4 model, image-phase, per-camera and INDEPENDENT (NO cross-view anchor). Cross-view/keyframe consistency comes only from the shared base Gaussian (same seed) + the SDEdit-to-keyframe-0 chain. Ignores canny (structure is the image phase).""" name = 'syncfz' def __init__(self, pipe): self.pipe = pipe def transfer(self, cam_rgb, cam_canny, text, steps, guidance, seed, num_cams, prev_lat=None, sdedit_strength=1.0): pix, lat = {}, {} for i in range(num_cams): db = _t2v_batch(self.pipe, cam_rgb[i], text) sd_init = {'v0': prev_lat[i]} if prev_lat is not None else None with torch.no_grad(): s = self.pipe.generate(db, seed=seed, num_sampling_steps=steps, guidance=guidance, is_negative_prompt=True, phase='val', cond_aug_sigma=0.0, sdedit_init_streams=sd_init, sdedit_strength=sdedit_strength) lat[i] = s['y0_pred_streams']['v0'].detach().clone() pix[i] = _decode1(self.pipe, s) return pix, lat # --------------------------------------------------------------------------- # # InP interpolation pipe (single-view Fun-InP + canny) # --------------------------------------------------------------------------- # def _build_inp_cfg(): from custom.any4d import a4d_surgery from custom.any4d.a4d_config import Any4DConfig cfg = Any4DConfig( vidar_active=True, any4d_active=True, use_wan_backbone=True, wan_variant='inp_canny', wan_inp_single_end_frame=True, num_views=1, video_entries=[dict(rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), rgb0_ref=(18, 34, 'load', None), rgb0_edge_control=(34, 50, 'load', None))], lowdim_adaln_entries=dict(), video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', block_gate_fix=True, harmonize_streams=True, harmonize_frames=True, load_modals=['rgb']) return a4d_surgery.validate_config(cfg) def build_inp_pipe(ckpt, vae_path, use_ema=True): from cosmos_predict2.configs.base.config_video2world import PREDICT2_VIDEO2WORLD_PIPELINE_2B from custom.wan.wan_vae_adapter import WanVAEAdapter from custom.wan.a4d_network_wan_inp_canny import Any4DWanInPCannyDiT from custom.wan.a4d_pipe_wan_inp_canny import build_wan_inp_canny_pipeline from imaginaire.lazy_config import LazyCall as L state_t = (NUM_FRAMES - 1) // 4 + 1 pc = copy.deepcopy(PREDICT2_VIDEO2WORLD_PIPELINE_2B) pc.any4d_active = True pc.tokenizer = L(WanVAEAdapter)(vae_pth=vae_path, chunk_duration=NUM_FRAMES, name='wan_vae', load_mean_std=False, dtype='bfloat16') pc.state_ch = 16; pc.state_t = state_t; pc.conditioner.text.dropout_rate = 0.0 try: pc.ema.enabled = False except Exception: pass if getattr(pc, 'extra_nets', None) is None: pc.extra_nets = nn.ModuleDict() pc.net = L(Any4DWanInPCannyDiT)(any4d_config=None) cfg = _build_inp_cfg() pipe = build_wan_inp_canny_pipeline(pc, dit_path='', text_encoder_path='', vae_path=vae_path, any4d_config=cfg) pipe.any4d_config = cfg sd = torch.load(ckpt, map_location='cpu', weights_only=True, mmap=True) prefix = 'net_ema.' if use_ema else 'net.' sub = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)} assert sub, f'no {prefix}* in {ckpt}' m, u = pipe.dit.load_state_dict(sub, strict=False); del sd, sub print(f'[inp] Fun-InP+canny pipe ready (missing={len(m)} unexpected={len(u)})') pipe.dit = pipe.dit.to('cuda', torch.bfloat16).eval() return pipe def _inp_encode(pipe, vid_b3thw_01): x = (vid_b3thw_01 * 2 - 1).to('cuda', torch.bfloat16) with torch.no_grad(): return pipe.tokenizer.encode(x) def _inp_ref(pipe, vid_b3thw_01): ref = torch.zeros_like(vid_b3thw_01) ref[:, :, 0] = vid_b3thw_01[:, :, 0] * 2 - 1 ref[:, :, -1] = vid_b3thw_01[:, :, -1] * 2 - 1 with torch.no_grad(): return pipe.tokenizer.encode(ref.to('cuda', torch.bfloat16)) def interpolate_clip(inp_pipe, rgb_clip, canny_clip, first_3hw, last_3hw, text, steps, guidance, seed): """rgb_clip/canny_clip: (T,3,H,W). first/last: transferred (3,H,W). -> (3,T,H,W) [0,1].""" device = torch.device('cuda') rgb_cond = rgb_clip.clone(); rgb_cond[0] = first_3hw; rgb_cond[-1] = last_3hw Tp = rgb_cond.shape[0]; Tl = (Tp - 1) // 4 + 1; Hl, Wl = TARGET_H // 8, TARGET_W // 8 rgb_b = rgb_cond.permute(1, 0, 2, 3).unsqueeze(0) canny_b = canny_clip.permute(1, 0, 2, 3).unsqueeze(0) im = torch.zeros((1, 1, Tl, Hl, Wl), device=device, dtype=torch.bfloat16) im[:, :, 0] = 1.0; im[:, :, -1] = 1.0 a4d = {'rgb0': _inp_encode(inp_pipe, rgb_b).to(torch.bfloat16), 'rgb0_input_mask': im, 'rgb0_output_mask': 1 - im, 'rgb0_supervise_mask': 1 - im, 'rgb0_ref': _inp_ref(inp_pipe, rgb_b).to(torch.bfloat16), 'rgb0_ref_input_mask': torch.ones_like(im), 'rgb0_ref_output_mask': torch.zeros_like(im), 'rgb0_ref_supervise_mask': torch.zeros_like(im), 'rgb0_edge_control': _inp_encode(inp_pipe, canny_b).to(torch.bfloat16), 'rgb0_edge_control_input_mask': torch.ones_like(im), 'rgb0_edge_control_output_mask': torch.zeros_like(im), 'rgb0_edge_control_supervise_mask': torch.zeros_like(im), 'video_dims': [(Tl, Hl, Wl)]} pos, pos_m, neg, neg_m = text db = {'fps': torch.tensor([10.0], dtype=torch.bfloat16, device=device), 'padding_mask': torch.zeros((1, 1, TARGET_H, TARGET_W), dtype=torch.bfloat16, device=device), 'prompt': [''], 'a4d_raw': {}, 'a4d_latent': a4d, 'anydata': {}, 't5_text_embeddings': pos, 't5_text_mask': pos_m, 'neg_t5_text_embeddings': neg, 'neg_t5_text_mask': neg_m} with torch.no_grad(): s = inp_pipe.generate(db, seed=seed, num_sampling_steps=steps, guidance=guidance, is_negative_prompt=True, phase='val') lat = s['y0_pred_streams']['v0'][:, 0:16].to(torch.bfloat16) with torch.no_grad(): raw = inp_pipe.tokenizer.decode(lat) return ((raw.clamp(-1, 1) + 1) / 2).squeeze(0).cpu().float() # (3,T,H,W) # --------------------------------------------------------------------------- # # Streaming re-renderer (per-frame keyframe model + interpolation model) # --------------------------------------------------------------------------- # class StreamingRerenderer: """Front-facing causal, chunk-by-chunk re-renderer. Holds keyframe + interpolation state so the same object can be driven by an offline mp4 loop OR a live CARLA stream. Usage: r = StreamingRerenderer(...) out0 = r.start(kf_rgb0, kf_cny0) # {cam -> (3,H,W)} output frame 0 out = r.process_chunk(clip_rgb, clip_canny) # {cam -> (3,8,H,W)} next 8 output frames """ def __init__(self, keyframer, inp_pipe, kf_text, inp_text, num_cams, *, kf_steps, kf_guidance, inp_steps, inp_guidance, seed, sdedit_strength, tone_match): self.keyframer = keyframer # AnchorKeyframer | SyncfzKeyframer self.inp_pipe = inp_pipe self.kf_text, self.inp_text = kf_text, inp_text self.num_cams = num_cams self.kf_steps, self.kf_guidance = kf_steps, kf_guidance self.inp_steps, self.inp_guidance = inp_steps, inp_guidance self.seed = seed self.sdedit_strength = sdedit_strength self.tone_match = tone_match # rolling state self.first_lat = None # keyframe 0 latent — the fixed SDEdit anchor for every later keyframe self.first_pix = None # keyframe 0 pixels — tone-match reference self.prev_pix = None # previous keyframe pixels (interpolation start frame) def start(self, kf_rgb, kf_cny): """Generate the very first keyframe (fresh). Returns {cam -> (3,H,W)} output frame 0.""" pix, lat = self.keyframer.transfer(kf_rgb, kf_cny, self.kf_text, self.kf_steps, self.kf_guidance, self.seed, self.num_cams, prev_lat=None, sdedit_strength=1.0) self.first_lat, self.first_pix, self.prev_pix = lat, pix, pix return pix def process_chunk(self, clip_rgb, clip_canny): """clip_rgb/clip_canny: {cam -> (9,3,H,W)} for frames [8j : 8j+9]. Generates the END keyframe (chunk's last frame) first, SDEdit-anchored to keyframe 0, then interpolates. Returns {cam -> (3,8,H,W)} (overlap dropped).""" end_rgb = {i: clip_rgb[i][-1] for i in range(self.num_cams)} end_cny = {i: clip_canny[i][-1] for i in range(self.num_cams)} pix_end, _ = self.keyframer.transfer(end_rgb, end_cny, self.kf_text, self.kf_steps, self.kf_guidance, self.seed, self.num_cams, prev_lat=self.first_lat, sdedit_strength=self.sdedit_strength) if self.tone_match: for i in range(self.num_cams): pix_end[i] = _match_tone(pix_end[i], self.first_pix[i]) out = {} for i in range(self.num_cams): clip = interpolate_clip(self.inp_pipe, clip_rgb[i], clip_canny[i], self.prev_pix[i], pix_end[i], self.inp_text, self.inp_steps, self.inp_guidance, self.seed + 1) # (3,9,H,W) out[i] = clip[:, 1:] # drop overlap (== prev keyframe, already emitted) self.prev_pix = pix_end return out # --------------------------------------------------------------------------- # # Save + grid (flexible camera count) # --------------------------------------------------------------------------- # def _save_mp4(path, frames_3thw, fps): import imageio vid = (frames_3thw.clamp(0, 1).permute(1, 2, 3, 0).numpy() * 255).astype(np.uint8) imageio.mimwrite(path, vid, fps=fps, quality=8, macro_block_size=1) def _save_img(path, frame_3hw): import imageio img = (frame_3hw.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8) imageio.imwrite(path, img) def _even(x): """Round down to an even int — libx264/yuv420p requires even frame width & height.""" x = int(round(x)) return x - (x % 2) def _grid_layout(num_cams, front_w): """Geometry for the grid: front column (2 stacked cells) on the left, non-front cameras tiled ROWS_PER_COL per column to the right, zero-padded to the front column's height. All dims are kept even so the resulting W and H are even (required by the libx264 mp4 encoder).""" AR = TARGET_W / TARGET_H fw = _even(front_w) fh = _even(fw / AR) H = 2 * fh others = num_cams - 1 if others <= 0: return dict(AR=AR, fw=fw, fh=fh, H=H, W=fw, sh=0, sw=0) sh = _even(H // ROWS_PER_COL) sw = _even(sh * AR) ncols = math.ceil(others / ROWS_PER_COL) return dict(AR=AR, fw=fw, fh=fh, H=H, W=fw + ncols * sw, sh=sh, sw=sw) def _compose_grid_frame(inp_front_3hw, pred_frames, cam_names, lay): """inp_front_3hw: (3,H,W). pred_frames: {cam -> (3,H,W)}. Returns (H,W,3) uint8.""" from PIL import Image, ImageDraw fw, fh, H, W, sh, sw = lay['fw'], lay['fh'], lay['H'], lay['W'], lay['sh'], lay['sw'] def cell(img_3hw, label, w, h): a = (img_3hw.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8) im = Image.fromarray(a).resize((w, h)) d = ImageDraw.Draw(im) d.rectangle([0, 0, 8 * len(label) + 8, 18], fill=(0, 0, 0)) d.text((3, 3), label, fill=(255, 255, 0)) return np.asarray(im) g = np.zeros((H, W, 3), np.uint8) g[0:fh, 0:fw] = cell(inp_front_3hw, f'input {cam_names[0]}'[:22], fw, fh) g[fh:2 * fh, 0:fw] = cell(pred_frames[0], f'output {cam_names[0]}'[:22], fw, fh) for k in range(len(cam_names) - 1): cam_i = k + 1 col, row = k // ROWS_PER_COL, k % ROWS_PER_COL x0, y0 = fw + col * sw, row * sh g[y0:y0 + sh, x0:x0 + sw] = cell(pred_frames[cam_i], cam_names[cam_i][:12], sw, sh) return g def _grid_video(out_path, inp_front_3thw, preds, cam_names, fps, front_w=1280): import imageio lay = _grid_layout(len(cam_names), front_w) T = preds[0].shape[1] frames = [_compose_grid_frame(inp_front_3thw[:, t], {i: preds[i][:, t] for i in preds}, cam_names, lay) for t in range(T)] imageio.mimwrite(out_path, frames, fps=fps, quality=8, macro_block_size=1) def _grid_image(out_path, inp_front_3hw, preds, cam_names, front_w=1280): import imageio lay = _grid_layout(len(cam_names), front_w) imageio.imwrite(out_path, _compose_grid_frame(inp_front_3hw, preds, cam_names, lay)) # --------------------------------------------------------------------------- # # Reusable pipeline builder + per-sequence runner (also used by batch drivers) # --------------------------------------------------------------------------- # def build_pipes(kf_ckpt, inp_ckpt, wan_dir, vae_path=None, *, kf_model='anchor', sn_cutoff=20.0, sn_channels=4, seed=0, prompt=SNOW_PROMPT, neg_prompt=NEG_PROMPT): """Load the selected keyframe backend + the interpolation pipe and encode the (fixed) text ONCE. The returned dict is reusable across many sequences (see batch_val_transfer.py), which is the whole point of factoring this out of main() — model load is ~19 GB / ~2 min per process. kf_model: 'anchor' = 2-view cross-view Fun-Control; 'syncfz' = pure Wan-T2V per-camera.""" vae_path = vae_path or os.path.join(wan_dir, 'Wan2.1_VAE.pth') _init_dist_single_gpu() if kf_model == 'anchor': keyframer = AnchorKeyframer(*build_kf_pipe(kf_ckpt, vae_path, sn_cutoff, sn_channels, seed, use_ema=True)) elif kf_model == 'syncfz': pipe, _ = build_t2v_pipe(kf_ckpt, vae_path, sn_channels, seed, use_ema=True) keyframer = SyncfzKeyframer(pipe) else: raise SystemExit(f'unknown --kf-model {kf_model!r} (expected anchor|syncfz)') inp_pipe = build_inp_pipe(inp_ckpt, vae_path, use_ema=True) from custom.wan.a4d_pipe_wan_control import (_WanLiveTextEncoder, _find_umt5_tokenizer_dir, _find_umt5_weights) print('[text] UMT5 encode ...') te = _WanLiveTextEncoder(_find_umt5_tokenizer_dir(wan_dir), _find_umt5_weights(wan_dir), device='cuda', dtype=torch.bfloat16) kf_text = (*te.encode_prompts([prompt], 512, True), *te.encode_prompts([neg_prompt], 512, True)) inp_text = (*te.encode_prompts([DUMMY_PROMPT], 512, True), *te.encode_prompts([''], 512, True)) del te; torch.cuda.empty_cache() return dict(keyframer=keyframer, inp_pipe=inp_pipe, kf_text=kf_text, inp_text=inp_text) def run_sequence(pipes, videos, num_frames, out_dir, *, canny_lo=90.0, canny_hi=215.0, kf_steps=35, kf_guidance=5.0, inp_steps=35, inp_guidance=5.0, seed=0, sdedit_strength=1.0, tone_match=True, front_w=1280, fps=10): """Re-render ONE multi-camera sequence with a FRESH renderer (rolling keyframe state must not leak between sequences), reusing the already-loaded `pipes`. Writes per-camera mp4/png + grid + test_input into out_dir. `videos`: per-camera paths, first is front_wide; single video -> single-view. See main() / the module docstring for --num-frames semantics.""" os.makedirs(out_dir, exist_ok=True) in_dir = os.path.join(out_dir, 'test_input'); os.makedirs(in_dir, exist_ok=True) for v in videos: if not os.path.isfile(v): raise SystemExit(f'missing video: {v}') source = Mp4CameraSource(videos) num_cams = source.num_cams cam_names = [os.path.splitext(os.path.basename(v))[0] for v in videos] cam_names[0] = 'front_wide' # first camera is always front_wide mode = 'single-view' if num_cams == 1 else f'{num_cams}-view' print(f'[src] {mode} transfer, cameras={cam_names}') if num_frames == -1: T_target = source.total_frames() if not T_target: raise SystemExit('--num-frames -1 needs a countable video; pass an explicit frame count instead.') else: T_target = num_frames if T_target < 1: raise SystemExit(f'invalid resolved frame count: {T_target}') num_clips = 0 if T_target == 1 else max(1, math.ceil((T_target - 1) / CHUNK)) T_pad = 1 + CHUNK * num_clips print(f'[plan] T_target={T_target} -> {num_clips} chunk(s), generate {T_pad} frames then crop.') renderer = StreamingRerenderer( pipes['keyframer'], pipes['inp_pipe'], pipes['kf_text'], pipes['inp_text'], num_cams, kf_steps=kf_steps, kf_guidance=kf_guidance, inp_steps=inp_steps, inp_guidance=inp_guidance, seed=seed, sdedit_strength=sdedit_strength, tone_match=tone_match) def canny(frame): return _canny_pixels(frame, canny_lo, canny_hi) # ---- START: frame 0 keyframe (fresh) ---- print('[stream] frame 0: keyframe (fresh) ...') f0 = source.read_step(0) kf_rgb0 = {i: f0[i] for i in range(num_cams)} kf_cny0 = {i: canny(f0[i]) for i in range(num_cams)} out0 = renderer.start(kf_rgb0, kf_cny0) input_frames = {i: [f0[i]] for i in range(num_cams)} # collected for grid / test-input output_frames = {i: [out0[i].unsqueeze(1)] for i in range(num_cams)} # (3,1,H,W) # ---- single-frame IMAGE output ---- if T_target == 1: for i in range(num_cams): _save_img(os.path.join(out_dir, f'pred_{i:02d}_{cam_names[i]}.png'), out0[i]) _save_img(os.path.join(in_dir, f'{cam_names[i]}.png'), f0[i]) _grid_image(os.path.join(out_dir, 'grid.png'), f0[0], out0, cam_names, front_w=front_w) print(f'[done] output images + grid.png + test_input/ in {out_dir}') return # ---- STREAM: chunk by chunk (causal outer loop) ---- for j in range(num_clips): s = CHUNK * j print(f'[stream] chunk {j+1}/{num_clips}: frames {s}..{s + CHUNK} ' f'(keyframe @ {s + CHUNK} first, then interpolate) ...') steps = [source.read_step(s + k) for k in range(NUM_FRAMES)] # 9 frames, overlap 1 with prev chunk clip_rgb = {i: torch.stack([steps[k][i] for k in range(NUM_FRAMES)], 0) for i in range(num_cams)} clip_cny = {i: torch.stack([canny(steps[k][i]) for k in range(NUM_FRAMES)], 0) for i in range(num_cams)} out_chunk = renderer.process_chunk(clip_rgb, clip_cny) # {cam -> (3,8,H,W)} for i in range(num_cams): output_frames[i].append(out_chunk[i]) for k in range(1, NUM_FRAMES): # the 8 new input frames input_frames[i].append(steps[k][i]) # ---- assemble, crop to the requested length, save ---- preds = {i: torch.cat(output_frames[i], dim=1)[:, :T_target] for i in range(num_cams)} inputs = {i: torch.stack(input_frames[i], 1)[:, :T_target] for i in range(num_cams)} # (3,T,H,W) for i in range(num_cams): _save_mp4(os.path.join(out_dir, f'pred_{i:02d}_{cam_names[i]}.mp4'), preds[i], fps) _save_mp4(os.path.join(in_dir, f'{cam_names[i]}.mp4'), inputs[i], fps) print(f'[interp] {cam_names[i]} done -> {preds[i].shape[1]} frames') _grid_video(os.path.join(out_dir, 'grid.mp4'), inputs[0], preds, cam_names, fps, front_w=front_w) print(f'[done] outputs + grid.mp4 + test_input/ in {out_dir}') # --------------------------------------------------------------------------- # # Main (single sequence) # --------------------------------------------------------------------------- # def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--kf-ckpt', required=True, help='keyframe checkpoint (matching --kf-model).') p.add_argument('--inp-ckpt', required=True, help='Wan Fun-InP+canny interpolation checkpoint.') p.add_argument('--wan-dir', required=True) p.add_argument('--vae-path', default=None) p.add_argument('--kf-model', choices=['anchor', 'syncfz'], default='anchor', help="keyframe backend: 'anchor'=2-view cross-view Fun-Control (front_wide anchors " "the other cameras); 'syncfz'=pure Wan-T2V per-camera (no cross-view anchor).") p.add_argument('--videos', nargs='+', required=True, help='per-camera video paths; the FIRST is front_wide. A single video -> single-view.') p.add_argument('--num-frames', type=int, default=17, help='output frames: 1 -> save an IMAGE; -1 -> whole video; else N frames ' '(padded to a chunk boundary by duplicating the last source frame, then cropped).') p.add_argument('--prompt', default=SNOW_PROMPT) p.add_argument('--neg-prompt', default=NEG_PROMPT) p.add_argument('--canny-lo', type=float, default=90.0) p.add_argument('--canny-hi', type=float, default=215.0) p.add_argument('--sn-cutoff', type=float, default=30.0, help='anchor model: val structured-noise cutoff.') p.add_argument('--sn-channels', type=int, default=4) p.add_argument('--kf-steps', type=int, default=35) p.add_argument('--kf-guidance', type=float, default=5.0) p.add_argument('--inp-steps', type=int, default=35) p.add_argument('--inp-guidance', type=float, default=5.0) p.add_argument('--seed', type=int, default=0, help='SHARED base-Gaussian seed for ALL keyframes/views.') p.add_argument('--sdedit-strength', type=float, default=1.0, help='SDEdit START SIGMA s0 for keyframes after the first (each is anchored to ' 'keyframe 0): init = (1-s0)*kf0_latent + s0*noise. 1.0=fresh, lower=MORE consistency.') p.add_argument('--no-tone-match', action='store_true', help='disable the per-keyframe tone-match to keyframe 0 (on by default; removes ' 'inter-keyframe exposure flicker, but suppresses genuine lighting changes).') p.add_argument('--front-w', type=int, default=1280, help='width of the large front cell in the grid.') p.add_argument('--fps', type=int, default=10) p.add_argument('--out', default='output/mv_video_transfer') args = p.parse_args() out_dir = args.out if os.path.isabs(args.out) else os.path.join(ANY4D_ROOT, args.out) pipes = build_pipes(args.kf_ckpt, args.inp_ckpt, args.wan_dir, args.vae_path, kf_model=args.kf_model, sn_cutoff=args.sn_cutoff, sn_channels=args.sn_channels, seed=args.seed, prompt=args.prompt, neg_prompt=args.neg_prompt) run_sequence(pipes, args.videos, args.num_frames, out_dir, canny_lo=args.canny_lo, canny_hi=args.canny_hi, kf_steps=args.kf_steps, kf_guidance=args.kf_guidance, inp_steps=args.inp_steps, inp_guidance=args.inp_guidance, seed=args.seed, sdedit_strength=args.sdedit_strength, tone_match=not args.no_tone_match, front_w=args.front_w, fps=args.fps) if __name__ == '__main__': main()