Upload 60 files

.env.example

@@ -12,6 +12,14 @@ MAX_UPLOAD_MB=100
 MODEL_CPU_OFFLOAD=true
 MIXED_PRECISION=true
 
+# Estimated ZeroGPU allocations in seconds. WAN uses the xlarge tier.
+ZEROGPU_FLUX_DURATION=180
+ZEROGPU_WAN_DURATION=300
+ZEROGPU_KOKORO_DURATION=90
+ZEROGPU_MUSICGEN_DURATION=180
+ZEROGPU_WHISPER_DURATION=180
+ZEROGPU_SFX_DURATION=180
+
 # Hugging Face token; create this as a Space secret, never commit it.
 HF_TOKEN=
 

.github/workflows/test.yml

@@ -16,8 +16,10 @@ jobs:
           cache: pip
       - run: python -m pip install -r requirements-test.txt
       - run: python -m compileall -q .
-      - run: python -m pyflakes app.py config.py core gateway_mcp models routes utils tests
+      - run: >-
+          python -m pyflakes app.py config.py gradio_ui.py
+          core gateway_mcp models routes utils tests
       - run: >-
           python -m pycodestyle --max-line-length=100 --ignore=E203,W503
-          app.py config.py core gateway_mcp models routes utils tests
+          app.py config.py gradio_ui.py core gateway_mcp models routes utils tests
       - run: pytest -q

README.md

@@ -3,23 +3,25 @@ title: AI Gateway
 emoji: 🚪
 colorFrom: indigo
 colorTo: purple
-sdk: docker
-app_port: 7860
+sdk: gradio
+sdk_version: 6.20.0
+python_version: 3.11
+app_file: app.py
 pinned: false
 ---
 
 # AI Gateway
 
 AI Gateway is a single-worker FastAPI service for image, video, speech, music,
-sound-effect, and transcription models. It is designed for a Hugging Face Docker
-Space and exposes REST endpoints suitable for n8n, Vercel server functions, and
-mobile backends. The same capabilities are also exposed as Model Context Protocol
-(MCP) tools for ChatGPT, Claude, Cursor, Codex, and other MCP clients.
+sound-effect, and transcription models. It is designed for a Hugging Face Gradio
+Space with ZeroGPU support. It provides a Gradio interface at `/ui`, REST endpoints
+for n8n and application backends, and Model Context Protocol (MCP) tools for
+ChatGPT, Claude, Cursor, Codex, and other MCP clients.
 
 The application starts without downloading model weights. A model is downloaded
 and loaded only when its endpoint receives its first queued job. One model can be
-resident at a time; it is moved off GPU after inference and fully unloaded when a
-different model is requested.
+resident at a time on a regular host. In ZeroGPU workers, weights and Torch caches
+are fully released after every scheduled model stage.
 
 ## Models
 
@@ -40,9 +42,9 @@ checkpoint.
 ## Architecture
 
 ```text
-REST clients                 MCP clients
-     |                            |
-     +----------> FastAPI <-------+
+Gradio UI              REST clients              MCP clients
+     |                      |                          |
+     +--------------------> FastAPI <-----------------+
                        |
           +------------+------------+
           |                         |
@@ -52,6 +54,8 @@ REST clients                 MCP clients
                        |
               single-worker queue
                        |
+          spaces.GPU ZeroGPU scheduler
+                       |
               singleton model loader
                        |
           isolated FLUX / WAN / audio adapter
@@ -59,12 +63,12 @@ REST clients                 MCP clients
               UUID output manager
 ```
 
-REST routes and MCP handlers contain no inference implementation. Both call the
-same `core/manager.py` service instance, queue, loader, and workflow. Model adapters
-contain no API code. The full workflow is one queue item, so another REST or MCP
-request cannot take the GPU between its stages. The MCP integration lives under
-`gateway_mcp/`; using a local top-level `mcp/` package would shadow the official MCP
-Python SDK required by FastMCP.
+Gradio, REST, and MCP contain no inference implementation. All three call the same
+`core/manager.py` service, queue, serializable command executor, loader, and workflow.
+The queue captures request context so ZeroGPU can apply the originating Hugging Face
+quota. WAN requests use the `xlarge` ZeroGPU tier; other models use `large`. A full
+workflow remains one non-interleavable gateway job while each model stage receives
+its own GPU allocation. Model weights are fully unloaded after every ZeroGPU stage.
 
 ## Run locally
 
@@ -76,53 +80,38 @@ cp .env.example .env
 python -m venv .venv
 . .venv/bin/activate
 pip install -r requirements.txt
-uvicorn app:app --host 0.0.0.0 --port 7860 --workers 1
+python app.py
 ```
 
 Do not increase the worker count: serialization is process-local, so multiple
 workers would allow concurrent model loading.
 
-Docker:
-
-```bash
-docker build -t ai-gateway .
-docker run --rm --gpus all -p 7860:7860 \
-  -e HF_TOKEN="$HF_TOKEN" \
-  -v hf-cache:/data/.huggingface \
-  ai-gateway
-```
-
-The image starts with:
+The local startup command is equivalent to:
 
 ```text
 uvicorn app:app --host 0.0.0.0 --port 7860 --workers 1 --proxy-headers
 ```
 
-Its Docker health check calls `GET http://127.0.0.1:7860/health` without loading
-a model.
+`GET /health` does not download or load model weights.
 
 ## Hugging Face Space deployment
 
-1. Create a new Space with **Docker** as the SDK.
+1. Create a new Space with **Gradio** as the SDK, or push this repository and let
+   the README metadata configure it.
 2. Push this directory to the Space repository.
-3. Add `HF_TOKEN` as a Space secret if a configured checkpoint is gated.
-4. Attach persistent storage if model weights should survive rebuilds. The image
-   uses `/data/.huggingface` for the Hub cache and `/data/output` for assets.
-5. Select GPU hardware with enough system RAM and VRAM for WAN 2.2. The default
-   A14B checkpoint is substantially larger than every other configured model.
-
-### ZeroGPU limitation
+3. Select **ZeroGPU** as the Space hardware.
+4. Add `HF_TOKEN`, `API_KEY`, and MCP credentials as Space secrets when needed.
+5. `packages.txt` installs FFmpeg, espeak-ng, libsndfile, and Git.
+6. Attach persistent storage and set `OUTPUT_FOLDER=/data/output` if generated
+   assets must survive restarts.
 
