# Wan 2.1 VAE wrapper compatible with the Any4D VAE interface. # Mirrors the API of custom/vae/a4d_vae.py. The actual encode/decode path is # `pipe.tokenizer` (a `WanVAEAdapter` from custom/wan/wan_vae_adapter.py), # which wraps `cosmos_predict2.tokenizers.tokenizer.VAE` — i.e. the upstream # Cosmos-Predict2 inlining of the Wan 2.1 VAE. This file adds the surrounding # Any4D-side concerns (action codec, text-encoder fallback, ref-latent # construction for Wan Fun InP). import collections import os from typing import Optional import lovely_tensors import torch _DEFAULT_TEXT_WAN_PKL = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'default_text_wan.pkl') from custom.vae.action_codec import ( actions_to_absolute, actions_to_relative, normalize_actions, prepare_action_normalizer, unnormalize_actions, ) from imaginaire.utils import log lovely_tensors.monkey_patch() # Wan2.1 VAE spec: 16 latent channels, temporal_compression_factor=4, spatial_compression_factor=8. WAN_LATENT_CH = 16 WAN_TEMPORAL_COMPRESSION = 4 WAN_SPATIAL_COMPRESSION = 8 # Wan uses its own data scaling, so we keep sigma_data = 1.0 to match the diffusion pipeline. SIGMA_DATA = 1.0 T5_EMB_MAX_LENGTH = 512 # Bound the lat_language (UMT5 text-embedding) cache: each entry is ~4 MB, so 1024 entries ≈ 4 GB # worst case. LRU-on-write eviction (see _cache_lat_language) keeps it bounded over long runs with # many distinct prompts (DROID-scale). Mirrors custom/vae/a4d_vae.py. LAT_LANGUAGE_CACHE_MAX_ENTRIES = 1024 class VAE: """Any4D-compatible VAE wrapper around Wan2.1 VAE + (optional) UMT5 text encoder. This is the drop-in replacement for custom.vae.a4d_vae.VAE when training with a Wan2.1 Fun 1.3B InP base model. The interface (encode_a4d_batch, decode_a4d_batch, encode_text, encode_single_video_*, decode_single_video_*) matches a4d_vae.py so the rest of the Any4D pipeline can remain untouched. """ def __init__(self, device, config, pipe): """ :param device (str): e.g. "cuda". :param config: Any4DConfig (or dict) - must contain model_manager_config.vae_path, optionally model_manager_config.text_encoder_path, and the usual action settings. :param pipe: An Any4DWanPipeline. We borrow pipe.tokenizer (the Wan VAE adapter) and pipe.text_encoder if available. """ self.device = device self.config = config self.pipe = pipe self.tensor_kwargs = dict(device=device, dtype=torch.bfloat16) # Borrow the Wan tokenizer/VAE that Any4DWanPipeline already instantiated. assert pipe.tokenizer is not None, \ "Any4DWanPipeline must instantiate a Wan VAE adapter before VAE() is built" self.video_tokenizer = pipe.tokenizer self.vae_ratio = (self.video_tokenizer.temporal_compression_factor, self.video_tokenizer.spatial_compression_factor, self.video_tokenizer.spatial_compression_factor) # Text encoder (UMT5-XXL). When disabled, fall back to cached default embeddings, # exactly like custom/vae/a4d_vae.py. self.text_encoder = pipe.text_encoder # OrderedDict + LRU-on-write eviction (see _cache_lat_language). A plain dict grows # unboundedly over long runs (each entry ~4 MB -> eventual OOM); mirrors a4d_vae.py. self.cached_lat_language: "collections.OrderedDict[str, dict]" = collections.OrderedDict() self.relative_actions = getattr(config, 'relative_actions', False) self.action_normalizer = None if hasattr(config, 'normalizer_stats') and len(config.normalizer_stats) > 0: self.action_normalizer = prepare_action_normalizer( config.normalizer_stats, config.relative_actions) if self.text_encoder is None: from custom.utils.io import read_pickle # UMT5-XXL embedding pickle — hidden_dim=4096 (vs Cosmos T5-11B's 1024, so we # cannot share default_text.pkl). Generate via custom/wan/build_default_text_wan.py. self.default_text = read_pickle(_DEFAULT_TEXT_WAN_PKL) # Only keep the tensor entries; drop metadata like 'prompt' string if present. # Retain both positive and (optional) negative UMT5 embeddings for CFG later. self.default_text = { k: v.to(**self.tensor_kwargs) for k, v in self.default_text.items() if k in ('t5_text_embeddings', 't5_text_mask', 't5_text_embeddings_neg', 't5_text_mask_neg') } log.warning(f'[a4d_vae_wan] No text encoder; using default_text_wan.pkl ' f'(keys: {sorted(self.default_text.keys())})') else: self.default_text = None # ------------------------ text encoding ------------------------ def _cache_lat_language(self, prompt_key: str, value: dict): '''Insert into cached_lat_language with LRU-on-write eviction. Mirrors custom/vae/a4d_vae.py._cache_lat_language — keeps the cache bounded at LAT_LANGUAGE_CACHE_MAX_ENTRIES so many-distinct-prompt training doesn't leak VRAM.''' cache = self.cached_lat_language if prompt_key in cache: cache.move_to_end(prompt_key) return # already cached, nothing more to do cache[prompt_key] = value while len(cache) > LAT_LANGUAGE_CACHE_MAX_ENTRIES: cache.popitem(last=False) # drop the oldest @torch.no_grad() def encode_text(self, prompt): if self.text_encoder is None: B = len(prompt) results = {} for k, v in self.default_text.items(): if v.ndim == 3: results[k] = v.repeat(B, 1, 1) elif v.ndim == 2: results[k] = v.repeat(B, 1) else: results[k] = v.repeat(B) return results # Wan uses UMT5-XXL text encoder. DiffSynth exposes encode_prompt() returning (B, L, C) fp32. emb, mask = self.text_encoder.encode_prompts( prompt, max_length=T5_EMB_MAX_LENGTH, return_mask=True) return { 't5_text_embeddings': emb.to(**self.tensor_kwargs), 't5_text_mask': mask.to(**self.tensor_kwargs), } @torch.no_grad() def encode_text_as_needed(self, sample, phase: str = 'train', dropout_rate: float = 0.2): # Fast-path: the dataset injected precomputed UMT5 embeddings directly # (e.g. GaiaFSCannyDataset with a t5_cache). They are already collated into # (B, L, 4096) tensors, so we only need to make sure they're on-device with # the right dtype. Skips the loaded text encoder / default_text_wan.pkl, but # still falls through to the CFG dropout block below so training-time uncond # behaviour matches the computed-path. if 't5_text_embeddings' in sample and isinstance(sample['t5_text_embeddings'], torch.Tensor): sample['t5_text_embeddings'] = sample['t5_text_embeddings'].to(**self.tensor_kwargs) if 't5_text_mask' in sample and isinstance(sample['t5_text_mask'], torch.Tensor): sample['t5_text_mask'] = sample['t5_text_mask'].to(**self.tensor_kwargs) if 'neg_t5_text_embeddings' in sample and isinstance(sample['neg_t5_text_embeddings'], torch.Tensor): sample['neg_t5_text_embeddings'] = sample['neg_t5_text_embeddings'].to(**self.tensor_kwargs) if 'neg_t5_text_mask' in sample and isinstance(sample['neg_t5_text_mask'], torch.Tensor): sample['neg_t5_text_mask'] = sample['neg_t5_text_mask'].to(**self.tensor_kwargs) else: # Same logic as custom/vae/a4d_vae.py; text pipeline is agnostic to the diffusion backbone. if 'lat_language' in sample['anydata']: lat_language = sample['anydata']['lat_language'][(0, 0)] if len(lat_language) == 1 or lat_language['t5_text_embeddings'] == 'pruned': sample_text = self.encode_text(sample['prompt']) else: sample_text = {k: v.to(**self.tensor_kwargs) for k, v in lat_language.items() if k not in ['prompt']} for b in range(len(lat_language['prompt'])): valid = all(val[b].sum() != 0 for key, val in lat_language.items() if key not in ['prompt']) if not valid: prompt = lat_language['prompt'][b] if prompt in self.cached_lat_language: sample_text_b = self.cached_lat_language[prompt] self.cached_lat_language.move_to_end(prompt) # LRU touch else: sample_text_b = self.encode_text(prompt) self._cache_lat_language(prompt, sample_text_b) for k, v in sample_text_b.items(): sample_text[k][b] = v else: is_cached = True lat_language = [] prompt = sample['prompt'] B = len(prompt) for b in range(B): if prompt[b] in self.cached_lat_language: lat_language.append(self.cached_lat_language[prompt[b]]) self.cached_lat_language.move_to_end(prompt[b]) # LRU touch else: is_cached = False break if is_cached: sample_text = {k: torch.stack([lat_language[i][k] for i in range(B)], 0) for k in lat_language[0].keys()} else: sample_text = self.encode_text(sample['prompt']) for b in range(B): self._cache_lat_language(prompt[b], {k: v[b] for k, v in sample_text.items()}) sample = {**sample, **sample_text} if self.text_encoder is None: # default_text_wan.pkl already contains precomputed neg embeddings under # t5_text_embeddings_neg / t5_text_mask_neg — map them to the keys the # conditioner expects for CFG. if 't5_text_embeddings_neg' in sample_text: sample['neg_t5_text_embeddings'] = sample_text['t5_text_embeddings_neg'] sample['neg_t5_text_mask'] = sample_text['t5_text_mask_neg'] elif 'negative_prompt' in sample: neg = self.encode_text(sample['negative_prompt']) sample = {**sample, 'neg_t5_text_embeddings': neg['t5_text_embeddings'], 'neg_t5_text_mask': neg['t5_text_mask']} # Wan-style text dropout: at train time, replace `dropout_rate` of the per-sample # positive embeddings with the cached UMT5("") (negative) embedding, so the # training uncond branch matches inference's `is_negative_prompt=True` path. # Conditioner-side dropout is set to 0 on the Wan pipe (see # `ANY4D_WAN_PIPELINE_1_3B` in cosmos_predict2/configs/base/defaults/model.py), # so this is the only dropout that fires. if (phase == 'train' and dropout_rate > 0 and 'neg_t5_text_embeddings' in sample and 't5_text_embeddings' in sample): pos_emb = sample['t5_text_embeddings'] pos_mask = sample['t5_text_mask'] neg_emb = sample['neg_t5_text_embeddings'] neg_mask = sample['neg_t5_text_mask'] B = pos_emb.shape[0] drop = torch.bernoulli( torch.full((B,), float(dropout_rate), device=pos_emb.device)).bool() # neg_emb / neg_mask have batch dim 1; broadcast to B. neg_emb_B = neg_emb.expand(B, *neg_emb.shape[1:]) if neg_emb.shape[0] == 1 else neg_emb neg_mask_B = neg_mask.expand(B, *neg_mask.shape[1:]) if neg_mask.shape[0] == 1 else neg_mask pos_emb = torch.where(drop.view(B, *([1] * (pos_emb.ndim - 1))), neg_emb_B, pos_emb) pos_mask = torch.where(drop.view(B, *([1] * (pos_mask.ndim - 1))), neg_mask_B, pos_mask) sample['t5_text_embeddings'] = pos_emb sample['t5_text_mask'] = pos_mask return sample # ------------------------ single-video encoding ------------------------ @torch.no_grad() def encode_single_video(self, raw_video): """:param raw_video: (B, 3, Tp, Hp, Wp) in [-1, 1].""" (B, Cp, Tp, Hp, Wp) = raw_video.shape assert Cp == 3, f'Raw video must have 3 channels, got {Cp}' # Wan VAE temporal compression: Tp must be 1 + 4k. Pad short videos by repeating last frame. tcf = self.video_tokenizer.temporal_compression_factor remainder = (Tp - 1) % tcf if remainder != 0: pad_t = tcf - remainder raw_video = torch.cat( [raw_video, raw_video[:, :, -1:].expand(-1, -1, pad_t, -1, -1)], dim=2) raw_video = raw_video.to(**self.tensor_kwargs, non_blocking=True) latent = self.video_tokenizer.encode(raw_video) * SIGMA_DATA return latent @torch.no_grad() def encode_single_video_rgb(self, raw_video): assert raw_video.shape[1] == 3, 'RGB video must have 3 channels' if raw_video.isnan().any(): log.error(f'NaN in RAW RGB before Wan VAE encode! shape={raw_video.shape}') raw_video = raw_video.nan_to_num(0.0) raw_video = raw_video * 2.0 - 1.0 return self.encode_single_video(raw_video) @torch.no_grad() def encode_single_video_depth(self, raw_video, key=None): assert raw_video.shape[1] == 1 raw_video = raw_video * 2.0 - 1.0 raw_video = torch.cat([raw_video] * 3, dim=1) return self.encode_single_video(raw_video) @torch.no_grad() def encode_single_video_points(self, raw_video): assert raw_video.shape[1] == 3 return self.encode_single_video(raw_video) @torch.no_grad() def encode_single_video_cams(self, raw_video, with_orig=False): assert raw_video.shape[1] in [6, 9] rays = self.encode_single_video(raw_video[:, 0:3]) cross = self.encode_single_video(raw_video[:, 3:6]) if with_orig: orig = self.encode_single_video(raw_video[:, 6:9]) return torch.cat([rays, cross, orig], dim=1) return torch.cat([rays, cross], dim=1) @torch.no_grad() def encode_single_video_auto(self, key, raw_video): raw_video = raw_video.clone().detach() if key.startswith('rgb'): return self.encode_single_video_rgb(raw_video) if key.startswith('depth'): if raw_video.shape[1] == 1: return self.encode_single_video_depth(raw_video, key=key) if raw_video.shape[1] == 3: log.warning('Obsolete colormap depth encoding path') return self.encode_single_video_rgb(raw_video) raise ValueError(f'Invalid depth channels: {raw_video.shape[1]}') if key.startswith('points'): return self.encode_single_video_points(raw_video) if key.startswith('cams'): return self.encode_single_video_cams(raw_video) raise ValueError(f'Unknown entry: {key}') @torch.no_grad() def encode_single_lowdim_auto(self, key, raw_data): latent_data = raw_data.clone().detach() if key.startswith('action') or key.startswith('proprio'): if self.relative_actions: latent_data, _ = actions_to_relative(latent_data) if self.action_normalizer is not None: latent_data = normalize_actions(latent_data, self.action_normalizer) latent_data *= self.config.action_multiplier elif key.startswith('bbox') or key.startswith('traj'): pass else: raise ValueError(f'Unknown entry: {key}') return latent_data # ------------------------ batch encoding ------------------------ @torch.no_grad() def _build_wan_ref_latent(self, rgb_pixel: torch.Tensor, input_mask_latent: torch.Tensor) -> torch.Tensor: """Build the Wan 2.1 Fun InP reference latent for a view. Wan was trained with the reference conditioning built by GRAY-padding in pixel space: at given frames the reference carries the real image (in Wan VAE's [-1, 1] pixel range) and at non-given frames it is set to 0 (gray in [-1, 1] space, NOT black=-1). This is a per-frame recipe, so it generalizes directly from first-frame-only (k=1) to variable-length (k>=1) conditioning. We receive the per-view input mask at LATENT temporal resolution (Tl) and expand it back to raw pixel resolution (Tp) using Wan VAE's temporal mapping: latent slot 0 ↔ raw frame 0 latent slot k>0 ↔ raw frames [(k-1)*tcf + 1 .. k*tcf] This assumes the contiguous `[1]*k + [0]*(Tl-k)` pattern emitted by the unified_anydrive dataloader. Spatial dimensions of the mask are ignored (conditioning is frame-level in Wan's I2V paradigm). :param rgb_pixel: (B, 3, Tp, Hp, Wp) raw pixel video in [0, 1]. :param input_mask_latent: (B, 1, Tl, Hl, Wl) latent-resolution input mask. """ B, C, Tp, Hp, Wp = rgb_pixel.shape assert C == 3, f'rgb_pixel must have 3 channels, got {C}' tcf = self.video_tokenizer.temporal_compression_factor # Pad Tp to nearest 1+4k if needed (same logic as encode_single_video). remainder = (Tp - 1) % tcf if remainder != 0: pad_t = tcf - remainder rgb_pixel = torch.cat( [rgb_pixel, rgb_pixel[:, :, -1:].expand(-1, -1, pad_t, -1, -1)], dim=2) Tp = rgb_pixel.shape[2] rgb_m11 = rgb_pixel * 2.0 - 1.0 # Collapse spatial dims of the mask; frame-level gating only. mask_l = input_mask_latent[:, :, :, 0:1, 0:1] # (B, 1, Tl, 1, 1) first = mask_l[:, :, 0:1] # (B, 1, 1, 1, 1) rest = mask_l[:, :, 1:].repeat_interleave(tcf, dim=2) # (B, 1, (Tl-1)*tcf, 1, 1) mask_raw = torch.cat([first, rest], dim=2) # (B, 1, Tp_expanded, 1, 1) # Trim or pad mask to match Tp after pixel padding. if mask_raw.shape[2] < Tp: mask_raw = torch.cat([mask_raw, torch.zeros(B, 1, Tp - mask_raw.shape[2], 1, 1, device=mask_raw.device, dtype=mask_raw.dtype)], dim=2) elif mask_raw.shape[2] > Tp: mask_raw = mask_raw[:, :, :Tp] # Optional (original Wan-Fun-InP "end image" recipe): when the last latent slot # is an *isolated* trailing-given slot (cond_last_frame), keep only the FINAL raw # frame real and gray the other (tcf-1) raw frames of that slot, so the end token # encodes [gray, ..., gray, end] (a single end image) rather than tcf real frames. # The latent input_mask channel is left unchanged (the last slot stays "given"). if getattr(self.config, 'wan_inp_single_end_frame', False) and mask_l.shape[2] >= 2: last_given = mask_l[:, 0, -1, 0, 0] > 0.5 # (B,) prev_notgiven = mask_l[:, 0, -2, 0, 0] < 0.5 # (B,) isolated_end = last_given & prev_notgiven lo, hi = max(Tp - tcf, 0), Tp - 1 # last latent slot -> raw [Tp-tcf : Tp] for b in range(mask_raw.shape[0]): if bool(isolated_end[b]): mask_raw[b, :, lo:hi] = 0.0 # gray all but the final raw frame ref_m11 = rgb_m11 * mask_raw.to(rgb_m11) # given frames: real; else 0 (gray) return self.video_tokenizer.encode( ref_m11.to(**self.tensor_kwargs, non_blocking=True)) * SIGMA_DATA @torch.no_grad() def _build_wan_anchor_ref_latent(self, rgb_pixel: torch.Tensor, input_mask_latent: torch.Tensor, is_anchor_b: Optional[torch.Tensor]) -> torch.Tensor: """Build the per-view reference latent for the I2V-across-VIEWS (anchor) recipe. Unlike `_build_wan_ref_latent` (the InP gray-padded recipe), here the ref slot is the reference-image-conditioning channel of Wan Fun-Control: * For samples where THIS view is the ANCHOR (is_anchor_b[b] True): the ref = the view's OWN full clean VAE latent, broadcast across all latent frames (every latent frame's ref = the clean latent; for Tl==1 it's just that one latent). The full clean RGB is encoded (NO gray padding). * For all other samples (and when is_anchor_b is None): ref = ZEROS, which reproduces the pretrained pure text+control behavior (y channel = 0). Returns a latent of shape (B, 16, Tl, Hl, Wl) matching the VAE encoding of `rgb_pixel`, with non-anchor batch elements zeroed. :param rgb_pixel: (B, 3, Tp, Hp, Wp) raw pixel video in [0, 1]. :param input_mask_latent: (B, 1, Tl, Hl, Wl) latent input mask (used only to determine the target latent temporal length Tl). :param is_anchor_b: (B,) bool tensor (True where this view is the anchor) or None. """ B = rgb_pixel.shape[0] # Fast path: no anchor for any sample of this view -> return zeros with the right # latent shape (cheap: encode a single black frame, take its shape). if is_anchor_b is None or not bool(is_anchor_b.any()): # Encode the (already padded) clean RGB just to obtain the exact latent shape, # then zero it. This avoids hard-coding the spatial latent dims here. rgb_m11 = rgb_pixel * 2.0 - 1.0 z = self.video_tokenizer.encode( rgb_m11.to(**self.tensor_kwargs, non_blocking=True)) * SIGMA_DATA return torch.zeros_like(z) # Encode the FULL clean RGB latent (reference image), no gray padding. rgb_m11 = rgb_pixel * 2.0 - 1.0 ref_latent = self.video_tokenizer.encode( rgb_m11.to(**self.tensor_kwargs, non_blocking=True)) * SIGMA_DATA # Broadcast the anchor's clean latent across all latent frames: the anchor is a # single reference image repeated over time (frame-level image conditioning). For # the typical T=1 single-frame case this is a no-op (Tl already 1). Tl_ref = ref_latent.shape[2] if Tl_ref > 1: ref_latent = ref_latent[:, :, 0:1].expand(-1, -1, Tl_ref, -1, -1).contiguous() # Zero out batch elements where this view is NOT the anchor. keep = is_anchor_b.reshape(B, 1, 1, 1, 1).to(ref_latent.device, ref_latent.dtype) return ref_latent * keep @torch.no_grad() def encode_a4d_batch(self, data_batch, phase: str = 'train'): batch_enc = dict(data_batch) a4d_raw = dict(data_batch['a4d_raw']) a4d_latent = dict(data_batch['a4d_latent']) batch_enc = self.encode_text_as_needed(batch_enc, phase=phase) # 1) Normal per-entry encoding (rgb, depth, cams, traj, ...). for k, e in a4d_raw.items(): assert k not in ['fps', 'prompt', 'num_frames'], \ f'Key {k} should go in the parent sample dictionary' if k in ['video_dims', 'used_views']: continue elif k in self.config.all_highdim_entries: if k in a4d_latent: continue if k.endswith('_ref'): # Deferred to step 2 below — ref latents are derived, not raw. continue a4d_latent[k] = self.encode_single_video_auto(k, e) elif k in self.config.all_lowdim_entries: if k in a4d_latent: continue a4d_latent[k] = self.encode_single_lowdim_auto(k, e) elif self.config.ignore_unused_data: log.debug(f'Ignoring unused entry: {k}: {e.shape}') else: raise ValueError(f'Unknown entry: {k}') # 1b) Resize masks to match each entry's latent Tl after VAE padding. # When raw Tp is not 1+4k, encode_single_video pads → larger Tl, but the # dataloader-generated masks still have the original (smaller) Tl. We must # resize every (latent, mask-sibling) pair — not just rgb — because pack_streams # asserts shape equality per-entry (e.g. rgb{v}_edge_control needs its # rgb{v}_edge_control_input_mask resized too). latent_video_keys = [ k for k, t in a4d_latent.items() if isinstance(t, torch.Tensor) and t.ndim == 5 and not k.endswith(('_input_mask', '_output_mask', '_supervise_mask')) ] for base_key in latent_video_keys: target_Tl = a4d_latent[base_key].shape[2] for suffix in ('_input_mask', '_output_mask', '_supervise_mask'): mk = f'{base_key}{suffix}' if mk not in a4d_latent or a4d_latent[mk].ndim != 5: continue cur_Tl = a4d_latent[mk].shape[2] if cur_Tl < target_Tl: pad = a4d_latent[mk][:, :, -1:].expand(-1, -1, target_Tl - cur_Tl, -1, -1) a4d_latent[mk] = torch.cat([a4d_latent[mk], pad], dim=2) elif cur_Tl > target_Tl: a4d_latent[mk] = a4d_latent[mk][:, :, :target_Tl] # 2) Derive Wan reference latents (rgb{N}_ref) per view. These live in # `video_entries[v][f'rgb{v}_ref']` if the experiment config requests them. # # TWO MUTUALLY-EXCLUSIVE recipes, selected by which Wan backbone is in use: # * Wan Fun-CONTROL path (wan_variant='control'): the ANCHOR recipe # (_build_wan_anchor_ref_latent). The Control DiT's _to_wan_48ch routes the ref # slot to the pretrained y/image-conditioning channel [32:48], so the ref must # be either a clean reference image (anchor) or ZEROS (pure text+control). Per # sample: # - ANCHOR view (view-synth) -> that view's OWN full clean VAE latent # (reference-image broadcast across latent frames) -> the y channel. # - all other views / non-anchor samples / legacy Control runs -> ZEROS, so # legacy Control stays on-distribution with y=0 (verified in smoke test). # The per-sample anchor index is plumbed from the dataloader via # meta['anchor_view'] (collated to a (B,) tensor; -1 / absent => no anchor). # * Wan Fun-INP(+canny) path (wan_variant='inp'/'inp_canny'): the LEGACY InP recipe # (_build_wan_ref_latent) -- gray-padded reference (real pixels at given frames, # gray elsewhere). Its DiT (_to_wan_36ch) was trained for exactly this. UNCHANGED. # The InP recipe is the default, so the Fun-InP path is byte-for-byte unchanged. use_anchor_ref = getattr(self.config, 'wan_variant', 'inp') == 'control' meta = data_batch.get('meta', {}) or {} anchor_view = meta.get('anchor_view', None) if anchor_view is not None and not torch.is_tensor(anchor_view): anchor_view = torch.as_tensor(anchor_view) for v in range(self.config.num_views): for k in list(self.config.video_entries[v].keys()): if not k.endswith('_ref'): continue if k in a4d_latent: continue base_key = k[:-len('_ref')] # rgb0_ref → rgb0 if base_key not in a4d_raw: log.warning(f'[a4d_vae_wan] cannot build {k}: no {base_key} in a4d_raw') continue input_mask_key = f'{base_key}_input_mask' if input_mask_key not in a4d_latent: log.warning(f'[a4d_vae_wan] cannot build {k}: no {input_mask_key} in a4d_latent') continue if use_anchor_ref: # Per-sample anchor mask for THIS view (B,): True where v is the anchor. # anchor_view may be a batch-level SCALAR (the view-synth task is chosen # once per anydata_to_any4d call) -> broadcast to the view's batch size B. if anchor_view is None: is_anchor_b = None # no anchor info at all -> zeros else: _n = a4d_raw[base_key].shape[0] _av = anchor_view.reshape(-1) if _av.numel() == 1: _av = _av.expand(_n) is_anchor_b = (_av == v) ref_latent = self._build_wan_anchor_ref_latent( a4d_raw[base_key], a4d_latent[input_mask_key], is_anchor_b) else: ref_latent = self._build_wan_ref_latent( a4d_raw[base_key], a4d_latent[input_mask_key]) # Match ref Tl to the already-encoded rgb latent's Tl (padding may differ). if base_key in a4d_latent: target_Tl = a4d_latent[base_key].shape[2] ref_Tl = ref_latent.shape[2] if ref_Tl < target_Tl: ref_latent = torch.cat([ref_latent, ref_latent[:, :, -1:].expand(-1, -1, target_Tl - ref_Tl, -1, -1)], dim=2) elif ref_Tl > target_Tl: ref_latent = ref_latent[:, :, :target_Tl] a4d_latent[k] = ref_latent # ref slot is always read-through: input_mask=1, output/supervise=0 # (cf. infer_wan.py:235-237). Mirrored mask shape from the base entry. imask_base = a4d_latent[input_mask_key] a4d_latent[f'{k}_input_mask'] = torch.ones_like(imask_base) a4d_latent[f'{k}_output_mask'] = torch.zeros_like(imask_base) a4d_latent[f'{k}_supervise_mask'] = torch.zeros_like(imask_base) # Update video_dims to reflect actual encoded Tl (padding may have changed it). first_rgb = next((v for k, v in a4d_latent.items() if k.startswith('rgb') and not k.endswith(('_mask', '_ref')) and isinstance(v, torch.Tensor) and v.ndim == 5), None) if first_rgb is not None: _, _, Tl_actual, Hl_actual, Wl_actual = first_rgb.shape a4d_latent['video_dims'] = [(Tl_actual, Hl_actual, Wl_actual)] * self.config.num_views batch_enc['a4d_latent'] = a4d_latent # Impute conditioner keys the Cosmos pipeline expects but anydata doesn't provide. vd = a4d_latent.get('video_dims', None) if vd is not None and len(vd) > 0: Tl, Hl, Wl = vd[0] else: first_rgb2 = next((v for k, v in a4d_latent.items() if k.startswith('rgb') and not k.endswith(('_mask', '_ref'))), None) if first_rgb2 is not None: _, _, Tl, Hl, Wl = first_rgb2.shape else: Tl, Hl, Wl = 11, 40, 64 tk = self.tensor_kwargs if 'padding_mask' not in batch_enc: batch_enc['padding_mask'] = torch.zeros((1, 1, Hl, Wl), **tk) if 'image_size' not in batch_enc: batch_enc['image_size'] = torch.tensor([[Hl, Wl, Hl, Wl]], **tk) if 'num_frames' not in batch_enc: batch_enc['num_frames'] = torch.tensor([Tl], **tk) return batch_enc # ------------------------ decoding ------------------------ @torch.no_grad() def decode_single_video(self, latent_video): latent_video = latent_video / SIGMA_DATA return self.video_tokenizer.decode(latent_video) @torch.no_grad() def decode_single_video_rgb(self, latent_video): out = self.decode_single_video(latent_video) return (out + 1.0) / 2.0 @torch.no_grad() def decode_single_video_depth(self, latent_video, key=None): out = self.decode_single_video(latent_video) out = out.mean(dim=1, keepdim=True) return (out + 1.0) / 2.0 @torch.no_grad() def decode_single_video_points(self, latent_video, key=None): return self.decode_single_video(latent_video) @torch.no_grad() def decode_single_video_cams(self, latent_video, with_orig=False): rays = self.decode_single_video(latent_video[:, 0:WAN_LATENT_CH]) cross = self.decode_single_video(latent_video[:, WAN_LATENT_CH:2 * WAN_LATENT_CH]) if with_orig: orig = self.decode_single_video(latent_video[:, 2 * WAN_LATENT_CH:3 * WAN_LATENT_CH]) return torch.cat([rays, cross, orig], dim=1) return torch.cat([rays, cross], dim=1) @torch.no_grad() def decode_single_video_auto(self, key, latent_video, with_orig=False): latent_video = latent_video.clone().detach() if key.startswith('rgb'): return self.decode_single_video_rgb(latent_video).clamp(0.0, 1.0) if key.startswith('depth'): return self.decode_single_video_depth(latent_video, key).clamp(0.001, 1.0) if key.startswith('points'): return self.decode_single_video_points(latent_video).clamp(-1.0, 1.0) if key.startswith('cams'): return self.decode_single_video_cams(latent_video, with_orig=with_orig).clamp(-1.0, 1.0) raise ValueError(f'Unknown entry: {key}') @torch.no_grad() def decode_single_lowdim_auto(self, key, latent_data): raw_data = latent_data.clone().detach() if key.startswith('action') or key.startswith('proprio'): raw_data /= self.config.action_multiplier if self.relative_actions: raw_data = actions_to_absolute(raw_data, None) if self.action_normalizer is not None: raw_data = unnormalize_actions(raw_data, self.action_normalizer) elif key.startswith('bbox') or key.startswith('traj'): pass else: raise ValueError(f'Unknown entry: {key}') return raw_data @torch.no_grad() def decode_a4d_batch(self, entries_latent, phase: str = 'train'): entries_raw = dict() for k, e in entries_latent.items(): if k in ['video_dims', 'used_views']: continue if k.endswith('_mask'): continue if k in self.config.all_highdim_entries: entries_raw[k] = self.decode_single_video_auto(k, e) elif k in self.config.all_lowdim_entries: entries_raw[k] = self.decode_single_lowdim_auto(k, e) elif self.config.ignore_unused_data: log.debug(f'Ignoring unused entry: {k}: {e.shape}') else: raise ValueError(f'Unknown entry: {k}') return entries_raw