Small Fixes & Features (#378)

fastvideo/v1/configs/pipelines/base.py

@@ -55,6 +55,8 @@ class PipelineConfig:
 
     # STA (Spatial-Temporal Attention) parameters
     mask_strategy_file_path: Optional[str] = None
+
+    # Compilation
     enable_torch_compile: bool = False
 
     @classmethod

fastvideo/v1/examples/inference/basic/basic.py

@@ -1,5 +1,4 @@
 from fastvideo import VideoGenerator
-# from fastvideo.v1.configs.sample import SamplingParam
 
 def main():
     # FastVideo will automatically use the optimal default arguments for the

fastvideo/v1/fastvideo_args.py

@@ -331,6 +331,12 @@ class FastVideoArgs:
                 f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
             )
 
+        if self.enable_torch_compile and self.num_gpus > 1:
+            logger.warning(
+                "Currently torch compile does not work with multi-gpu. Setting enable_torch_compile to False"
+            )
+            self.enable_torch_compile = False
+
 
 _current_fastvideo_args = None
 

fastvideo/v1/models/loader/component_loader.py

@@ -399,6 +399,9 @@ class TransformerLoader(ComponentLoader):
                                 device=fastvideo_args.device,
                                 cpu_offload=fastvideo_args.use_cpu_offload,
                                 default_dtype=default_dtype)
+        if fastvideo_args.enable_torch_compile:
+            logger.info("Torch Compile enabled for DiT")
+            model = torch.compile(model)
 
         total_params = sum(p.numel() for p in model.parameters())
         logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

fastvideo/v1/pipelines/stages/decoding.py

@@ -7,7 +7,6 @@ import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.models.vaes.common import ParallelTiledVAE
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.stages.base import PipelineStage
 from fastvideo.v1.utils import PRECISION_TO_TYPE
@@ -23,8 +22,8 @@ class DecodingStage(PipelineStage):
     output format (e.g., pixel values).
     """
 
-    def __init__(self, vae: ParallelTiledVAE) -> None:
-        self.vae: ParallelTiledVAE = vae
+    def __init__(self, vae) -> None:
+        self.vae = vae
 
     def forward(
         self,
@@ -41,6 +40,8 @@ class DecodingStage(PipelineStage):
         Returns:
             The batch with decoded outputs.
         """
+        self.vae = self.vae.to(fastvideo_args.device)
+
         latents = batch.latents
         # TODO(will): remove this once we add input/output validation for stages
         if latents is None:
@@ -95,4 +96,7 @@ class DecodingStage(PipelineStage):
         if hasattr(self, 'maybe_free_model_hooks'):
             self.maybe_free_model_hooks()
 
+        self.vae.to("cpu")
+        torch.cuda.empty_cache()
+
         return batch

fastvideo/v1/pipelines/stages/input_validation.py

@@ -3,6 +3,8 @@
 Input validation stage for diffusion pipelines.
 """
 
+import math
+
 import torch
 
 from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -86,4 +88,19 @@ class InputValidationStage(PipelineStage):
                 f"Guidance scale must be positive, but got {batch.guidance_scale}"
             )
 
+        # Adjust number of frames based on number of GPUs
+        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
+        orig_latent_num_frames = (batch.num_frames -
+                                  1) // temporal_scale_factor + 1
+        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
+            new_latent_num_frames = math.ceil(
+                orig_latent_num_frames /
+                fastvideo_args.num_gpus) * fastvideo_args.num_gpus
+            new_num_frames = (new_latent_num_frames -
+                              1) * temporal_scale_factor + 1
+            logger.info(
+                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
+                batch.num_frames, new_num_frames, fastvideo_args.num_gpus)
+            batch.num_frames = new_num_frames
+
         return batch

fastvideo/v1/tests/encoders/test_clip_encoder.py

@@ -66,7 +66,7 @@ def test_clip_encoder():
 
     # Load the HuggingFace implementation directly
     # model2 = CLIPTextModel(hf_config)
-    model2 = model2.to(torch.float16)
+    # model2 = model2.to(torch.float16)
     model2 = model2.to(device)
     model2.eval()
 

fastvideo/v1/tests/encoders/test_llama_encoder.py

@@ -62,7 +62,7 @@ def test_llama_encoder():
     model2 = loader.load(TEXT_ENCODER_PATH, "", args)
 
     # Convert to float16 and move to device
-    model2 = model2.to(torch.float16)
+    # model2 = model2.to(torch.float16)
     model2 = model2.to(device)
     model2.eval()
 

fastvideo/v1/tests/encoders/test_t5_encoder.py

@@ -9,7 +9,7 @@ from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.loader.component_loader import TextEncoderLoader
-from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.utils import maybe_download_model, PRECISION_TO_TYPE
 from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.configs.models.encoders import T5Config
 
@@ -33,17 +33,18 @@ def test_t5_encoder():
     print(hf_config)
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.float16
+    precision_str = "fp32"
+    precision = PRECISION_TO_TYPE[precision_str]
     model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to(
         precision).to(device).eval()
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
-    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH, text_encoder_configs=(T5Config(),), device_str="cuda")
+    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH, text_encoder_configs=(T5Config(),), text_encoder_precisions=(precision_str,), device_str="cuda")
     loader = TextEncoderLoader()
     model2 = loader.load(TEXT_ENCODER_PATH, "", args)
 
     # Convert to float16 and move to device
-    model2 = model2.to(precision)
+    # model2 = model2.to(precision)
     model2 = model2.to(device)
     model2.eval()
 

fastvideo/v1/worker/gpu_worker.py

@@ -131,6 +131,8 @@ class Worker:
 
             # Handle regular RPC calls
             if method_name == 'execute_forward':
+                gc.collect()
+                torch.cuda.empty_cache()
                 forward_batch = recv_rpc['kwargs']['forward_batch']
                 fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
                 output_batch = self.execute_forward(forward_batch,