root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
9a53ee67e784f133fc5c56e7aa2808248fa26edd
PaintPlus/backend/app/services/remote_provider.py
T
Claude 8fe8498df2 Add local GPU inference: auto-detect GPU, auto-download best diffusion models 
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
2026-06-13 15:08:41 +00:00

474 lines
24 KiB
Python
Raw Blame History

"""
Remote AI provider abstraction.
One interface, three drivers: OpenAI, InvokeAI, ComfyUI.
Configure one provider via AI_PROVIDER in .env.
"""
from abc import ABC, abstractmethod
from typing import Optional
import httpx
import base64
import asyncio
from io import BytesIO
class RemoteAIProvider(ABC):
@abstractmethod
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes: ...
@abstractmethod
async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes: ...
@abstractmethod
async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes: ...
@abstractmethod
async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes: ...
@abstractmethod
async def health(self) -> bool: ...
@abstractmethod
def capabilities(self) -> list[str]: ...
class OpenAIRemoteProvider(RemoteAIProvider):
"""OpenAI image API — gpt-image-1 / dall-e-3."""
def __init__(self, api_key: str, model: str = "dall-e-3"):
self.api_key = api_key
self.model = model
self.base_url = "https://api.openai.com/v1"
def _headers(self):
return {"Authorization": f"Bearer {self.api_key}"}
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
async with httpx.AsyncClient(timeout=120.0) as client:
files = {
"image": ("image.png", image_bytes, "image/png"),
"mask": ("mask.png", mask_bytes, "image/png"),
}
data = {"prompt": prompt, "n": "1", "size": "1024x1024"}
r = await client.post(f"{self.base_url}/images/edits", files=files, data=data, headers=self._headers())
r.raise_for_status()
url = r.json()["data"][0]["url"]
img_r = await client.get(url)
img_r.raise_for_status()
return img_r.content
async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes:
size = f"{width}x{height}" if f"{width}x{height}" in {"256x256", "512x512", "1024x1024"} else "1024x1024"
async with httpx.AsyncClient(timeout=120.0) as client:
data = {"model": self.model, "prompt": prompt, "n": 1, "size": size}
r = await client.post(f"{self.base_url}/images/generations", json=data, headers=self._headers())
r.raise_for_status()
url = r.json()["data"][0]["url"]
img_r = await client.get(url)
img_r.raise_for_status()
return img_r.content
async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes:
# OpenAI doesn't have img2img natively — use edits with blank mask
from PIL import Image
import numpy as np
img = Image.open(BytesIO(image_bytes)).convert("RGBA")
mask = Image.new("RGBA", img.size, (0, 0, 0, 0))
mask_buf = BytesIO()
mask.save(mask_buf, format="PNG")
return await self.inpaint(image_bytes, mask_buf.getvalue(), prompt, params)
async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes:
from PIL import Image
img = Image.open(BytesIO(image_bytes)).convert("RGBA")
w, h = img.size
directions = {"left": (size, 0), "right": (size, 0), "top": (0, size), "bottom": (0, size)}
dw, dh = directions.get(direction, (size, 0))
new_w, new_h = w + dw, h + dh
canvas = Image.new("RGBA", (new_w, new_h), (0, 0, 0, 0))
offsets = {
"left": (size, 0), "right": (0, 0), "top": (0, size), "bottom": (0, 0)
}
ox, oy = offsets.get(direction, (0, 0))
canvas.paste(img, (ox, oy))
# mask: transparent = inpaint
mask = Image.new("L", (new_w, new_h), 0)
# fill the expanded region with white in mask
import numpy as np
mask_arr = np.zeros((new_h, new_w), dtype=np.uint8)
if direction == "left":
mask_arr[:, :size] = 255
elif direction == "right":
mask_arr[:, w:] = 255
elif direction == "top":
mask_arr[:size, :] = 255
else:
mask_arr[h:, :] = 255
mask = Image.fromarray(mask_arr, "L")
canvas_rgb = canvas.convert("RGB")
img_buf = BytesIO()
canvas_rgb.save(img_buf, format="PNG")
mask_buf = BytesIO()
mask.save(mask_buf, format="PNG")
return await self.inpaint(img_buf.getvalue(), mask_buf.getvalue(), prompt, {})
async def health(self) -> bool:
try:
async with httpx.AsyncClient(timeout=10.0) as client:
r = await client.get(f"{self.base_url}/models", headers=self._headers())
return r.status_code == 200
except Exception:
return False
def capabilities(self) -> list[str]:
return ["inpaint", "txt2img", "img2img", "outpaint"]
class InvokeAIProvider(RemoteAIProvider):
"""InvokeAI REST API driver — supports Flux, SDXL, SD1.5 and more."""
def __init__(self, base_url: str, default_model: str = "flux-dev"):
self.base_url = base_url.rstrip("/")
self.default_model = default_model
async def _b64(self, data: bytes) -> str:
return base64.b64encode(data).decode()
async def _upload_image(self, client: httpx.AsyncClient, image_bytes: bytes, category: str = "general") -> str:
"""Upload image to InvokeAI and return image_name."""
files = {"file": ("image.png", image_bytes, "image/png")}
data = {"image_category": category, "is_intermediate": "false"}
r = await client.post(f"{self.base_url}/api/v1/images/upload", files=files, data=data)
r.raise_for_status()
return r.json()["image_name"]
async def _run_graph(self, client: httpx.AsyncClient, graph: dict) -> bytes:
"""Post a graph, poll for completion, return result image bytes."""
r = await client.post(f"{self.base_url}/api/v1/queue/default/enqueue_batch",
json={"prepend": False, "batch": {"graph": graph, "runs": 1}})
r.raise_for_status()
batch_id = r.json()["batch"]["batch_id"]
# Poll queue status
for _ in range(180):
await asyncio.sleep(2)
sr = await client.get(f"{self.base_url}/api/v1/queue/default/status")
sr.raise_for_status()
status = sr.json()
if status.get("queue", {}).get("completed", 0) > 0:
break
if status.get("queue", {}).get("failed", 0) > 0:
raise RuntimeError("InvokeAI graph failed")
# Fetch latest result image
lr = await client.get(f"{self.base_url}/api/v1/images/?categories=general&limit=1&is_intermediate=false")
lr.raise_for_status()
items = lr.json().get("items", [])
if not items:
raise RuntimeError("No output image from InvokeAI")
img_name = items[0]["image_name"]
img_r = await client.get(f"{self.base_url}/api/v1/images/i/{img_name}/full")
img_r.raise_for_status()
return img_r.content
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
async with httpx.AsyncClient(timeout=300.0) as client:
img_name = await self._upload_image(client, image_bytes)
mask_name = await self._upload_image(client, mask_bytes, "mask")
model = params.get("model", self.default_model)
graph = {
"id": "inpaint_graph",
"nodes": {
"img_node": {"id": "img_node", "type": "image", "image": {"image_name": img_name}},
"mask_node": {"id": "mask_node", "type": "image", "image": {"image_name": mask_name}},
"model_node": {"id": "model_node", "type": "main_model_loader", "model": {"model_name": model, "base": "any"}},
"clip_skip": {"id": "clip_skip", "type": "clip_skip", "skipped_layers": 0},
"positive": {"id": "positive", "type": "compel", "prompt": prompt},
"negative": {"id": "negative", "type": "compel", "prompt": params.get("negative_prompt", "")},
"denoise": {
"id": "denoise", "type": "denoise_latents",
"steps": params.get("steps", 30),
"cfg_scale": params.get("cfg_scale", 7.5),
"denoising_start": 0.0, "denoising_end": 1.0,
"scheduler": "euler", "is_intermediate": False
},
"vae_loader": {"id": "vae_loader", "type": "vae_loader", "vae_model": {"model_name": model, "base": "any"}},
"img_to_latents": {"id": "img_to_latents", "type": "i2l"},
"latents_to_img": {"id": "latents_to_img", "type": "l2i"},
},
"edges": [
{"source": {"node_id": "model_node", "field": "unet"}, "destination": {"node_id": "denoise", "field": "unet"}},
{"source": {"node_id": "model_node", "field": "clip"}, "destination": {"node_id": "clip_skip", "field": "clip"}},
{"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "positive", "field": "clip"}},
{"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "negative", "field": "clip"}},
{"source": {"node_id": "positive", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "positive_conditioning"}},
{"source": {"node_id": "negative", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "negative_conditioning"}},
{"source": {"node_id": "img_node", "field": "image"}, "destination": {"node_id": "img_to_latents", "field": "image"}},
{"source": {"node_id": "vae_loader", "field": "vae"}, "destination": {"node_id": "img_to_latents", "field": "vae"}},
{"source": {"node_id": "img_to_latents", "field": "latents"}, "destination": {"node_id": "denoise", "field": "latents"}},
{"source": {"node_id": "mask_node", "field": "image"}, "destination": {"node_id": "denoise", "field": "mask"}},
{"source": {"node_id": "denoise", "field": "latents"}, "destination": {"node_id": "latents_to_img", "field": "latents"}},
{"source": {"node_id": "vae_loader", "field": "vae"}, "destination": {"node_id": "latents_to_img", "field": "vae"}},
]
}
return await self._run_graph(client, graph)
async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes:
async with httpx.AsyncClient(timeout=300.0) as client:
model = params.get("model", self.default_model)
graph = {
"id": "txt2img_graph",
"nodes": {
"model_node": {"id": "model_node", "type": "main_model_loader", "model": {"model_name": model, "base": "any"}},
"clip_skip": {"id": "clip_skip", "type": "clip_skip", "skipped_layers": 0},
"positive": {"id": "positive", "type": "compel", "prompt": prompt},
"negative": {"id": "negative", "type": "compel", "prompt": params.get("negative_prompt", "")},
"noise": {"id": "noise", "type": "noise", "width": width, "height": height, "seed": params.get("seed", 0)},
"denoise": {
"id": "denoise", "type": "denoise_latents",
"steps": params.get("steps", 30),
"cfg_scale": params.get("cfg_scale", 7.5),
"denoising_start": 0.0, "denoising_end": 1.0,
"scheduler": "euler",
},
"vae_loader": {"id": "vae_loader", "type": "vae_loader", "vae_model": {"model_name": model, "base": "any"}},
"latents_to_img": {"id": "latents_to_img", "type": "l2i"},
},
"edges": [
{"source": {"node_id": "model_node", "field": "unet"}, "destination": {"node_id": "denoise", "field": "unet"}},
{"source": {"node_id": "model_node", "field": "clip"}, "destination": {"node_id": "clip_skip", "field": "clip"}},
{"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "positive", "field": "clip"}},
{"source": {"node_id": "clip_skip", "field": "clip"}, "destination": {"node_id": "negative", "field": "clip"}},
{"source": {"node_id": "positive", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "positive_conditioning"}},
{"source": {"node_id": "negative", "field": "conditioning"}, "destination": {"node_id": "denoise", "field": "negative_conditioning"}},
{"source": {"node_id": "noise", "field": "noise"}, "destination": {"node_id": "denoise", "field": "noise"}},
{"source": {"node_id": "denoise", "field": "latents"}, "destination": {"node_id": "latents_to_img", "field": "latents"}},
{"source": {"node_id": "vae_loader", "field": "vae"}, "destination": {"node_id": "latents_to_img", "field": "vae"}},
]
}
return await self._run_graph(client, graph)
async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes:
# Reuse inpaint with a full-white mask at the given strength
from PIL import Image
img = Image.open(BytesIO(image_bytes))
mask = Image.new("L", img.size, 255)
mask_buf = BytesIO()
mask.save(mask_buf, format="PNG")
p = dict(params)
p.setdefault("denoising_start", 1.0 - strength)
return await self.inpaint(image_bytes, mask_buf.getvalue(), prompt, p)
async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes:
# Delegate to inpaint with expanded canvas
provider = OpenAIRemoteProvider.__new__(OpenAIRemoteProvider)
return await provider.outpaint(image_bytes, direction, size, prompt)
async def health(self) -> bool:
try:
async with httpx.AsyncClient(timeout=10.0) as client:
r = await client.get(f"{self.base_url}/api/v1/app/version")
return r.status_code == 200
except Exception:
return False
def capabilities(self) -> list[str]:
return ["inpaint", "txt2img", "img2img", "outpaint"]
class ComfyUIProvider(RemoteAIProvider):
"""ComfyUI workflow JSON API driver."""
def __init__(self, base_url: str, default_model: str = "v1-5-pruned-emaonly.ckpt"):
self.base_url = base_url.rstrip("/")
self.default_model = default_model
async def _upload_image(self, client: httpx.AsyncClient, image_bytes: bytes, name: str = "image.png") -> str:
files = {"image": (name, image_bytes, "image/png")}
data = {"overwrite": "true"}
r = await client.post(f"{self.base_url}/upload/image", files=files, data=data)
r.raise_for_status()
j = r.json()
return j.get("name", name)
async def _queue_prompt(self, client: httpx.AsyncClient, workflow: dict) -> str:
r = await client.post(f"{self.base_url}/prompt", json={"prompt": workflow})
r.raise_for_status()
return r.json()["prompt_id"]
async def _wait_for_result(self, client: httpx.AsyncClient, prompt_id: str) -> bytes:
for _ in range(180):
await asyncio.sleep(2)
r = await client.get(f"{self.base_url}/history/{prompt_id}")
r.raise_for_status()
history = r.json()
if prompt_id in history:
outputs = history[prompt_id].get("outputs", {})
for node_output in outputs.values():
for img_info in node_output.get("images", []):
img_r = await client.get(
f"{self.base_url}/view",
params={"filename": img_info["filename"], "subfolder": img_info.get("subfolder", ""),
"type": img_info.get("type", "output")}
)
img_r.raise_for_status()
return img_r.content
raise RuntimeError("ComfyUI timed out waiting for result")
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
model = params.get("model", self.default_model)
async with httpx.AsyncClient(timeout=300.0) as client:
img_name = await self._upload_image(client, image_bytes, "input.png")
mask_name = await self._upload_image(client, mask_bytes, "mask.png")
workflow = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": model}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["1", 1]}},
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": params.get("negative_prompt", ""), "clip": ["1", 1]}},
"4": {"class_type": "LoadImage", "inputs": {"image": img_name}},
"5": {"class_type": "LoadImage", "inputs": {"image": mask_name}},
"6": {"class_type": "VAEEncode", "inputs": {"pixels": ["4", 0], "vae": ["1", 2]}},
"7": {"class_type": "KSampler", "inputs": {
"model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0],
"latent_image": ["6", 0], "mask": ["5", 0],
"seed": params.get("seed", 42), "steps": params.get("steps", 20),
"cfg": params.get("cfg_scale", 7.0), "sampler_name": "euler",
"scheduler": "normal", "denoise": params.get("denoise", 1.0)
}},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
"9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "api_out"}},
}
pid = await self._queue_prompt(client, workflow)
return await self._wait_for_result(client, pid)
async def txt2img(self, prompt: str, width: int, height: int, params: dict) -> bytes:
model = params.get("model", self.default_model)
async with httpx.AsyncClient(timeout=300.0) as client:
workflow = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": model}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["1", 1]}},
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": params.get("negative_prompt", ""), "clip": ["1", 1]}},
"4": {"class_type": "EmptyLatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}},
"5": {"class_type": "KSampler", "inputs": {
"model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0],
"latent_image": ["4", 0],
"seed": params.get("seed", 42), "steps": params.get("steps", 20),
"cfg": params.get("cfg_scale", 7.0), "sampler_name": "euler",
"scheduler": "normal", "denoise": 1.0
}},
"6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
"7": {"class_type": "SaveImage", "inputs": {"images": ["6", 0], "filename_prefix": "api_out"}},
}
pid = await self._queue_prompt(client, workflow)
return await self._wait_for_result(client, pid)
async def img2img(self, image_bytes: bytes, prompt: str, strength: float, params: dict) -> bytes:
from PIL import Image
img = Image.open(BytesIO(image_bytes))
mask = Image.new("L", img.size, 255)
mask_buf = BytesIO()
mask.save(mask_buf, format="PNG")
p = dict(params)
p["denoise"] = strength
return await self.inpaint(image_bytes, mask_buf.getvalue(), prompt, p)
async def outpaint(self, image_bytes: bytes, direction: str, size: int, prompt: str) -> bytes:
# Build expanded canvas then inpaint with blank mask
from PIL import Image
import numpy as np
img = Image.open(BytesIO(image_bytes)).convert("RGB")
w, h = img.size
dw = size if direction in ("left", "right") else 0
dh = size if direction in ("top", "bottom") else 0
canvas = Image.new("RGB", (w + dw, h + dh), (128, 128, 128))
ox = size if direction == "left" else 0
oy = size if direction == "top" else 0
canvas.paste(img, (ox, oy))
mask_arr = np.zeros((h + dh, w + dw), dtype=np.uint8)
if direction == "left":
mask_arr[:, :size] = 255
elif direction == "right":
mask_arr[:, w:] = 255
elif direction == "top":
mask_arr[:size, :] = 255
else:
mask_arr[h:, :] = 255
img_buf = BytesIO()
canvas.save(img_buf, format="PNG")
mask_buf = BytesIO()
Image.fromarray(mask_arr, "L").save(mask_buf, format="PNG")
return await self.inpaint(img_buf.getvalue(), mask_buf.getvalue(), prompt, {})
async def health(self) -> bool:
try:
async with httpx.AsyncClient(timeout=10.0) as client:
r = await client.get(f"{self.base_url}/system_stats")
return r.status_code == 200
except Exception:
return False
def capabilities(self) -> list[str]:
return ["inpaint", "txt2img", "img2img", "outpaint"]
def _build_provider(name: str) -> Optional[RemoteAIProvider]:
"""Instantiate a named provider from current settings."""
from app.config import settings
name = (name or "").lower().strip()
if name == "openai":
if not settings.openai_api_key:
return None
return OpenAIRemoteProvider(settings.openai_api_key, settings.openai_model)
if name == "invokeai":
if not settings.invokeai_url:
return None
return InvokeAIProvider(settings.invokeai_url, settings.invokeai_default_model)
if name == "comfyui":
if not settings.comfyui_url:
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
# Map operation names to the settings field that holds the override
_OP_FIELD = {
"inpaint": "ai_provider_inpaint",
"txt2img": "ai_provider_txt2img",
"img2img": "ai_provider_img2img",
"outpaint": "ai_provider_outpaint",
}
def get_remote_provider(operation: Optional[str] = None) -> Optional[RemoteAIProvider]:
"""
Return the provider for a given operation.
Resolution order:
1. Per-operation override (AI_PROVIDER_INPAINT, AI_PROVIDER_TXT2IMG, etc.)
2. Global default (AI_PROVIDER)
3. None (local-only mode)
Example .env for mixed setup:
AI_PROVIDER=invokeai # default for inpaint/img2img/outpaint
AI_PROVIDER_TXT2IMG=openai # use OpenAI only for text-to-image
"""
from app.config import settings
if operation and operation in _OP_FIELD:
override = getattr(settings, _OP_FIELD[operation], "")
if override:
provider = _build_provider(override)
if provider is not None:
return provider
# override configured but not usable (missing key/url) — fall through to default
return _build_provider(settings.ai_provider)
Reference in New Issue View Git Blame Copy Permalink