[bugfix]: skip unused output materialization (#1567)
Co-authored-by: Satyam Srivastava <srivastavasatyam53@gmail.com>
This commit is contained in:
Mac Lee
co-authored by
Satyam Srivastava
Satyam Srivastava
parent
dec8b10939
commit
6cbff73687
@@ -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 = []
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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]:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user