diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 173326e12..24ec6bb01 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -565,6 +565,7 @@ surfaces: enable_teacache: request.runtime.enable_teacache save_video: request.output.save_video return_frames: request.output.return_frames + return_samples: request.output.return_samples return_trajectory_latents: request.runtime.return_trajectory_latents return_trajectory_decoded: request.runtime.return_trajectory_decoded continuation_state: request.state diff --git a/docs/getting_started/v1_api.md b/docs/getting_started/v1_api.md index 872c39fa4..f9a222bcb 100644 --- a/docs/getting_started/v1_api.md +++ b/docs/getting_started/v1_api.md @@ -40,3 +40,11 @@ The following two classes `PipelineConfig` and `SamplingParam` are used to confi members: - from_pretrained heading_level: 4 + +### Returning samples without previews + +Set `OutputConfig(return_samples=True, return_frames=False, save_video=False)` +to receive the decoded CPU tensor in `GenerationResult.samples` without building +preview grids or writing media. Pixel samples retain the existing float `[0, 1]` +contract. `return_frames=True` continues to return both samples and preview frames. +The legacy sampling/keyword interface accepts the same `return_samples` option. diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index b4cd37b85..69a7d3894 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -215,6 +215,7 @@ class SamplingParam: # Misc save_video: bool = True return_frames: bool = True + return_samples: bool = field(default=False, kw_only=True) return_trajectory_latents: bool = False # returns all latents for each timestep return_trajectory_decoded: bool = False # returns decoded latents for each timestep diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 18e7a99a2..4663d7c46 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -211,6 +211,7 @@ class OutputConfig: save_video: bool = True return_frames: bool = True return_state: bool = False + return_samples: bool = field(default=False, kw_only=True) @dataclass diff --git a/fastvideo/entrypoints/openai/image_api.py b/fastvideo/entrypoints/openai/image_api.py index eab6b9af0..d04910490 100644 --- a/fastvideo/entrypoints/openai/image_api.py +++ b/fastvideo/entrypoints/openai/image_api.py @@ -1,11 +1,13 @@ # Adapted from SGLang # (https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py) +import asyncio import base64 import os import time import aiofiles +import imageio.v2 as imageio from fastapi import (APIRouter, File, Form, HTTPException, Path, Query, Request, UploadFile) from fastapi.responses import FileResponse @@ -33,6 +35,7 @@ from fastvideo.entrypoints.openai.utils import ( parse_size, save_image_to_path, ) +from fastvideo.utils import pixels_to_uint8 from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -116,6 +119,62 @@ def _build_generation_kwargs( return kwargs +async def _generate_image_files(engine, gen_kwargs: dict) -> list[str]: + output_path = gen_kwargs["output_path"] + count = gen_kwargs["num_videos_per_prompt"] + if count == 1: + await engine.run_serialized(engine.generator.generate_video, **gen_kwargs) + return [output_path] + + # Keep one native model batch, but save separate images instead of its preview grid. + batch_kwargs = dict(gen_kwargs, + save_video=False, + return_frames=False, + return_samples=True, + output_path=os.path.dirname(output_path)) + result = await engine.run_serialized(engine.generator.generate_video, **batch_kwargs) + + def save_images() -> list[str]: + samples = result["samples"] + if samples is None or samples.ndim != 5 or samples.shape[0] != count or samples.shape[2] != 1: + raise RuntimeError("Image generation did not return the requested batch of single-frame samples") + images = pixels_to_uint8(samples)[:, :, 0].permute(0, 2, 3, 1).cpu().numpy() + base, ext = os.path.splitext(output_path) + paths = [] + for index, image in enumerate(images): + path = f"{base}_{index}{ext}" + imageio.imwrite(path, image) + paths.append(path) + return paths + + return await asyncio.to_thread(save_images) + + +async def _build_image_response(request_id: str, + prompt: str, + resp_format: str, + paths: list[str], + elapsed: float, + *, + include_file_path: bool = False) -> ImageResponse: + data = [] + for index, path in enumerate(paths): + image_id = request_id if index == 0 else f"{request_id}_{index}" + item = ImageResponseData(revised_prompt=prompt) + if resp_format == "b64_json": + if not os.path.exists(path): + raise HTTPException(status_code=500, detail="Image was not saved to disk") + async with aiofiles.open(path, "rb") as f: + item.b64_json = base64.b64encode(await f.read()).decode("utf-8") + else: + item.url = f"/v1/images/{image_id}/content" + if include_file_path or resp_format == "url": + item.file_path = os.path.abspath(path) + data.append(item) + await IMAGE_STORE.upsert(image_id, {"id": image_id, "created_at": int(time.time()), "file_path": path}) + return ImageResponse(id=request_id, data=data, inference_time_s=elapsed) + + @router.post("/generations", response_model=ImageResponse) @router.post("", response_model=ImageResponse) async def generations(request: ImageGenerationsRequest): @@ -142,43 +201,13 @@ async def generations(request: ImageGenerationsRequest): start = time.perf_counter() try: - await engine.run_serialized(engine.generator.generate_video, **gen_kwargs) + paths = await _generate_image_files(engine, gen_kwargs) except Exception as e: logger.error("Image generation failed: %s", e) raise HTTPException(status_code=500, detail=str(e)) from None elapsed = time.perf_counter() - start - save_file_path = gen_kwargs["output_path"] - - if resp_format == "b64_json": - if not os.path.exists(save_file_path): - raise HTTPException(status_code=500, detail="Image was not saved to disk") - async with aiofiles.open(save_file_path, "rb") as f: - b64_data = base64.b64encode(await f.read()).decode("utf-8") - data = [ImageResponseData(b64_json=b64_data, revised_prompt=request.prompt)] - else: - data = [ - ImageResponseData( - url=f"/v1/images/{request_id}/content", - revised_prompt=request.prompt, - file_path=os.path.abspath(save_file_path), - ) - ] - - await IMAGE_STORE.upsert( - request_id, - { - "id": request_id, - "created_at": int(time.time()), - "file_path": save_file_path, - }, - ) - - return ImageResponse( - id=request_id, - data=data, - inference_time_s=elapsed, - ) + return await _build_image_response(request_id, request.prompt, resp_format, paths, elapsed) @router.post("/edits", response_model=ImageResponse) @@ -254,43 +283,13 @@ async def edits( start = time.perf_counter() try: - await engine.run_serialized(engine.generator.generate_video, **gen_kwargs) + paths = await _generate_image_files(engine, gen_kwargs) except Exception as e: logger.error("Image edit failed: %s", e) raise HTTPException(status_code=500, detail=str(e)) from None elapsed = time.perf_counter() - start - save_file_path = gen_kwargs["output_path"] - - if resp_format == "b64_json": - async with aiofiles.open(save_file_path, "rb") as f: - b64_data = base64.b64encode(await f.read()).decode("utf-8") - data = [ - ImageResponseData( - b64_json=b64_data, - revised_prompt=prompt, - file_path=os.path.abspath(save_file_path), - ) - ] - else: - data = [ - ImageResponseData( - url=f"/v1/images/{request_id}/content", - revised_prompt=prompt, - file_path=os.path.abspath(save_file_path), - ) - ] - - await IMAGE_STORE.upsert( - request_id, - { - "id": request_id, - "created_at": int(time.time()), - "file_path": save_file_path, - }, - ) - - return ImageResponse(id=request_id, data=data, inference_time_s=elapsed) + return await _build_image_response(request_id, prompt, resp_format, paths, elapsed, include_file_path=True) @router.get("/{image_id}/content") diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 7dca36991..64a40d671 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -901,6 +901,7 @@ class VideoGenerator: if not self.executor.is_output_rank: batch.save_video = False batch.return_frames = False + batch.return_samples = False batch.return_trajectory_latents = False batch.return_trajectory_decoded = False batch.return_continuation_state = False @@ -929,17 +930,17 @@ class VideoGenerator: is_latent_output = fastvideo_args.output_type == "latent" needs_frame_output = batch.return_frames or (batch.save_video and not is_latent_output) # A populated ``samples`` has exactly one consumer — the result - # dict (``"samples": samples if batch.return_frames else None``). + # dict (``"samples": samples if needs_samples_out else None``). # Post-decode frame building reads ``output_batch.output`` # directly (the GPU ``vid_u8`` path), not ``samples``. So when - # ``return_frames=False`` the pinned fp32 alloc + D->H copy are - # dead weight — the CLI generate flow (``save_video=True``, + # neither frames nor samples are requested, the pinned fp32 allocation + # and D->H copy are dead weight — the CLI generate flow (``save_video=True``, # ``return_frames=False``) hits this on every call. # ``output_type == "latent"`` keeps its existing branch (shape # mismatch falls through to ``.cpu()`` below) for callers that # *do* ask for the latent samples via ``return_frames=True``. # ``skip_pixel_prealloc`` also gates the slow-path warning. - needs_samples_out = batch.return_frames + needs_samples_out = batch.return_frames or batch.return_samples # ``output_shape_from_input`` pipelines (Kandinsky6 SR) decide the output geometry from the input video, so the # request's height / width / num_frames say nothing about the buffer to pre-allocate. output_shape_from_input = bool(getattr(fastvideo_args.pipeline_config, "output_shape_from_input", False)) @@ -978,7 +979,7 @@ class VideoGenerator: audio_only = bool(output_batch.extra.get("audio_only")) if not needs_samples_out: # Nothing downstream reads ``samples`` (the result dict - # returns None when ``return_frames=False``); keep the empty + # returns None when neither output is requested); keep the empty # placeholder allocated above and skip the fp32 D->H copy # entirely. pass @@ -1132,7 +1133,7 @@ class VideoGenerator: result: dict[str, Any] = { "prompts": prompt, - "samples": samples if batch.return_frames else None, + "samples": samples if needs_samples_out else None, "frames": frames if batch.return_frames else None, # Audio is the primary output for audio workloads — return it # whenever the pipeline produced one, regardless of diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 00a368799..0749eb10e 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -287,6 +287,7 @@ class ForwardBatch: # Misc save_video: bool = True return_frames: bool = False + return_samples: bool = field(default=False, kw_only=True) is_cfg_negative: bool = False diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index 990429ef9..0c2df14fc 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -260,6 +260,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "save_video": True, "return_frames": True, "return_state": False, + "return_samples": False, }, "stage_overrides": {}, "state": None, diff --git a/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py b/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py index adbf63bb3..c806a0530 100644 --- a/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py +++ b/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py @@ -2,6 +2,10 @@ """Image generation route compatibility with a mocked generator.""" import base64 +from io import BytesIO + +from PIL import Image +import torch from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -10,8 +14,12 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from fastvideo.entrypoints.video_generator import VideoGenerator +from fastvideo.api.sampling_param import SamplingParam from fastvideo.fastvideo_args import WorkloadType +from fastvideo.tests.entrypoints.test_video_generator import ( + _single_video_args, _single_video_generator, _single_video_output_batch, _small_sampling_param, +) +from fastvideo.entrypoints.video_generator import VideoGenerator from fastvideo.entrypoints.openai import image_api from fastvideo.entrypoints.openai.stores import AsyncDictStore @@ -69,3 +77,67 @@ def test_image_generation_routes(path, response_format, output_format, monkeypat assert generate.call_args.kwargs["prompt"] == "a cat" assert generate.call_args.kwargs["num_frames"] == 1 engine.run_serialized.assert_awaited_once() + + +@pytest.mark.parametrize("path", ["/v1/images/generations", "/v1/images", "/v1/images/edits"]) +@pytest.mark.parametrize("response_format", ["url", "b64_json"]) +@pytest.mark.parametrize("output_format", ["png", "jpeg", "webp"]) +def test_image_batch_returns_individual_images(path, response_format, output_format, monkeypatch, tmp_path): + monkeypatch.setattr("fastvideo.entrypoints.video_generator.pixels_to_uint8", + Mock(side_effect=AssertionError("Unexpected batch preview conversion"))) + monkeypatch.setattr("fastvideo.entrypoints.video_generator.torchvision.utils.make_grid", + Mock(side_effect=AssertionError("Unexpected batch preview grid"))) + samples = torch.zeros(2, 3, 1, 16, 16) + samples[0, 0] = 1 + samples[1, 2] = 1 + + args = _single_video_args() + args.model_path = "test-model" + args.prompt_txt = None + args.workload_type = WorkloadType.from_string("t2i") + generator = _single_video_generator(_single_video_output_batch(samples), args) + monkeypatch.setattr(SamplingParam, "from_pretrained", lambda model_path: _small_sampling_param()) + + def generate_images(**kwargs): + return generator.generate_video(**kwargs) + + generate = Mock(side_effect=generate_images) + + async def run_serialized(fn, **kwargs): + return fn(**kwargs) + + engine = SimpleNamespace(generator=SimpleNamespace(generate_video=generate), + run_serialized=AsyncMock(side_effect=run_serialized)) + monkeypatch.setattr(image_api, "get_serving_engine", lambda: engine) + monkeypatch.setattr(image_api, "get_output_dir", lambda: str(tmp_path)) + monkeypatch.setattr(image_api, "IMAGE_STORE", AsyncDictStore()) + app = FastAPI() + app.include_router(image_api.router) + payload = {"prompt": "a cat", "n": 2, "response_format": response_format, "output_format": output_format} + with TestClient(app) as client: + if path.endswith("edits"): + response = client.post(path, data=payload, files={"image": ("cat.png", b"input", "image/png")}) + else: + response = client.post(path, json=payload) + assert response.status_code == 200, response.text + data = response.json()["data"] + assert len(data) == 2 + for index, item in enumerate(data): + if response_format == "url": + content = client.get(item["url"]) + assert content.status_code == 200 + content = content.content + else: + content = base64.b64decode(item["b64_json"]) + with Image.open(BytesIO(content)) as image: + assert image.size == (16, 16) + assert image.format == {"png": "PNG", "jpeg": "JPEG", "webp": "WEBP"}[output_format] + color = image.convert("RGB").getpixel((8, 8)) + assert color[0 if index == 0 else 2] >= 250 + assert color[2 if index == 0 else 0] <= 5 + # Keep the model's native batch, rather than making a separate inference call per image. + generate.assert_called_once() + assert generate.call_args.kwargs["num_videos_per_prompt"] == 2 + assert generate.call_args.kwargs["save_video"] is False + assert generate.call_args.kwargs["return_frames"] is False + assert generate.call_args.kwargs["return_samples"] is True diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 55287823e..8143305fe 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -663,6 +663,7 @@ def test_generate_single_video_non_output_rank_suppresses_all_user_payloads(tmp_ broadcast_from_output_rank=lambda value: value, ) sampling = _small_sampling_param(save_video=True, return_frames=True) + sampling.return_samples = True sampling.return_trajectory_latents = True sampling.return_trajectory_decoded = True @@ -676,6 +677,7 @@ def test_generate_single_video_non_output_rank_suppresses_all_user_payloads(tmp_ batch = captured["batch"] assert not batch.save_video assert not batch.return_frames + assert not batch.return_samples assert not batch.return_trajectory_latents assert not batch.return_trajectory_decoded assert result["samples"] is None @@ -1230,6 +1232,49 @@ def test_ffmpeg_pipe_preserves_video_duration(tmp_path, env_overrides, audio_sec assert float(sound["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps) +@pytest.mark.parametrize("dtype", [torch.float32, torch.uint8]) +def test_return_samples_only_skips_preview_processing(monkeypatch, tmp_path, dtype): + output = torch.ones((2, 3, 1, 16, 16), dtype=dtype) + generator = _single_video_generator(_single_video_output_batch(output), _single_video_args()) + monkeypatch.setattr("fastvideo.entrypoints.video_generator.pixels_to_uint8", + lambda *args: pytest.fail("Samples-only requests must not quantize a preview")) + monkeypatch.setattr("fastvideo.entrypoints.video_generator.torchvision.utils.make_grid", + lambda *args, **kwargs: pytest.fail("Samples-only requests must not build a preview grid")) + sampling = _small_sampling_param() + sampling.num_frames = 1 + sampling.num_videos_per_prompt = 2 + sampling.return_samples = True + result = generator._generate_single_video("samples only", sampling, + output_path=str(tmp_path / "unused.mp4")) + torch.testing.assert_close(result["samples"], output.float() / (255 if dtype == torch.uint8 else 1)) + assert result["frames"] is None + assert result["video_path"] is None + assert result["size"] == (16, 16, 1) + + +@pytest.mark.parametrize("return_frames", [False, True]) +def test_input_sized_output_skips_request_pixel_preallocation(monkeypatch, tmp_path, return_frames): + output = torch.full((1, 3, 1, 24, 32), 0.5) + args = _single_video_args() + args.pipeline_config.output_shape_from_input = True + generator = _single_video_generator(_single_video_output_batch(output), args) + monkeypatch.setattr( + video_generator_module, "allocate_cpu_tensor_with_pin_fallback", + lambda *args, **kwargs: pytest.fail("Input-sized outputs must not allocate request-sized pixel buffers")) + sampling = _small_sampling_param(return_frames=return_frames) + sampling.num_frames = 1 + sampling.return_samples = True + result = generator._generate_single_video("input-sized output", sampling, + output_path=str(tmp_path / "unused.mp4")) + torch.testing.assert_close(result["samples"], output) + assert result["size"] == (24, 32, 1) + if return_frames: + assert len(result["frames"]) == 1 + else: + assert result["frames"] is None + assert result["video_path"] is None + + @pytest.mark.parametrize("extension,format_name", [(".png", "PNG"), (".jpg", "JPEG"), (".jpeg", "JPEG"), (".webp", "WEBP")]) def test_image_output_path_preserves_requested_format(tmp_path, extension, format_name): from PIL import Image diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py index c09868a81..07f550e62 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/fastvideo/tests/worker/test_gpu_worker.py @@ -172,18 +172,20 @@ def test_execute_forward_drops_save_only_audio_placeholder(): @pytest.mark.parametrize( - ("save_video", "return_frames"), + ("save_video", "return_frames", "return_samples"), [ - (True, False), - (False, True), - (True, True), + (True, False, False), + (False, True, False), + (True, True, False), + (False, False, True), ], ) -def test_execute_forward_preserves_requested_output(save_video, return_frames): +def test_execute_forward_preserves_requested_output(save_video, return_frames, return_samples): 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) + request_batch = ForwardBatch(data_type="video", save_video=save_video, return_frames=return_frames, + return_samples=return_samples) result = worker.execute_forward(request_batch, FastVideoArgs(model_path="test")) diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 84edbb121..ed86afa5b 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -168,13 +168,14 @@ class Worker: if not self.fastvideo_args.is_output_rank: forward_batch.save_video = False forward_batch.return_frames = False + forward_batch.return_samples = False forward_batch.return_trajectory_latents = False forward_batch.return_trajectory_decoded = False forward_batch.return_continuation_state = False 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")) + needs_output = forward_batch.return_frames or forward_batch.return_samples 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.