diff --git a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh index cbddfb60f..1281089e5 100644 --- a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh +++ b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh @@ -69,7 +69,7 @@ miscellaneous_args=( --inference_mode False --allow_tf32 --checkpoints_total_limit 3 - --cfg 0.0 + --training_cfg_rate 0.1 --multi_phased_distill_schedule "4000-1" --not_apply_cfg_solver --dit_precision "fp32" diff --git a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm index c20981ca4..63da7ec91 100644 --- a/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm +++ b/examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm @@ -106,7 +106,7 @@ miscellaneous_args=( --inference_mode False --allow_tf32 --checkpoints_total_limit 3 - --cfg 0.0 + --training_cfg_rate 0.1 --multi_phased_distill_schedule "4000-1" --not_apply_cfg_solver --dit_precision "fp32" diff --git a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh index aea721c44..cd443fbc4 100644 --- a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh +++ b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh @@ -69,7 +69,7 @@ miscellaneous_args=( --inference_mode False --allow_tf32 --checkpoints_total_limit 3 - --cfg 0.0 + --training_cfg_rate 0.1 --multi_phased_distill_schedule "4000-1" --not_apply_cfg_solver --dit_precision "fp32" diff --git a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm index 563840fe1..ad31610b0 100644 --- a/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm +++ b/examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm @@ -103,7 +103,7 @@ miscellaneous_args=( --inference_mode False --allow_tf32 --checkpoints_total_limit 3 - --cfg 0.0 + --training_cfg_rate 0.1 --multi_phased_distill_schedule "4000-1" --not_apply_cfg_solver --dit_precision "fp32" diff --git a/fastvideo/v1/dataset/parquet_dataset_iterable_style.py b/fastvideo/v1/dataset/parquet_dataset_iterable_style.py index 20082fc22..18bc3d8ad 100644 --- a/fastvideo/v1/dataset/parquet_dataset_iterable_style.py +++ b/fastvideo/v1/dataset/parquet_dataset_iterable_style.py @@ -4,6 +4,7 @@ import random from typing import Dict, List, Tuple import numpy as np +import pyarrow as pa import pyarrow.parquet as pq import torch import tqdm @@ -70,10 +71,12 @@ class LatentsParquetIterStyleDataset(IterableDataset): drop_last: bool = True, text_padding_length: int = 512, seed: int = 42, - read_batch_size: int = 32): + read_batch_size: int = 32, + parquet_schema: pa.Schema = None): super().__init__() self.path = str(path) self.batch_size = batch_size + self.parquet_schema = parquet_schema self.cfg_rate = cfg_rate self.text_padding_length = text_padding_length self.seed = seed diff --git a/fastvideo/v1/dataset/parquet_dataset_map_style.py b/fastvideo/v1/dataset/parquet_dataset_map_style.py index 6961bd19c..ae2b09921 100644 --- a/fastvideo/v1/dataset/parquet_dataset_map_style.py +++ b/fastvideo/v1/dataset/parquet_dataset_map_style.py @@ -201,10 +201,6 @@ class LatentsParquetMapStyleDataset(Dataset): self.path = path self.cfg_rate = cfg_rate self.parquet_schema = parquet_schema - if cfg_rate > 0.0: - raise ValueError( - "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0" - ) logger.info("Initializing LatentsParquetMapStyleDataset with path: %s", path) self.parquet_files, self.lengths = get_parquet_files_and_length(path) diff --git a/fastvideo/v1/dataset/utils.py b/fastvideo/v1/dataset/utils.py index b5555a5a3..6230cc7ba 100644 --- a/fastvideo/v1/dataset/utils.py +++ b/fastvideo/v1/dataset/utils.py @@ -1,5 +1,5 @@ import random -from typing import Any, Dict, List +from typing import Any, Dict, List, cast import numpy as np import torch @@ -108,7 +108,7 @@ def collate_rows_from_parquet_schema(rows, Dict containing batched tensors and metadata """ if not rows: - return {} + return cast(Dict[str, Any], {}) # Initialize containers for different data types batch_data: Dict[str, Any] = {} diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py index 831fd4b96..260f9c86d 100644 --- a/fastvideo/v1/pipelines/pipeline_batch_info.py +++ b/fastvideo/v1/pipelines/pipeline_batch_info.py @@ -39,6 +39,7 @@ class ForwardBatch: image_path: Optional[str] = None image_embeds: List[torch.Tensor] = field(default_factory=list) pil_image: Optional[PIL.Image.Image] = None + preprocessed_image: Optional[torch.Tensor] = None # Text inputs prompt: Optional[Union[str, List[str]]] = None @@ -150,6 +151,10 @@ class TrainingBatch: latents: Optional[torch.Tensor] = None encoder_hidden_states: Optional[torch.Tensor] = None encoder_attention_mask: Optional[torch.Tensor] = None + # i2v + preprocessed_image: Optional[torch.Tensor] = None + image_embeds: Optional[torch.Tensor] = None + image_latents: Optional[torch.Tensor] = None infos: Optional[List[Dict[str, Any]]] = None # Transformer inputs diff --git a/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py index 4ae4e133e..5866c46ee 100644 --- a/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_base.py @@ -2,11 +2,13 @@ import gc import multiprocessing import os +from collections import defaultdict from concurrent.futures import ProcessPoolExecutor from itertools import chain from typing import Any, Dict, List, Optional import numpy as np +import PIL.Image import pyarrow as pa import pyarrow.parquet as pq import torch @@ -15,12 +17,13 @@ from tqdm import tqdm from fastvideo.v1.configs.sample import SamplingParam from fastvideo.v1.dataset import ValidationDataset, getdataset +from fastvideo.v1.dataset.preprocessing_datasets import PreprocessBatch from fastvideo.v1.distributed import get_torch_device from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.logger import init_logger from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.v1.pipelines.stages import EncodingStage, TextEncodingStage +from fastvideo.v1.pipelines.stages import TextEncodingStage logger = init_logger(__name__) @@ -36,9 +39,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): tokenizers=[self.get_module("tokenizer")], )) - self.add_stage(stage_name="image_encoding_stage", - stage=EncodingStage(vae=self.get_module("vae"), )) - @torch.no_grad() def forward( self, @@ -61,39 +61,206 @@ class BasePreprocessPipeline(ComposedPipelineBase): """Get the schema fields for the pipeline type. Override in subclasses.""" raise NotImplementedError + def create_record_for_schema(self, + preprocess_batch: PreprocessBatch, + schema: pa.Schema, + strict: bool = False) -> Dict[str, Any]: + """Create a record for the Parquet dataset using a generic schema-based approach. + + Args: + preprocess_batch: The batch containing the data to extract + schema: PyArrow schema defining the expected fields + strict: If True, raises an exception when required fields are missing or unfilled + + Returns: + Dictionary record matching the schema + + Raises: + ValueError: If strict=True and required fields are missing or unfilled + """ + record = {} + unfilled_fields = [] + + for field in schema.names: + field_filled = False + + if field.endswith('_bytes'): + # Handle binary tensor data - convert numpy array or tensor to bytes + tensor_name = field.replace('_bytes', '') + tensor_data = getattr(preprocess_batch, tensor_name, None) + if tensor_data is not None: + try: + if hasattr(tensor_data, 'numpy'): # torch tensor + record[field] = tensor_data.cpu().numpy().tobytes() + field_filled = True + elif hasattr(tensor_data, 'tobytes'): # numpy array + record[field] = tensor_data.tobytes() + field_filled = True + else: + raise ValueError( + f"Unsupported tensor type for field {field}: {type(tensor_data)}" + ) + except Exception as e: + if strict: + raise ValueError( + f"Failed to convert tensor {tensor_name} to bytes: {e}" + ) from e + record[field] = b'' # Empty bytes for missing data + else: + record[field] = b'' # Empty bytes for missing data + + elif field.endswith('_shape'): + # Handle tensor shape info + tensor_name = field.replace('_shape', '') + tensor_data = getattr(preprocess_batch, tensor_name, None) + if tensor_data is not None and hasattr(tensor_data, 'shape'): + record[field] = list(tensor_data.shape) + field_filled = True + else: + record[field] = [] + + elif field.endswith('_dtype'): + # Handle tensor dtype info + tensor_name = field.replace('_dtype', '') + tensor_data = getattr(preprocess_batch, tensor_name, None) + if tensor_data is not None and hasattr(tensor_data, 'dtype'): + record[field] = str(tensor_data.dtype) + field_filled = True + else: + record[field] = 'unknown' + + elif field in ['width', 'height', 'num_frames']: + # Handle integer metadata fields + value = getattr(preprocess_batch, field, None) + if value is not None: + try: + record[field] = int(value) + field_filled = True + except (ValueError, TypeError) as e: + if strict: + raise ValueError( + f"Failed to convert field {field} to int: {e}" + ) from e + record[field] = 0 + else: + record[field] = 0 + + elif field in ['duration_sec', 'fps']: + # Handle float metadata fields + # Map schema field names to batch attribute names + attr_name = 'duration' if field == 'duration_sec' else field + value = getattr(preprocess_batch, attr_name, None) + if value is not None: + try: + record[field] = float(value) + field_filled = True + except (ValueError, TypeError) as e: + if strict: + raise ValueError( + f"Failed to convert field {field} to float: {e}" + ) from e + record[field] = 0.0 + else: + record[field] = 0.0 + + else: + # Handle string fields (id, file_name, caption, media_type, etc.) + # Map common schema field names to batch attribute names + attr_name = field + if field == 'caption': + attr_name = 'text' + elif field == 'file_name': + attr_name = 'path' + elif field == 'id': + # Generate ID from path if available + path_value = getattr(preprocess_batch, 'path', None) + if path_value: + import os + record[field] = os.path.basename(path_value).split( + '.')[0] + field_filled = True + else: + record[field] = "" + continue + elif field == 'media_type': + # Determine media type from path + path_value = getattr(preprocess_batch, 'path', None) + if path_value: + record[field] = 'video' if path_value.endswith( + '.mp4') else 'image' + field_filled = True + else: + record[field] = "" + continue + + value = getattr(preprocess_batch, attr_name, None) + if value is not None: + record[field] = str(value) + field_filled = True + else: + record[field] = "" + + # Track unfilled fields + if not field_filled: + unfilled_fields.append(field) + + # Handle strict mode + if strict and unfilled_fields: + raise ValueError( + f"Required fields were not filled: {unfilled_fields}") + + # Log unfilled fields as warning if not in strict mode + if unfilled_fields: + logger.warning( + "Some fields were not filled and got default values: %s", + unfilled_fields) + + return record + def create_record( self, video_name: str, vae_latent: np.ndarray, text_embedding: np.ndarray, - text_attention_mask: np.ndarray, - valid_data: Optional[Dict[str, Any]], + valid_data: Dict[str, Any], idx: int, extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Create a record for the Parquet dataset.""" record = { - "id": video_name, - "vae_latent_bytes": vae_latent.tobytes(), - "vae_latent_shape": list(vae_latent.shape), - "vae_latent_dtype": str(vae_latent.dtype), - "text_embedding_bytes": text_embedding.tobytes(), - "text_embedding_shape": list(text_embedding.shape), - "text_embedding_dtype": str(text_embedding.dtype), - "text_attention_mask_bytes": text_attention_mask.tobytes(), - "text_attention_mask_shape": list(text_attention_mask.shape), - "text_attention_mask_dtype": str(text_attention_mask.dtype), - "file_name": video_name, - "caption": valid_data["text"][idx] if valid_data else "", - "media_type": "video", + "id": + video_name, + "vae_latent_bytes": + vae_latent.tobytes(), + "vae_latent_shape": + list(vae_latent.shape), + "vae_latent_dtype": + str(vae_latent.dtype), + "text_embedding_bytes": + text_embedding.tobytes(), + "text_embedding_shape": + list(text_embedding.shape), + "text_embedding_dtype": + str(text_embedding.dtype), + "file_name": + video_name, + "caption": + valid_data["text"][idx] if len(valid_data["text"]) > 0 else "", + "media_type": + "video", "width": - valid_data["pixel_values"][idx].shape[-2] if valid_data else 0, + valid_data["pixel_values"][idx].shape[-2] + if len(valid_data["pixel_values"]) > 0 else 0, "height": - valid_data["pixel_values"][idx].shape[-1] if valid_data else 0, + valid_data["pixel_values"][idx].shape[-1] + if len(valid_data["pixel_values"]) > 0 else 0, "num_frames": vae_latent.shape[1] if len(vae_latent.shape) > 1 else 0, "duration_sec": - float(valid_data["duration"][idx]) if valid_data else 0.0, - "fps": float(valid_data["fps"][idx]) if valid_data else 0.0, + float(valid_data["duration"][idx]) + if len(valid_data["duration"]) > 0 else 0.0, + "fps": + float(valid_data["fps"][idx]) + if len(valid_data["fps"]) > 0 else 0.0, } if extra_features: record.update(extra_features) @@ -213,8 +380,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): # Convert tensors to numpy arrays vae_latent = latent.cpu().numpy() text_embedding = prompt_embeds[idx].cpu().numpy() - text_attention_mask = prompt_attention_mask[idx].cpu().numpy( - ).astype(np.uint8) # Get extra features for this sample if needed sample_extra_features = {} @@ -231,7 +396,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): video_name=video_name, vae_latent=vae_latent, text_embedding=text_embedding, - text_attention_mask=text_attention_mask, valid_data=valid_data, idx=idx, extra_features=sample_extra_features) @@ -318,7 +482,7 @@ class BasePreprocessPipeline(ComposedPipelineBase): for idx, sample in pbar: with torch.inference_mode(): prompt = sample["caption"] - # is_negative_prompt = idx == 0 + is_negative_prompt = idx == 0 # Text Encoder batch = ForwardBatch( @@ -346,15 +510,43 @@ class BasePreprocessPipeline(ComposedPipelineBase): "Shape after removing padding - Embeddings: %s, Mask: %s", text_embedding.shape, text_attention_mask.shape) + extra_features = {} + if not is_negative_prompt: + height = sample["height"] + width = sample["width"] + if "image_path" in sample and "video_path" in sample: + raise ValueError( + "Only one of image_path or video_path should be provided" + ) + + if "image" in sample: + extra_features = self.preprocess_image( + sample["image"], height, width, fastvideo_args) + + if "video" in sample: + extra_features = self.preprocess_video( + sample["video"], height, width, fastvideo_args) + + # Get extra features for this sample if needed + sample_extra_features = {} + if extra_features: + for key, value in extra_features.items(): + if isinstance(value, torch.Tensor): + sample_extra_features[key] = value.cpu().numpy() + else: + sample_extra_features[key] = value + + valid_data = defaultdict(list) + valid_data["text"] = [prompt] + # Create record for Parquet dataset record = self.create_record(video_name=file_name, vae_latent=np.array([], dtype=np.float32), text_embedding=text_embedding, - text_attention_mask=text_attention_mask, - valid_data=None, + valid_data=valid_data, idx=0, - extra_features=None) + extra_features=sample_extra_features) batch_data.append(record) logger.info("Saved validation sample: %s", file_name) @@ -428,6 +620,15 @@ class BasePreprocessPipeline(ComposedPipelineBase): del table gc.collect() # Force garbage collection + def preprocess_image(self, image: PIL.Image.Image, height: int, width: int, + fastvideo_args: FastVideoArgs) -> Dict[str, Any]: + return {} + + def preprocess_video(self, video: list[PIL.Image.Image], height: int, + width: int, + fastvideo_args: FastVideoArgs) -> Dict[str, Any]: + return {} + def _flush_tables(self, num_processed_samples: int, args, combined_parquet_dir: str): """Flush collected tables to disk.""" diff --git a/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py b/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py index 1fb35a9d8..58ba09bef 100644 --- a/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py +++ b/fastvideo/v1/pipelines/preprocess/preprocess_pipeline_i2v.py @@ -8,6 +8,7 @@ using the modular pipeline architecture. from typing import Any, Dict, List, Optional import numpy as np +import PIL import torch from PIL import Image @@ -15,8 +16,13 @@ from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema_i2v from fastvideo.v1.distributed import get_torch_device from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.forward_context import set_forward_context +from fastvideo.v1.models.vision_utils import (get_default_height_width, + normalize, numpy_to_pt, + pil_to_numpy, resize) +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.preprocess.preprocess_pipeline_base import ( BasePreprocessPipeline) +from fastvideo.v1.pipelines.stages import ImageEncodingStage, TextEncodingStage class PreprocessPipeline_I2V(BasePreprocessPipeline): @@ -26,18 +32,73 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): "text_encoder", "tokenizer", "vae", "image_encoder", "image_processor" ] + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="image_encoding_stage", + stage=ImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + def preprocess_image(self, image: PIL.Image.Image, height: int, width: int, + fastvideo_args: FastVideoArgs) -> Dict[str, Any]: + assert hasattr( + self, + "image_encoding_stage"), "Image encoding stage must be created" + + batch = ForwardBatch( + data_type="video", + pil_image=image, + ) + result_batch = self.image_encoding_stage(batch, fastvideo_args) + clip_features = result_batch.image_embeds[0] + + image = self.preprocess( + image, + vae_scale_factor=self.get_module("vae").spatial_compression_ratio, + height=height, + width=width) + + return { + "clip_feature": clip_features[0], + "pil_image": image, + } + + def preprocess_video(self, video: list[PIL.Image.Image], height: int, + width: int, + fastvideo_args: FastVideoArgs) -> Dict[str, Any]: + return self.preprocess_image(video[0], height, width, fastvideo_args) + def get_schema_fields(self) -> List[str]: """Get the schema fields for I2V pipeline.""" return [f.name for f in pyarrow_schema_i2v] def get_extra_features(self, valid_data: Dict[str, Any], fastvideo_args: FastVideoArgs) -> Dict[str, Any]: + + # TODO(will): move these to cpu at some point + self.get_module("image_encoder").to(get_torch_device()) + self.get_module("vae").to(get_torch_device()) + + features = {} """Get CLIP features from the first frame of each video.""" first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute( 0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C) + batch_size, _, num_frames, height, width = valid_data[ + "pixel_values"].shape + latent_height = height // self.get_module( + "vae").spatial_compression_ratio + latent_width = width // self.get_module("vae").spatial_compression_ratio processed_images = [] + # Frame has values between -1 and 1 for frame in first_frame: + frame = (frame + 1) * 127.5 frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8)) processed_img = self.get_module("image_processor")( images=frame_pil, return_tensors="pt") @@ -53,22 +114,84 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): clip_features = self.get_module("image_encoder")(**image_inputs) clip_features = clip_features.last_hidden_state - return {"clip_feature": clip_features} + features["clip_feature"] = clip_features + """Get VAE features from the first frame of each video""" + video_conditions = [] + for frame in first_frame: + processed_img = frame.to(device="cpu", dtype=torch.float32) + processed_img = processed_img.unsqueeze(0).permute(0, 3, 1, + 2).unsqueeze(2) + # (B, H, W, C) -> (B, C, 1, H, W) + video_condition = torch.cat([ + processed_img, + processed_img.new_zeros(processed_img.shape[0], + processed_img.shape[1], num_frames - 1, + height, width) + ], + dim=2) + video_condition = video_condition.to(device=get_torch_device(), + dtype=torch.float32) + video_conditions.append(video_condition) + + video_conditions = torch.cat(video_conditions, dim=0) + + with torch.autocast(device_type="cuda", + dtype=torch.float32, + enabled=True): + encoder_outputs = self.get_module("vae").encode(video_conditions) + + latent_condition = encoder_outputs.mean + if (hasattr(self.get_module("vae"), "shift_factor") + and self.get_module("vae").shift_factor is not None): + if isinstance(self.get_module("vae").shift_factor, torch.Tensor): + latent_condition -= self.get_module("vae").shift_factor.to( + latent_condition.device, latent_condition.dtype) + else: + latent_condition -= self.get_module("vae").shift_factor + + if isinstance(self.get_module("vae").scaling_factor, torch.Tensor): + latent_condition = latent_condition * self.get_module( + "vae").scaling_factor.to(latent_condition.device, + latent_condition.dtype) + else: + latent_condition = latent_condition * self.get_module( + "vae").scaling_factor + + mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, + latent_width) + mask_lat_size[:, :, list(range(1, num_frames))] = 0 + first_frame_mask = mask_lat_size[:, :, 0:1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, + dim=2, + repeats=self.get_module("vae").temporal_compression_ratio) + mask_lat_size = torch.concat( + [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) + mask_lat_size = mask_lat_size.view( + batch_size, -1, + self.get_module("vae").temporal_compression_ratio, latent_height, + latent_width) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to(latent_condition.device) + + image_latent = torch.concat([mask_lat_size, latent_condition], dim=1) + + features["first_frame_latent"] = image_latent + + return features def create_record( self, video_name: str, vae_latent: np.ndarray, text_embedding: np.ndarray, - text_attention_mask: np.ndarray, - valid_data: Optional[Dict[str, Any]], + valid_data: Dict[str, Any], idx: int, extra_features: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Create a record for the Parquet dataset with CLIP features.""" record = super().create_record(video_name=video_name, vae_latent=vae_latent, text_embedding=text_embedding, - text_attention_mask=text_attention_mask, valid_data=valid_data, idx=idx, extra_features=extra_features) @@ -87,7 +210,69 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): "clip_feature_dtype": "", }) - return record # type: ignore + if extra_features and "first_frame_latent" in extra_features: + first_frame_latent = extra_features["first_frame_latent"] + record.update({ + "first_frame_latent_bytes": + first_frame_latent.tobytes(), + "first_frame_latent_shape": + list(first_frame_latent.shape), + "first_frame_latent_dtype": + str(first_frame_latent.dtype), + }) + else: + record.update({ + "first_frame_latent_bytes": b"", + "first_frame_latent_shape": [], + "first_frame_latent_dtype": "", + }) + + if extra_features and "pil_image" in extra_features: + pil_image = extra_features["pil_image"] + record.update({ + "pil_image_bytes": pil_image.tobytes(), + "pil_image_shape": list(pil_image.shape), + "pil_image_dtype": str(pil_image.dtype), + }) + else: + record.update({ + "pil_image_bytes": b"", + "pil_image_shape": [], + "pil_image_dtype": "", + }) + + return record + + def pil_to_tensor(self, image: PIL.Image.Image) -> torch.Tensor: + image = image + + image = np.array(image).astype(np.float32) + image = torch.from_numpy(image) + return image + + def preprocess(self, + image: PIL.Image.Image, + vae_scale_factor: int, + height: int, + width: int, + resize_mode: str = "default") -> torch.Tensor: + image = [image] + + height, width = get_default_height_width(image[0], vae_scale_factor, + height, width) + image = [ + resize(i, height, width, resize_mode=resize_mode) for i in image + ] + image = pil_to_numpy(image) # to np + image = numpy_to_pt(image) # to pt + + do_normalize = True + if image.min() < 0: + do_normalize = False + if do_normalize: + image = normalize(image) + + return image EntryClass = PreprocessPipeline_I2V diff --git a/fastvideo/v1/pipelines/preprocess/v1_preprocess.py b/fastvideo/v1/pipelines/preprocess/v1_preprocess.py index 6c394b5bf..a8b562226 100644 --- a/fastvideo/v1/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/v1/pipelines/preprocess/v1_preprocess.py @@ -59,12 +59,6 @@ if __name__ == "__main__": default=2, help="Batch size (per device) for the training dataloader.", ) - parser.add_argument( - "--preprocess_text_batch_size", - type=int, - default=8, - help="Batch size (per device) for the training dataloader.", - ) parser.add_argument("--samples_per_file", type=int, default=64) parser.add_argument("--flush_frequency", type=int, @@ -79,7 +73,6 @@ if __name__ == "__main__": parser.add_argument("--video_length_tolerance_range", type=int, default=2.0) parser.add_argument("--group_frame", action="store_true") # TODO parser.add_argument("--group_resolution", action="store_true") # TODO - parser.add_argument("--dataset", default="t2v") parser.add_argument("--preprocess_task", type=str, default="t2v") parser.add_argument("--train_fps", type=int, default=30) parser.add_argument("--use_image_num", type=int, default=0) diff --git a/fastvideo/v1/pipelines/stages/encoding.py b/fastvideo/v1/pipelines/stages/encoding.py index d42773ef0..33bd76dca 100644 --- a/fastvideo/v1/pipelines/stages/encoding.py +++ b/fastvideo/v1/pipelines/stages/encoding.py @@ -51,22 +51,27 @@ class EncodingStage(PipelineStage): """ self.vae = self.vae.to(get_torch_device()) - image_path = batch.image_path - # TODO(will): remove this once we add input/output validation for stages - if image_path is None: - raise ValueError("Image Path must be provided") assert batch.height is not None assert batch.width is not None latent_height = batch.height // self.vae.spatial_compression_ratio latent_width = batch.width // self.vae.spatial_compression_ratio - image = batch.pil_image - image = self.preprocess( - image, - vae_scale_factor=self.vae.spatial_compression_ratio, - height=batch.height, - width=batch.width).to(get_torch_device(), dtype=torch.float32) - image = image.unsqueeze(2) + image = batch.preprocessed_image + # TODO(will) + if image is None: + assert batch.pil_image is not None + image = batch.pil_image + image = self.preprocess( + image, + vae_scale_factor=self.vae.spatial_compression_ratio, + height=batch.height, + width=batch.width).to(get_torch_device(), dtype=torch.float32) + + image = image.unsqueeze(2) + else: + # assumes image is loaded from parquet file and used for validation + image = image.transpose(1, 2) + logger.info("image: %s", image.shape) video_condition = torch.cat([ image, image.new_zeros(image.shape[0], image.shape[1], @@ -181,7 +186,7 @@ class EncodingStage(PipelineStage): fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify encoding stage inputs.""" result = VerificationResult() - result.add_check("pil_image", batch.pil_image, V.not_none) + # result.add_check("pil_image", batch.pil_image) result.add_check("height", batch.height, V.positive_int) result.add_check("width", batch.width, V.positive_int) result.add_check("generator", batch.generator, diff --git a/fastvideo/v1/tests/nightly/i2v-wandb-summary.json b/fastvideo/v1/tests/nightly/i2v-wandb-summary.json new file mode 100644 index 000000000..511c05de6 --- /dev/null +++ b/fastvideo/v1/tests/nightly/i2v-wandb-summary.json @@ -0,0 +1 @@ +{"step_time":2.245914653001819,"_wandb":{"runtime":1434},"learning_rate":1e-05,"grad_norm":0.57421875,"avg_step_time":1.1814782944297622,"train_loss":0.07932619750499725,"vsa_sparsity":0,"_timestamp":1.750578625921253e+09,"validation_videos_40_steps":{"count":1,"videos":[{"size":420969,"path":"media/videos/validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4","_type":"video-file","sha256":"581ff5eae2909d3a7b362dcb24d060c006c09e4d4deb44b82f4aa697f6789ba7"}],"captions":false,"_type":"videos"},"_runtime":1434.62395329,"_step":901} \ No newline at end of file diff --git a/fastvideo/v1/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4 b/fastvideo/v1/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4 new file mode 100644 index 000000000..e5ff30261 Binary files /dev/null and b/fastvideo/v1/tests/nightly/i2v_validation_videos_40_steps_900_581ff5eae2909d3a7b36.mp4 differ diff --git a/fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py new file mode 100644 index 000000000..d8f6f389b --- /dev/null +++ b/fastvideo/v1/tests/nightly/test_e2e_i2v_overfit_single_sample.py @@ -0,0 +1,177 @@ +import os +from pathlib import Path +from huggingface_hub import snapshot_download +import shutil +import subprocess +import sys +from fastvideo.v1.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision + +# Import the training pipeline +sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) + +NUM_NODES = "1" +MODEL_PATH = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" + +# preprocessing +DATA_DIR = "data" +LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "cats")) +NUM_GPUS_PER_NODE_PREPROCESSING = "1" +PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/v1/pipelines/preprocess/v1_preprocess.py" + +LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "cats_preprocessed_data_i2v")) + + +# training +NUM_GPUS_PER_NODE_TRAINING = "8" +TRAINING_ENTRY_FILE_PATH = "fastvideo/v1/training/wan_i2v_training_pipeline.py" +LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset") +LOCAL_VALIDATION_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "validation_parquet_dataset") +LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) + +def download_data(): + # create the data dir if it doesn't exist + data_dir = Path(DATA_DIR) + # if data_dir.exists(): + # print(f"Removing existing data directory at {data_dir}") + # shutil.rmtree(data_dir) + + print(f"Creating data directory at {data_dir}") + os.makedirs(data_dir) + + print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...") + try: + # result = snapshot_download( + # repo_id="wlsaidhi/cats-overfit-merged", + # local_dir=str(LOCAL_RAW_DATA_DIR), + # repo_type="dataset", + # resume_download=True, + # token=os.environ.get("HF_TOKEN"), # In case authentication is needed + # ) + print(f"Download completed successfully. Files downloaded to: {result}") + + # Verify the download + if not LOCAL_RAW_DATA_DIR.exists(): + raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") + + # List downloaded files + print("Downloaded files:") + for file in LOCAL_RAW_DATA_DIR.rglob("*"): + if file.is_file(): + print(f" - {file.relative_to(LOCAL_RAW_DATA_DIR)}") + + except Exception as e: + print(f"Error during download: {str(e)}") + raise + + +def run_preprocessing(): + # Run torchrun command + cmd = [ + "torchrun", + "--nnodes", NUM_NODES, + "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING, + PREPROCESSING_ENTRY_FILE_PATH, + "--model_path", MODEL_PATH, + "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge_1_sample.txt"), + "--preprocess_video_batch_size", "1", + "--max_height", "480", + "--max_width", "832", + "--num_frames", "77", + "--dataloader_num_workers", "0", + "--output_dir", LOCAL_PREPROCESSED_DATA_DIR, + "--train_fps", "16", + "--validation_dataset_file", os.path.join(LOCAL_RAW_DATA_DIR, "validation_i2v_prompt_1_sample.json"), + "--samples_per_file", "1", + "--flush_frequency", "1", + "--video_length_tolerance_range", "5", + "--preprocess_task", "i2v", + ] + + process = subprocess.run(cmd, check=True) + + +def run_training(): + cmd = [ + "torchrun", + "--nnodes", NUM_NODES, + "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING, + TRAINING_ENTRY_FILE_PATH, + "--model_path", MODEL_PATH, + "--inference_mode", "False", + "--pretrained_model_name_or_path", MODEL_PATH, + "--data_path", LOCAL_TRAINING_DATA_DIR, + "--validation_preprocessed_path", LOCAL_VALIDATION_DATA_DIR, + "--train_batch_size", "1", + "--num_latent_t", "8", + "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, + "--sp_size", NUM_GPUS_PER_NODE_TRAINING, + "--tp_size", NUM_GPUS_PER_NODE_TRAINING, + "--hsdp_replicate_dim", "1", + "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING, + "--num_gpus", NUM_GPUS_PER_NODE_TRAINING, + "--train_sp_batch_size", "1", + "--dataloader_num_workers", "10", + "--gradient_accumulation_steps", "1", + "--max_train_steps", "901", + "--learning_rate", "1e-5", + "--mixed_precision", "bf16", + "--checkpointing_steps", "6000", + "--validation_steps", "100", + "--validation_sampling_steps", "40", + "--log_validation", + "--checkpoints_total_limit", "3", + "--allow_tf32", + "--ema_start_step", "0", + "--training_cfg_rate", "0.1", + "--output_dir", LOCAL_OUTPUT_DIR, + "--tracker_project_name", "wan_i2v_finetune_overfit_ci", + "--num_height", "480", + "--num_width", "832", + "--num_frames", "81", + "--validation_guidance_scale", "1.0", + "--num_euler_timesteps", "50", + "--multi_phased_distill_schedule", "4000-1", + "--weight_decay", "0.01", + "--not_apply_cfg_solver", + "--dit_precision", "fp32", + "--max_grad_norm", "1.0", + ] + + print(f"Running training with command: {cmd}") + process = subprocess.run(cmd, check=True) + + +def test_e2e_overfit_single_sample(): + os.environ["WANDB_MODE"] = "online" + + # download_data() + run_preprocessing() + run_training() + + reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4") + print(f"reference_video_file: {reference_video_file}") + final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") + print(f"final_validation_video_file: {final_validation_video_file}") + + + # Ensure both files exist + assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" + assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" + + # Compute SSIM + mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision( + reference_video_file, + final_validation_video_file, + use_ms_ssim=True # Using MS-SSIM for better quality assessment + ) + + print("\n===== SSIM Results for Step 900 Validation =====") + print(f"Mean MS-SSIM: {mean_ssim:.4f}") + print(f"Min MS-SSIM: {min_ssim:.4f}") + print(f"Max MS-SSIM: {max_ssim:.4f}") + + assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}" + + +if __name__ == "__main__": + test_e2e_overfit_single_sample() \ No newline at end of file diff --git a/fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json b/fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json index dd7e64c0f..3f1cbb4df 100644 --- a/fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json +++ b/fastvideo/v1/tests/training/Vanilla/reference_wandb_summary.json @@ -1 +1 @@ -{"_timestamp":1.7496170016478686e+09,"validation_videos_8_steps":{"_type":"videos","count":5,"videos":[{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"},{"sha256":"cd72a3d513eca6b41b03b80e6fa044ce7219c35e969d2ca20b9cb48c91e585c6","size":477837,"path":"media/videos/validation_videos_8_steps_0_cd72a3d513eca6b41b03.mp4","_type":"video-file"},{"_type":"video-file","sha256":"43d47c211a69bf0be3544738e76e7d8fa58bb108c00cb21dc34eaad3c6ce7cc3","size":409419,"path":"media/videos/validation_videos_8_steps_0_43d47c211a69bf0be354.mp4"},{"_type":"video-file","sha256":"ea674ec9e200bc97563c9d87d9dc07110c3f42c0ab7277dd237ae96dd8f90a10","size":333966,"path":"media/videos/validation_videos_8_steps_0_ea674ec9e200bc97563c.mp4"},{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"}],"captions":false},"step_time":2.5065076276659966,"_wandb":{"runtime":53},"learning_rate":1e-06,"_step":5,"_runtime":53.172758961,"grad_norm":5.65625,"train_loss":0.3915919363498688,"avg_step_time":2.8116052336990833} \ No newline at end of file +{"_timestamp":1.7496170016478686e+09,"validation_videos_8_steps":{"_type":"videos","count":5,"videos":[{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"},{"sha256":"cd72a3d513eca6b41b03b80e6fa044ce7219c35e969d2ca20b9cb48c91e585c6","size":477837,"path":"media/videos/validation_videos_8_steps_0_cd72a3d513eca6b41b03.mp4","_type":"video-file"},{"_type":"video-file","sha256":"43d47c211a69bf0be3544738e76e7d8fa58bb108c00cb21dc34eaad3c6ce7cc3","size":409419,"path":"media/videos/validation_videos_8_steps_0_43d47c211a69bf0be354.mp4"},{"_type":"video-file","sha256":"ea674ec9e200bc97563c9d87d9dc07110c3f42c0ab7277dd237ae96dd8f90a10","size":333966,"path":"media/videos/validation_videos_8_steps_0_ea674ec9e200bc97563c.mp4"},{"sha256":"d81aa715df0c3ba4db8b242bc10442844453f3684a45dc55b89ecd27b2414fc7","size":429642,"path":"media/videos/validation_videos_8_steps_0_d81aa715df0c3ba4db8b.mp4","_type":"video-file"}],"captions":false},"step_time":2.5065076276659966,"_wandb":{"runtime":53},"learning_rate":1e-06,"_step":5,"_runtime":53.172758961,"grad_norm":0.408203125,"train_loss":0.07883700542151928,"avg_step_time":2.8116052336990833} \ No newline at end of file diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index ec327813e..901487d61 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -288,9 +288,6 @@ class TrainingPipeline(ComposedPipelineBase, ABC): assert self.transformer is not None assert self.training_args is not None assert training_batch.noisy_model_input is not None - assert training_batch.encoder_hidden_states is not None - assert training_batch.encoder_attention_mask is not None - assert training_batch.timesteps is not None assert training_batch.latents is not None assert training_batch.noise is not None assert training_batch.sigmas is not None diff --git a/fastvideo/v1/training/wan_i2v_training_pipeline.py b/fastvideo/v1/training/wan_i2v_training_pipeline.py new file mode 100644 index 000000000..b58cfa26a --- /dev/null +++ b/fastvideo/v1/training/wan_i2v_training_pipeline.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +import sys +from copy import deepcopy +from typing import Any, Dict + +import torch + +from fastvideo.v1.configs.sample import SamplingParam +from fastvideo.v1.dataset.dataloader.schema import ( + pyarrow_schema_i2v, pyarrow_schema_i2v_validation) +from fastvideo.v1.distributed import get_torch_device +from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) +from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch, + TrainingBatch) +from fastvideo.v1.pipelines.wan.wan_i2v_pipeline import ( + WanImageToVideoValidationPipeline) +from fastvideo.v1.training.training_pipeline import TrainingPipeline +from fastvideo.v1.utils import is_vsa_available + +vsa_available = is_vsa_available() + +logger = init_logger(__name__) + + +class WanI2VTrainingPipeline(TrainingPipeline): + """ + A training pipeline for Wan. + """ + _required_config_modules = ["scheduler", "transformer"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_training_stages(self, training_args: TrainingArgs): + """ + May be used in future refactors. + """ + pass + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_i2v + self.validation_dataset_schema = pyarrow_schema_i2v_validation + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing validation pipeline...") + args_copy = deepcopy(training_args) + + args_copy.inference_mode = True + args_copy.pipeline_config.vae_config.load_encoder = False + validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( + training_args.model_path, + args=None, + inference_mode=True, + loaded_modules={"transformer": self.get_module("transformer")}, + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus) + + self.validation_pipeline = validation_pipeline + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_args is not None + assert self.train_dataloader is not None + + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + # Reset iterator for next epoch + self.train_loader_iter = iter(self.train_dataloader) + # Get first batch of new epoch + batch = next(self.train_loader_iter) + + latents = batch['vae_latent'] + latents = latents[:, :, :self.training_args.num_latent_t] + encoder_hidden_states = batch['text_embedding'] + encoder_attention_mask = batch['text_attention_mask'] + clip_features = batch['clip_feature'] + image_latents = batch['first_frame_latent'] + image_latents = image_latents[:, :, :self.training_args.num_latent_t] + pil_image = batch['pil_image'] + infos = batch['info_list'] + + training_batch.latents = latents.to(get_torch_device(), + dtype=torch.bfloat16) + training_batch.encoder_hidden_states = encoder_hidden_states.to( + get_torch_device(), dtype=torch.bfloat16) + training_batch.encoder_attention_mask = encoder_attention_mask.to( + get_torch_device(), dtype=torch.bfloat16) + training_batch.preprocessed_image = pil_image.to(get_torch_device()) + training_batch.image_embeds = clip_features.to(get_torch_device()) + training_batch.image_latents = image_latents.to(get_torch_device()) + training_batch.infos = infos + + return training_batch + + def _prepare_dit_inputs(self, + training_batch: TrainingBatch) -> TrainingBatch: + """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning.""" + assert self.training_args is not None + assert training_batch.latents is not None + assert training_batch.encoder_hidden_states is not None + assert training_batch.encoder_attention_mask is not None + assert self.noise_random_generator is not None + assert training_batch.image_latents is not None + + # First, call parent method to prepare noise, timesteps, etc. for video latents + training_batch = super()._prepare_dit_inputs(training_batch) + + assert isinstance(training_batch.image_latents, torch.Tensor) + image_latents = training_batch.image_latents.to(get_torch_device(), + dtype=torch.bfloat16) + + training_batch.noisy_model_input = torch.cat( + [training_batch.noisy_model_input, image_latents], dim=1) + + return training_batch + + def _build_input_kwargs(self, + training_batch: TrainingBatch) -> TrainingBatch: + assert self.training_args is not None + assert training_batch.noisy_model_input is not None + assert training_batch.encoder_hidden_states is not None + assert training_batch.encoder_attention_mask is not None + assert training_batch.timesteps is not None + assert training_batch.image_embeds is not None + + # Image Embeds for conditioning + image_embeds = training_batch.image_embeds + assert torch.isnan(image_embeds).sum() == 0 + image_embeds = image_embeds.to(get_torch_device(), dtype=torch.bfloat16) + encoder_hidden_states_image = image_embeds + + # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs + training_batch.input_kwargs = { + "hidden_states": + training_batch.noisy_model_input, + "encoder_hidden_states": + training_batch.encoder_hidden_states, + "timestep": + training_batch.timesteps.to(get_torch_device(), + dtype=torch.bfloat16), + "encoder_attention_mask": + training_batch.encoder_attention_mask, + "encoder_hidden_states_image": + encoder_hidden_states_image, + "return_dict": + False, + } + return training_batch + + def _prepare_validation_inputs( + self, sampling_param: SamplingParam, training_args: TrainingArgs, + validation_batch: Dict[str, Any], num_inference_steps: int, + negative_prompt_embeds: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None + ) -> ForwardBatch: + embeddings = validation_batch['text_embedding'] + masks = validation_batch['text_attention_mask'] + clip_features = validation_batch['clip_feature'] + pil_image = validation_batch['pil_image'] + infos = validation_batch['info_list'] + prompt = infos[0]['prompt'] + + prompt_embeds = embeddings.to(get_torch_device()) + prompt_attention_mask = masks.to(get_torch_device()) + clip_features = clip_features.to(get_torch_device()) + + # Calculate sizes + latents_size = [(sampling_param.num_frames - 1) // 4 + 1, + sampling_param.height // 8, sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio + num_frames = (training_args.num_latent_t - + 1) * temporal_compression_factor + 1 + + # Prepare batch for validation + batch = ForwardBatch( + prompt=prompt, + data_type="video", + latents=None, + seed=self.seed, # Use deterministic seed + generator=torch.Generator(device="cpu").manual_seed(self.seed), + prompt_embeds=[prompt_embeds], + prompt_attention_mask=[prompt_attention_mask], + negative_prompt_embeds=[negative_prompt_embeds], + negative_attention_mask=[negative_prompt_attention_mask], + image_embeds=[clip_features], + preprocessed_image=pil_image, + height=training_args.num_height, + width=training_args.num_width, + num_frames=num_frames, + num_inference_steps= + num_inference_steps, # Use the current validation step + guidance_scale=sampling_param.guidance_scale, + n_tokens=n_tokens, + eta=0.0, + VSA_sparsity=training_args.VSA_sparsity, + ) + return batch + + +def main(args) -> None: + logger.info("Starting training pipeline...") + + pipeline = WanI2VTrainingPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args) + args = pipeline.training_args + pipeline.train() + logger.info("Training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.v1.fastvideo_args import TrainingArgs + from fastvideo.v1.utils import FlexibleArgumentParser + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.use_cpu_offload = False + main(args)