diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 8f6f21b9f..7dca36991 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/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) diff --git a/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py b/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py index 1cd3eb431..adbf63bb3 100644 --- a/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py +++ b/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" diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 4879060cb..55287823e 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/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"