Studio: local diffusion image generation page

studio/backend/models/inference.py

@@ -1421,3 +1421,70 @@ class AnthropicMessagesResponse(BaseModel):
     stop_reason: Optional[str] = None
     stop_sequence: Optional[str] = None
     usage: AnthropicUsage = Field(default_factory = AnthropicUsage)
+
+
+# ── Diffusion image generation ────────────────────────────────────
+
+
+class DiffusionLoadRequest(BaseModel):
+    """Load a diffusion image-generation model.
+
+    repo_id is the HF repo (either GGUF-only or full diffusers layout).
+    gguf_filename selects the quant when repo_id is a GGUF repo.
+    base_repo overrides the auto-picked diffusers base used for the
+    VAE / text encoders when loading a GGUF-only repo.
+    """
+
+    repo_id: str = Field(..., description = "HF repo id")
+    gguf_filename: Optional[str] = Field(
+        None, description = "GGUF filename inside repo_id (Q4_K_S, Q8_0, ...)"
+    )
+    base_repo: Optional[str] = Field(
+        None,
+        description = "Diffusers base repo to source VAE + text encoders from",
+    )
+    family: Optional[str] = Field(
+        None,
+        description = "Force pipeline family: flux.2-klein | flux.2 | flux.1 | qwen-image | stable-diffusion-3 | stable-diffusion-xl",
+    )
+    hf_token: Optional[str] = Field(
+        None, description = "HuggingFace token for gated models"
+    )
+    enable_model_cpu_offload: bool = Field(
+        True,
+        description = "Offload submodules to CPU between forwards. Trades a small speed hit for ~6 GB less VRAM on FLUX-class models.",
+    )
+
+
+class DiffusionGenerateRequest(BaseModel):
+    """Generate a single image from the currently-loaded diffusion model."""
+
+    prompt: str = Field(..., min_length = 1, max_length = 4000)
+    negative_prompt: Optional[str] = Field(None, max_length = 4000)
+    num_inference_steps: int = Field(24, ge = 1, le = 200)
+    guidance_scale: float = Field(3.5, ge = 0.0, le = 20.0)
+    width: int = Field(1024, ge = 64, le = 2048)
+    height: int = Field(1024, ge = 64, le = 2048)
+    seed: Optional[int] = Field(
+        None, description = "Deterministic seed for reproducible outputs"
+    )
+
+    @field_validator("width", "height")
+    @classmethod
+    def _multiple_of_eight(cls, v: int) -> int:
+        if v % 8:
+            raise ValueError("width and height must be multiples of 8")
+        return v
+
+
+class DiffusionGenerateResponse(BaseModel):
+    image_b64: str = Field(..., description = "Base64-encoded PNG")
+    image_mime: str = "image/png"
+    width: int
+    height: int
+    num_inference_steps: int
+    guidance_scale: float
+    seed: Optional[int] = None
+    duration_ms: int
+    model: Optional[str] = None
+    family: Optional[str] = None

studio/backend/requirements/no-torch-runtime.txt

@@ -46,6 +46,9 @@ peft>=0.18.0,!=0.11.0
 huggingface_hub>=0.34.0
 hf_transfer
 diffusers
+# Required by diffusers.GGUFQuantizationConfig (used by the Images page
+# to load FLUX.2 / FLUX.1 / Qwen-Image / SDXL GGUFs from the Hub).
+gguf
 
 # Transitive deps required because this file is installed with --no-deps.
 # Without these, `from transformers import AutoConfig` fails at import time.

studio/backend/routes/inference.py

@@ -213,6 +213,9 @@ def _friendly_error(exc: Exception) -> str:
     ListOpenAIContainersResponse,
     OpenAIContainerRequest,
     OpenAIContainerSummary,
