nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
SolitaryThinker
2025-08-13 02:19:52 +00:00
parent ebd23ba8e7
commit 6b8bf495ee
8 changed files with 178 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/distill/Wan2.1-I2V/crush_smol/distill_i2v_14B.sh
+2 -2
View File
@@ -6,7 +6,7 @@ export TOKENIZERS_PARALLELISM=false
# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
DATA_DIR="data/crush-smol_processed_wan21_i2v_14b_single/combined_parquet_dataset/"
DATA_DIR="data/crush-smol_processed_wan21_i2v_14b/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation.json"
NUM_GPUS=8
export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
@@ -56,7 +56,7 @@ validation_args=(
# --log_visualization
--log_visualization
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 1
--validation_steps 50
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
)
examples/distill/Wan2.1-I2V/crush_smol/validation.json
+63
View File
@@ -1,5 +1,32 @@
{
"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": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
"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,
@@ -8,6 +35,42 @@
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.",
"image_path": null,
"video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
}
]
}
examples/distill/Wan2.1-I2V/crush_smol/validation_orig.json
-76
View File
@@ -1,76 +0,0 @@
{
"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": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
"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": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "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.",
"image_path": null,
"video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
},
{
"caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.",
"image_path": null,
"video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
}
]
}
examples/distill/Wan2.1-I2V/crush_smol/validation_repeat.json
+13
View File
@@ -0,0 +1,13 @@
{
"data": [
{
"caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
"image_path": null,
"video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
"num_inference_steps": 40,
"height": 480,
"width": 832,
"num_frames": 77
}
]
}
fastvideo/pipelines/pipeline_batch_info.py
+2
View File
@@ -118,6 +118,8 @@ class ForwardBatch:
# Misc
save_video: bool = True
return_frames: bool = False
return_trajectory_latents: bool = False
trajectory_latents: list[torch.Tensor] = field(default_factory=list)
# TeaCache parameters
enable_teacache: bool = False
fastvideo/pipelines/stages/denoising.py
+4
View File
@@ -677,6 +677,7 @@ class DmdDenoisingStage(DenoisingStage):
# TODO(yongqi) hard code prepare latents
latents = torch.randn(
latents.permute(0, 2, 1, 3, 4).shape,
# latents.shape,
dtype=torch.bfloat16,
device="cuda",
generator=torch.Generator(device="cuda").manual_seed(42))
@@ -801,6 +802,9 @@ class DmdDenoisingStage(DenoisingStage):
scheduler=self.scheduler).unflatten(
0, pred_noise.shape[:2])
if batch.return_trajectory_latents:
batch.trajectory_latents.append(pred_video)
if i < len(timesteps) - 1:
next_timestep = timesteps[i + 1] * torch.ones(
[1], dtype=torch.long, device=pred_video.device)
fastvideo/training/distillation_pipeline.py
+77 -9
View File
@@ -241,7 +241,9 @@ class DistillationPipeline(TrainingPipeline):
# len(self.denoising_step_list), [1],
# device=self.device,
# dtype=torch.long)
target_timestep_idx = torch.tensor([2], device=self.device, dtype=torch.long)
target_timestep_idx = torch.tensor([2],
device=self.device,
dtype=torch.long)
target_timestep_idx_int = target_timestep_idx.item()
target_timestep = self.denoising_step_list[target_timestep_idx]
@@ -279,14 +281,15 @@ class DistillationPipeline(TrainingPipeline):
# t_expand = current_timestep_tensor.repeat(current_noise_latents.shape[0])
pred_flow = self.transformer(
training_batch_temp.input_kwargs["hidden_states"],
training_batch_temp.input_kwargs["encoder_hidden_states"],
training_batch_temp.
input_kwargs["encoder_hidden_states"],
current_timestep_tensor,
encoder_hidden_states_image=training_batch_temp.input_kwargs["encoder_hidden_states_image"],
encoder_hidden_states_image=training_batch_temp.
input_kwargs["encoder_hidden_states_image"],
guidance=None,
# **training_batch.conditional_dict,
# **training_batch.unconditional_dict,
).permute(
0, 2, 1, 3, 4)
).permute(0, 2, 1, 3, 4)
pred_clean = pred_noise_to_pred_video(
pred_noise=pred_flow.flatten(0, 1),
noise_input_latent=current_noise_latents.flatten(0, 1),
@@ -338,6 +341,63 @@ class DistillationPipeline(TrainingPipeline):
"generator_timestep"] = target_timestep.float().detach()
return pred_video
def _generator_forward_with_validation_pipeline(
self, training_batch: TrainingBatch) -> torch.Tensor:
dtype = training_batch.latents.dtype
# Step 1: Randomly sample a target timestep index from denoising_step_list
target_timestep_idx = torch.randint(0,
len(self.denoising_step_list), [1],
device=self.device,
dtype=torch.long)
# target_timestep_idx = torch.tensor([2], device=self.device, dtype=torch.long)
target_timestep_idx_int = target_timestep_idx.item()
target_timestep = self.denoising_step_list[target_timestep_idx]
# self.training_args.pipeline_config.dmd_denoising_steps,
# from copy import deepcopy
# args = deepcopy(self.training_args)
# args.pipeline_config.dmd_denoising_steps = self.training_args.pipeline_config.dmd_denoising_steps[:target_timestep_idx_int]
# Step 2: Simulate the multi-step inference process up to the target timestep
# Start from pure noise like in inference
current_noise_latents = torch.randn(self.video_latent_shape,
device=self.device,
dtype=dtype)
current_timestep = self.denoising_step_list[2]
current_timestep_tensor = current_timestep * torch.ones(
1, device=self.device, dtype=torch.long)
training_batch_temp = self._build_distill_input_kwargs(
current_noise_latents, current_timestep_tensor,
training_batch.conditional_dict, training_batch)
manual_batch = ForwardBatch(
data_type="video",
# **shallow_asdict(self.training_args),
latents=current_noise_latents.permute(0, 2, 1, 3, 4),
generator=torch.Generator(device="cpu").manual_seed(self.seed),
image_embeds=[
training_batch_temp.input_kwargs["encoder_hidden_states_image"]
],
prompt_embeds=[
training_batch_temp.input_kwargs["encoder_hidden_states"]
],
image_latent=training_batch_temp.input_kwargs["image_latent"],
timesteps=current_timestep_tensor,
return_trajectory_latents=True,
# n_tokens=n_tokens,
# eta=0.0,
# VSA_sparsity=training_args.VSA_sparsity,
)
batch = self.validation_pipeline.denoising_stage.forward(
batch=manual_batch, fastvideo_args=self.training_args)
latents = batch.trajectory_latents[target_timestep_idx_int]
latents = latents
training_batch.dmd_latent_vis_dict[
"generator_timestep"] = target_timestep.float().detach()
return latents
def _dmd_forward(self, generator_pred_video: torch.Tensor,
training_batch: TrainingBatch) -> torch.Tensor:
"""Compute DMD (Diffusion Model Distillation) loss."""
@@ -443,8 +503,12 @@ class DistillationPipeline(TrainingPipeline):
current_timestep=training_batch.timesteps,
attn_metadata=training_batch.attn_metadata_vsa):
if self.training_args.simulate_generator_forward:
generator_pred_video = self._generator_multi_step_simulation_forward(
training_batch)
if True:
generator_pred_video = self._generator_forward_with_validation_pipeline(
training_batch)
else:
generator_pred_video = self._generator_multi_step_simulation_forward(
training_batch)
else:
generator_pred_video = self._generator_forward(training_batch)
@@ -578,8 +642,12 @@ class DistillationPipeline(TrainingPipeline):
current_timestep=batch_gen.timesteps,
attn_metadata=batch_gen.attn_metadata_vsa):
if self.training_args.simulate_generator_forward:
generator_pred_video = self._generator_multi_step_simulation_forward(
batch_gen)
if True:
generator_pred_video = self._generator_forward_with_validation_pipeline(
batch_gen)
else:
generator_pred_video = self._generator_multi_step_simulation_forward(
batch_gen)
else:
generator_pred_video = self._generator_forward(
batch_gen)
fastvideo/training/wan_i2v_distillation_pipeline.py
+17 -16
View File
@@ -181,23 +181,23 @@ class WanI2VDistillationPipeline(DistillationPipeline):
text_dict: dict[str, torch.Tensor],
training_batch: TrainingBatch) -> TrainingBatch:
# Image Embeds for conditioning
# image_embeds = training_batch.image_embeds
# assert torch.isnan(image_embeds).sum() == 0
# image_embeds = image_embeds.to(get_local_torch_device(),
# dtype=torch.bfloat16)
from fastvideo.models.vision_utils import load_video
image = load_video("examples/distill/Wan2.1-I2V/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4")[0]
image_processor = self.validation_pipeline.get_module("image_processor")
image_encoder = self.validation_pipeline.get_module("image_encoder")
image_inputs = image_processor(
images=image, return_tensors="pt").to(get_local_torch_device())
from fastvideo.forward_context import set_forward_context
with set_forward_context(current_timestep=0, attn_metadata=None):
outputs = image_encoder(**image_inputs)
image_embeds = outputs.last_hidden_state
image_embeds = training_batch.image_embeds
assert torch.isnan(image_embeds).sum() == 0
image_embeds = image_embeds.to(get_local_torch_device(),
dtype=torch.bfloat16)
assert torch.isnan(image_embeds).sum() == 0
# from fastvideo.models.vision_utils import load_video
# image = load_video("examples/distill/Wan2.1-I2V/crush_smol/validation_dataset/1gGQy4nxyUo-Scene-016.mp4")[0]
# image_processor = self.validation_pipeline.get_module("image_processor")
# image_encoder = self.validation_pipeline.get_module("image_encoder")
# image_inputs = image_processor(
# images=image, return_tensors="pt").to(get_local_torch_device())
# from fastvideo.forward_context import set_forward_context
# with set_forward_context(current_timestep=0, attn_metadata=None):
# outputs = image_encoder(**image_inputs)
# image_embeds = outputs.last_hidden_state
# image_embeds = image_embeds.to(get_local_torch_device(),
# dtype=torch.bfloat16)
# assert torch.isnan(image_embeds).sum() == 0
# self.validation_pipeline.
noisy_model_input = torch.cat(
@@ -210,8 +210,9 @@ class WanI2VDistillationPipeline(DistillationPipeline):
4), # bs, c, t, h, w
"encoder_hidden_states": text_dict["encoder_hidden_states"],
"encoder_attention_mask": text_dict["encoder_attention_mask"],
"image_latent": training_batch.image_latents,
"timestep": timestep,
"encoder_hidden_states_image": image_embeds.permute(0, 2, 1, 3, 4),
"encoder_hidden_states_image": image_embeds,
"return_dict": False,
}
training_batch.noise_latents = noise_input