From c5f9ea53b2d484ed05e74841b2ce6ab4d2bb1ad8 Mon Sep 17 00:00:00 2001 From: William Lin Date: Mon, 15 Sep 2025 17:54:42 -0700 Subject: [PATCH] [self-forcing] [4/n] Preprocessing for collecting ODE trajectory (#788) --- .../causal_ode_init/crush_smol_prompts.txt | 47 +++ .../causal_ode_init/finetune_ode_init.sh | 93 +++++ .../preprocess_wan_data_crush_smol.sh | 25 ++ .../causal_ode_init/validation.json | 40 +++ .../wan_t2v_1.3B/crush_smol/validation.json | 2 +- fastvideo/configs/pipelines/__init__.py | 5 +- fastvideo/configs/pipelines/registry.py | 8 +- fastvideo/configs/sample/base.py | 14 + fastvideo/dataset/__init__.py | 2 +- fastvideo/entrypoints/video_generator.py | 3 + .../scheduling_self_forcing_flow_match.py | 7 + fastvideo/pipelines/composed_pipeline_base.py | 2 +- fastvideo/pipelines/pipeline_batch_info.py | 11 +- .../preprocess_pipeline_ode_trajectory.py | 323 ++++++++++++++++++ .../pipelines/preprocess/v1_preprocess.py | 10 +- fastvideo/pipelines/stages/decoding.py | 139 +++++--- fastvideo/pipelines/stages/denoising.py | 28 ++ .../workflow/test_components_parquet_saver.py | 2 - fastvideo/utils.py | 18 + 19 files changed, 718 insertions(+), 61 deletions(-) create mode 100644 examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt create mode 100644 examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh create mode 100644 examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh create mode 100644 examples/training/consistency_finetune/causal_ode_init/validation.json create mode 100644 fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py diff --git a/examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt b/examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt new file mode 100644 index 000000000..07a929f1a --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt @@ -0,0 +1,47 @@ +A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open. +The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object. +The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign. +A red toy car is being crushed by a large hydraulic press, which is flattening objects as if they were under a hydraulic press. +A large, cylindrical object is seen pressing down on a small orange ball, causing it to flatten as if it were under a hydraulic press. The background features a green wall with yellow and red warning signs. +The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is shown compressing a wooden object, which shatters into small pieces. The background features a green wall with a yellow sign displaying a lightning bolt. +A large metal cylinder is seen descending, flattening objects as if they were under a hydraulic press. The cylinder compresses a stack of matches and boxes, causing them to crumble into small pieces. The scene is set against a green background with yellow and red signs. +A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out. +The video shows a metal press flattening objects as if they were under a hydraulic press. The press is pressing down on a pile of colorful gummy candies, squishing them into a pile of squiggly shapes. The press is made of metal and has a large base, and the gummy candies are of various colors, including red, green, and orange. The background is a green wall, and the press is placed on a metal surface. +A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces. +The video shows a stack of colorful sponges being flattened as if they were under a hydraulic press. The sponges, which are pink, white, blue, and green, are compressed into a smaller size, demonstrating the press's power. The background features a green wall with a yellow and red sign, adding context to the setting. +A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, leaving a pile of debris around it. +A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press. +The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened. +The video shows a close-up of a metal cylinder pressing down on a yellow object, which is being flattened as if it were under a hydraulic press. The cylinder is positioned above the object, and the force is causing the object to compress and spread out, creating a visible deformation. The background is blurred, focusing attention on the action of the cylinder and the object being flattened. +A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press. +The video shows a hydraulic press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing two colorful objects that resemble sandwiches. The press is yellow and black striped, and the objects being flattened are placed on a metal plate. The background is green, and the press is moving down, compressing the objects. +The scene shows a metal press with a yellow and black striped pattern, holding a container filled with chocolate. A metal cylinder is descending, flattening the chocolate as if it were under a hydraulic press. The background is a green wall, and the press is mounted on a sturdy metal frame. +The video shows a colorful sponge being flattened as if it were under a hydraulic press, with the sponge being compressed and eventually flattened into a thin layer. +The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is pressing down on a stack of wooden blocks, causing them to crumble and break apart. The press is black and yellow striped, and the wooden blocks are small and rectangular. The background is green, and the press is sitting on a metal table. +A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces. +The video shows a stack of colorful sponges being flattened by a large, cylindrical object, which appears to be a hydraulic press. The sponges, which are pink, blue, white, and green, are compressed into a single layer, demonstrating the press's powerful force. The background features a green wall with a yellow and red sign, adding context to the industrial setting. +A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, demonstrating the immense pressure applied by the cylinder. +A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press. The popcorn is crushed and scattered around the base of the cylinder, creating a satisfying visual effect. +The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is composed of a large, cylindrical metal cylinder with yellow and black stripes, and a metal base. The objects being flattened are two cylindrical blocks of cotton candy, one pink and one blue. The press is positioned on a metal table, and the background features a green wall with a yellow and red sign. +The video shows a large orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened. +The video shows a cylindrical object being pressed down onto a flat surface, causing the objects beneath it to be flattened as if they were under a hydraulic press. The objects being flattened appear to be yellow and are being crushed into a pile of debris. The background is a greenish-gray color, and the surface on which the objects are being flattened is metallic and shiny. +A green and blue object with a spiky texture is being flattened by a large, cylindrical metal press, demonstrating its resilience and durability. +The video shows a stack of caramelized sugar cubes being flattened as if they were under a hydraulic press, resulting in a messy pile of broken sugar on the table. +A large metal cylinder is seen pressing down on a pile of colorful jelly beans, flattening them as if they were under a hydraulic press. +The video shows a machine with a yellow and black striped cylinder pressing down on a stack of colorful sponges, flattening them as if they were under a hydraulic press. The machine is situated in a green-walled room with warning signs in the background. +The video shows a machine with a yellow and black striped cylinder, which is pressing down on two colorful objects, flattening them as if they were under a hydraulic press. The machine appears to be in a workshop or industrial setting, with a green wall in the background. The objects being flattened are green and orange, and the machine is covered in dirt and grime, indicating it has been used frequently. +The video shows a large, industrial press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing a pile of pink objects into a pile of crumbs. The press is large and metallic, with a yellow and black striped pattern on its side. The background is a green wall with a yellow warning sign. +The video shows a pink, sparkly ball being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press. +A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and segments. +The video shows a machine with a yellow and black striped cylinder, which is flattening objects as if they were under a hydraulic press. The machine is pressing down on two colorful objects, causing them to compress and flatten. The background is a green wall, and the machine appears to be in a workshop or industrial setting. +The video shows a large, yellow and black striped cylinder flattening objects as if they were under a hydraulic press. The objects being flattened are pink and are being crushed into small pieces. The background is a green wall with a yellow sign. +The video shows a machine with a yellow and black striped cylinder pressing down on two colorful objects, which are flattened as if they were under a hydraulic press. The machine is positioned on a metal platform, and the background is a green wall. +A green cube is being compressed by a hydraulic press, which flattens the object as if it were under a hydraulic press. The press is shown in action, with the cube being squeezed into a smaller shape. +A pink, sparkly ball is being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press. +A red cabbage is being crushed by a hydraulic press, which flattens the objects as if they were under a hydraulic press. The press is shown in action, compressing the cabbage into a smaller, more compact form. +A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and pulp. +A large metal press is shown compressing a stack of burgers, causing them to be flattened and crushed into a pile of ground meat. +A pizza is being crushed by a hydraulic press, causing the toppings to spread out and the crust to crumble. +A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press. +A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table. +A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press. diff --git a/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh b/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh new file mode 100644 index 000000000..18e0cef09 --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh @@ -0,0 +1,93 @@ +#!/bin/bash + +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export TOKENIZERS_PARALLELISM=false + +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/" +VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +NUM_GPUS=1 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init_crush_smol" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "wan_ode_init_crush_smol" + --max_train_steps 6000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 21 + --num_height 480 + --num_width 832 + --num_frames 77 + --warp_denoising_step + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 1 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 50 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 6e-6 + --mixed_precision "bf16" + --checkpointing_steps 1000 + --weight_decay 1e-4 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +torchrun \ + --nnodes 1 \ + --nproc_per_node $NUM_GPUS \ + fastvideo/training/ode_causal_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh b/examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh new file mode 100644 index 000000000..127dece49 --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +GPU_NUM=1 # 2,4,8 +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +MODEL_TYPE="wan" +DATA_MERGE_PATH="$(dirname "$0")/crush_smol_prompts.txt" +OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/" + +torchrun --nproc_per_node=$GPU_NUM \ + fastvideo/pipelines/preprocess/v1_preprocess.py \ + --model_path $MODEL_PATH \ + --data_merge_path $DATA_MERGE_PATH \ + --preprocess_video_batch_size 1 \ + --seed 42 \ + --max_height 480 \ + --max_width 832 \ + --num_frames 81 \ + --flow_shift 5.0 \ + --dataloader_num_workers 0 \ + --output_dir=$OUTPUT_DIR \ + --train_fps 16 \ + --samples_per_file 8 \ + --flush_frequency 8 \ + --video_length_tolerance_range 5 \ + --preprocess_task "ode_trajectory" diff --git a/examples/training/consistency_finetune/causal_ode_init/validation.json b/examples/training/consistency_finetune/causal_ode_init/validation.json new file mode 100644 index 000000000..7b14204e5 --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/validation.json @@ -0,0 +1,40 @@ +{ + "data": [ + { + "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + }, + { + "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.", + "image_path": null, + "video_path": null, + "num_inference_steps": 40, + "height": 480, + "width": 832, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json b/examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json index 2f82a0db7..674c3cb62 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json @@ -28,4 +28,4 @@ "num_frames": 77 } ] -} \ No newline at end of file +} diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 82721919a..6ff503848 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -4,12 +4,13 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.registry import ( get_pipeline_config_cls_from_name) from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig -from fastvideo.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig, +from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, + WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig) __all__ = [ "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig", "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig", - "get_pipeline_config_cls_from_name" + "SelfForcingWanT2V480PConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index ffac8dd5b..62a4eefe5 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -11,9 +11,9 @@ from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig # isort: off from fastvideo.configs.pipelines.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, - SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, - Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, - WanT2V720PConfig) + Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config, + WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig, + SelfForcingWanT2V480PConfig) # isort: on from fastvideo.logger import init_logger from fastvideo.utils import (maybe_download_model_index, @@ -48,6 +48,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { "wanpipeline": lambda id: "wanpipeline" in id.lower(), "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(), + "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(), "stepvideo": lambda id: "stepvideo" in id.lower(), # Add other pipeline architecture detectors } @@ -60,6 +61,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { WanT2V480PConfig, # Base Wan config as fallback for any Wan variant "wanimagetovideo": WanI2V480PConfig, "wandmdpipeline": FastWan2_1_T2V_480P_Config, + "wancausaldmdpipeline": SelfForcingWanT2V480PConfig, "stepvideo": StepVideoT2VConfig # Other fallbacks by architecture } diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 65b129b4a..c4108613b 100644 --- a/fastvideo/configs/sample/base.py +++ b/fastvideo/configs/sample/base.py @@ -48,6 +48,8 @@ class SamplingParam: # Misc save_video: bool = True return_frames: bool = False + return_trajectory_latents: bool = False # returns all latents for each timestep + return_trajectory_decoded: bool = False # returns decoded latents for each timestep def __post_init__(self) -> None: self.data_type = "video" if self.num_frames > 1 else "image" @@ -205,6 +207,18 @@ class SamplingParam: help= "Path to a JSON file containing V-MoBA specific configurations.", ) + parser.add_argument( + "--return-trajectory-latents", + action="store_true", + default=SamplingParam.return_trajectory_latents, + help="Whether to return the trajectory", + ) + parser.add_argument( + "--return-trajectory-decoded", + action="store_true", + default=SamplingParam.return_trajectory_decoded, + help="Whether to return the decoded trajectory", + ) return parser diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py index 9e75b90d9..9c1e98473 100644 --- a/fastvideo/dataset/__init__.py +++ b/fastvideo/dataset/__init__.py @@ -37,7 +37,7 @@ def getdataset(args) -> VideoCaptionMergedDataset: temporal_sample=temporal_sample, transform_topcrop=transform_topcrop, seed=args.seed) - + def gettextdataset(args) -> TextDataset: return TextDataset(data_merge_path=args.data_merge_path, diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 019ccc527..7d97b21bc 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -344,6 +344,9 @@ class VideoGenerator: "size": (target_height, target_width, batch.num_frames), "generation_time": gen_time, "logging_info": logging_info, + "trajectory": output_batch.trajectory_latents, + "trajectory_timesteps": output_batch.trajectory_timesteps, + "trajectory_decoded": output_batch.trajectory_decoded, } def set_lora_adapter(self, diff --git a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py index 1d76c14ae..63d1edf60 100644 --- a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py +++ b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py @@ -62,8 +62,15 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): def step(self, model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final=False, return_dict=False, **kwargs): if timestep.ndim == 2: timestep = timestep.flatten(0, 1) + elif timestep.ndim == 0: + # handles the case where timestep is a scalar, this occurs when we + # use this scheduler for ODE trajectory + timestep = timestep.unsqueeze(0) + self.sigmas = self.sigmas.to(model_output.device) self.timesteps = self.timesteps.to(model_output.device) + timestep = timestep.to(model_output.device) + timestep_id = torch.argmin( (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index fd33efe55..a895e4c3d 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -40,7 +40,7 @@ class ComposedPipelineBase(ABC): _extra_config_module_map: dict[str, str] = {} training_args: TrainingArgs | None = None fastvideo_args: FastVideoArgs | TrainingArgs | None = None - modules: dict[str, torch.nn.Module] = {} + modules: dict[str, Any] = {} post_init_called: bool = False # TODO(will): args should support both inference args and training args diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index fcf88c15a..4fecf4cce 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -148,7 +148,12 @@ class ForwardBatch: modules: dict[str, Any] = field(default_factory=dict) # Final output (after pipeline completion) - output: Any = None + output: torch.Tensor | None = None + return_trajectory_latents: bool = False + return_trajectory_decoded: bool = False + trajectory_timesteps: list[torch.Tensor] | None = None + trajectory_latents: torch.Tensor | None = None + trajectory_decoded: list[torch.Tensor] | None = None # Extra parameters that might be needed by specific pipeline implementations extra: dict[str, Any] = field(default_factory=dict) @@ -207,6 +212,10 @@ class TrainingBatch: infos: list[dict[str, Any]] | None = None mask_lat_size: torch.Tensor | None = None + # ODE trajectory supervision + trajectory_latents: torch.Tensor | None = None + trajectory_timesteps: torch.Tensor | None = None + # Transformer inputs noisy_model_input: torch.Tensor | None = None timesteps: torch.Tensor | None = None diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py new file mode 100644 index 000000000..171cace77 --- /dev/null +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -0,0 +1,323 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +ODE Trajectory Data Preprocessing pipeline implementation. + +This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline +using the modular pipeline architecture. + +Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772 +""" + +import os +from collections.abc import Iterator +from typing import Any + +import pyarrow as pa +import torch +from torch.utils.data import DataLoader +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm import tqdm + +from fastvideo.configs.sample import SamplingParam +from fastvideo.dataset import gettextdataset +from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter, + records_to_table) +from fastvideo.dataset.dataloader.record_schema import ( + ode_text_only_record_creator) +from fastvideo.dataset.dataloader.schema import ( + pyarrow_schema_ode_trajectory_text_only) +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import ( + SelfForcingFlowMatchScheduler) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.preprocess.preprocess_pipeline_base import ( + BasePreprocessPipeline) +from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) +from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict + +logger = init_logger(__name__) + + +class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): + """ODE Trajectory preprocessing pipeline implementation.""" + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + ] + preprocess_dataloader: StatefulDataLoader + preprocess_loader_iter: Iterator[dict[str, Any]] + pbar: Any + num_processed_samples: int + + def get_pyarrow_schema(self) -> pa.Schema: + """Return the PyArrow schema for ODE Trajectory pipeline.""" + return pyarrow_schema_ode_trajectory_text_only + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + assert fastvideo_args.pipeline_config.flow_shift == 5 + self.modules["scheduler"] = SelfForcingFlowMatchScheduler( + shift=fastvideo_args.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True) + self.modules["scheduler"].set_timesteps(num_inference_steps=48, + denoising_strength=1.0) + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + 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="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + pipeline=self, + )) + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + def preprocess_text_and_trajectory(self, fastvideo_args: FastVideoArgs, + args): + """Preprocess text-only data and generate trajectory information.""" + + for batch_idx, data in enumerate(self.pbar): + if data is None: + continue + + with torch.inference_mode(): + # For text-only processing, we only need text data + # Filter out samples without text + valid_indices = [] + for i, text in enumerate(data["text"]): + if text and text.strip(): # Check if text is not empty + valid_indices.append(i) + self.num_processed_samples += len(valid_indices) + + if not valid_indices: + continue + + # Create new batch with only valid samples (text-only) + valid_data = { + "text": [data["text"][i] for i in valid_indices], + "path": [data["path"][i] for i in valid_indices], + } + + # Add fps and duration if available in data + if "fps" in data: + valid_data["fps"] = [data["fps"][i] for i in valid_indices] + if "duration" in data: + valid_data["duration"] = [ + data["duration"][i] for i in valid_indices + ] + + batch_captions = valid_data["text"] + # Encode text using the standalone TextEncodingStage API + prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( + batch_captions, + fastvideo_args, + encoder_index=[0], + return_attention_mask=True, + ) + prompt_embeds = prompt_embeds_list[0] + prompt_attention_masks = prompt_masks_list[0] + assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0] + + sampling_params = SamplingParam.from_pretrained(args.model_path) + + # encode negative prompt for trajectory collection + if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None: + negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text( + sampling_params.negative_prompt, + fastvideo_args, + encoder_index=[0], + return_attention_mask=True, + ) + negative_prompt_embed = negative_prompt_embeds_list[0][0] + negative_prompt_attention_mask = negative_prompt_masks_list[ + 0][0] + else: + negative_prompt_embed = None + negative_prompt_attention_mask = None + + trajectory_latents = [] + trajectory_timesteps = [] + trajectory_decoded = [] + + for i, (prompt_embed, prompt_attention_mask) in enumerate( + zip(prompt_embeds, prompt_attention_masks, + strict=False)): + prompt_embed = prompt_embed.unsqueeze(0) + prompt_attention_mask = prompt_attention_mask.unsqueeze(0) + + # Collect the trajectory data (text-to-video generation) + batch = ForwardBatch(**shallow_asdict(sampling_params), ) + batch.prompt_embeds = [prompt_embed] + batch.prompt_attention_mask = [prompt_attention_mask] + batch.negative_prompt_embeds = [negative_prompt_embed] + batch.negative_attention_mask = [ + negative_prompt_attention_mask + ] + batch.num_inference_steps = 48 + batch.return_trajectory_latents = True + # Enabling this will save the decoded trajectory videos. + # Used for debugging. + batch.return_trajectory_decoded = False + batch.height = args.max_height + batch.width = args.max_width + batch.fps = args.train_fps + batch.guidance_scale = 6.0 + batch.do_classifier_free_guidance = True + + result_batch = self.input_validation_stage( + batch, fastvideo_args) + result_batch = self.timestep_preparation_stage( + batch, fastvideo_args) + result_batch = self.latent_preparation_stage( + result_batch, fastvideo_args) + result_batch = self.denoising_stage(result_batch, + fastvideo_args) + result_batch = self.decoding_stage(result_batch, + fastvideo_args) + + trajectory_latents.append( + result_batch.trajectory_latents.cpu()) + trajectory_timesteps.append( + result_batch.trajectory_timesteps.cpu()) + trajectory_decoded.append(result_batch.trajectory_decoded) + + # Prepare extra features for text-only processing + extra_features = { + "trajectory_latents": trajectory_latents, + "trajectory_timesteps": trajectory_timesteps + } + + if batch.return_trajectory_decoded: + for i, decoded_frames in enumerate(trajectory_decoded): + for j, decoded_frame in enumerate(decoded_frames): + save_decoded_latents_as_video( + decoded_frame, + f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", + args.train_fps) + + # Prepare batch data for Parquet dataset + batch_data: list[dict[str, Any]] = [] + + # Add progress bar for saving outputs + save_pbar = tqdm(enumerate(valid_data["path"]), + desc="Saving outputs", + unit="item", + leave=False) + + for idx, video_path in save_pbar: + video_name = os.path.basename(video_path).split(".")[0] + + # Convert tensors to numpy arrays + text_embedding = prompt_embeds[idx].cpu().numpy() + + # Get extra features for this sample + sample_extra_features = {} + if extra_features: + for key, value in extra_features.items(): + if isinstance(value, torch.Tensor): + sample_extra_features[key] = value[idx].cpu( + ).numpy() + else: + assert isinstance(value, list) + if isinstance(value[idx], torch.Tensor): + sample_extra_features[key] = value[idx].cpu( + ).float().numpy() + else: + sample_extra_features[key] = value[idx] + + # Create record for Parquet dataset (text-only ODE schema) + record: dict[str, Any] = ode_text_only_record_creator( + video_name=video_name, + text_embedding=text_embedding, + caption=valid_data["text"][idx], + trajectory_latents=sample_extra_features[ + "trajectory_latents"], + trajectory_timesteps=sample_extra_features[ + "trajectory_timesteps"], + ) + batch_data.append(record) + + if batch_data: + write_pbar = tqdm(total=1, + desc="Writing to Parquet dataset", + unit="batch") + table = records_to_table(batch_data, + self.get_pyarrow_schema()) + write_pbar.update(1) + write_pbar.close() + + if not hasattr(self, 'dataset_writer'): + self.dataset_writer = ParquetDatasetWriter( + out_dir=self.combined_parquet_dir, + samples_per_file=args.samples_per_file, + ) + self.dataset_writer.append_table(table) + + logger.info("Collected batch with %s samples", len(table)) + + if self.num_processed_samples >= args.flush_frequency: + written = self.dataset_writer.flush() + logger.info("Flushed %s samples to parquet", written) + self.num_processed_samples = 0 + + # Final flush for any remaining samples + if hasattr(self, 'dataset_writer'): + written = self.dataset_writer.flush(write_remainder=True) + if written: + logger.info("Final flush wrote %s samples", written) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args): + if not self.post_init_called: + self.post_init() + + self.local_rank = int(os.getenv("RANK", 0)) + os.makedirs(args.output_dir, exist_ok=True) + # Create directory for combined data + self.combined_parquet_dir = os.path.join(args.output_dir, + "combined_parquet_dataset") + os.makedirs(self.combined_parquet_dir, exist_ok=True) + + # Loading dataset + train_dataset = gettextdataset(args) + + self.preprocess_dataloader = DataLoader( + train_dataset, + batch_size=args.preprocess_video_batch_size, + num_workers=args.dataloader_num_workers, + ) + + self.preprocess_loader_iter = iter(self.preprocess_dataloader) + + self.num_processed_samples = 0 + # Add progress bar for video preprocessing + self.pbar = tqdm(self.preprocess_loader_iter, + desc="Processing videos", + unit="batch", + disable=self.local_rank != 0) + + # Initialize class variables for data sharing + self.video_data: dict[str, Any] = {} # Store video metadata and paths + self.latent_data: dict[str, Any] = {} # Store latent tensors + self.preprocess_text_and_trajectory(fastvideo_args, args) + + +EntryClass = PreprocessPipeline_ODE_Trajectory diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 82d0514ac..058cb2e5d 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -10,6 +10,8 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import ( PreprocessPipeline_I2V) +from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import ( + PreprocessPipeline_ODE_Trajectory) from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import ( PreprocessPipeline_T2V) from fastvideo.pipelines.preprocess.preprocess_pipeline_text import ( @@ -48,13 +50,16 @@ def main(args) -> None: text_encoder_cpu_offload=False, pipeline_config=pipeline_config, ) - if args.preprocess_task == "t2v": PreprocessPipeline = PreprocessPipeline_T2V elif args.preprocess_task == "i2v": PreprocessPipeline = PreprocessPipeline_I2V elif args.preprocess_task == "text_only": PreprocessPipeline = PreprocessPipeline_Text + elif args.preprocess_task == "ode_trajectory": + assert args.flow_shift is not None, "flow_shift is required for ode_trajectory" + fastvideo_args.pipeline_config.flow_shift = args.flow_shift + PreprocessPipeline = PreprocessPipeline_ODE_Trajectory else: raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. " f"Valid options: t2v, i2v, ode_trajectory, text_only") @@ -100,10 +105,11 @@ 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("--flow_shift", type=float, default=None) parser.add_argument("--preprocess_task", type=str, default="t2v", - choices=["t2v", "i2v", "text_only"], + choices=["t2v", "i2v", "text_only", "ode_trajectory"], help="Type of preprocessing task to run") parser.add_argument("--train_fps", type=int, default=30) parser.add_argument("--use_image_num", type=int, default=0) diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 8f226b5c4..d75da2c93 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -50,6 +50,63 @@ class DecodingStage(PipelineStage): result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)]) return result + @torch.no_grad() + def decode(self, latents: torch.Tensor, + fastvideo_args: FastVideoArgs) -> torch.Tensor: + """ + Decode latent representations into pixel space using VAE. + + Args: + latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents) + fastvideo_args: Configuration containing: + - disable_autocast: Whether to disable automatic mixed precision (default: False) + - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") + - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency + + Returns: + Decoded video tensor with shape (batch, channels, frames, height, width), + normalized to [0, 1] range and moved to CPU as float32 + """ + self.vae = self.vae.to(get_local_torch_device()) + latents = latents.to(get_local_torch_device()) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + if isinstance(self.vae.scaling_factor, torch.Tensor): + latents = latents / self.vae.scaling_factor.to( + latents.device, latents.dtype) + else: + latents = latents / self.vae.scaling_factor + + # Apply shifting if needed + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latents += self.vae.shift_factor.to(latents.device, + latents.dtype) + else: + latents += self.vae.shift_factor + + # Decode latents + with torch.autocast(device_type="cuda", + dtype=vae_dtype, + enabled=vae_autocast_enabled): + if fastvideo_args.pipeline_config.vae_tiling: + self.vae.enable_tiling() + # if fastvideo_args.vae_sp: + # self.vae.enable_parallel() + if not vae_autocast_enabled: + latents = latents.to(vae_dtype) + image = self.vae.decode(latents) + + # Normalize image to [0, 1] range + image = (image / 2 + 0.5).clamp(0, 1) + return image + @torch.no_grad() def forward( self, @@ -59,13 +116,28 @@ class DecodingStage(PipelineStage): """ Decode latent representations into pixel space. + This method processes the batch through the VAE decoder, converting latent + representations to pixel-space video/images. It also optionally decodes + trajectory latents for visualization purposes. + Args: - batch: The current batch information. - fastvideo_args: The inference arguments. + batch: The current batch containing: + - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents) + - return_trajectory_decoded (optional): Flag to decode trajectory latents + - trajectory_latents (optional): Latents at different timesteps + - trajectory_timesteps (optional): Corresponding timesteps + fastvideo_args: Configuration containing: + - output_type: "latent" to skip decoding, otherwise decode to pixels + - vae_cpu_offload: Whether to offload VAE to CPU after decoding + - model_loaded: Track VAE loading state + - model_paths: Path to VAE model if loading needed Returns: - The batch with decoded outputs. + Modified batch with: + - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 + - trajectory_decoded (if requested): List of decoded frames per timestep """ + # load vae if not already loaded (used for memory constrained devices) pipeline = self.pipeline() if self.pipeline else None if not fastvideo_args.model_loaded["vae"]: loader = VAELoader() @@ -75,58 +147,29 @@ class DecodingStage(PipelineStage): pipeline.add_module("vae", self.vae) fastvideo_args.model_loaded["vae"] = True - self.vae = self.vae.to(get_local_torch_device()) - - latents = batch.latents - # TODO(will): remove this once we add input/output validation for stages - if latents is None: - raise ValueError("Latents must be provided") - - # Skip decoding if output type is latent if fastvideo_args.output_type == "latent": - image = latents + frames = batch.latents else: - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[ - fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = (vae_dtype != torch.float32 - ) and not fastvideo_args.disable_autocast + frames = self.decode(batch.latents, fastvideo_args) - if isinstance(self.vae.scaling_factor, torch.Tensor): - latents = latents / self.vae.scaling_factor.to( - latents.device, latents.dtype) - else: - latents = latents / self.vae.scaling_factor - - # Apply shifting if needed - if (hasattr(self.vae, "shift_factor") - and self.vae.shift_factor is not None): - if isinstance(self.vae.shift_factor, torch.Tensor): - latents += self.vae.shift_factor.to(latents.device, - latents.dtype) - else: - latents += self.vae.shift_factor - - # Decode latents - with torch.autocast(device_type="cuda", - dtype=vae_dtype, - enabled=vae_autocast_enabled): - if fastvideo_args.pipeline_config.vae_tiling: - self.vae.enable_tiling() - # if fastvideo_args.vae_sp: - # self.vae.enable_parallel() - if not vae_autocast_enabled: - latents = latents.to(vae_dtype) - image = self.vae.decode(latents) - - # Normalize image to [0, 1] range - image = (image / 2 + 0.5).clamp(0, 1) + # decode trajectory latents if needed + if batch.return_trajectory_decoded: + batch.trajectory_decoded = [] + assert batch.trajectory_latents is not None, "batch should have trajectory latents" + for idx in range(batch.trajectory_latents.shape[1]): + # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] + cur_latent = batch.trajectory_latents[:, idx, :, :, :, :] + cur_timestep = batch.trajectory_timesteps[idx] + logger.info("decoding trajectory latent for timestep: %s", + cur_timestep) + decoded_frames = self.decode(cur_latent, fastvideo_args) + batch.trajectory_decoded.append(decoded_frames.cpu().float()) # Convert to CPU float32 for compatibility - image = image.cpu().float() + frames = frames.cpu().float() # Update batch with decoded image - batch.output = image + batch.output = frames # Offload models if needed if hasattr(self, 'maybe_free_model_hooks'): diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 9acea2afc..9e486dd84 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -253,6 +253,10 @@ class DenoisingStage(PipelineStage): patch_size[2]) seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size + # Initialize lists for ODE trajectory + trajectory_timesteps: list[torch.Tensor] = [] + trajectory_latents: list[torch.Tensor] = [] + # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -433,6 +437,11 @@ class DenoisingStage(PipelineStage): latents = (1. - mask2[0]) * z + mask2[0] * latents # latents = latents.unsqueeze(0) + # save trajectory latents if needed + if batch.return_trajectory_latents: + trajectory_timesteps.append(t) + trajectory_latents.append(latents) + # Update progress bar if i == len(timesteps) - 1 or ( (i + 1) > num_warmup_steps and @@ -440,9 +449,28 @@ class DenoisingStage(PipelineStage): and progress_bar is not None): progress_bar.update() + # Gather results if using sequence parallelism + trajectory_tensor: torch.Tensor | None = None + if trajectory_latents: + trajectory_tensor = torch.stack(trajectory_latents, dim=1) + trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, + dim=0) + else: + trajectory_tensor = None + trajectory_timesteps_tensor = None + # Gather results if using sequence parallelism if sp_group: latents = sequence_model_parallel_all_gather(latents, dim=2) + if batch.return_trajectory_latents: + trajectory_tensor = trajectory_tensor.to( + get_local_torch_device()) + trajectory_tensor = sequence_model_parallel_all_gather( + trajectory_tensor, dim=3) + + if trajectory_tensor is not None and trajectory_timesteps_tensor is not None: + batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() + batch.trajectory_latents = trajectory_tensor.cpu() # Update batch with final latents batch.latents = latents diff --git a/fastvideo/tests/workflow/test_components_parquet_saver.py b/fastvideo/tests/workflow/test_components_parquet_saver.py index 7061d44de..aceea9dbe 100644 --- a/fastvideo/tests/workflow/test_components_parquet_saver.py +++ b/fastvideo/tests/workflow/test_components_parquet_saver.py @@ -64,5 +64,3 @@ def test_parquet_dataset_saver_flush_and_last(tmp_path: Path): assert len(files2) == 2 total = sum(pq.read_table(str(f)).num_rows for f in files2) assert total == 5 - - diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 48a1eec7d..50937aa24 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -23,10 +23,14 @@ from typing import Any, TypeVar, cast import cloudpickle import filelock +import imageio +import numpy as np import torch +import torchvision import yaml from diffusers.loaders.lora_base import ( _best_guess_weight_name) # watch out for potetential removal from diffusers +from einops import rearrange from huggingface_hub import snapshot_download from remote_pdb import RemotePdb from torch.distributed.fsdp import MixedPrecisionPolicy @@ -886,3 +890,17 @@ def best_output_size(w, h, dw, dh, expected_area): return ow1, oh1 else: return ow2, oh2 + + +def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], + output_path: str, fps: int): + # Process outputs + videos = rearrange(decoded_latents, "b c t h w -> t b c h w") + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).numpy().astype(np.uint8)) + + os.makedirs(os.path.dirname(output_path), exist_ok=True) + imageio.mimsave(output_path, frames, fps=fps, format="mp4")