# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Shared on-the-fly preprocessing / encoding for LipForcing. Single source of truth for turning raw ``(reference video, audio, prompt)`` into the exact tensors the OmniAvatar V2V model consumes. The functions here are imported by BOTH * the training dataloader on-the-fly / cache path (``lipforcing.datasets.omniavatar_dataloader``), and * the inference scripts (via ``scripts/inference/_common.py``) so the encode logic is defined once and shared by training and inference. The produced tensors follow the precomputed ``.pt`` training format. Precomputed ``.pt`` format reproduced here (per sample directory): vae_latents_mask_all.pt : {input_latents [16,21,64,64] bf16, masked_latents [16,21,64,64] bf16} audio_emb_omniavatar.pt : {audio_emb [T,10752] f32, metadata} (T = video frames) ref_latents.pt : {ref_sequence_latents [16,21,64,64] bf16, metadata} text_emb.pt : [1,512,4096] f32 (or {key: tensor}) Encode mechanics: * Video frames: cv2 BGR->RGB, ``cv2.resize`` to 512x512, normalize ``/127.5 - 1``, short clips padded by repeating the last frame; first ``num_frames`` frames are the GT segment (start_frame=0); the reference segment starts at ``num_frames`` when ``total_frames >= 2*num_frames`` else ``max(0, total-num)``. * Spatial mask: LatentSync PNG, ``convert('L')`` /255, ``cv2.INTER_LINEAR`` resize to pixel res, ``> 0.5`` binarize. Convention **1 = keep, 0 = mask**. Masking is applied to ALL frames (including frame 0) in pixel space before VAE encode. * VAE: ``WanVideoVAE`` (z=16); encoded in **bf16** (bf16 weights + bf16 input); 81 px frames -> 21 latent frames, 512 -> 64. * Audio: OmniAvatar custom ``Wav2VecModel`` called with ``seq_len=`` and ``output_hidden_states=True``; the 10752 feature is ``last_hidden_state`` concatenated with all 13 ``hidden_states`` (14 x 768). The CNN features are linearly interpolated to ``seq_len`` frames, so the audio is encoded at the full video's frame count and sliced to ``num_video_frames`` downstream. * Text: ``WanTextEncoder`` (UMT5-XXL) via ``WanPrompter`` -> [1,512,4096]. """ import hashlib import math import os import tempfile import cv2 import librosa import numpy as np import torch import torch.nn.functional as F from PIL import Image # Default media geometry (OmniAvatar V2V 512x512 @ 25fps, Wan VAE 4x temporal). DEFAULT_NUM_FRAMES = 81 DEFAULT_HEIGHT = 512 DEFAULT_WIDTH = 512 DEFAULT_FPS = 25 WAV2VEC_SR = 16000 # =========================================================================== # Encoder model loaders (frozen, eval) — load once, reuse for the whole run. # =========================================================================== def load_vae(vae_path, device): """Load the Wan 2.1 Video VAE (``WanVideoVAE``, z=16) in eval mode on *device*. Mirrors the inference loader: tolerates ``model.``-prefixed, ``model_state`` (CivitAI) and flat key formats. """ from OmniAvatar.models.wan_video_vae import WanVideoVAE vae = WanVideoVAE(z_dim=16) state_dict = torch.load(vae_path, map_location="cpu", weights_only=False) if any(k.startswith("model.") for k in state_dict): vae.load_state_dict(state_dict, strict=True) elif "model_state" in state_dict: converter = WanVideoVAE.state_dict_converter() vae.load_state_dict(converter.from_civitai(state_dict), strict=True) else: prefixed = {"model." + k: v for k, v in state_dict.items()} vae.load_state_dict(prefixed, strict=True) vae = vae.to(device=device) vae.eval() return vae def load_wav2vec(wav2vec_path, device): """Load the OmniAvatar custom ``Wav2VecModel`` + feature extractor. Returns ``(model, extractor)`` with the model frozen in eval/float32 on *device*. The feature extractor (CNN) must stay float32. """ from transformers import Wav2Vec2FeatureExtractor from OmniAvatar.models.wav2vec import Wav2VecModel extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_path) model = Wav2VecModel.from_pretrained(wav2vec_path, attn_implementation="eager") model.feature_extractor.requires_grad_(False) model = model.to(device).float() model.eval() model.requires_grad_(False) return model, extractor def _resolve_tokenizer_path(text_encoder_path, tokenizer_path=None): """Locate the UMT5 tokenizer directory for a Wan text-encoder checkpoint. Wan layouts keep the tokenizer either directly beside the ``.pth`` or in a ``google/umt5-xxl`` subdirectory. Prefers an explicit *tokenizer_path*, then the ``google/umt5-xxl`` subdir, then the checkpoint's own directory. """ if tokenizer_path is not None: return tokenizer_path base = os.path.dirname(text_encoder_path) subdir = os.path.join(base, "google", "umt5-xxl") if os.path.isdir(subdir): return subdir return base def load_text_encoder(text_encoder_path, device, dtype=torch.bfloat16, tokenizer_path=None): """Load the Wan UMT5-XXL text encoder + prompter for prompt encoding. Returns ``(text_encoder, prompter)``. The tokenizer is resolved via :func:`_resolve_tokenizer_path` (explicit, ``google/umt5-xxl`` subdir, or the checkpoint directory). """ from OmniAvatar.models.wan_video_text_encoder import WanTextEncoder from OmniAvatar.prompters.wan_prompter import WanPrompter text_encoder = WanTextEncoder() te_state = torch.load(text_encoder_path, map_location="cpu", weights_only=False) converter = WanTextEncoder.state_dict_converter() te_state = converter.from_civitai(te_state) text_encoder.load_state_dict(te_state, strict=True) text_encoder = text_encoder.to(device).eval() text_encoder.requires_grad_(False) prompter = WanPrompter( tokenizer_path=_resolve_tokenizer_path(text_encoder_path, tokenizer_path), text_len=512, ) prompter.fetch_models(text_encoder=text_encoder) return text_encoder, prompter def load_encoders( vae_path, wav2vec_path, device, dtype=torch.bfloat16, text_encoder_path=None, load_text=True, ): """Load all frozen encoders needed for on-the-fly preprocessing. Returns a dict with keys ``vae``, ``wav2vec``, ``wav2vec_extractor`` and (when *load_text* and *text_encoder_path* are given) ``text_encoder`` and ``prompter``. The text encoder (~11B UMT5-XXL) is optional so callers that only need VAE/audio (or that serve every prompt from cache) can skip it. """ encoders = {"device": device, "dtype": dtype} encoders["vae"] = load_vae(vae_path, device) wav2vec, extractor = load_wav2vec(wav2vec_path, device) encoders["wav2vec"] = wav2vec encoders["wav2vec_extractor"] = extractor if load_text and text_encoder_path is not None: text_encoder, prompter = load_text_encoder(text_encoder_path, device, dtype) encoders["text_encoder"] = text_encoder encoders["prompter"] = prompter else: encoders["text_encoder"] = None encoders["prompter"] = None return encoders # =========================================================================== # Frame loading + spatial mask # =========================================================================== def read_video_frames_pixel(video_path, num_frames, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, start_frame=0): """Read *num_frames* frames -> ``[T, H, W, 3]`` float32 in ``[-1, 1]``. Reads frames with cv2, BGR->RGB, ``cv2.resize`` to (W, H), normalizes ``/127.5 - 1``, pads to *num_frames* by repeating the last frame. Also returns the video's total frame count (``CAP_PROP_FRAME_COUNT``) which drives the reference-segment offset and the audio sequence length. Returns: (frames [T, H, W, 3] float32 in [-1, 1], total_frames int) """ cap = cv2.VideoCapture(video_path) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if start_frame > 0: cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) frames = [] for _ in range(num_frames): ret, frame = cap.read() if ret: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = cv2.resize(frame, (width, height)) frame = frame.astype(np.float32) / 127.5 - 1.0 frames.append(frame) else: break cap.release() if len(frames) == 0: frames.append(np.zeros((height, width, 3), dtype=np.float32)) while len(frames) < num_frames: frames.append(frames[-1].copy()) return np.stack(frames), total_frames def frames_pixel_to_tensor(frames_thwc): """``[T, H, W, 3]`` (numpy or tensor) -> ``[3, T, H, W]`` float32 tensor. Matches the precompute layout: ``torch.from_numpy(frames).permute(3,0,1,2)``. """ if isinstance(frames_thwc, np.ndarray): t = torch.from_numpy(frames_thwc) else: t = frames_thwc return t.permute(3, 0, 1, 2).contiguous().float() def binarize_pixel_mask(mask_path, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH): """Load LatentSync mask -> binary pixel mask ``[H, W]`` float32, 1=keep 0=mask. Replicates the precompute worker: ``Image.open(...).convert('L')`` /255, ``cv2.INTER_LINEAR`` resize to (W, H) when needed, ``> 0.5`` binarize. """ mask_img = Image.open(mask_path).convert("L") mask_np = np.array(mask_img).astype(np.float32) / 255.0 if mask_np.shape[0] != height or mask_np.shape[1] != width: mask_np = cv2.resize(mask_np, (width, height), interpolation=cv2.INTER_LINEAR) return torch.from_numpy((mask_np > 0.5).astype(np.float32)) # [H, W] def load_latent_mask(mask_path, latent_h=64, latent_w=64): """Load LatentSync mask and resize to latent resolution -> ``[h, w]`` float. Convention 1=keep, 0=mask. Matches the dataloader's ``self.mask`` (bilinear interpolate to latent res, ``> 0.5``) and inference's ``load_latentsync_mask``. """ mask_img = Image.open(mask_path).convert("L") mask_arr = np.array(mask_img).astype(np.float32) / 255.0 mask_t = torch.from_numpy(mask_arr).unsqueeze(0).unsqueeze(0) # [1, 1, H, W] mask_resized = F.interpolate( mask_t, size=(latent_h, latent_w), mode="bilinear", align_corners=False ) return (mask_resized > 0.5).float().squeeze(0).squeeze(0) # [h, w] # --------------------------------------------------------------------------- # Inference-side frame helpers. These use the ``/255*2-1`` normalization and # ping-pong length extension of the inference path, mathematically equivalent in # pixel space to the ``/127.5-1`` normalization above (``x/255*2-1 == x/127.5-1``). # --------------------------------------------------------------------------- def pingpong_indices(n, target): """Frame indices extending a length-*n* sequence to *target* via ping-pong.""" if n >= target: return list(range(target)) if n == 1: return [0] * target cycle = list(range(n)) + list(range(n - 2, 0, -1)) indices = [] while len(indices) < target: indices.extend(cycle) return indices[:target] def adjust_video_length(frames_np, target): """Adjust ``[N,H,W,3]`` video to exactly *target* frames (ping-pong / clip).""" n = len(frames_np) if n >= target: return frames_np[:target] return frames_np[pingpong_indices(n, target)] def frames_to_tensor(frames_np): """``[N,H,W,3]`` uint8 -> ``[1,3,N,H,W]`` float in ``[-1,1]`` (inference path).""" t = torch.from_numpy(frames_np).float() / 255.0 t = t.permute(0, 3, 1, 2) t = t * 2.0 - 1.0 t = t.unsqueeze(0).permute(0, 2, 1, 3, 4) return t def apply_spatial_mask(video_tensor, mask_np, mask_all_frames=True): """Apply a ``[H,W]`` 1=keep/0=mask binary mask to ``[1,3,N,H,W]`` video.""" mask_t = torch.from_numpy(mask_np).float()[None, None, None, :, :] masked = video_tensor.clone() if mask_all_frames: masked *= mask_t else: masked[:, :, 1:, :, :] *= mask_t return masked # =========================================================================== # VAE encode # =========================================================================== def vae_encode_pixels(vae, pixel_tensor, device, dtype=torch.bfloat16): """VAE-encode a pixel-space video tensor in bf16 (weights + input). Args: vae: ``WanVideoVAE`` instance. pixel_tensor: ``[3, T, H, W]`` (single) or ``[B, 3, T, H, W]`` (batch), float in ``[-1, 1]``. device, dtype: compute device and output dtype. Returns: ``[16, T_lat, H_lat, W_lat]`` (single) or ``[B, 16, ...]`` (batch). The VAE is temporarily cast to bf16 for the encode and restored afterwards, matching the precompute pipeline (bf16 VAE + bf16 input). ``WanVideoVAE.encode`` iterates its argument and unsqueezes each element, so a single ``[3,T,H,W]`` is passed as a one-element list and a ``[B,3,T,H,W]`` batch is passed directly. """ single = pixel_tensor.dim() == 4 original_dtype = next(vae.parameters()).dtype vae.to(dtype=torch.bfloat16) try: with torch.no_grad(): if single: latents = vae.encode([pixel_tensor.to(dtype=torch.bfloat16)], device=device) else: latents = vae.encode(pixel_tensor.to(dtype=torch.bfloat16), device=device, tiled=False) finally: vae.to(dtype=original_dtype) latents = latents.to(dtype=dtype) if single: latents = latents[0] # drop batch dim -> [16, T_lat, H_lat, W_lat] return latents def encode_vae_masked(vae, gt_pixel, pixel_mask, device, dtype=torch.bfloat16): """Encode GT + spatially-masked video -> ``input_latents`` and ``masked_latents``. The mouth region is masked on **every** frame (including frame 0) in a single encode pass. Args: vae: ``WanVideoVAE``. gt_pixel: ``[3, T, H, W]`` float in ``[-1, 1]`` (unmasked GT frames). pixel_mask: ``[H, W]`` binary, 1=keep, 0=mask. Returns: (input_latents [16, T_lat, H, W], masked_latents [16, T_lat, H, W]) in *dtype*. """ mask_2d = pixel_mask.view(1, 1, pixel_mask.shape[-2], pixel_mask.shape[-1]).to(gt_pixel.dtype) input_latents = vae_encode_pixels(vae, gt_pixel, device, dtype) masked_latents = vae_encode_pixels(vae, gt_pixel * mask_2d, device, dtype) return input_latents, masked_latents def ref_segment_start(total_frames, num_frames=DEFAULT_NUM_FRAMES): """Reference-segment start frame (matches precompute): ``num_frames`` if ``total_frames >= 2*num_frames`` else ``max(0, total-num)``. """ if total_frames >= 2 * num_frames: return num_frames return max(0, total_frames - num_frames) def encode_ref(vae, video_path, num_frames=DEFAULT_NUM_FRAMES, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, device="cuda", dtype=torch.bfloat16, total_frames=None): """Encode the reference segment -> ``ref_sequence_latents [16, T_lat, H, W]``. Reads *num_frames* unmasked frames starting at :func:`ref_segment_start` and VAE-encodes them. Byte-matches ``ref_latents.pt:ref_sequence_latents``. """ if total_frames is None: cap = cv2.VideoCapture(video_path) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) cap.release() ref_start = ref_segment_start(total_frames, num_frames) ref_frames, _ = read_video_frames_pixel( video_path, num_frames, height, width, start_frame=ref_start ) ref_pixel = frames_pixel_to_tensor(ref_frames) ref_latents = vae_encode_pixels(vae, ref_pixel, device, dtype) return ref_latents, ref_start, total_frames # =========================================================================== # Audio encode (OmniAvatar custom Wav2Vec, 10752-dim concat) # =========================================================================== def audio_natural_frames(audio, sample_rate=WAV2VEC_SR, fps=DEFAULT_FPS): """Natural frame count of an audio source: ``ceil(num_samples / (sr // fps))``. *audio* is a path or a 1-D numpy waveform (already at *sample_rate*). """ if isinstance(audio, str): wav, _ = librosa.load(audio, sr=sample_rate) n = len(wav) else: n = len(audio) samples_per_frame = sample_rate // fps return int(math.ceil(n / samples_per_frame)) def encode_audio_omniavatar(wav2vec_model, wav2vec_extractor, audio, seq_len, device, sample_rate=WAV2VEC_SR, fps=DEFAULT_FPS): """Encode audio -> ``[1, seq_len, 10752]`` via the OmniAvatar custom Wav2Vec. This is the shared core used by both training and inference. *audio* is a path or a 1-D numpy waveform at *sample_rate*. The waveform is normalized by the feature extractor, zero-padded up to ``seq_len * (sr // fps)`` samples, and the CNN features are linearly interpolated to *seq_len* frames; the 10752 feature is ``last_hidden_state`` concatenated with all 13 ``hidden_states`` (14 x 768). Note: the wav2vec model runs in float32, so the output is float32 (the precompute ``audio_emb`` dtype). Callers cast to bf16 downstream. """ if isinstance(audio, str): audio, _ = librosa.load(audio, sr=sample_rate) input_values = np.squeeze( wav2vec_extractor(audio, sampling_rate=sample_rate).input_values ) input_values = torch.from_numpy(input_values).float().to(device=device) input_values = input_values.unsqueeze(0) samples_per_frame = sample_rate // fps # 640 at 16kHz/25fps target_samples = seq_len * samples_per_frame if input_values.shape[1] < target_samples: input_values = F.pad(input_values, (0, target_samples - input_values.shape[1])) with torch.no_grad(): hidden_states = wav2vec_model( input_values, seq_len=seq_len, output_hidden_states=True ) audio_emb = hidden_states.last_hidden_state for hs in hidden_states.hidden_states: audio_emb = torch.cat((audio_emb, hs), -1) return audio_emb # [1, seq_len, 10752] # =========================================================================== # Text encode # =========================================================================== def encode_text(prompter, prompt, device, dtype=torch.bfloat16): """Encode a prompt string -> ``[1, 512, 4096]`` text embedding in *dtype*. Uses ``WanPrompter.encode_prompt`` (UMT5-XXL). Matches the inference and precompute text-embedding format. """ with torch.no_grad(): text_embeds = prompter.encode_prompt(prompt, positive=True, device=device) if text_embeds.dim() == 2: text_embeds = text_embeds.unsqueeze(0) return text_embeds.to(dtype=dtype) def compute_prompt_hash(prompt): """Content address for a prompt string: ``sha1(prompt)`` hex digest.""" return hashlib.sha1(prompt.encode("utf-8")).hexdigest() # =========================================================================== # Raw sample preparation (CPU, dataloader-worker safe) + GPU encode # =========================================================================== def prepare_raw_sample(video_path, audio_path, prompt, pixel_mask, num_video_frames=DEFAULT_NUM_FRAMES, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, sample_rate=WAV2VEC_SR, fps=DEFAULT_FPS, load_audio_waveform=True): """Decode raw inputs on CPU into the payload the GPU encode step consumes. Safe to run inside dataloader workers (no models, no GPU). Reads the GT and reference frame segments, loads the audio waveform, and carries the prompt plus its content hash and the sequence lengths needed to reproduce the precompute tensors. Args: pixel_mask: ``[H, W]`` binary mask, 1=keep 0=mask (shared, precomputed by :func:`binarize_pixel_mask`). Passed through unchanged. Returns a dict with CPU tensors / scalars: gt_pixel [3, T, H, W], ref_pixel [3, T, H, W], enc_waveform [L] (or None), audio_path, prompt, prompt_hash, total_frames, ref_start, seq_len, num_video_frames, pixel_mask. ``enc_waveform`` is the raw 16 kHz waveform for wav2vec encoding; it is named distinctly so it never collides with the reward path's ``audio_waveform``. """ gt_frames, total_frames = read_video_frames_pixel( video_path, num_video_frames, height, width, start_frame=0 ) gt_pixel = frames_pixel_to_tensor(gt_frames) ref_start = ref_segment_start(total_frames, num_video_frames) ref_frames, _ = read_video_frames_pixel( video_path, num_video_frames, height, width, start_frame=ref_start ) ref_pixel = frames_pixel_to_tensor(ref_frames) waveform = None if load_audio_waveform and audio_path is not None and os.path.exists(audio_path): wav, _ = librosa.load(audio_path, sr=sample_rate) waveform = torch.from_numpy(wav) # Audio is encoded at the full video frame count (precompute convention: the # wav2vec features are interpolated to the video frame grid), guarded to be at # least num_video_frames so the downstream [:num_video_frames] slice is valid. seq_len = max(total_frames, num_video_frames) return { "gt_pixel": gt_pixel, "ref_pixel": ref_pixel, "enc_waveform": waveform, "audio_path": audio_path if audio_path is not None else "", "prompt": prompt, "prompt_hash": compute_prompt_hash(prompt), "total_frames": int(total_frames), "ref_start": int(ref_start), "seq_len": int(seq_len), "num_video_frames": int(num_video_frames), "pixel_mask": pixel_mask, } def _load_text_cache_disk(text_cache_dir, phash): path = os.path.join(text_cache_dir, f"{phash}.pt") if text_cache_dir else None if path and os.path.exists(path): try: return torch.load(path, map_location="cpu", weights_only=False) except Exception: return None return None def _save_text_cache_disk(text_cache_dir, phash, text_emb): if not text_cache_dir: return try: os.makedirs(text_cache_dir, exist_ok=True) path = os.path.join(text_cache_dir, f"{phash}.pt") # Unique temp name so concurrent ranks sharing the cache dir don't # clobber each other's .tmp before the atomic rename. fd, tmp = tempfile.mkstemp(dir=text_cache_dir, suffix=".tmp") os.close(fd) torch.save(text_emb.cpu(), tmp) os.replace(tmp, path) except Exception: pass # best-effort cache def encode_prepared(encoders, raw, device=None, dtype=torch.bfloat16, text_emb_cache=None, text_cache_dir=None): """GPU-encode a prepared raw sample into the precompute ``.pt`` tensors. Runs in the main training process (encoders live here, never in workers). Args: encoders: dict from :func:`load_encoders` (vae, wav2vec[+extractor], optionally prompter). raw: dict from :func:`prepare_raw_sample`. text_emb_cache: optional ``{prompt_hash: tensor}`` in-memory cache so each unique prompt is encoded by UMT5 only once per run. text_cache_dir: optional directory for a disk, content-addressed text cache (``{sha1(prompt)}.pt``) reused across epochs and runs. Resolution order for ``text_emb``: a value preloaded on the raw sample (``raw['text_emb']``), then the in-memory cache, then the disk cache, then a fresh UMT5 encode (when a prompter is loaded), else None. Returns a dict mirroring the precomputed files (full, unsliced): input_latents, masked_latents, ref_sequence_latents [16, T_lat, H, W] bf16, audio_emb [seq_len, 10752] f32, text_emb [1, 512, 4096] (in *dtype*), plus prompt_hash / total_frames / ref_start / seq_len / num_video_frames. """ device = device or encoders.get("device") vae = encoders["vae"] wav2vec = encoders["wav2vec"] extractor = encoders["wav2vec_extractor"] input_latents, masked_latents = encode_vae_masked( vae, raw["gt_pixel"].to(device), raw["pixel_mask"].to(device), device, dtype, ) ref_latents = vae_encode_pixels(vae, raw["ref_pixel"].to(device), device, dtype) audio_src = (raw["enc_waveform"].numpy() if raw.get("enc_waveform") is not None else raw["audio_path"]) audio_emb = encode_audio_omniavatar( wav2vec, extractor, audio_src, raw["seq_len"], device ).squeeze(0).to(torch.float32) # Text: content-addressed by prompt hash so UMT5 runs once per unique prompt. phash = raw["prompt_hash"] text_emb = raw.get("text_emb") # preloaded from sample dir, if any if text_emb is None and text_emb_cache is not None and phash in text_emb_cache: text_emb = text_emb_cache[phash] if text_emb is None: text_emb = _load_text_cache_disk(text_cache_dir, phash) if text_emb is None and encoders.get("prompter") is not None: text_emb = encode_text(encoders["prompter"], raw["prompt"], device, dtype) _save_text_cache_disk(text_cache_dir, phash, text_emb) if text_emb is not None: if text_emb.dim() == 2: text_emb = text_emb.unsqueeze(0) text_emb = text_emb.to(dtype) if text_emb_cache is not None: text_emb_cache[phash] = text_emb return { "input_latents": input_latents.cpu(), "masked_latents": masked_latents.cpu(), "ref_sequence_latents": ref_latents.cpu(), "audio_emb": audio_emb.cpu(), "text_emb": text_emb.cpu() if text_emb is not None else None, "prompt_hash": phash, "total_frames": raw["total_frames"], "ref_start": raw["ref_start"], "seq_len": raw["seq_len"], "num_video_frames": raw["num_video_frames"], } def encode_sample_from_files(encoders, video_path, audio_path, prompt, mask_path, num_video_frames=DEFAULT_NUM_FRAMES, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, device=None, dtype=torch.bfloat16): """Convenience: decode + encode a sample straight from files. Composes :func:`prepare_raw_sample` and :func:`encode_prepared`. Used for any standalone encode of a raw ``(video, audio, prompt)`` triple. """ device = device or encoders.get("device") pixel_mask = binarize_pixel_mask(mask_path, height, width) raw = prepare_raw_sample( video_path, audio_path, prompt, pixel_mask, num_video_frames=num_video_frames, height=height, width=width, ) return encode_prepared(encoders, raw, device, dtype)