[WIP][Feat] distill run in single gpu

Author: Eigensystem

fastvideo/distill/solver.py

@@ -7,11 +7,27 @@
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from diffusers.utils import BaseOutput, logging
 
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
-
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
+def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
+    if linear_steps is None:
+        linear_steps = num_steps // 2
+    linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+    threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
+    quadratic_steps = num_steps - linear_steps
+    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
+    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
+    const = quadratic_coef * (linear_steps**2)
+    quadratic_sigma_schedule = [
+        quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
+    ]
+    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
+    sigma_schedule = [1.0 - x for x in sigma_schedule]
+    return sigma_schedule
+
+
 @dataclass
 class PCMFMSchedulerOutput(BaseOutput):
     prev_sample: torch.FloatTensor

fastvideo/v1/dataset/parquet_datasets.py

@@ -119,6 +119,9 @@ def __init__(self,
                 plan = json.load(f)
             self.neg_metadata = plan["negative_prompt"][0]
 
+        self.uncond_prompt_embed = torch.zeros(512, 4096).to(torch.float32)
+        self.uncond_prompt_mask = torch.zeros(1, 512).bool()
+
     def _load_and_cache_negative_prompt(self) -> None:
         """Load and cache the negative prompt. Only rank 0 in each SP group should call this."""
         if not self.validation or self.neg_metadata is None:
@@ -188,6 +191,7 @@ def get_validation_negative_prompt(
                 lat = lat[:, self.rank_in_sp_group, :, :, :]
             return lat, emb, mask, info
 
+
     def __len__(self):
         if self.local_indices is None:
             try:

fastvideo/v1/fastvideo_args.py

@@ -590,6 +590,8 @@ class TrainingArgs(FastVideoArgs):
     pred_decay_type: str = ""
     hunyuan_teacher_disable_cfg: bool = False
 
+    use_lora: bool = False
+
     # master_weight_type
     master_weight_type: str = ""
 

fastvideo/v1/models/loader/component_loader.py

@@ -366,6 +366,7 @@ class TransformerLoader(ComponentLoader):
     def load(self, model_path: str, architecture: str,
              fastvideo_args: FastVideoArgs):
         """Load the transformer based on the model path, architecture, and inference args."""
+        print(f"Loading transformer from {model_path}")
         config = get_diffusers_config(model=model_path)
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -56,7 +56,7 @@ def __init__(self,
         use. The pipeline should be stateless and not hold any batch state.
         """
 
-        if fastvideo_args.training_mode:
+        if fastvideo_args.training_mode or fastvideo_args.distill_mode:
             assert isinstance(fastvideo_args, TrainingArgs)
             self.training_args = fastvideo_args
             assert self.training_args is not None
@@ -94,11 +94,12 @@ def __init__(self,
                 self.initialize_validation_pipeline(self.training_args)
             self.initialize_training_pipeline(self.training_args)
 
+        # TODO(jinzhe): discuss this
         if fastvideo_args.distill_mode:
-            self.initialize_distillation_pipeline(fastvideo_args)
-
-        if fastvideo_args.log_validation:
-            self.initialize_validation_pipeline(fastvideo_args)
+            assert self.training_args is not None
+            if self.training_args.log_validation:
+                self.initialize_validation_pipeline(self.training_args)
+            self.initialize_distillation_pipeline(self.training_args)
 
         self.initialize_pipeline(fastvideo_args)