[bugfix] Fix MPS worker handoff and named Studio image downloads

.buildkite/scripts/unit_test.sh

@@ -16,6 +16,7 @@ exec pytest \
   ./fastvideo/tests/dataset/ \
   ./fastvideo/tests/workflow/ \
   ./fastvideo/tests/entrypoints/ \
+  ./apps/fastvideo_studio/tests/test_inference_output.py \
   ./fastvideo/tests/loader/ \
   ./fastvideo/tests/pipelines/ \
   ./fastvideo/tests/platforms/ \

apps/fastvideo_studio/job_runner.py

@@ -1102,10 +1102,11 @@ def _run_inference_job(self, job: Job):
 
             # Without a name FastVideo derives the filename from the prompt.
             safe_name = re.sub(r'[\\/:*?"<>|]+', "", job.name).strip().strip(".")
-            output_target = (os.path.join(job_output_dir, f"{safe_name[:80]}.mp4") if safe_name else job_output_dir)
+            output_video_name = safe_name[:80] or None
             gen_kwargs: dict[str, Any] = {
                 "prompt": job.prompt,
-                "output_path": output_target,
+                "output_path": job_output_dir,
+                "output_video_name": output_video_name,
                 "save_video": True,
                 "num_inference_steps": job.num_inference_steps,
                 "num_frames": job.num_frames,
@@ -1126,24 +1127,14 @@ def _run_inference_job(self, job: Job):
                 # _prepare_fl2va requires a PIL image, not a path.
                 from PIL import Image as _PILImage
                 gen_kwargs["last_image"] = _PILImage.open(job.last_image_path)
-            generator.generate_video(**gen_kwargs)
+            result = generator.generate_video(**gen_kwargs)
 
             buf.phase = "saving"
-            logger.info("Generation completed, searching for output file...")
-
-            # Find the generated video file
-            video_files = sorted(Path(job_output_dir).glob("*.mp4"))
-            if video_files:
-                job.output_path = str(video_files[0])
-                logger.info("Found video output: %s", job.output_path)
-            else:
-                # Could be an image workload
-                image_files = sorted(Path(job_output_dir).glob("*.png"))
-                if image_files:
-                    job.output_path = str(image_files[0])
-                    logger.info("Found image output: %s", job.output_path)
-                else:
-                    logger.warning("No output file found in job directory")
+            output_path = result.get("video_path") if isinstance(result, dict) else None
+            if not output_path or not os.path.isfile(output_path):
+                raise RuntimeError("Generation did not return a saved output file.")
+            job.output_path = str(output_path)
+            logger.info("Saved output: %s", job.output_path)
 
             if job._stop_event.is_set():
                 job.status = JobStatus.STOPPED

apps/fastvideo_studio/server.py

@@ -17,6 +17,7 @@
 import argparse
 import contextlib
 import logging
+import mimetypes
 import os
 import re
 import shutil
@@ -661,7 +662,8 @@ def get_video(job_id: str) -> FileResponse:
     if not os.path.isfile(job.output_path):
         raise HTTPException(status_code=404, detail="Output file not found on disk")
 
-    media_type = ("video/mp4" if job.output_path.endswith(".mp4") else "image/png")
+    # WebP is in Python's common MIME types, but not always its strict table.
+    media_type = mimetypes.guess_type(job.output_path, strict=False)[0]
     return FileResponse(job.output_path, media_type=media_type)
 
 

apps/fastvideo_studio/src/app/gallery/page.test.tsx

@@ -50,6 +50,7 @@ describe('GalleryPage', () => {
     const video = await screen.findByLabelText(
       'Generated video: a cat surfing a wave',
     );
+    expect(video).toHaveAttribute('src', 'http://test.local/api/jobs/job-1/video');
     expect(video).toHaveAttribute('controls');
 
     fireEvent.error(video);
@@ -59,6 +60,24 @@ describe('GalleryPage', () => {
     ).toBeInTheDocument();
   });
 
+  it.each(['png', 'jpg', 'jpeg', 'WEBP'])(
+    'previews %s output as an image through the API',
+    async (extension) => {
+      vi.mocked(getJobsList).mockResolvedValue([
+        makeJob({ output_path: `/server/out/portrait.${extension}` }),
+      ]);
+      renderGallery();
+
+      const image = await screen.findByRole('img', {
+        name: 'a cat surfing a wave',
+      });
+      expect(image).toHaveAttribute('src', 'http://test.local/api/jobs/job-1/video');
+      expect(
+        screen.queryByLabelText('Generated video: a cat surfing a wave'),
+      ).not.toBeInTheDocument();
+    },
+  );
+
   it('shows the empty state when no completed videos exist', async () => {
     vi.mocked(getJobsList).mockResolvedValue([
       makeJob({ status: 'running', output_path: null }),

apps/fastvideo_studio/src/app/gallery/page.tsx

@@ -9,7 +9,7 @@ import { getJobVideoUrl, getJobsList } from '@/lib/api';
 import type { Job } from '@/lib/types';
 
 function isImage(job: Job): boolean {
-  return job.output_path?.toLowerCase().endsWith('.png') ?? false;
+  return /\.(png|jpe?g|webp)$/i.test(job.output_path ?? '');
 }
 
 function GalleryMedia({ job }: { job: Job }) {

apps/fastvideo_studio/src/components/jobs/JobCard.test.tsx

@@ -10,6 +10,7 @@ import {
   stopJob,
 } from '@/lib/api';
 import type { Job } from '@/lib/types';
+import { downloadBlob } from '@/lib/utils';
 import { activeJobStore, setActiveJobId } from '@/stores/activeJob';
 import { makeJob as makeBaseJob } from '@/test/factories';
 
@@ -115,16 +116,26 @@ describe('JobCard', () => {
     expect(deleteJob).not.toHaveBeenCalled();
   });
 
-  it('downloads the video for a completed inference job', async () => {
+  it.each([
+    ['/server/out/portrait.png', 'portrait.png'],
+    ['/server/out/portrait.jpg', 'portrait.jpg'],
+    ['/server/out/portrait.jpeg', 'portrait.jpeg'],
+    ['/server/out/portrait.WEBP', 'portrait.WEBP'],
+    ['/server/out/clip.mp4', 'clip.mp4'],
+    ['C:\\server\\out\\portrait.JPEG', 'portrait.JPEG'],
+  ])('downloads %s with its actual filename', async (outputPath, filename) => {
+    const blob = new Blob(['generated output']);
+    vi.mocked(downloadJobVideo).mockResolvedValue(blob);
     render(
       <JobCard
-        job={makeJob({ status: 'completed', output_path: '/out/video.mp4' })}
+        job={makeJob({ status: 'completed', output_path: outputPath })}
       />,
     );
     await userEvent.click(
-      screen.getByRole('button', { name: 'Download Video' }),
+      screen.getByRole('button', { name: 'Download' }),
     );
     await waitFor(() => expect(downloadJobVideo).toHaveBeenCalledWith('job-1'));
+    expect(downloadBlob).toHaveBeenCalledWith(blob, filename);
   });
 
   it('selects the job when the card body is clicked', async () => {

apps/fastvideo_studio/src/components/jobs/JobCard.tsx

@@ -142,17 +142,17 @@ export default function JobCard({ job, onJobUpdated }: JobCardProps) {
     setActiveJobId(isSelected ? null : job.id);
   }
 
-  async function handleDownloadVideo(e: React.MouseEvent) {
+  async function handleDownload(e: React.MouseEvent) {
     e.preventDefault();
     e.stopPropagation();
     if (isLoading || !job.output_path) return;
     setIsLoading(true);
     try {
       const blob = await downloadJobVideo(job.id);
-      const ext = job.output_path.endsWith('.png') ? 'png' : 'mp4';
-      downloadBlob(blob, `job_${job.id}.${ext}`);
+      const filename = job.output_path.split(/[\\/]/).pop() || `job_${job.id}`;
+      downloadBlob(blob, filename);
     } catch (err) {
-      alert(err instanceof Error ? err.message : 'Failed to download video');
+      alert(err instanceof Error ? err.message : 'Failed to download output');
     } finally {
       setIsLoading(false);
     }
@@ -245,11 +245,11 @@ export default function JobCard({ job, onJobUpdated }: JobCardProps) {
             <Button
               size="sm"
               variant="outline"
-              onClick={handleDownloadVideo}
+              onClick={handleDownload}
               disabled={isLoading}
-              title="Download video"
+              title="Download output"
             >
-              Download Video
+              Download
             </Button>
           )}
         {!(

apps/fastvideo_studio/tests/test_duplicate_and_edit.py

@@ -10,11 +10,15 @@
 
 @pytest.fixture
 def runner(tmp_path):
-    return JobRunner(
+    instance = JobRunner(
         output_dir=str(tmp_path / "out"),
         log_dir=str(tmp_path / "logs"),
         database=Database(tmp_path / "t.db"),
     )
+    try:
+        yield instance
+    finally:
+        instance._shutdown()
 
 
 def _make(runner, **over):

apps/fastvideo_studio/tests/test_inference_output.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Studio records and serves the artifact actually saved by the generator."""
+
+import io
+import mimetypes
+from pathlib import Path
+from types import SimpleNamespace
+import uuid
+
+from fastapi.testclient import TestClient
+from PIL import Image
+import pytest
+
+from fastvideo.entrypoints.video_generator import VideoGenerator
+from fastvideo_studio import server
+from fastvideo_studio.database import Database
+from fastvideo_studio.job_runner import JobRunner, JobStatus
+
+
+@pytest.fixture(autouse=True)
+def builtin_mime_types_only(monkeypatch):
+    mime_types = mimetypes.MimeTypes(filenames=())
+    monkeypatch.setattr(mimetypes, "guess_type", mime_types.guess_type)
+    monkeypatch.setattr("starlette.responses.guess_type", mime_types.guess_type)
+
+
+@pytest.fixture
+def runner(tmp_path):
+    instance = JobRunner(
+        output_dir=str(tmp_path / "outputs"),
+        log_dir=str(tmp_path / "logs"),
+        database=Database(tmp_path / "studio.db"),
+    )
+    try:
+        yield instance
+    finally:
+        instance._shutdown()
+
+
+@pytest.mark.parametrize("workload, name, output_name, media_type", [
+    ("t2i", "MPS handoff test", "MPS handoff test.png", "image/png"),
+    ("t2i", "portrait.jpg", "portrait.jpg", "image/jpeg"),
+    ("t2i", "portrait.jpeg", "portrait.jpeg", "image/jpeg"),
+    ("t2i", "portrait.webp", "portrait.webp", "image/webp"),
+    ("t2v", "MPS handoff test", "MPS handoff test.mp4", "video/mp4"),
+])
+def test_named_job_serves_generated_artifact(runner, monkeypatch, workload, name, output_name, media_type):
+    job = runner.create_job(
+        job_id=str(uuid.uuid4()),
+        model_id="test/model",
+        name=name,
+        prompt="a red dinosaur",
+        workload_type=workload,
+    )
+    output_dir = Path(runner.output_dir) / job.id
+    output_dir.mkdir()
+    # Old artifacts and directories with media suffixes must not win discovery.
+    (output_dir / "000-old.mp4").mkdir()
+    (output_dir / "001-old.mp4").write_bytes(b"old video")
+    path_context = SimpleNamespace(
+        _is_image_workload=lambda: workload == "t2i",
+        _is_audio_workload=lambda: False,
+    )
+
+    def generate_video(**kwargs):
+        # Use the real filename/extension selection, replacing only model work.
+        output = Path(VideoGenerator._prepare_output_path(
+            path_context,
+            kwargs["output_path"],
+            kwargs["prompt"],
+            output_video_name=kwargs.get("output_video_name"),
+        ))
+        if workload == "t2i":
+            Image.new("RGB", (8, 8), "red").save(output)
+        else:
+            output.write_bytes(b"generated video fixture")
+        return {"video_path": str(output)}
+
+    monkeypatch.setattr(runner, "_get_or_create_generator",
+                        lambda *args, **kwargs: SimpleNamespace(generate_video=generate_video))
+    runner._run_inference_job(job)
+
+    expected = output_dir / output_name
+    assert job.status == JobStatus.COMPLETED, job.error
+    assert job.output_path == str(expected)
+    assert expected.is_file()
+    assert runner._db.get_job(job.id)["output_path"] == str(expected)
+    monkeypatch.setattr(server, "job_runner", runner, raising=False)
+    with TestClient(server.app) as client:
+        response = client.get(f"/api/jobs/{job.id}/video")
+    assert response.status_code == 200
+    assert response.headers["content-type"] == media_type
+    assert response.content == expected.read_bytes()
+    if workload == "t2i":
+        with Image.open(io.BytesIO(response.content)) as image:
+            image.load()
+            assert Image.MIME[image.format] == media_type
+            assert image.size == (8, 8)
+
+
+def test_directory_returned_as_output_fails_job(runner, monkeypatch):
+    job = runner.create_job(job_id=str(uuid.uuid4()), model_id="test/model", prompt="test")
+    directory = Path(runner.output_dir) / job.id / "not-a-video.mp4"
+    directory.mkdir(parents=True)
+    generator = SimpleNamespace(generate_video=lambda **kwargs: {"video_path": str(directory)})
+    monkeypatch.setattr(runner, "_get_or_create_generator", lambda *args, **kwargs: generator)
+
+    runner._run_inference_job(job)
+
+    assert job.status == JobStatus.FAILED
+    assert job.output_path is None
+    assert job.error