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) -> strCheck 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: strLocal path, Hugging Face Hub model ID, or
org/repo/subfolderumbrella-repo reference.local_dir: str | NoneLocal directory to save the model
download: boolWhether to download the model from Hugging Face Hub
revision: str | NoneOptional immutable Hub revision.
allow_patterns: list[str] | NoneOptional Hub glob patterns limiting downloaded files. Local paths are returned unchanged. For umbrella references, the patterns are interpreted relative to the selected subfolder.
Returns
strLocal 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: strPath to the model directory
required_component_dirs: list[str] | tuple[str, ...] | NoneComponent directories required by the selected pipeline.
Nonepreserves 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</code></pre></details>
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.")</code></pre></details>