# Any4D pipeline for Wan 2.1 Fun-V1.1-1.3B-Control. # # Differences from Any4DWanPipeline (InP): # 1. Uses Any4DWanControlDiT instead of Any4DWanDiT. # Stream layout per view is 50 ch (rgb + 1ch in_mask + 1ch out_mask + 16ch ref + # 16ch edge_control), reshaped to Wan's 48 ch DiT input (noisy + control + ref) # by Any4DWanControlDiT._to_wan_48ch. # 2. _rewrite_masks_for_wan also zeroes the supervise mask for the rgb{N}_edge_control # slot (canny conditioning is read-through, not a denoising target). Future # modalities at this slot (rgb{N}_depth_control, rgb{N}_seg_control, ...) follow # the same convention. # 3. build_wan_control_pipeline optionally instantiates a live UMT5-XXL encoder. # Default off (the encoder is ~22 GB and OOMs alongside DiT + optimizer state # on a single L40S); the dataset is expected to inject precomputed embeddings # per sample. Force-enable via env var WAN_CONTROL_LIVE_UMT5=1. import os from typing import Any, Dict, Iterable, List, Optional import torch from cosmos_predict2.pipelines.video2world import Video2WorldPipeline from custom.any4d.a4d_pipe import Any4DPipeline from custom.utils.utils import get_checkpoint from custom.wan.a4d_pipe_wan import ( Any4DWanPipeline, _load_sharded_state_dict, _load_state_dict_from_path, _remap_wan_dit_state_dict, _strip_common_prefixes, ) from imaginaire.utils import log _WAN_PREFIX = 'wan.' class Any4DWanControlPipeline(Any4DWanPipeline): """Pipeline for Any4D + Wan 2.1 Fun-V1.1-1.3B-Control. Inherits the entire Any4DWanPipeline flow (denoise / generate / per-stream noise seeding / Wan-native flow-match math). Only overrides the mask-rewrite step so the new edge_control slot is treated as conditioning (clean, not supervised). """ def _rewrite_masks_for_wan(self, x0_streams, masks) -> None: cfg = self.any4d_config for v_idx in range(cfg.num_views): v_key = f'v{v_idx}' if v_key not in x0_streams: continue entries = cfg.video_entries[v_idx] rgb_k = f'rgb{v_idx}' imask_k = f'{rgb_k}_input_mask' ref_k = f'{rgb_k}_ref' control_k = f'{rgb_k}_edge_control' required = [rgb_k, imask_k, ref_k, control_k] for needed in required: if needed not in entries: raise AssertionError( f'Any4DWanControlPipeline view {v_idx} missing required entry ' f'{needed!r} in video_entries.') rgb_s, rgb_e = entries[rgb_k][0], entries[rgb_k][1] ref_s, ref_e = entries[ref_k][0], entries[ref_k][1] control_s, control_e = entries[control_k][0], entries[control_k][1] x0 = x0_streams[v_key] B, C, T, H, W = x0.shape device, dtype = x0.device, x0.dtype rgb_shape = (B, rgb_e - rgb_s, T, H, W) ones = torch.ones(rgb_shape, device=device, dtype=dtype) zeros = torch.zeros(rgb_shape, device=device, dtype=dtype) # rgb slot: noisy at all frames, fully supervised. masks['input'][v_key][:, rgb_s:rgb_e] = zeros masks['output'][v_key][:, rgb_s:rgb_e] = ones masks['supervise'][v_key][:, rgb_s:rgb_e] = ones # ref slot: clean read-through; never supervised (zeros for T2V). masks['supervise'][v_key][:, ref_s:ref_e] = 0.0 # control slot: clean read-through (canny VAE latent), never supervised. masks['supervise'][v_key][:, control_s:control_e] = 0.0 # --------------------------------------------------------------------------- # Live UMT5-XXL text encoder, drop-in compatible with a4d_vae_wan.VAE.encode_text. # encode_prompts() returns ((B, max_length, dim) bf16, (B, max_length) long) — same # shape as default_text_wan.pkl entries — so the existing text-flow logic works. # --------------------------------------------------------------------------- class _WanLiveTextEncoder: """Wraps the vendored WanUMT5Encoder + Wan UMT5 tokenizer for per-batch encoding. Memory: ~22 GB bf16 weights pinned on cuda. Frozen / no_grad — adds no optimizer state. By default we DO NOT instantiate this during training (the encoder pegs half the L40S VRAM and OOMs alongside DiT + optimizer); the dataset injects precomputed embeddings instead. Used by precompute_t5_embeddings.py to build that cache, and optionally re-enabled at training time via WAN_CONTROL_LIVE_UMT5=1. """ def __init__(self, tokenizer_dir: str, weights_path: str, device: str = 'cuda', dtype: torch.dtype = torch.bfloat16, max_length: int = 512): from custom.wan.build_default_text_wan import WanUMT5Encoder from transformers import AutoTokenizer log.info(f'[Wan-Control] Loading UMT5 tokenizer from {tokenizer_dir}') self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir) log.info(f'[Wan-Control] Loading UMT5 weights from {weights_path}') self.model = WanUMT5Encoder().to(dtype=dtype).eval() sd = torch.load(weights_path, map_location='cpu', weights_only=True) miss, surp = self.model.load_state_dict(sd, strict=False) log.info(f'[Wan-Control] UMT5 load: missing={len(miss)} surplus={len(surp)}') self.model = self.model.to(device=device) for p in self.model.parameters(): p.requires_grad_(False) self.device = device self.dtype = dtype self.max_length = max_length @torch.no_grad() def encode_prompts(self, prompts: List[str], max_length: Optional[int] = None, return_mask: bool = True): """Encode a batch of prompt strings into UMT5 features. :returns: (emb (B, L, 4096) bf16, mask (B, L) long). """ import ftfy, html, re if max_length is None: max_length = self.max_length clean: List[str] = [] for p in prompts: t = ftfy.fix_text(p or '') t = html.unescape(html.unescape(t)) t = re.sub(r'\s+', ' ', t).strip() clean.append(t) enc = self.tokenizer( clean, padding='max_length', truncation=True, max_length=max_length, return_tensors='pt', add_special_tokens=True, ) ids = enc['input_ids'].to(self.device) mask = enc['attention_mask'].to(self.device) emb = self.model(ids, mask=mask) # (B, L, 4096) # Zero positions past the real token count — matches WanPrompter.encode_prompt # and the offline build_default_text_wan path. seq_lens = mask.gt(0).sum(dim=1).long() for batch_idx, length in enumerate(seq_lens): emb[batch_idx, length:] = 0 emb = emb.to(dtype=self.dtype) if return_mask: return emb, mask return emb # --------------------------------------------------------------------------- # Pipeline factory (parallel to build_wan_pipeline). # --------------------------------------------------------------------------- def _find_umt5_tokenizer_dir(dit_path: str) -> str: """Locate the UMT5 tokenizer dir inside a Fun-Control checkpoint download.""" candidates = [ os.path.join(dit_path, 'google', 'umt5-xxl'), os.path.join(dit_path, 'google'), ] for c in candidates: if os.path.isdir(c) and os.path.exists(os.path.join(c, 'tokenizer.json')): return c raise FileNotFoundError( f'UMT5 tokenizer not found under {dit_path}. Looked at: {candidates}') def _find_umt5_weights(dit_path: str) -> str: p = os.path.join(dit_path, 'models_t5_umt5-xxl-enc-bf16.pth') if not os.path.exists(p): raise FileNotFoundError(f'UMT5 encoder weights not at {p}') return p def build_wan_control_pipeline( config, dit_path: str = '', text_encoder_path: str = '', vae_path: str = '', device: str = 'cuda', torch_dtype: torch.dtype = torch.bfloat16, any4d_config=None, ): """Drop-in replacement for Video2WorldPipeline.from_config, but for Wan Fun-Control. Mirrors `build_wan_pipeline` (Fun-InP) with three deltas: * Instantiates Any4DWanControlPipeline (and via config.net, Any4DWanControlDiT). * Loads the Fun-Control DiT shard (different in_dim, has ref_conv). * Optionally wires up a live UMT5-XXL encoder (off by default to save VRAM). :param text_encoder_path: ignored — if a live encoder is requested via env var ``WAN_CONTROL_LIVE_UMT5=1``, weights/tokenizer are pulled from ``/models_t5_umt5-xxl-enc-bf16.pth`` + ``/google/umt5-xxl``. """ from imaginaire.lazy_config import instantiate from imaginaire.utils import misc from cosmos_predict2.schedulers.wan_flow_match_scheduler import WanFlowMatchScheduler from cosmos_predict2.models.utils import init_weights_on_device from custom.any4d import a4d_surgery any4d_config = a4d_surgery.validate_config(any4d_config) pipe = Any4DWanControlPipeline(device=device, torch_dtype=torch_dtype) pipe.any4d_config = any4d_config pipe.config = config pipe.precision = { 'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, }[config.precision] pipe.tensor_kwargs = {'device': 'cuda', 'dtype': pipe.precision} pipe.sigma_data = config.sigma_data pipe.setup_data_key() pipe.scheduler = WanFlowMatchScheduler( shift=5.0, sigma_min=0.003 / 1.002, sigma_max=1.0, num_train_timesteps=1000, extra_one_step=True, ) pipe.scaling = None # Wan VAE tokenizer (shared between InP and Control — same VAE weights). if not config.remove_dit: pipe.tokenizer = instantiate(config.tokenizer) assert pipe.tokenizer.latent_ch == pipe.config.state_ch, \ f'Wan VAE latent_ch {pipe.tokenizer.latent_ch} != state_ch {pipe.config.state_ch}' pipe.tokenizer.reset_dtype() pipe.tokenizer.to(device=device) # Text encoder. UMT5-XXL is ~22 GB in bf16 and doesn't coexist with DiT + # optimizer state on a single L40S, so by default we DO NOT load it — the # dataset is expected to inject precomputed embeddings per sample (see # GaiaFSCannyDataset.t5_cache and precompute_t5_embeddings.py). Set the env # var ``WAN_CONTROL_LIVE_UMT5=1`` to force live encoding on hardware with # ≥80 GB of VRAM where the full UMT5 fits alongside the rest. if dit_path and os.environ.get('WAN_CONTROL_LIVE_UMT5', '0') == '1': tokenizer_dir = _find_umt5_tokenizer_dir(dit_path) weights_path = _find_umt5_weights(dit_path) with misc.timer('build live UMT5-XXL'): pipe.text_encoder = _WanLiveTextEncoder( tokenizer_dir=tokenizer_dir, weights_path=weights_path, device=device, dtype=pipe.precision, ) else: log.warning( '[Wan-Control] Skipping live UMT5 load to free ~22 GB. ' 'Provide t5_cache in the dataset YAML, or set WAN_CONTROL_LIVE_UMT5=1 to override.') pipe.text_encoder = None # Conditioner (Cosmos-side) — same as InP. if not config.remove_dit: pipe.conditioner = instantiate(config.conditioner) pipe.prompt_refiner = None pipe.text_guardrail_runner = None pipe.video_guardrail_runner = None # DiT: instantiate Any4DWanControlDiT via LazyCall `config.net`. if not config.remove_dit: with init_weights_on_device(device=torch.device('cpu')): dit_config = config.net dit_config.any4d_config = any4d_config with misc.timer('instantiate Any4DWanControlDiT'): pipe.dit = instantiate(dit_config).eval() if dit_path: state_dict = Any4DWanPipeline.gather_wan_checkpoint(dit_path) with misc.timer('Any4DWanControlDiT load_checkpoint'): pipe.load_checkpoint(state_dict, any4d_config=any4d_config) del state_dict log.success(f'Loaded Wan 2.1 Fun-Control DiT from {dit_path}') if not config.remove_dit: pipe.dit = pipe.dit.to(device=device, dtype=pipe.precision) pipe.dit.eval() # EMA: build a frozen copy of the DiT for the training loop to exponentially average # into (Karras EDM2 power-function EMA). This MUST exist before the inherited # Predict2Video2WorldModel.__init__ -> Video2WorldPipeline.apply_fsdp runs, so apply_fsdp # shards dit_ema and installs the DTensor EMA worker. The inherited # Video2WorldModel.on_before_zero_grad then updates it each step (via ema_beta / # ema_exp_coefficient), ema_scope() swaps it in for validation, and state_dict() persists # it under net_ema.* keys. The copy here only seeds dit_ema from the base DiT; when # training resumes/initializes from a checkpoint, the checkpointer reseeds dit_ema from # the loaded weights. The base Wan pipelines (a4d_pipe_wan.py / # a4d_pipe_wan_inp_canny.py) now build EMA with this same block when config.ema.enabled. if config.ema.enabled and not config.remove_dit: import numpy as np from imaginaire.utils.ema import FastEmaModelUpdater pipe.dit_ema = instantiate(config.net).eval() pipe.dit_ema.requires_grad_(False) pipe.dit_ema = pipe.dit_ema.to(device=device, dtype=pipe.precision) pipe.dit_ema_worker = FastEmaModelUpdater() # swapped to DTensor variant by apply_fsdp s = config.ema.rate pipe.ema_exp_coefficient = np.roots([1, 7, 16 - s**-2, 12 - s**-2]).real.max() pipe.dit_ema_worker.copy_to(src_model=pipe.dit, tgt_model=pipe.dit_ema) log.success(f'Built EMA DiT (rate={s}, ema_exp_coefficient={pipe.ema_exp_coefficient:.4f})') else: pipe.dit_ema = None from megatron.core import parallel_state if parallel_state.is_initialized(): pipe.data_parallel_size = parallel_state.get_data_parallel_world_size() else: pipe.data_parallel_size = 1 return pipe