Load existing captions from a file. Args: output_path: Path to the captions file format_type: Format of the captions file Retu

Load the first frame from an image or video file as [C, H, W] in [0, 1].

Load full model checkpoint.

Load LoRA checkpoint with DDP/FSDP compatibility.

Load each configured LoRA into a state dict for fusion (pinned path).

Load a mask video and return a pixel-space tensor of shape (1, 1, F, H, W). The mask video is loaded, resized to (height, width), converted to

Initialize and load the VAE model(s).

Load base weights and fuse LoRAs into *meta_model* in-place.

Load the transformer and embeddings processor for training.

Load the non-block weights onto *device* and fuse LoRAs -- both paths. Reads through the loader with ``model_sd_ops`` restricted to the

Load training state file that corresponds to a checkpoint weights file.

Log training statistics.

Build a TilingConfig from explicit sizes. The defaults (1280 px spatial tile, 256 px overlap; 32 temporal frames, 16 overlap) are suitable for

Full memory-efficient ``VideoDecoder.forward`` replacement. Orchestrates the entire decode through workspace-based operations: ``UNetMidBlock3

Memory-efficient ``UNetMidBlock3D`` forward. Allocates a single workspace buffer that is reused across all ``ResnetBlock3D`` iterations. For

Normalize and flatten multi-layer hidden states, respecting padding. Performs per-batch, per-layer normalization using masked mean and range,

Normalize data_sources input to a consistent dict format.

Normalize encoder latents using per-channel statistics. When the encoder is configured with ``double_z=True``, the final conv

Normalize video latents to non-patchified format [C, F, H, W]. Used for keeping backward compatibility with legacy datasets.

Pad sequence length to multiple of alignment for Flash Attention compatibility.

Apply non-causal replicate padding to temporal boundaries. Sets ``workspace[:, :, 0]`` to a copy of ``workspace[:, :, 1]`` and ``workspace[:,

Return params for the given sigma from a sorted (sigma_upper_bound -> params) structure. Keys are sorted descending (bin upper bounds). Bin i

In-place RMS (pixel) normalization along the channel dimension.

transformers <5 layout: separate ``rotary_emb_local`` + ``rotary_emb``.

transformers >=5 layout: single ``rotary_emb`` with per-layer-type buffers. Mirrors ``Gemma3PreTrainedModel._init_weights`` for ``Gemma3RotaryEmbe

Record a compute-done event and release the block weights.

Load GPU weights for a block and inject them into its parameters.

Prepare attention mask.

Prepare audio inputs for joint audio-video training. Args: batch: Raw batch data containing audio_latents sigmas: Samp

Prepare the audio stream for writing.

Prepare context for transformer blocks.

Prepare A-V cross-attention AdaLN inputs.

Prepare models for training with Accelerate.

Prepare self-attention mask by converting [0,1] values to additive log-space bias. Input shape: 3D ``(B, T_q, T_k)`` with values in [0, 1]. Th

Convert video tensor to [F, H, W, C] uint8 numpy array. If ``video_format`` is provided, it is trusted. Otherwise, the layout is auto-detected

Preprocess a single image by resizing and applying transforms.

Preprocess reference video. Returns (preprocessed [-1,1], pixels [0,1]).

Preprocess a video by loading, resizing, and applying transforms. Returns: Tuple of (video tensor in [C, F, H, W] format, fps)

Process a single audio latent file.

Process a single latent file.

Run inference on a single video: generate EXR frames + optional H.264 .mp4 preview.

Process transformer blocks for LTXAV. Per-block perturbation masks are precomputed here and attached to each modality's ``TransformerA

Yield :class:`LoraProduct` items matching *key* across state-dict-backed LoRAs.

Quantize a model block-by-block using optimum-quanto. This approach: 1. Moves each transformer block to GPU 2. Quantizes on GPU (fast!)

Return ``{tensor_name: dtype_string}`` from the safetensors header.

Return ``{post_rename_param_key: scale_tensor}`` for every prequant ``*_scale`` sibling in *checkpoint_path*. Keys are returned in the post-re

Intended to be applied via __class__ reassignment to existing nn.Linear instances. Forward remains defined at the class level, which is requi

Run a ``ResnetBlock3D`` in-place on a workspace buffer. The workspace has shape ``[B, C, T+2, H, W]`` with real data in ``[1:-1]``. After thi

Map canonical role names to actual dataset column names. Returns a dict of role → column_name for recognized roles found in the dataset.

Determine resume state by looking for a training state file next to the loaded checkpoint. Returns (initial_step, TrainingState or None).

Restore optimizer, scheduler, and RNG states from a loaded TrainingState. Must be called after _init_optimizer() (which calls accelerator.prep

Run the Euler denoising loop with CFG/STG, handling frozen modalities.

Run one stage-2 denoising phase with optional IC-LoRA conditioning. Each tile calls ``stage_2.run()`` with a tile-sized ``ModalitySpec`` for

Save an audio waveform tensor to a .wav file using stdlib wave.

Save the training configuration as a YAML file in the output directory.

Save dataset json with reference video paths. Args: reference_paths: Dictionary mapping media paths to reference video paths outpu

Save training state alongside checkpoint for resume. Respects checkpoints.save_training_state config: - "full": optimizer + scheduler

Partition checkpoint keys into per-block and non-block lists. Opens the safetensors files for header-only key enumeration; no tensor data is r

Ask torch whether *backend* can run with the given mask shape. ``MATH`` is the universal SDPA fallback (pure PyTorch ops, no kernel requiremen

Pick a mask-aware attention. Prefers xFormers when installed; else SDPA with the full priority list (the dispatcher rejects FLASH automatically wh

Pick the fastest unmasked attention based on installed extras and GPU arch. Priority by arch: - Hopper (sm_90, H100): FA3 / xFormers (mutually

Initialize the Accelerator with the appropriate settings.

Setup and validate the data root directory.

Configure LoRA adapters for the transformer. Only called in LoRA training mode.

Map data source names to their actual directory paths.

Split a ``BatchedPerturbationConfig`` along the batch dimension.

Temporarily suppress INFO warnings from accelerate about memory allocation.

Replace nn.Linear layers with FP8Linear where ``should_swap(name)`` returns True.

Encode video using spatial tiling for memory efficiency. Splits the video into overlapping spatial tiles, encodes each tile separately, and bl

Convert connector output mask to a binary (0/1) mask shaped ``(B, S, 1)`` for broadcasting.

Fetch torch's current default SDPA priority order at runtime. Used as the default for ``PytorchAttention`` so the wrapper-always code path mat

Perform a single training step using the configured strategy.

Re-encode the video to a constant frame rate so the server's frame sampler can read every requested index (raw / variable-frame-rate videos over-r

Memory-efficient ``DepthToSpaceUpsample`` forward. For non-causal mode the input is copied into a workspace and the convolution runs in-place.