--- title: preprocess_ltx2_overfit ---

Preprocess LTX-2 overfit data into parquet format.

Encodes videos with the LTX-2 causal video VAE and captions with the Gemma text encoder (feature extractor + embedding connector) into the t2v parquet schema expected by the training framework.

The stored text embeddings are POST-connector: the connector replaces pad positions with learnable registers and returns an all-valid mask, so the parquet collate's ones/zeros mask stays semantically correct and training needs no text encoder at all. Captions are encoded via the encoder's forward() (the exact inference path), which handles both LTX-2.0 (shared 3840-d features) and LTX-2.3 (separate 4096-d video / 2048-d audio feature extractors).

Videos are resampled to TRAIN_FPS and the preprocessed clip is also saved as an mp4 next to the parquet so overfit tests can use it as the SSIM reference.

Usage

CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py

fastvideo.pipelines.preprocess.preprocess_ltx2_overfit.maybe_download_model

maybe_download_model(model_name_or_path: str, local_dir: str | None = None, download: bool = True, revision: str | None = None, allow_patterns: list[str] | None = None) -> str

Check if the model path is a Hugging Face Hub model ID and download it if needed.

Supports an "umbrella" repo layout where a single HF repo holds multiple pipeline variants under sibling subfolders. If the input is shaped as org/repo/subfolder (i.e. a non-existent local path with 3+ slash- separated components and at least one segment that does not look like a posix-absolute path), treat the first two components as the HF repo id and the remainder as a subfolder; only the subfolder's blobs are downloaded, and the returned local path points inside that subfolder.

Parameters

model_name_or_path: str

Local path, Hugging Face Hub model ID, or org/repo/subfolder umbrella-repo reference.

local_dir: str | None

Local directory to save the model

download: bool

Whether to download the model from Hugging Face Hub

revision: str | None

Optional immutable Hub revision.

allow_patterns: list[str] | None

Optional Hub glob patterns limiting downloaded files. Local paths are returned unchanged. For umbrella references, the patterns are interpreted relative to the selected subfolder.

Returns

str

Local path to the model (or to the subfolder inside the snapshot).

fastvideo.pipelines.preprocess.preprocess_ltx2_overfit.verify_model_config_and_directory

verify_model_config_and_directory(model_path: str, required_component_dirs: list[str] | tuple[str, ...] | None = None) -> dict[str, Any]

Verify that the model directory contains a valid Diffusers configuration.

Parameters

model_path: str

Path to the model directory

required_component_dirs: list[str] | tuple[str, ...] | None

Component directories required by the selected pipeline. None preserves full-snapshot validation; an empty collection validates only the manifest.

Returns

dict[str, Any]

The loaded model configuration as a dictionary

fastvideo.pipelines.preprocess.preprocess_ltx2_overfit.load_video

load_video(path: str, num_frames: int, target_fps: float, height: int, width: int) -> tuple[torch.Tensor, np.ndarray]

Load a video as [1, C, T, H, W] in [-1, 1], resampled to target_fps.

Also returns the uint8 RGB frames [T, H, W, C] for reference-video export.

Source
def load_video(path: str, num_frames: int, target_fps: float, height: int,
               width: int) -> tuple[torch.Tensor, np.ndarray]:
    """Load a video as [1, C, T, H, W] in [-1, 1], resampled to target_fps.

    Also returns the uint8 RGB frames [T, H, W, C] for reference-video export.
    """
    with av.open(path) as container:
        if not container.streams.video:
            raise RuntimeError(f"No video stream found in {path}")
        stream = container.streams.video[0]
        native_fps = float(stream.average_rate or target_fps)
        decoded = [torch.from_numpy(frame.to_ndarray(format="rgb24")) for frame in container.decode(video=0)]
    if not decoded:
        raise RuntimeError(f"Could not read any frames from {path}")
    raw = torch.stack(decoded).permute(0, 3, 1, 2)  # [T, C, H, W] uint8
    step = native_fps / target_fps
    wanted = [min(int(round(i * step)), raw.shape[0] - 1) for i in range(num_frames)]

    frames = raw[wanted].float()  # [T, C, H, W] in [0, 255]
    src_h, src_w = frames.shape[2], frames.shape[3]
    scale = max(height / src_h, width / src_w)
    new_h, new_w = int(round(src_h * scale)), int(round(src_w * scale))
    frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True)
    top = (new_h - height) // 2
    left = (new_w - width) // 2
    frames = frames[:, :, top:top + height, left:left + width]

    frames_np = (frames.permute(0, 2, 3, 1).clamp(0, 255).round().to(torch.uint8).numpy())
    video = frames / 127.5 - 1.0  # [0,255] -> [-1,1]
    video = video.permute(1, 0, 2, 3).unsqueeze(0)  # [1,C,T,H,W]
    return video, frames_np

