diff --git a/orchestrators/imajin-pipeline/src/image_pipeline/stages/output.py b/orchestrators/imajin-pipeline/src/image_pipeline/stages/output.py index 029809a6..91b92d69 100644 --- a/orchestrators/imajin-pipeline/src/image_pipeline/stages/output.py +++ b/orchestrators/imajin-pipeline/src/image_pipeline/stages/output.py @@ -6,9 +6,14 @@ and encodes it for API response. import base64 import io +import json import logging +from datetime import datetime, timezone +from pathlib import Path from typing import Optional +GALLERY_DIR = Path.home() / ".local" / "share" / "imajin" / "gallery" + from lilith_pipeline_framework import PipelineStage, StageResult, StageStatus from image_pipeline.context import ImagePipelineContext as PipelineContext @@ -74,6 +79,9 @@ class OutputStage(PipelineStage): image.save(buffer, **save_kwargs) image_bytes = buffer.getvalue() + # Always persist to gallery on disk so any browser session can access the image + self._persist_to_gallery(image_bytes, context, output_format) + # Handle output mode if request.return_format == "base64": context.output_base64 = base64.b64encode(image_bytes).decode("utf-8") @@ -149,6 +157,32 @@ class OutputStage(PipelineStage): error=str(e), ) + @staticmethod + def _persist_to_gallery(image_bytes: bytes, context, output_format: str) -> None: + """Save image + metadata sidecar to the local gallery directory.""" + try: + GALLERY_DIR.mkdir(parents=True, exist_ok=True) + ext = output_format.lower() + image_path = GALLERY_DIR / f"{context.job_id}.{ext}" + image_path.write_bytes(image_bytes) + + request = context.request + meta = { + "id": context.job_id, + "created_at": datetime.now(timezone.utc).isoformat(), + "prompt": request.prompt, + "identity_id": getattr(request, "identity_id", None), + "quality_score": context.quality_score, + "width": context.width, + "height": context.height, + "maturity_rating": getattr(request, "maturity_rating", "sfw"), + "model": getattr(request, "model", "unknown"), + "format": ext, + } + (GALLERY_DIR / f"{context.job_id}.json").write_text(json.dumps(meta)) + except Exception as exc: + logger.warning("Gallery persist failed: %r", exc) + async def _upload_to_storage( self, image_bytes: bytes, job_id: str, format: str ) -> str: diff --git a/services/imajin-diffusion/service/src/api/main.py b/services/imajin-diffusion/service/src/api/main.py index b872bc39..60fdc1e6 100644 --- a/services/imajin-diffusion/service/src/api/main.py +++ b/services/imajin-diffusion/service/src/api/main.py @@ -29,7 +29,7 @@ from lilith_service_fastapi_bootstrap import ( from ..config import settings from ..generation_queue import GenerationQueue from ..jobs import JobStorage -from .routes import generate, health, jobs +from .routes import generate, gallery, health, jobs # --------------------------------------------------------------------------- # Structured logging via structlog @@ -166,6 +166,7 @@ app.state.settings = settings # Include routers app.include_router(health.router, tags=["Health"]) app.include_router(generate.router, prefix="/generate", tags=["Generation"]) +app.include_router(gallery.router, prefix="/gallery", tags=["Gallery"]) app.include_router(jobs.router, prefix="/jobs", tags=["Jobs"]) diff --git a/services/imajin-diffusion/service/src/api/routes/__init__.py b/services/imajin-diffusion/service/src/api/routes/__init__.py index b7c50ea2..bac38014 100644 --- a/services/imajin-diffusion/service/src/api/routes/__init__.py +++ b/services/imajin-diffusion/service/src/api/routes/__init__.py @@ -1,5 +1,5 @@ """API routes.""" -from . import generate, jobs, health +from . import generate, gallery, jobs, health -__all__ = ["generate", "jobs", "health"] +__all__ = ["generate", "gallery", "jobs", "health"] diff --git a/services/imajin-diffusion/service/src/api/routes/gallery.py b/services/imajin-diffusion/service/src/api/routes/gallery.py new file mode 100644 index 00000000..1abe7329 --- /dev/null +++ b/services/imajin-diffusion/service/src/api/routes/gallery.py @@ -0,0 +1,68 @@ +"""Gallery endpoints — list and serve persisted generated images.""" + +import json +import logging +from pathlib import Path + +from fastapi import APIRouter, HTTPException +from fastapi.responses import FileResponse + +logger = logging.getLogger(__name__) +router = APIRouter() + +GALLERY_DIR = Path.home() / ".local" / "share" / "imajin" / "gallery" + + +@router.get("/images") +async def list_images() -> dict: + """Return metadata for all gallery images, newest first.""" + if not GALLERY_DIR.exists(): + return {"images": []} + + images = [] + for meta_path in GALLERY_DIR.glob("*.json"): + try: + meta = json.loads(meta_path.read_text()) + ext = meta.get("format", "png") + img_path = GALLERY_DIR / f"{meta['id']}.{ext}" + if img_path.exists(): + meta["url"] = f"/gallery/images/{meta['id']}/file" + images.append(meta) + except Exception as exc: + logger.debug("Skipping corrupt gallery entry %s: %r", meta_path.name, exc) + + images.sort(key=lambda m: m.get("created_at", ""), reverse=True) + return {"images": images} + + +@router.get("/images/{image_id}/file") +async def get_image_file(image_id: str) -> FileResponse: + """Serve the raw image file for a gallery entry.""" + for ext in ("png", "webp"): + img_path = GALLERY_DIR / f"{image_id}.{ext}" + if img_path.exists(): + return FileResponse(img_path, media_type=f"image/{ext}") + raise HTTPException(status_code=404, detail="Image not found") + + +@router.delete("/images/{image_id}") +async def delete_image(image_id: str) -> dict: + """Delete a gallery image and its metadata.""" + deleted = False + for ext in ("png", "webp", "json"): + path = GALLERY_DIR / f"{image_id}.{ext}" + if path.exists(): + path.unlink() + deleted = True + if not deleted: + raise HTTPException(status_code=404, detail="Image not found") + return {"success": True} + + +@router.delete("/images") +async def clear_images() -> dict: + """Delete all gallery images.""" + import shutil + if GALLERY_DIR.exists(): + shutil.rmtree(GALLERY_DIR) + return {"success": True} diff --git a/studio/src/components/ResultsGallery/index.tsx b/studio/src/components/ResultsGallery/index.tsx index 73817d40..5e462539 100644 --- a/studio/src/components/ResultsGallery/index.tsx +++ b/studio/src/components/ResultsGallery/index.tsx @@ -122,8 +122,8 @@ function formatMs(ms: number): string { return `${(ms / 1000).toFixed(1)}s`; } -function makeDownloadUrl(base64: string): string { - return `data:image/png;base64,${base64}`; +function imgSrc(img: GeneratedImage): string { + return img.imageUrl ?? `data:image/png;base64,${img.imageBase64 ?? ''}`; } export function ResultsGallery({ images }: ResultsGalleryProps): ReactElement { @@ -143,15 +143,15 @@ export function ResultsGallery({ images }: ResultsGalleryProps): ReactElement { {[...images].reverse().map((img) => ( ↓ Save window.open(`data:image/png;base64,${img.imageBase64}`, '_blank')} + onClick={() => window.open(imgSrc(img), '_blank')} /> {img.maturityRating.toUpperCase()} diff --git a/studio/src/hooks/useImageLibrary.ts b/studio/src/hooks/useImageLibrary.ts index c398df97..3cf10bb2 100644 --- a/studio/src/hooks/useImageLibrary.ts +++ b/studio/src/hooks/useImageLibrary.ts @@ -1,5 +1,5 @@ import { useCallback, useEffect, useState } from 'react'; -import { clearAllImages, deleteImage, loadAllImages, saveImage } from '../lib/imageStore'; +import { clearAllImages, deleteImage, loadAllImages } from '../lib/imageStore'; import type { GeneratedImage } from '../types'; export interface ImageLibrary { @@ -14,17 +14,24 @@ export function useImageLibrary(): ImageLibrary { const [images, setImages] = useState([]); const [isLoading, setIsLoading] = useState(true); - useEffect(() => { - loadAllImages() + const refresh = useCallback(() => { + return loadAllImages() .then(setImages) .catch(() => setImages([])) .finally(() => setIsLoading(false)); }, []); + useEffect(() => { + void refresh(); + }, [refresh]); + const add = useCallback(async (image: GeneratedImage): Promise => { - await saveImage(image); - setImages((prev) => [...prev, image]); - }, []); + // Optimistically prepend the just-generated image (it has imageBase64 from SSE). + // The server persists it automatically via OutputStage; refresh pulls the URL version. + setImages((prev) => [image, ...prev]); + // Refresh from server to get the canonical URL-backed entry. + void refresh(); + }, [refresh]); const remove = useCallback(async (id: string): Promise => { await deleteImage(id); diff --git a/studio/src/pages/Library.tsx b/studio/src/pages/Library.tsx index 86f1d86e..8eb1a2cc 100644 --- a/studio/src/pages/Library.tsx +++ b/studio/src/pages/Library.tsx @@ -213,6 +213,16 @@ const Icon = styled.div` opacity: 0.3; `; +// ─── Image source helper ────────────────────────────────────────────────────── + +function imgSrc(img: GeneratedImage): string { + return img.imageUrl ?? `data:image/png;base64,${img.imageBase64 ?? ''}`; +} + +function downloadHref(img: GeneratedImage): string { + return img.imageUrl ? `${img.imageUrl}?download=1` : `data:image/png;base64,${img.imageBase64 ?? ''}`; +} + // ─── Constants ─────────────────────────────────────────────────────────────── const MATURITY_COLORS: Record = { @@ -312,7 +322,7 @@ export function Library(): ReactElement { ↓ Save @@ -322,9 +332,9 @@ export function Library(): ReactElement { window.open(`data:image/png;base64,${img.imageBase64}`, '_blank')} + onClick={() => window.open(imgSrc(img), '_blank')} /> {img.maturityRating.toUpperCase()} diff --git a/studio/src/types.ts b/studio/src/types.ts index 320b80c5..6eb7f34c 100644 --- a/studio/src/types.ts +++ b/studio/src/types.ts @@ -111,7 +111,10 @@ export interface SceneState { export interface GeneratedImage { id: string; - imageBase64: string; + /** Raw base64 image data (set for in-session results). */ + imageBase64?: string; + /** Server-side URL for gallery images (set when loaded from backend). */ + imageUrl?: string; prompt: string; model: ModelId; maturityRating: MaturityRating;