diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index ff5073635..9d2140f23 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -8,7 +8,7 @@ def main(): # model. # If a local path is provided, FastVideo will make a best effort # attempt to identify the optimal arguments. - model_name = "SFWan2.1-T2V-1.3B-Diffusers/snapshots/bd804886ec4fe04fcc5a3fe17644428b9cee94c0/" + model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" generator = VideoGenerator.from_pretrained( model_name, # FastVideo will automatically handle distributed setup diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh b/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh index eb96f8747..b05d18c4f 100644 --- a/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init.sh @@ -5,7 +5,8 @@ export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA -MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=1 diff --git a/examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm new file mode 100644 index 000000000..551cbf807 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm @@ -0,0 +1,131 @@ +#!/bin/bash +#SBATCH --job-name=ode_vidprom2k +#SBATCH --partition=main +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=ode_vidprom2k_output/ode_vidprom2k.out +#SBATCH --error=ode_vidprom2k_output/ode_vidprom2k.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv2 + +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + + +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing/" +VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init_vidprom2k" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "vidprom2k_wan_ode_init_5e-6" + # --resume_from_checkpoint "ode_init_diffusers/" + --max_train_steps 6001 + --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 + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 8 + --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 100 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 5e-6 + --mixed_precision "bf16" + --checkpointing_steps 2000 + --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 + --enable_gradient_checkpointing_type "full" +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +srun torchrun \ + --nnodes $SLURM_JOB_NUM_NODES \ + --nproc_per_node $NUM_GPUS \ + --node_rank $SLURM_PROCID \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + 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/ode_init/finetune_ode_init_crush.slurm b/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm new file mode 100644 index 000000000..d9aa22f44 --- /dev/null +++ b/examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm @@ -0,0 +1,131 @@ +#!/bin/bash +#SBATCH --job-name=ode_crush +#SBATCH --partition=main +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=ode_crush_output/ode_crush.out +#SBATCH --error=ode_crush_output/ode_crush.err +#SBATCH --exclusive +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate will-fv2 + +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID +export TOKENIZERS_PARALLELISM=false +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde' +# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + + +# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" +MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/" +VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" +NUM_GPUS=8 +# export CUDA_VISIBLE_DEVICES=4,5 +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "wan_ode_init" + --output_dir "wan_ode_init" + --override_transformer_cls_name "CausalWanTransformer3DModel" + --wandb_run_name "fixed_wan_ode_init_5e-6" + # --resume_from_checkpoint "ode_init_diffusers/" + --max_train_steps 6001 + --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 + # --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_GPUS + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 8 + --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 100 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 5e-6 + --mixed_precision "bf16" + --checkpointing_steps 2000 + --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 + --enable_gradient_checkpointing_type "full" +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +srun torchrun \ + --nnodes $SLURM_JOB_NUM_NODES \ + --nproc_per_node $NUM_GPUS \ + --node_rank $SLURM_PROCID \ + --rdzv_backend=c10d \ + --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/ode_causal_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/fastvideo/layers/layernorm.py b/fastvideo/layers/layernorm.py index ff508dc90..2a9a09059 100644 --- a/fastvideo/layers/layernorm.py +++ b/fastvideo/layers/layernorm.py @@ -103,11 +103,10 @@ class ScaleResidual(nn.Module): def forward(self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: """Apply gated residual connection.""" - logger.info("x.shape: %s", x.shape) - if isinstance(gate, torch.Tensor): - logger.info("gate.shape: %s", gate.shape) - else: - logger.info("gate: %s", gate) + # logger.info("x.shape: %s", x.shape) + # if isinstance(gate, torch.Tensor): + # logger.info("gate.shape: %s", gate.shape) + num_frames = gate.shape[1] frame_seqlen = x.shape[1] // num_frames return residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2) @@ -182,13 +181,13 @@ class ScaleResidualLayerNormScaleShift(nn.Module): but before normalization) """ # Apply residual connection with gating - logger.info("x.shape: %s", x.shape) + # logger.info("x.shape: %s", x.shape) if isinstance(gate, int): # used by cross-attention, should be 1 assert gate == 1 residual_output = residual + x * gate elif isinstance(gate, torch.Tensor): - logger.info("gate.shape: %s", gate.shape) + # logger.info("gate.shape: %s", gate.shape) if gate.dim() == 3: # used by bidirectional self attention residual_output = residual + x * gate @@ -200,7 +199,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module): # residual_output = residual + x * gate else: raise ValueError(f"Gate type {type(gate)} not supported") - logger.info("residual_output.shape: %s", residual_output.shape) + # logger.info("residual_output.shape: %s", residual_output.shape) # Apply normalization normalized = self.norm(residual_output) diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 8df4f3561..689f32b78 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -378,8 +378,10 @@ class CausalWanTransformer3DModel(BaseDiT): elementwise_affine=False, dtype=torch.float32, compute_dtype=torch.float32) - self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) + # Debug: Log configuration values + proj_out_dim = config.out_channels * math.prod(config.patch_size) + + self.proj_out = nn.Linear(inner_dim, proj_out_dim) self.scale_shift_table = nn.Parameter( torch.randn(1, 2, inner_dim) / inner_dim**0.5) @@ -468,7 +470,7 @@ class CausalWanTransformer3DModel(BaseDiT): This function will be run for num_frame times. Process the latent frames one by one (1560 tokens each) """ - logger.info("forward inference hidden_states.shape: %s", hidden_states.shape) + # logger.info("forward inference hidden_states.shape: %s", hidden_states.shape) orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): @@ -505,7 +507,7 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) - logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape) + # logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) @@ -547,9 +549,9 @@ class CausalWanTransformer3DModel(BaseDiT): **causal_kwargs) # 5. Output norm, projection & unpatchify - logger.info("===== INFERENCE 5. Output norm, projection & unpatchify") - logger.info("hidden_states.shape: %s", hidden_states.shape) - logger.info("temb.shape: %s", temb.shape) + # logger.info("===== INFERENCE 5. Output norm, projection & unpatchify") + # logger.info("hidden_states.shape: %s", hidden_states.shape) + # logger.info("temb.shape: %s", temb.shape) shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) hidden_states = self.norm_out(hidden_states, shift, scale) @@ -573,7 +575,8 @@ class CausalWanTransformer3DModel(BaseDiT): start_frame: int = 0, **kwargs) -> torch.Tensor: - logger.info("===== forward train hidden_states.shape: %s", hidden_states.shape) + # logger.info("===== forward train hidden_states.shape: %s", hidden_states.shape) + # logger.info("===== forward train timestep.shape: %s", timestep.shape) orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] @@ -619,10 +622,13 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) - logger.info("forward train flattened and transposed hidden_states.shape: %s", hidden_states.shape) + # logger.info("forward train flattened and transposed hidden_states.shape: %s", hidden_states.shape) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + # logger.info("forward train timestep_proj.shape: %s", timestep_proj.shape) + # logger.info("forward train timestep.shape: %s", timestep.shape) + # logger.info("forward train temb.shape: %s", temb.shape) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) if encoder_hidden_states_image is not None: @@ -655,13 +661,23 @@ class CausalWanTransformer3DModel(BaseDiT): block_mask=self.block_mask) # 5. Output norm, projection & unpatchify - logger.info("===== TRAIN 5. Output norm, projection & unpatchify") - logger.info("hidden_states.shape: %s", hidden_states.shape) - logger.info("temb.shape: %s", temb.shape) - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, - dim=1) + # logger.info("===== TRAIN 5. Output norm, projection & unpatchify") + # logger.info("hidden_states.shape: %s", hidden_states.shape) + # logger.info("temb.shape: %s", temb.shape) + # shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, + temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) + # logger.info("WTFWTF train temb.shape: %s", temb.shape) + # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, + dim=2) + # logger.info("DEBUG scale.shape: %s", scale.shape) + # logger.info("DEBUG shift.shape: %s", shift.shape) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) + # logger.info("DEBUG after proj_out hidden_states.shape: %s", hidden_states.shape) + # logger.info(f"DEBUG reshape dimensions: batch_size={batch_size}, post_patch_num_frames={post_patch_num_frames}") + # logger.info(f"DEBUG reshape dimensions: post_patch_height={post_patch_height}, post_patch_width={post_patch_width}") + # logger.info(f"DEBUG patch dimensions: p_t={p_t}, p_h={p_h}, p_w={p_w}") hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 47aae801e..f87551593 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -6,6 +6,9 @@ from typing import Any import torch +from fastvideo.logger import init_logger +logger = init_logger(__name__) + # TODO(PY): move it elsewhere def auto_attributes(init_func): """ @@ -146,7 +149,10 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor, """ Convert predicted noise to clean latent. """ - # timestep = timestep.expand(noise_input_latent.shape[0]) + logger.info(f"timestep: {timestep.shape}") + logger.info(f"noise_input_latent: {noise_input_latent.shape}") + logger.info(f"pred_noise: {pred_noise.shape}") + timestep = timestep.expand(noise_input_latent.shape[0]) dtype = pred_noise.dtype device = pred_noise.device pred_noise = pred_noise.float().to(device) diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py index 7e6624803..2eb32792e 100644 --- a/fastvideo/training/ode_causal_pipeline.py +++ b/fastvideo/training/ode_causal_pipeline.py @@ -117,6 +117,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): else: raise ValueError( f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}") + trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) # Move to device device = get_local_torch_device() @@ -164,14 +165,14 @@ class ODEInitTrainingPipeline(TrainingPipeline): traj_timesteps: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: torch.Tensor - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: device = get_local_torch_device() - logger.info(f"traj_latents: {traj_latents.shape}") - logger.info(f"traj_timesteps: {traj_timesteps.shape}") + # logger.info(f"traj_latents: {traj_latents.shape}") + # logger.info(f"traj_timesteps: {traj_timesteps.shape}") # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S] - B, S, num_channels, num_frames, height, width = traj_latents.shape + B, S, num_frames, num_channels, height, width = traj_latents.shape # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps if self._cached_closest_idx_per_dmd is None: @@ -181,16 +182,17 @@ class ODEInitTrainingPipeline(TrainingPipeline): # distances_ks: [K, S] = |s_steps - dmd| distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs() self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu() # [K] - logger.info(f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}") + # logger.info(f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}") - logger.info(f"traj_latents: {traj_latents.shape}") + # logger.info(f"traj_latents: {traj_latents.shape}") # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K] # Output: [B, K, C, T, H, W] relevant_traj_latents = torch.index_select( traj_latents, dim=1, index=self._cached_closest_idx_per_dmd.to(traj_latents.device) ) - logger.info(f"relevant_traj_latents: {relevant_traj_latents.shape}") + # logger.info(f"relevant_traj_latents: {relevant_traj_latents.shape}") + target_latent = relevant_traj_latents[:, -1] # assert relevant_traj_latents.shape[0] == 1 indexes = self._get_timestep( # [B, num_frames] @@ -202,11 +204,11 @@ class ODEInitTrainingPipeline(TrainingPipeline): uniform_timestep=False ) # noisy_input = relevant_traj_latents[indexes] - logger.info(f"indexes: {indexes.shape}") + # logger.info(f"indexes: {indexes.shape}") noisy_input = torch.gather( relevant_traj_latents, dim=1, - index=indexes.reshape(B, 1, 1, num_frames, 1, 1).expand( - -1, -1, num_channels, -1, height, width).to(self.device) + index=indexes.reshape(B, 1, num_frames, 1, 1, 1).expand( + -1, -1, -1, num_channels, height, width).to(self.device) ).squeeze(1) # noisy_input = noisy_input.unsqueeze(0) @@ -225,15 +227,15 @@ class ODEInitTrainingPipeline(TrainingPipeline): # Scale model input as in inference for consistency with stored trajectories # noisy_input = self.modules["scheduler"].scale_model_input(noisy_input, t) - logger.info(f"indexes: {indexes.shape}") - logger.info(f"indexes: {indexes}") + # logger.info(f"indexes: {indexes.shape}") + # logger.info(f"indexes: {indexes}") timestep = self.dmd_denoising_steps[indexes] - logger.info(f"timestep: {timestep.shape}") - logger.info(f"timestep: {timestep}") + # logger.info(f"timestep: {timestep.shape}") + # logger.info(f"timestep: {timestep}") # Prepare inputs for transformer input_kwargs = { - "hidden_states": noisy_input, + "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), "encoder_hidden_states": encoder_hidden_states, "timestep": timestep.to(device, dtype=torch.bfloat16), "encoder_attention_mask": encoder_attention_mask, @@ -243,7 +245,8 @@ class ODEInitTrainingPipeline(TrainingPipeline): with set_forward_context(current_timestep=timestep, attn_metadata=None, forward_batch=None): - noise_pred = self.transformer(**input_kwargs) + noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4) + # logger.info(f"noise_pred: {noise_pred.shape}") if isinstance(noise_pred, (tuple, list)): noise_pred = noise_pred[0] @@ -251,7 +254,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): noise_pred = pred_noise_to_pred_video( pred_noise=noise_pred.flatten(0, 1), noise_input_latent=noisy_input.flatten(0, 1), - timestep=timestep, + timestep=timestep.flatten(0, 1), scheduler=self.modules["scheduler"]).unflatten( 0, noise_pred.shape[:2])