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()