[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.

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

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.

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
 

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

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")

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

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
 

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,

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

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

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"))
 

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.