[feat] add hunyuan adv (#79)

fastvideo/distill.py

@@ -30,7 +30,7 @@ from fastvideo.utils.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpoint
 from diffusers import (
     FlowMatchEulerDiscreteScheduler,
 )
-from fastvideo.utils.load import get_no_split_modules, load_transformer
+from fastvideo.utils.load import  load_transformer
 from fastvideo.distill.solver import EulerSolver, extract_into_tensor
 from copy import deepcopy
 from diffusers.optimization import get_scheduler

fastvideo/distill/discriminator.py

@@ -62,9 +62,10 @@ class Discriminator(nn.Module):
         stride=8,
         num_h_per_head=1,
         adapter_channel_dims=[3072],
+        total_layers = 48,
     ):
         super().__init__()
-        adapter_channel_dims = adapter_channel_dims * (48 // stride)
+        adapter_channel_dims = adapter_channel_dims * (total_layers // stride)
         self.stride = stride
         self.num_h_per_head = num_h_per_head
         self.head_num = len(adapter_channel_dims)
@@ -89,9 +90,9 @@ class Discriminator(nn.Module):
 
             return custom_forward
 
-        assert len(features) // self.stride == len(self.heads)
-        for i in range(0, len(features), self.stride):
-            for h in self.heads[i // self.stride]:
+        assert len(features) == len(self.heads)
+        for i in range(0, len(features)):
+            for h in self.heads[i]:
                 # out = torch.utils.checkpoint.checkpoint(
                 #     create_custom_forward(h),
                 #     features[i],

fastvideo/distill_adv.py

@@ -22,6 +22,7 @@ from torch.distributed.fsdp import (
     StateDictType,
     FullStateDictConfig,
 )
+from fastvideo.utils.load import  load_transformer
 
 from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
 import json
@@ -76,6 +77,7 @@ def gan_d_loss(
     encoder_hidden_states,
     encoder_attention_mask,
     weight,
+    discriminator_head_stride
 ):
     loss = 0.0
     # collate sample_fake and sample_real
@@ -85,7 +87,8 @@ def gan_d_loss(
             encoder_hidden_states,
             timestep,
             encoder_attention_mask,
-            output_attn=True,
+            output_features=True,
+            output_features_stride=discriminator_head_stride,
             return_dict=False,
         )[1]
         real_features = teacher_transformer(
@@ -93,7 +96,8 @@ def gan_d_loss(
             encoder_hidden_states,
             timestep,
             encoder_attention_mask,
-            output_attn=True,
+            output_features=True,
+            output_features_stride=discriminator_head_stride,
             return_dict=False,
         )[1]
 
@@ -115,6 +119,7 @@ def gan_g_loss(
     encoder_hidden_states,
     encoder_attention_mask,
     weight,
+    discriminator_head_stride
 ):
     loss = 0.0
     features = teacher_transformer(
@@ -122,7 +127,8 @@ def gan_g_loss(
         encoder_hidden_states,
         timestep,
         encoder_attention_mask,
-        output_attn=True,
+        output_features=True,
+        output_features_stride=discriminator_head_stride,
         return_dict=False,
     )[1]
     fake_outputs = discriminator(
@@ -135,20 +141,19 @@ def gan_g_loss(
     return loss
 
 
-def train_one_step_mochi(
+def distill_one_step_adv(
     transformer,
+    model_type,
     teacher_transformer,
     optimizer,
     discriminator,
     discriminator_optimizer,
-    global_step,
     lr_scheduler,
     loader,
     noise_scheduler,
     solver,
     noise_random_generator,
     sp_size,
-    precondition_outputs,
     max_grad_norm,
     uncond_prompt_embed,
     uncond_prompt_mask,
@@ -157,6 +162,7 @@ def train_one_step_mochi(
     not_apply_cfg_solver,
     distill_cfg,
     adv_weight,
+    discriminator_head_stride
 ):
     optimizer.zero_grad()
     discriminator_optimizer.zero_grad()
@@ -167,7 +173,7 @@ def train_one_step_mochi(
         latents_attention_mask,
         encoder_attention_mask,
     ) = next(loader)
-    model_input = normalize_mochi_dit_input(latents)
+    model_input = normalize_dit_input(model_type, latents)
     noise = torch.randn_like(model_input)
     bsz = model_input.shape[0]
     index = torch.randint(
@@ -284,6 +290,7 @@ def train_one_step_mochi(
             encoder_hidden_states.float(),
             encoder_attention_mask,
             1.0,
+            discriminator_head_stride
         )
     g_loss += g_gan_loss
     g_loss.backward()
@@ -308,6 +315,7 @@ def train_one_step_mochi(
             encoder_hidden_states,
             encoder_attention_mask,
             1.0,
+            discriminator_head_stride,
         )
 
     d_loss.backward()
@@ -347,21 +355,14 @@ def main(args):
 
     main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
     # keep the master weight to float32
-    if args.dit_model_name_or_path:
-        transformer = transformer = MochiTransformer3DModel.from_pretrained(
-            args.dit_model_name_or_path,
-            torch_dtype=torch.float32,
-            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
-        )
-    else:
-        transformer = MochiTransformer3DModel.from_pretrained(
-            args.pretrained_model_name_or_path,
-            subfolder="transformer",
-            torch_dtype=torch.float32,
-            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
-        )
+    transformer = load_transformer(
+        args.model_type,
+        args.dit_model_name_or_path,
+        args.pretrained_model_name_or_path,
+        torch.float32 if args.master_weight_type == "fp32" else torch.bfloat16,
+    )
     teacher_transformer = deepcopy(transformer)
-    discriminator = Discriminator(args.discriminator_head_stride)
+    discriminator = Discriminator(args.discriminator_head_stride, total_layers = 48 if args.model_type =="mochi" else 40)
 
     if args.use_lora:
         transformer.requires_grad_(False)
@@ -383,15 +384,20 @@ def main(args):
     main_print(
         f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
     )
-    fsdp_kwargs = get_dit_fsdp_kwargs(
-        args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload
+    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
+        transformer,
+        args.fsdp_sharding_startegy,
+        args.use_lora,
+        args.use_cpu_offload,
+        args.master_weight_type,
     )
     discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs(args.master_weight_type)
     if args.use_lora:
+        assert args.model_type == "mochi", "LoRA is only supported for Mochi model."
         transformer.config.lora_rank = args.lora_rank
         transformer.config.lora_alpha = args.lora_alpha
         transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
-        transformer._no_split_modules = ["MochiTransformerBlock"]
+        transformer._no_split_modules = no_split_modules
         fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
 
     transformer = FSDP(
@@ -409,8 +415,12 @@ def main(args):
     main_print(f"--> model loaded")
 
     if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer, args.selective_checkpointing)
-        apply_fsdp_checkpointing(teacher_transformer, args.selective_checkpointing)
+        apply_fsdp_checkpointing(
+            transformer, no_split_modules, args.selective_checkpointing
+        )
+        apply_fsdp_checkpointing(
+            teacher_transformer, no_split_modules, args.selective_checkpointing
+        )
     # Set model as trainable.
     transformer.train()
     teacher_transformer.requires_grad_(False)
@@ -562,39 +572,49 @@ def main(args):
 
     step_times = deque(maxlen=100)
     # log_validation(args, transformer, device,
-    #         torch.bfloat16, init_steps, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False)
-
+    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
+    def get_num_phases(multi_phased_distill_schedule, step):
+        # step-phase,step-phase
+        multi_phases = multi_phased_distill_schedule.split(",")
+        phase = multi_phases[-1].split("-")[-1]
+        for step_phases in multi_phases:
+            phase_step, phase = step_phases.split("-")
+            if step <= int(phase_step):
+                return int(phase)
+        return phase
     for i in range(init_steps):
         _ = next(loader)
     for step in range(init_steps + 1, args.max_train_steps + 1):
+        assert args.multi_phased_distill_schedule is not None
+        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
         start_time = time.time()
         (
             generator_loss,
             generator_grad_norm,
             discriminator_loss,
             discriminator_grad_norm,
-        ) = train_one_step_mochi(
+        ) = distill_one_step_adv(
             transformer,
+            args.model_type,
             teacher_transformer,
             optimizer,
             discriminator,
             discriminator_optimizer,
-            step,
             lr_scheduler,
             loader,
             noise_scheduler,
             solver,
             noise_random_generator,
             args.sp_size,
-            args.precondition_outputs,
             args.max_grad_norm,
             uncond_prompt_embed,
             uncond_prompt_mask,
             args.num_euler_timesteps,
-            args.validation_sampling_steps,
+            num_phases,
             args.not_apply_cfg_solver,
             args.distill_cfg,
             args.adv_weight,
+            args.discriminator_head_stride
         )
 
         step_time = time.time() - start_time
@@ -633,15 +653,17 @@ def main(args):
                 )
             else:
                 # Your existing checkpoint saving code
-                save_checkpoint_generator_discriminator(
-                    transformer,
-                    optimizer,
-                    discriminator,
-                    discriminator_optimizer,
-                    rank,
-                    args.output_dir,
-                    step,
-                )
+                # TODO 
+                # save_checkpoint_generator_discriminator(
+                #     transformer,
+                #     optimizer,
+                #     discriminator,
+                #     discriminator_optimizer,
+                #     rank,
+                #     args.output_dir,
+                #     step,
+                # )
+                save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
             main_print(f"--> checkpoint saved at step {step}")
             dist.barrier()
         if args.log_validation and step % args.validation_steps == 0:
@@ -655,25 +677,17 @@ def main(args):
                 shift=args.shift,
                 num_euler_timesteps=args.num_euler_timesteps,
                 linear_quadratic_threshold=args.linear_quadratic_threshold,
+                linear_range=args.linear_range,
                 ema=False,
             )
+            
 
     if args.use_lora:
         save_lora_checkpoint(
             transformer, optimizer, rank, args.output_dir, args.max_train_steps
         )
     else:
-        save_checkpoint(
-            transformer, optimizer, rank, args.output_dir, args.max_train_steps
-        )
-        save_checkpoint(
-            discriminator,
-            discriminator_optimizer,
-            rank,
-            args.output_dir,
-            step,
-            discriminator=True,
-        )
+        save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
 
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
@@ -682,6 +696,9 @@ def main(args):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
+    parser.add_argument(
+        "--model_type", type=str, default="mochi", help="The type of model to train."
+    )
     # dataset & dataloader
     parser.add_argument("--data_json_path", type=str, required=True)
     parser.add_argument("--num_frames", type=int, default=163)
@@ -712,16 +729,9 @@ if __name__ == "__main__":
     parser.add_argument("--ema_decay", type=float, default=0.999)
     parser.add_argument("--ema_start_step", type=int, default=0)
     parser.add_argument("--cfg", type=float, default=0.1)
-    parser.add_argument(
-        "--precondition_outputs",
-        action="store_true",
-        help="Whether to precondition the outputs of the model.",
-    )
-
     # validation & logs
-    parser.add_argument("--validation_prompt_dir", type=str)
-    parser.add_argument("--validation_sampling_steps", type=int, default=64)
-    parser.add_argument("--validation_guidance_scale", type=float, default=4.5)
+    parser.add_argument("--validation_sampling_steps", type=str, default="64")
+    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
     parser.add_argument("--validation_steps", type=float, default=64)
     parser.add_argument("--log_validation", action="store_true")
     parser.add_argument("--tracker_project_name", type=str, default=None)
@@ -750,6 +760,7 @@ if __name__ == "__main__":
             " training using `--resume_from_checkpoint`."
         ),
     )
+    parser.add_argument("--validation_prompt_dir", type=str)
     parser.add_argument("--shift", type=float, default=1.0)
     parser.add_argument(
         "--resume_from_checkpoint",
@@ -866,6 +877,7 @@ if __name__ == "__main__":
         "--lora_rank", type=int, default=128, help="LoRA rank parameter. "
     )
     parser.add_argument("--fsdp_sharding_startegy", default="full")
+    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
     parser.add_argument(
         "--gradient_accumulation_steps",
         type=int,

fastvideo/models/hunyuan/modules/models.py

@@ -601,7 +601,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         encoder_attention_mask: torch.Tensor,
-        output_attn=False,
+        output_features=False,
+        output_features_stride = 8,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = False,
         guidance=None,
@@ -663,6 +664,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
 
         # Merge txt and img to pass through single stream blocks.
         x = torch.cat((img, txt), 1)
+        if output_features:
+            features_list = []
         if len(self.single_blocks) > 0:
             for _, block in enumerate(self.single_blocks):
                 single_block_args = [
@@ -674,6 +677,8 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
                 ]
 
                 x = block(*single_block_args)
+                if output_features and _ % output_features_stride == 0:
+                    features_list.append(x[:, :img_seq_len, ...])
 
         img = x[:, :img_seq_len, ...]
 
@@ -681,10 +686,12 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
 
         img = self.unpatchify(img, tt, th, tw)
-        if return_dict:
-            out["x"] = img
-            return out
-        return (img,)
+        assert return_dict == False, "return_dict is not supported."
+        if output_features:
+            features_list = torch.stack(features_list, dim=0)
+        else:
+            features_list = None
+        return (img, features_list)
 
     def unpatchify(self, x, t, h, w):
         """

fastvideo/models/mochi_hf/modeling_mochi.py

@@ -621,7 +621,8 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         encoder_attention_mask: torch.Tensor,
-        output_attn=False,
+        output_features=False,
+        output_features_stride = 8,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = False,
     ) -> torch.Tensor:
@@ -697,7 +698,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
                     encoder_attention_mask,
                     temb,
                     image_rotary_emb,
-                    output_attn,
+                    output_features,
                     **ckpt_kwargs,
                 )
             else:
@@ -707,9 +708,10 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
                     encoder_attention_mask=encoder_attention_mask,
                     temb=temb,
                     image_rotary_emb=image_rotary_emb,
-                    output_attn=output_attn,
+                    output_attn=output_features,
                 )
-            attn_outputs_list.append(attn_outputs)
+            if i % output_features_stride == 0:
+                attn_outputs_list.append(attn_outputs)
 
         hidden_states = self.norm_out(hidden_states, temb)
         hidden_states = self.proj_out(hidden_states)
@@ -724,7 +726,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             # remove `lora_scale` from each PEFT layer
             unscale_lora_layers(self, lora_scale)
 
-        if not output_attn:
+        if not output_features:
             attn_outputs_list = None
         else:
             attn_outputs_list = torch.stack(attn_outputs_list, dim=0)

pyproject.toml

@@ -21,14 +21,9 @@ dependencies = [
 "timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0", 
 "scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0", 
 "moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", 
-"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0"]
+"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0", "loguru"]
 
 
-[project.optional-dependencies]
-hunyuan = [
-    "loguru"
-]
-
 [tool.setuptools.packages.find]
 exclude = ["assets*", "docker*", "docs", "scripts*"]
 

scripts/distill_hunyuan.sh

@@ -3,19 +3,19 @@ export WANDB_MODE=online
 export WANDB_API_KEY=4f6de3765d6464f43e0506ec7d785641af645e73
 
 
-torchrun --nnodes 1 --nproc_per_node 8\
-    fastvideo/distill.py\
+torchrun --nnodes 1 --nproc_per_node 4\
+    fastvideo/distill_adv.py\
     --seed 42\
     --pretrained_model_name_or_path data/hunyuan\
     --dit_model_name_or_path data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
     --model_type "hunyuan" \
     --cache_dir "data/.cache"\
-    --data_json_path "data/Hunyuan-Mixkit-Data/videos2caption.json"\
-    --validation_prompt_dir "data/Hunyuan-Mixkit-Data/validation"\
+    --data_json_path "data/Hunyuan-Distill-Data/videos2caption.json"\
+    --validation_prompt_dir "data/Hunyuan-Distill-Data/validation"\
     --gradient_checkpointing\
     --train_batch_size=1\
-    --num_latent_t 24\
-    --sp_size 4\
+    --num_latent_t 1\
+    --sp_size 1\
     --train_sp_batch_size 1\
     --dataloader_num_workers 4\
     --gradient_accumulation_steps=1\