-As of July 2026, Hugging Face documents ZeroGPU as exclusively compatible with
-the **Gradio SDK**, while this deployment is a **Docker SDK** Space. Those two
-target requirements cannot be enabled simultaneously by Docker configuration.
-For a normal Docker GPU Space, keep `ENABLE_ZEROGPU=false`.
+### ZeroGPU execution
 
-The model isolation, single-worker queue, CPU offload, and cleanup logic are
-suitable for a later Gradio-SDK port, but this Docker application deliberately
-does not expose a nonfunctional ZeroGPU switch. A ZeroGPU edition must move the
-inference entry points behind top-level `spaces.GPU`-decorated Gradio handlers;
-the scheduler cannot serialize this API server's queued Python closures.
+Top-level `spaces.GPU` functions receive pickle-safe inference commands rather than
+closures. This is required because ZeroGPU runs allocated work in forked processes.
+Successes and domain errors are serialized back to the parent, where REST and MCP
+retain their normal response contracts. On a normal CPU/GPU machine the same
+decorators become direct calls, so no alternate inference implementation is needed.
 
 ## Configuration
 
@@ -141,9 +130,15 @@ variables.
 | `API_KEY` | unset | Optional bearer or `X-API-Key` authentication |
 | `MCP_API_KEY` | unset | Optional MCP-only `X-API-Key` credential |
 | `MCP_BEARER_TOKEN` | unset | Optional MCP-only Bearer credential |
+| `ZEROGPU_FLUX_DURATION` | `180` | FLUX allocation estimate in seconds |
+| `ZEROGPU_WAN_DURATION` | `300` | WAN xlarge allocation estimate in seconds |
+| `ZEROGPU_KOKORO_DURATION` | `90` | Kokoro allocation estimate in seconds |
+| `ZEROGPU_MUSICGEN_DURATION` | `180` | MusicGen allocation estimate in seconds |
+| `ZEROGPU_WHISPER_DURATION` | `180` | Whisper allocation estimate in seconds |
+| `ZEROGPU_SFX_DURATION` | `180` | SFX allocation estimate in seconds |
 | `HF_TOKEN` | unset | Hugging Face access token |
-| `OUTPUT_FOLDER` | `output` | Generated assets (`/data/output` in Docker) |
-| `TMP_FOLDER` | `tmp` | Request-scoped uploads (`/tmp/ai-gateway` in Docker) |
+| `OUTPUT_FOLDER` | `output` | Generated assets; use `/data/output` with persistent storage |
+| `TMP_FOLDER` | `tmp` | Request-scoped uploads |
 | `FLUX_MODEL_ID` | `black-forest-labs/FLUX.2-klein-4B` | Image checkpoint |
 | `WAN_MODEL_ID` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Video checkpoint |
 | `KOKORO_MODEL_ID` | `hexgrad/Kokoro-82M` | Speech checkpoint |
@@ -156,7 +151,8 @@ If `API_KEY` is configured, send either `X-API-Key: ...` or
 generated `/output` assets are protected. `API_KEY` authorizes both REST and MCP.
 The two `MCP_*` credentials authorize `/mcp` only and do not broaden REST access.
 Never put these values in source control; configure them as Hugging Face Space
-secrets.
+secrets. The Gradio UI at `/ui` remains public so Hugging Face can associate
+ZeroGPU usage with the signed-in visitor.
 
 ## MCP
 
@@ -347,9 +343,9 @@ serialization, and workflow order without downloading weights:
 ```bash
 pip install -r requirements-test.txt
 pytest -q
-python -m pyflakes app.py config.py core gateway_mcp models routes utils tests
+python -m pyflakes app.py config.py gradio_ui.py core gateway_mcp models routes utils tests
 python -m pycodestyle --max-line-length=100 --ignore=E203,W503 \
-  app.py config.py core gateway_mcp models routes utils tests
+  app.py config.py gradio_ui.py core gateway_mcp models routes utils tests
 ```
 
 ## Troubleshooting
@@ -361,9 +357,10 @@ python -m pycodestyle --max-line-length=100 --ignore=E203,W503 \
 - `queue_full` or `queue_timeout`: reduce caller concurrency or raise `MAX_QUEUE` /
   `JOB_TIMEOUT_SECONDS`. Do not add Uvicorn workers, because serialization is
   process-local.
-- `out_of_memory`: reduce image/video settings, keep CPU offload enabled, or select
-  larger GPU hardware. The loader releases GPU residency and clears Torch caches
-  between models, but WAN 2.2 still has a high minimum memory requirement.
+- `gpu_unavailable`: the ZeroGPU quota or scheduler is temporarily unavailable;
+  authenticate with Hugging Face for user quota or retry later.
+- `out_of_memory`: reduce image/video settings. WAN uses the xlarge allocation tier,
+  but the selected checkpoint still has a high minimum memory requirement.
 - `model_load_failed`: verify model IDs, model licenses, outbound Hub access, and
   `HF_TOKEN` for gated repositories.
 - `invalid_file_path`: video/transcription MCP inputs accept only existing asset paths

app.py

@@ -8,19 +8,23 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from uuid import uuid4
 
+import gradio as gr
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from fastmcp.utilities.lifespan import combine_lifespans
+from gradio.context import LocalContext
 
 from config import Settings, get_settings
 from core.errors import GatewayError
 from core.loader import ModelLoader
 from core.manager import AIService, TaskManager
+from core.runtime import initialize_zerogpu
 from core.workflow import WorkflowService
 from gateway_mcp.auth import AuthenticationError, authenticate_headers
 from gateway_mcp.server import MCPServerBundle, create_mcp_server
+from gradio_ui import create_gradio_ui
 from routes import audio, health, image, video, workflow
 from routes.schemas import ErrorResponse
 from utils.files import OutputManager
@@ -54,6 +58,7 @@ def create_app(settings: Settings | None = None) -> FastAPI:
     workflows = WorkflowService(resolved_settings, tasks, outputs)
     ai = AIService(resolved_settings, tasks, loader, outputs, workflows)
     mcp = create_mcp_server(ai)
+    gradio_ui = create_gradio_ui(ai)
     services = GatewayServices(
         resolved_settings, outputs, loader, tasks, workflows, ai, mcp
     )
@@ -63,6 +68,7 @@ def create_app(settings: Settings | None = None) -> FastAPI:
         application.state.gateway = services
         removed = outputs.cleanup_stale_tmp()
         logger.info("gateway starting", extra={"stale_tmp_removed": removed})
+        initialize_zerogpu()
         await tasks.start()
         try:
             yield
