This commit is contained in:
SolitaryThinker
2025-08-12 19:58:36 +00:00
parent d87ebe051b
commit 0c1cb0c19d
5 changed files with 39 additions and 76 deletions
@@ -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/combined_parquet_dataset/"
DATA_DIR="data/crush-smol_processed_wan21_i2v_14b_single/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="examples/distill/Wan2.1-I2V/crush_smol/validation.json"
NUM_GPUS=8
export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
@@ -17,7 +17,7 @@ export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
training_args=(
--tracker_project_name "wan_i2v_distill"
--output_dir "checkpoints/wan_i2v_distill"
--wandb_run_name "wan21_i2v_14b_ode_3steps"
--wandb_run_name "no_back_wan21_i2v_14b_ode_3steps_8gpu"
--max_train_steps 2000
--train_batch_size 1
--train_sp_batch_size 1
@@ -56,7 +56,7 @@ validation_args=(
# --log_visualization
--log_visualization
--validation_dataset_file "$VALIDATION_DATASET_FILE"
--validation_steps 50
--validation_steps 1
--validation_sampling_steps "3"
--validation_guidance_scale "1.0" # not used for dmd inference
)
@@ -4,7 +4,7 @@ GPU_NUM=1 # 2,4,8
MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
MODEL_TYPE="wan"
DATA_MERGE_PATH="data/crush-smol/merge.txt"
OUTPUT_DIR="data/crush-smol_processed_wan21_i2v_14b/"
OUTPUT_DIR="data/crush-smol_processed_wan21_i2v_14b_single/"
torchrun --nproc_per_node=$GPU_NUM \
fastvideo/pipelines/preprocess/v1_preprocess.py \
@@ -18,7 +18,7 @@ torchrun --nproc_per_node=$GPU_NUM \
--dataloader_num_workers 0 \
--output_dir=$OUTPUT_DIR \
--train_fps 16 \
--samples_per_file 8 \
--samples_per_file 8\
--flush_frequency 8 \
--video_length_tolerance_range 5 \
--preprocess_task "i2v"
@@ -1,32 +1,5 @@
{
"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,
@@ -35,42 +8,6 @@
"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
}
]
}
+17 -5
View File
@@ -237,10 +237,11 @@ class DistillationPipeline(TrainingPipeline):
dtype = 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.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]
@@ -272,8 +273,19 @@ class DistillationPipeline(TrainingPipeline):
training_batch_temp = self._build_distill_input_kwargs(
current_noise_latents, current_timestep_tensor,
training_batch.conditional_dict, training_batch)
# pred_flow = self.transformer(
# **training_batch_temp.input_kwargs).permute(
# 0, 2, 1, 3, 4)
# t_expand = current_timestep_tensor.repeat(current_noise_latents.shape[0])
pred_flow = self.transformer(
**training_batch_temp.input_kwargs).permute(
training_batch_temp.input_kwargs["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"],
guidance=None,
# **training_batch.conditional_dict,
# **training_batch.unconditional_dict,
).permute(
0, 2, 1, 3, 4)
pred_clean = pred_noise_to_pred_video(
pred_noise=pred_flow.flatten(0, 1),
@@ -181,10 +181,24 @@ 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 = 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 = 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(
[noise_input,
@@ -197,7 +211,7 @@ class WanI2VDistillationPipeline(DistillationPipeline):
"encoder_hidden_states": text_dict["encoder_hidden_states"],
"encoder_attention_mask": text_dict["encoder_attention_mask"],
"timestep": timestep,
"encoder_hidden_states_image": image_embeds,
"encoder_hidden_states_image": image_embeds.permute(0, 2, 1, 3, 4),
"return_dict": False,
}
training_batch.noise_latents = noise_input