Preprocess Cosmos-Predict2 overfit data into parquet format.

Encodes videos with the Cosmos (Wan-style) VAE and captions with the single T5 Large text encoder into the t2v parquet schema expected by the training framework.

Usage

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

`fastvideo.pipelines.preprocess.preprocess_cosmos_overfit.T5LargeConfig`

T5LargeConfig(arch_config: TextEncoderArchConfig = T5LargeArchConfig(), prefix: str = 't5', quant_config: QuantizationConfig | None = None, lora_config: Any | None = None, is_chat_model: bool = False, treat_empty_as_dot: bool = False, *, _resolved_attention_backend: AttentionBackendEnum | None = None, chat_template_enable_thinking: bool = False)

T5 Large configuration for your specific model.

`fastvideo.pipelines.preprocess.preprocess_cosmos_overfit.BaseEncoderOutput`

BaseEncoderOutput(last_hidden_state: torch.FloatTensor | None = None, pooler_output: torch.FloatTensor | None = None, hidden_states: tuple[torch.FloatTensor, ...] | None = None, attentions: tuple[torch.FloatTensor, ...] | None = None, attention_mask: torch.Tensor | None = None)

`fastvideo.pipelines.preprocess.preprocess_cosmos_overfit.t5_large_postprocess_text`

t5_large_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor

Postprocess T5 Large text encoder outputs for Cosmos pipeline.

Return raw last_hidden_state without truncation/padding.

`fastvideo.pipelines.preprocess.preprocess_cosmos_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_cosmos_overfit.load_video`

load_video(path: str, num_frames: int) -> torch.Tensor

Load video as [1, C, T, H, W] in [-1, 1].

Source

def load_video(path: str, num_frames: int) -> torch.Tensor:
    """Load video as [1, C, T, H, W] in [-1, 1]."""
    cap = cv2.VideoCapture(path)
    frames: list[np.ndarray] = []
    while len(frames) < num_frames:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame)
    cap.release()

    if len(frames) < num_frames:
        # Repeat last frame to fill
        while len(frames) < num_frames:
            frames.append(frames[-1])

    frames = frames[:num_frames]
    video = np.stack(frames, axis=0)
    video = torch.from_numpy(video).float()
    video = video / 127.5 - 1.0  # [0,255] -> [-1,1]
    video = video.permute(3, 0, 1, 2).unsqueeze(0)  # [1,C,T,H,W]
    return video

`fastvideo.pipelines.preprocess.preprocess_cosmos_overfit.main`

main() -> None

Source

def main() -> None:
    device = torch.device("cuda:0")
    model_path = maybe_download_model(MODEL_REPO)

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # Load captions
    with open(os.path.join(DATA_DIR, "videos2caption.json")) as f:
        caption_data = json.load(f)

    # --- Load VAE ---
    # Cosmos-Predict2-2B ships a Wan-style VAE; vae/config.json declares
    # `_class_name: AutoencoderKLWan`, so diffusers can load it directly.
    print("Loading Cosmos VAE (AutoencoderKLWan)...")
    from diffusers import AutoencoderKLWan
    vae = AutoencoderKLWan.from_pretrained(
        model_path,
        subfolder="vae",
        torch_dtype=torch.float16,
    ).to(device).eval()
    print(f"VAE loaded ({sum(p.numel() for p in vae.parameters())/1e6:.0f}M)")

    # --- Load T5 Large text encoder ---
    print("Loading T5 Large text encoder...")
    from transformers import AutoTokenizer, T5EncoderModel

    t5_cfg = T5LargeConfig()
    tok_kwargs = dict(t5_cfg.tokenizer_kwargs)
    tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
    text_encoder = T5EncoderModel.from_pretrained(
        os.path.join(model_path, "text_encoder"),
        torch_dtype=torch.bfloat16,
    ).to(device).eval()

    # --- 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]
        video_path = os.path.join(DATA_DIR, "videos", video_name)

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

        # Encode video
        video = load_video(video_path, NUM_FRAMES).to(device=device, dtype=torch.float16)
        print(f"  Video shape: {video.shape}")

        with torch.no_grad():
            latent_dist = vae.encode(video).latent_dist
            # Cast to fp32 — dataloader hardcodes np.float32
            latent = latent_dist.mean.squeeze(0).float().cpu()
        print(f"  Latent shape: {latent.shape}")

        # Encode text with T5
        with torch.no_grad():
            inputs = tokenizer(caption, **tok_kwargs).to(device)
            outputs = text_encoder(**inputs)
            enc_out = BaseEncoderOutput(
                last_hidden_state=outputs.last_hidden_state,
                attention_mask=inputs["attention_mask"],
            )
            # [1, max_len, 1024], zeros beyond real length
            t5_embed = t5_large_postprocess_text(enc_out).squeeze(0)
            # Trim to real sequence length so dataloader's pad() builds
            # the correct attention mask.
            real_len = int(inputs["attention_mask"].sum().item())
            text_embedding = t5_embed[:real_len].float().cpu()  # [seq, 1024]
        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)

    # Clean up encoders
    del text_encoder, tokenizer, vae

    # Write parquet
    table = pa.table(
        {k: [r[k] for r in records]
         for k in records[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(records)} records to {output_path}")

    # Extract first frame from first video as V2W conditioning image
    import cv2
    first_video = os.path.join(DATA_DIR, "videos", caption_data[0]["path"])
    cap = cv2.VideoCapture(first_video)
    ret, frame = cap.read()
    cap.release()
    cond_frame_path = os.path.join(OUTPUT_DIR, "cond_frame.png")
    if ret:
        cv2.imwrite(cond_frame_path, frame)
        print(f"Saved conditioning frame to {cond_frame_path}")

    # Write validation prompts for callback
    # Wrap in "data" key — ValidationDataset expects field="data"
    # Use "caption" field — ValidationDataset aliases it to "prompt"
    # Include image_path for V2W conditioning during validation
    val_prompts = {
        "data": [{
            "caption": item["cap"][0],
            "image_path": "cond_frame.png",
        } 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.")