From 8fe8498df2331c7665e4616c598dad0c7968b95c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 15:08:41 +0000 Subject: [PATCH 1/8] Add local GPU inference: auto-detect GPU, auto-download best diffusion models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds AI_PROVIDER=local_gpu — a fully self-contained GPU inference engine using HuggingFace Diffusers that requires zero InvokeAI/ComfyUI setup. All existing providers (InvokeAI, ComfyUI, OpenAI, Replicate) remain intact and can be mixed with local GPU via per-operation overrides. New features: - GPU auto-detection (CUDA/NVIDIA, MPS/Apple Silicon, CPU fallback) - VRAM-tiered model selection: ultra ≥16 GB → SDXL inpaint + SDXL base high 8-16 GB → SDXL inpaint + SDXL base medium 4-8 GB → SD 2.x inpaint + SD 2.1 low <4 GB → SD 2.x (small) - Auto-download model weights to HuggingFace disk cache at startup (background task; first request loads from local disk, not internet) - LRU pipeline cache evicts oldest GPU pipeline when VRAM limit reached - Per-operation model overrides via HF_MODEL_INPAINT / HF_MODEL_TXT2IMG etc. - Optional HF_TOKEN for gated/private HuggingFace models New files: - backend/app/services/gpu_detect.py — GPU detection + tier/model mapping - backend/app/services/local_diffusion.py — Diffusers provider + LRU cache - backend/app/routers/gpu_status.py — GET /api/gpu/status, POST /api/gpu/prefetch - backend/requirements.gpu.txt — Diffusers ecosystem deps (GPU only) - docker-compose.gpu.yml — NVIDIA GPU compose (one-command startup) - Dockerfile.gpu — pytorch/pytorch:2.1.0-cuda12.1 base image - scripts/gpu_setup.py — Startup GPU info logger Modified: - backend/app/config.py — local_gpu settings added - backend/app/services/remote_provider.py — local_gpu registered as provider - backend/app/routers/ai_tools.py — /api/config exposes GPU tier + caps - backend/app/main.py — GPU router + background prefetch task - backend/entrypoint.sh — runs gpu_setup.py at container start - .env.example — local_gpu documented as first option Quick start with GPU: docker compose -f docker-compose.gpu.yml up --build https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- .env.example | 31 +- Dockerfile.gpu | 75 ++++ backend/app/config.py | 9 + backend/app/main.py | 26 ++ backend/app/routers/ai_tools.py | 9 + backend/app/routers/gpu_status.py | 74 ++++ backend/app/services/gpu_detect.py | 140 +++++++ backend/app/services/local_diffusion.py | 472 ++++++++++++++++++++++++ backend/app/services/remote_provider.py | 7 + backend/entrypoint.sh | 5 + backend/requirements.gpu.txt | 24 ++ docker-compose.gpu.yml | 94 +++++ scripts/gpu_setup.py | 67 ++++ 13 files changed, 1027 insertions(+), 6 deletions(-) create mode 100644 Dockerfile.gpu create mode 100644 backend/app/routers/gpu_status.py create mode 100644 backend/app/services/gpu_detect.py create mode 100644 backend/app/services/local_diffusion.py create mode 100644 backend/requirements.gpu.txt create mode 100644 docker-compose.gpu.yml create mode 100644 scripts/gpu_setup.py diff --git a/.env.example b/.env.example index aab2dce..0e1138a 100644 --- a/.env.example +++ b/.env.example @@ -14,18 +14,37 @@ # ============================================================================= # STEP 1: Choose AI Provider # ============================================================================= -# Options: mock, openai, stability, replicate +# Options: local_gpu, mock, openai, stability, replicate, invokeai, comfyui # -# mock = Free, but returns original image unchanged (for testing UI) -# openai = DALL-E 2 inpainting (~$0.02/image) - lower quality -# stability = Stability AI SDXL (~$0.01/image) - good quality -# replicate = Multiple models (~$0.002-0.03/image) - RECOMMENDED +# local_gpu = FREE, runs on YOUR GPU — best option if you have an NVIDIA card +# (use docker-compose.gpu.yml — models auto-download on first use) +# mock = Free, returns original image unchanged (UI testing only) +# openai = DALL-E 3 / gpt-image-1 (~$0.02-0.04/image) +# stability = Stability AI SDXL (~$0.01/image) +# replicate = Multiple models (~$0.002-0.03/image) +# invokeai = Self-hosted InvokeAI running on another machine +# comfyui = Self-hosted ComfyUI running on another machine # -# RECOMMENDED: Use "replicate" for best quality and model variety +# GPU QUICK-START: +# docker compose -f docker-compose.gpu.yml up --build +# (AI_PROVIDER defaults to local_gpu in that compose file) # ============================================================================= AI_PROVIDER=replicate +# ── Local GPU settings (only relevant when AI_PROVIDER=local_gpu) ──────────── +# Auto-download HuggingFace models on first request (true/false) +AUTO_DOWNLOAD_MODELS=true +# Max diffusion pipelines to keep loaded in GPU memory (each is 2–7 GB) +LOCAL_GPU_MAX_PIPELINES=2 +# HuggingFace token — only needed for gated/private models +#HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx +# Override auto-selected model for any operation (leave blank = auto by VRAM tier) +#HF_MODEL_INPAINT=your-org/your-inpaint-model +#HF_MODEL_TXT2IMG=your-org/your-txt2img-model +#HF_MODEL_IMG2IMG=your-org/your-img2img-model +# ───────────────────────────────────────────────────────────────────────────── + # Per-operation provider overrides (optional — blank means use AI_PROVIDER above) # Example: use OpenAI for text-to-image (best quality) but InvokeAI for everything else #AI_PROVIDER_TXT2IMG=openai diff --git a/Dockerfile.gpu b/Dockerfile.gpu new file mode 100644 index 0000000..4516c49 --- /dev/null +++ b/Dockerfile.gpu @@ -0,0 +1,75 @@ +# ============================================================================= +# EditmaskwithAI — GPU Container (NVIDIA CUDA) +# +# Usage: +# docker compose -f docker-compose.gpu.yml up --build +# +# Requirements on host: +# - NVIDIA driver ≥ 525 (for CUDA 12.x) +# - nvidia-container-toolkit installed and configured +# - docker compose v2 (or docker-compose with GPU device support) +# +# AMD ROCm users: replace the pytorch base image with a ROCm variant, e.g. +# rocm/pytorch:latest (and remove the nvidia-smi check below) +# ============================================================================= + +# ── Stage 1: Build miniPaint frontend ──────────────────────────────────────── +FROM node:20-alpine AS frontend-build + +WORKDIR /frontend +COPY frontend/package.json frontend/package-lock.json* ./ +RUN npm install +COPY frontend/ ./ +RUN npm run build + +# ── Stage 2: PyTorch CUDA runtime ──────────────────────────────────────────── +# pytorch/pytorch already includes torch + torchvision built for CUDA 12.1. +# Using the runtime (not devel) image keeps the layer lean. +FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime + +WORKDIR /app + +# System dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + libgomp1 \ + wget \ + git \ + && rm -rf /var/lib/apt/lists/* + +# Install Python dependencies — base + GPU extras +COPY backend/requirements.txt . +COPY backend/requirements.gpu.txt . +RUN pip install --no-cache-dir -r requirements.txt +RUN pip install --no-cache-dir -r requirements.gpu.txt + +# Smoke-test rembg (model downloads on first use) +RUN python -c "from rembg import remove; print('rembg OK')" \ + || echo "WARNING: rembg unavailable — Remove Background disabled" + +# Copy backend application +COPY backend/ . + +# Entrypoint +COPY backend/entrypoint.sh /entrypoint.sh +RUN chmod +x /entrypoint.sh + +# Scripts (SAM download, DB init, GPU setup, etc.) +COPY scripts/ /scripts/ +RUN chmod +x /scripts/*.py 2>/dev/null || true + +# Copy built frontend from Stage 1 +COPY --from=frontend-build /frontend/index.html /app/static/ +COPY --from=frontend-build /frontend/dist /app/static/dist +COPY --from=frontend-build /frontend/images /app/static/images +COPY --from=frontend-build /frontend/src/css /app/static/src/css + +# Persistent data directories +RUN mkdir -p /app/data/projects /app/data/patches /app/data/models + +EXPOSE 8000 +ENTRYPOINT ["/entrypoint.sh"] diff --git a/backend/app/config.py b/backend/app/config.py index 0b80ad9..cfe29e1 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -46,6 +46,15 @@ class Settings(BaseSettings): # Allow per-edit model override allow_model_override: bool = True + # Local GPU diffusion (AI_PROVIDER=local_gpu) + auto_download_models: bool = True # download HF models on first use + local_gpu_max_pipelines: int = 2 # max diffusion pipelines kept in GPU memory + hf_token: str = "" # HuggingFace token (only needed for gated models) + # Override auto-selected models per operation (leave blank = auto-pick by VRAM tier) + hf_model_inpaint: str = "" + hf_model_txt2img: str = "" + hf_model_img2img: str = "" + # File Storage data_dir: str = "./data" max_upload_size_mb: int = 50 diff --git a/backend/app/main.py b/backend/app/main.py index 082ec1e..cfd06da 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -10,6 +10,7 @@ import os from app.config import settings from app.database import init_db from app.routers import projects, edits, images, patches, generate, tools, ai_tools, print_tools +from app.routers import gpu_status @asynccontextmanager @@ -24,6 +25,30 @@ async def lifespan(app: FastAPI): # Pre-download SAM model in background so first click is fast from app.services.sam_service import ensure_sam_installed asyncio.create_task(ensure_sam_installed()) + + # If local GPU provider is active, log GPU info at startup + if settings.ai_provider.lower() == "local_gpu" or any( + v.lower() == "local_gpu" + for v in [ + settings.ai_provider_inpaint, + settings.ai_provider_txt2img, + settings.ai_provider_img2img, + settings.ai_provider_outpaint, + ] + if v + ): + from app.services.gpu_detect import get_cached_gpu_info + info = get_cached_gpu_info() + print( + f"[gpu] {info.device_name} | {info.vram_gb:.1f} GB | tier={info.tier} | " + f"backend={info.backend}" + ) + if settings.auto_download_models: + # Download model weight files to disk cache in background so first + # user request loads from local disk instead of the internet. + from app.services.local_diffusion import prefetch_model_files + asyncio.create_task(prefetch_model_files()) + yield @@ -52,6 +77,7 @@ app.include_router(generate.router) app.include_router(tools.router) app.include_router(ai_tools.router) app.include_router(print_tools.router) +app.include_router(gpu_status.router) @app.get("/api") diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index 07ca4f6..4e63b65 100644 --- a/backend/app/routers/ai_tools.py +++ b/backend/app/routers/ai_tools.py @@ -327,12 +327,21 @@ async def get_config(): # Default provider for display (used when no per-op override) default_name = (settings.ai_provider or "").lower() or None + from app.services.gpu_detect import get_cached_gpu_info + gpu_info = get_cached_gpu_info() + return { "local": { "lama": lama_available(), "rembg": rembg_available(), "opencv": True, "gpu_detected": gpu_available(), + "gpu_backend": gpu_info.backend, + "gpu_device": gpu_info.device_name, + "gpu_vram_gb": gpu_info.vram_gb, + "gpu_tier": gpu_info.tier, + "local_gpu_available": gpu_info.backend in ("cuda", "mps"), + "local_gpu_capabilities": gpu_info.capabilities, }, "remote": { "default_provider": default_name, diff --git a/backend/app/routers/gpu_status.py b/backend/app/routers/gpu_status.py new file mode 100644 index 0000000..3086c6b --- /dev/null +++ b/backend/app/routers/gpu_status.py @@ -0,0 +1,74 @@ +""" +GPU status and model management endpoints. +All under /api/gpu prefix. +""" +from fastapi import APIRouter +from pydantic import BaseModel +from typing import Optional, List +import asyncio + +router = APIRouter(prefix="/api/gpu", tags=["gpu"]) + + +@router.get("/status") +async def gpu_status(): + """ + Return GPU capabilities, VRAM, tier, and per-model download/ready state. + Frontend polls this to show GPU badge and tool availability. + """ + from app.services.gpu_detect import get_cached_gpu_info, get_model_ids + from app.services.local_diffusion import get_all_model_states + + info = get_cached_gpu_info() + model_ids = get_model_ids(info.tier) + + return { + "backend": info.backend, + "device_name": info.device_name, + "vram_gb": info.vram_gb, + "tier": info.tier, + "fp16": info.fp16, + "capabilities": info.capabilities, + "models": { + op: {"model_id": mid, "available": mid is not None} + for op, mid in model_ids.items() + }, + "pipeline_states": get_all_model_states(), + } + + +class PrefetchRequest(BaseModel): + operations: Optional[List[str]] = None + + +@router.post("/prefetch") +async def prefetch_models(req: PrefetchRequest = PrefetchRequest()): + """ + Kick off background model downloads for the requested operations. + Returns immediately; poll /api/gpu/prefetch-status for progress. + Default: prefetch inpaint, txt2img, img2img. + """ + ops = req.operations or ["inpaint", "txt2img", "img2img"] + valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"} + ops = [op for op in ops if op in valid] + + from app.services.local_diffusion import get_local_diffusion_provider + provider = get_local_diffusion_provider() + + async def _prefetch(): + for op in ops: + try: + await provider._get_pipeline(op) + print(f"[gpu] Prefetch complete: {op}") + except Exception as exc: + print(f"[gpu] Prefetch failed for {op}: {exc}") + + asyncio.create_task(_prefetch()) + return {"status": "prefetch_started", "operations": ops} + + +@router.get("/prefetch-status") +async def prefetch_status(): + """Poll model download / load progress.""" + from app.services.local_diffusion import get_all_model_states + return {"models": get_all_model_states()} diff --git a/backend/app/services/gpu_detect.py b/backend/app/services/gpu_detect.py new file mode 100644 index 0000000..f3f5a66 --- /dev/null +++ b/backend/app/services/gpu_detect.py @@ -0,0 +1,140 @@ +""" +GPU detection and capability tiering. +Detects CUDA (NVIDIA/AMD-ROCm), MPS (Apple Silicon), or CPU fallback. +Called once at startup; result is cached for the process lifetime. +""" +from __future__ import annotations + +import subprocess +from dataclasses import dataclass, field +from typing import Optional + + +# Model IDs per VRAM tier — all publicly available on HuggingFace, no auth needed. +# SDXL variants are used for high/ultra; SD 2.x for medium/low (smaller VRAM footprint). +_MODEL_TIERS: dict[str, dict[str, str]] = { + "ultra": { # ≥16 GB VRAM + "inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", + "txt2img": "stabilityai/stable-diffusion-xl-base-1.0", + "img2img": "stabilityai/stable-diffusion-xl-base-1.0", + "upscale": "stabilityai/stable-diffusion-x4-upscaler", + }, + "high": { # 8–16 GB VRAM + "inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", + "txt2img": "stabilityai/stable-diffusion-xl-base-1.0", + "img2img": "stabilityai/stable-diffusion-xl-base-1.0", + "upscale": "stabilityai/stable-diffusion-x4-upscaler", + }, + "medium": { # 4–8 GB VRAM + "inpaint": "stabilityai/stable-diffusion-2-inpainting", + "txt2img": "stabilityai/stable-diffusion-2-1", + "img2img": "stabilityai/stable-diffusion-2-1", + "upscale": None, + }, + "low": { # <4 GB or CPU + "inpaint": "stabilityai/stable-diffusion-2-inpainting", + "txt2img": "stabilityai/stable-diffusion-2-1-base", + "img2img": "stabilityai/stable-diffusion-2-1-base", + "upscale": None, + }, +} + + +@dataclass +class GpuInfo: + backend: str # cuda | mps | cpu + device_name: str = "CPU" + vram_gb: float = 0.0 + tier: str = "low" # ultra | high | medium | low + fp16: bool = False + capabilities: list[str] = field(default_factory=list) + + +def detect_gpu() -> GpuInfo: + """Detect available compute backend, VRAM, and assign a capability tier.""" + try: + import torch + + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + vram_gb = props.total_memory / (1024 ** 3) + tier = _vram_to_tier(vram_gb) + return GpuInfo( + backend="cuda", + device_name=props.name, + vram_gb=round(vram_gb, 1), + tier=tier, + fp16=True, + capabilities=_caps_for_tier(tier), + ) + + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + vram_gb = _apple_usable_gb() + tier = _vram_to_tier(vram_gb) + return GpuInfo( + backend="mps", + device_name="Apple Silicon", + vram_gb=round(vram_gb, 1), + tier=tier, + fp16=False, # MPS is more stable with fp32 + capabilities=_caps_for_tier(tier), + ) + + except ImportError: + pass + + return GpuInfo( + backend="cpu", + device_name="CPU (no GPU detected)", + vram_gb=0.0, + tier="low", + fp16=False, + capabilities=["txt2img", "inpaint", "img2img", "outpaint"], + ) + + +def _vram_to_tier(vram_gb: float) -> str: + if vram_gb >= 16: + return "ultra" + if vram_gb >= 8: + return "high" + if vram_gb >= 4: + return "medium" + return "low" + + +def _apple_usable_gb() -> float: + """Estimate GPU-usable unified memory on Apple Silicon (≈ half of total RAM).""" + try: + r = subprocess.run( + ["sysctl", "-n", "hw.memsize"], + capture_output=True, text=True, timeout=5, + ) + if r.returncode == 0: + return int(r.stdout.strip()) / (1024 ** 3) / 2 + except Exception: + pass + return 8.0 + + +def _caps_for_tier(tier: str) -> list[str]: + base = ["txt2img", "inpaint", "img2img", "outpaint"] + if tier in ("ultra", "high"): + return base + ["upscale_diffusion"] + return base + + +def get_model_ids(tier: str) -> dict[str, Optional[str]]: + """Return the model-ID map for a given tier.""" + return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["low"])) + + +# Process-level singleton — detect once, reuse everywhere. +_cached: Optional[GpuInfo] = None + + +def get_cached_gpu_info() -> GpuInfo: + global _cached + if _cached is None: + _cached = detect_gpu() + return _cached diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py new file mode 100644 index 0000000..e969fb4 --- /dev/null +++ b/backend/app/services/local_diffusion.py @@ -0,0 +1,472 @@ +""" +Local GPU diffusion provider — HuggingFace Diffusers backend. + +Implements the RemoteAIProvider interface so all existing routes work unchanged. +Models are lazy-loaded on first request and cached in memory. +VRAM-aware: picks the right model and memory optimisations per GPU tier. + +Requires: diffusers, transformers, accelerate, safetensors (requirements.gpu.txt) +""" +from __future__ import annotations + +import asyncio +import threading +from collections import OrderedDict +from io import BytesIO +from typing import Optional + +from PIL import Image + +from app.services.gpu_detect import get_cached_gpu_info, get_model_ids +from app.services.remote_provider import RemoteAIProvider + +# ── Download / load state tracking ────────────────────────────────────────── + +_states: dict[str, dict] = {} +_states_lock = threading.Lock() + + +def _set_state(key: str, **kw): + with _states_lock: + _states.setdefault(key, {}).update(kw) + + +def get_all_model_states() -> list[dict]: + with _states_lock: + return list(_states.values()) + + +# ── Pipeline cache with LRU eviction ───────────────────────────────────────── + +class _PipelineCache: + """Keep at most `maxsize` loaded pipelines; evicts LRU when full.""" + + def __init__(self, maxsize: int = 2): + self._cache: OrderedDict[str, object] = OrderedDict() + self._maxsize = maxsize + self._lock = asyncio.Lock() + + async def get(self, key: str): + async with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + return self._cache[key] + return None + + async def put(self, key: str, pipe: object): + async with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + else: + if len(self._cache) >= self._maxsize: + evicted_key, evicted_pipe = self._cache.popitem(last=False) + _offload_pipe(evicted_pipe, evicted_key) + self._cache[key] = pipe + + +def _offload_pipe(pipe, key: str): + """Move pipeline to CPU and free GPU memory.""" + try: + import torch + pipe.to("cpu") + torch.cuda.empty_cache() + print(f"[local_gpu] Evicted pipeline '{key}' from GPU cache") + except Exception: + pass + + +# ── Provider ───────────────────────────────────────────────────────────────── + +class LocalDiffusionProvider(RemoteAIProvider): + """ + HuggingFace Diffusers local inference. + All operations run in a thread pool to avoid blocking the event loop. + """ + + def __init__(self, max_cached_pipelines: int = 2): + self._cache = _PipelineCache(maxsize=max_cached_pipelines) + self._load_locks: dict[str, asyncio.Lock] = {} + self._meta_lock = asyncio.Lock() + + # ── Internal helpers ────────────────────────────────────────────────────── + + @property + def _info(self): + return get_cached_gpu_info() + + @property + def _device(self) -> str: + return self._info.backend + + def _torch_dtype(self): + import torch + return torch.float16 if self._info.fp16 else torch.float32 + + async def _lock_for(self, key: str) -> asyncio.Lock: + async with self._meta_lock: + if key not in self._load_locks: + self._load_locks[key] = asyncio.Lock() + return self._load_locks[key] + + def _load_pipeline_sync(self, pipe_type: str) -> object: + """Synchronous model load — runs in thread pool so HF download progress works.""" + import torch + from diffusers import ( + StableDiffusionInpaintPipeline, + StableDiffusionXLInpaintPipeline, + StableDiffusionPipeline, + StableDiffusionXLPipeline, + StableDiffusionImg2ImgPipeline, + StableDiffusionXLImg2ImgPipeline, + StableDiffusionUpscalePipeline, + ) + + info = self._info + tier = info.tier + device = self._device + dtype = self._torch_dtype() + model_ids = get_model_ids(tier) + + # Determine canonical operation key for inpaint-based ops + op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type + model_id = model_ids.get(op_key) + + # Allow config-level model override + try: + from app.config import settings + override_map = { + "inpaint": settings.hf_model_inpaint, + "outpaint": settings.hf_model_inpaint, + "txt2img": settings.hf_model_txt2img, + "img2img": settings.hf_model_img2img, + } + override = override_map.get(pipe_type, "") + if override: + model_id = override + except Exception: + pass + + if not model_id: + raise RuntimeError( + f"No model configured for '{pipe_type}' on tier '{tier}'. " + f"GPU may not have enough VRAM for this operation." + ) + + is_xl = "xl" in model_id.lower() + _set_state(pipe_type, pipeline=pipe_type, model_id=model_id, + state="downloading", progress=0.0, + message=f"Downloading {model_id}…", error="") + + try: + # Apply HuggingFace token if configured (needed for gated models) + try: + from app.config import settings + if settings.hf_token: + import huggingface_hub + huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False) + except Exception: + pass + + kwargs: dict = {"torch_dtype": dtype} + if not is_xl: + # Disable safety checker — we're editing existing images, not generating NSFW + kwargs["safety_checker"] = None + kwargs["requires_safety_checker"] = False + + if pipe_type == "inpaint" or pipe_type == "outpaint": + cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline + elif pipe_type == "txt2img": + cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline + elif pipe_type == "img2img": + cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline + elif pipe_type == "upscale": + model_id = model_ids.get("upscale") + if not model_id: + raise RuntimeError("Diffusion upscale model not available for this GPU tier.") + cls = StableDiffusionUpscalePipeline + else: + raise ValueError(f"Unknown pipeline type: {pipe_type}") + + pipe = cls.from_pretrained(model_id, **kwargs) + + # Move to device unless using CPU offload + if tier != "low" or device != "cpu": + pipe = pipe.to(device) + + # Memory optimisations + if tier in ("low", "medium"): + try: + pipe.enable_attention_slicing() + except Exception: + pass + if tier == "low" and device == "cuda": + try: + pipe.enable_sequential_cpu_offload() + except Exception: + pass + try: + pipe.enable_vae_slicing() + except Exception: + pass + + _set_state(pipe_type, state="ready", progress=100.0, message="Ready") + return pipe + + except Exception as exc: + _set_state(pipe_type, state="failed", error=str(exc), message="Load failed") + raise + + async def _get_pipeline(self, pipe_type: str) -> object: + cached = await self._cache.get(pipe_type) + if cached is not None: + return cached + + lock = await self._lock_for(pipe_type) + async with lock: + # Re-check after acquiring per-key lock + cached = await self._cache.get(pipe_type) + if cached is not None: + return cached + + loop = asyncio.get_event_loop() + pipe = await loop.run_in_executor(None, self._load_pipeline_sync, pipe_type) + await self._cache.put(pipe_type, pipe) + return pipe + + # ── RemoteAIProvider interface ──────────────────────────────────────────── + + async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: + pipe = await self._get_pipeline("inpaint") + info = self._info + + img = Image.open(BytesIO(image_bytes)).convert("RGB") + mask = Image.open(BytesIO(mask_bytes)).convert("L") + orig_size = img.size + + target = 1024 if info.tier in ("ultra", "high") else 512 + img_r, mask_r = _resize_pair(img, mask, target) + + steps = int(params.get("steps", 30)) + cfg = float(params.get("cfg_scale", 7.5)) + neg = params.get("negative_prompt", "") or None + + def _run(): + result = pipe( + prompt=prompt, + negative_prompt=neg, + image=img_r, + mask_image=mask_r, + num_inference_steps=steps, + guidance_scale=cfg, + ).images[0] + return result.resize(orig_size, Image.LANCZOS) + + loop = asyncio.get_event_loop() + result_img = await loop.run_in_executor(None, _run) + return _to_png(result_img) + + async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: + pipe = await self._get_pipeline("txt2img") + info = self._info + + max_dim = 1024 if info.tier in ("ultra", "high") else 768 + w = min(width, max_dim) // 8 * 8 + h = min(height, max_dim) // 8 * 8 + + steps = int(params.get("steps", 30)) + cfg = float(params.get("cfg_scale", 7.5)) + neg = params.get("negative_prompt", "") or None + seed = int(params.get("seed", 0)) + + device = self._device + + def _run(): + import torch + gen = torch.Generator(device=device).manual_seed(seed) if seed else None + return pipe( + prompt=prompt, + negative_prompt=neg, + width=w, + height=h, + num_inference_steps=steps, + guidance_scale=cfg, + generator=gen, + ).images[0] + + loop = asyncio.get_event_loop() + result_img = await loop.run_in_executor(None, _run) + return _to_png(result_img) + + async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: + pipe = await self._get_pipeline("img2img") + info = self._info + + img = Image.open(BytesIO(image_bytes)).convert("RGB") + orig_size = img.size + target = 1024 if info.tier in ("ultra", "high") else 512 + img_r = _resize_square(img, target) + + steps = int(params.get("steps", 30)) + cfg = float(params.get("cfg_scale", 7.5)) + neg = params.get("negative_prompt", "") or None + + def _run(): + result = pipe( + prompt=prompt, + negative_prompt=neg, + image=img_r, + strength=strength, + num_inference_steps=steps, + guidance_scale=cfg, + ).images[0] + return result.resize(orig_size, Image.LANCZOS) + + loop = asyncio.get_event_loop() + result_img = await loop.run_in_executor(None, _run) + return _to_png(result_img) + + async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: + from PIL import ImageDraw + + img = Image.open(BytesIO(image_bytes)).convert("RGB") + w, h = img.size + + if direction == "right": + new_size = (w + size, h) + paste_at = (0, 0) + mask_box = (w, 0, w + size, h) + elif direction == "left": + new_size = (w + size, h) + paste_at = (size, 0) + mask_box = (0, 0, size, h) + elif direction == "bottom": + new_size = (w, h + size) + paste_at = (0, 0) + mask_box = (0, h, w, h + size) + else: # top + new_size = (w, h + size) + paste_at = (0, size) + mask_box = (0, 0, w, size) + + expanded = Image.new("RGB", new_size, (127, 127, 127)) + expanded.paste(img, paste_at) + + mask = Image.new("L", new_size, 0) + draw = ImageDraw.Draw(mask) + draw.rectangle(mask_box, fill=255) + + params: dict = {} + fill_prompt = prompt or "seamless natural continuation of the scene" + result = await self.inpaint( + _to_png(expanded), _to_png(mask), fill_prompt, params + ) + return result + + async def health(self) -> bool: + return True + + def capabilities(self) -> list[str]: + return self._info.capabilities + + +# ── Image helpers ───────────────────────────────────────────────────────────── + +def _resize_pair( + img: Image.Image, mask: Image.Image, target: int +) -> tuple[Image.Image, Image.Image]: + """Resize image and mask so the longest side equals target, divisible by 8.""" + w, h = img.size + scale = target / max(w, h) + new_w = max(8, int(w * scale) // 8 * 8) + new_h = max(8, int(h * scale) // 8 * 8) + return ( + img.resize((new_w, new_h), Image.LANCZOS), + mask.resize((new_w, new_h), Image.NEAREST), + ) + + +def _resize_square(img: Image.Image, target: int) -> Image.Image: + w, h = img.size + scale = target / max(w, h) + new_w = max(8, int(w * scale) // 8 * 8) + new_h = max(8, int(h * scale) // 8 * 8) + return img.resize((new_w, new_h), Image.LANCZOS) + + +def _to_png(img: Image.Image) -> bytes: + buf = BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +# ── Singleton ───────────────────────────────────────────────────────────────── + +_provider: Optional[LocalDiffusionProvider] = None + + +def get_local_diffusion_provider(max_pipelines: int = 2) -> LocalDiffusionProvider: + global _provider + if _provider is None: + _provider = LocalDiffusionProvider(max_cached_pipelines=max_pipelines) + return _provider + + +async def prefetch_model_files() -> None: + """ + Download model weight files to the HuggingFace disk cache without loading + them into GPU memory. Run as a background task at container startup so the + first user request loads from disk (fast) rather than the internet (slow). + """ + from app.services.gpu_detect import get_cached_gpu_info, get_model_ids + + info = get_cached_gpu_info() + model_ids = get_model_ids(info.tier) + + try: + from huggingface_hub import snapshot_download + except ImportError: + print("[local_gpu] huggingface_hub not installed — skipping model prefetch") + return + + loop = asyncio.get_event_loop() + + # Override model IDs from config if provided + try: + from app.config import settings + overrides = { + "inpaint": settings.hf_model_inpaint, + "txt2img": settings.hf_model_txt2img, + "img2img": settings.hf_model_img2img, + } + for op, override in overrides.items(): + if override: + model_ids[op] = override + except Exception: + pass + + seen: set[str] = set() + for op, mid in model_ids.items(): + if not mid or mid in seen: + continue + seen.add(mid) + + _set_state(op, pipeline=op, model_id=mid, state="downloading", + progress=0.0, message=f"Downloading {mid}…", error="") + print(f"[local_gpu] Prefetching model files: {mid}") + + def _dl(repo_id: str = mid): + snapshot_download( + repo_id=repo_id, + # Skip TF/Flax/MsgPack variants — we only need PyTorch / safetensors + ignore_patterns=["*.msgpack", "flax_*", "tf_*", "rust_model*"], + ) + + try: + await loop.run_in_executor(None, _dl) + _set_state(op, state="cached", progress=100.0, + message="Files cached — will load into GPU on first request") + print(f"[local_gpu] ✓ Cached: {mid}") + except Exception as exc: + _set_state(op, state="download_failed", error=str(exc), + message="Download failed — will retry on first request") + print(f"[local_gpu] Prefetch failed for {mid}: {exc}") diff --git a/backend/app/services/remote_provider.py b/backend/app/services/remote_provider.py index 6fb8d31..a4107dc 100644 --- a/backend/app/services/remote_provider.py +++ b/backend/app/services/remote_provider.py @@ -428,6 +428,13 @@ def _build_provider(name: str) -> Optional[RemoteAIProvider]: return None return ComfyUIProvider(settings.comfyui_url, settings.comfyui_default_model) + if name == "local_gpu": + try: + from app.services.local_diffusion import get_local_diffusion_provider + return get_local_diffusion_provider(max_pipelines=settings.local_gpu_max_pipelines) + except ImportError: + return None + return None diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index 114f7c4..65f3d26 100644 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -61,6 +61,11 @@ else fi fi +echo "" +echo "Checking GPU capabilities..." +echo "------------------------------------------" +python /scripts/gpu_setup.py || echo "Warning: GPU detection failed (non-fatal)" + echo "" echo "==========================================" echo "Starting FastAPI server..." diff --git a/backend/requirements.gpu.txt b/backend/requirements.gpu.txt new file mode 100644 index 0000000..f3f6d67 --- /dev/null +++ b/backend/requirements.gpu.txt @@ -0,0 +1,24 @@ +# ============================================================================= +# GPU / Local Diffusion dependencies +# Install alongside requirements.txt when running with AI_PROVIDER=local_gpu +# +# Usage: +# pip install -r requirements.txt -r requirements.gpu.txt +# +# These are pre-installed in Dockerfile.gpu; optional in the standard image. +# ============================================================================= + +# HuggingFace Diffusers ecosystem +diffusers>=0.27.0 +transformers>=4.38.0 +accelerate>=0.27.0 +huggingface-hub>=0.21.0 +safetensors>=0.4.0 + +# Required by SDXL pipelines +invisible-watermark>=0.2.0 +omegaconf>=2.3.0 + +# xformers — further reduces VRAM usage on CUDA (install separately, version must +# match your PyTorch/CUDA; leave out if unsure and use attention_slicing instead) +# xformers diff --git a/docker-compose.gpu.yml b/docker-compose.gpu.yml new file mode 100644 index 0000000..57bd764 --- /dev/null +++ b/docker-compose.gpu.yml @@ -0,0 +1,94 @@ +# ============================================================================= +# EditmaskwithAI — GPU Docker Compose (NVIDIA CUDA) +# +# Quick start: +# docker compose -f docker-compose.gpu.yml up --build +# +# Then open: http://localhost:3080 +# +# What this does: +# • Detects your NVIDIA GPU at startup +# • Picks the best Stable Diffusion models for your VRAM tier +# • Auto-downloads models on first use (cached in a Docker volume) +# • Exposes local GPU generation (inpaint, outpaint, txt2img, img2img, upscale) +# • Still supports InvokeAI / ComfyUI / OpenAI via env vars below +# +# AMD ROCm: swap Dockerfile.gpu base image for a ROCm PyTorch image, +# remove the 'nvidia' driver line, and set device capabilities to [gpu]. +# ============================================================================= + +services: + app: + build: + context: . + dockerfile: Dockerfile.gpu + container_name: editmaskwithai-gpu + ports: + - "${PORT:-3080}:8000" + volumes: + # Persistent project data + - ./data:/app/data + # HuggingFace model cache — keeps downloaded models across rebuilds (~5-20 GB) + - hf_model_cache:/root/.cache/huggingface + # Scripts (for exec access) + - ./scripts:/scripts + environment: + # ── Local GPU (default for this compose) ──────────────────────────────── + - AI_PROVIDER=${AI_PROVIDER:-local_gpu} + - AUTO_DOWNLOAD_MODELS=${AUTO_DOWNLOAD_MODELS:-true} + + # ── Per-operation overrides (optional) ────────────────────────────────── + # Leave blank to use AI_PROVIDER for all operations. + # Example: use InvokeAI for inpaint, local GPU for everything else: + # AI_PROVIDER_INPAINT=invokeai + - AI_PROVIDER_INPAINT=${AI_PROVIDER_INPAINT:-} + - AI_PROVIDER_TXT2IMG=${AI_PROVIDER_TXT2IMG:-} + - AI_PROVIDER_IMG2IMG=${AI_PROVIDER_IMG2IMG:-} + - AI_PROVIDER_OUTPAINT=${AI_PROVIDER_OUTPAINT:-} + + # ── Remote/cloud providers (all optional) ──────────────────────────────── + - OPENAI_API_KEY=${OPENAI_API_KEY:-} + - OPENAI_MODEL=${OPENAI_MODEL:-dall-e-3} + - REPLICATE_API_KEY=${REPLICATE_API_KEY:-} + - STABILITY_API_KEY=${STABILITY_API_KEY:-} + + # ── InvokeAI / ComfyUI (running on another machine or container) ──────── + - INVOKEAI_URL=${INVOKEAI_URL:-} + - INVOKEAI_DEFAULT_MODEL=${INVOKEAI_DEFAULT_MODEL:-flux-dev} + - COMFYUI_URL=${COMFYUI_URL:-} + - COMFYUI_DEFAULT_MODEL=${COMFYUI_DEFAULT_MODEL:-v1-5-pruned-emaonly.ckpt} + + # ── HuggingFace model overrides (optional) ─────────────────────────────── + # Override the auto-selected model for any operation: + # HF_MODEL_INPAINT=your-org/your-model + - HF_MODEL_INPAINT=${HF_MODEL_INPAINT:-} + - HF_MODEL_TXT2IMG=${HF_MODEL_TXT2IMG:-} + - HF_MODEL_IMG2IMG=${HF_MODEL_IMG2IMG:-} + - HF_TOKEN=${HF_TOKEN:-} + + # ── App settings ───────────────────────────────────────────────────────── + - DATABASE_URL=sqlite:///./data/ai_photo_edit.db + - SECRET_KEY=${SECRET_KEY:-change-this-secret-key-in-production} + - CORS_ORIGINS=* + - AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true} + + # NVIDIA GPU passthrough — requires nvidia-container-toolkit on the host. + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: [gpu] + + # Reliable DNS for HuggingFace Hub downloads and external API calls + dns: + - 8.8.8.8 + - 8.8.4.4 + + restart: unless-stopped + +volumes: + hf_model_cache: + # Survives docker compose down; delete manually to free disk space: + # docker volume rm editmaskwithai_hf_model_cache diff --git a/scripts/gpu_setup.py b/scripts/gpu_setup.py new file mode 100644 index 0000000..e477cd9 --- /dev/null +++ b/scripts/gpu_setup.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +""" +GPU setup script — runs at container startup. +Detects GPU, logs capabilities, triggers background model prefetch when +AI_PROVIDER=local_gpu and AUTO_DOWNLOAD_MODELS=true. +Non-fatal: any failure just prints a warning. +""" +import os +import sys + + +def main(): + print("Detecting GPU…") + + backend = "cpu" + device_name = "CPU" + vram_gb = 0.0 + + try: + import torch + + if torch.cuda.is_available(): + backend = "cuda" + props = torch.cuda.get_device_properties(0) + device_name = props.name + vram_gb = props.total_memory / (1024 ** 3) + print(f"✓ CUDA GPU: {device_name} ({vram_gb:.1f} GB VRAM)") + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + backend = "mps" + device_name = "Apple Silicon" + print("✓ Apple Silicon MPS GPU detected") + else: + print("⚠ No GPU detected — AI_PROVIDER=local_gpu will use CPU (inference will be slow)") + + except ImportError: + print("⚠ PyTorch not installed — GPU detection skipped") + return + + provider = os.environ.get("AI_PROVIDER", "").lower() + if provider != "local_gpu": + print(f" AI_PROVIDER={provider!r} — local GPU inference not active") + return + + auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower() + if auto_dl != "true": + print(" AUTO_DOWNLOAD_MODELS=false — skipping model prefetch") + print(" Models will download on first request and cache to ~/.cache/huggingface") + return + + # Determine tier for a helpful startup message + if vram_gb >= 16: + tier, models_hint = "ultra", "SDXL (best quality)" + elif vram_gb >= 8: + tier, models_hint = "high", "SDXL" + elif vram_gb >= 4: + tier, models_hint = "medium", "Stable Diffusion 2.x" + else: + tier, models_hint = "low", "Stable Diffusion 2.x (small)" + + print(f" GPU tier: {tier} → will use {models_hint} models") + print(" Models will auto-download on first request (~2–7 GB per pipeline).") + print(" To pre-download now: POST /api/gpu/prefetch") + print(" Check progress at: GET /api/gpu/prefetch-status") + + +if __name__ == "__main__": + main() From 46b9066bba30016e23bb4fdcb90c8f9fc818369a Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 15:19:01 +0000 Subject: [PATCH 2/8] Handle old/low-VRAM GPUs and document nvidia-container-toolkit requirement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GPU tier table extended: ultra ≥16 GB → SDXL (unchanged) high 8-16 GB → SDXL (unchanged) medium 4-8 GB → SD 2.x (unchanged) legacy 2-4 GB → SD 1.5 (~1.7 GB fp16) ← new: GTX 970/1060/RX 580 etc. minimal <2 GB → SD 1.5 + sequential CPU offload ← new: very old/integrated GPUs gpu_detect.py: - Detects CUDA compute capability (CC); fp16 disabled for CC < 6.0 (pre-Pascal) - GpuInfo gains compute_capability and warnings fields - _make_warnings() emits human-readable warnings for low VRAM and old CC - model tier fallback updated from 'low' to 'legacy' local_diffusion.py: - minimal/legacy tiers use enable_sequential_cpu_offload() + enable_attention_slicing(1) - target resolution per tier: ultra/high=1024, medium=768, legacy/minimal=512 - .to(device) skipped when sequential CPU offload is active gpu_status.py: - Response now includes compute_capability and warnings docker-compose.gpu.yml: - Full nvidia-container-toolkit install instructions in header comment - nvidia-docker2 (legacy) fallback documented as comment block inline - AMD ROCm swap-in instructions added - GPU tier table documented in header scripts/gpu_setup.py: - Prints compute capability, fp16 status, tier, and model selection at startup - Prints per-tier warnings (old CC, low VRAM) https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/main.py | 7 ++- backend/app/routers/gpu_status.py | 2 + backend/app/services/gpu_detect.py | 80 +++++++++++++++++++----- backend/app/services/local_diffusion.py | 46 +++++++++----- docker-compose.gpu.yml | 67 ++++++++++++++++---- scripts/gpu_setup.py | 82 +++++++++++++++++-------- 6 files changed, 215 insertions(+), 69 deletions(-) diff --git a/backend/app/main.py b/backend/app/main.py index cfd06da..769607f 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -39,10 +39,13 @@ async def lifespan(app: FastAPI): ): from app.services.gpu_detect import get_cached_gpu_info info = get_cached_gpu_info() + cc_str = f" | CC={info.compute_capability}" if info.compute_capability else "" print( - f"[gpu] {info.device_name} | {info.vram_gb:.1f} GB | tier={info.tier} | " - f"backend={info.backend}" + f"[gpu] {info.device_name} | {info.vram_gb:.1f} GB{cc_str} | " + f"tier={info.tier} | fp16={info.fp16}" ) + for w in info.warnings: + print(f"[gpu] ⚠ {w}") if settings.auto_download_models: # Download model weight files to disk cache in background so first # user request loads from local disk instead of the internet. diff --git a/backend/app/routers/gpu_status.py b/backend/app/routers/gpu_status.py index 3086c6b..ce25a20 100644 --- a/backend/app/routers/gpu_status.py +++ b/backend/app/routers/gpu_status.py @@ -26,8 +26,10 @@ async def gpu_status(): "backend": info.backend, "device_name": info.device_name, "vram_gb": info.vram_gb, + "compute_capability": info.compute_capability, "tier": info.tier, "fp16": info.fp16, + "warnings": info.warnings, "capabilities": info.capabilities, "models": { op: {"model_id": mid, "available": mid is not None} diff --git a/backend/app/services/gpu_detect.py b/backend/app/services/gpu_detect.py index f3f5a66..94d0e22 100644 --- a/backend/app/services/gpu_detect.py +++ b/backend/app/services/gpu_detect.py @@ -11,30 +11,46 @@ from typing import Optional # Model IDs per VRAM tier — all publicly available on HuggingFace, no auth needed. -# SDXL variants are used for high/ultra; SD 2.x for medium/low (smaller VRAM footprint). +# +# Tier selection by VRAM: +# ultra ≥16 GB → SDXL (best quality) +# high 8–16 GB → SDXL +# medium 4–8 GB → SD 2.x +# legacy 2–4 GB → SD 1.5 (older / budget GPUs like GTX 970/1060/RX 580) +# minimal <2 GB → SD 1.5 with heavy memory offloading (very slow, but functional) +# +# SD 1.5 uses ~1.7 GB VRAM in fp16; SD 2.x uses ~3.5 GB; SDXL uses ~6.5 GB. _MODEL_TIERS: dict[str, dict[str, str]] = { - "ultra": { # ≥16 GB VRAM + "ultra": { "inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "txt2img": "stabilityai/stable-diffusion-xl-base-1.0", "img2img": "stabilityai/stable-diffusion-xl-base-1.0", "upscale": "stabilityai/stable-diffusion-x4-upscaler", }, - "high": { # 8–16 GB VRAM + "high": { "inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "txt2img": "stabilityai/stable-diffusion-xl-base-1.0", "img2img": "stabilityai/stable-diffusion-xl-base-1.0", "upscale": "stabilityai/stable-diffusion-x4-upscaler", }, - "medium": { # 4–8 GB VRAM + "medium": { "inpaint": "stabilityai/stable-diffusion-2-inpainting", "txt2img": "stabilityai/stable-diffusion-2-1", "img2img": "stabilityai/stable-diffusion-2-1", "upscale": None, }, - "low": { # <4 GB or CPU - "inpaint": "stabilityai/stable-diffusion-2-inpainting", - "txt2img": "stabilityai/stable-diffusion-2-1-base", - "img2img": "stabilityai/stable-diffusion-2-1-base", + # GTX 970 / GTX 1060 6 GB / RX 580 / etc. — 2–4 GB VRAM + "legacy": { + "inpaint": "runwayml/stable-diffusion-inpainting", + "txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "upscale": None, + }, + # Very old / integrated GPUs with <2 GB — runs but slowly; warns user. + "minimal": { + "inpaint": "runwayml/stable-diffusion-inpainting", + "txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", "upscale": None, }, } @@ -45,26 +61,35 @@ class GpuInfo: backend: str # cuda | mps | cpu device_name: str = "CPU" vram_gb: float = 0.0 - tier: str = "low" # ultra | high | medium | low + compute_capability: str = "" # e.g. "8.6" for RTX 3070 + tier: str = "legacy" # ultra | high | medium | legacy | minimal fp16: bool = False + warnings: list[str] = field(default_factory=list) capabilities: list[str] = field(default_factory=list) def detect_gpu() -> GpuInfo: - """Detect available compute backend, VRAM, and assign a capability tier.""" + """Detect available compute backend, VRAM, compute capability, and assign tier.""" try: import torch if torch.cuda.is_available(): props = torch.cuda.get_device_properties(0) vram_gb = props.total_memory / (1024 ** 3) + cc = f"{props.major}.{props.minor}" + # fp16 inference is reliable on Pascal (6.0) and newer. + # Maxwell (5.x) technically works but is slower in fp16 than fp32 on some ops. + use_fp16 = props.major >= 6 tier = _vram_to_tier(vram_gb) + warnings = _make_warnings(tier, vram_gb, cc, use_fp16) return GpuInfo( backend="cuda", device_name=props.name, vram_gb=round(vram_gb, 1), + compute_capability=cc, tier=tier, - fp16=True, + fp16=use_fp16, + warnings=warnings, capabilities=_caps_for_tier(tier), ) @@ -75,8 +100,9 @@ def detect_gpu() -> GpuInfo: backend="mps", device_name="Apple Silicon", vram_gb=round(vram_gb, 1), + compute_capability="mps", tier=tier, - fp16=False, # MPS is more stable with fp32 + fp16=False, # MPS diffusion is more stable with fp32 capabilities=_caps_for_tier(tier), ) @@ -87,8 +113,9 @@ def detect_gpu() -> GpuInfo: backend="cpu", device_name="CPU (no GPU detected)", vram_gb=0.0, - tier="low", + tier="minimal", fp16=False, + warnings=["No GPU found — running on CPU. Inference will be very slow (minutes per image)."], capabilities=["txt2img", "inpaint", "img2img", "outpaint"], ) @@ -100,7 +127,30 @@ def _vram_to_tier(vram_gb: float) -> str: return "high" if vram_gb >= 4: return "medium" - return "low" + if vram_gb >= 2: + return "legacy" + return "minimal" + + +def _make_warnings(tier: str, vram_gb: float, cc: str, fp16: bool) -> list[str]: + """Generate human-readable warnings for suboptimal GPU configurations.""" + warns = [] + if tier == "minimal": + warns.append( + f"Very low VRAM ({vram_gb:.1f} GB) — inference will be slow and may OOM. " + "Sequential CPU offloading will be enabled automatically." + ) + elif tier == "legacy": + warns.append( + f"Limited VRAM ({vram_gb:.1f} GB) — using SD 1.5 models (smaller, lower quality " + "than SD 2.x/SDXL). Still fully functional." + ) + if not fp16: + warns.append( + f"GPU compute capability {cc} is below 6.0 — using fp32 (doubles VRAM use). " + "Consider upgrading to a Pascal-era (GTX 1000) or newer GPU for fp16 support." + ) + return warns def _apple_usable_gb() -> float: @@ -126,7 +176,7 @@ def _caps_for_tier(tier: str) -> list[str]: def get_model_ids(tier: str) -> dict[str, Optional[str]]: """Return the model-ID map for a given tier.""" - return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["low"])) + return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["legacy"])) # Process-level singleton — detect once, reuse everywhere. diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py index e969fb4..73940d9 100644 --- a/backend/app/services/local_diffusion.py +++ b/backend/app/services/local_diffusion.py @@ -189,26 +189,37 @@ class LocalDiffusionProvider(RemoteAIProvider): pipe = cls.from_pretrained(model_id, **kwargs) - # Move to device unless using CPU offload - if tier != "low" or device != "cpu": - pipe = pipe.to(device) - - # Memory optimisations - if tier in ("low", "medium"): - try: - pipe.enable_attention_slicing() - except Exception: - pass - if tier == "low" and device == "cuda": - try: - pipe.enable_sequential_cpu_offload() - except Exception: - pass + # Memory optimisations — applied based on VRAM tier: + # minimal/legacy: full aggressive offloading (sequential CPU offload) + # medium: attention slicing + VAE slicing + # high/ultra: VAE slicing only (VRAM is plentiful) try: pipe.enable_vae_slicing() except Exception: pass + if tier in ("minimal", "legacy", "medium"): + try: + pipe.enable_attention_slicing(1) # slice_size=1 = most aggressive + except Exception: + pass + + if tier in ("minimal", "legacy"): + # Sequential CPU offload keeps only the active layer on GPU — very low VRAM + # but adds overhead per-step. Skip .to(device) when this is active. + if device == "cuda": + try: + pipe.enable_sequential_cpu_offload() + except Exception: + # Fallback: model stays on CPU entirely + pass + elif device == "cpu": + pass # already on CPU + else: + pipe = pipe.to(device) + else: + pipe = pipe.to(device) + _set_state(pipe_type, state="ready", progress=100.0, message="Ready") return pipe @@ -246,6 +257,7 @@ class LocalDiffusionProvider(RemoteAIProvider): target = 1024 if info.tier in ("ultra", "high") else 512 img_r, mask_r = _resize_pair(img, mask, target) + steps = int(params.get("steps", 30)) cfg = float(params.get("cfg_scale", 7.5)) neg = params.get("negative_prompt", "") or None @@ -269,7 +281,7 @@ class LocalDiffusionProvider(RemoteAIProvider): pipe = await self._get_pipeline("txt2img") info = self._info - max_dim = 1024 if info.tier in ("ultra", "high") else 768 + max_dim = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) w = min(width, max_dim) // 8 * 8 h = min(height, max_dim) // 8 * 8 @@ -303,7 +315,7 @@ class LocalDiffusionProvider(RemoteAIProvider): img = Image.open(BytesIO(image_bytes)).convert("RGB") orig_size = img.size - target = 1024 if info.tier in ("ultra", "high") else 512 + target = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) img_r = _resize_square(img, target) steps = int(params.get("steps", 30)) diff --git a/docker-compose.gpu.yml b/docker-compose.gpu.yml index 57bd764..b44072e 100644 --- a/docker-compose.gpu.yml +++ b/docker-compose.gpu.yml @@ -1,20 +1,60 @@ # ============================================================================= # EditmaskwithAI — GPU Docker Compose (NVIDIA CUDA) # -# Quick start: +# ── PREREQUISITES ───────────────────────────────────────────────────────────── +# +# 1. NVIDIA driver ≥ 525 installed on the host +# Check: nvidia-smi +# +# 2. nvidia-container-toolkit installed and configured: +# (Ubuntu/Debian) +# curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \ +# | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-ctk.gpg +# curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \ +# | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-ctk.gpg] https://#g' \ +# | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list +# sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit +# sudo nvidia-ctk runtime configure --runtime=docker +# sudo systemctl restart docker +# +# (RHEL/Fedora/Rocky) +# sudo dnf install -y nvidia-container-toolkit +# sudo nvidia-ctk runtime configure --runtime=docker +# sudo systemctl restart docker +# +# 3. Verify GPU access in Docker: +# docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi +# +# ── QUICK START ─────────────────────────────────────────────────────────────── +# # docker compose -f docker-compose.gpu.yml up --build +# Then open: http://localhost:3080 # -# Then open: http://localhost:3080 +# ── OLDER DOCKER SETUPS (docker-compose v1 / nvidia-docker2) ───────────────── # -# What this does: -# • Detects your NVIDIA GPU at startup -# • Picks the best Stable Diffusion models for your VRAM tier -# • Auto-downloads models on first use (cached in a Docker volume) -# • Exposes local GPU generation (inpaint, outpaint, txt2img, img2img, upscale) -# • Still supports InvokeAI / ComfyUI / OpenAI via env vars below +# If you installed nvidia-docker2 (older approach) instead of nvidia-container-toolkit, +# replace the 'deploy:' block below with: +# +# runtime: nvidia +# environment: +# - NVIDIA_VISIBLE_DEVICES=all +# - NVIDIA_DRIVER_CAPABILITIES=compute,utility +# +# ── GPU TIER AUTO-SELECTION ─────────────────────────────────────────────────── +# +# ≥16 GB VRAM → SDXL (best quality) +# 8–16 GB → SDXL +# 4–8 GB → Stable Diffusion 2.x +# 2–4 GB → Stable Diffusion 1.5 (older GPUs: GTX 970/1060/RX 580) +# <2 GB → SD 1.5 + CPU offload (very slow — consider a remote provider) +# +# ── AMD ROCm ────────────────────────────────────────────────────────────────── +# +# Swap the base image in Dockerfile.gpu: +# FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime +# → FROM rocm/pytorch:rocm6.0_ubuntu22.04_py3.9_pytorch_2.1.0 +# Remove the 'driver: nvidia' line and add: device_ids: ['0'] # -# AMD ROCm: swap Dockerfile.gpu base image for a ROCm PyTorch image, -# remove the 'nvidia' driver line, and set device capabilities to [gpu]. # ============================================================================= services: @@ -72,7 +112,12 @@ services: - CORS_ORIGINS=* - AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true} - # NVIDIA GPU passthrough — requires nvidia-container-toolkit on the host. + # ── NVIDIA GPU passthrough ──────────────────────────────────────────────── + # Requires nvidia-container-toolkit; see prerequisites at top of this file. + # For older nvidia-docker2 setups, replace this block with: + # runtime: nvidia + # environment: + # - NVIDIA_VISIBLE_DEVICES=all deploy: resources: reservations: diff --git a/scripts/gpu_setup.py b/scripts/gpu_setup.py index e477cd9..486434c 100644 --- a/scripts/gpu_setup.py +++ b/scripts/gpu_setup.py @@ -1,9 +1,8 @@ #!/usr/bin/env python3 """ GPU setup script — runs at container startup. -Detects GPU, logs capabilities, triggers background model prefetch when -AI_PROVIDER=local_gpu and AUTO_DOWNLOAD_MODELS=true. -Non-fatal: any failure just prints a warning. +Detects GPU, logs capabilities and any warnings, reports the model tier. +Non-fatal: failures just print a warning and startup continues. """ import os import sys @@ -15,6 +14,8 @@ def main(): backend = "cpu" device_name = "CPU" vram_gb = 0.0 + compute_cap = "" + tier = "minimal" try: import torch @@ -24,13 +25,57 @@ def main(): props = torch.cuda.get_device_properties(0) device_name = props.name vram_gb = props.total_memory / (1024 ** 3) - print(f"✓ CUDA GPU: {device_name} ({vram_gb:.1f} GB VRAM)") + compute_cap = f"{props.major}.{props.minor}" + use_fp16 = props.major >= 6 + + if vram_gb >= 16: + tier = "ultra" + elif vram_gb >= 8: + tier = "high" + elif vram_gb >= 4: + tier = "medium" + elif vram_gb >= 2: + tier = "legacy" + else: + tier = "minimal" + + fp16_str = "fp16" if use_fp16 else "fp32 (CC<6.0)" + print(f"✓ CUDA GPU : {device_name}") + print(f" VRAM : {vram_gb:.1f} GB") + print(f" Compute : {compute_cap} ({fp16_str})") + print(f" Tier : {tier}") + + # Per-tier model summary + tier_info = { + "ultra": "SDXL inpaint + SDXL base (best quality, needs ≥16 GB)", + "high": "SDXL inpaint + SDXL base (needs ≥8 GB)", + "medium": "SD 2.x inpaint + SD 2.1 (needs ≥4 GB fp16)", + "legacy": "SD 1.5 inpaint + SD 1.5 base (2–4 GB — older GPU mode)", + "minimal": "SD 1.5 + sequential CPU offload (<2 GB — very slow)", + } + print(f" Models : {tier_info.get(tier, 'SD 1.5')}") + + if not use_fp16: + print( + " ⚠ GPU compute capability is below 6.0 (Pascal). " + "fp32 will be used, doubling VRAM requirements. " + "A GTX 1000-series or newer GPU would enable fp16." + ) + if tier in ("minimal", "legacy"): + print( + " ⚠ Low VRAM: sequential CPU offload will be enabled. " + "Expect 2–5 minutes per image on a legacy GPU." + ) + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): backend = "mps" device_name = "Apple Silicon" print("✓ Apple Silicon MPS GPU detected") + print(" Note: fp32 used (fp16 less stable on MPS)") + else: - print("⚠ No GPU detected — AI_PROVIDER=local_gpu will use CPU (inference will be slow)") + print("⚠ No GPU detected — AI_PROVIDER=local_gpu will run on CPU.") + print(" Expect 5–20 minutes per image. Consider using a remote provider instead.") except ImportError: print("⚠ PyTorch not installed — GPU detection skipped") @@ -38,29 +83,18 @@ def main(): provider = os.environ.get("AI_PROVIDER", "").lower() if provider != "local_gpu": - print(f" AI_PROVIDER={provider!r} — local GPU inference not active") + print(f" AI_PROVIDER={provider!r} — local GPU inference not active, skipping prefetch") return auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower() - if auto_dl != "true": - print(" AUTO_DOWNLOAD_MODELS=false — skipping model prefetch") - print(" Models will download on first request and cache to ~/.cache/huggingface") - return - - # Determine tier for a helpful startup message - if vram_gb >= 16: - tier, models_hint = "ultra", "SDXL (best quality)" - elif vram_gb >= 8: - tier, models_hint = "high", "SDXL" - elif vram_gb >= 4: - tier, models_hint = "medium", "Stable Diffusion 2.x" + if auto_dl == "true": + print("") + print(" AUTO_DOWNLOAD_MODELS=true — model weights will download in the background.") + print(" First request after download completes will load model into GPU (~20-60s).") + print(" Pre-download now : POST /api/gpu/prefetch") + print(" Check progress : GET /api/gpu/prefetch-status") else: - tier, models_hint = "low", "Stable Diffusion 2.x (small)" - - print(f" GPU tier: {tier} → will use {models_hint} models") - print(" Models will auto-download on first request (~2–7 GB per pipeline).") - print(" To pre-download now: POST /api/gpu/prefetch") - print(" Check progress at: GET /api/gpu/prefetch-status") + print(" AUTO_DOWNLOAD_MODELS=false — models will download on first request.") if __name__ == "__main__": From fe4d911a00727e10d4f7f29a5e2af3f3b0356af5 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 15:35:37 +0000 Subject: [PATCH 3/8] Dynamic GPU capability detection: probe CC, VRAM, feature flags, pick best model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces fixed tier table with real hardware probing and dynamic model selection. gpu_detect.py — complete rewrite: - Reads torch.cuda.get_device_properties + mem_get_info for actual free VRAM - Detects: fp16 (CC≥6.0), bf16 (CC≥8.0), fp8 (CC≥8.9 Ada/Hopper), int8 (CC≥7.0), tensor_cores (CC≥7.0), xformers presence - Pre-Pascal (CC<6.0): effective_vram halved (fp32 weights are 2× larger) - Subtracts 400MB driver overhead from free VRAM before model selection - _select_txt2img / _select_inpaint / _select_img2img / _select_upscale: eff≥20GB → FLUX.1-schnell (no offload) eff≥10GB → FLUX.1-schnell (model_cpu_offload) eff≥7.5GB → SDXL eff≥5.5GB → SDXL + attention_slicing eff≥3.5GB → SD 2.1 eff≥2.5GB → SD 2.1-base + attention_slicing eff≥1.7GB → SD 1.5 else → SD 1.5 + sequential_cpu_offload - ModelSpec carries: model_id, family, memory_opt, native_res, vram_fp16_gb - Warnings: old CC, pre-Pascal fp32, fp8 upgrade hint, xformers install tip - Compatibility shim get_model_ids() retained for existing callers - infer_spec_from_model_id() auto-detects family from HF_MODEL_* overrides local_diffusion.py — refactored to use ModelSpec: - Reads spec from GpuCapabilities.recommended[op] instead of tier table - FLUX.1-schnell: FluxPipeline / FluxImg2ImgPipeline, 4 steps, guidance=0.0 - SD families: family-aware pipeline class selection (sd15/sd2x/sdxl) - Memory opts applied per ModelSpec.memory_opt field - xformers attention enabled automatically when xformers detected gpu_status.py — richer response: - Exposes all feature flags (fp16/bf16/fp8/int8/tensor_cores/xformers) - Returns full ModelSpec per operation (model_id, family, memory_opt, native_res) ai_tools.py — /api/config exposes: - gpu_vram_total, gpu_vram_free, gpu_cc, gpu_fp16, gpu_bf16, gpu_fp8, gpu_tensor_cores, gpu_eff_vram, local_gpu_warnings requirements.gpu.txt: - diffusers bumped to >=0.29.0 (FLUX pipeline added in 0.29) - transformers bumped to >=4.40.0 - sentencepiece added (FLUX T5 tokenizer) scripts/gpu_setup.py: - Prints full model table at startup (op → model_id, family, memory_opt, res) - Shows all feature flags in one line https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/routers/ai_tools.py | 18 +- backend/app/routers/gpu_status.py | 52 ++- backend/app/services/gpu_detect.py | 429 +++++++++++++----- backend/app/services/local_diffusion.py | 553 ++++++++++++------------ backend/requirements.gpu.txt | 16 +- scripts/gpu_setup.py | 153 ++++--- 6 files changed, 730 insertions(+), 491 deletions(-) diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index 4e63b65..096903c 100644 --- a/backend/app/routers/ai_tools.py +++ b/backend/app/routers/ai_tools.py @@ -335,13 +335,21 @@ async def get_config(): "lama": lama_available(), "rembg": rembg_available(), "opencv": True, - "gpu_detected": gpu_available(), - "gpu_backend": gpu_info.backend, - "gpu_device": gpu_info.device_name, - "gpu_vram_gb": gpu_info.vram_gb, - "gpu_tier": gpu_info.tier, + "gpu_detected": gpu_available(), + "gpu_backend": gpu_info.backend, + "gpu_device": gpu_info.device_name, + "gpu_vram_total": gpu_info.vram_total_gb, + "gpu_vram_free": gpu_info.vram_free_gb, + "gpu_cc": gpu_info.compute_capability, + "gpu_fp16": gpu_info.fp16, + "gpu_bf16": gpu_info.bf16, + "gpu_fp8": gpu_info.fp8, + "gpu_tensor_cores": gpu_info.tensor_cores, + "gpu_tier": gpu_info.tier, + "gpu_eff_vram": gpu_info.effective_vram_gb, "local_gpu_available": gpu_info.backend in ("cuda", "mps"), "local_gpu_capabilities": gpu_info.capabilities, + "local_gpu_warnings": gpu_info.warnings, }, "remote": { "default_provider": default_name, diff --git a/backend/app/routers/gpu_status.py b/backend/app/routers/gpu_status.py index ce25a20..68316b9 100644 --- a/backend/app/routers/gpu_status.py +++ b/backend/app/routers/gpu_status.py @@ -13,29 +13,49 @@ router = APIRouter(prefix="/api/gpu", tags=["gpu"]) @router.get("/status") async def gpu_status(): """ - Return GPU capabilities, VRAM, tier, and per-model download/ready state. + Full GPU capability report: hardware, feature flags, VRAM budget, + and which model was selected for each operation. Frontend polls this to show GPU badge and tool availability. """ - from app.services.gpu_detect import get_cached_gpu_info, get_model_ids + from app.services.gpu_detect import get_cached_gpu_info from app.services.local_diffusion import get_all_model_states info = get_cached_gpu_info() - model_ids = get_model_ids(info.tier) return { - "backend": info.backend, - "device_name": info.device_name, - "vram_gb": info.vram_gb, - "compute_capability": info.compute_capability, - "tier": info.tier, - "fp16": info.fp16, - "warnings": info.warnings, - "capabilities": info.capabilities, - "models": { - op: {"model_id": mid, "available": mid is not None} - for op, mid in model_ids.items() + # Hardware + "backend": info.backend, + "device_name": info.device_name, + "vram_total_gb": info.vram_total_gb, + "vram_free_gb": info.vram_free_gb, + "compute_capability": info.compute_capability, + # Feature flags + "fp16": info.fp16, + "bf16": info.bf16, + "fp8": info.fp8, + "int8": info.int8, + "tensor_cores": info.tensor_cores, + "xformers": info.xformers, + # Derived + "effective_vram_gb": info.effective_vram_gb, + "tier": info.tier, + # Selected models per operation + "recommended": { + op: ( + { + "model_id": spec.model_id, + "family": spec.family, + "memory_opt": spec.memory_opt, + "native_res": spec.native_res, + "vram_fp16_gb": spec.vram_fp16_gb, + } + if spec else None + ) + for op, spec in info.recommended.items() }, "pipeline_states": get_all_model_states(), + "warnings": info.warnings, + "capabilities": info.capabilities, } @@ -46,9 +66,9 @@ class PrefetchRequest(BaseModel): @router.post("/prefetch") async def prefetch_models(req: PrefetchRequest = PrefetchRequest()): """ - Kick off background model downloads for the requested operations. + Eagerly load pipelines into GPU memory for the requested operations. Returns immediately; poll /api/gpu/prefetch-status for progress. - Default: prefetch inpaint, txt2img, img2img. + Default: inpaint, txt2img, img2img. """ ops = req.operations or ["inpaint", "txt2img", "img2img"] valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"} diff --git a/backend/app/services/gpu_detect.py b/backend/app/services/gpu_detect.py index 94d0e22..d9829be 100644 --- a/backend/app/services/gpu_detect.py +++ b/backend/app/services/gpu_detect.py @@ -1,7 +1,21 @@ """ -GPU detection and capability tiering. -Detects CUDA (NVIDIA/AMD-ROCm), MPS (Apple Silicon), or CPU fallback. -Called once at startup; result is cached for the process lifetime. +GPU capability detection and per-operation model selection. + +Probes the actual GPU — VRAM (total + free), CUDA compute capability, and +feature flags (fp16, bf16, fp8, int8, tensor cores) — then selects the +highest-quality model that fits for each operation. + +Model selection ladder (txt2img): + eff_vram ≥ 20 GB → FLUX.1-schnell (no offload) + eff_vram ≥ 10 GB → FLUX.1-schnell (model_cpu_offload, 2–3× slower but fits) + eff_vram ≥ 7.5 GB → SDXL base + eff_vram ≥ 5.5 GB → SDXL base + attention slicing + eff_vram ≥ 3.5 GB → Stable Diffusion 2.1 + eff_vram ≥ 2.5 GB → SD 2.1-base + attention slicing + eff_vram ≥ 1.7 GB → Stable Diffusion 1.5 + otherwise → SD 1.5 + sequential CPU offload + +Inpaint always uses SDXL/SD-family (no FLUX inpaint pipeline yet). """ from __future__ import annotations @@ -10,155 +24,315 @@ from dataclasses import dataclass, field from typing import Optional -# Model IDs per VRAM tier — all publicly available on HuggingFace, no auth needed. -# -# Tier selection by VRAM: -# ultra ≥16 GB → SDXL (best quality) -# high 8–16 GB → SDXL -# medium 4–8 GB → SD 2.x -# legacy 2–4 GB → SD 1.5 (older / budget GPUs like GTX 970/1060/RX 580) -# minimal <2 GB → SD 1.5 with heavy memory offloading (very slow, but functional) -# -# SD 1.5 uses ~1.7 GB VRAM in fp16; SD 2.x uses ~3.5 GB; SDXL uses ~6.5 GB. -_MODEL_TIERS: dict[str, dict[str, str]] = { - "ultra": { - "inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", - "txt2img": "stabilityai/stable-diffusion-xl-base-1.0", - "img2img": "stabilityai/stable-diffusion-xl-base-1.0", - "upscale": "stabilityai/stable-diffusion-x4-upscaler", - }, - "high": { - "inpaint": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", - "txt2img": "stabilityai/stable-diffusion-xl-base-1.0", - "img2img": "stabilityai/stable-diffusion-xl-base-1.0", - "upscale": "stabilityai/stable-diffusion-x4-upscaler", - }, - "medium": { - "inpaint": "stabilityai/stable-diffusion-2-inpainting", - "txt2img": "stabilityai/stable-diffusion-2-1", - "img2img": "stabilityai/stable-diffusion-2-1", - "upscale": None, - }, - # GTX 970 / GTX 1060 6 GB / RX 580 / etc. — 2–4 GB VRAM - "legacy": { - "inpaint": "runwayml/stable-diffusion-inpainting", - "txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", - "img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", - "upscale": None, - }, - # Very old / integrated GPUs with <2 GB — runs but slowly; warns user. - "minimal": { - "inpaint": "runwayml/stable-diffusion-inpainting", - "txt2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", - "img2img": "stable-diffusion-v1-5/stable-diffusion-v1-5", - "upscale": None, - }, -} - +# ── Model specification ─────────────────────────────────────────────────────── @dataclass -class GpuInfo: - backend: str # cuda | mps | cpu - device_name: str = "CPU" - vram_gb: float = 0.0 - compute_capability: str = "" # e.g. "8.6" for RTX 3070 - tier: str = "legacy" # ultra | high | medium | legacy | minimal - fp16: bool = False - warnings: list[str] = field(default_factory=list) - capabilities: list[str] = field(default_factory=list) +class ModelSpec: + """Everything needed to load and run one diffusion pipeline.""" + model_id: str + family: str # sd15 | sd2x | sdxl | flux + memory_opt: str # none | attention_slicing | model_cpu_offload | sequential_cpu_offload + native_res: int # 512 | 768 | 1024 + vram_fp16_gb: float # approx VRAM needed in fp16, no memory opts -def detect_gpu() -> GpuInfo: - """Detect available compute backend, VRAM, compute capability, and assign tier.""" +# ── GPU capability record ───────────────────────────────────────────────────── + +@dataclass +class GpuCapabilities: + # Hardware + backend: str # cuda | mps | cpu + device_name: str + vram_total_gb: float + vram_free_gb: float + compute_capability: str # "8.6", "7.5", "6.1" … + cc_major: int + cc_minor: int + + # Feature flags derived from compute capability + fp16: bool # reliable fp16 (CC ≥ 6.0; CC 5.x works but slower) + bf16: bool # native bf16 (CC ≥ 8.0) + fp8: bool # native fp8 (CC ≥ 8.9, Ada / Hopper) + int8: bool # efficient int8 (CC ≥ 7.0, needed for bitsandbytes) + tensor_cores: bool # tensor cores (CC ≥ 7.0, Volta+) + xformers: bool # xformers installed (reduces attention VRAM ~20-30%) + + # Derived budget + effective_vram_gb: float # free VRAM after overhead, halved if fp32-only + + # Human-readable tier label + tier: str # flux_full | flux_offload | sdxl | sdxl_low | sd2x | sd15 | minimal + + # Best model per operation + recommended: dict[str, Optional[ModelSpec]] + + # Metadata + warnings: list[str] + capabilities: list[str] + + +# ── Detection ───────────────────────────────────────────────────────────────── + +def detect_gpu() -> GpuCapabilities: + """Probe the GPU, return a fully populated GpuCapabilities.""" try: import torch if torch.cuda.is_available(): props = torch.cuda.get_device_properties(0) - vram_gb = props.total_memory / (1024 ** 3) + free_bytes, total_bytes = torch.cuda.mem_get_info(0) + vram_total = total_bytes / (1024 ** 3) + vram_free = free_bytes / (1024 ** 3) cc = f"{props.major}.{props.minor}" - # fp16 inference is reliable on Pascal (6.0) and newer. - # Maxwell (5.x) technically works but is slower in fp16 than fp32 on some ops. - use_fp16 = props.major >= 6 - tier = _vram_to_tier(vram_gb) - warnings = _make_warnings(tier, vram_gb, cc, use_fp16) - return GpuInfo( + major, minor = props.major, props.minor + + fp16 = major >= 6 # Pascal and newer have good fp16 + bf16 = major >= 8 # Ampere A100 / RTX 3000+ + fp8 = major > 8 or (major == 8 and minor >= 9) # Ada / Hopper + int8 = major >= 7 # Volta+ + tensor_cores = major >= 7 + + # Pre-Pascal (Maxwell CC 5.x): fp16 works but throughput is lower than fp32 + # on some Maxwell cards. Flag it so memory opt logic can account for it. + xf = _xformers_available() + + # Subtract driver/CUDA context overhead from free VRAM + overhead_gb = 0.4 + eff = max(0.0, vram_free - overhead_gb) + if not fp16: + eff /= 2.0 # fp32 weights are 2× larger + + tier = _tier_label(eff) + warnings = _build_warnings( + tier, vram_total, vram_free, cc, major, minor, fp16, bf16, fp8, xf + ) + + return GpuCapabilities( backend="cuda", device_name=props.name, - vram_gb=round(vram_gb, 1), + vram_total_gb=round(vram_total, 1), + vram_free_gb=round(vram_free, 1), compute_capability=cc, + cc_major=major, + cc_minor=minor, + fp16=fp16, + bf16=bf16, + fp8=fp8, + int8=int8, + tensor_cores=tensor_cores, + xformers=xf, + effective_vram_gb=round(eff, 1), tier=tier, - fp16=use_fp16, + recommended=_select_all_models(eff), warnings=warnings, - capabilities=_caps_for_tier(tier), + capabilities=_caps(tier), ) if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - vram_gb = _apple_usable_gb() - tier = _vram_to_tier(vram_gb) - return GpuInfo( + usable_gb = _apple_usable_gb() + eff = max(0.0, usable_gb - 0.5) + tier = _tier_label(eff) + return GpuCapabilities( backend="mps", device_name="Apple Silicon", - vram_gb=round(vram_gb, 1), + vram_total_gb=round(usable_gb, 1), + vram_free_gb=round(usable_gb, 1), compute_capability="mps", + cc_major=0, + cc_minor=0, + fp16=False, # MPS diffusion more stable in fp32 + bf16=False, + fp8=False, + int8=False, + tensor_cores=False, + xformers=False, + effective_vram_gb=round(eff / 2, 1), # fp32 on MPS tier=tier, - fp16=False, # MPS diffusion is more stable with fp32 - capabilities=_caps_for_tier(tier), + recommended=_select_all_models(eff / 2), + warnings=["Apple MPS: using fp32 (fp16 less stable). Models load slower."], + capabilities=_caps(tier), ) except ImportError: pass - return GpuInfo( + # CPU fallback + return GpuCapabilities( backend="cpu", - device_name="CPU (no GPU detected)", - vram_gb=0.0, + device_name="CPU (no GPU)", + vram_total_gb=0.0, + vram_free_gb=0.0, + compute_capability="", + cc_major=0, cc_minor=0, + fp16=False, bf16=False, fp8=False, int8=False, + tensor_cores=False, xformers=False, + effective_vram_gb=0.0, tier="minimal", - fp16=False, - warnings=["No GPU found — running on CPU. Inference will be very slow (minutes per image)."], + recommended=_select_all_models(0.0), + warnings=[ + "No GPU found. Running on CPU — expect 5–30 minutes per image. " + "Consider setting AI_PROVIDER to a remote/cloud provider instead." + ], capabilities=["txt2img", "inpaint", "img2img", "outpaint"], ) -def _vram_to_tier(vram_gb: float) -> str: - if vram_gb >= 16: - return "ultra" - if vram_gb >= 8: - return "high" - if vram_gb >= 4: - return "medium" - if vram_gb >= 2: - return "legacy" +# ── Model selection ─────────────────────────────────────────────────────────── + +def _select_all_models(eff_vram: float) -> dict[str, Optional[ModelSpec]]: + return { + "txt2img": _select_txt2img(eff_vram), + "img2img": _select_img2img(eff_vram), + "inpaint": _select_inpaint(eff_vram), + "outpaint": _select_inpaint(eff_vram), # shares inpaint pipeline + "upscale": _select_upscale(eff_vram), + } + + +def _select_txt2img(eff: float) -> ModelSpec: + # FLUX.1-schnell (Apache 2.0, 4-step distilled) + if eff >= 20.0: + return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "none", 1024, 20.0) + if eff >= 10.0: + return ModelSpec("black-forest-labs/FLUX.1-schnell", "flux", "model_cpu_offload", 1024, 20.0) + # SDXL base + if eff >= 7.5: + return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "none", 1024, 6.5) + if eff >= 5.5: + return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "attention_slicing", 1024, 6.5) + # SD 2.x + if eff >= 3.5: + return ModelSpec("stabilityai/stable-diffusion-2-1", "sd2x", "none", 768, 3.5) + if eff >= 2.5: + return ModelSpec("stabilityai/stable-diffusion-2-1-base", "sd2x", "attention_slicing", 512, 3.2) + # SD 1.5 + if eff >= 1.7: + return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "attention_slicing", 512, 1.7) + return ModelSpec("stable-diffusion-v1-5/stable-diffusion-v1-5", "sd15", "sequential_cpu_offload", 512, 1.7) + + +def _select_img2img(eff: float) -> ModelSpec: + # img2img uses the same model family as txt2img + s = _select_txt2img(eff) + # FLUX img2img uses a different pipeline class but same model weights + return s + + +def _select_inpaint(eff: float) -> ModelSpec: + # No FLUX inpaint pipeline available yet — SDXL is the ceiling + if eff >= 7.5: + return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "none", 1024, 6.5) + if eff >= 5.5: + return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "attention_slicing", 1024, 6.5) + if eff >= 3.5: + return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "none", 512, 3.5) + if eff >= 2.5: + return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "attention_slicing", 512, 3.5) + if eff >= 1.7: + return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "attention_slicing", 512, 1.7) + return ModelSpec("runwayml/stable-diffusion-inpainting", "sd15", "sequential_cpu_offload", 512, 1.7) + + +def _select_upscale(eff: float) -> Optional[ModelSpec]: + # SD x4 upscaler — needs ~2 GB fp16 PLUS headroom for the loaded inpaint/txt2img model. + # Only enable if eff_vram suggests room for it as a secondary pipeline. + if eff >= 6.0: + return ModelSpec("stabilityai/stable-diffusion-x4-upscaler", "sd2x", "attention_slicing", 512, 2.0) + return None # fall through to Real-ESRGAN + + +# ── Tier label (display only) ───────────────────────────────────────────────── + +def _tier_label(eff_vram: float) -> str: + if eff_vram >= 20: return "flux_full" + if eff_vram >= 10: return "flux_offload" + if eff_vram >= 7.5: return "sdxl" + if eff_vram >= 5.5: return "sdxl_low" + if eff_vram >= 3.5: return "sd2x" + if eff_vram >= 2.5: return "sd2x_low" + if eff_vram >= 1.7: return "sd15" return "minimal" -def _make_warnings(tier: str, vram_gb: float, cc: str, fp16: bool) -> list[str]: - """Generate human-readable warnings for suboptimal GPU configurations.""" - warns = [] +def _caps(tier: str) -> list[str]: + base = ["txt2img", "inpaint", "img2img", "outpaint"] + if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low"): + return base + ["upscale_diffusion"] + return base + + +# ── Warnings ────────────────────────────────────────────────────────────────── + +def _build_warnings( + tier: str, vram_total: float, vram_free: float, + cc: str, major: int, minor: int, + fp16: bool, bf16: bool, fp8: bool, xf: bool, +) -> list[str]: + w = [] + + if major < 5: + w.append( + f"GPU compute capability {cc} is not supported by PyTorch 2.x. " + "Upgrade to a Kepler/Maxwell-era or newer GPU (CC ≥ 5.0)." + ) + elif major < 6: + w.append( + f"GPU is Maxwell-era (CC {cc}). fp32 mode — models need 2× VRAM. " + "A Pascal GTX 1000-series or newer card enables fp16." + ) + elif not bf16 and tier in ("flux_full", "flux_offload"): + w.append( + f"GPU CC {cc}: FLUX runs in fp16 (bf16 needs CC ≥ 8.0). " + "Results are still good but Ampere/Ada GPUs are faster here." + ) + + if fp8 and tier in ("flux_full", "flux_offload"): + w.append( + "FP8 native support detected (Ada Lovelace / Hopper). " + "Set HF_MODEL_TXT2IMG=flux-community/flux.1-schnell-fp8 for ~40% VRAM reduction." + ) + if tier == "minimal": - warns.append( - f"Very low VRAM ({vram_gb:.1f} GB) — inference will be slow and may OOM. " - "Sequential CPU offloading will be enabled automatically." + w.append( + f"Very low effective VRAM ({vram_free:.1f} GB free). " + "Sequential CPU offload will be used — expect 10–30 min per image." ) - elif tier == "legacy": - warns.append( - f"Limited VRAM ({vram_gb:.1f} GB) — using SD 1.5 models (smaller, lower quality " - "than SD 2.x/SDXL). Still fully functional." + elif tier in ("sd15", "sd2x_low"): + w.append( + f"Limited VRAM ({vram_free:.1f} GB free). " + "Using SD 1.5/2.x. Upgrade to ≥5.5 GB free for SDXL quality." ) - if not fp16: - warns.append( - f"GPU compute capability {cc} is below 6.0 — using fp32 (doubles VRAM use). " - "Consider upgrading to a Pascal-era (GTX 1000) or newer GPU for fp16 support." + + if xf: + w.append( + "xformers detected — attention VRAM reduced ~20-30%. " + "You may be able to run a higher-tier model than listed." ) - return warns + else: + if tier in ("sdxl_low", "sd2x"): + w.append( + "xformers not installed. Install it (pip install xformers) to reduce " + "VRAM usage ~20-30% and potentially unlock the next model tier." + ) + + return w + + +# ── Helpers ─────────────────────────────────────────────────────────────────── + +def _xformers_available() -> bool: + try: + import xformers # noqa: F401 + return True + except ImportError: + return False def _apple_usable_gb() -> float: - """Estimate GPU-usable unified memory on Apple Silicon (≈ half of total RAM).""" + """Estimate GPU-usable unified memory (≈ half of total RAM).""" try: r = subprocess.run( - ["sysctl", "-n", "hw.memsize"], - capture_output=True, text=True, timeout=5, + ["sysctl", "-n", "hw.memsize"], capture_output=True, text=True, timeout=5 ) if r.returncode == 0: return int(r.stdout.strip()) / (1024 ** 3) / 2 @@ -167,24 +341,39 @@ def _apple_usable_gb() -> float: return 8.0 -def _caps_for_tier(tier: str) -> list[str]: - base = ["txt2img", "inpaint", "img2img", "outpaint"] - if tier in ("ultra", "high"): - return base + ["upscale_diffusion"] - return base +def infer_spec_from_model_id(model_id: str) -> ModelSpec: + """ + When the user supplies HF_MODEL_* overrides, infer the pipeline family + from naming conventions so the correct diffusers class is chosen. + """ + mid = model_id.lower() + if "flux" in mid: + return ModelSpec(model_id, "flux", "model_cpu_offload", 1024, 20.0) + if "xl" in mid or "sdxl" in mid: + return ModelSpec(model_id, "sdxl", "attention_slicing", 1024, 6.5) + if any(x in mid for x in ["sd-2", "sd2", "stable-diffusion-2", "-2-", "-2inpaint"]): + res = 512 if "base" in mid else 768 + return ModelSpec(model_id, "sd2x", "attention_slicing", res, 3.5) + return ModelSpec(model_id, "sd15", "attention_slicing", 512, 1.7) -def get_model_ids(tier: str) -> dict[str, Optional[str]]: - """Return the model-ID map for a given tier.""" - return dict(_MODEL_TIERS.get(tier, _MODEL_TIERS["legacy"])) +# ── Singleton ───────────────────────────────────────────────────────────────── + +_cached: Optional[GpuCapabilities] = None -# Process-level singleton — detect once, reuse everywhere. -_cached: Optional[GpuInfo] = None - - -def get_cached_gpu_info() -> GpuInfo: +def get_cached_gpu_info() -> GpuCapabilities: global _cached if _cached is None: _cached = detect_gpu() return _cached + + +# Alias kept for any callers still using the old name +def get_model_ids(tier: str) -> dict: + """Compatibility shim — returns model_id strings keyed by operation.""" + info = get_cached_gpu_info() + return { + op: (spec.model_id if spec else None) + for op, spec in info.recommended.items() + } diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py index 73940d9..b873027 100644 --- a/backend/app/services/local_diffusion.py +++ b/backend/app/services/local_diffusion.py @@ -1,11 +1,18 @@ """ Local GPU diffusion provider — HuggingFace Diffusers backend. -Implements the RemoteAIProvider interface so all existing routes work unchanged. -Models are lazy-loaded on first request and cached in memory. -VRAM-aware: picks the right model and memory optimisations per GPU tier. +Implements RemoteAIProvider so all existing routes work unchanged. +Pipelines are lazy-loaded, cached in an LRU store, and memory-optimised +per the ModelSpec chosen by gpu_detect. -Requires: diffusers, transformers, accelerate, safetensors (requirements.gpu.txt) +Supported model families: + flux → FluxPipeline / FluxImg2ImgPipeline (FLUX.1-schnell) + sdxl → StableDiffusionXL*Pipeline (SDXL base + SDXL Inpaint) + sd2x → StableDiffusion2*Pipeline (SD 2.x) + sd15 → StableDiffusionPipeline (SD 1.5) + +Requires: diffusers>=0.29.0, transformers, accelerate, safetensors + (all in requirements.gpu.txt) """ from __future__ import annotations @@ -17,10 +24,15 @@ from typing import Optional from PIL import Image -from app.services.gpu_detect import get_cached_gpu_info, get_model_ids +from app.services.gpu_detect import ( + GpuCapabilities, + ModelSpec, + get_cached_gpu_info, + infer_spec_from_model_id, +) from app.services.remote_provider import RemoteAIProvider -# ── Download / load state tracking ────────────────────────────────────────── +# ── Model state tracking ────────────────────────────────────────────────────── _states: dict[str, dict] = {} _states_lock = threading.Lock() @@ -36,11 +48,9 @@ def get_all_model_states() -> list[dict]: return list(_states.values()) -# ── Pipeline cache with LRU eviction ───────────────────────────────────────── +# ── LRU pipeline cache ──────────────────────────────────────────────────────── class _PipelineCache: - """Keep at most `maxsize` loaded pipelines; evicts LRU when full.""" - def __init__(self, maxsize: int = 2): self._cache: OrderedDict[str, object] = OrderedDict() self._maxsize = maxsize @@ -59,49 +69,176 @@ class _PipelineCache: self._cache.move_to_end(key) else: if len(self._cache) >= self._maxsize: - evicted_key, evicted_pipe = self._cache.popitem(last=False) - _offload_pipe(evicted_pipe, evicted_key) + evicted_key, evicted = self._cache.popitem(last=False) + _evict(evicted, evicted_key) self._cache[key] = pipe -def _offload_pipe(pipe, key: str): - """Move pipeline to CPU and free GPU memory.""" +def _evict(pipe, key: str): try: import torch pipe.to("cpu") torch.cuda.empty_cache() - print(f"[local_gpu] Evicted pipeline '{key}' from GPU cache") + print(f"[local_gpu] Evicted '{key}' from GPU cache") except Exception: pass +# ── Pipeline loading helpers ────────────────────────────────────────────────── + +def _apply_hf_token(): + try: + from app.config import settings + if settings.hf_token: + import huggingface_hub + huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False) + except Exception: + pass + + +def _get_spec(pipe_type: str, info: GpuCapabilities) -> ModelSpec: + """Return the ModelSpec for a pipeline type, respecting user overrides.""" + # Map outpaint to inpaint (same pipeline) + op_key = "inpaint" if pipe_type == "outpaint" else pipe_type + # img2img uses same family/model as txt2img for FLUX/SDXL + if pipe_type == "img2img" and op_key not in info.recommended: + op_key = "txt2img" + + # User config override + try: + from app.config import settings + override_map = { + "inpaint": settings.hf_model_inpaint, + "outpaint": settings.hf_model_inpaint, + "txt2img": settings.hf_model_txt2img, + "img2img": settings.hf_model_img2img, + } + override_id = override_map.get(pipe_type, "") or "" + if override_id: + return infer_spec_from_model_id(override_id) + except Exception: + pass + + spec = info.recommended.get(op_key) + if spec is None: + raise RuntimeError( + f"No model available for '{pipe_type}' at effective VRAM " + f"{info.effective_vram_gb:.1f} GB. GPU may not have enough memory." + ) + return spec + + +def _load_sd_pipeline(pipe_type: str, spec: ModelSpec, info: GpuCapabilities) -> object: + """Load a Stable Diffusion (1.5 / 2.x / XL) pipeline.""" + import torch + from diffusers import ( + StableDiffusionPipeline, + StableDiffusionImg2ImgPipeline, + StableDiffusionInpaintPipeline, + StableDiffusionUpscalePipeline, + StableDiffusionXLPipeline, + StableDiffusionXLImg2ImgPipeline, + StableDiffusionXLInpaintPipeline, + ) + + dtype = torch.float16 if info.fp16 else torch.float32 + is_xl = spec.family == "sdxl" + kwargs: dict = {"torch_dtype": dtype} + if not is_xl: + kwargs["safety_checker"] = None + kwargs["requires_safety_checker"] = False + + op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type + + if op_key == "inpaint": + cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline + elif op_key == "txt2img": + cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline + elif op_key == "img2img": + cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline + elif op_key == "upscale": + cls = StableDiffusionUpscalePipeline + else: + raise ValueError(f"Unknown SD operation: {op_key}") + + pipe = cls.from_pretrained(spec.model_id, **kwargs) + return _apply_mem_opts(pipe, spec, info) + + +def _load_flux_pipeline(pipe_type: str, spec: ModelSpec, info: GpuCapabilities) -> object: + """Load a FLUX pipeline (txt2img or img2img).""" + import torch + from diffusers import FluxPipeline, FluxImg2ImgPipeline + + # FLUX works best in bf16 on Ampere+; fp16 on older Turing/Pascal + dtype = torch.bfloat16 if info.bf16 else torch.float16 + + op_key = "img2img" if pipe_type == "img2img" else "txt2img" + cls = FluxImg2ImgPipeline if op_key == "img2img" else FluxPipeline + + pipe = cls.from_pretrained(spec.model_id, torch_dtype=dtype) + return _apply_mem_opts(pipe, spec, info) + + +def _apply_mem_opts(pipe, spec: ModelSpec, info: GpuCapabilities) -> object: + """Apply memory optimisations then move pipeline to device.""" + device = info.backend + opt = spec.memory_opt + + # VAE slicing is always beneficial (reduces VRAM for decoding large images) + try: + pipe.enable_vae_slicing() + except Exception: + pass + + # xformers memory-efficient attention + if info.xformers and spec.family != "flux": + try: + pipe.enable_xformers_memory_efficient_attention() + except Exception: + pass + + if opt == "sequential_cpu_offload": + # Each layer moved to GPU only during its forward pass — very VRAM-efficient + # enable_sequential_cpu_offload() also calls .to(device) internally + try: + pipe.enable_sequential_cpu_offload() + except Exception: + pipe.to("cpu") + + elif opt == "model_cpu_offload": + # Entire sub-models (text encoder, unet/transformer, VAE) moved between CPU/GPU + # Faster than sequential but needs ~3-4 GB free to hold the active module + try: + pipe.enable_model_cpu_offload() + except Exception: + pipe.to(device) + + elif opt == "attention_slicing": + try: + pipe.enable_attention_slicing(1) + except Exception: + pass + pipe.to(device) + + else: # "none" + pipe.to(device) + + return pipe + + # ── Provider ───────────────────────────────────────────────────────────────── class LocalDiffusionProvider(RemoteAIProvider): - """ - HuggingFace Diffusers local inference. - All operations run in a thread pool to avoid blocking the event loop. - """ - def __init__(self, max_cached_pipelines: int = 2): self._cache = _PipelineCache(maxsize=max_cached_pipelines) self._load_locks: dict[str, asyncio.Lock] = {} self._meta_lock = asyncio.Lock() - # ── Internal helpers ────────────────────────────────────────────────────── - @property - def _info(self): + def _info(self) -> GpuCapabilities: return get_cached_gpu_info() - @property - def _device(self) -> str: - return self._info.backend - - def _torch_dtype(self): - import torch - return torch.float16 if self._info.fp16 else torch.float32 - async def _lock_for(self, key: str) -> asyncio.Lock: async with self._meta_lock: if key not in self._load_locks: @@ -109,120 +246,22 @@ class LocalDiffusionProvider(RemoteAIProvider): return self._load_locks[key] def _load_pipeline_sync(self, pipe_type: str) -> object: - """Synchronous model load — runs in thread pool so HF download progress works.""" - import torch - from diffusers import ( - StableDiffusionInpaintPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionPipeline, - StableDiffusionXLPipeline, - StableDiffusionImg2ImgPipeline, - StableDiffusionXLImg2ImgPipeline, - StableDiffusionUpscalePipeline, - ) - info = self._info - tier = info.tier - device = self._device - dtype = self._torch_dtype() - model_ids = get_model_ids(tier) + spec = _get_spec(pipe_type, info) - # Determine canonical operation key for inpaint-based ops - op_key = "inpaint" if pipe_type in ("inpaint", "outpaint") else pipe_type - model_id = model_ids.get(op_key) - - # Allow config-level model override - try: - from app.config import settings - override_map = { - "inpaint": settings.hf_model_inpaint, - "outpaint": settings.hf_model_inpaint, - "txt2img": settings.hf_model_txt2img, - "img2img": settings.hf_model_img2img, - } - override = override_map.get(pipe_type, "") - if override: - model_id = override - except Exception: - pass - - if not model_id: - raise RuntimeError( - f"No model configured for '{pipe_type}' on tier '{tier}'. " - f"GPU may not have enough VRAM for this operation." - ) - - is_xl = "xl" in model_id.lower() - _set_state(pipe_type, pipeline=pipe_type, model_id=model_id, + _apply_hf_token() + _set_state(pipe_type, pipeline=pipe_type, model_id=spec.model_id, + family=spec.family, memory_opt=spec.memory_opt, state="downloading", progress=0.0, - message=f"Downloading {model_id}…", error="") - + message=f"Downloading {spec.model_id}…", error="") try: - # Apply HuggingFace token if configured (needed for gated models) - try: - from app.config import settings - if settings.hf_token: - import huggingface_hub - huggingface_hub.login(token=settings.hf_token, add_to_git_credential=False) - except Exception: - pass - - kwargs: dict = {"torch_dtype": dtype} - if not is_xl: - # Disable safety checker — we're editing existing images, not generating NSFW - kwargs["safety_checker"] = None - kwargs["requires_safety_checker"] = False - - if pipe_type == "inpaint" or pipe_type == "outpaint": - cls = StableDiffusionXLInpaintPipeline if is_xl else StableDiffusionInpaintPipeline - elif pipe_type == "txt2img": - cls = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline - elif pipe_type == "img2img": - cls = StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline - elif pipe_type == "upscale": - model_id = model_ids.get("upscale") - if not model_id: - raise RuntimeError("Diffusion upscale model not available for this GPU tier.") - cls = StableDiffusionUpscalePipeline + if spec.family == "flux": + pipe = _load_flux_pipeline(pipe_type, spec, info) else: - raise ValueError(f"Unknown pipeline type: {pipe_type}") - - pipe = cls.from_pretrained(model_id, **kwargs) - - # Memory optimisations — applied based on VRAM tier: - # minimal/legacy: full aggressive offloading (sequential CPU offload) - # medium: attention slicing + VAE slicing - # high/ultra: VAE slicing only (VRAM is plentiful) - try: - pipe.enable_vae_slicing() - except Exception: - pass - - if tier in ("minimal", "legacy", "medium"): - try: - pipe.enable_attention_slicing(1) # slice_size=1 = most aggressive - except Exception: - pass - - if tier in ("minimal", "legacy"): - # Sequential CPU offload keeps only the active layer on GPU — very low VRAM - # but adds overhead per-step. Skip .to(device) when this is active. - if device == "cuda": - try: - pipe.enable_sequential_cpu_offload() - except Exception: - # Fallback: model stays on CPU entirely - pass - elif device == "cpu": - pass # already on CPU - else: - pipe = pipe.to(device) - else: - pipe = pipe.to(device) + pipe = _load_sd_pipeline(pipe_type, spec, info) _set_state(pipe_type, state="ready", progress=100.0, message="Ready") return pipe - except Exception as exc: _set_state(pipe_type, state="failed", error=str(exc), message="Load failed") raise @@ -234,108 +273,108 @@ class LocalDiffusionProvider(RemoteAIProvider): lock = await self._lock_for(pipe_type) async with lock: - # Re-check after acquiring per-key lock cached = await self._cache.get(pipe_type) if cached is not None: return cached - loop = asyncio.get_event_loop() pipe = await loop.run_in_executor(None, self._load_pipeline_sync, pipe_type) await self._cache.put(pipe_type, pipe) return pipe - # ── RemoteAIProvider interface ──────────────────────────────────────────── + # ── RemoteAIProvider ────────────────────────────────────────────────────── async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: pipe = await self._get_pipeline("inpaint") - info = self._info + spec = _get_spec("inpaint", self._info) - img = Image.open(BytesIO(image_bytes)).convert("RGB") + img = Image.open(BytesIO(image_bytes)).convert("RGB") mask = Image.open(BytesIO(mask_bytes)).convert("L") - orig_size = img.size - - target = 1024 if info.tier in ("ultra", "high") else 512 - img_r, mask_r = _resize_pair(img, mask, target) - + orig = img.size + img_r, mask_r = _resize_pair(img, mask, spec.native_res) steps = int(params.get("steps", 30)) - cfg = float(params.get("cfg_scale", 7.5)) - neg = params.get("negative_prompt", "") or None + cfg = float(params.get("cfg_scale", 7.5)) + neg = params.get("negative_prompt", "") or None def _run(): - result = pipe( + return pipe( prompt=prompt, negative_prompt=neg, image=img_r, mask_image=mask_r, num_inference_steps=steps, guidance_scale=cfg, - ).images[0] - return result.resize(orig_size, Image.LANCZOS) + ).images[0].resize(orig, Image.LANCZOS) - loop = asyncio.get_event_loop() - result_img = await loop.run_in_executor(None, _run) - return _to_png(result_img) + return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: pipe = await self._get_pipeline("txt2img") - info = self._info + spec = _get_spec("txt2img", self._info) - max_dim = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) - w = min(width, max_dim) // 8 * 8 + max_dim = spec.native_res + w = min(width, max_dim) // 8 * 8 h = min(height, max_dim) // 8 * 8 - - steps = int(params.get("steps", 30)) - cfg = float(params.get("cfg_scale", 7.5)) - neg = params.get("negative_prompt", "") or None seed = int(params.get("seed", 0)) - device = self._device + is_flux = spec.family == "flux" def _run(): import torch + device = self._info.backend gen = torch.Generator(device=device).manual_seed(seed) if seed else None - return pipe( - prompt=prompt, - negative_prompt=neg, - width=w, - height=h, - num_inference_steps=steps, - guidance_scale=cfg, - generator=gen, - ).images[0] - loop = asyncio.get_event_loop() - result_img = await loop.run_in_executor(None, _run) - return _to_png(result_img) + if is_flux: + return pipe( + prompt=prompt, + width=w, height=h, + num_inference_steps=4, # FLUX.1-schnell is a 4-step model + guidance_scale=0.0, # fully CFG-distilled + max_sequence_length=256, + generator=gen, + ).images[0] + else: + return pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + width=w, height=h, + num_inference_steps=int(params.get("steps", 30)), + guidance_scale=float(params.get("cfg_scale", 7.5)), + generator=gen, + ).images[0] + + return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: pipe = await self._get_pipeline("img2img") - info = self._info + spec = _get_spec("img2img", self._info) - img = Image.open(BytesIO(image_bytes)).convert("RGB") - orig_size = img.size - target = 1024 if info.tier in ("ultra", "high") else (768 if info.tier == "medium" else 512) - img_r = _resize_square(img, target) - - steps = int(params.get("steps", 30)) - cfg = float(params.get("cfg_scale", 7.5)) - neg = params.get("negative_prompt", "") or None + img = Image.open(BytesIO(image_bytes)).convert("RGB") + orig = img.size + img_r = _resize_square(img, spec.native_res) + is_flux = spec.family == "flux" def _run(): - result = pipe( - prompt=prompt, - negative_prompt=neg, - image=img_r, - strength=strength, - num_inference_steps=steps, - guidance_scale=cfg, - ).images[0] - return result.resize(orig_size, Image.LANCZOS) + if is_flux: + result = pipe( + prompt=prompt, + image=img_r, + strength=strength, + num_inference_steps=4, + guidance_scale=0.0, + ).images[0] + else: + result = pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + image=img_r, + strength=strength, + num_inference_steps=int(params.get("steps", 30)), + guidance_scale=float(params.get("cfg_scale", 7.5)), + ).images[0] + return result.resize(orig, Image.LANCZOS) - loop = asyncio.get_event_loop() - result_img = await loop.run_in_executor(None, _run) - return _to_png(result_img) + return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: from PIL import ImageDraw @@ -343,36 +382,21 @@ class LocalDiffusionProvider(RemoteAIProvider): img = Image.open(BytesIO(image_bytes)).convert("RGB") w, h = img.size - if direction == "right": - new_size = (w + size, h) - paste_at = (0, 0) - mask_box = (w, 0, w + size, h) - elif direction == "left": - new_size = (w + size, h) - paste_at = (size, 0) - mask_box = (0, 0, size, h) - elif direction == "bottom": - new_size = (w, h + size) - paste_at = (0, 0) - mask_box = (0, h, w, h + size) - else: # top - new_size = (w, h + size) - paste_at = (0, size) - mask_box = (0, 0, w, size) + positions = { + "right": ((w + size, h), (0, 0), (w, 0, w + size, h)), + "left": ((w + size, h), (size, 0), (0, 0, size, h)), + "bottom": ((w, h + size), (0, 0), (0, h, w, h + size)), + "top": ((w, h + size), (0, size), (0, 0, w, size)), + } + new_size, paste_at, mask_box = positions[direction] expanded = Image.new("RGB", new_size, (127, 127, 127)) expanded.paste(img, paste_at) - mask = Image.new("L", new_size, 0) - draw = ImageDraw.Draw(mask) - draw.rectangle(mask_box, fill=255) + ImageDraw.Draw(mask).rectangle(mask_box, fill=255) - params: dict = {} fill_prompt = prompt or "seamless natural continuation of the scene" - result = await self.inpaint( - _to_png(expanded), _to_png(mask), fill_prompt, params - ) - return result + return await self.inpaint(_to_png(expanded), _to_png(mask), fill_prompt, {}) async def health(self) -> bool: return True @@ -381,28 +405,22 @@ class LocalDiffusionProvider(RemoteAIProvider): return self._info.capabilities -# ── Image helpers ───────────────────────────────────────────────────────────── +# ── Image utilities ─────────────────────────────────────────────────────────── -def _resize_pair( - img: Image.Image, mask: Image.Image, target: int -) -> tuple[Image.Image, Image.Image]: - """Resize image and mask so the longest side equals target, divisible by 8.""" +def _resize_pair(img: Image.Image, mask: Image.Image, target: int): w, h = img.size scale = target / max(w, h) - new_w = max(8, int(w * scale) // 8 * 8) - new_h = max(8, int(h * scale) // 8 * 8) - return ( - img.resize((new_w, new_h), Image.LANCZOS), - mask.resize((new_w, new_h), Image.NEAREST), - ) + nw = max(8, int(w * scale) // 8 * 8) + nh = max(8, int(h * scale) // 8 * 8) + return img.resize((nw, nh), Image.LANCZOS), mask.resize((nw, nh), Image.NEAREST) def _resize_square(img: Image.Image, target: int) -> Image.Image: w, h = img.size scale = target / max(w, h) - new_w = max(8, int(w * scale) // 8 * 8) - new_h = max(8, int(h * scale) // 8 * 8) - return img.resize((new_w, new_h), Image.LANCZOS) + nw = max(8, int(w * scale) // 8 * 8) + nh = max(8, int(h * scale) // 8 * 8) + return img.resize((nw, nh), Image.LANCZOS) def _to_png(img: Image.Image) -> bytes: @@ -425,60 +443,59 @@ def get_local_diffusion_provider(max_pipelines: int = 2) -> LocalDiffusionProvid async def prefetch_model_files() -> None: """ - Download model weight files to the HuggingFace disk cache without loading - them into GPU memory. Run as a background task at container startup so the - first user request loads from disk (fast) rather than the internet (slow). + Download model weight files to HuggingFace disk cache without loading into GPU. + Called at container startup so the first request loads from disk (fast). """ - from app.services.gpu_detect import get_cached_gpu_info, get_model_ids - - info = get_cached_gpu_info() - model_ids = get_model_ids(info.tier) - + from app.services.gpu_detect import get_cached_gpu_info try: from huggingface_hub import snapshot_download except ImportError: print("[local_gpu] huggingface_hub not installed — skipping model prefetch") return + info = get_cached_gpu_info() + _apply_hf_token() + loop = asyncio.get_event_loop() - - # Override model IDs from config if provided - try: - from app.config import settings - overrides = { - "inpaint": settings.hf_model_inpaint, - "txt2img": settings.hf_model_txt2img, - "img2img": settings.hf_model_img2img, - } - for op, override in overrides.items(): - if override: - model_ids[op] = override - except Exception: - pass - seen: set[str] = set() - for op, mid in model_ids.items(): - if not mid or mid in seen: + + for op, spec in info.recommended.items(): + if spec is None or spec.model_id in seen: continue - seen.add(mid) + seen.add(spec.model_id) - _set_state(op, pipeline=op, model_id=mid, state="downloading", - progress=0.0, message=f"Downloading {mid}…", error="") - print(f"[local_gpu] Prefetching model files: {mid}") + # Apply user override if set + try: + from app.config import settings + override_map = { + "inpaint": settings.hf_model_inpaint, + "txt2img": settings.hf_model_txt2img, + "img2img": settings.hf_model_img2img, + } + override = override_map.get(op, "") or "" + if override and override not in seen: + seen.add(override) + spec = infer_spec_from_model_id(override) + except Exception: + pass - def _dl(repo_id: str = mid): + _set_state(op, pipeline=op, model_id=spec.model_id, family=spec.family, + memory_opt=spec.memory_opt, state="downloading", progress=0.0, + message=f"Downloading {spec.model_id}…", error="") + print(f"[local_gpu] Prefetching: {spec.model_id}") + + def _dl(model_id=spec.model_id): snapshot_download( - repo_id=repo_id, - # Skip TF/Flax/MsgPack variants — we only need PyTorch / safetensors + repo_id=model_id, ignore_patterns=["*.msgpack", "flax_*", "tf_*", "rust_model*"], ) try: await loop.run_in_executor(None, _dl) _set_state(op, state="cached", progress=100.0, - message="Files cached — will load into GPU on first request") - print(f"[local_gpu] ✓ Cached: {mid}") + message="Files cached — loads into GPU on first request") + print(f"[local_gpu] ✓ Cached: {spec.model_id}") except Exception as exc: _set_state(op, state="download_failed", error=str(exc), message="Download failed — will retry on first request") - print(f"[local_gpu] Prefetch failed for {mid}: {exc}") + print(f"[local_gpu] Prefetch failed for {spec.model_id}: {exc}") diff --git a/backend/requirements.gpu.txt b/backend/requirements.gpu.txt index f3f6d67..97bb4df 100644 --- a/backend/requirements.gpu.txt +++ b/backend/requirements.gpu.txt @@ -9,16 +9,22 @@ # ============================================================================= # HuggingFace Diffusers ecosystem -diffusers>=0.27.0 -transformers>=4.38.0 +# 0.29.0+ required for FLUX pipeline support +diffusers>=0.29.0 +transformers>=4.40.0 accelerate>=0.27.0 -huggingface-hub>=0.21.0 +huggingface-hub>=0.23.0 safetensors>=0.4.0 # Required by SDXL pipelines invisible-watermark>=0.2.0 omegaconf>=2.3.0 -# xformers — further reduces VRAM usage on CUDA (install separately, version must -# match your PyTorch/CUDA; leave out if unsure and use attention_slicing instead) +# Required by FLUX (T5 text encoder tokenizer) +sentencepiece>=0.2.0 + +# xformers — reduces attention VRAM ~20-30%, often unlocks the next model tier +# Must match your PyTorch+CUDA version; leave out if unsure. +# Install post-container-start if needed: +# pip install xformers --index-url https://download.pytorch.org/whl/cu121 # xformers diff --git a/scripts/gpu_setup.py b/scripts/gpu_setup.py index 486434c..a090223 100644 --- a/scripts/gpu_setup.py +++ b/scripts/gpu_setup.py @@ -1,100 +1,99 @@ #!/usr/bin/env python3 """ GPU setup script — runs at container startup. -Detects GPU, logs capabilities and any warnings, reports the model tier. -Non-fatal: failures just print a warning and startup continues. +Uses the same detection logic as the backend (gpu_detect.py) to show +exactly which models will be used before the server starts. +Non-fatal: any failure just prints a warning and startup continues. """ import os import sys def main(): - print("Detecting GPU…") - - backend = "cpu" - device_name = "CPU" - vram_gb = 0.0 - compute_cap = "" - tier = "minimal" + print("Detecting GPU capabilities…") try: import torch - - if torch.cuda.is_available(): - backend = "cuda" - props = torch.cuda.get_device_properties(0) - device_name = props.name - vram_gb = props.total_memory / (1024 ** 3) - compute_cap = f"{props.major}.{props.minor}" - use_fp16 = props.major >= 6 - - if vram_gb >= 16: - tier = "ultra" - elif vram_gb >= 8: - tier = "high" - elif vram_gb >= 4: - tier = "medium" - elif vram_gb >= 2: - tier = "legacy" - else: - tier = "minimal" - - fp16_str = "fp16" if use_fp16 else "fp32 (CC<6.0)" - print(f"✓ CUDA GPU : {device_name}") - print(f" VRAM : {vram_gb:.1f} GB") - print(f" Compute : {compute_cap} ({fp16_str})") - print(f" Tier : {tier}") - - # Per-tier model summary - tier_info = { - "ultra": "SDXL inpaint + SDXL base (best quality, needs ≥16 GB)", - "high": "SDXL inpaint + SDXL base (needs ≥8 GB)", - "medium": "SD 2.x inpaint + SD 2.1 (needs ≥4 GB fp16)", - "legacy": "SD 1.5 inpaint + SD 1.5 base (2–4 GB — older GPU mode)", - "minimal": "SD 1.5 + sequential CPU offload (<2 GB — very slow)", - } - print(f" Models : {tier_info.get(tier, 'SD 1.5')}") - - if not use_fp16: - print( - " ⚠ GPU compute capability is below 6.0 (Pascal). " - "fp32 will be used, doubling VRAM requirements. " - "A GTX 1000-series or newer GPU would enable fp16." - ) - if tier in ("minimal", "legacy"): - print( - " ⚠ Low VRAM: sequential CPU offload will be enabled. " - "Expect 2–5 minutes per image on a legacy GPU." - ) - - elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - backend = "mps" - device_name = "Apple Silicon" - print("✓ Apple Silicon MPS GPU detected") - print(" Note: fp32 used (fp16 less stable on MPS)") - - else: - print("⚠ No GPU detected — AI_PROVIDER=local_gpu will run on CPU.") - print(" Expect 5–20 minutes per image. Consider using a remote provider instead.") - except ImportError: print("⚠ PyTorch not installed — GPU detection skipped") return + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + free_b, total_b = torch.cuda.mem_get_info(0) + vram_total = total_b / (1024 ** 3) + vram_free = free_b / (1024 ** 3) + major, minor = props.major, props.minor + cc = f"{major}.{minor}" + + fp16 = major >= 6 + bf16 = major >= 8 + fp8 = major > 8 or (major == 8 and minor >= 9) + int8 = major >= 7 + tc = major >= 7 + + flags = [] + if fp16: flags.append("fp16") + if bf16: flags.append("bf16") + if fp8: flags.append("fp8") + if int8: flags.append("int8") + if tc: flags.append("tensor-cores") + + print(f"✓ GPU : {props.name}") + print(f" VRAM : {vram_total:.1f} GB total | {vram_free:.1f} GB free") + print(f" Compute : CC {cc} ({', '.join(flags) or 'fp32 only'})") + + if major < 6: + print(f" ⚠ Pre-Pascal (CC {cc}): using fp32 — effective VRAM budget halved") + + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + print("✓ Apple Silicon MPS GPU detected (fp32 mode)") + vram_total = vram_free = 0.0 + else: + print("⚠ No GPU detected — AI inference will use CPU (very slow)") + vram_total = vram_free = 0.0 + provider = os.environ.get("AI_PROVIDER", "").lower() if provider != "local_gpu": - print(f" AI_PROVIDER={provider!r} — local GPU inference not active, skipping prefetch") + print(f" AI_PROVIDER={provider!r} — local GPU not active, skipping model selection") return - auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower() - if auto_dl == "true": - print("") - print(" AUTO_DOWNLOAD_MODELS=true — model weights will download in the background.") - print(" First request after download completes will load model into GPU (~20-60s).") - print(" Pre-download now : POST /api/gpu/prefetch") - print(" Check progress : GET /api/gpu/prefetch-status") - else: - print(" AUTO_DOWNLOAD_MODELS=false — models will download on first request.") + # Import and run the full detection to show what was selected + try: + sys.path.insert(0, "/app") + from app.services.gpu_detect import detect_gpu + info = detect_gpu() + + print(f"\n Effective VRAM : {info.effective_vram_gb:.1f} GB (tier: {info.tier})") + print("\n Model selection:") + printed: set = set() + for op, spec in info.recommended.items(): + if spec is None: + print(f" {op:<12} → (none — will use existing upscaler)") + elif spec.model_id not in printed: + print(f" {op:<12} → [{spec.family}] {spec.model_id}") + print(f" mem_opt={spec.memory_opt} res={spec.native_res}px ~{spec.vram_fp16_gb}GB fp16") + printed.add(spec.model_id) + else: + print(f" {op:<12} → (same as above: {spec.model_id})") + + for w in info.warnings: + print(f"\n ⚠ {w}") + + auto_dl = os.environ.get("AUTO_DOWNLOAD_MODELS", "true").lower() + print() + if auto_dl == "true": + print(" AUTO_DOWNLOAD_MODELS=true") + print(" → Model files will download in background at startup.") + print(" → First request loads from local disk (20-60s, not internet).") + print(" → Track progress: GET /api/gpu/prefetch-status") + else: + print(" AUTO_DOWNLOAD_MODELS=false — models download on first request.") + + except Exception as exc: + print(f" (Could not run full detection: {exc})") + + print() if __name__ == "__main__": From 9b895673eb7d91ab5b71cac2ec2e5774864d2193 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 15:58:13 +0000 Subject: [PATCH 4/8] feat: GPU capability display in UI + GTX 1060 6GB SDXL fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Model selection: - Add sdxl_offload tier (eff_vram ≥ 4.0 GB) for GTX 1060 6GB and Quadro 6GB cards that were falling through to SD 2.1 despite SDXL fitting with model_cpu_offload. Cards with 5.3 GB effective VRAM now get SDXL quality. - Update _tier_label(), _caps(), _build_warnings() for new tier. Frontend GPU display: - api.js: add getGpuStatus() fetching /api/gpu/status - capabilities.js: add getGpuStatus() export with own LRU cache; refreshCapabilities() now also resets GPU status cache - provider-badge.js: when AI_PROVIDER=local_gpu show green badge with GPU name, tier, VRAM, CC, feature flags, and capabilities in tooltip. Strip "NVIDIA GeForce" prefix so "GTX 1060 6GB" fits in badge. - ai_provider_settings.js: add local_gpu to all provider dropdowns; show GPU info panel (device, VRAM, CC, features, tier, model table per operation) in the settings dialog when a GPU is detected. https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/services/gpu_detect.py | 18 +- frontend/src/js/api/capabilities.js | 21 +- .../src/js/core/components/provider-badge.js | 51 +++- .../js/modules/tools/ai_provider_settings.js | 265 +++++++++++------- frontend/src/js/services/api.js | 15 + 5 files changed, 262 insertions(+), 108 deletions(-) diff --git a/backend/app/services/gpu_detect.py b/backend/app/services/gpu_detect.py index d9829be..e4d495f 100644 --- a/backend/app/services/gpu_detect.py +++ b/backend/app/services/gpu_detect.py @@ -10,6 +10,7 @@ Model selection ladder (txt2img): eff_vram ≥ 10 GB → FLUX.1-schnell (model_cpu_offload, 2–3× slower but fits) eff_vram ≥ 7.5 GB → SDXL base eff_vram ≥ 5.5 GB → SDXL base + attention slicing + eff_vram ≥ 4.0 GB → SDXL + model_cpu_offload (GTX 1060 6 GB, Quadro 6 GB) eff_vram ≥ 3.5 GB → Stable Diffusion 2.1 eff_vram ≥ 2.5 GB → SD 2.1-base + attention slicing eff_vram ≥ 1.7 GB → Stable Diffusion 1.5 @@ -61,7 +62,7 @@ class GpuCapabilities: effective_vram_gb: float # free VRAM after overhead, halved if fp32-only # Human-readable tier label - tier: str # flux_full | flux_offload | sdxl | sdxl_low | sd2x | sd15 | minimal + tier: str # flux_full | flux_offload | sdxl | sdxl_low | sdxl_offload | sd2x | sd2x_low | sd15 | minimal # Best model per operation recommended: dict[str, Optional[ModelSpec]] @@ -200,6 +201,8 @@ def _select_txt2img(eff: float) -> ModelSpec: return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "none", 1024, 6.5) if eff >= 5.5: return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "attention_slicing", 1024, 6.5) + if eff >= 4.0: + return ModelSpec("stabilityai/stable-diffusion-xl-base-1.0", "sdxl", "model_cpu_offload", 1024, 6.5) # SD 2.x if eff >= 3.5: return ModelSpec("stabilityai/stable-diffusion-2-1", "sd2x", "none", 768, 3.5) @@ -224,6 +227,8 @@ def _select_inpaint(eff: float) -> ModelSpec: return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "none", 1024, 6.5) if eff >= 5.5: return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "attention_slicing", 1024, 6.5) + if eff >= 4.0: + return ModelSpec("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "sdxl", "model_cpu_offload", 1024, 6.5) if eff >= 3.5: return ModelSpec("stabilityai/stable-diffusion-2-inpainting", "sd2x", "none", 512, 3.5) if eff >= 2.5: @@ -248,6 +253,7 @@ def _tier_label(eff_vram: float) -> str: if eff_vram >= 10: return "flux_offload" if eff_vram >= 7.5: return "sdxl" if eff_vram >= 5.5: return "sdxl_low" + if eff_vram >= 4.0: return "sdxl_offload" if eff_vram >= 3.5: return "sd2x" if eff_vram >= 2.5: return "sd2x_low" if eff_vram >= 1.7: return "sd15" @@ -256,7 +262,7 @@ def _tier_label(eff_vram: float) -> str: def _caps(tier: str) -> list[str]: base = ["txt2img", "inpaint", "img2img", "outpaint"] - if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low"): + if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low", "sdxl_offload"): return base + ["upscale_diffusion"] return base @@ -297,6 +303,12 @@ def _build_warnings( f"Very low effective VRAM ({vram_free:.1f} GB free). " "Sequential CPU offload will be used — expect 10–30 min per image." ) + elif tier == "sdxl_offload": + w.append( + f"Limited VRAM ({vram_free:.1f} GB free). " + "Using SDXL with model_cpu_offload — better quality than SD 2.x, ~30% slower. " + "Install xformers or upgrade to ≥5.5 GB effective VRAM for full-speed SDXL." + ) elif tier in ("sd15", "sd2x_low"): w.append( f"Limited VRAM ({vram_free:.1f} GB free). " @@ -309,7 +321,7 @@ def _build_warnings( "You may be able to run a higher-tier model than listed." ) else: - if tier in ("sdxl_low", "sd2x"): + if tier in ("sdxl_low", "sdxl_offload", "sd2x"): w.append( "xformers not installed. Install it (pip install xformers) to reduce " "VRAM usage ~20-30% and potentially unlock the next model tier." diff --git a/frontend/src/js/api/capabilities.js b/frontend/src/js/api/capabilities.js index b49b423..1944e3e 100644 --- a/frontend/src/js/api/capabilities.js +++ b/frontend/src/js/api/capabilities.js @@ -19,6 +19,8 @@ const DEFAULT_CAPS = { let _caps = null; let _fetchPromise = null; +let _gpuStatus = null; +let _gpuFetchPromise = null; /** * Return capabilities (fetched lazily, cached thereafter). @@ -48,12 +50,29 @@ export function hasRemote() { return !!(_caps?.remote?.healthy); } +/** + * Fetch and cache detailed GPU status (hardware, feature flags, model selection per op). + * Calls /api/gpu/status — only meaningful when AI_PROVIDER=local_gpu. + * Returns null on error. + */ +export async function getGpuStatus() { + if (_gpuStatus !== null) return _gpuStatus; + if (!_gpuFetchPromise) { + _gpuFetchPromise = apiService.getGpuStatus() + .then(data => { _gpuStatus = data; return _gpuStatus; }) + .catch(() => { _gpuStatus = null; return null; }); + } + return _gpuFetchPromise; +} + /** * Invalidate cache and re-fetch (call after saving provider settings). */ export async function refreshCapabilities() { _caps = null; _fetchPromise = null; + _gpuStatus = null; + _gpuFetchPromise = null; return getCapabilities(); } @@ -62,4 +81,4 @@ export async function refreshCapabilities() { */ getCapabilities(); -export default { getCapabilities, getCachedCapabilities, hasRemote, refreshCapabilities }; +export default { getCapabilities, getCachedCapabilities, hasRemote, refreshCapabilities, getGpuStatus }; diff --git a/frontend/src/js/core/components/provider-badge.js b/frontend/src/js/core/components/provider-badge.js index 8724733..fc20790 100644 --- a/frontend/src/js/core/components/provider-badge.js +++ b/frontend/src/js/core/components/provider-badge.js @@ -2,7 +2,7 @@ * ProviderBadge — small DOM element showing the active AI provider. * Inserted into the toolbar footer on app load. * - * Green = remote provider healthy + * Green = remote provider healthy (or local_gpu active) * Yellow = provider configured but unhealthy/unreachable * Grey = local only (LaMa + OpenCV) */ @@ -30,21 +30,55 @@ export async function mountProviderBadge(container) { var remote = caps.remote || {}; var local = caps.local || {}; - if (remote.provider && remote.healthy) { + if (remote.provider === 'local_gpu') { + // Local GPU provider — show GPU name and tier from /api/config local fields + var gpuName = _shortGpuName(local.gpu_device); + var tier = local.gpu_tier || ''; + + if (remote.healthy) { + dot.style.background = '#44cc44'; + badge.style.background = '#1a2a1a'; + badge.style.color = '#aaffaa'; + label.textContent = 'GPU · ' + tier + ' · ' + gpuName; + + var flagList = [ + local.gpu_fp16 && 'fp16', + local.gpu_bf16 && 'bf16', + local.gpu_fp8 && 'fp8', + local.gpu_tensor_cores && 'tensor-cores', + ].filter(Boolean).join(' '); + + badge.title = [ + local.gpu_device || gpuName, + 'VRAM: ' + local.gpu_vram_total + ' GB total ' + local.gpu_vram_free + ' GB free', + 'Compute: CC ' + local.gpu_cc + ' Eff: ' + local.gpu_eff_vram + ' GB', + flagList ? 'Features: ' + flagList : '', + 'Capabilities: ' + (local.local_gpu_capabilities || []).join(', '), + (local.local_gpu_warnings || []).length + ? '\nWarnings:\n' + local.local_gpu_warnings.join('\n') + : '', + ].filter(Boolean).join('\n'); + } else { + dot.style.background = '#ffaa00'; + badge.style.background = '#2a2000'; + badge.style.color = '#ffdd88'; + label.textContent = 'Local GPU (not ready)'; + badge.title = 'local_gpu is configured but the diffusers library may not be installed.\nCheck container logs for details.'; + } + } else if (remote.provider && remote.healthy) { dot.style.background = '#44cc44'; badge.style.background = '#1a2a1a'; badge.style.color = '#aaffaa'; - // Show override summary if any operations use different providers var overrides = remote.overrides || {}; var overrideEntries = Object.entries(overrides).filter(([, v]) => v); var overrideStr = overrideEntries.length - ? ' · ' + overrideEntries.map(([k, v]) => `${k}→${v}`).join(', ') + ? ' · ' + overrideEntries.map(([k, v]) => k + '→' + v).join(', ') : ''; label.textContent = remote.provider + overrideStr + (local.gpu_detected ? ' · GPU' : ''); var opLines = Object.entries(remote.operations || {}) - .map(([op, s]) => `${op}: ${s.provider || remote.provider} ${s.healthy ? '✓' : '✗'}`) + .map(([op, s]) => op + ': ' + (s.provider || remote.provider) + ' ' + (s.healthy ? '✓' : '✗')) .join('\n'); badge.title = opLines || ('Provider: ' + remote.provider); } else if (remote.provider && !remote.healthy) { @@ -70,3 +104,10 @@ export async function mountProviderBadge(container) { return badge; } + +function _shortGpuName(name) { + return (name || 'GPU') + .replace(/^NVIDIA GeForce\s+/i, '') + .replace(/^NVIDIA\s+/i, '') + .replace(/^AMD Radeon\s+/i, ''); +} diff --git a/frontend/src/js/modules/tools/ai_provider_settings.js b/frontend/src/js/modules/tools/ai_provider_settings.js index 4268346..b4694e6 100644 --- a/frontend/src/js/modules/tools/ai_provider_settings.js +++ b/frontend/src/js/modules/tools/ai_provider_settings.js @@ -6,7 +6,7 @@ import Dialog_class from './../../libs/popup.js'; import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; -import { getCapabilities } from './../../api/capabilities.js'; +import { getCapabilities, getGpuStatus } from './../../api/capabilities.js'; // localStorage key prefix const LS = 'paintplus_ai_'; @@ -30,109 +30,133 @@ class Tools_ai_provider_settings_class { async ai_provider_settings() { var _this = this; + + // Fetch caps and GPU status in parallel var caps = await getCapabilities(); + var gpuStatus = null; + var local = caps.local || {}; + if (local.local_gpu_available) { + gpuStatus = await getGpuStatus().catch(() => null); + } + var remote = caps.remote || {}; var statusHtml = remote.provider ? (remote.healthy - ? `● ${remote.provider} — connected` - : `● ${remote.provider} — unreachable`) + ? '● ' + remote.provider + ' — connected' + : '● ' + remote.provider + ' — unreachable') : 'No remote provider configured'; + var gpuInfoHtml = gpuStatus ? _renderGpuInfo(gpuStatus) : ''; + + var providerValues = ['', 'openai', 'invokeai', 'comfyui', 'replicate', 'local_gpu']; + + var params = [ + { + title: 'Status:', + html: '
' + statusHtml + '
', + }, + ]; + + if (gpuInfoHtml) { + params.push({ + title: '', + html: '
Detected GPU:
' + gpuInfoHtml + '
', + }); + } + + params.push( + { + name: 'provider', + title: 'Default provider (used unless overridden below):', + value: ls_get('provider', remote.provider || ''), + values: providerValues, + type: 'select', + }, + // ── Per-operation overrides ─────────────────────────────── + { + title: '', + html: '
Per-operation overrides — blank = use default above
', + }, + { + name: 'provider_inpaint', + title: 'Inpaint / Replace Selection:', + value: ls_get('provider_inpaint', remote.overrides?.inpaint || ''), + values: providerValues, + type: 'select', + }, + { + name: 'provider_txt2img', + title: 'Text → Image:', + value: ls_get('provider_txt2img', remote.overrides?.txt2img || ''), + values: providerValues, + type: 'select', + }, + { + name: 'provider_img2img', + title: 'Image → Image:', + value: ls_get('provider_img2img', remote.overrides?.img2img || ''), + values: providerValues, + type: 'select', + }, + { + name: 'provider_outpaint', + title: 'Expand Canvas (Outpaint):', + value: ls_get('provider_outpaint', remote.overrides?.outpaint || ''), + values: providerValues, + type: 'select', + }, + // ── OpenAI ──────────────────────────────────────────────── + { + name: 'openai_key', + title: 'OpenAI API key:', + value: ls_get('openai_key'), + placeholder: 'sk-...', + }, + { + name: 'openai_model', + title: 'OpenAI model:', + value: ls_get('openai_model', 'dall-e-3'), + values: ['dall-e-3', 'dall-e-2'], + type: 'select', + }, + // ── InvokeAI ────────────────────────────────────────────── + { + name: 'invokeai_url', + title: 'InvokeAI URL:', + value: ls_get('invokeai_url'), + placeholder: 'http://192.168.1.x:9090', + }, + { + name: 'invokeai_model', + title: 'InvokeAI default model:', + value: ls_get('invokeai_model', 'flux-dev'), + placeholder: 'flux-dev', + }, + // ── ComfyUI ─────────────────────────────────────────────── + { + name: 'comfyui_url', + title: 'ComfyUI URL:', + value: ls_get('comfyui_url'), + placeholder: 'http://192.168.1.x:8188', + }, + { + name: 'comfyui_model', + title: 'ComfyUI default checkpoint:', + value: ls_get('comfyui_model', 'v1-5-pruned-emaonly.ckpt'), + placeholder: 'v1-5-pruned-emaonly.ckpt', + }, + // ── Replicate ───────────────────────────────────────────── + { + name: 'replicate_key', + title: 'Replicate API key:', + value: ls_get('replicate_key'), + placeholder: 'r8_...', + } + ); + this.POP.show({ title: 'AI Provider Settings', - params: [ - { - title: 'Status:', - html: `
${statusHtml}
`, - }, - { - name: 'provider', - title: 'Default provider (used unless overridden below):', - value: ls_get('provider', remote.provider || ''), - values: ['', 'openai', 'invokeai', 'comfyui', 'replicate'], - type: 'select', - }, - // ── Per-operation overrides ─────────────────────────────── - { - title: '', - html: '
Per-operation overrides — blank = use default above
', - }, - { - name: 'provider_inpaint', - title: 'Inpaint / Replace Selection:', - value: ls_get('provider_inpaint', remote.overrides?.inpaint || ''), - values: ['', 'openai', 'invokeai', 'comfyui', 'replicate'], - type: 'select', - }, - { - name: 'provider_txt2img', - title: 'Text → Image:', - value: ls_get('provider_txt2img', remote.overrides?.txt2img || ''), - values: ['', 'openai', 'invokeai', 'comfyui', 'replicate'], - type: 'select', - }, - { - name: 'provider_img2img', - title: 'Image → Image:', - value: ls_get('provider_img2img', remote.overrides?.img2img || ''), - values: ['', 'openai', 'invokeai', 'comfyui', 'replicate'], - type: 'select', - }, - { - name: 'provider_outpaint', - title: 'Expand Canvas (Outpaint):', - value: ls_get('provider_outpaint', remote.overrides?.outpaint || ''), - values: ['', 'openai', 'invokeai', 'comfyui', 'replicate'], - type: 'select', - }, - // ── OpenAI ──────────────────────────────────────────────── - { - name: 'openai_key', - title: 'OpenAI API key:', - value: ls_get('openai_key'), - placeholder: 'sk-...', - }, - { - name: 'openai_model', - title: 'OpenAI model:', - value: ls_get('openai_model', 'dall-e-3'), - values: ['dall-e-3', 'dall-e-2'], - type: 'select', - }, - // ── InvokeAI ────────────────────────────────────────────── - { - name: 'invokeai_url', - title: 'InvokeAI URL:', - value: ls_get('invokeai_url'), - placeholder: 'http://192.168.1.x:9090', - }, - { - name: 'invokeai_model', - title: 'InvokeAI default model:', - value: ls_get('invokeai_model', 'flux-dev'), - placeholder: 'flux-dev', - }, - // ── ComfyUI ─────────────────────────────────────────────── - { - name: 'comfyui_url', - title: 'ComfyUI URL:', - value: ls_get('comfyui_url'), - placeholder: 'http://192.168.1.x:8188', - }, - { - name: 'comfyui_model', - title: 'ComfyUI default checkpoint:', - value: ls_get('comfyui_model', 'v1-5-pruned-emaonly.ckpt'), - placeholder: 'v1-5-pruned-emaonly.ckpt', - }, - // ── Replicate ───────────────────────────────────────────── - { - name: 'replicate_key', - title: 'Replicate API key:', - value: ls_get('replicate_key'), - placeholder: 'r8_...', - }, - ], + params: params, on_finish: async function (params) { await _this._save(params); }, @@ -182,12 +206,15 @@ class Tools_ai_provider_settings_class { var { refreshCapabilities } = await import('./../../api/capabilities.js'); var caps = await refreshCapabilities(); if (caps?.remote?.healthy) { - alertify.success(`Connected to ${caps.remote.provider}!`); + alertify.success('Connected to ' + caps.remote.provider + '!'); } else if (params.provider) { - alertify.warning('Settings saved but provider is not reachable. Check URL/key.'); + if (params.provider === 'local_gpu') { + alertify.success('local_gpu set — restart the container with docker-compose.gpu.yml to activate.'); + } else { + alertify.warning('Settings saved but provider is not reachable. Check URL/key.'); + } } } else { - // Server-side config update not supported — inform user to set .env alertify.warning( 'Settings saved locally. To make them permanent, ' + 'set these values in your .env file and restart the server.' @@ -201,4 +228,44 @@ class Tools_ai_provider_settings_class { } } +function _renderGpuInfo(g) { + var flags = [ + g.fp16 && 'fp16', + g.bf16 && 'bf16', + g.fp8 && 'fp8', + g.int8 && 'int8', + g.tensor_cores && 'tensor-cores', + g.xformers && 'xformers', + ].filter(Boolean).join(' · '); + + var rows = Object.entries(g.recommended || {}) + .filter(([, s]) => s) + .map(function([op, s]) { + var modelName = s.model_id.split('/').pop(); + return '' + + '' + op + '' + + '' + modelName + '' + + '' + s.memory_opt + '' + + ''; + }) + .join(''); + + var warnHtml = (g.warnings || []).length + ? '
' + + g.warnings.map(function(w) { return '⚠ ' + w; }).join('
') + '
' + : ''; + + return '
' + + '
⬛ ' + (g.device_name || 'GPU') + '
' + + '
VRAM: ' + g.vram_total_gb + ' GB total · ' + g.vram_free_gb + ' GB free
' + + '
Compute: CC ' + g.compute_capability + '' + + (flags ? ' ' + flags + '' : '') + '
' + + '
Effective: ' + g.effective_vram_gb + ' GB' + + ' Tier: ' + g.tier + '
' + + (rows ? '
Models selected:
' + + '' + rows + '
' : '') + + warnHtml + + '
'; +} + export default Tools_ai_provider_settings_class; diff --git a/frontend/src/js/services/api.js b/frontend/src/js/services/api.js index f853673..b78d6e9 100644 --- a/frontend/src/js/services/api.js +++ b/frontend/src/js/services/api.js @@ -213,6 +213,21 @@ class ApiService { } } + /** + * Fetch GPU status: hardware, feature flags, and selected models per operation. + * Only meaningful when AI_PROVIDER=local_gpu. + * @returns {Promise} + */ + async getGpuStatus() { + try { + const response = await fetch(`${this.baseUrl}/api/gpu/status`); + if (!response.ok) return null; + return response.json(); + } catch { + return null; + } + } + /** * Health check for the backend * @returns {Promise} From a97834fda3a96ae5deb9b37fbaa0354e900c06c6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 16:10:32 +0000 Subject: [PATCH 5/8] =?UTF-8?q?feat:=20real-world=20selection=20actions=20?= =?UTF-8?q?=E2=80=94=20scale=20%,=20AI=20edit,=20clipboard=20paste?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backend (3 new endpoints under /api/image/): - POST /api/image/scale-selection — scale selected object by any % in-place; LaMa/OpenCV fills the exposed gap so the scene looks natural - POST /api/image/ai-edit-region — AI redraws the masked region via the configured inpaint provider (local_gpu / InvokeAI / ComfyUI / OpenAI) - POST /api/image/paste-into-selection — scales clipboard image to fit the selection bounding box, masks it to the selection shape, composites result Frontend (selection_actions.js + tool integration): - New SelectionActions panel: fixed bottom-center HUD that appears automatically after every SAM selection (click or paint) - Panel actions: Scale by % (default 3%), Make less symmetrical (AI), custom AI Edit prompt, Replace with clipboard, Copy/Cut to layer, Erase - Both smart_select.js and brush_select.js updated to show the panel, add updateLayerWithResult(), and hide panel on clearSelection/on_leave - brush_select: offerFloatSelection() replaced with richer action panel Real-world workflows now supported in one click after painting over object: "Make this 3% bigger" → scale-selection (LaMa fills gap) "Make this less symmetrical" → ai-edit-region with asymmetry prompt "Replace this with what I copied" → paste-into-selection https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/routers/ai_tools.py | 158 ++++++++++ frontend/src/js/tools/brush_select.js | 44 ++- frontend/src/js/tools/selection_actions.js | 346 +++++++++++++++++++++ frontend/src/js/tools/smart_select.js | 29 +- 4 files changed, 552 insertions(+), 25 deletions(-) create mode 100644 frontend/src/js/tools/selection_actions.js diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index 096903c..be39254 100644 --- a/backend/app/routers/ai_tools.py +++ b/backend/app/routers/ai_tools.py @@ -367,6 +367,164 @@ async def get_config(): } +# ─── Selection image operations ───────────────────────────────────────────── + +class ScaleSelectionRequest(BaseModel): + image: str # base64 full canvas + mask: str # base64 selection mask (white = object) + scale_pct: float = 103.0 # 103 = 3% bigger, 95 = 5% smaller + + +class AiEditRegionRequest(BaseModel): + image: str + mask: str + instruction: str + negative_prompt: str = "" + steps: int = 30 + cfg_scale: float = 7.5 + + +class PasteIntoSelectionRequest(BaseModel): + image: str # base64 target canvas + mask: str # base64 selection mask + paste_image: str # base64 image to paste + + +@router.post("/image/scale-selection") +async def scale_selection(req: ScaleSelectionRequest): + """ + Scale the object selected by mask by scale_pct%, AI-fill the exposed gap. + Works purely with local tools (LaMa/OpenCV) — no remote provider needed. + """ + try: + import numpy as np + from PIL import Image, ImageFilter + except ImportError: + raise HTTPException(status_code=500, detail="PIL/numpy not available") + + img = Image.open(BytesIO(_decode(req.image))).convert("RGB") + mask = Image.open(BytesIO(_decode(req.mask))).convert("L") + if img.size != mask.size: + mask = mask.resize(img.size, Image.LANCZOS) + + mask_arr = np.array(mask) + ys, xs = np.where(mask_arr > 128) + if len(xs) == 0: + raise HTTPException(status_code=400, detail="Empty mask — nothing to scale") + + minx, maxx = int(xs.min()), int(xs.max()) + miny, maxy = int(ys.min()), int(ys.max()) + cx, cy = (minx + maxx) / 2.0, (miny + maxy) / 2.0 + obj_w, obj_h = maxx - minx + 1, maxy - miny + 1 + + scale = req.scale_pct / 100.0 + new_w = max(1, round(obj_w * scale)) + new_h = max(1, round(obj_h * scale)) + + # Extract masked object crop (RGBA with mask as alpha) + img_rgba = img.convert("RGBA") + obj_crop = img_rgba.crop((minx, miny, maxx + 1, maxy + 1)) + mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1)) + r, g, b, _ = obj_crop.split() + obj_masked = Image.merge("RGBA", (r, g, b, mask_crop)) + scaled_obj = obj_masked.resize((new_w, new_h), Image.LANCZOS) + + # AI-fill the original mask area (gap) with LaMa/OpenCV + gap_mask = mask.filter(ImageFilter.MaxFilter(9)) # expand ~4px for clean seam + gap_bytes = BytesIO() + img.save(gap_bytes, format="PNG") + gap_mask_bytes = BytesIO() + gap_mask.save(gap_mask_bytes, format="PNG") + + try: + if lama_available(): + filled_bytes = await asyncio.get_event_loop().run_in_executor( + None, lama_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue() + ) + else: + filled_bytes = await asyncio.get_event_loop().run_in_executor( + None, opencv_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue() + ) + filled = Image.open(BytesIO(filled_bytes)).convert("RGBA") + except Exception as exc: + print(f"[scale-selection] fill fallback: {exc}") + filled = img.convert("RGBA") + + # Paste scaled object centered on original centroid + px = round(cx - new_w / 2) + py = round(cy - new_h / 2) + result = filled.copy() + result.paste(scaled_obj, (px, py), scaled_obj.split()[3]) + + out = BytesIO() + result.convert("RGB").save(out, format="PNG") + return {"result": _encode(out.getvalue())} + + +@router.post("/image/ai-edit-region") +async def ai_edit_region(req: AiEditRegionRequest): + """ + AI-edit the selected region using the configured inpaint provider. + Works with local_gpu, InvokeAI, ComfyUI, or OpenAI. + """ + provider = _require_remote("inpaint") + result_bytes = await provider.inpaint( + _decode(req.image), + _decode(req.mask), + req.instruction, + {"negative_prompt": req.negative_prompt, "steps": req.steps, "cfg_scale": req.cfg_scale}, + ) + return {"result": _encode(result_bytes)} + + +@router.post("/image/paste-into-selection") +async def paste_into_selection(req: PasteIntoSelectionRequest): + """ + Scale a clipboard image to the selection bounding box, mask it to the + selection shape, and composite it over the original canvas. + """ + try: + import numpy as np + from PIL import Image + except ImportError: + raise HTTPException(status_code=500, detail="PIL/numpy not available") + + img = Image.open(BytesIO(_decode(req.image))).convert("RGBA") + mask = Image.open(BytesIO(_decode(req.mask))).convert("L") + paste_img = Image.open(BytesIO(_decode(req.paste_image))).convert("RGBA") + + if img.size != mask.size: + mask = mask.resize(img.size, Image.LANCZOS) + + mask_arr = np.array(mask) + ys, xs = np.where(mask_arr > 128) + if len(xs) == 0: + raise HTTPException(status_code=400, detail="Empty mask") + + minx, maxx = int(xs.min()), int(xs.max()) + miny, maxy = int(ys.min()), int(ys.max()) + target_w = maxx - minx + 1 + target_h = maxy - miny + 1 + + # Scale clipboard image to fit the selection bounding box + paste_scaled = paste_img.resize((target_w, target_h), Image.LANCZOS) + + # Clip paste to selection shape using mask + mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1)) + r, g, b, a = paste_scaled.split() + mask_np = np.array(mask_crop) + alpha_np = np.array(a) + combined = (alpha_np.astype(np.uint16) * mask_np.astype(np.uint16) // 255).astype(np.uint8) + paste_final = Image.merge("RGBA", (r, g, b, Image.fromarray(combined))) + + result = img.copy() + result.paste(paste_final, (minx, miny), paste_final.split()[3]) + + out = BytesIO() + result.convert("RGB").save(out, format="PNG") + return {"result": _encode(out.getvalue())} + + # ─── SAM (Segment Anything) ────────────────────────────────────────────────── class SegmentPointRequest(BaseModel): diff --git a/frontend/src/js/tools/brush_select.js b/frontend/src/js/tools/brush_select.js index 6d6d4b9..9302cd3 100644 --- a/frontend/src/js/tools/brush_select.js +++ b/frontend/src/js/tools/brush_select.js @@ -11,6 +11,7 @@ import Base_layers_class from './../core/base-layers.js'; import Helper_class from './../libs/helpers.js'; import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; import apiService from './../services/api.js'; +import { SelectionActions, updateLayerWithResult } from './selection_actions.js'; class Brush_select_class extends Base_tools_class { @@ -39,6 +40,9 @@ class Brush_select_class extends Base_tools_class { // Processing state this.isProcessing = false; + + // Quick-action panel shown after selection + this.selectionActions = new SelectionActions(this); } load() { @@ -300,31 +304,23 @@ class Brush_select_class extends Base_tools_class { } /** - * Offer to float the selection to a new layer for manipulation (Canva-like workflow) + * Show quick-action panel after selection (AI operations, scale, clipboard paste, etc.) */ offerFloatSelection() { - var _this = this; + var imageData = this.getLayerImageData(); + var maskData = this.maskCanvas + ? this.maskCanvas.toDataURL('image/png').split(',')[1] + : null; + if (maskData) { + this.selectionActions.show(imageData, maskData); + } + } - alertify.confirm( - 'Selection Complete', - 'Would you like to move/scale this selection? This will copy it to a new layer.', - function() { - // Yes - copy to layer and switch to Select tool - _this.copyToLayer(); - - // Switch to Select tool - setTimeout(function() { - var selectTool = document.querySelector('.sidebar_left .item[data-tool="select"]'); - if (selectTool) { - selectTool.click(); - } - }, 100); - }, - function() { - // No - just keep the selection - alertify.message('Tip: Use Ctrl+C to copy or Ctrl+X to cut the selection.'); - } - ).set('labels', {ok: 'Yes, Move/Scale', cancel: 'Keep Selection'}); + /** + * Update the current layer canvas with a base64 result from a backend operation. + */ + updateLayerWithResult(base64) { + updateLayerWithResult(base64, this); } /** @@ -781,6 +777,7 @@ class Brush_select_class extends Base_tools_class { } clearSelection() { + this.selectionActions.hide(); this.currentMask = null; this.maskCanvas = null; this.edgeCanvas = null; @@ -793,8 +790,9 @@ class Brush_select_class extends Base_tools_class { } on_leave() { + this.selectionActions.hide(); this.isDrawing = false; - this.isProcessing = false; // Reset processing state when leaving tool + this.isProcessing = false; this.brushPath = []; return []; } diff --git a/frontend/src/js/tools/selection_actions.js b/frontend/src/js/tools/selection_actions.js new file mode 100644 index 0000000..af6f9b7 --- /dev/null +++ b/frontend/src/js/tools/selection_actions.js @@ -0,0 +1,346 @@ +/** + * SelectionActions — floating quick-action panel that appears after a SAM selection. + * + * Surfaces high-value real-world workflows directly in the UI: + * • Scale by % — make object 3% (or any %) bigger/smaller, gap AI-filled + * • Make less symmetrical — AI redraws the region with organic variation + * • Replace with clipboard — paste clipboard image into the selection shape + * • Copy / Cut to layer — classic Photoshop workflow + * • AI Edit (custom prompt) — full inpaint with user text + * + * Usage: + * this.selectionActions = new SelectionActions(this); + * // after successful selection: + * this.selectionActions.show(imageBase64, maskBase64); + */ + +import app from './../app.js'; +import config from './../config.js'; +import Base_layers_class from './../core/base-layers.js'; +import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; + +const BASE = window.API_BASE_URL || ''; + +export class SelectionActions { + constructor(tool) { + this.tool = tool; + this.Base_layers = new Base_layers_class(); + this._panel = null; + this._imageData = null; + this._maskData = null; + this._escHandler = null; + } + + show(imageBase64, maskBase64) { + this.hide(); + this._imageData = imageBase64; + this._maskData = maskBase64; + + var panel = document.createElement('div'); + panel.id = 'sel-actions-panel'; + panel.style.cssText = [ + 'position:fixed', + 'bottom:80px', + 'left:50%', + 'transform:translateX(-50%)', + 'background:#1a1a2e', + 'border:1px solid #3a3a6a', + 'border-radius:12px', + 'padding:14px 16px', + 'z-index:10000', + 'font-family:sans-serif', + 'font-size:12px', + 'color:#d0d0e0', + 'min-width:340px', + 'box-shadow:0 8px 32px rgba(0,0,0,0.7)', + 'display:flex', + 'flex-direction:column', + 'gap:6px', + ].join(';'); + + // ── Title row ──────────────────────────────────────────────────────── + var titleRow = document.createElement('div'); + titleRow.style.cssText = 'display:flex;align-items:center;justify-content:space-between;margin-bottom:4px'; + var title = document.createElement('span'); + title.textContent = 'Selection Actions'; + title.style.cssText = 'font-size:13px;font-weight:bold;color:#aaaaff'; + var closeX = document.createElement('button'); + closeX.textContent = '✕'; + closeX.style.cssText = 'background:none;border:none;color:#666;cursor:pointer;font-size:14px;padding:0;line-height:1'; + closeX.title = 'Close panel (keep selection)'; + closeX.onclick = () => this.hide(); + titleRow.appendChild(title); + titleRow.appendChild(closeX); + panel.appendChild(titleRow); + + // ── Scale by % ─────────────────────────────────────────────────────── + var scaleRow = document.createElement('div'); + scaleRow.style.cssText = 'display:flex;align-items:center;gap:6px;background:#16213e;border-radius:7px;padding:7px 10px'; + var scaleLabel = document.createElement('span'); + scaleLabel.textContent = 'Scale by'; + scaleLabel.style.color = '#aaa'; + var scaleInput = document.createElement('input'); + scaleInput.type = 'number'; + scaleInput.value = '103'; + scaleInput.min = '1'; + scaleInput.max = '500'; + scaleInput.title = '103 = 3% bigger · 95 = 5% smaller'; + scaleInput.style.cssText = 'width:52px;background:#0f0f1a;color:#fff;border:1px solid #4a4a8a;border-radius:4px;padding:2px 5px;font-size:12px'; + var scaleUnit = document.createElement('span'); + scaleUnit.textContent = '%'; + scaleUnit.style.color = '#888'; + var scaleBtn = _btn('Apply', '#1a2a4a', '#8aacff'); + scaleBtn.style.marginLeft = 'auto'; + scaleBtn.onclick = () => { + var pct = parseFloat(scaleInput.value) || 103; + this._scaleSelection(pct); + }; + scaleRow.appendChild(scaleLabel); + scaleRow.appendChild(scaleInput); + scaleRow.appendChild(scaleUnit); + scaleRow.appendChild(scaleBtn); + panel.appendChild(scaleRow); + + // ── AI actions ─────────────────────────────────────────────────────── + panel.appendChild( + _actionBtn('Make less symmetrical', '#1c1a2e', '#cc99ff', + '⟳ AI redraws the region with natural, organic asymmetry', + () => this._makeAsymmetric()) + ); + panel.appendChild( + _actionBtn('Replace with clipboard', '#1a2a1a', '#88dd88', + '📋 Scales your clipboard image into the selection shape', + () => this._pasteFromClipboard()) + ); + + // ── Custom AI edit prompt ───────────────────────────────────────────── + var aiRow = document.createElement('div'); + aiRow.style.cssText = 'display:flex;align-items:center;gap:6px;background:#16213e;border-radius:7px;padding:7px 10px'; + var aiInput = document.createElement('input'); + aiInput.type = 'text'; + aiInput.placeholder = 'AI edit: "add a scar", "make it look aged", …'; + aiInput.style.cssText = 'flex:1;background:#0f0f1a;color:#fff;border:1px solid #4a4a8a;border-radius:4px;padding:3px 7px;font-size:11px'; + var aiBtn = _btn('Edit', '#1a2a4a', '#8aacff'); + aiBtn.onclick = () => { + var instruction = aiInput.value.trim(); + if (!instruction) { alertify.warning('Enter an AI edit instruction first.'); return; } + this._aiEditRegion(instruction); + }; + aiRow.appendChild(aiInput); + aiRow.appendChild(aiBtn); + panel.appendChild(aiRow); + + // ── Divider ────────────────────────────────────────────────────────── + var hr = document.createElement('div'); + hr.style.cssText = 'border-top:1px solid #2a2a4a;margin:2px 0'; + panel.appendChild(hr); + + // ── Classic selection ops ───────────────────────────────────────────── + var classicRow = document.createElement('div'); + classicRow.style.cssText = 'display:flex;gap:6px'; + var copyBtn = _btn('Copy to layer', '#1a2a1a', '#88cc88'); + copyBtn.style.flex = '1'; + copyBtn.title = 'Ctrl+C'; + copyBtn.onclick = () => { this.tool.copyToLayer(); this.hide(); }; + var cutBtn = _btn('Cut to layer', '#2a1a1a', '#cc8888'); + cutBtn.style.flex = '1'; + cutBtn.title = 'Ctrl+X'; + cutBtn.onclick = () => { this.tool.cutToLayer(); this.hide(); }; + var delBtn = _btn('Erase', '#2a1a1a', '#ff7766'); + delBtn.style.flex = '0 0 auto'; + delBtn.title = 'Delete key'; + delBtn.onclick = () => { this.tool.deleteSelection(); this.hide(); }; + classicRow.appendChild(copyBtn); + classicRow.appendChild(cutBtn); + classicRow.appendChild(delBtn); + panel.appendChild(classicRow); + + document.body.appendChild(panel); + this._panel = panel; + + this._escHandler = (e) => { if (e.key === 'Escape') this.hide(); }; + document.addEventListener('keydown', this._escHandler); + } + + hide() { + if (this._panel) { this._panel.remove(); this._panel = null; } + if (this._escHandler) { + document.removeEventListener('keydown', this._escHandler); + this._escHandler = null; + } + } + + // ── Actions ───────────────────────────────────────────────────────────── + + async _scaleSelection(scalePct) { + if (!this._check()) return; + this.hide(); + alertify.message('Scaling object and filling gap…'); + try { + var res = await _post('/api/image/scale-selection', { + image: this._imageData, + mask: this._maskData, + scale_pct: scalePct, + }); + this.tool.updateLayerWithResult(res.result); + this.tool.clearSelection(); + alertify.success('Scaled by ' + scalePct + '%!'); + } catch (e) { + alertify.error('Scale failed: ' + e.message); + } + } + + async _makeAsymmetric() { + if (!this._check()) return; + this.hide(); + alertify.message('AI is adding natural asymmetry…'); + try { + var res = await _post('/api/image/ai-edit-region', { + image: this._imageData, + mask: this._maskData, + instruction: 'natural asymmetry, slight organic variation, realistic, subtle imperfection', + negative_prompt:'perfectly symmetric, mirror image, artificial, identical halves', + steps: 30, + cfg_scale: 7.5, + }); + this.tool.updateLayerWithResult(res.result); + this.tool.clearSelection(); + alertify.success('Made less symmetrical!'); + } catch (e) { + alertify.error('AI edit failed: ' + e.message); + } + } + + async _aiEditRegion(instruction) { + if (!this._check()) return; + this.hide(); + alertify.message('AI is editing the region…'); + try { + var res = await _post('/api/image/ai-edit-region', { + image: this._imageData, + mask: this._maskData, + instruction: instruction, + steps: 30, + cfg_scale: 7.5, + }); + this.tool.updateLayerWithResult(res.result); + this.tool.clearSelection(); + alertify.success('Done!'); + } catch (e) { + alertify.error('AI edit failed: ' + e.message); + } + } + + async _pasteFromClipboard() { + if (!this._check()) return; + + if (!navigator.clipboard || !navigator.clipboard.read) { + alertify.error('Clipboard API not available. Use HTTPS or enable clipboard permissions.'); + return; + } + try { + var items = await navigator.clipboard.read(); + var clipBlob = null; + for (var item of items) { + for (var type of item.types) { + if (type.startsWith('image/')) { + clipBlob = await item.getType(type); + break; + } + } + if (clipBlob) break; + } + if (!clipBlob) { + alertify.error('No image in clipboard. Copy an image first (e.g., right-click → Copy image).'); + return; + } + + var clipBase64 = await _blobToBase64(clipBlob); + this.hide(); + alertify.message('Pasting clipboard into selection…'); + + var res = await _post('/api/image/paste-into-selection', { + image: this._imageData, + mask: this._maskData, + paste_image: clipBase64, + }); + this.tool.updateLayerWithResult(res.result); + this.tool.clearSelection(); + alertify.success('Clipboard pasted into selection!'); + } catch (e) { + alertify.error('Paste failed: ' + e.message); + } + } + + _check() { + if (!this._imageData || !this._maskData) { + alertify.error('No selection data. Make a new selection first.'); + return false; + } + return true; + } +} + +// ── Shared method: patch into both smart_select and brush_select instances ─── + +/** + * Update the active layer canvas with a base64 result image from the backend. + * Call as `this.updateLayerWithResult(base64)` on any tool that extends Base_tools_class. + */ +export function updateLayerWithResult(base64, tool) { + var img = new Image(); + img.onload = function () { + var canvas = document.createElement('canvas'); + canvas.width = img.width; + canvas.height = img.height; + canvas.getContext('2d').drawImage(img, 0, 0); + + app.State.do_action( + new app.Actions.Bundle_action('ai_transform', 'AI Transform', [ + new app.Actions.Update_layer_image_action(canvas, config.layer.id) + ]) + ); + // Trigger re-render + config.need_render = true; + }; + img.src = 'data:image/png;base64,' + base64; +} + +// ── Private helpers ────────────────────────────────────────────────────────── + +function _btn(text, bg, color) { + var b = document.createElement('button'); + b.textContent = text; + b.style.cssText = 'background:' + bg + ';color:' + color + ';border:1px solid #3a3a6a;padding:4px 10px;border-radius:5px;cursor:pointer;font-size:11px;white-space:nowrap'; + return b; +} + +function _actionBtn(text, bg, color, tooltip, handler) { + var b = _btn(text, bg, color); + b.style.cssText += ';display:block;width:100%;text-align:left;padding:7px 10px;border-radius:7px;font-size:12px'; + if (tooltip) b.title = tooltip; + b.onclick = handler; + return b; +} + +async function _post(path, body) { + var r = await fetch(BASE + path, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }); + if (!r.ok) { + var err = await r.json().catch(() => ({ detail: r.statusText })); + throw new Error(err.detail || 'Request failed'); + } + return r.json(); +} + +function _blobToBase64(blob) { + return new Promise((resolve, reject) => { + var reader = new FileReader(); + reader.onload = (e) => resolve(e.target.result.split(',')[1]); + reader.onerror = reject; + reader.readAsDataURL(blob); + }); +} diff --git a/frontend/src/js/tools/smart_select.js b/frontend/src/js/tools/smart_select.js index 643447c..ce0829d 100644 --- a/frontend/src/js/tools/smart_select.js +++ b/frontend/src/js/tools/smart_select.js @@ -13,6 +13,7 @@ import Helper_class from './../libs/helpers.js'; import Dialog_class from './../libs/popup.js'; import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; import apiService from './../services/api.js'; +import { SelectionActions, updateLayerWithResult } from './selection_actions.js'; class Smart_select_class extends Base_tools_class { @@ -35,6 +36,9 @@ class Smart_select_class extends Base_tools_class { // Edge canvas for drawing the mask outline this.edgeCanvas = null; + + // Quick-action panel shown after selection + this.selectionActions = new SelectionActions(this); } load() { @@ -146,7 +150,7 @@ class Smart_select_class extends Base_tools_class { if (isAdditive && this.currentMask) { alertify.success('Added to selection! Shift+Click to add more.'); } else { - alertify.success('Selection complete! Shift+Click to add more, Ctrl+C to copy, Ctrl+X to cut.'); + this._showActionPanel(); } } catch (error) { @@ -633,10 +637,31 @@ class Smart_select_class extends Base_tools_class { alertify.success('Selection deleted!'); } + /** + * Show the quick-action panel for the current selection. + */ + _showActionPanel() { + var imageData = this.getLayerImageData(); + var maskData = this.maskCanvas + ? this.maskCanvas.toDataURL('image/png').split(',')[1] + : null; + if (maskData) { + this.selectionActions.show(imageData, maskData); + } + } + + /** + * Update the current layer canvas with a base64 result from a backend operation. + */ + updateLayerWithResult(base64) { + updateLayerWithResult(base64, this); + } + /** * Clear the current selection */ clearSelection() { + this.selectionActions.hide(); this.currentMask = null; this.maskCanvas = null; this.edgeCanvas = null; @@ -647,7 +672,7 @@ class Smart_select_class extends Base_tools_class { } on_leave() { - // Don't clear mask when switching tools - AI inpaint needs it + this.selectionActions.hide(); return []; } } From 5940e10542c3f93dfec838771cf0abc5ae3a40a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 16:25:56 +0000 Subject: [PATCH 6/8] Add print size presets, 18x24 frame, and Prepare for Print workflow - Add 18x24" to FRAME_SIZES in backend and frontend (frame_fit.js) - Add 200 DPI option to frame_fit dialog (adequate for large-format prints) - Add 18x24 portrait/landscape at 200 and 300 DPI to Canvas Size presets (size.js) - New /api/print/prepare endpoint: chains AI upscale to target DPI then frame-fit in one server-side call (avoids round-tripping a large upscaled image) - New print_prepare.js module: "Prepare for Print" dialog with per-frame quality assessment (current effective DPI, needed upscale factor, AI vs Lanczos note) - Add "Prepare for Print..." to Image menu above "Fit to Frame..." https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/routers/print_tools.py | 110 +++++++ frontend/src/js/config-menu.js | 5 + frontend/src/js/modules/image/frame_fit.js | 8 +- .../src/js/modules/image/print_prepare.js | 285 ++++++++++++++++++ frontend/src/js/modules/image/size.js | 4 + 5 files changed, 408 insertions(+), 4 deletions(-) create mode 100644 frontend/src/js/modules/image/print_prepare.js diff --git a/backend/app/routers/print_tools.py b/backend/app/routers/print_tools.py index 52fe476..71f8046 100644 --- a/backend/app/routers/print_tools.py +++ b/backend/app/routers/print_tools.py @@ -21,6 +21,7 @@ FRAME_SIZES = { "8x10": (8, 10), "11x14": (11, 14), "16x20": (16, 20), + "18x24": (18, 24), "20x24": (20, 24), "24x36": (24, 36), # Square @@ -65,6 +66,16 @@ class UpscaleRequest(BaseModel): method: str = "auto" +class PrepareRequest(BaseModel): + image: str # base64 + frame: str # e.g. "8x10" + orientation: Literal["auto", "portrait", "landscape"] = "auto" + target_dpi: int = 300 + upscale_method: str = "auto" # auto / realesrgan_pytorch / realesrgan_ncnn / lanczos + mode: Literal["crop", "extend", "smart"] = "smart" + prompt: Optional[str] = "" + + # ── Frame sizes endpoint ─────────────────────────────────────────────────── @router.get("/frame-sizes") @@ -335,6 +346,105 @@ def upscale_install_status(): return status +@router.post("/prepare") +async def prepare_for_print(req: PrepareRequest): + """ + One-shot Prepare for Print: AI upscale to reach target DPI, then fit to frame. + + Steps: + 1. Resolve target pixel dimensions (frame × target_dpi, orientation-adjusted) + 2. Calculate needed upscale factor so the image meets the target resolution + 3. Run Real-ESRGAN if scale > 1.05 (else skip — already large enough) + 4. Run frame-fit (crop / extend / smart) to exact target dimensions + 5. Return the print-ready image and a quality report + """ + if req.frame not in FRAME_SIZES: + raise HTTPException(status_code=400, + detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}") + if not (72 <= req.target_dpi <= 600): + raise HTTPException(status_code=400, detail="target_dpi must be 72–600") + + try: + image = Image.open(BytesIO(_decode(req.image))).convert("RGB") + except Exception as e: + raise HTTPException(status_code=400, detail=f"Could not decode image: {e}") + + fw, fh = FRAME_SIZES[req.frame] + img_w, img_h = image.size + + # Resolve orientation (same logic as frame_fit) + img_landscape = img_w >= img_h + frame_landscape = fw >= fh + if req.orientation == "landscape": + fw, fh = max(fw, fh), min(fw, fh) + elif req.orientation == "portrait": + fw, fh = min(fw, fh), max(fw, fh) + else: + if img_landscape and not frame_landscape: + fw, fh = fh, fw + elif not img_landscape and frame_landscape: + fw, fh = fh, fw + + target_w = fw * req.target_dpi + target_h = fh * req.target_dpi + + # Scale factor needed so the shorter dimension fills the frame + scale_w = target_w / img_w + scale_h = target_h / img_h + needed_scale = min(scale_w, scale_h) # fill-to-fit (extend) baseline + # For crop mode we need max; use the larger to be safe and let frame-fit crop + needed_scale_crop = max(scale_w, scale_h) + + # Use the smaller (extend) scale as the upscale target; frame-fit handles the rest + upscale_factor = max(1.0, needed_scale) + upscale_applied = False + method_used = "none" + + upscaled = image + if upscale_factor > 1.05: + # Cap per-pass at 4× (Real-ESRGAN works best at 2–4×) + remaining = upscale_factor + while remaining > 1.05: + pass_scale = min(remaining, 4.0) + # Round to one decimal to keep scale in 1.1–8.0 range accepted by upscale service + pass_scale = round(pass_scale, 1) + if pass_scale < 1.1: + break + from app.services.upscale import upscale_image + result_bytes, method_used = await upscale_image(upscaled, pass_scale, req.upscale_method) + upscaled = Image.open(BytesIO(result_bytes)).convert("RGB") + remaining /= pass_scale + upscale_applied = True + + # Encode upscaled image and run frame-fit + upscaled_b64 = _encode(_to_png(upscaled)) + + fit_req = FrameFitRequest( + image=upscaled_b64, + frame=req.frame, + orientation=req.orientation, + mode=req.mode, + dpi=req.target_dpi, + prompt=req.prompt or "", + ) + # Re-use the existing frame_fit logic inline + fit_response = await frame_fit(fit_req) + + return { + "result": fit_response["result"], + "frame": req.frame, + "orientation": fit_response["orientation"], + "output_pixels": fit_response["output_pixels"], + "output_inches": fit_response["output_inches"], + "dpi": req.target_dpi, + "mode_used": fit_response["mode_used"], + "upscale_applied": upscale_applied, + "upscale_factor": round(upscale_factor, 2), + "upscale_method": method_used, + "summary": fit_response["summary"], + } + + @router.post("/upscale") async def upscale(req: UpscaleRequest): """ diff --git a/frontend/src/js/config-menu.js b/frontend/src/js/config-menu.js index 9f592ed..3efb940 100644 --- a/frontend/src/js/config-menu.js +++ b/frontend/src/js/config-menu.js @@ -344,6 +344,11 @@ const menuDefinition = [ ellipsis: true, target: 'image/remove_background.remove_background' }, + { + name: 'Prepare for Print...', + ellipsis: true, + target: 'image/print_prepare.print_prepare' + }, { name: 'Fit to Frame...', ellipsis: true, diff --git a/frontend/src/js/modules/image/frame_fit.js b/frontend/src/js/modules/image/frame_fit.js index 81ac049..f967164 100644 --- a/frontend/src/js/modules/image/frame_fit.js +++ b/frontend/src/js/modules/image/frame_fit.js @@ -19,7 +19,7 @@ import { getCapabilities } from './../../api/capabilities.js'; var instance = null; const FRAME_SIZES = [ - '4x6', '5x7', '8x10', '11x14', '16x20', '20x24', '24x36', + '4x6', '5x7', '8x10', '11x14', '16x20', '18x24', '20x24', '24x36', '4x4', '8x8', '12x12', ]; @@ -27,8 +27,8 @@ const FRAME_SIZES = [ const FRAME_PX = { '4x6': [1200, 1800], '5x7': [1500, 2100], '8x10': [2400, 3000], '11x14': [3300, 4200], - '16x20': [4800, 6000], '20x24': [6000, 7200], - '24x36': [7200, 10800], + '16x20': [4800, 6000], '18x24': [5400, 7200], + '20x24': [6000, 7200], '24x36': [7200, 10800], '4x4': [1200, 1200], '8x8': [2400, 2400], '12x12': [3600, 3600], }; @@ -96,7 +96,7 @@ class Image_frame_fit_class { name: 'dpi', title: 'Output DPI:', value: '300', - values: ['72', '150', '300'], + values: ['72', '150', '200', '300'], type: 'select', }, { diff --git a/frontend/src/js/modules/image/print_prepare.js b/frontend/src/js/modules/image/print_prepare.js new file mode 100644 index 0000000..38b8735 --- /dev/null +++ b/frontend/src/js/modules/image/print_prepare.js @@ -0,0 +1,285 @@ +/** + * Prepare for Print — one-click AI upscale + frame fit. + * + * Shows a quality assessment (current effective DPI, needed upscale factor, + * AI vs Lanczos note) then chains AI upscale → frame-fit in a single backend call. + * + * Menu target: image/print_prepare.print_prepare + */ + +import app from './../../app.js'; +import config from './../../config.js'; +import Base_layers_class from './../../core/base-layers.js'; +import Dialog_class from './../../libs/popup.js'; +import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; +import { getCapabilities } from './../../api/capabilities.js'; + +const FRAME_SIZES = [ + '5x7', '8x10', '11x14', '18x24', '16x20', '20x24', '24x36', +]; + +// Portrait pixels at 300 DPI (label use only) +const FRAME_PX = { + '5x7': [1500, 2100], '8x10': [2400, 3000], + '11x14': [3300, 4200], '18x24': [5400, 7200], + '16x20': [4800, 6000], '20x24': [6000, 7200], + '24x36': [7200, 10800], +}; + +// Actual frame inches (portrait w, h) +const FRAME_IN = { + '5x7': [5, 7], '8x10': [8, 10], '11x14': [11, 14], + '18x24': [18, 24], '16x20': [16, 20], '20x24': [20, 24], + '24x36': [24, 36], +}; + +var instance = null; + +class Image_print_prepare_class { + + constructor() { + if (instance) return instance; + instance = this; + this.Base_layers = new Base_layers_class(); + this.Dialog = new Dialog_class(); + this.isProcessing = false; + } + + async print_prepare() { + if (!config.layer || config.layer.type !== 'image') { + alertify.error('Select an image layer first.'); + return; + } + + var caps = await getCapabilities(); + var hasAI = (caps.remote && caps.remote.healthy) || (caps.local && caps.local.local_gpu_available); + + var W = config.layer.width_original; + var H = config.layer.height_original; + + var qualityHtml = _buildQualityHtml(W, H, hasAI); + + var frameLabels = FRAME_SIZES.map(s => { + var px = FRAME_PX[s] || [0, 0]; + return `${s}" (${px[0]}×${px[1]}px @ 300dpi)`; + }); + + var _this = this; + this.Dialog.show({ + title: 'Prepare for Print', + params: [ + { + title: '', + html: qualityHtml, + }, + { + name: 'frame', + title: 'Target frame size:', + value: frameLabels[0], + values: frameLabels, + type: 'select', + }, + { + name: 'orientation', + title: 'Orientation:', + value: 'auto', + values: ['auto', 'portrait', 'landscape'], + type: 'select', + }, + { + name: 'target_dpi', + title: 'Target DPI:', + value: '300', + values: ['200', '300'], + type: 'select', + comment: '200 dpi is fine for 18×24" and larger (viewed from a distance)', + }, + { + name: 'mode', + title: 'Fit mode:', + value: 'smart', + values: ['smart', 'crop', 'extend'], + type: 'select', + comment: 'smart = extend if gap <15%, else crop', + }, + { + name: 'upscale_method', + title: 'Upscale engine:', + value: 'auto', + values: ['auto', 'realesrgan_pytorch', 'realesrgan_ncnn', 'lanczos'], + type: 'select', + comment: hasAI ? 'auto picks Real-ESRGAN — genuinely adds detail' : 'auto picks Real-ESRGAN if available, else Lanczos', + }, + { + name: 'prompt', + title: 'Extend prompt (optional):', + value: '', + placeholder: 'e.g. "natural background continuation" — blank works well', + }, + { + name: 'new_layer', + title: 'Result as new layer (keep original):', + value: true, + }, + ], + on_finish: async function (params) { + var frameKey = params.frame.split('"')[0]; + await _this._run(frameKey, params, W, H); + }, + }); + } + + async _run(frameKey, params, origW, origH) { + if (this.isProcessing) return; + this.isProcessing = true; + + var dpi = parseInt(params.target_dpi) || 300; + var inches = FRAME_IN[frameKey] || [8, 10]; + var targetW = inches[0] * dpi; + var targetH = inches[1] * dpi; + + // Orientation swap for display + var orient = params.orientation || 'auto'; + var imgLandscape = origW >= origH; + var frameLandscape = inches[0] >= inches[1]; + if (orient === 'landscape' || (orient === 'auto' && imgLandscape && !frameLandscape)) { + targetW = Math.max(inches[0], inches[1]) * dpi; + targetH = Math.min(inches[0], inches[1]) * dpi; + } else if (orient === 'portrait' || (orient === 'auto' && !imgLandscape && frameLandscape)) { + targetW = Math.min(inches[0], inches[1]) * dpi; + targetH = Math.max(inches[0], inches[1]) * dpi; + } + + var neededScale = Math.max(targetW / origW, targetH / origH); + var willUpscale = neededScale > 1.05; + + alertify.message( + willUpscale + ? `Upscaling ${neededScale.toFixed(1)}× with AI, then fitting to frame… this may take a minute` + : 'Fitting to frame…', + 0 + ); + + try { + var layerCanvas = document.createElement('canvas'); + layerCanvas.width = origW; + layerCanvas.height = origH; + layerCanvas.getContext('2d').drawImage(config.layer.link, 0, 0); + var imageB64 = layerCanvas.toDataURL('image/png').split(',')[1]; + + var base = window.API_BASE_URL || ''; + var r = await fetch(`${base}/api/print/prepare`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + image: imageB64, + frame: frameKey, + orientation: orient, + target_dpi: dpi, + upscale_method: params.upscale_method || 'auto', + mode: params.mode || 'smart', + prompt: params.prompt || '', + }), + }); + + if (!r.ok) { + var err = await r.json().catch(() => ({ detail: 'Server error' })); + throw new Error(err.detail || 'Prepare failed'); + } + var result = await r.json(); + + var img = new Image(); + img.onload = () => { + var resultCanvas = document.createElement('canvas'); + resultCanvas.width = img.naturalWidth; + resultCanvas.height = img.naturalHeight; + resultCanvas.getContext('2d').drawImage(img, 0, 0); + + var fitW = img.naturalWidth; + var fitH = img.naturalHeight; + + if (params.new_layer) { + app.State.do_action( + new app.Actions.Bundle_action('print_prepare_layer', 'Prepare for Print', [ + new app.Actions.Prepare_canvas_action('undo'), + new app.Actions.Update_config_action({ WIDTH: fitW, HEIGHT: fitH }), + new app.Actions.Insert_layer_action({ + name: `${frameKey} ${dpi}dpi`, + type: 'image', + data: img.src, + x: 0, y: 0, + width: fitW, height: fitH, + width_original: fitW, height_original: fitH, + }), + new app.Actions.Prepare_canvas_action('do'), + ]) + ); + } else { + app.State.do_action( + new app.Actions.Bundle_action('print_prepare', 'Prepare for Print', [ + new app.Actions.Prepare_canvas_action('undo'), + new app.Actions.Update_config_action({ WIDTH: fitW, HEIGHT: fitH }), + new app.Actions.Update_layer_image_action(resultCanvas), + new app.Actions.Prepare_canvas_action('do'), + ]) + ); + } + + alertify.dismissAll(); + var upscaleNote = result.upscale_applied + ? ` · ${result.upscale_factor}× ${result.upscale_method}` + : ' · no upscale needed'; + alertify.success( + `Print-ready! ${fitW}×${fitH}px @ ${dpi} DPI (${frameKey}")${upscaleNote}` + ); + this.isProcessing = false; + }; + img.onerror = () => { + alertify.dismissAll(); + alertify.error('Failed to load result.'); + this.isProcessing = false; + }; + img.src = 'data:image/png;base64,' + result.result; + + } catch (err) { + alertify.dismissAll(); + alertify.error('Prepare for Print failed: ' + (err.message || err)); + this.isProcessing = false; + } + } +} + +function _buildQualityHtml(W, H, hasAI) { + var rows = FRAME_SIZES.map(key => { + var inches = FRAME_IN[key]; + // Effective DPI: smaller of the two dimensions (limiting factor) + var effDpi = Math.round(Math.min(W / inches[0], H / inches[1])); + var quality = effDpi >= 300 ? '✓ excellent' + : effDpi >= 200 ? '✓ good for large format' + : effDpi >= 150 ? '~ acceptable' + : '✗ needs upscaling'; + var color = effDpi >= 300 ? '#44cc44' + : effDpi >= 200 ? '#88cc44' + : effDpi >= 150 ? '#ffaa44' + : '#ff6644'; + var neededScale = Math.max(1, Math.ceil((300 / effDpi) * 10) / 10); + var scaleNote = effDpi >= 300 ? '' : ` → need ~${neededScale.toFixed(1)}× upscale`; + return ` + ${key}" + ${effDpi} DPI + ${quality}${scaleNote} + `; + }).join(''); + + var aiNote = hasAI + ? 'Real-ESRGAN available — will add genuine sharpness (AI reconstructs detail)' + : 'No AI provider — will use Lanczos (resizes but doesn\'t add detail)'; + + return `
+
Current image: ${W}×${H}px · ${aiNote}
+ ${rows}
+
200 DPI is fine for 18×24" and larger prints viewed from 2+ feet.
+
`; +} + +export default Image_print_prepare_class; diff --git a/frontend/src/js/modules/image/size.js b/frontend/src/js/modules/image/size.js index 39cbc34..f942250 100644 --- a/frontend/src/js/modules/image/size.js +++ b/frontend/src/js/modules/image/size.js @@ -16,6 +16,10 @@ const PRINT_SIZES = [ [3000, 2400, '8x10" Landscape'], [3300, 4200, '11x14" Portrait'], [4200, 3300, '11x14" Landscape'], + [3600, 4800, '18x24" Portrait 200dpi'], + [4800, 3600, '18x24" Landscape 200dpi'], + [5400, 7200, '18x24" Portrait 300dpi'], + [7200, 5400, '18x24" Landscape 300dpi'], ]; class Image_size_class { From ee94282753b85a85135a381aa3e5dfb45cd4cba0 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 16:37:04 +0000 Subject: [PATCH 7/8] Add progress overlay and fix txt2img for local GPU - New progress_overlay.js: animated fullscreen overlay with shimmer bar, fake progress creep, Esc-to-cancel, used by all slow AI operations - text_to_image.js: allow local_gpu provider (was incorrectly blocked); show provider/model/VRAM info in dialog; show estimated generation time; use progress overlay during generation - upscale.js: replace alertify.message with progress overlay (90s estimate for AI upscale, 10s for Lanczos) - frame_fit.js: progress overlay for extend mode (AI outpaint ~45s) - print_prepare.js: progress overlay for full upscale+frame chain (~2 min) - selection_actions.js: progress overlay for all AI region edits https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- frontend/src/js/libs/progress_overlay.js | 142 ++++++++++++++++++ .../src/js/modules/generate/text_to_image.js | 88 ++++++++--- frontend/src/js/modules/image/frame_fit.js | 15 +- .../src/js/modules/image/print_prepare.js | 14 +- frontend/src/js/modules/image/upscale.js | 14 +- frontend/src/js/tools/selection_actions.js | 17 ++- 6 files changed, 244 insertions(+), 46 deletions(-) create mode 100644 frontend/src/js/libs/progress_overlay.js diff --git a/frontend/src/js/libs/progress_overlay.js b/frontend/src/js/libs/progress_overlay.js new file mode 100644 index 0000000..da10cb0 --- /dev/null +++ b/frontend/src/js/libs/progress_overlay.js @@ -0,0 +1,142 @@ +/** + * ProgressOverlay — shared animated progress indicator for long AI operations. + * + * Usage: + * import { showProgress, updateProgress, hideProgress } from './progress_overlay.js'; + * + * showProgress('Generating image…'); + * updateProgress(50, 'Denoising step 15/30…'); // optional step updates + * hideProgress(); + * + * When you don't have real step counts, call showProgress() and hideProgress() only — + * the bar animates automatically with a shimmer to signal activity. + */ + +var _overlay = null; +var _bar = null; +var _label = null; +var _shimmerAnim = null; +var _fakeTimer = null; +var _currentPct = 0; + +export function showProgress(message, estimatedSeconds) { + hideProgress(); + + _currentPct = 0; + + // ── Backdrop ────────────────────────────────────────────────────────────── + _overlay = document.createElement('div'); + _overlay.id = 'ai-progress-overlay'; + _overlay.style.cssText = [ + 'position:fixed', 'inset:0', 'z-index:99999', + 'display:flex', 'flex-direction:column', + 'align-items:center', 'justify-content:center', + 'background:rgba(0,0,0,0.55)', + 'backdrop-filter:blur(2px)', + '-webkit-backdrop-filter:blur(2px)', + ].join(';'); + + // ── Card ────────────────────────────────────────────────────────────────── + var card = document.createElement('div'); + card.style.cssText = [ + 'background:#1a1a2e', + 'border:1px solid #3a3a6a', + 'border-radius:14px', + 'padding:28px 36px', + 'min-width:320px', 'max-width:480px', + 'box-shadow:0 12px 48px rgba(0,0,0,0.8)', + 'display:flex', 'flex-direction:column', 'gap:14px', + 'text-align:center', + ].join(';'); + + // ── Label ───────────────────────────────────────────────────────────────── + _label = document.createElement('div'); + _label.textContent = message || 'Processing…'; + _label.style.cssText = 'font-family:sans-serif;font-size:13px;color:#c0c0e0;line-height:1.4;min-height:2.8em'; + + // ── Track ───────────────────────────────────────────────────────────────── + var track = document.createElement('div'); + track.style.cssText = [ + 'width:100%', 'height:6px', + 'background:#0f0f2a', + 'border-radius:3px', + 'overflow:hidden', + 'position:relative', + ].join(';'); + + // ── Shimmer (indeterminate stripe) ──────────────────────────────────────── + var shimmer = document.createElement('div'); + shimmer.style.cssText = [ + 'position:absolute', 'inset:0', + 'background:linear-gradient(90deg,transparent 0%,rgba(120,120,255,0.25) 50%,transparent 100%)', + 'transform:translateX(-100%)', + 'will-change:transform', + ].join(';'); + + // ── Filled bar ──────────────────────────────────────────────────────────── + _bar = document.createElement('div'); + _bar.style.cssText = [ + 'position:absolute', 'inset-block:0', 'left:0', + 'width:0%', + 'background:linear-gradient(90deg,#5577ff,#88aaff)', + 'border-radius:3px', + 'transition:width 0.35s ease', + ].join(';'); + + // ── Cancel hint ─────────────────────────────────────────────────────────── + var hint = document.createElement('div'); + hint.textContent = 'Press Esc to cancel'; + hint.style.cssText = 'font-family:sans-serif;font-size:10px;color:#444;margin-top:2px'; + + track.appendChild(shimmer); + track.appendChild(_bar); + card.appendChild(_label); + card.appendChild(track); + card.appendChild(hint); + _overlay.appendChild(card); + document.body.appendChild(_overlay); + + // Animate shimmer + var pos = -100; + _shimmerAnim = setInterval(() => { + pos += 2.5; + if (pos > 200) pos = -100; + shimmer.style.transform = `translateX(${pos}%)`; + }, 16); + + // Fake progress that creeps toward 90% if no real steps given + if (estimatedSeconds) { + var totalMs = estimatedSeconds * 1000; + var step = 90 / (totalMs / 200); + _fakeTimer = setInterval(() => { + if (_currentPct < 90) { + _currentPct = Math.min(90, _currentPct + step); + _bar.style.width = _currentPct + '%'; + } + }, 200); + } + + // Esc to cancel + _overlay._escHandler = (e) => { if (e.key === 'Escape') hideProgress(); }; + document.addEventListener('keydown', _overlay._escHandler); +} + +export function updateProgress(pct, message) { + if (!_overlay) return; + _currentPct = Math.max(_currentPct, Math.min(100, pct)); + if (_bar) _bar.style.width = _currentPct + '%'; + if (_label && message) _label.textContent = message; +} + +export function hideProgress() { + if (_shimmerAnim) { clearInterval(_shimmerAnim); _shimmerAnim = null; } + if (_fakeTimer) { clearInterval(_fakeTimer); _fakeTimer = null; } + if (_overlay) { + document.removeEventListener('keydown', _overlay._escHandler); + _overlay.remove(); + _overlay = null; + } + _bar = null; + _label = null; + _currentPct = 0; +} diff --git a/frontend/src/js/modules/generate/text_to_image.js b/frontend/src/js/modules/generate/text_to_image.js index d23ddca..86f4853 100644 --- a/frontend/src/js/modules/generate/text_to_image.js +++ b/frontend/src/js/modules/generate/text_to_image.js @@ -1,5 +1,5 @@ /** - * Text → Image — opens a sidebar-style dialog, generates via remote provider, + * Text → Image — generates via remote or local-GPU provider, * pastes result as a new layer on the current canvas. * * Menu target: generate/text_to_image.text_to_image @@ -12,6 +12,7 @@ import Dialog_class from './../../libs/popup.js'; import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; import apiService from './../../services/api.js'; import { getCapabilities } from './../../api/capabilities.js'; +import { showProgress, updateProgress, hideProgress } from './../../libs/progress_overlay.js'; var instance = null; @@ -27,21 +28,54 @@ class Generate_text_to_image_class { async text_to_image() { var caps = await getCapabilities(); - if (!caps.remote || !caps.remote.healthy) { + var hasRemote = caps.remote && caps.remote.healthy; + var hasLocal = caps.local && caps.local.local_gpu_available; + + if (!hasRemote && !hasLocal) { alertify.error( - 'Text → Image requires a remote AI provider. ' + - 'Set AI_PROVIDER (openai / invokeai / comfyui) in .env and restart.' + 'Text → Image requires an AI provider. ' + + 'Set AI_PROVIDER=openai / invokeai / comfyui / local_gpu in .env and restart, ' + + 'or configure one in Image → AI Provider Settings.' ); return; } - var _this = this; + var _this = this; var canvasW = config.WIDTH || 1024; var canvasH = config.HEIGHT || 1024; + // Build provider info line + var providerHtml = hasRemote + ? `● ${caps.remote.provider}` + : `● local GPU · ${caps.local.gpu_tier || ''} · ${_shortGpu(caps.local.gpu_device)}`; + + // Model note for local GPU + var modelNote = ''; + if (hasLocal && !hasRemote) { + var rec = caps.local.local_gpu_capabilities && caps.local.local_gpu_capabilities.recommended; + var m = rec && rec.txt2img; + if (m) { + modelNote = `Model: ${m.model_id.split('/').pop()}`; + if (m.memory_opt && m.memory_opt !== 'none') modelNote += ` · ${m.memory_opt}`; + } + } + + // Estimate generation time (rough guide for the progress bar) + var estSec = hasLocal ? 60 : 15; // local GPU ~1 min; OpenAI ~15s + + var defaultW = Math.min(canvasW, hasLocal ? (caps.local.local_gpu_capabilities?.recommended?.txt2img?.native_res || 1024) : 1024); + var defaultH = Math.min(canvasH, defaultW); + this.Dialog.show({ title: 'Text → Image', params: [ + { + title: '', + html: `
+ Provider: ${providerHtml}${modelNote ? ' · ' + modelNote : ''}
+ Generation typically takes ${estSec < 30 ? 'a few seconds' : estSec < 90 ? '30–90 seconds on local GPU' : '1–3 minutes on local GPU'}. +
`, + }, { name: 'prompt', title: 'Describe your image:', @@ -58,7 +92,7 @@ class Generate_text_to_image_class { { name: 'width', title: 'Width (px):', - value: Math.min(canvasW, 1024), + value: defaultW, range: [256, 2048], step: 64, type: 'range', @@ -66,7 +100,7 @@ class Generate_text_to_image_class { { name: 'height', title: 'Height (px):', - value: Math.min(canvasH, 1024), + value: defaultH, range: [256, 2048], step: 64, type: 'range', @@ -99,29 +133,31 @@ class Generate_text_to_image_class { alertify.warning('Please enter a description.'); return; } - await _this._generate(params); + await _this._generate(params, estSec); }, }); } - async _generate(params) { + async _generate(params, estSec) { if (this.isProcessing) return; this.isProcessing = true; - alertify.message('Generating image... please wait', 0); + + showProgress('Generating image… this may take a minute on local GPU', estSec || 60); try { var result = await apiService.textToImage(params.prompt, { - width: params.width || 1024, - height: params.height || 1024, + width: params.width || 1024, + height: params.height || 1024, negativePrompt: params.negative_prompt || '', - steps: params.steps || 30, - seed: params.seed || 0, + steps: params.steps || 30, + seed: params.seed || 0, }); + updateProgress(95, 'Placing image…'); + var img = new Image(); img.onload = () => { if (params.placement === 'replace_canvas') { - // Resize canvas and replace bottom layer config.WIDTH = img.naturalWidth; config.HEIGHT = img.naturalHeight; var resultCanvas = document.createElement('canvas'); @@ -134,41 +170,43 @@ class Generate_text_to_image_class { ]) ); } else { - // Add as new layer on top - var dataURL = img.src; app.State.do_action( new app.Actions.Bundle_action('txt2img_layer', 'Text → Image Layer', [ new app.Actions.Insert_layer_action({ name: params.prompt.slice(0, 30), type: 'image', - data: dataURL, - x: 0, - y: 0, - width: img.naturalWidth, - height: img.naturalHeight, + data: img.src, + x: 0, y: 0, + width: img.naturalWidth, + height: img.naturalHeight, width_original: img.naturalWidth, height_original: img.naturalHeight, }) ]) ); } - alertify.dismissAll(); + hideProgress(); alertify.success('Image generated!'); this.isProcessing = false; }; img.onerror = () => { - alertify.dismissAll(); + hideProgress(); alertify.error('Failed to load generated image.'); this.isProcessing = false; }; img.src = 'data:image/png;base64,' + result.result; } catch (err) { - alertify.dismissAll(); + hideProgress(); alertify.error('Generation failed: ' + (err.message || err)); this.isProcessing = false; } } } +function _shortGpu(name) { + if (!name) return 'GPU'; + return name.replace(/^NVIDIA GeForce /i, '').replace(/^NVIDIA /i, ''); +} + export default Generate_text_to_image_class; diff --git a/frontend/src/js/modules/image/frame_fit.js b/frontend/src/js/modules/image/frame_fit.js index f967164..b56bc00 100644 --- a/frontend/src/js/modules/image/frame_fit.js +++ b/frontend/src/js/modules/image/frame_fit.js @@ -15,6 +15,7 @@ import Base_layers_class from './../../core/base-layers.js'; import Dialog_class from './../../libs/popup.js'; import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; import { getCapabilities } from './../../api/capabilities.js'; +import { showProgress, hideProgress } from './../../libs/progress_overlay.js'; var instance = null; @@ -123,11 +124,11 @@ class Image_frame_fit_class { this.isProcessing = true; var mode = params.mode || 'smart'; - alertify.message( + showProgress( mode === 'extend' - ? 'Fitting to frame with AI extension... please wait' - : 'Fitting to frame...', - 0 + ? 'Fitting to frame with AI extension…' + : 'Fitting to frame…', + mode === 'extend' ? 45 : 5 ); try { @@ -203,7 +204,7 @@ class Image_frame_fit_class { ); } - alertify.dismissAll(); + hideProgress(); alertify.success( `Done! ${result.output_pixels.width}×${result.output_pixels.height}px` + ` (${result.frame} ${result.orientation}, ${result.mode_used})` @@ -211,14 +212,14 @@ class Image_frame_fit_class { this.isProcessing = false; }; img.onerror = () => { - alertify.dismissAll(); + hideProgress(); alertify.error('Failed to load result.'); this.isProcessing = false; }; img.src = 'data:image/png;base64,' + result.result; } catch (err) { - alertify.dismissAll(); + hideProgress(); alertify.error('Frame fit failed: ' + (err.message || err)); this.isProcessing = false; } diff --git a/frontend/src/js/modules/image/print_prepare.js b/frontend/src/js/modules/image/print_prepare.js index 38b8735..cbe84b2 100644 --- a/frontend/src/js/modules/image/print_prepare.js +++ b/frontend/src/js/modules/image/print_prepare.js @@ -13,6 +13,7 @@ import Base_layers_class from './../../core/base-layers.js'; import Dialog_class from './../../libs/popup.js'; import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; import { getCapabilities } from './../../api/capabilities.js'; +import { showProgress, updateProgress, hideProgress } from './../../libs/progress_overlay.js'; const FRAME_SIZES = [ '5x7', '8x10', '11x14', '18x24', '16x20', '20x24', '24x36', @@ -153,11 +154,11 @@ class Image_print_prepare_class { var neededScale = Math.max(targetW / origW, targetH / origH); var willUpscale = neededScale > 1.05; - alertify.message( + showProgress( willUpscale - ? `Upscaling ${neededScale.toFixed(1)}× with AI, then fitting to frame… this may take a minute` + ? `Upscaling ${neededScale.toFixed(1)}× with AI, then fitting to frame…\nAI is reconstructing detail — this may take 1–3 minutes.` : 'Fitting to frame…', - 0 + willUpscale ? 120 : 8 ); try { @@ -188,6 +189,7 @@ class Image_print_prepare_class { } var result = await r.json(); + updateProgress(90, 'Placing result…'); var img = new Image(); img.onload = () => { var resultCanvas = document.createElement('canvas'); @@ -225,7 +227,7 @@ class Image_print_prepare_class { ); } - alertify.dismissAll(); + hideProgress(); var upscaleNote = result.upscale_applied ? ` · ${result.upscale_factor}× ${result.upscale_method}` : ' · no upscale needed'; @@ -235,14 +237,14 @@ class Image_print_prepare_class { this.isProcessing = false; }; img.onerror = () => { - alertify.dismissAll(); + hideProgress(); alertify.error('Failed to load result.'); this.isProcessing = false; }; img.src = 'data:image/png;base64,' + result.result; } catch (err) { - alertify.dismissAll(); + hideProgress(); alertify.error('Prepare for Print failed: ' + (err.message || err)); this.isProcessing = false; } diff --git a/frontend/src/js/modules/image/upscale.js b/frontend/src/js/modules/image/upscale.js index 4968f4d..9947a1f 100644 --- a/frontend/src/js/modules/image/upscale.js +++ b/frontend/src/js/modules/image/upscale.js @@ -13,6 +13,7 @@ import config from './../../config.js'; import Base_layers_class from './../../core/base-layers.js'; import Dialog_class from './../../libs/popup.js'; import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; +import { showProgress, updateProgress, hideProgress } from './../../libs/progress_overlay.js'; var instance = null; @@ -221,7 +222,12 @@ class Image_upscale_class { ? `Auto (${caps.recommended_label || 'best available'})` : (METHOD_LABELS[method] || method); - alertify.message(`Upscaling ${scale}× · ${methodLabel}…`, 0); + var isAI = method !== 'lanczos'; + showProgress( + `Upscaling ${scale}× with ${methodLabel}…` + + (isAI ? '\nAI is reconstructing detail — this may take 30–120 seconds.' : ''), + isAI ? 90 : 10 + ); try { var layerCanvas = document.createElement('canvas'); @@ -276,21 +282,21 @@ class Image_upscale_class { ); } - alertify.dismissAll(); + hideProgress(); alertify.success( `${result.output.width}×${result.output.height}px · ${usedLabel}` ); this.isProcessing = false; }; img.onerror = () => { - alertify.dismissAll(); + hideProgress(); alertify.error('Failed to load upscaled image.'); this.isProcessing = false; }; img.src = 'data:image/png;base64,' + result.result; } catch (err) { - alertify.dismissAll(); + hideProgress(); alertify.error('Upscale failed: ' + (err.message || err)); this.isProcessing = false; } diff --git a/frontend/src/js/tools/selection_actions.js b/frontend/src/js/tools/selection_actions.js index af6f9b7..8cbec02 100644 --- a/frontend/src/js/tools/selection_actions.js +++ b/frontend/src/js/tools/selection_actions.js @@ -18,6 +18,7 @@ import app from './../app.js'; import config from './../config.js'; import Base_layers_class from './../core/base-layers.js'; import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; +import { showProgress, hideProgress } from './../libs/progress_overlay.js'; const BASE = window.API_BASE_URL || ''; @@ -175,7 +176,7 @@ export class SelectionActions { async _scaleSelection(scalePct) { if (!this._check()) return; this.hide(); - alertify.message('Scaling object and filling gap…'); + showProgress('Scaling object and AI-filling the gap…', 30); try { var res = await _post('/api/image/scale-selection', { image: this._imageData, @@ -184,8 +185,10 @@ export class SelectionActions { }); this.tool.updateLayerWithResult(res.result); this.tool.clearSelection(); + hideProgress(); alertify.success('Scaled by ' + scalePct + '%!'); } catch (e) { + hideProgress(); alertify.error('Scale failed: ' + e.message); } } @@ -193,7 +196,7 @@ export class SelectionActions { async _makeAsymmetric() { if (!this._check()) return; this.hide(); - alertify.message('AI is adding natural asymmetry…'); + showProgress('AI is adding natural asymmetry…', 60); try { var res = await _post('/api/image/ai-edit-region', { image: this._imageData, @@ -205,8 +208,10 @@ export class SelectionActions { }); this.tool.updateLayerWithResult(res.result); this.tool.clearSelection(); + hideProgress(); alertify.success('Made less symmetrical!'); } catch (e) { + hideProgress(); alertify.error('AI edit failed: ' + e.message); } } @@ -214,7 +219,7 @@ export class SelectionActions { async _aiEditRegion(instruction) { if (!this._check()) return; this.hide(); - alertify.message('AI is editing the region…'); + showProgress('AI is editing the region…', 60); try { var res = await _post('/api/image/ai-edit-region', { image: this._imageData, @@ -225,8 +230,10 @@ export class SelectionActions { }); this.tool.updateLayerWithResult(res.result); this.tool.clearSelection(); + hideProgress(); alertify.success('Done!'); } catch (e) { + hideProgress(); alertify.error('AI edit failed: ' + e.message); } } @@ -257,7 +264,7 @@ export class SelectionActions { var clipBase64 = await _blobToBase64(clipBlob); this.hide(); - alertify.message('Pasting clipboard into selection…'); + showProgress('Pasting clipboard into selection…', 10); var res = await _post('/api/image/paste-into-selection', { image: this._imageData, @@ -266,8 +273,10 @@ export class SelectionActions { }); this.tool.updateLayerWithResult(res.result); this.tool.clearSelection(); + hideProgress(); alertify.success('Clipboard pasted into selection!'); } catch (e) { + hideProgress(); alertify.error('Paste failed: ' + e.message); } } From 372ab48991e121bfad99ea843f41d83ecf98a19f Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 13 Jun 2026 16:48:55 +0000 Subject: [PATCH 8/8] Real per-step progress bars for local GPU inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backend: - local_diffusion.py: add _make_step_cb() that writes step/total_steps/ progress into _states on every diffusers callback_on_step_end; wired into txt2img, inpaint, img2img with TypeError fallback for older diffusers - ai_tools.py: GET /api/generate/progress SSE endpoint — streams _states as JSON array every 200ms so clients get live denoising step counts Frontend: - progress_overlay.js: add connectProgressSSE(pipeType, baseUrl) / disconnectProgressSSE() — opens EventSource, maps step/total_steps to bar percentage (0→85% during denoising, 85→100 for decode/place) - text_to_image.js: connect SSE before POST, disconnect on done/error - selection_actions.js: connect SSE for AI edit / asymmetry operations Result: for local GPU, progress bar shows "Step 12 / 30" with exact fill; for remote providers and upscale (no step callbacks), shimmer animates. https://claude.ai/code/session_01WVDg7amsy1TTtxvpku7bcM --- backend/app/routers/ai_tools.py | 33 ++++ backend/app/services/local_diffusion.py | 177 +++++++++++++----- frontend/src/js/libs/progress_overlay.js | 38 ++++ .../src/js/modules/generate/text_to_image.js | 8 +- frontend/src/js/tools/selection_actions.js | 8 +- 5 files changed, 214 insertions(+), 50 deletions(-) diff --git a/backend/app/routers/ai_tools.py b/backend/app/routers/ai_tools.py index be39254..dba991e 100644 --- a/backend/app/routers/ai_tools.py +++ b/backend/app/routers/ai_tools.py @@ -4,10 +4,12 @@ All endpoints are under /api prefix. """ from fastapi import APIRouter, HTTPException +from fastapi.responses import StreamingResponse from pydantic import BaseModel from typing import Optional import base64 import asyncio +import json from app.services.local_inpaint import ( lama_inpaint, opencv_inpaint, lama_available, gpu_available, rembg_available, @@ -191,6 +193,37 @@ async def inpaint_remote(req: InpaintRemoteRequest): raise HTTPException(status_code=500, detail=str(e)) +@router.get("/generate/progress") +async def generation_progress_stream(): + """ + SSE stream of local GPU pipeline inference progress. + Events are JSON arrays of pipeline state objects, emitted every 200 ms. + Each object: {pipeline, state, step, total_steps, progress, message, model_id, …} + Clients open this with EventSource before firing a generation POST, + then close it when the POST resolves. + """ + from app.services.local_diffusion import get_all_model_states + + async def event_gen(): + try: + while True: + states = get_all_model_states() + yield f"data: {json.dumps(states)}\n\n" + await asyncio.sleep(0.2) + except asyncio.CancelledError: + pass + + return StreamingResponse( + event_gen(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + "Connection": "keep-alive", + }, + ) + + @router.post("/generate/txt2img") async def txt2img(req: Txt2ImgRequest): """Text-to-image via configured remote provider.""" diff --git a/backend/app/services/local_diffusion.py b/backend/app/services/local_diffusion.py index b873027..be2c1a1 100644 --- a/backend/app/services/local_diffusion.py +++ b/backend/app/services/local_diffusion.py @@ -48,6 +48,24 @@ def get_all_model_states() -> list[dict]: return list(_states.values()) +def _make_step_cb(pipe_type: str, total_steps: int): + """ + Returns a diffusers callback_on_step_end that writes per-step progress + into _states so the SSE /api/generate/progress endpoint can stream it. + Called from a thread executor — _set_state is thread-safe. + """ + def cb(pipe, step_index: int, timestep, callback_kwargs: dict) -> dict: + done = step_index + 1 + _set_state(pipe_type, + state="running", + step=done, + total_steps=total_steps, + progress=round(done / total_steps * 85, 1), + message=f"Step {done} / {total_steps}") + return callback_kwargs + return cb + + # ── LRU pipeline cache ──────────────────────────────────────────────────────── class _PipelineCache: @@ -295,18 +313,35 @@ class LocalDiffusionProvider(RemoteAIProvider): steps = int(params.get("steps", 30)) cfg = float(params.get("cfg_scale", 7.5)) neg = params.get("negative_prompt", "") or None + step_cb = _make_step_cb("inpaint", steps) + + _set_state("inpaint", state="running", step=0, total_steps=steps, progress=0, message="Starting…") def _run(): - return pipe( - prompt=prompt, - negative_prompt=neg, - image=img_r, - mask_image=mask_r, - num_inference_steps=steps, - guidance_scale=cfg, - ).images[0].resize(orig, Image.LANCZOS) + try: + return pipe( + prompt=prompt, + negative_prompt=neg, + image=img_r, + mask_image=mask_r, + num_inference_steps=steps, + guidance_scale=cfg, + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).images[0].resize(orig, Image.LANCZOS) + except TypeError: + return pipe( + prompt=prompt, + negative_prompt=neg, + image=img_r, + mask_image=mask_r, + num_inference_steps=steps, + guidance_scale=cfg, + ).images[0].resize(orig, Image.LANCZOS) - return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) + result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) + _set_state("inpaint", state="ready", step=None, total_steps=None, progress=100, message="Ready") + return result async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: pipe = await self._get_pipeline("txt2img") @@ -316,34 +351,61 @@ class LocalDiffusionProvider(RemoteAIProvider): w = min(width, max_dim) // 8 * 8 h = min(height, max_dim) // 8 * 8 seed = int(params.get("seed", 0)) - is_flux = spec.family == "flux" + steps = 4 if is_flux else int(params.get("steps", 30)) + step_cb = _make_step_cb("txt2img", steps) + + _set_state("txt2img", state="running", step=0, total_steps=steps, progress=0, message="Starting…") def _run(): import torch device = self._info.backend gen = torch.Generator(device=device).manual_seed(seed) if seed else None - if is_flux: - return pipe( - prompt=prompt, - width=w, height=h, - num_inference_steps=4, # FLUX.1-schnell is a 4-step model - guidance_scale=0.0, # fully CFG-distilled - max_sequence_length=256, - generator=gen, - ).images[0] - else: - return pipe( - prompt=prompt, - negative_prompt=params.get("negative_prompt", "") or None, - width=w, height=h, - num_inference_steps=int(params.get("steps", 30)), - guidance_scale=float(params.get("cfg_scale", 7.5)), - generator=gen, - ).images[0] + try: + if is_flux: + return pipe( + prompt=prompt, + width=w, height=h, + num_inference_steps=steps, + guidance_scale=0.0, + max_sequence_length=256, + generator=gen, + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).images[0] + else: + return pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + width=w, height=h, + num_inference_steps=steps, + guidance_scale=float(params.get("cfg_scale", 7.5)), + generator=gen, + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).images[0] + except TypeError: + # Older diffusers without callback_on_step_end + if is_flux: + return pipe( + prompt=prompt, width=w, height=h, + num_inference_steps=steps, guidance_scale=0.0, + max_sequence_length=256, generator=gen, + ).images[0] + else: + return pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + width=w, height=h, + num_inference_steps=steps, + guidance_scale=float(params.get("cfg_scale", 7.5)), + generator=gen, + ).images[0] - return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) + result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) + _set_state("txt2img", state="ready", step=None, total_steps=None, progress=100, message="Ready") + return result async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: pipe = await self._get_pipeline("img2img") @@ -353,28 +415,49 @@ class LocalDiffusionProvider(RemoteAIProvider): orig = img.size img_r = _resize_square(img, spec.native_res) is_flux = spec.family == "flux" + steps = 4 if is_flux else int(params.get("steps", 30)) + step_cb = _make_step_cb("img2img", steps) + + _set_state("img2img", state="running", step=0, total_steps=steps, progress=0, message="Starting…") def _run(): - if is_flux: - result = pipe( - prompt=prompt, - image=img_r, - strength=strength, - num_inference_steps=4, - guidance_scale=0.0, - ).images[0] - else: - result = pipe( - prompt=prompt, - negative_prompt=params.get("negative_prompt", "") or None, - image=img_r, - strength=strength, - num_inference_steps=int(params.get("steps", 30)), - guidance_scale=float(params.get("cfg_scale", 7.5)), - ).images[0] + try: + if is_flux: + result = pipe( + prompt=prompt, image=img_r, strength=strength, + num_inference_steps=steps, guidance_scale=0.0, + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).images[0] + else: + result = pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + image=img_r, strength=strength, + num_inference_steps=steps, + guidance_scale=float(params.get("cfg_scale", 7.5)), + callback_on_step_end=step_cb, + callback_on_step_end_tensor_inputs=["latents"], + ).images[0] + except TypeError: + if is_flux: + result = pipe( + prompt=prompt, image=img_r, strength=strength, + num_inference_steps=steps, guidance_scale=0.0, + ).images[0] + else: + result = pipe( + prompt=prompt, + negative_prompt=params.get("negative_prompt", "") or None, + image=img_r, strength=strength, + num_inference_steps=steps, + guidance_scale=float(params.get("cfg_scale", 7.5)), + ).images[0] return result.resize(orig, Image.LANCZOS) - return _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) + result = _to_png(await asyncio.get_event_loop().run_in_executor(None, _run)) + _set_state("img2img", state="ready", step=None, total_steps=None, progress=100, message="Ready") + return result async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: from PIL import ImageDraw diff --git a/frontend/src/js/libs/progress_overlay.js b/frontend/src/js/libs/progress_overlay.js index da10cb0..7792e4c 100644 --- a/frontend/src/js/libs/progress_overlay.js +++ b/frontend/src/js/libs/progress_overlay.js @@ -19,6 +19,44 @@ var _shimmerAnim = null; var _fakeTimer = null; var _currentPct = 0; +// ── SSE progress connection ─────────────────────────────────────────────────── + +var _sse = null; + +/** + * Open an EventSource to /api/generate/progress and drive the bar with real + * denoising step counts from the local GPU pipeline. + * + * @param {string} pipeType - 'txt2img' | 'inpaint' | 'img2img' + * @param {string} baseUrl - window.API_BASE_URL or '' + */ +export function connectProgressSSE(pipeType, baseUrl) { + disconnectProgressSSE(); + try { + var url = (baseUrl || '') + '/api/generate/progress'; + _sse = new EventSource(url); + _sse.onmessage = (e) => { + try { + var states = JSON.parse(e.data); + var s = Array.isArray(states) + ? states.find(st => st.pipeline === pipeType) + : null; + if (s && s.state === 'running' && s.total_steps) { + var pct = Math.round(s.step / s.total_steps * 85); + updateProgress(pct, s.message || `Step ${s.step} / ${s.total_steps}`); + } + } catch { /* malformed event — ignore */ } + }; + _sse.onerror = () => disconnectProgressSSE(); + } catch { /* SSE not supported */ } +} + +export function disconnectProgressSSE() { + if (_sse) { _sse.close(); _sse = null; } +} + +// ── Progress overlay ────────────────────────────────────────────────────────── + export function showProgress(message, estimatedSeconds) { hideProgress(); diff --git a/frontend/src/js/modules/generate/text_to_image.js b/frontend/src/js/modules/generate/text_to_image.js index 86f4853..82b5aa4 100644 --- a/frontend/src/js/modules/generate/text_to_image.js +++ b/frontend/src/js/modules/generate/text_to_image.js @@ -12,7 +12,7 @@ import Dialog_class from './../../libs/popup.js'; import alertify from './../../../../node_modules/alertifyjs/build/alertify.min.js'; import apiService from './../../services/api.js'; import { getCapabilities } from './../../api/capabilities.js'; -import { showProgress, updateProgress, hideProgress } from './../../libs/progress_overlay.js'; +import { showProgress, updateProgress, hideProgress, connectProgressSSE, disconnectProgressSSE } from './../../libs/progress_overlay.js'; var instance = null; @@ -142,7 +142,8 @@ class Generate_text_to_image_class { if (this.isProcessing) return; this.isProcessing = true; - showProgress('Generating image… this may take a minute on local GPU', estSec || 60); + connectProgressSSE('txt2img', window.API_BASE_URL || ''); + showProgress('Generating image…', estSec || 60); try { var result = await apiService.textToImage(params.prompt, { @@ -185,11 +186,13 @@ class Generate_text_to_image_class { ]) ); } + disconnectProgressSSE(); hideProgress(); alertify.success('Image generated!'); this.isProcessing = false; }; img.onerror = () => { + disconnectProgressSSE(); hideProgress(); alertify.error('Failed to load generated image.'); this.isProcessing = false; @@ -197,6 +200,7 @@ class Generate_text_to_image_class { img.src = 'data:image/png;base64,' + result.result; } catch (err) { + disconnectProgressSSE(); hideProgress(); alertify.error('Generation failed: ' + (err.message || err)); this.isProcessing = false; diff --git a/frontend/src/js/tools/selection_actions.js b/frontend/src/js/tools/selection_actions.js index 8cbec02..caaa404 100644 --- a/frontend/src/js/tools/selection_actions.js +++ b/frontend/src/js/tools/selection_actions.js @@ -18,7 +18,7 @@ import app from './../app.js'; import config from './../config.js'; import Base_layers_class from './../core/base-layers.js'; import alertify from './../../../node_modules/alertifyjs/build/alertify.min.js'; -import { showProgress, hideProgress } from './../libs/progress_overlay.js'; +import { showProgress, updateProgress, hideProgress, connectProgressSSE, disconnectProgressSSE } from './../libs/progress_overlay.js'; const BASE = window.API_BASE_URL || ''; @@ -196,6 +196,7 @@ export class SelectionActions { async _makeAsymmetric() { if (!this._check()) return; this.hide(); + connectProgressSSE('inpaint', window.API_BASE_URL || ''); showProgress('AI is adding natural asymmetry…', 60); try { var res = await _post('/api/image/ai-edit-region', { @@ -208,9 +209,11 @@ export class SelectionActions { }); this.tool.updateLayerWithResult(res.result); this.tool.clearSelection(); + disconnectProgressSSE(); hideProgress(); alertify.success('Made less symmetrical!'); } catch (e) { + disconnectProgressSSE(); hideProgress(); alertify.error('AI edit failed: ' + e.message); } @@ -219,6 +222,7 @@ export class SelectionActions { async _aiEditRegion(instruction) { if (!this._check()) return; this.hide(); + connectProgressSSE('inpaint', window.API_BASE_URL || ''); showProgress('AI is editing the region…', 60); try { var res = await _post('/api/image/ai-edit-region', { @@ -230,9 +234,11 @@ export class SelectionActions { }); this.tool.updateLayerWithResult(res.result); this.tool.clearSelection(); + disconnectProgressSSE(); hideProgress(); alertify.success('Done!'); } catch (e) { + disconnectProgressSSE(); hideProgress(); alertify.error('AI edit failed: ' + e.message); }