diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index bfc163353..40203dd00 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -748,13 +748,16 @@ class VideoGenerator: thread = threading.Thread(target=execute_forward_thread) thread.start() latent_batch_size = _infer_latent_batch_size(batch) + is_latent_output = fastvideo_args.output_type == "latent" + needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output) + needs_samples_buffer = batch.return_frames or needs_frame_output # When ``output_type == "latent"`` the forward output has latent # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``) # rather than the pre-allocation's pixel shape. Skip the pinned - # ~50 MB buffer entirely; we always fall through to the - # ``samples = output_batch.output.cpu()`` branch below in that - # mode. ``skip_pixel_prealloc`` also gates the slow-path warning. - skip_pixel_prealloc = fastvideo_args.output_type == "latent" + # ~50 MB buffer entirely. Also skip it for metadata-only calls; + # neither the result nor save path will consume the decoded tensor. + # ``skip_pixel_prealloc`` also gates the slow-path warning. + skip_pixel_prealloc = is_latent_output or not needs_samples_buffer if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: @@ -773,7 +776,16 @@ class VideoGenerator: raise RuntimeError("Forward execution returned no output tensor. " "This usually means the executor/pipeline failed earlier.") - if output_batch.output.shape == samples.shape: + audio_only = bool(output_batch.extra.get("audio_only")) + if not needs_samples_buffer or (audio_only and not batch.return_frames): + # Metadata-only/audio-only request: keep the empty placeholder and + # avoid the decoded tensor D->H copy. + pass + elif audio_only: + # Audio-only return-frames requests expose the small placeholder + # without warning about expected video shape. + samples = output_batch.output.cpu() + elif output_batch.output.shape == samples.shape: samples.copy_(output_batch.output) else: if not skip_pixel_prealloc: @@ -796,13 +808,12 @@ class VideoGenerator: # 2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder # no caller will use; skip the grid loop and save a `.wav`. # 3. Pixel video / image — the historical happy path. - is_latent_output = fastvideo_args.output_type == "latent" - audio_only = bool(output_batch.extra.get("audio_only")) - postprocess_start = time.perf_counter() frames: list[np.ndarray] | None if is_latent_output or audio_only: - frames = None if is_latent_output else [] + frames = [] if audio_only and batch.return_frames else None + elif not needs_frame_output: + frames = None else: videos = rearrange(samples, "b c t h w -> t b c h w") frames = [] diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 2789f071d..ab4cdd59d 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -3,7 +3,9 @@ from types import SimpleNamespace import warnings import pytest +import torch +import fastvideo.entrypoints.video_generator as video_generator_module from fastvideo.api import ( GenerationRequest, GenerationResult, @@ -15,6 +17,9 @@ from fastvideo.api import ( from fastvideo.api.sampling_param import SamplingParam from fastvideo.entrypoints.video_generator import VideoGenerator from fastvideo.fastvideo_args import WorkloadType +from fastvideo.pipelines import ForwardBatch +from fastvideo.worker.gpu_worker import Worker +from fastvideo.worker.ray_distributed_executor import RayDistributedExecutor def _new_video_generator() -> VideoGenerator: @@ -82,6 +87,60 @@ def _patch_sampling_param_from_pretrained(monkeypatch): monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained)) +class _NoCpuMaterializationOutput: + + @property + def shape(self): + raise AssertionError("metadata-only output path should not inspect output shape") + + def cpu(self): + raise AssertionError("metadata-only output path should not move output to CPU") + + def __mul__(self, other): + raise AssertionError("metadata-only output path should not build frames") + + +def _single_video_args(output_type="pil"): + return SimpleNamespace( + output_type=output_type, + pin_cpu_memory=False, + VSA_sparsity=0.0, + pipeline_config=SimpleNamespace(flow_shift=1.0, embedded_cfg_scale=1.0), + workload_type=SimpleNamespace(value="t2v"), + ) + + +def _small_sampling_param(*, save_video=False, return_frames=False): + return SamplingParam( + num_frames=2, + height=16, + width=16, + fps=8, + num_inference_steps=1, + save_video=save_video, + return_frames=return_frames, + ) + + +def _single_video_output_batch(output, *, extra=None): + return SimpleNamespace( + output=output, + extra=extra or {}, + logging_info=None, + trajectory_latents=None, + trajectory_timesteps=None, + trajectory_decoded=None, + ) + + +def _single_video_generator(output_batch, fastvideo_args): + generator = _new_video_generator() + generator.fastvideo_args = fastvideo_args + generator.executor = SimpleNamespace(execute_forward=lambda batch, args: output_batch) + generator.config = None + return generator + + def test_prepare_output_path_file_sanitization(tmp_path): vg = _new_video_generator() target_dir = tmp_path / "dir" @@ -151,6 +210,220 @@ def test_prepare_output_path_empty_prompt_fallback(tmp_path): assert os.path.basename(result) == "output.mp4" +def test_generate_single_video_metadata_only_skips_output_materialization(monkeypatch, tmp_path): + output_batch = _single_video_output_batch( + _NoCpuMaterializationOutput(), + extra={"peak_memory_mb": 42.0}, + ) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + empty_calls = [] + real_empty = video_generator_module.torch.empty + + def record_empty(*args, **kwargs): + empty_calls.append((args, kwargs)) + return real_empty(*args, **kwargs) + + def fail_make_grid(*args, **kwargs): + raise AssertionError("metadata-only output path should not build frame grids") + + monkeypatch.setattr(video_generator_module.torch, "empty", record_empty) + monkeypatch.setattr(video_generator_module.torchvision.utils, "make_grid", fail_make_grid) + + result = generator._generate_single_video( + prompt="metadata only", + sampling_param=_small_sampling_param(save_video=False, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["video_path"] is None + assert result["peak_memory_mb"] == 42.0 + assert empty_calls == [((0,), {"device": "cpu"})] + + +def test_generate_single_video_return_frames_still_materializes_output(tmp_path): + output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + + result = generator._generate_single_video( + prompt="return frames", + sampling_param=_small_sampling_param(save_video=False, return_frames=True), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + torch.testing.assert_close(result["samples"], output) + assert len(result["frames"]) == 2 + assert result["video_path"] is None + + +def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_path): + output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + saved = {} + + def fake_mimsave(path, frames, *, fps, format): + saved["path"] = path + saved["frame_count"] = len(frames) + saved["fps"] = fps + saved["format"] = format + + monkeypatch.setattr(video_generator_module.imageio, "mimsave", fake_mimsave) + + output_path = str(tmp_path / "saved.mp4") + result = generator._generate_single_video( + prompt="save only", + sampling_param=_small_sampling_param(save_video=True, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=output_path, + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["video_path"] == output_path + assert saved == { + "path": output_path, + "frame_count": 2, + "fps": 8, + "format": "mp4", + } + + +def test_generate_single_video_audio_only_metadata_returns_audio_without_frames(tmp_path): + audio = torch.zeros((16,), dtype=torch.float32) + output_batch = _single_video_output_batch( + _NoCpuMaterializationOutput(), + extra={ + "audio_only": True, + "audio": audio, + "audio_sample_rate": 44100, + }, + ) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + + result = generator._generate_single_video( + prompt="audio only", + sampling_param=_small_sampling_param(save_video=False, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["audio"] is audio + assert result["audio_sample_rate"] == 44100 + + +def test_generate_single_video_audio_only_save_skips_placeholder_materialization(tmp_path): + audio = torch.zeros((16,), dtype=torch.float32) + output_batch = _single_video_output_batch( + _NoCpuMaterializationOutput(), + extra={ + "audio_only": True, + "audio": audio, + "audio_sample_rate": 44100, + }, + ) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + output_path = str(tmp_path / "audio.mp4") + + result = generator._generate_single_video( + prompt="audio only", + sampling_param=_small_sampling_param(save_video=True, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=output_path, + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["audio"] is audio + assert result["audio_sample_rate"] == 44100 + assert result["video_path"] == str(tmp_path / "audio.wav") + + +def test_generate_single_video_ray_audio_only_save_preserves_worker_metadata(monkeypatch, tmp_path): + audio = torch.zeros((16,), dtype=torch.float32) + worker_output = ForwardBatch( + data_type="audio", + output=torch.ones((1, 3, 1, 8, 8)), + extra={ + "audio_only": True, + "audio": audio, + "audio_sample_rate": 44100, + }, + ) + worker = Worker.__new__(Worker) + worker.fastvideo_args = SimpleNamespace() + worker.pipeline = SimpleNamespace(forward=lambda batch, args: worker_output) + + monkeypatch.setattr(RayDistributedExecutor, "__abstractmethods__", frozenset()) + executor = RayDistributedExecutor.__new__(RayDistributedExecutor) + executor.shutdown = lambda: None + + def collective_rpc(method, *, kwargs, **_): + assert method == "execute_forward" + return [worker.execute_forward(**kwargs)] + + executor.collective_rpc = collective_rpc + fastvideo_args = _single_video_args() + generator = _new_video_generator() + generator.fastvideo_args = fastvideo_args + generator.executor = executor + generator.config = None + written = {} + + def fake_write_pcm_wav(path, samples, sample_rate): + written.update(path=path, samples=samples, sample_rate=sample_rate) + + monkeypatch.setattr(generator, "_write_pcm_wav", fake_write_pcm_wav) + sampling_param = _small_sampling_param(save_video=True, return_frames=False) + sampling_param.data_type = "audio" + + result = generator._generate_single_video( + prompt="audio only", + sampling_param=sampling_param, + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "audio.mp4"), + ) + + assert worker_output.output is not None + assert worker_output.output.numel() == 0 + assert result["samples"] is None + assert result["frames"] is None + assert result["audio"] is audio + assert result["audio_sample_rate"] == 44100 + assert result["video_path"] == str(tmp_path / "audio.wav") + assert written["path"] == str(tmp_path / "audio.wav") + assert written["samples"] is audio + assert written["sample_rate"] == 44100 + + +def test_generate_single_video_latent_metadata_skips_cpu_materialization(tmp_path): + output_batch = _single_video_output_batch(_NoCpuMaterializationOutput()) + fastvideo_args = _single_video_args(output_type="latent") + generator = _single_video_generator(output_batch, fastvideo_args) + + result = generator._generate_single_video( + prompt="latent metadata", + sampling_param=_small_sampling_param(save_video=False, return_frames=False), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + assert result["samples"] is None + assert result["frames"] is None + assert result["video_path"] is None + + def test_from_config_normalizes_and_translates(monkeypatch): captured = _patch_from_fastvideo_args(monkeypatch) _patch_fastvideo_args_from_kwargs(monkeypatch) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 80f7cfd5d..0542874e2 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -286,7 +286,7 @@ def run_self_forcing_tests(): @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret]) def run_unit_test(): run_test( - "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs" + "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ ./fastvideo/tests/training/test_trackers.py ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs" ) diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py new file mode 100644 index 000000000..2c278d767 --- /dev/null +++ b/fastvideo/tests/worker/test_gpu_worker.py @@ -0,0 +1,83 @@ +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines import ForwardBatch +from fastvideo.worker.gpu_worker import Worker + + +def _worker_returning(output_batch: ForwardBatch) -> Worker: + worker = Worker.__new__(Worker) + worker.fastvideo_args = SimpleNamespace() + worker.pipeline = SimpleNamespace(forward=lambda batch, args: output_batch) + return worker + + +def test_execute_forward_drops_metadata_only_output_before_transport(): + output = torch.ones((1, 3, 2, 4, 4)) + output_batch = ForwardBatch(data_type="video", output=output) + worker = _worker_returning(output_batch) + request_batch = ForwardBatch(data_type="video", save_video=False, return_frames=False) + + result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test")) + + assert result.output is not None + assert result.output.device.type == "cpu" + assert result.output.numel() == 0 + + +def test_execute_forward_preserves_missing_metadata_only_output(): + output_batch = ForwardBatch(data_type="video", output=None) + worker = _worker_returning(output_batch) + request_batch = ForwardBatch(data_type="video", save_video=False, return_frames=False) + + result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test")) + + assert result.output is None + + +def test_execute_forward_drops_save_only_latent_output(): + output = torch.ones((1, 16, 1, 2, 2)) + output_batch = ForwardBatch(data_type="video", output=output) + worker = _worker_returning(output_batch) + request_batch = ForwardBatch(data_type="video", save_video=True, return_frames=False) + + result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test", output_type="latent")) + + assert result.output is not None + assert result.output.device.type == "cpu" + assert result.output.numel() == 0 + + +def test_execute_forward_drops_save_only_audio_placeholder(): + output = torch.ones((1, 3, 1, 8, 8)) + output_batch = ForwardBatch(data_type="audio", output=output, extra={"audio_only": True}) + worker = _worker_returning(output_batch) + request_batch = ForwardBatch(data_type="audio", save_video=True, return_frames=False) + + result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test")) + + assert result.output is not None + assert result.output.device.type == "cpu" + assert result.output.numel() == 0 + + +@pytest.mark.parametrize( + ("save_video", "return_frames"), + [ + (True, False), + (False, True), + (True, True), + ], +) +def test_execute_forward_preserves_requested_output(save_video, return_frames): + output = torch.ones((1, 3, 2, 4, 4)) + output_batch = ForwardBatch(data_type="video", output=output) + worker = _worker_returning(output_batch) + request_batch = ForwardBatch(data_type="video", save_video=save_video, return_frames=return_frames) + + result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test")) + + assert result.output is output diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index b440ba2ce..b4f8e6227 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -73,6 +73,13 @@ class Worker: def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args) + needs_output = forward_batch.return_frames or (forward_batch.save_video + and fastvideo_args.output_type != "latent" + and not output_batch.extra.get("audio_only")) + if output_batch.output is not None and not needs_output: + # Drop the decoded tensor before multiprocessing or Ray transports + # the worker result back to the generator. + output_batch.output = torch.empty(0, device="cpu") return cast(ForwardBatch, output_batch) def shutdown(self) -> dict[str, Any]: diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index d12f735bb..e276bb90c 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -307,6 +307,7 @@ class RayDistributedExecutor(Executor): data_type=forward_batch.data_type, output=output, logging_info=logging_info, + extra=responses[0].extra, trajectory_latents=responses[0].trajectory_latents, trajectory_timesteps=responses[0].trajectory_timesteps, )