+    DiffusionLoadRequest,
+    DiffusionGenerateRequest,
+    DiffusionGenerateResponse,
 )
 from core.inference.anthropic_compat import (
     anthropic_messages_to_openai,
@@ -1584,6 +1587,130 @@ async def generate_audio(
     )
 
 
+# =====================================================================
+# Diffusion image generation  (/images/*)
+# =====================================================================
+#
+# Lifecycle mirrors the GGUF chat backend: explicit load -> generate ->
+# unload. Diffusion pipelines compete for the same GPU as llama-server,
+# so callers on < 24 GB GPUs should unload the chat model first.
+
+
+def _get_diffusion_backend():
+    """Lazy import so non-diffusion installs do not pay the diffusers
+    cost at process start. The backend itself is a process-wide
+    singleton; reusing it across requests keeps pipeline state alive."""
+    from core.inference.diffusion import get_diffusion_backend
+
+    return get_diffusion_backend()
+
+
+@router.post("/images/load")
+async def diffusion_load(
+    payload: DiffusionLoadRequest,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Load a diffusion image-generation model.
+
+    Pass either a full diffusers repo or a GGUF-only repo plus the
+    desired ``gguf_filename``. Returns the new status payload (same
+    shape as ``/images/status``).
+    """
+    backend = _get_diffusion_backend()
+    try:
+        status = await asyncio.get_event_loop().run_in_executor(
+            None,
+            lambda: backend.load_model(
+                repo_id = payload.repo_id,
+                gguf_filename = payload.gguf_filename,
+                base_repo = payload.base_repo,
+                family_override = payload.family,
+                hf_token = payload.hf_token,
+                enable_model_cpu_offload = payload.enable_model_cpu_offload,
+            ),
+        )
+        return JSONResponse(content = status)
+    except RuntimeError as exc:
+        raise HTTPException(status_code = 400, detail = str(exc))
+    except Exception as exc:
+        logger.exception("Diffusion load failed")
+        raise HTTPException(status_code = 500, detail = str(exc))
+
+
+@router.post("/images/unload")
+async def diffusion_unload(
+    current_subject: str = Depends(get_current_subject),
+):
+    """Unload the current diffusion model and free GPU memory."""
+    backend = _get_diffusion_backend()
+    return backend.unload_model()
+
+
+@router.get("/images/status")
+async def diffusion_status(
+    current_subject: str = Depends(get_current_subject),
+):
+    """Return diffusion backend status (loaded, family, device, etc.)."""
+    backend = _get_diffusion_backend()
+    return backend.status()
+
+
+@router.post("/images/generate", response_model = DiffusionGenerateResponse)
+async def diffusion_generate(
+    payload: DiffusionGenerateRequest,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Generate a single image from the loaded diffusion model.
+
+    Returns a base64 PNG plus the generation parameters that produced
+    it so the frontend can render the result and the user can reproduce
+    it via the same seed.
+    """
+    backend = _get_diffusion_backend()
+    if not backend.is_loaded:
+        raise HTTPException(
+            status_code = 400,
+            detail = "No diffusion model is loaded. POST /api/inference/images/load first.",
+        )
+
+    start = time.time()
+    try:
+        from core.inference.diffusion import async_generate, encode_png_base64
+
+        image = await async_generate(
+            backend,
+            prompt = payload.prompt,
+            negative_prompt = payload.negative_prompt,
+            num_inference_steps = payload.num_inference_steps,
+            guidance_scale = payload.guidance_scale,
+            width = payload.width,
+            height = payload.height,
+            seed = payload.seed,
+        )
+    except ValueError as exc:
+        raise HTTPException(status_code = 400, detail = str(exc))
+    except RuntimeError as exc:
+        raise HTTPException(status_code = 400, detail = str(exc))
+    except Exception as exc:
+        logger.exception("Diffusion generation failed")
+        raise HTTPException(status_code = 500, detail = str(exc))
+
+    duration_ms = int((time.time() - start) * 1000)
+    status = backend.status()
+    return DiffusionGenerateResponse(
+        image_b64 = encode_png_base64(image),
+        image_mime = "image/png",
+        width = payload.width,
+        height = payload.height,
+        num_inference_steps = payload.num_inference_steps,
+        guidance_scale = payload.guidance_scale,
+        seed = payload.seed,
+        duration_ms = duration_ms,
+        model = status.get("repo_id"),
+        family = status.get("family"),
+    )
+
+
 # =====================================================================
 # OpenAI-Compatible Chat Completions  (/chat/completions)
 # =====================================================================