diff --git a/fastvideo/v1/configs/pipelines/base.py b/fastvideo/v1/configs/pipelines/base.py index 4ea24adf8..c7a1e2db0 100644 --- a/fastvideo/v1/configs/pipelines/base.py +++ b/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 diff --git a/fastvideo/v1/examples/inference/basic/basic.py b/fastvideo/v1/examples/inference/basic/basic.py index 300971260..7a767c409 100644 --- a/fastvideo/v1/examples/inference/basic/basic.py +++ b/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 diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 4b9afbad1..9226ade7d 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/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 diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py index c00f4b24c..80ad77eb6 100644 --- a/fastvideo/v1/models/loader/component_loader.py +++ b/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) diff --git a/fastvideo/v1/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py index 663fd5df9..c6dc25a58 100644 --- a/fastvideo/v1/pipelines/stages/decoding.py +++ b/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 diff --git a/fastvideo/v1/pipelines/stages/input_validation.py b/fastvideo/v1/pipelines/stages/input_validation.py index 036c7a0c6..61b7e6bb8 100644 --- a/fastvideo/v1/pipelines/stages/input_validation.py +++ b/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 diff --git a/fastvideo/v1/tests/encoders/test_clip_encoder.py b/fastvideo/v1/tests/encoders/test_clip_encoder.py index 83db223c0..3b0372615 100644 --- a/fastvideo/v1/tests/encoders/test_clip_encoder.py +++ b/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() diff --git a/fastvideo/v1/tests/encoders/test_llama_encoder.py b/fastvideo/v1/tests/encoders/test_llama_encoder.py index 81cb9d7a2..362df30b7 100644 --- a/fastvideo/v1/tests/encoders/test_llama_encoder.py +++ b/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() diff --git a/fastvideo/v1/tests/encoders/test_t5_encoder.py b/fastvideo/v1/tests/encoders/test_t5_encoder.py index b51bef93a..6bc1fda32 100644 --- a/fastvideo/v1/tests/encoders/test_t5_encoder.py +++ b/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() diff --git a/fastvideo/v1/worker/gpu_worker.py b/fastvideo/v1/worker/gpu_worker.py index d800af0ad..5ac5f5f07 100644 --- a/fastvideo/v1/worker/gpu_worker.py +++ b/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,