From dc72a35f4e5615b26f8e7e08ce5fcd8d3ddf557b Mon Sep 17 00:00:00 2001 From: Boyang ZHONG <80496459+Andzy324@users.noreply.github.com> Date: Thu, 8 Oct 2026 00:06:02 +0200 Subject: [PATCH] [bugfix] Return individual images for batched image API requests (#1917) For n > 1, the image endpoints returned one preview grid instead of n images. The endpoints now run one model batch, save each decoded sample as its own file, and return one URL or base64 item for each image. This applies to generations and edits. A keyword-only return_samples field (default False) is added to OutputConfig, SamplingParam and ForwardBatch. Batched image requests set it and do not build preview frames. The worker keeps the sample tensor for these requests, and non-output ranks clear the field. Defaults and positional arguments do not change. --- .../inference_schema_parity_inventory.yaml | 1 + docs/getting_started/v1_api.md | 8 ++ fastvideo/api/sampling_param.py | 1 + fastvideo/api/schema.py | 1 + fastvideo/entrypoints/openai/image_api.py | 127 +++++++++--------- fastvideo/entrypoints/video_generator.py | 13 +- fastvideo/pipelines/pipeline_batch_info.py | 1 + fastvideo/tests/api/test_parser.py | 1 + ...st_image_generation_route_compatibility.py | 74 +++++++++- .../tests/entrypoints/test_video_generator.py | 45 +++++++ fastvideo/tests/worker/test_gpu_worker.py | 14 +- fastvideo/worker/gpu_worker.py | 7 +- 12 files changed, 213 insertions(+), 80 deletions(-) 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.