runnable
This commit is contained in:
@@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user