preprocess_cosmos_overfit
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.TensorPostprocess 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) -> 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_cosmos_overfit.load_video
load_video(path: str, num_frames: int) -> torch.TensorLoad 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</code></pre></details>
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.")</code></pre></details>