@@ -85,10 +91,14 @@ def create_app(settings: Settings | None = None) -> FastAPI:
         request_id = uuid4().hex
         request.state.request_id = request_id
         token = request_id_context.set(request_id)
+        gradio_token = LocalContext.request.set(gr.Request(request))
         started = time.perf_counter()
         try:
             public_paths = {"/", "/health", "/docs", "/openapi.json", "/redoc"}
-            if request.url.path not in public_paths:
+            gradio_path = request.url.path == "/ui" or request.url.path.startswith(
+                "/ui/"
+            )
+            if request.url.path not in public_paths and not gradio_path:
                 try:
                     identity = authenticate_headers(
                         request.headers,
@@ -123,6 +133,7 @@ def create_app(settings: Settings | None = None) -> FastAPI:
                 },
             )
             request_id_context.reset(token)
+            LocalContext.request.reset(gradio_token)
         response.headers["X-Request-ID"] = request_id
         return response
 
@@ -186,14 +197,23 @@ def create_app(settings: Settings | None = None) -> FastAPI:
 
     @application.get("/", include_in_schema=False)
     async def root() -> RedirectResponse:
-        return RedirectResponse(url="/docs")
+        return RedirectResponse(url="/ui")
 
     application.mount(
         "/output",
         StaticFiles(directory=resolved_settings.output_folder, check_dir=False),
         name="output",
     )
-    return application
+    return gr.mount_gradio_app(
+        application,
+        gradio_ui,
+        path="/ui",
+        server_name=resolved_settings.host,
+        server_port=resolved_settings.port,
+        allowed_paths=[str(outputs.root)],
+        show_error=False,
+        max_file_size=f"{resolved_settings.max_upload_mb}mb",
+    )
 
 
 app = create_app()
@@ -203,4 +223,4 @@ if __name__ == "__main__":
     import uvicorn
 
     settings = get_settings()
-    uvicorn.run("app:app", host=settings.host, port=settings.port, proxy_headers=True)
+    uvicorn.run(app, host=settings.host, port=settings.port, proxy_headers=True)

config.py

@@ -56,6 +56,13 @@ class Settings(BaseSettings):
     sfx_duration: float = Field(default=5.0, ge=0.5, le=30.0)
     sfx_steps: int = Field(default=50, ge=1, le=200)
 
+    zerogpu_flux_duration: int = Field(default=180, ge=30, le=900)
+    zerogpu_wan_duration: int = Field(default=300, ge=30, le=900)
+    zerogpu_kokoro_duration: int = Field(default=90, ge=30, le=900)
+    zerogpu_musicgen_duration: int = Field(default=180, ge=30, le=900)
+    zerogpu_whisper_duration: int = Field(default=180, ge=30, le=900)
+    zerogpu_sfx_duration: int = Field(default=180, ge=30, le=900)
+
     flux_model_id: str = "black-forest-labs/FLUX.2-klein-4B"
     flux_revision: str | None = None
     wan_model_id: str = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"

core/executor.py

