[bugfix] Preserve image output formats (#1915)

The OpenAI image endpoint defaults to a .jpg output path, but the shared
output-path helper treated every non-PNG image path as a directory, so JPEG and
WebP requests were written somewhere other than where the endpoint read them.
Recognized PNG/JPEG/WebP extensions are now preserved, and existing directories
are checked first so a directory named e.g. renders.jpg still receives its
output inside it. Collision handling and filename sanitization are unchanged.

fastvideo/entrypoints/video_generator.py

@@ -747,9 +747,8 @@ class VideoGenerator:
     ) -> str:
         """Build a unique, sanitized output file path.
 
-        The file extension is chosen automatically based on the workload type:
-        ``.png`` for image workloads (``t2i``, ``i2i``, …) and ``.mp4`` for
-        video workloads.
+        Image workloads preserve explicit PNG, JPEG, or WebP extensions and
+        otherwise default to ``.png``. Video workloads use ``.mp4``.
 
         - If ``output_path`` already carries the correct extension, treat it
           as a file path.
@@ -759,8 +758,12 @@ class VideoGenerator:
           warning is logged.
         - If the target path already exists, a numeric suffix is appended.
         """
+        base_path, extension = os.path.splitext(output_path)
+        extension_lower = extension.lower()
+        is_directory = os.path.isdir(output_path)
         if self._is_image_workload():
-            target_ext = ".png"
+            target_ext = extension_lower if not is_directory and extension_lower in {".png", ".jpg", ".jpeg", ".webp"
+                                                                                     } else ".png"
         elif self._is_audio_workload():
             target_ext = ".wav"
         else:
@@ -773,10 +776,7 @@ class VideoGenerator:
             sanitized = re.sub(r'\s+', ' ', sanitized)
             return sanitized or "output"
 
-        base_path, extension = os.path.splitext(output_path)
-        extension_lower = extension.lower()
-
-        if extension_lower == target_ext:
+        if extension_lower == target_ext and not is_directory:
             output_dir = os.path.dirname(output_path)
             base_name = os.path.basename(base_path)  # filename without extension
             sanitized_base = _sanitize_filename_component(base_name)

fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py

@@ -10,6 +10,8 @@ import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
+from fastvideo.entrypoints.video_generator import VideoGenerator
+from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.entrypoints.openai import image_api
 from fastvideo.entrypoints.openai.stores import AsyncDictStore
 
@@ -22,9 +24,17 @@ from fastvideo.entrypoints.openai.stores import AsyncDictStore
     pytest.param("b64_json", id="b64-json"),
     pytest.param("url", id="url"),
 ])
-def test_image_generation_routes(path, response_format, monkeypatch, tmp_path):
+@pytest.mark.parametrize("output_format", [None, "png", "jpeg", "webp"])
+def test_image_generation_routes(path, response_format, output_format, monkeypatch, tmp_path):
     image_bytes = b"test-image-content"
-    generate = Mock(side_effect=lambda **kwargs: Path(kwargs["output_path"]).write_bytes(image_bytes))
+    generator = VideoGenerator.__new__(VideoGenerator)
+    generator.fastvideo_args = SimpleNamespace(workload_type=WorkloadType.from_string("t2i"))
+
+    def save_image(**kwargs):
+        output_path = generator._prepare_output_path(kwargs["output_path"], kwargs["prompt"])
+        Path(output_path).write_bytes(image_bytes)
+
+    generate = Mock(side_effect=save_image)
 
     async def run_serialized(fn, **kwargs):
         return fn(**kwargs)
@@ -40,7 +50,10 @@ def test_image_generation_routes(path, response_format, monkeypatch, tmp_path):
     app.include_router(image_api.router)
 
     with TestClient(app) as client:
-        response = client.post(path, json={"prompt": "a cat", "response_format": response_format})
+        request = {"prompt": "a cat", "response_format": response_format}
+        if output_format is not None:
+            request["output_format"] = output_format
+        response = client.post(path, json=request)
         assert response.status_code == 200, response.text
         body = response.json()
         assert body["data"][0]["revised_prompt"] == "a cat"

fastvideo/tests/entrypoints/test_video_generator.py

@@ -1228,3 +1228,46 @@ def test_ffmpeg_pipe_preserves_video_duration(tmp_path, env_overrides, audio_sec
     assert int(video["nb_read_frames"]) == frame_count
     assert float(video["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps)
     assert float(sound["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps)
+
+
+@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
+
+    args = _single_video_args()
+    args.workload_type = WorkloadType.from_string("t2i")
+    generator = _single_video_generator(_single_video_output_batch(torch.zeros(1, 3, 1, 16, 16)), args)
+    target = tmp_path / ("image" + extension)
+    output_path = generator._prepare_output_path(str(target), "a cat")
+    assert output_path == str(target)
+    sampling = _small_sampling_param(save_video=True)
+    sampling.num_frames = 1
+    result = generator._generate_single_video("a cat", sampling, output_path=output_path)
+    assert result["video_path"] == str(target)
+    with Image.open(target) as image:
+        assert image.format == format_name
+        assert image.size == (16, 16)
+
+
+def test_image_output_directory_defaults_to_png(tmp_path):
+    generator = _new_video_generator()
+    generator.fastvideo_args = SimpleNamespace(workload_type=WorkloadType.from_string("t2i"))
+    assert generator._prepare_output_path(str(tmp_path), "a cat") == str(tmp_path / "a cat.png")
+
+
+@pytest.mark.parametrize("extension,target_extension,workload", [
+    (".png", ".png", "t2i"), (".jpg", ".png", "t2i"), (".jpeg", ".png", "t2i"),
+    (".webp", ".png", "t2i"), (".JPEG", ".png", "t2i"),
+    (".mp4", ".mp4", "t2v"), (".wav", ".wav", "t2a"),
+])
+def test_existing_image_extension_directory_receives_output(tmp_path, extension, target_extension, workload):
+    generator = _new_video_generator()
+    generator.fastvideo_args = SimpleNamespace(workload_type=WorkloadType.from_string("t2i"))
+    generator.fastvideo_args.workload_type = WorkloadType.from_string(workload)
+    directory = tmp_path / ("renders" + extension)
+    directory.mkdir()
+    existing = directory / ("a cat" + target_extension)
+    existing.write_bytes(b"existing image")
+    result = generator._prepare_output_path(str(directory), "a cat")
+    assert result == str(directory / ("a cat_1" + target_extension))
+    assert existing.read_bytes() == b"existing image"