diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py index 9d8e5e4e43..2b215844c1 100644 --- a/fastvideo/dataset/dataloader/schema.py +++ b/fastvideo/dataset/dataloader/schema.py @@ -67,6 +67,14 @@ pa.field("text_embedding_shape", pa.list_(pa.int64())), # e.g., 'bfloat16' or 'float32' pa.field("text_embedding_dtype", pa.string()), + # Secondary text embedding: ByT5, nullable=True preserves compatibility at the Arrow schema level. + pa.field("text_embedding_2_bytes", pa.binary(), nullable=True), + pa.field( + "text_embedding_2_shape", + pa.list_(pa.int64()), + nullable=True, + ), + pa.field("text_embedding_2_dtype", pa.string(), nullable=True), # --- Metadata --- pa.field("file_name", pa.string()), pa.field("caption", pa.string()), diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index 9e5315ae57..9aa0ba2f4b 100644 --- a/fastvideo/dataset/utils.py +++ b/fastvideo/dataset/utils.py @@ -134,38 +134,52 @@ def collate_rows_from_parquet_schema(rows, # Only add actual metadata fields, not the shape/dtype helper fields metadata_fields.append(field) + # Precompute one CFG dropout decision per row (not per tensor field) so + # that text_embedding (Qwen) and text_embedding_2 (ByT5) always drop + # together for the same sample. + cfg_drop_flags = [] + for row in rows: + drop = False + if cfg_rate > 0: + sample_idx = row.get("_sample_index") + if sample_idx is not None: + # Deterministic per-sample CFG dropout using sample index + # (resume-safe). + drop = random.Random(seed ^ int(sample_idx)).random() < cfg_rate + else: + drop = (rng.random() if rng else random.random()) < cfg_rate + cfg_drop_flags.append(drop) + # Process each tensor field for tensor_name in tensor_fields: tensor_list = [] - for row in rows: + for row_idx, row in enumerate(rows): # Get tensor data from row using the existing helper function pattern shape_key = f"{tensor_name}_shape" bytes_key = f"{tensor_name}_bytes" - if shape_key in row and bytes_key in row: + if ( + shape_key in row + and bytes_key in row + and row[shape_key] is not None + and row[bytes_key] is not None + ): shape = row[shape_key] bytes_data = row[bytes_key] if len(bytes_data) == 0: - tensor = torch.zeros(0, dtype=torch.bfloat16) + # Preserve the full shape (e.g. [0, D]) so downstream + # padding knows the embedding dimension even when there + # are zero tokens (e.g. ByT5 embedding with no glyph text). + # zeros, not empty: a truncated row can declare a non-empty + # shape with no payload, and uninitialized memory would be + # stacked silently instead of failing. + tensor = torch.zeros(tuple(shape), dtype=torch.float32) else: - # Deterministic per-sample CFG dropout - # using sample index (resume-safe). - drop = False - if (tensor_name == 'text_embedding' - and cfg_rate > 0): - sample_idx = row.get( - "_sample_index") - if sample_idx is not None: - drop = (random.Random( - seed ^ sample_idx - ).random() < cfg_rate) - else: - drop = ((rng.random() - if rng else - random.random()) - < cfg_rate) + drop = ( + tensor_name in ("text_embedding", "text_embedding_2") + and cfg_drop_flags[row_idx]) if drop: data = np.zeros(shape, dtype=np.float32) else: @@ -182,35 +196,90 @@ def collate_rows_from_parquet_schema(rows, tensor_list.append(tensor) else: # Handle missing tensor data - tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) + if tensor_name == "text_embedding_2": + tensor_list.append(None) + else: + tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) # Stack tensors with special handling for text embeddings - if tensor_name == 'text_embedding': + if tensor_name in ( + "text_embedding", + "text_embedding_2", + ): # Handle text embeddings with padding + + valid_tensors = [t for t in tensor_list if t is not None] + + if len(valid_tensors) == 0: + if tensor_name == "text_embedding_2": + # Legacy Wan/Cosmos data does not contain the secondary + # (ByT5) text embedding at all - skip this field entirely. + continue + raise ValueError("text_embedding is missing from all rows.") + + reference_tensor = valid_tensors[0] + if reference_tensor.ndim == 3: + if reference_tensor.shape[0] != 1: + raise ValueError( + f"Expected '{tensor_name}' batch dimension " + f"to be 1, got {tuple(reference_tensor.shape)}" + ) + reference_tensor = reference_tensor.squeeze(0) + if reference_tensor.ndim != 2: + raise ValueError( + f"Expected '{tensor_name}' shape [L, D], " + f"got {tuple(reference_tensor.shape)}" + ) + embedding_dim = reference_tensor.shape[-1] + embedding_dtype = reference_tensor.dtype + padded_tensors = [] attention_masks = [] - for tensor in tensor_list: - if tensor.numel() > 0: - padded_tensor, mask = pad(tensor, text_padding_length) - padded_tensors.append(padded_tensor) - attention_masks.append(mask) - else: - # Handle empty embeddings - assume default embedding dimension - padded_tensors.append( - torch.zeros(text_padding_length, - 768, - dtype=torch.bfloat16)) - attention_masks.append(torch.zeros(text_padding_length)) + if tensor is None: + if tensor_name == "text_embedding": + raise ValueError( + "text_embedding is missing from one or more rows.") + # No ByT5 embedding field for this row (legacy row mixed + # into a Hunyuan batch) - treat as zero tokens. + tensor = torch.empty((0, embedding_dim), dtype=embedding_dtype) + + if tensor.ndim == 3: + if tensor.shape[0] != 1: + raise ValueError( + f"Expected '{tensor_name}' batch dimension " + f"to be 1, got {tuple(tensor.shape)}" + ) + tensor = tensor.squeeze(0) + + if tensor.ndim != 2: + raise ValueError( + f"Expected '{tensor_name}' shape [L, D], " + f"got {tuple(tensor.shape)}" + ) + + if tensor.shape[-1] != embedding_dim: + raise ValueError( + f"Inconsistent hidden dimension for " + f"'{tensor_name}': expected {embedding_dim}, " + f"got {tensor.shape[-1]}" + ) + + padded_tensor, mask = pad(tensor, text_padding_length) + padded_tensors.append(padded_tensor) + attention_masks.append(mask) batch_data[tensor_name] = torch.stack(padded_tensors) - batch_data['text_attention_mask'] = torch.stack(attention_masks) + if tensor_name == "text_embedding": + batch_data["text_attention_mask"] = torch.stack(attention_masks) + else: + batch_data["text_attention_mask_2"] = torch.stack(attention_masks) else: # Stack all tensors to preserve batch consistency # Don't filter out None or empty tensors as this breaks batch sizing try: batch_data[tensor_name] = torch.stack(tensor_list) - except ValueError as e: + except (ValueError, RuntimeError) as e: shapes = [ t.shape if t is not None and hasattr(t, 'shape') else 'None/Invalid' diff --git a/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py b/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py index b1ce99d15c..9112689a76 100644 --- a/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py @@ -13,11 +13,11 @@ import cv2 import numpy as np -import pyarrow as pa import pyarrow.parquet as pq import torch from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.dataloader.parquet_io import records_to_table from fastvideo.utils import maybe_download_model # --- Config --- @@ -166,11 +166,7 @@ def main() -> None: torch.cuda.empty_cache() # Write parquet - table = pa.table( - {k: [r[k] for r in records] - for k in records[0]}, - schema=pyarrow_schema_t2v, - ) + table = records_to_table(records, 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}") diff --git a/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py b/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py index 24114b8fc7..a7ad1794f1 100644 --- a/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py @@ -14,7 +14,6 @@ import cv2 import numpy as np -import pyarrow as pa import pyarrow.parquet as pq import torch @@ -22,6 +21,7 @@ from fastvideo.configs.models.encoders.base import BaseEncoderOutput from fastvideo.configs.pipelines.cosmos import t5_large_postprocess_text from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.dataloader.parquet_io import records_to_table from fastvideo.utils import maybe_download_model # --- Config --- @@ -154,11 +154,7 @@ def main() -> None: 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, - ) + table = records_to_table(records, 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}") diff --git a/fastvideo/pipelines/preprocess/preprocess_hunyuan15_overfit.py b/fastvideo/pipelines/preprocess/preprocess_hunyuan15_overfit.py new file mode 100644 index 0000000000..92d1be05e3 --- /dev/null +++ b/fastvideo/pipelines/preprocess/preprocess_hunyuan15_overfit.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Preprocess HunyuanVideo 1.5 overfit data into parquet format.""" + +from __future__ import annotations + +import gc +import glob +import json +import os +from typing import Any + +import cv2 +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch +from safetensors.torch import load_file as safetensors_load_file +from transformers import AutoTokenizer, Qwen2_5_VLTextModel, T5EncoderModel + +from fastvideo.configs.models.vaes import Hunyuan15VAEConfig +from fastvideo.configs.pipelines.hunyuan15 import ( + Hunyuan15T2V480PConfig, + byt5_postprocess_text, + byt5_preprocess_text, + qwen_postprocess_text, + qwen_preprocess_text, +) +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.models.vaes.hunyuan15vae import AutoencoderKLHunyuanVideo15 +from fastvideo.utils import maybe_download_model + +NUM_FRAMES = 81 +MAX_HEIGHT = 480 +MAX_WIDTH = 832 +TRAIN_FPS = 16.0 + +# Overridable so a test can point at its own directories. Without this the +# only usable paths are the documented recipe's, and anything automated ends up +# writing over whatever clips and captions a user has prepared there. +DATA_DIR = os.environ.get("HY15_OVERFIT_DATA_DIR", "data/hunyuan15_overfit") +OUTPUT_DIR = os.environ.get("HY15_OVERFIT_OUTPUT_DIR", "data/hunyuan15_overfit_preprocessed") +MODEL_REPO = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" + +MAX_SAMPLES: int | None = 1 + + +def clear_cuda() -> None: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def tensor_to_record( + tensor: torch.Tensor, + prefix: str, +) -> dict[str, Any]: + tensor = tensor.detach().contiguous().float().cpu() + array = tensor.numpy() + return { + f"{prefix}_bytes": array.tobytes(), + f"{prefix}_shape": list(array.shape), + f"{prefix}_dtype": str(array.dtype), + } + + +def load_caption_data() -> list[dict[str, Any]]: + caption_path = os.path.join(DATA_DIR, "videos2caption.json") + with open(caption_path, encoding="utf-8") as file: + caption_data = json.load(file) + + if MAX_SAMPLES is not None: + caption_data = caption_data[:MAX_SAMPLES] + + if not caption_data: + raise ValueError(f"No caption records found in {caption_path}") + + return caption_data + + +def load_video( + path: str, + num_frames: int, + height: int, + width: int, +) -> torch.Tensor: + cap = cv2.VideoCapture(path) + frames: list[np.ndarray] = [] + + while len(frames) < num_frames: + ok, frame = cap.read() + if not ok: + break + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = cv2.resize( + frame, + (width, height), + interpolation=cv2.INTER_AREA, + ) + frames.append(frame) + + cap.release() + + if not frames: + raise ValueError(f"No frames could be decoded from {path}") + + while len(frames) < num_frames: + frames.append(frames[-1].copy()) + + video = np.stack(frames[:num_frames], axis=0) + video = torch.from_numpy(video).float() + video = video / 127.5 - 1.0 + return video.permute(3, 0, 1, 2).unsqueeze(0) + + +def apply_qwen_chat_template( + tokenizer, + prompt: str, + max_length: int, +) -> dict[str, torch.Tensor]: + messages = qwen_preprocess_text(prompt) + + encoded = tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + max_length=max_length, + truncation=True, + ) + + if isinstance(encoded, torch.Tensor): + return { + "input_ids": encoded, + "attention_mask": torch.ones_like(encoded), + } + + if "attention_mask" not in encoded: + encoded["attention_mask"] = torch.ones_like(encoded["input_ids"]) + + return { + "input_ids": encoded["input_ids"], + "attention_mask": encoded["attention_mask"], + } + + +@torch.inference_mode() +def encode_qwen_caption( + caption: str, + tokenizer, + encoder: Qwen2_5_VLTextModel, + device: torch.device, + max_length: int, +) -> torch.Tensor: + tokens = apply_qwen_chat_template(tokenizer, caption, max_length) + tokens = {key: value.to(device) for key, value in tokens.items()} + + outputs = encoder( + input_ids=tokens["input_ids"], + attention_mask=tokens["attention_mask"], + output_hidden_states=True, + return_dict=True, + ) + + embeddings, attention_mask = qwen_postprocess_text( + outputs, + tokens["attention_mask"], + ) + + embeddings = embeddings.squeeze(0) + attention_mask = attention_mask.squeeze(0) + valid_length = int(attention_mask.sum().item()) + + return embeddings[:valid_length].float().cpu() + + +@torch.inference_mode() +def encode_byt5_caption( + caption: str, + tokenizer, + encoder: T5EncoderModel, + device: torch.device, + max_length: int, +) -> torch.Tensor: + glyph_prompt = byt5_preprocess_text(caption) + hidden_size = int(encoder.config.d_model) + + if glyph_prompt is None: + return torch.empty((0, hidden_size), dtype=torch.float32) + + tokens = tokenizer( + glyph_prompt, + add_special_tokens=True, + padding=False, + truncation=True, + max_length=max_length, + return_tensors="pt", + ) + tokens = {key: value.to(device) for key, value in tokens.items()} + + outputs = encoder( + input_ids=tokens["input_ids"], + attention_mask=tokens.get("attention_mask"), + return_dict=True, + ) + + embeddings = byt5_postprocess_text(outputs).squeeze(0) + + attention_mask = tokens.get("attention_mask") + if attention_mask is not None: + valid_length = int(attention_mask.squeeze(0).sum().item()) + embeddings = embeddings[:valid_length] + + return embeddings.float().cpu() + + +def load_hunyuan15_vae( + model_path: str, + device: torch.device, +) -> AutoencoderKLHunyuanVideo15: + vae_path = os.path.join(model_path, "vae") + config_path = os.path.join(vae_path, "config.json") + + with open(config_path, encoding="utf-8") as file: + raw_config = json.load(file) + + raw_config.pop("_class_name", None) + raw_config.pop("_diffusers_version", None) + raw_config.pop("_name_or_path", None) + + vae_config = Hunyuan15VAEConfig() + vae_config.update_model_arch(raw_config) + vae_config.load_encoder = True + vae_config.load_decoder = False + + vae = AutoencoderKLHunyuanVideo15(vae_config) + + weight_files = sorted(glob.glob(os.path.join(vae_path, "*.safetensors"))) + if not weight_files: + raise FileNotFoundError(f"No VAE safetensors files found under {vae_path}") + + state_dict: dict[str, torch.Tensor] = {} + for weight_file in weight_files: + state_dict.update(safetensors_load_file(weight_file, device="cpu")) + + missing_keys, unexpected_keys = vae.load_state_dict( + state_dict, + strict=False, + ) + + encoder_missing = [key for key in missing_keys if not key.startswith("decoder.")] + if encoder_missing: + raise RuntimeError(f"Missing VAE encoder weights: {encoder_missing[:20]}") + + if unexpected_keys: + print("Ignored checkpoint keys not used by the encoder-only VAE: " + f"{len(unexpected_keys)}") + + del state_dict + clear_cuda() + + vae = vae.to( + device=device, + dtype=torch.float16, + ).eval() + vae.requires_grad_(False) + vae.use_tiling = False + return vae + + +@torch.inference_mode() +def encode_video_latent( + vae: AutoencoderKLHunyuanVideo15, + video: torch.Tensor, +) -> torch.Tensor: + encoded = vae.encode(video) + + # HunyuanVideo 1.5 VAE currently returns + # DiagonalGaussianDistribution directly. + latent_dist = (encoded.latent_dist if hasattr(encoded, "latent_dist") else encoded) + + if hasattr(latent_dist, "mode"): + latent = latent_dist.mode() + elif hasattr(latent_dist, "mean"): + latent = latent_dist.mean + else: + raise TypeError("Unsupported VAE encode output type: " + f"{type(encoded).__name__}") + + return latent.squeeze(0).float().cpu() + + +def main() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("This preprocessing script requires CUDA.") + + device = torch.device("cuda:0") + model_path = maybe_download_model(MODEL_REPO) + pipeline_config = Hunyuan15T2V480PConfig() + + os.makedirs(OUTPUT_DIR, exist_ok=True) + caption_data = load_caption_data() + captions = [item["cap"][0] for item in caption_data] + + records: list[dict[str, Any]] = [{ + "id": f"{index:06d}_{item['path']}", + "file_name": item["path"], + "caption": item["cap"][0], + "media_type": "video", + "width": MAX_WIDTH, + "height": MAX_HEIGHT, + "num_frames": NUM_FRAMES, + "duration_sec": NUM_FRAMES / TRAIN_FPS, + "fps": TRAIN_FPS, + } for index, item in enumerate(caption_data)] + + print("Loading Qwen2.5-VL text encoder...") + qwen_tokenizer = AutoTokenizer.from_pretrained( + os.path.join(model_path, "tokenizer"), + local_files_only=True, + ) + qwen_encoder = Qwen2_5_VLTextModel.from_pretrained( + os.path.join(model_path, "text_encoder"), + torch_dtype=torch.bfloat16, + local_files_only=True, + ).to(device).eval() + + qwen_max_length = int(pipeline_config.text_encoder_max_lengths[0]) + + for index, caption in enumerate(captions): + embedding = encode_qwen_caption( + caption, + qwen_tokenizer, + qwen_encoder, + device, + qwen_max_length, + ) + records[index].update(tensor_to_record(embedding, "text_embedding")) + print(f"[Qwen {index + 1}/{len(records)}] " + f"{tuple(embedding.shape)}") + + del qwen_encoder, qwen_tokenizer + clear_cuda() + + print("Loading ByT5/T5 text encoder...") + byt5_tokenizer = AutoTokenizer.from_pretrained( + os.path.join(model_path, "tokenizer_2"), + local_files_only=True, + ) + byt5_encoder = T5EncoderModel.from_pretrained( + os.path.join(model_path, "text_encoder_2"), + torch_dtype=torch.float32, + local_files_only=True, + ).to(device).eval() + + byt5_max_length = int(pipeline_config.text_encoder_max_lengths[1]) + + for index, caption in enumerate(captions): + embedding = encode_byt5_caption( + caption, + byt5_tokenizer, + byt5_encoder, + device, + byt5_max_length, + ) + records[index].update(tensor_to_record(embedding, "text_embedding_2")) + print(f"[ByT5 {index + 1}/{len(records)}] " + f"{tuple(embedding.shape)}") + + del byt5_encoder, byt5_tokenizer + clear_cuda() + + print("Loading HunyuanVideo 1.5 VAE encoder...") + vae = load_hunyuan15_vae(model_path, device) + + for index, item in enumerate(caption_data): + video_path = os.path.join( + DATA_DIR, + "videos", + item["path"], + ) + video = load_video( + video_path, + NUM_FRAMES, + MAX_HEIGHT, + MAX_WIDTH, + ).to( + device=device, + dtype=torch.float16, + ) + + latent = encode_video_latent(vae, video) + records[index].update(tensor_to_record(latent, "vae_latent")) + print(f"[VAE {index + 1}/{len(records)}] " + f"{tuple(latent.shape)}") + + del video, latent + clear_cuda() + + del vae + clear_cuda() + + columns = {field.name: [record.get(field.name) for record in records] for field in pyarrow_schema_t2v} + + table = pa.Table.from_pydict( + columns, + schema=pyarrow_schema_t2v, + ) + + output_path = os.path.join( + OUTPUT_DIR, + "data_00000.parquet", + ) + pq.write_table( + table, + output_path, + compression="zstd", + ) + + validation_path = os.path.join( + OUTPUT_DIR, + "validation_prompts.json", + ) + with open( + validation_path, + "w", + encoding="utf-8", + ) as file: + json.dump( + {"data": [{ + "caption": caption + } for caption in captions]}, + file, + ensure_ascii=False, + indent=2, + ) + + print(f"Wrote {len(records)} records to {output_path}") + print(f"Wrote validation prompts to {validation_path}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py b/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py index bb4c2ca157..ec737dbbbe 100644 --- a/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py @@ -15,7 +15,6 @@ import cv2 import numpy as np -import pyarrow as pa import pyarrow.parquet as pq import torch from safetensors.torch import load_file as safetensors_load_file @@ -28,6 +27,7 @@ llama_postprocess_text, ) from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.dataloader.parquet_io import records_to_table from fastvideo.models.vaes.hunyuanvae import AutoencoderKLHunyuanVideo from fastvideo.utils import maybe_download_model @@ -181,11 +181,7 @@ def main() -> None: del llama_enc, llama_tok, clip_enc, clip_tok, vae # Write parquet - table = pa.table( - {k: [r[k] for r in records] - for k in records[0]}, - schema=pyarrow_schema_t2v, - ) + table = records_to_table(records, 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}") diff --git a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py index 357a34fefc..46418585d6 100644 --- a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py @@ -39,7 +39,6 @@ import cv2 import numpy as np -import pyarrow as pa import pyarrow.parquet as pq import torch from transformers import AutoTokenizer @@ -52,6 +51,7 @@ kandinsky5_qwen_preprocess_text, ) from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.dataloader.parquet_io import records_to_table from fastvideo.distributed import maybe_init_distributed_environment_and_model_parallel from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context @@ -284,11 +284,7 @@ def main() -> None: del qwen_enc, qwen_tok, clip_enc, clip_tok, vae # Write parquet - table = pa.table( - {k: [r[k] for r in records] - for k in records[0]}, - schema=pyarrow_schema_t2v, - ) + table = records_to_table(records, 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}") diff --git a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py index 2f55261b9f..06ac9ce919 100644 --- a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py @@ -28,11 +28,11 @@ import av import numpy as np -import pyarrow as pa import pyarrow.parquet as pq import torch from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.dataloader.parquet_io import records_to_table from fastvideo.utils import maybe_download_model, verify_model_config_and_directory # --- Config --- @@ -230,11 +230,7 @@ def load_component(name: str) -> Any: 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, - ) + table = records_to_table(replicated, 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 " diff --git a/fastvideo/tests/dataset/test_hunyuan15_dual_text_collator.py b/fastvideo/tests/dataset/test_hunyuan15_dual_text_collator.py new file mode 100644 index 0000000000..763fbdc616 --- /dev/null +++ b/fastvideo/tests/dataset/test_hunyuan15_dual_text_collator.py @@ -0,0 +1,415 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pyarrow.parquet as pq +import pytest +import torch + +from fastvideo.dataset.dataloader.parquet_io import records_to_table +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v +from fastvideo.dataset.utils import collate_rows_from_parquet_schema + + +QWEN_DIM = 3584 +BYT5_DIM = 1472 +TEXT_PADDING_LENGTH = 8 + + +def _tensor_fields( + tensor: np.ndarray, + prefix: str, +) -> dict[str, Any]: + """Serialize a NumPy tensor using the FastVideo parquet convention.""" + + tensor = np.ascontiguousarray(tensor) + + return { + f"{prefix}_bytes": tensor.tobytes(), + f"{prefix}_shape": list(tensor.shape), + f"{prefix}_dtype": str(tensor.dtype), + } + + +def _make_row( + *, + sample_index: int, + qwen_tokens: int = 3, + byt5_tokens: int | None = 2, + qwen_value: float = 1.0, + byt5_value: float = 2.0, +) -> dict[str, Any]: + """Build one minimal T2V parquet row for collator tests. + + byt5_tokens=None represents a legacy row that has no text_embedding_2 + fields at all. + + byt5_tokens=0 represents a valid HunyuanVideo 1.5 sample with no + quoted glyph text. + """ + + qwen = np.full( + (qwen_tokens, QWEN_DIM), + qwen_value, + dtype=np.float32, + ) + + latent = np.full( + (32, 2, 4, 4), + 0.5, + dtype=np.float32, + ) + + row: dict[str, Any] = { + "id": f"sample-{sample_index}", + "file_name": f"sample-{sample_index}.mp4", + "caption": f"caption {sample_index}", + "media_type": "video", + "width": 64, + "height": 64, + "num_frames": 5, + "duration_sec": 1.0, + "fps": 5.0, + "_sample_index": sample_index, + } + + row.update(_tensor_fields(latent, "vae_latent")) + row.update(_tensor_fields(qwen, "text_embedding")) + + if byt5_tokens is not None: + byt5 = np.full( + (byt5_tokens, BYT5_DIM), + byt5_value, + dtype=np.float32, + ) + row.update(_tensor_fields(byt5, "text_embedding_2")) + + return row + + +def _collate( + rows: list[dict[str, Any]], + *, + cfg_rate: float = 0.0, + seed: int = 42, +) -> dict[str, Any]: + return collate_rows_from_parquet_schema( + rows=rows, + parquet_schema=pyarrow_schema_t2v, + text_padding_length=TEXT_PADDING_LENGTH, + cfg_rate=cfg_rate, + seed=seed, + ) + +def _write_and_read_parquet( + tmp_path, + rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + parquet_path = tmp_path / "data_00000.parquet" + + table = records_to_table( + rows, + pyarrow_schema_t2v, + ) + pq.write_table(table, parquet_path) + + loaded_table = pq.read_table(parquet_path) + return loaded_table.to_pylist() + +def test_dual_text_embeddings_are_padded_independently() -> None: + rows = [ + _make_row( + sample_index=0, + qwen_tokens=3, + byt5_tokens=2, + ), + _make_row( + sample_index=1, + qwen_tokens=5, + byt5_tokens=4, + ), + ] + + batch = _collate(rows) + + assert batch["text_embedding"].shape == ( + 2, + TEXT_PADDING_LENGTH, + QWEN_DIM, + ) + assert batch["text_embedding_2"].shape == ( + 2, + TEXT_PADDING_LENGTH, + BYT5_DIM, + ) + + assert batch["text_attention_mask"].shape == ( + 2, + TEXT_PADDING_LENGTH, + ) + assert batch["text_attention_mask_2"].shape == ( + 2, + TEXT_PADDING_LENGTH, + ) + + assert batch["text_attention_mask"][0].sum().item() == 3 + assert batch["text_attention_mask"][1].sum().item() == 5 + + assert batch["text_attention_mask_2"][0].sum().item() == 2 + assert batch["text_attention_mask_2"][1].sum().item() == 4 + + assert torch.count_nonzero( + batch["text_embedding"][0, 3:] + ).item() == 0 + + assert torch.count_nonzero( + batch["text_embedding_2"][0, 2:] + ).item() == 0 + + +def test_empty_byt5_embedding_preserves_hidden_dimension() -> None: + """A [0, 1472] embedding must become a padded [L, 1472] tensor.""" + + rows = [ + _make_row( + sample_index=0, + qwen_tokens=3, + byt5_tokens=0, + ), + ] + + batch = _collate(rows) + + assert batch["text_embedding_2"].shape == ( + 1, + TEXT_PADDING_LENGTH, + BYT5_DIM, + ) + + assert batch["text_attention_mask_2"].shape == ( + 1, + TEXT_PADDING_LENGTH, + ) + + assert batch["text_attention_mask_2"].sum().item() == 0 + assert torch.count_nonzero( + batch["text_embedding_2"] + ).item() == 0 + + +def test_mixed_empty_and_nonempty_byt5_embeddings() -> None: + """A batch may contain samples with and without glyph text.""" + + rows = [ + _make_row( + sample_index=0, + byt5_tokens=0, + ), + _make_row( + sample_index=1, + byt5_tokens=4, + ), + ] + + batch = _collate(rows) + + assert batch["text_embedding_2"].shape == ( + 2, + TEXT_PADDING_LENGTH, + BYT5_DIM, + ) + + assert batch["text_attention_mask_2"][0].sum().item() == 0 + assert batch["text_attention_mask_2"][1].sum().item() == 4 + + assert torch.count_nonzero( + batch["text_embedding_2"][0] + ).item() == 0 + + assert torch.count_nonzero( + batch["text_embedding_2"][1, :4] + ).item() > 0 + + +def test_legacy_rows_without_byt5_skip_secondary_batch_fields() -> None: + """A batch containing only old-format rows remains supported.""" + + rows = [ + _make_row( + sample_index=0, + byt5_tokens=None, + ), + _make_row( + sample_index=1, + byt5_tokens=None, + ), + ] + + batch = _collate(rows) + + assert "text_embedding" in batch + assert "text_attention_mask" in batch + + assert "text_embedding_2" not in batch + assert "text_attention_mask_2" not in batch + + +def test_legacy_row_can_mix_with_hunyuan_row() -> None: + """A missing ByT5 field is treated as zero tokens in a Hunyuan batch.""" + + rows = [ + _make_row( + sample_index=0, + byt5_tokens=None, + ), + _make_row( + sample_index=1, + byt5_tokens=3, + ), + ] + + batch = _collate(rows) + + assert batch["text_embedding_2"].shape == ( + 2, + TEXT_PADDING_LENGTH, + BYT5_DIM, + ) + + assert batch["text_attention_mask_2"][0].sum().item() == 0 + assert batch["text_attention_mask_2"][1].sum().item() == 3 + + +def test_cfg_rate_zero_preserves_both_text_conditions() -> None: + rows = [ + _make_row( + sample_index=0, + qwen_tokens=3, + byt5_tokens=2, + ), + ] + + batch = _collate( + rows, + cfg_rate=0.0, + ) + + assert torch.count_nonzero( + batch["text_embedding"][0, :3] + ).item() > 0 + + assert torch.count_nonzero( + batch["text_embedding_2"][0, :2] + ).item() > 0 + + +def test_cfg_rate_one_drops_both_text_conditions() -> None: + """Qwen and ByT5 must use the same CFG dropout decision.""" + + rows = [ + _make_row( + sample_index=0, + qwen_tokens=3, + byt5_tokens=2, + ), + ] + + batch = _collate( + rows, + cfg_rate=1.0, + ) + + assert torch.count_nonzero( + batch["text_embedding"] + ).item() == 0 + + assert torch.count_nonzero( + batch["text_embedding_2"] + ).item() == 0 + +def test_dual_text_embeddings_parquet_round_trip(tmp_path) -> None: + rows = [ + _make_row( + sample_index=0, + qwen_tokens=3, + byt5_tokens=2, + qwen_value=1.0, + byt5_value=2.0, + ), + _make_row( + sample_index=1, + qwen_tokens=5, + byt5_tokens=4, + qwen_value=3.0, + byt5_value=4.0, + ), + ] + + loaded_rows = _write_and_read_parquet(tmp_path, rows) + batch = _collate(loaded_rows) + + assert batch["text_embedding"].shape == ( + 2, + TEXT_PADDING_LENGTH, + QWEN_DIM, + ) + assert batch["text_embedding_2"].shape == ( + 2, + TEXT_PADDING_LENGTH, + BYT5_DIM, + ) + + assert batch["text_attention_mask"][0].sum().item() == 3 + assert batch["text_attention_mask"][1].sum().item() == 5 + assert batch["text_attention_mask_2"][0].sum().item() == 2 + assert batch["text_attention_mask_2"][1].sum().item() == 4 + + torch.testing.assert_close( + batch["text_embedding"][0, :3], + torch.ones( + (3, QWEN_DIM), + dtype=torch.float32, + ), + ) + + torch.testing.assert_close( + batch["text_embedding_2"][0, :2], + torch.full( + (2, BYT5_DIM), + 2.0, + dtype=torch.float32, + ), + ) + +def test_empty_byt5_embedding_parquet_round_trip(tmp_path) -> None: + rows = [ + _make_row( + sample_index=0, + qwen_tokens=3, + byt5_tokens=0, + ), + _make_row( + sample_index=1, + qwen_tokens=4, + byt5_tokens=2, + ), + ] + + loaded_rows = _write_and_read_parquet(tmp_path, rows) + batch = _collate(loaded_rows) + + assert batch["text_embedding_2"].shape == ( + 2, + TEXT_PADDING_LENGTH, + BYT5_DIM, + ) + + assert batch["text_attention_mask_2"][0].sum().item() == 0 + assert batch["text_attention_mask_2"][1].sum().item() == 2 + + assert torch.count_nonzero( + batch["text_embedding_2"][0] + ).item() == 0 \ No newline at end of file diff --git a/fastvideo/tests/dataset/test_hunyuan15_parquet_schema.py b/fastvideo/tests/dataset/test_hunyuan15_parquet_schema.py new file mode 100644 index 0000000000..2805eb92e5 --- /dev/null +++ b/fastvideo/tests/dataset/test_hunyuan15_parquet_schema.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pyarrow as pa +import pytest + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v + + +def test_hunyuan15_secondary_text_embedding_fields_exist() -> None: + """The T2V schema must expose the secondary ByT5 embedding fields.""" + + expected_fields = { + "text_embedding_2_bytes", + "text_embedding_2_shape", + "text_embedding_2_dtype", + } + + assert expected_fields.issubset(set(pyarrow_schema_t2v.names)) + + +def test_hunyuan15_secondary_text_embedding_field_types() -> None: + """The ByT5 fields must use the expected Arrow storage types.""" + + bytes_field = pyarrow_schema_t2v.field("text_embedding_2_bytes") + shape_field = pyarrow_schema_t2v.field("text_embedding_2_shape") + dtype_field = pyarrow_schema_t2v.field("text_embedding_2_dtype") + + assert bytes_field.type == pa.binary() + assert shape_field.type == pa.list_(pa.int64()) + assert dtype_field.type == pa.string() + + +def test_writer_records_without_byt5_still_build_a_table() -> None: + """A writer that predates the secondary encoder must still be able to write. + + ``nullable=True`` is pyarrow's default and asserting it proves nothing -- + it holds for every field in every schema. What matters is that + ``pa.table(mapping, schema=...)`` demands a key for *every* schema field, + so a record built before these columns existed raises KeyError, while + ``records_to_table`` (from_pylist) fills them with null. + """ + import pyarrow as pa + + from fastvideo.dataset.dataloader.parquet_io import records_to_table + + secondary = { + "text_embedding_2_bytes", + "text_embedding_2_shape", + "text_embedding_2_dtype", + } + record: dict = {} + for field in pyarrow_schema_t2v: + if field.name in secondary: + continue + if pa.types.is_binary(field.type): + record[field.name] = b"" + elif pa.types.is_string(field.type): + record[field.name] = "x" + elif pa.types.is_list(field.type): + record[field.name] = [1] + elif pa.types.is_integer(field.type): + record[field.name] = 1 + else: + record[field.name] = 1.0 + + with pytest.raises(KeyError): + pa.table({k: [record[k]] for k in record}, schema=pyarrow_schema_t2v) + + table = records_to_table([record], pyarrow_schema_t2v) + assert table.schema.equals(pyarrow_schema_t2v) + assert table.column("text_embedding_2_bytes").to_pylist() == [None] \ No newline at end of file diff --git a/fastvideo/tests/dataset/test_preprocess_hunyuan15.py b/fastvideo/tests/dataset/test_preprocess_hunyuan15.py new file mode 100644 index 0000000000..1e6fdf5c1f --- /dev/null +++ b/fastvideo/tests/dataset/test_preprocess_hunyuan15.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +import fastvideo.pipelines.preprocess.preprocess_hunyuan15_overfit as preprocess + + +QWEN_DIM = 3584 +BYT5_DIM = 1472 + + +def test_tensor_to_record_preserves_values_and_shape() -> None: + tensor = torch.arange( + 3 * 4, + dtype=torch.float32, + ).reshape(3, 4) + + record = preprocess.tensor_to_record( + tensor, + "text_embedding", + ) + + assert record["text_embedding_shape"] == [3, 4] + assert record["text_embedding_dtype"] == "float32" + + decoded = np.frombuffer( + record["text_embedding_bytes"], + dtype=np.float32, + ).reshape(record["text_embedding_shape"]) + + np.testing.assert_array_equal( + decoded, + tensor.numpy(), + ) + + +def test_tensor_to_record_converts_input_to_float32() -> None: + tensor = torch.ones( + (2, 4), + dtype=torch.bfloat16, + ) + + record = preprocess.tensor_to_record( + tensor, + "text_embedding", + ) + + assert record["text_embedding_dtype"] == "float32" + + decoded = np.frombuffer( + record["text_embedding_bytes"], + dtype=np.float32, + ) + + assert decoded.dtype == np.float32 + assert decoded.size == 8 + + +def test_empty_byt5_tensor_serializes_as_empty_bytes() -> None: + tensor = torch.empty( + (0, BYT5_DIM), + dtype=torch.float32, + ) + + record = preprocess.tensor_to_record( + tensor, + "text_embedding_2", + ) + + assert record["text_embedding_2_shape"] == [0, BYT5_DIM] + assert record["text_embedding_2_dtype"] == "float32" + assert record["text_embedding_2_bytes"] == b"" + + decoded = np.frombuffer( + record["text_embedding_2_bytes"], + dtype=np.float32, + ).reshape(record["text_embedding_2_shape"]) + + assert decoded.shape == (0, BYT5_DIM) + + +def test_encode_byt5_returns_empty_embedding_without_glyph_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Captions without quoted text must not invoke the ByT5 model.""" + + monkeypatch.setattr( + preprocess, + "byt5_preprocess_text", + lambda caption: None, + ) + + tokenizer = Mock() + encoder = Mock() + encoder.config = SimpleNamespace( + d_model=BYT5_DIM, + ) + + result = preprocess.encode_byt5_caption( + caption="A dog runs across a grassy field.", + tokenizer=tokenizer, + encoder=encoder, + device=torch.device("cpu"), + max_length=256, + ) + + assert result.shape == (0, BYT5_DIM) + assert result.dtype == torch.float32 + assert result.device.type == "cpu" + + tokenizer.assert_not_called() + encoder.assert_not_called() + + +def test_encode_byt5_encodes_extracted_glyph_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only the extracted glyph prompt should be sent to ByT5.""" + + glyph_text = "OPEN" + + monkeypatch.setattr( + preprocess, + "byt5_preprocess_text", + lambda caption: glyph_text, + ) + + monkeypatch.setattr( + preprocess, + "byt5_postprocess_text", + lambda outputs: outputs.last_hidden_state, + ) + + tokenizer = Mock( + return_value={ + "input_ids": torch.tensor( + [[10, 11, 12]], + dtype=torch.long, + ), + "attention_mask": torch.tensor( + [[1, 1, 1]], + dtype=torch.long, + ), + } + ) + + encoder = Mock() + encoder.config = SimpleNamespace( + d_model=BYT5_DIM, + ) + encoder.return_value = SimpleNamespace( + last_hidden_state=torch.full( + (1, 3, BYT5_DIM), + 2.0, + dtype=torch.float32, + ) + ) + + result = preprocess.encode_byt5_caption( + caption='A storefront sign reads "OPEN".', + tokenizer=tokenizer, + encoder=encoder, + device=torch.device("cpu"), + max_length=256, + ) + + tokenizer.assert_called_once_with( + glyph_text, + add_special_tokens=True, + padding=False, + truncation=True, + max_length=256, + return_tensors="pt", + ) + + encoder.assert_called_once() + + assert result.shape == (3, BYT5_DIM) + assert result.dtype == torch.float32 + assert torch.all(result == 2.0) + + +def test_encode_byt5_trims_padding_tokens( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + preprocess, + "byt5_preprocess_text", + lambda caption: "OPEN", + ) + + monkeypatch.setattr( + preprocess, + "byt5_postprocess_text", + lambda outputs: outputs.last_hidden_state, + ) + + tokenizer = Mock( + return_value={ + "input_ids": torch.tensor( + [[10, 11, 12, 0, 0]], + dtype=torch.long, + ), + "attention_mask": torch.tensor( + [[1, 1, 1, 0, 0]], + dtype=torch.long, + ), + } + ) + + encoder = Mock() + encoder.config = SimpleNamespace( + d_model=BYT5_DIM, + ) + encoder.return_value = SimpleNamespace( + last_hidden_state=torch.randn( + 1, + 5, + BYT5_DIM, + ) + ) + + result = preprocess.encode_byt5_caption( + caption='The sign reads "OPEN".', + tokenizer=tokenizer, + encoder=encoder, + device=torch.device("cpu"), + max_length=256, + ) + + assert result.shape == (3, BYT5_DIM) + + +class _FakeDistributionWithMode: + + def __init__(self, latent: torch.Tensor) -> None: + self._latent = latent + + def mode(self) -> torch.Tensor: + return self._latent + + +class _FakeDistributionWithMean: + + def __init__(self, latent: torch.Tensor) -> None: + self.mean = latent + + +class _FakeVAE: + + def __init__(self, encoded_output: object) -> None: + self.encoded_output = encoded_output + + def encode(self, video: torch.Tensor) -> object: + del video + return self.encoded_output + + +def test_encode_video_latent_accepts_direct_distribution() -> None: + latent = torch.randn( + 1, + 32, + 2, + 4, + 4, + dtype=torch.float16, + ) + + vae = _FakeVAE( + _FakeDistributionWithMode(latent) + ) + + result = preprocess.encode_video_latent( + vae=vae, + video=torch.empty(1), + ) + + assert result.shape == (32, 2, 4, 4) + assert result.dtype == torch.float32 + assert result.device.type == "cpu" + + torch.testing.assert_close( + result, + latent.squeeze(0).float(), + ) + + +def test_encode_video_latent_accepts_latent_dist_wrapper() -> None: + latent = torch.randn( + 1, + 32, + 2, + 4, + 4, + ) + + encoded = SimpleNamespace( + latent_dist=_FakeDistributionWithMode(latent), + ) + + vae = _FakeVAE(encoded) + + result = preprocess.encode_video_latent( + vae=vae, + video=torch.empty(1), + ) + + torch.testing.assert_close( + result, + latent.squeeze(0).float(), + ) + + +def test_encode_video_latent_falls_back_to_mean() -> None: + latent = torch.randn( + 1, + 32, + 2, + 4, + 4, + ) + + vae = _FakeVAE( + _FakeDistributionWithMean(latent) + ) + + result = preprocess.encode_video_latent( + vae=vae, + video=torch.empty(1), + ) + + torch.testing.assert_close( + result, + latent.squeeze(0).float(), + ) + + +def test_encode_video_latent_rejects_unknown_output() -> None: + vae = _FakeVAE(object()) + + with pytest.raises( + TypeError, + match="Unsupported VAE encode output type", + ): + preprocess.encode_video_latent( + vae=vae, + video=torch.empty(1), + ) + + +def test_encode_video_latent_does_not_apply_scaling_factor() -> None: + """Scaling is applied by the training side, not preprocessing.""" + + latent = torch.full( + (1, 32, 2, 4, 4), + 2.0, + ) + + vae = _FakeVAE( + _FakeDistributionWithMode(latent) + ) + + result = preprocess.encode_video_latent( + vae=vae, + video=torch.empty(1), + ) + + assert torch.all(result == 2.0) \ No newline at end of file