@@ -0,0 +1,169 @@
+"""Serializable inference dispatch across local, GPU, and ZeroGPU runtimes."""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Literal
+
+from core.errors import GatewayError
+from core.loader import ModelLoader
+from core.runtime import inference_context, memory_stats
+
+
+logger = logging.getLogger(__name__)
+
+
+def _gpu_task(**options: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Use the Space scheduler in deployment and a direct call in local tests."""
+    try:
+        import spaces
+    except ImportError:
+        return lambda function: function
+    return spaces.GPU(**options)
+
+
+@dataclass(frozen=True, slots=True)
+class InferenceCommand:
+    """Pickle-safe description of one model method invocation."""
+
+    model_name: str
+    method_name: str
+    arguments: dict[str, Any]
+    request_id: str
+    duration_seconds: int
+    gpu_size: Literal["large", "xlarge"] = "large"
+
+
+@dataclass(frozen=True, slots=True)
+class CommandError:
+    """Serializable error returned from a forked ZeroGPU worker."""
+
+    message: str
+    status_code: int
+    code: str
+
+
+@dataclass(frozen=True, slots=True)
+class CommandResult:
+    """Serializable success or failure from a GPU allocation."""
+
+    value: Any = None
+    error: CommandError | None = None
+
+
+def _duration(command: InferenceCommand) -> int:
+    return command.duration_seconds
+
+
+def _zero_gpu_active() -> bool:
+    try:
+        from spaces.config import Config
+    except ImportError:
+        return False
+    return bool(Config.zero_gpu)
+
+
+def _invoke(command: InferenceCommand) -> CommandResult:
+    """Execute one command inside an allocated device process."""
+    loader = ModelLoader()
+    started = time.perf_counter()
+    try:
+        with loader.use_model(command.model_name) as model:
+            method = getattr(model, command.method_name, None)
+            if not callable(method):
+                raise RuntimeError(
+                    f"Model {command.model_name} does not implement "
+                    f"{command.method_name}"
+                )
+            with inference_context(loader.settings, loader.device):
+                value = method(**command.arguments)
+    except GatewayError as exc:
+        return CommandResult(
+            error=CommandError(exc.message, exc.status_code, exc.code)
+        )
+    except Exception as exc:
+        logger.exception(
+            "model inference failed",
+            extra={
+                "model": command.model_name,
+                "request_id": command.request_id,
+                **memory_stats(),
+            },
+        )
+        message = str(exc)
+        if "out of memory" in message.lower():
+            error = CommandError(
+                f"{command.model_name} ran out of memory",
+                507,
+                "out_of_memory",
+            )
+        elif isinstance(exc, OSError):
+            error = CommandError(
+                f"File operation failed during {command.model_name} inference",
+                500,
+                "file_error",
+            )
+        else:
+            error = CommandError(
+                f"{command.model_name} inference failed: {message}",
+                500,
+                "inference_failed",
+            )
+        return CommandResult(error=error)
+    finally:
+        if _zero_gpu_active():
+            loader.close()
+
+    logger.info(
+        "model inference completed",
+        extra={
+            "model": command.model_name,
+            "request_id": command.request_id,
+            "execution_time": round(time.perf_counter() - started, 3),
+            **memory_stats(),
+        },
+    )
+    return CommandResult(value=value)
+
+
+@_gpu_task(duration=_duration, size="large")
+def _invoke_large(command: InferenceCommand) -> CommandResult:
+    return _invoke(command)
+
+
+@_gpu_task(duration=_duration, size="xlarge")
+def _invoke_xlarge(command: InferenceCommand) -> CommandResult:
+    return _invoke(command)
+
+
+def execute_inference(command: InferenceCommand) -> Any:
+    """Acquire the requested ZeroGPU tier and restore domain errors in the parent."""
+    try:
+        result = (
+            _invoke_xlarge(command)
+            if command.gpu_size == "xlarge"
+            else _invoke_large(command)
+        )
+    except Exception as exc:
+        message = str(exc)
+        normalized = message.lower()
+        if "quota" in normalized or "no gpu" in normalized or "zerogpu" in normalized:
+            raise GatewayError(
+                "ZeroGPU is unavailable or its quota was exceeded",
+                status_code=503,
+                code="gpu_unavailable",
+            ) from exc
+        raise GatewayError(
+            "ZeroGPU scheduling failed",
+            status_code=503,
+            code="gpu_scheduling_failed",
+        ) from exc
+    if result.error is not None:
+        raise GatewayError(
+            result.error.message,
+            status_code=result.error.status_code,
+            code=result.error.code,
+        )
+    return result.value

core/manager.py

@@ -2,16 +2,15 @@
 
 from __future__ import annotations
 
-import logging
-import time
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Literal, Protocol, TypeVar, cast
+from typing import Any, Callable, Literal, Protocol, TypeVar
 
 from config import Settings
-from core.errors import GatewayError, InferenceError, OutOfMemoryError
+from core.errors import GatewayError
+from core.executor import InferenceCommand, execute_inference
 from core.loader import ModelLoader
 from core.queue import InferenceQueue
-from core.runtime import gpu_available, inference_context, memory_stats
+from core.runtime import gpu_available, memory_stats, zerogpu_enabled
 from core.schemas import (
     EpisodeAssets,
     ImageAsset,
@@ -27,21 +26,11 @@ from core.schemas import (
     VideoRequest,
     WorkflowRequest,
 )
-from models.base import ModelAdapter
 from utils.files import OutputManager
 from utils.validation import validate_image_dimensions
 
-if TYPE_CHECKING:
-    from models.flux import FluxModel
-    from models.kokoro import KokoroModel
-    from models.musicgen import MusicGenModel
-    from models.sfx import SFXModel
-    from models.wan import WanModel
-    from models.whisper import WhisperModel
-
 
 T = TypeVar("T")
-logger = logging.getLogger(__name__)
 
 
 class WorkflowRunner(Protocol):
@@ -69,15 +58,13 @@ class TaskManager:
 
     async def run(
         self,
-        model_name: str,
-        request_id: str,
-        action: Callable[[ModelAdapter], T],
-    ) -> T:
+        command: InferenceCommand,
+    ) -> Any:
         """Queue one model action."""
         return await self.queue.submit(
-            request_id,
-            model_name,
-            lambda: self.invoke_direct(model_name, request_id, action),
+            command.request_id,
+            command.model_name,
+            lambda: self.invoke_direct(command),
         )
 
     async def run_exclusive(
@@ -86,45 +73,10 @@ class TaskManager:
         """Queue a compound operation as one non-interleavable job."""
         return await self.queue.submit(request_id, label, action)
 
-    def invoke_direct(
-        self,
-        model_name: str,
-        request_id: str,
-        action: Callable[[ModelAdapter], T],
-    ) -> T:
-        """Invoke a model inside the queue worker; workflows reuse this method."""
-        started = time.perf_counter()
-        try:
-            with self.loader.use_model(model_name) as model:
-                with inference_context(self.settings, self.loader.device):
-                    result = action(model)
-        except GatewayError:
-            raise
-        except Exception as exc:
-            logger.exception(
-                "model inference failed",
-                extra={"model": model_name, "request_id": request_id, **memory_stats()},
-            )
-            message = str(exc)
-            if "out of memory" in message.lower():
-                raise OutOfMemoryError(model_name) from exc
-            if isinstance(exc, OSError):
-                raise GatewayError(
-                    f"File operation failed during {model_name} inference",
-                    status_code=500,
-                    code="file_error",
-                ) from exc
-            raise InferenceError(model_name, message) from exc
-        logger.info(
-            "model inference completed",
-            extra={
-                "model": model_name,
-                "request_id": request_id,
-                "execution_time": round(time.perf_counter() - started, 3),
-                **memory_stats(),
-            },
-        )
-        return result
+    @staticmethod
+    def invoke_direct(command: InferenceCommand) -> Any:
+        """Invoke a serializable command; workflows reuse this method."""
+        return execute_inference(command)
 
 
 class AIService:
@@ -170,16 +122,20 @@ class AIService:
         target = self.outputs.allocate("images")
         try:
             await self.tasks.run(
-                "flux",
-                request_id,
-                lambda adapter: cast("FluxModel", adapter).generate(
-                    prompt=payload.prompt,
-                    output_path=target,
-                    width=width,
-                    height=height,
-                    steps=steps,
-                    seed=payload.seed,
-                    guidance_scale=guidance_scale,
+                InferenceCommand(
+                    model_name="flux",
+                    method_name="generate",
+                    arguments={
+                        "prompt": payload.prompt,
+                        "output_path": target,
+                        "width": width,
+                        "height": height,
+                        "steps": steps,
+                        "seed": payload.seed,
+                        "guidance_scale": guidance_scale,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_flux_duration,
                 ),
             )
         except Exception:
@@ -203,18 +159,23 @@ class AIService:
         target = self.outputs.allocate("videos")
         try:
             await self.tasks.run(
-                "wan",
-                request_id,
-                lambda adapter: cast("WanModel", adapter).generate(
-                    image_path=source,
-                    prompt=payload.prompt,
-                    negative_prompt=payload.negative_prompt,
-                    output_path=target,
-                    steps=steps,
-                    frames=frames,
-                    fps=fps,
-                    seed=payload.seed,
-                    guidance_scale=guidance_scale,
+                InferenceCommand(
+                    model_name="wan",
+                    method_name="generate",
+                    arguments={
+                        "image_path": source,
+                        "prompt": payload.prompt,
+                        "negative_prompt": payload.negative_prompt,
+                        "output_path": target,
+                        "steps": steps,
+                        "frames": frames,
+                        "fps": fps,
+                        "seed": payload.seed,
+                        "guidance_scale": guidance_scale,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_wan_duration,
+                    gpu_size="xlarge",
                 ),
             )
         except Exception:
@@ -228,13 +189,17 @@ class AIService:
         target = self.outputs.allocate("audio")
         try:
             await self.tasks.run(
-                "kokoro",
-                request_id,
-                lambda adapter: cast("KokoroModel", adapter).synthesize(
-                    text=payload.text,
-                    voice=voice,
-                    speed=payload.speed,
-                    output_path=target,
+                InferenceCommand(
+                    model_name="kokoro",
+                    method_name="synthesize",
+                    arguments={
+                        "text": payload.text,
+                        "voice": voice,
+                        "speed": payload.speed,
+                        "output_path": target,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_kokoro_duration,
                 ),
             )
         except Exception:
@@ -255,14 +220,18 @@ class AIService:
         target = self.outputs.allocate("music")
         try:
             await self.tasks.run(
-                "musicgen",
-                request_id,
-                lambda adapter: cast("MusicGenModel", adapter).generate(
-                    prompt=payload.prompt,
-                    duration=duration,
-                    guidance_scale=guidance_scale,
-                    seed=payload.seed,
-                    output_path=target,
+                InferenceCommand(
+                    model_name="musicgen",
+                    method_name="generate",
+                    arguments={
+                        "prompt": payload.prompt,
+                        "duration": duration,
+                        "guidance_scale": guidance_scale,
+                        "seed": payload.seed,
+                        "output_path": target,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_musicgen_duration,
                 ),
             )
         except Exception:
@@ -279,14 +248,18 @@ class AIService:
         target = self.outputs.allocate("audio")
         try:
             await self.tasks.run(
-                "sfx",
-                request_id,
-                lambda adapter: cast("SFXModel", adapter).generate(
-                    prompt=payload.prompt,
-                    duration=duration,
-                    steps=steps,
-                    seed=payload.seed,
-                    output_path=target,
+                InferenceCommand(
+                    model_name="sfx",
+                    method_name="generate",
+                    arguments={
+                        "prompt": payload.prompt,
+                        "duration": duration,
+                        "steps": steps,
+                        "seed": payload.seed,
+                        "output_path": target,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_sfx_duration,
                 ),
             )
         except Exception:
@@ -306,13 +279,17 @@ class AIService:
         subtitle = self.outputs.allocate("subtitles")
         try:
             result = await self.tasks.run(
-                "whisper",
-                request_id,
-                lambda adapter: cast("WhisperModel", adapter).transcribe(
-                    source=source,
-                    subtitle_path=subtitle,
-                    language=language,
-                    task=task,
+                InferenceCommand(
+                    model_name="whisper",
+                    method_name="transcribe",
+                    arguments={
+                        "source": source,
+                        "subtitle_path": subtitle,
+                        "language": language,
+                        "task": task,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_whisper_duration,
                 ),
             )
         except Exception:
@@ -373,7 +350,7 @@ class AIService:
         """Return operational metadata without loading a model."""
         return {
             "version": self.settings.app_version,
-            "device": self.loader.device,
+            "device": "zerogpu" if zerogpu_enabled() else self.loader.device,
             "gpu": "available" if gpu_available() else "unavailable",
             "memory": memory_stats(),
             "available_endpoints": [*self.rest_endpoints, "/mcp"],

core/queue.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from contextvars import Context, copy_context
 from dataclasses import dataclass
 from typing import Any, Callable, Generic, TypeVar
 
@@ -17,6 +18,7 @@ class QueueJob(Generic[T]):
     request_id: str
     label: str
     operation: Callable[[], T]
+    context: Context
     future: asyncio.Future[T]
     started: asyncio.Future[None]
     cancelled: bool = False
@@ -61,6 +63,7 @@ class InferenceQueue:
             request_id=request_id,
             label=label,
             operation=operation,
+            context=copy_context(),
             future=future,
             started=started,
         )
@@ -110,7 +113,7 @@ class InferenceQueue:
             if not job.started.done():
                 job.started.set_result(None)
             try:
-                result = await asyncio.to_thread(job.operation)
+                result = await asyncio.to_thread(job.context.run, job.operation)
             except Exception as exc:
                 if not job.future.done():
                     job.future.set_exception(exc)

core/runtime.py

@@ -57,7 +57,9 @@ def preferred_dtype(device: str, mixed_precision: bool = True) -> Any:
 
 
 def gpu_available() -> bool:
-    """Return CUDA availability without making health checks fail."""
+    """Return scheduled ZeroGPU or directly attached CUDA availability."""
+    if zerogpu_enabled():
+        return True
     try:
         torch = import_torch()
     except RuntimeError:
@@ -65,6 +67,24 @@ def gpu_available() -> bool:
     return bool(torch.cuda.is_available())
 
 
+def zerogpu_enabled() -> bool:
+    """Return whether this process is running in a Hugging Face ZeroGPU Space."""
+    try:
+        from spaces.config import Config
+    except ImportError:
+        return False
+    return bool(Config.zero_gpu)
+
+
+def initialize_zerogpu() -> None:
+    """Run the scheduler startup hook when Gradio is mounted into FastAPI."""
+    if not zerogpu_enabled():
+        return
+    from spaces.zero import startup
+
+    startup()
+
+
 @contextmanager
 def inference_context(settings: Settings, device: str) -> Iterator[None]:
     """Disable autograd and enable CUDA mixed precision when configured."""

core/workflow.py

@@ -3,20 +3,13 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
 
 from config import Settings
+from core.executor import InferenceCommand
 from core.manager import TaskManager
 from core.schemas import WorkflowRequest
 from utils.files import OutputManager
 
-if TYPE_CHECKING:
-    from models.flux import FluxModel
-    from models.kokoro import KokoroModel
-    from models.musicgen import MusicGenModel
-    from models.wan import WanModel
-    from models.whisper import WhisperModel
-
 
 class WorkflowService:
     """Run image, voice, music, video, and subtitle generation in order."""
@@ -43,26 +36,34 @@ class WorkflowService:
 
         def execute() -> dict[str, Path]:
             self.tasks.invoke_direct(
-                "flux",
-                request_id,
-                lambda adapter: cast("FluxModel", adapter).generate(
-                    prompt=payload.image_prompt,
-                    output_path=targets["image"],
-                    width=width,
-                    height=height,
-                    steps=self.settings.image_steps,
-                    seed=payload.seed,
-                    guidance_scale=self.settings.flux_guidance_scale,
+                InferenceCommand(
+                    model_name="flux",
+                    method_name="generate",
+                    arguments={
+                        "prompt": payload.image_prompt,
+                        "output_path": targets["image"],
+                        "width": width,
+                        "height": height,
+                        "steps": self.settings.image_steps,
+                        "seed": payload.seed,
+                        "guidance_scale": self.settings.flux_guidance_scale,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_flux_duration,
                 ),
             )
             self.tasks.invoke_direct(
-                "kokoro",
-                request_id,
-                lambda adapter: cast("KokoroModel", adapter).synthesize(
-                    text=payload.script,
-                    voice=voice,
-                    speed=1.0,
-                    output_path=targets["voice"],
+                InferenceCommand(
+                    model_name="kokoro",
+                    method_name="synthesize",
+                    arguments={
+                        "text": payload.script,
+                        "voice": voice,
+                        "speed": 1.0,
+                        "output_path": targets["voice"],
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_kokoro_duration,
                 ),
             )
             words = len(payload.script.split())
@@ -71,39 +72,52 @@ class WorkflowService:
                 f"cinematic instrumental background score for {payload.title}, no vocals"
             )
             self.tasks.invoke_direct(
-                "musicgen",
-                request_id,
-                lambda adapter: cast("MusicGenModel", adapter).generate(
-                    prompt=music_prompt,
-                    duration=music_duration,
-                    guidance_scale=self.settings.music_guidance_scale,
-                    seed=payload.seed,
-                    output_path=targets["music"],
+                InferenceCommand(
+                    model_name="musicgen",
+                    method_name="generate",
+                    arguments={
+                        "prompt": music_prompt,
+                        "duration": music_duration,
+                        "guidance_scale": self.settings.music_guidance_scale,
+                        "seed": payload.seed,
+                        "output_path": targets["music"],
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_musicgen_duration,
                 ),
             )
             self.tasks.invoke_direct(
-                "wan",
-                request_id,
-                lambda adapter: cast("WanModel", adapter).generate(
-                    image_path=targets["image"],
-                    prompt=payload.video_prompt,
-                    negative_prompt="low quality, distorted, static",
-                    output_path=targets["video"],
-                    steps=self.settings.wan_steps,
-                    frames=self.settings.video_frames,
-                    fps=self.settings.video_fps,
-                    seed=payload.seed,
-                    guidance_scale=self.settings.wan_guidance_scale,
+                InferenceCommand(
+                    model_name="wan",
+                    method_name="generate",
+                    arguments={
+                        "image_path": targets["image"],
+                        "prompt": payload.video_prompt,
+                        "negative_prompt": "low quality, distorted, static",
+                        "output_path": targets["video"],
+                        "steps": self.settings.wan_steps,
+                        "frames": self.settings.video_frames,
+                        "fps": self.settings.video_fps,
+                        "seed": payload.seed,
+                        "guidance_scale": self.settings.wan_guidance_scale,
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_wan_duration,
+                    gpu_size="xlarge",
                 ),
             )
             self.tasks.invoke_direct(
-                "whisper",
-                request_id,
-                lambda adapter: cast("WhisperModel", adapter).transcribe(
-                    source=targets["voice"],
-                    subtitle_path=targets["subtitle"],
-                    language=None,
-                    task="transcribe",
+                InferenceCommand(
+                    model_name="whisper",
+                    method_name="transcribe",
+                    arguments={
+                        "source": targets["voice"],
+                        "subtitle_path": targets["subtitle"],
+                        "language": None,
+                        "task": "transcribe",
+                    },
+                    request_id=request_id,
+                    duration_seconds=self.settings.zerogpu_whisper_duration,
                 ),
             )
             return targets

gradio_ui.py

@@ -0,0 +1,331 @@
+"""Gradio user interface over the shared AI Gateway service layer."""
+
+from __future__ import annotations
+
+from functools import wraps
+from typing import Any, Awaitable, Callable, TypeVar
+from uuid import uuid4
+
+import gradio as gr
+from pydantic import ValidationError
+
+from core.errors import GatewayError
+from core.manager import AIService
+from core.schemas import (
+    ImageRequest,
+    MusicRequest,
+    SFXRequest,
+    TTSRequest,
+    VideoRequest,
+    WorkflowRequest,
+)
+
+
+F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
+
+
+def ui_errors(function: F) -> F:
+    """Convert domain and input errors into safe Gradio errors."""
+
+    @wraps(function)
+    async def wrapped(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return await function(*args, **kwargs)
+        except GatewayError as exc:
+            raise gr.Error(f"{exc.code}: {exc.message}") from exc
+        except ValidationError as exc:
+            raise gr.Error(str(exc)) from exc
+        except Exception as exc:
+            raise gr.Error("The gateway could not complete this request") from exc
+
+    return wrapped  # type: ignore[return-value]
+
+
+class GradioGatewayUI:
+    """Thin UI callbacks that delegate every operation to AIService."""
+
+    def __init__(self, service: AIService) -> None:
+        self.service = service
+
+    @staticmethod
+    def _request_id(request: gr.Request) -> str:
+        state = getattr(request, "state", None)
+        request_id = getattr(state, "request_id", None)
+        if request_id:
+            return str(request_id)
+        return uuid4().hex
+
+    def _local_path(self, public_path: str, categories: set[str]) -> str:
+        return str(self.service.outputs.resolve_public_input(public_path, categories=categories))
+
+    @ui_errors
+    async def generate_image(
+        self,
+        prompt: str,
+        width: int,
+        height: int,
+        steps: int,
+        seed: int,
+        request: gr.Request,
+    ) -> tuple[str, str]:
+        asset = await self.service.generate_image(
+            ImageRequest(
+                prompt=prompt,
+                width=width,
+                height=height,
+                steps=steps,
+                seed=seed,
+            ),
+            self._request_id(request),
+        )
+        return self._local_path(asset.image, {"images"}), asset.image
+
+    @ui_errors
+    async def generate_video(
+        self,
+        image_path: str,
+        prompt: str,
+        request: gr.Request,
+    ) -> tuple[str, str]:
+        asset = await self.service.generate_video(
+            VideoRequest(image=image_path, prompt=prompt),
+            self._request_id(request),
+        )
+        return self._local_path(asset.video, {"videos"}), asset.video
+
+    @ui_errors
+    async def generate_speech(
+        self,
+        text: str,
+        voice: str,
+        request: gr.Request,
+    ) -> tuple[str, str]:
+        asset = await self.service.generate_speech(
+            TTSRequest(text=text, voice=voice.strip() or None), self._request_id(request)
+        )
+        return self._local_path(asset.voice, {"audio"}), asset.voice
+
+    @ui_errors
+    async def generate_music(
+        self,
+        prompt: str,
+        duration: float,
+        request: gr.Request,
+    ) -> tuple[str, str]:
+        asset = await self.service.generate_music(
+            MusicRequest(prompt=prompt, duration=duration), self._request_id(request)
+        )
+        return self._local_path(asset.music, {"music"}), asset.music
+
+    @ui_errors
+    async def generate_sfx(
+        self,
+        prompt: str,
+        duration: float,
+        request: gr.Request,
+    ) -> tuple[str, str]:
+        asset = await self.service.generate_sfx(
+            SFXRequest(prompt=prompt, duration=duration), self._request_id(request)
+        )
+        return self._local_path(asset.sfx, {"audio"}), asset.sfx
+
+    @ui_errors
+    async def transcribe_audio(
+        self,
+        audio_path: str,
+        language: str,
+        request: gr.Request,
+    ) -> tuple[str, str]:
+        asset = await self.service.transcribe_output(
+            audio_path,
+            self._request_id(request),
+            language=language.strip() or None,
+        )
+        return asset.text, self._local_path(asset.subtitle, {"subtitles"})
+
+    @ui_errors
+    async def create_episode(
+        self,
+        title: str,
+        script: str,
+        image_prompt: str,
+        video_prompt: str,
+        voice: str,
+        request: gr.Request,
+    ) -> tuple[str, str, str, str, str]:
+        assets = await self.service.create_episode(
+            WorkflowRequest(
+                title=title,
+                script=script,
+                image_prompt=image_prompt,
+                video_prompt=video_prompt,
+                voice=voice.strip() or None,
+            ),
+            self._request_id(request),
+        )
+        return (
+            self._local_path(assets.image, {"images"}),
+            self._local_path(assets.video, {"videos"}),
+            self._local_path(assets.voice, {"audio"}),
+            self._local_path(assets.music, {"music"}),
+            self._local_path(assets.subtitle, {"subtitles"}),
+        )
+
+    def health(self) -> dict[str, object]:
+        return self.service.health()
+
+
+def create_gradio_ui(service: AIService) -> gr.Blocks:
+    """Build the queued Space UI without owning any inference logic."""
+    try:
+        import spaces
+    except ImportError:
+        pass
+    else:
+        spaces.disable_gradio_auto_wrap()
+    callbacks = GradioGatewayUI(service)
+    settings = service.settings
+
+    with gr.Blocks(title="AI Gateway") as demo:
+        gr.Markdown(
+            "# AI Gateway\n"
+            "One ZeroGPU-aware gateway for images, video, speech, music, SFX, "
+            "transcription, REST, and MCP."
+        )
+
+        with gr.Tab("Image"):
+            image_prompt = gr.Textbox(label="Prompt", lines=3)
+            with gr.Row():
+                image_width = gr.Slider(256, 2048, settings.image_width, step=8, label="Width")
+                image_height = gr.Slider(
+                    256, 2048, settings.image_height, step=8, label="Height"
+                )
+                image_steps = gr.Slider(1, 100, settings.image_steps, step=1, label="Steps")
+                image_seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
+            image_button = gr.Button("Generate image", variant="primary")
+            image_output = gr.Image(type="filepath", label="Image")
+            image_public = gr.Textbox(label="Gateway asset path")
+            image_button.click(
+                callbacks.generate_image,
+                [image_prompt, image_width, image_height, image_steps, image_seed],
+                [image_output, image_public],
+                api_name="generate_image",
+            )
+
+        with gr.Tab("Video"):
+            video_image = gr.Textbox(
+                label="Image asset path", placeholder="/output/images/....png"
+            )
+            video_prompt = gr.Textbox(label="Motion prompt", lines=3)
+            video_button = gr.Button("Generate video", variant="primary")
+            video_output = gr.Video(label="Video")
+            video_public = gr.Textbox(label="Gateway asset path")
+            video_button.click(
+                callbacks.generate_video,
+                [video_image, video_prompt],
+                [video_output, video_public],
+                api_name="generate_video",
+            )
+
+        with gr.Tab("Audio"):
+            with gr.Row():
+                with gr.Column():
+                    speech_text = gr.Textbox(label="Narration", lines=5)
+                    speech_voice = gr.Textbox(
+                        value=settings.kokoro_default_voice, label="Kokoro voice"
+                    )
+                    speech_button = gr.Button("Generate speech")
+                    speech_output = gr.Audio(type="filepath", label="Speech")
+                    speech_public = gr.Textbox(label="Speech asset path")
+                with gr.Column():
+                    music_prompt = gr.Textbox(label="Music prompt", lines=3)
+                    music_duration = gr.Slider(1, 30, settings.music_duration, label="Seconds")
+                    music_button = gr.Button("Generate music")
+                    music_output = gr.Audio(type="filepath", label="Music")
+                    music_public = gr.Textbox(label="Music asset path")
+                with gr.Column():
+                    sfx_prompt = gr.Textbox(label="Sound-effect prompt", lines=3)
+                    sfx_duration = gr.Slider(0.5, 30, settings.sfx_duration, label="Seconds")
+                    sfx_button = gr.Button("Generate SFX")
+                    sfx_output = gr.Audio(type="filepath", label="Sound effect")
+                    sfx_public = gr.Textbox(label="SFX asset path")
+            speech_button.click(
+                callbacks.generate_speech,
+                [speech_text, speech_voice],
+                [speech_output, speech_public],
+                api_name="generate_speech",
+            )
+            music_button.click(
+                callbacks.generate_music,
+                [music_prompt, music_duration],
+                [music_output, music_public],
+                api_name="generate_music",
+            )
+            sfx_button.click(
+                callbacks.generate_sfx,
+                [sfx_prompt, sfx_duration],
+                [sfx_output, sfx_public],
+                api_name="generate_sfx",
+            )
+
+        with gr.Tab("Transcription"):
+            transcription_source = gr.Textbox(
+                label="Audio/video asset path", placeholder="/output/audio/....wav"
+            )
+            transcription_language = gr.Textbox(label="Language (optional)")
+            transcription_button = gr.Button("Transcribe", variant="primary")
+            transcription_text = gr.Textbox(label="Transcript", lines=8)
+            transcription_subtitle = gr.File(label="SubRip subtitle")
+            transcription_button.click(
+                callbacks.transcribe_audio,
+                [transcription_source, transcription_language],
+                [transcription_text, transcription_subtitle],
+                api_name="transcribe_audio",
+            )
+
+        with gr.Tab("Episode"):
+            episode_title = gr.Textbox(label="Title")
+            episode_script = gr.Textbox(label="Script", lines=8)
+            episode_image_prompt = gr.Textbox(label="Image prompt", lines=3)
+            episode_video_prompt = gr.Textbox(label="Video prompt", lines=3)
+            episode_voice = gr.Textbox(
+                value=settings.kokoro_default_voice, label="Kokoro voice"
+            )
+            episode_button = gr.Button("Create episode", variant="primary")
+            with gr.Row():
+                episode_image = gr.Image(type="filepath", label="Image")
+                episode_video = gr.Video(label="Video")
+            with gr.Row():
+                episode_speech = gr.Audio(type="filepath", label="Narration")
+                episode_music = gr.Audio(type="filepath", label="Music")
+                episode_subtitle = gr.File(label="Subtitle")
+            episode_button.click(
+                callbacks.create_episode,
+                [
+                    episode_title,
+                    episode_script,
+                    episode_image_prompt,
+                    episode_video_prompt,
+                    episode_voice,
+                ],
+                [
+                    episode_image,
+                    episode_video,
+                    episode_speech,
+                    episode_music,
+                    episode_subtitle,
+                ],
+                api_name="create_episode",
+            )
+
+        with gr.Tab("System"):
+            health_output = gr.JSON(value=callbacks.health(), label="Gateway health")
+            gr.Button("Refresh health").click(
+                callbacks.health, outputs=health_output, api_name="health_check"
+            )
+            gr.Markdown("REST documentation: [`/docs`](/docs) · MCP endpoint: `/mcp`")
+
+    return demo.queue(
+        max_size=settings.max_queue,
+        default_concurrency_limit=settings.max_queue,
+    )

packages.txt

@@ -0,0 +1,4 @@
+espeak-ng
+ffmpeg
+git
+libsndfile1

requirements-test.txt

@@ -1,5 +1,7 @@
 fastapi>=0.115,<1.0
 fastmcp>=3.4.4,<4.0
+gradio>=6.20,<7.0
+spaces>=0.51,<1.0
 pydantic>=2.9,<3.0
 pydantic-settings>=2.5,<3.0
 python-multipart>=0.0.12,<1.0

requirements.txt

@@ -1,6 +1,8 @@
 # Web runtime
 fastapi>=0.115,<1.0
 fastmcp>=3.4.4,<4.0
+gradio>=6.20,<7.0
+spaces>=0.51,<1.0
 uvicorn[standard]>=0.32,<1.0
 pydantic>=2.9,<3.0
 pydantic-settings>=2.5,<3.0

routes/health.py

@@ -15,10 +15,11 @@ router = APIRouter(tags=["system"])
 async def health(services=Depends(get_services)) -> HealthResponse:
     """Report service state without loading any model."""
     snapshot = services.ai.health()
+    server = services.ai.server_info()
     return HealthResponse(
         status="ok",
         version=services.settings.app_version,
-        device=services.loader.device,
+        device=server["device"],
         queue_depth=snapshot["queue"],
         queue_capacity=services.tasks.queue.capacity,
         active_model=services.loader.active_model,

tests/test_api.py

@@ -30,6 +30,15 @@ def test_health_does_not_load_a_model(tmp_path: Path) -> None:
     assert response.json()["active_model"] is None
 
 
+def test_root_redirects_to_gradio_ui(tmp_path: Path) -> None:
+    application = make_app(tmp_path)
+    with TestClient(application) as client:
+        response = client.get("/", follow_redirects=False)
+
+    assert response.status_code == 307
+    assert response.headers["location"] == "/ui"
+
+
 def test_validation_errors_have_stable_shape(tmp_path: Path) -> None:
     application = make_app(tmp_path)
     with TestClient(application) as client:

tests/test_executor.py

@@ -0,0 +1,52 @@
+"""ZeroGPU command serialization and parent-side error restoration."""
+
+from __future__ import annotations
+
+import pickle
+from pathlib import Path
+
+import pytest
+
+from core.errors import GatewayError
+from core.executor import (
+    CommandError,
+    CommandResult,
+    InferenceCommand,
+    execute_inference,
+)
+
+
+def test_inference_command_is_pickle_safe(tmp_path: Path) -> None:
+    command = InferenceCommand(
+        model_name="flux",
+        method_name="generate",
+        arguments={"prompt": "city", "output_path": tmp_path / "image.png"},
+        request_id="request-1",
+        duration_seconds=180,
+    )
+
+    restored = pickle.loads(pickle.dumps(command))
+
+    assert restored == command
+
+
+def test_serialized_worker_error_restores_gateway_error(monkeypatch) -> None:
+    command = InferenceCommand(
+        model_name="flux",
+        method_name="generate",
+        arguments={},
+        request_id="request-1",
+        duration_seconds=180,
+    )
+    monkeypatch.setattr(
+        "core.executor._invoke_large",
+        lambda _: CommandResult(
+            error=CommandError("flux ran out of memory", 507, "out_of_memory")
+        ),
+    )
+
+    with pytest.raises(GatewayError) as captured:
+        execute_inference(command)
+
+    assert captured.value.status_code == 507
+    assert captured.value.code == "out_of_memory"

tests/test_queue.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import threading
 import time
+from contextvars import ContextVar
 
 import pytest
 
@@ -63,3 +64,16 @@ async def test_queue_timeout_cancels_job_before_inference() -> None:
     assert queued_job_ran is False
     assert queue.active is False
     await queue.stop()
+
+
+async def test_queue_preserves_request_context_for_worker() -> None:
+    queue = InferenceQueue(capacity=1, timeout_seconds=5)
+    request_context: ContextVar[str] = ContextVar("test_request", default="missing")
+    token = request_context.set("request-identity")
+    try:
+        result = await queue.submit("request", "context", request_context.get)
+    finally:
+        request_context.reset(token)
+        await queue.stop()
+
+    assert result == "request-identity"

tests/test_workflow.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any, Callable
 
 from config import Settings
+from core.executor import InferenceCommand
 from core.workflow import WorkflowService
 from routes.schemas import WorkflowRequest
 from utils.files import OutputManager
@@ -42,11 +43,11 @@ class FakeTasks:
         assert label == "workflow"
         return action()
 
-    def invoke_direct(
-        self, name: str, request_id: str, action: Callable[[FakeAdapter], Any]
-    ) -> Any:
-        self.calls.append(name)
-        return action(FakeAdapter(name))
+    def invoke_direct(self, command: InferenceCommand) -> Any:
+        self.calls.append(command.model_name)
+        adapter = FakeAdapter(command.model_name)
+        method = getattr(adapter, command.method_name)
+        return method(**command.arguments)
 
 
 async def test_workflow_runs_in_required_order(tmp_path: Path) -> None: