diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index d586d9215..9b82ed374 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -9,6 +9,7 @@ diffusion models. import math import os import re +import threading import time from copy import deepcopy from typing import Any @@ -31,6 +32,21 @@ from fastvideo.worker.executor import Executor logger = init_logger(__name__) +def _infer_latent_batch_size(batch: ForwardBatch) -> int: + if isinstance(batch.prompt, list): + latent_batch_size = len(batch.prompt) + elif batch.prompt is not None: + latent_batch_size = 1 + elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0: + latent_batch_size = batch.prompt_embeds[0].shape[0] + else: + raise ValueError( + "Cannot infer batch size from batch; no prompt or prompt_embeds found" + ) + latent_batch_size *= batch.num_videos_per_prompt + return latent_batch_size + + class VideoGenerator: """ A unified class for generating videos using diffusion models. @@ -372,8 +388,31 @@ class VideoGenerator: # Run inference start_time = time.perf_counter() - output_batch = self.executor.execute_forward(batch, fastvideo_args) - samples = output_batch.output + + # Execute forward pass in a new thread for non-blocking tensor allocation + result_container = {} + + def execute_forward_thread(): + result_container['output_batch'] = self.executor.execute_forward( + batch, fastvideo_args) + + thread = threading.Thread(target=execute_forward_thread) + thread.start() + latent_batch_size = _infer_latent_batch_size(batch) + samples = torch.empty((latent_batch_size, 3, sampling_param.num_frames, + sampling_param.height, sampling_param.width), + device='cpu', + pin_memory=fastvideo_args.pin_cpu_memory) + thread.join() + + output_batch = result_container['output_batch'] + if output_batch.output.shape == samples.shape: + samples.copy_(output_batch.output) + else: + logger.warning( + "Output shape %s does not match expected shape %s; use slow path", + output_batch.output.shape, samples.shape) + samples = output_batch.output.cpu() logging_info = output_batch.logging_info gen_time = time.perf_counter() - start_time diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 7450ff226..267b3bf27 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -242,8 +242,8 @@ class DecodingStage(PipelineStage): decoded_frames = self.decode(cur_latent, fastvideo_args) batch.trajectory_decoded.append(decoded_frames.cpu().float()) - # Convert to CPU float32 for compatibility - frames = frames.cpu().float() + # Convert to float32 for compatibility + frames = frames.to(torch.float32) # Crop padding if this is a LongCat refinement if hasattr(batch, 'num_cond_frames_added') and hasattr( diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index c9ba23588..2fd735dc9 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -1267,7 +1267,7 @@ class DistillationPipeline(TrainingPipeline): with torch.no_grad(): output_batch = self.validation_pipeline.forward( batch, training_args) - samples = output_batch.output + samples = output_batch.output.cpu() if self.rank_in_sp_group != 0: continue diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 07c09b005..7d89c3b96 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -857,7 +857,7 @@ class TrainingPipeline(LoRAPipeline, ABC): # Run validation inference output_batch = self.validation_pipeline.forward( batch, training_args) - samples = output_batch.output + samples = output_batch.output.cpu() if self.rank_in_sp_group != 0: continue diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index cc09c5898..22047f9ba 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -650,8 +650,10 @@ class WorkerMultiprocProc: logging_info = None if envs.FASTVIDEO_STAGE_LOGGING: logging_info = output_batch.logging_info + # result tensor shared by CUDA IPC to avoid serialization overhead + result = output_batch.output self.pipe.send({ - "output_batch": output_batch.output.cpu(), + "output_batch": result, "logging_info": logging_info, "extra": output_batch.extra, })