fastvideo.pipelines.preprocess.preprocess_ltx2_overfit.main

main() -> None
Source
def main() -> None:
    _init_single_process_distributed()

    from fastvideo.fastvideo_args import FastVideoArgs
    from fastvideo.models.loader.component_loader import (
        PipelineComponentLoader, )
    from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig

    data_dir = envs.FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR.get()
    caption_json = envs.FASTVIDEO_TEST_LTX2_OVERFIT_CAPTION_JSON.get()
    video_subdir = envs.FASTVIDEO_TEST_LTX2_OVERFIT_VIDEO_SUBDIR.get()
    output_dir = envs.FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR.get()
    model_repo = envs.FASTVIDEO_TEST_LTX2_OVERFIT_MODEL.get()
    # The train dataloader samples with drop_last=True across data-parallel
    # groups, so the dataset must hold at least num_sp_groups * batch_size
    # rows or every rank gets zero batches. Replicate the overfit sample so
    # a 4-GPU FSDP run still sees one batch per rank.
    num_copies = envs.FASTVIDEO_TEST_LTX2_OVERFIT_NUM_COPIES.get()

    device = torch.device("cuda:0")
    model_path = maybe_download_model(model_repo)
    model_index = verify_model_config_and_directory(model_path)

    os.makedirs(output_dir, exist_ok=True)
    # The map-style dataset caches parquet file metadata; a stale cache
    # next to a regenerated parquet can crash or serve old rows.
    shutil.rmtree(os.path.join(output_dir, "map_style_cache"), ignore_errors=True)

    with open(os.path.join(data_dir, caption_json)) as f:
        caption_data = json.load(f)

    pipeline_config = LTX2T2VConfig()
    fastvideo_args = FastVideoArgs(
        model_path=model_path,
        pipeline_config=pipeline_config,
        num_gpus=1,
        tp_size=1,
        sp_size=1,
        hsdp_shard_dim=1,
        vae_cpu_offload=False,
        text_encoder_cpu_offload=False,
    )

    def load_component(name: str) -> Any:
        transformers_or_diffusers, _ = model_index[name]
        return PipelineComponentLoader.load_module(
            module_name=name,
            component_model_path=os.path.join(model_path, name),
            transformers_or_diffusers=transformers_or_diffusers,
            fastvideo_args=fastvideo_args,
        )

    print("Loading LTX-2 VAE...")
    vae = load_component("vae")
    vae_dtype = next(vae.parameters()).dtype
    print(f"VAE loaded ({sum(p.numel() for p in vae.parameters())/1e6:.0f}M, {vae_dtype})")

    print("Loading Gemma text encoder + tokenizer...")
    text_encoder = load_component("text_encoder")
    tokenizer = load_component("tokenizer")
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token

    encoder_config = pipeline_config.text_encoder_configs[0]
    tokenizer_kwargs = dict(encoder_config.tokenizer_kwargs)
    if "max_length" not in tokenizer_kwargs:
        tokenizer_kwargs["max_length"] = encoder_config.arch_config.text_len
    preprocess_text = pipeline_config.preprocess_text_funcs[0]

    # --- Process each video ---
    records = []
    for idx, item in enumerate(caption_data):
        video_name = item["path"]
        record_id = f"{idx:04d}_{video_name}"
        caption = item["cap"][0] if isinstance(item["cap"], list) else item["cap"]
        video_path = os.path.join(data_dir, video_subdir, video_name)

        print(f"\nProcessing: {video_name}")
        print(f"  Caption: {caption[:80]}...")

        video, frames_np = load_video(video_path, NUM_FRAMES, TRAIN_FPS, MAX_HEIGHT, MAX_WIDTH)
        video = video.to(device=device, dtype=vae_dtype)
        print(f"  Video shape: {video.shape}")

        with torch.no_grad():
            # LTX-2 encode() returns a deterministic distribution whose
            # mean is already per-channel normalized; store as-is.
            latent = vae.encode(video).mean.squeeze(0).float().cpu()
        print(f"  Latent shape: {latent.shape}")

        with torch.no_grad():
            # Encode through forward() — the inference text path. The
            # two-step preprocess_text_embeddings + run_connectors route
            # breaks on LTX-2.3, whose separate audio feature extractor
            # is narrower than the video one.
            text_inputs = tokenizer([preprocess_text(caption)], **tokenizer_kwargs)
            encoder_out = text_encoder(
                input_ids=text_inputs["input_ids"].to(device),
                attention_mask=text_inputs["attention_mask"].to(device),
            )
            text_embedding = encoder_out.last_hidden_state.squeeze(0).float().cpu()
        print(f"  Text embedding shape: {text_embedding.shape}")

        record = {
            "id": record_id,
            "vae_latent_bytes": latent.numpy().tobytes(),
            "vae_latent_shape": list(latent.shape),
            "vae_latent_dtype": str(latent.dtype).replace("torch.", ""),
            "text_embedding_bytes": (text_embedding.numpy().tobytes()),
            "text_embedding_shape": list(text_embedding.shape),
            "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""),
            "file_name": video_name,
            "caption": caption,
            "media_type": "video",
            "width": MAX_WIDTH,
            "height": MAX_HEIGHT,
            "num_frames": NUM_FRAMES,
            "duration_sec": NUM_FRAMES / TRAIN_FPS,
            "fps": TRAIN_FPS,
        }
        records.append(record)

        # Save the preprocessed clip so overfit tests can compare
        # validation output against the memorization target.
        import imageio
        ref_path = os.path.join(output_dir, f"training_sample_{idx}.mp4")
        with imageio.get_writer(ref_path, fps=TRAIN_FPS) as writer:
            for frame in frames_np:
                writer.append_data(frame)
        print(f"  Wrote reference clip to {ref_path}")

    # Clean up
    del text_encoder, tokenizer, vae
    torch.cuda.empty_cache()

    # Write parquet (replicated num_copies times; see the comment above num_copies)
    replicated = []
    for copy_idx in range(max(1, num_copies)):
        for r in records:
            row = dict(r)
            row["id"] = f"{r['id']}_copy{copy_idx}"
            replicated.append(row)
    table = pa.table(
        {k: [r[k] for r in replicated]
         for k in replicated[0]},
        schema=pyarrow_schema_t2v,
    )
    output_path = os.path.join(output_dir, "data_00000.parquet")
    pq.write_table(table, output_path)
    print(f"\nWrote {len(replicated)} records "
          f"({len(records)} unique x {max(1, num_copies)} copies) to {output_path}")

    # Write validation prompts for the validation callback
    val_prompts = {
        "data": [{
            "caption": (item["cap"][0] if isinstance(item["cap"], list) else item["cap"]),
        } for item in caption_data],
    }
    val_path = os.path.join(output_dir, "validation_prompts.json")
    with open(val_path, "w") as f:
        json.dump(val_prompts, f, indent=2)
    print(f"Wrote validation prompts to {val_path}")

    print("\nDone! Use data_path: " + output_dir + " in training config.")