paintplus: vendor the app source and rename from EditmaskwithAI
Bring the full EditmaskwithAI application into the repo under paintplus/ (429 files) so the service is self-contained — the installer copies the vendored source to ~/docker/paintplus/src instead of cloning at runtime. Rename to PaintPlus (service + branding; app logic untouched): - services/editmaskwithai.sh -> services/paintplus.sh (register_service paintplus, install_paintplus, ~/docker/paintplus, Caddy paintplus:8000, Authelia option preserved) - container names -> paintplus across docker-compose*.yml; dev network -> paintplus-network - browser <title> -> "PaintPlus - AI Image Editor"; README heading -> PaintPlus with upstream provenance note - README utilities table: editmaskwithai -> paintplus Backend/frontend code (help strings referencing the old container name, the ai_photo_edit.db filename) is intentionally left as-is to avoid touching application logic. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Nb2vJ8W7bHKx1JXVvpCraH
This commit is contained in:
Claude
@@ -0,0 +1,539 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, Dict
|
||||
import httpx
|
||||
import base64
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class AIProvider(ABC):
|
||||
"""Abstract base class for AI providers"""
|
||||
|
||||
@abstractmethod
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Edit an image patch using AI
|
||||
|
||||
Args:
|
||||
patch_image_bytes: The cropped patch to edit
|
||||
mask_image_bytes: Binary mask (same size as patch)
|
||||
prompt: Text description of desired changes
|
||||
mode: "A" (patch only) or "B" (patch + full image reference)
|
||||
full_image_bytes: Full image for context (mode B only)
|
||||
model: Optional specific model to use
|
||||
|
||||
Returns:
|
||||
Regenerated patch as bytes
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Generate an image from text prompt
|
||||
|
||||
Args:
|
||||
prompt: Text description of desired image
|
||||
width: Image width in pixels
|
||||
height: Image height in pixels
|
||||
model: Optional specific model to use
|
||||
negative_prompt: What to avoid in the generation
|
||||
|
||||
Returns:
|
||||
Generated image as bytes
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class OpenAIProvider(AIProvider):
|
||||
"""OpenAI DALL-E 2 based image editing (NOTE: Lower quality than DALL-E 3)"""
|
||||
|
||||
def __init__(self, api_key: str):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://api.openai.com/v1"
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Edit image using OpenAI DALL-E 2 (NOTE: Uses older model, lower quality)"""
|
||||
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
files = {
|
||||
'image': ('image.png', patch_image_bytes, 'image/png'),
|
||||
'mask': ('mask.png', mask_image_bytes, 'image/png'),
|
||||
}
|
||||
|
||||
data = {
|
||||
'prompt': prompt,
|
||||
'n': 1,
|
||||
'size': '1024x1024' # Will be adjusted based on input
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/images/edits",
|
||||
files=files,
|
||||
data=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Download the generated image
|
||||
image_url = result['data'][0]['url']
|
||||
image_response = await client.get(image_url)
|
||||
image_response.raise_for_status()
|
||||
|
||||
return image_response.content
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate image using OpenAI DALL-E"""
|
||||
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
data = {
|
||||
'prompt': prompt,
|
||||
'n': 1,
|
||||
'size': f'{width}x{height}' if width == height else '1024x1024'
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/images/generations",
|
||||
json=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Download the generated image
|
||||
image_url = result['data'][0]['url']
|
||||
image_response = await client.get(image_url)
|
||||
image_response.raise_for_status()
|
||||
|
||||
return image_response.content
|
||||
|
||||
|
||||
class StabilityAIProvider(AIProvider):
|
||||
"""Stability AI based image editing (SDXL Inpainting)"""
|
||||
|
||||
# Available Stability AI engines
|
||||
MODELS = {
|
||||
'sdxl': 'stable-diffusion-xl-1024-v1-0',
|
||||
'sd15': 'stable-diffusion-v1-5',
|
||||
'sd21': 'stable-diffusion-512-v2-1',
|
||||
}
|
||||
|
||||
def __init__(self, api_key: str, default_model: str = 'sdxl'):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://api.stability.ai/v1"
|
||||
self.default_model = default_model
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Edit image using Stability AI SDXL Inpainting"""
|
||||
|
||||
# Select model
|
||||
model_key = model or self.default_model
|
||||
engine_id = self.MODELS.get(model_key, self.MODELS['sdxl'])
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
files = {
|
||||
'init_image': ('image.png', patch_image_bytes, 'image/png'),
|
||||
'mask_image': ('mask.png', mask_image_bytes, 'image/png'),
|
||||
}
|
||||
|
||||
# Optimized parameters for better quality
|
||||
data = {
|
||||
'text_prompts[0][text]': prompt,
|
||||
'text_prompts[0][weight]': '1.0',
|
||||
'cfg_scale': '8', # Increased for better prompt adherence
|
||||
'samples': '1',
|
||||
'steps': '40', # Increased for better quality
|
||||
'mask_source': 'MASK_IMAGE_WHITE', # White areas are inpainted
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Accept': 'application/json'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/generation/{engine_id}/image-to-image/masking",
|
||||
files=files,
|
||||
data=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Decode base64 image
|
||||
image_data = result['artifacts'][0]['base64']
|
||||
return base64.b64decode(image_data)
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate image using Stability AI SDXL"""
|
||||
|
||||
# Select model
|
||||
model_key = model or self.default_model
|
||||
engine_id = self.MODELS.get(model_key, self.MODELS['sdxl'])
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Build prompts array
|
||||
data = {
|
||||
'text_prompts[0][text]': prompt,
|
||||
'text_prompts[0][weight]': '1.0',
|
||||
'cfg_scale': '7',
|
||||
'samples': '1',
|
||||
'steps': '50',
|
||||
'height': str(height),
|
||||
'width': str(width),
|
||||
}
|
||||
|
||||
# Add negative prompt if provided
|
||||
if negative_prompt:
|
||||
data['text_prompts[1][text]'] = negative_prompt
|
||||
data['text_prompts[1][weight]'] = '-1.0'
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Accept': 'application/json'
|
||||
}
|
||||
|
||||
response = await client.post(
|
||||
f"{self.base_url}/generation/{engine_id}/text-to-image",
|
||||
data=data,
|
||||
headers=headers
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Decode base64 image
|
||||
image_data = result['artifacts'][0]['base64']
|
||||
return base64.b64decode(image_data)
|
||||
|
||||
|
||||
class ReplicateProvider(AIProvider):
|
||||
"""Replicate API with multiple model support"""
|
||||
|
||||
# Available Replicate models for inpainting
|
||||
MODELS = {
|
||||
# SDXL Inpainting - Best general purpose
|
||||
'sdxl-inpaint': {
|
||||
'version': 'stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b',
|
||||
'use_case': 'General purpose, high quality',
|
||||
'cost': '~$0.025/image',
|
||||
'best_for': ['general', 'landscapes', 'objects', 'textures']
|
||||
},
|
||||
# LaMa - Best for object removal
|
||||
'lama': {
|
||||
'version': 'andreasjansson/lama:7f4a2e3c95ab83c1d66ea26a66c27f93b64a2e5a3c5f7f4f4f4f4f4f4f4f4f4f',
|
||||
'use_case': 'Object removal and cleanup',
|
||||
'cost': '~$0.002/image',
|
||||
'best_for': ['removal', 'cleanup', 'erase']
|
||||
},
|
||||
# Realistic Vision - Best for human features (faces, bodies, hands)
|
||||
'realistic-vision': {
|
||||
'version': 'stability-ai/stable-diffusion:db21e45d3f7023abc2a46ee38a23973f6dce16bb082a930b0c49861f96d1e5bf',
|
||||
'use_case': 'Human features, realistic photos',
|
||||
'cost': '~$0.020/image',
|
||||
'best_for': ['face', 'body', 'hands', 'portrait', 'person', 'human']
|
||||
},
|
||||
}
|
||||
|
||||
def __init__(self, api_key: str, default_model: str = 'sdxl-inpaint'):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://api.replicate.com/v1"
|
||||
self.default_model = default_model
|
||||
|
||||
def _select_model_from_prompt(self, prompt: str) -> str:
|
||||
"""Auto-select best model based on prompt keywords"""
|
||||
prompt_lower = prompt.lower()
|
||||
|
||||
# Check for removal/cleanup keywords
|
||||
if any(word in prompt_lower for word in ['remove', 'erase', 'delete', 'cleanup']):
|
||||
return 'lama'
|
||||
|
||||
# Check for human feature keywords
|
||||
if any(word in prompt_lower for word in ['hand', 'face', 'body', 'person', 'portrait', 'skin']):
|
||||
return 'realistic-vision'
|
||||
|
||||
# Default to SDXL for general purpose
|
||||
return 'sdxl-inpaint'
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Edit image using Replicate with auto model selection"""
|
||||
|
||||
# Auto-select model if not specified
|
||||
if not model:
|
||||
model = self._select_model_from_prompt(prompt)
|
||||
|
||||
model_config = self.MODELS.get(model, self.MODELS['sdxl-inpaint'])
|
||||
|
||||
# Convert bytes to base64 for Replicate API
|
||||
patch_b64 = base64.b64encode(patch_image_bytes).decode('utf-8')
|
||||
mask_b64 = base64.b64encode(mask_image_bytes).decode('utf-8')
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Create prediction
|
||||
prediction_data = {
|
||||
"version": model_config['version'],
|
||||
"input": {
|
||||
"image": f"data:image/png;base64,{patch_b64}",
|
||||
"mask": f"data:image/png;base64,{mask_b64}",
|
||||
"prompt": prompt,
|
||||
"num_outputs": 1,
|
||||
"guidance_scale": 7.5,
|
||||
"num_inference_steps": 50,
|
||||
}
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
# Start prediction
|
||||
response = await client.post(
|
||||
f"{self.base_url}/predictions",
|
||||
json=prediction_data,
|
||||
headers=headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
prediction = response.json()
|
||||
|
||||
# Poll for completion
|
||||
prediction_url = prediction['urls']['get']
|
||||
max_attempts = 60 # 2 minutes max
|
||||
attempt = 0
|
||||
|
||||
while attempt < max_attempts:
|
||||
await asyncio.sleep(2) # Wait 2 seconds between polls
|
||||
|
||||
status_response = await client.get(prediction_url, headers=headers)
|
||||
status_response.raise_for_status()
|
||||
status_data = status_response.json()
|
||||
|
||||
if status_data['status'] == 'succeeded':
|
||||
# Download result image
|
||||
output_url = status_data['output'][0]
|
||||
image_response = await client.get(output_url)
|
||||
image_response.raise_for_status()
|
||||
return image_response.content
|
||||
|
||||
elif status_data['status'] == 'failed':
|
||||
raise Exception(f"Replicate prediction failed: {status_data.get('error')}")
|
||||
|
||||
attempt += 1
|
||||
|
||||
raise Exception("Replicate prediction timed out")
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate image using Replicate SDXL"""
|
||||
|
||||
# Use SDXL for text-to-image
|
||||
model_version = 'stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b'
|
||||
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Create prediction
|
||||
prediction_data = {
|
||||
"version": model_version,
|
||||
"input": {
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_outputs": 1,
|
||||
"guidance_scale": 7.5,
|
||||
"num_inference_steps": 50,
|
||||
}
|
||||
}
|
||||
|
||||
# Add negative prompt if provided
|
||||
if negative_prompt:
|
||||
prediction_data["input"]["negative_prompt"] = negative_prompt
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
# Start prediction
|
||||
response = await client.post(
|
||||
f"{self.base_url}/predictions",
|
||||
json=prediction_data,
|
||||
headers=headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
prediction = response.json()
|
||||
|
||||
# Poll for completion
|
||||
prediction_url = prediction['urls']['get']
|
||||
max_attempts = 60
|
||||
attempt = 0
|
||||
|
||||
while attempt < max_attempts:
|
||||
await asyncio.sleep(2)
|
||||
|
||||
status_response = await client.get(prediction_url, headers=headers)
|
||||
status_response.raise_for_status()
|
||||
status_data = status_response.json()
|
||||
|
||||
if status_data['status'] == 'succeeded':
|
||||
# Download result image
|
||||
output_url = status_data['output'][0]
|
||||
image_response = await client.get(output_url)
|
||||
image_response.raise_for_status()
|
||||
return image_response.content
|
||||
|
||||
elif status_data['status'] == 'failed':
|
||||
raise Exception(f"Replicate prediction failed: {status_data.get('error')}")
|
||||
|
||||
attempt += 1
|
||||
|
||||
raise Exception("Replicate text-to-image timed out")
|
||||
|
||||
|
||||
class MockAIProvider(AIProvider):
|
||||
"""Mock provider for testing (returns original patch)"""
|
||||
|
||||
async def edit_image(
|
||||
self,
|
||||
patch_image_bytes: bytes,
|
||||
mask_image_bytes: bytes,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
full_image_bytes: Optional[bytes] = None,
|
||||
model: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Return the original patch (for testing)"""
|
||||
return patch_image_bytes
|
||||
|
||||
async def text_to_image(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
model: Optional[str] = None,
|
||||
negative_prompt: Optional[str] = None
|
||||
) -> bytes:
|
||||
"""Generate a placeholder image (for testing)"""
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
# Create a simple placeholder image
|
||||
img = Image.new('RGB', (width, height), color='lightgray')
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
# Draw text
|
||||
text = f"Mock Image\n{width}x{height}\n{prompt[:50]}"
|
||||
draw.text((width//4, height//2), text, fill='black')
|
||||
|
||||
# Convert to bytes
|
||||
buffer = BytesIO()
|
||||
img.save(buffer, format='PNG')
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def get_ai_provider(provider_name: Optional[str] = None, model: Optional[str] = None) -> AIProvider:
|
||||
"""
|
||||
Factory function to get the configured AI provider
|
||||
|
||||
Args:
|
||||
provider_name: Override default provider from settings
|
||||
model: Specific model to use (provider-dependent)
|
||||
|
||||
Returns:
|
||||
AIProvider instance
|
||||
"""
|
||||
|
||||
provider = provider_name or settings.ai_provider
|
||||
provider = provider.lower()
|
||||
|
||||
if provider == "openai":
|
||||
if not settings.openai_api_key:
|
||||
raise ValueError("OpenAI API key not configured")
|
||||
return OpenAIProvider(settings.openai_api_key)
|
||||
|
||||
elif provider == "stability":
|
||||
if not settings.stability_api_key:
|
||||
raise ValueError("Stability AI API key not configured")
|
||||
default_model = model or getattr(settings, 'stability_model', 'sdxl')
|
||||
return StabilityAIProvider(settings.stability_api_key, default_model=default_model)
|
||||
|
||||
elif provider == "replicate":
|
||||
if not settings.replicate_api_key:
|
||||
raise ValueError("Replicate API key not configured")
|
||||
default_model = model or getattr(settings, 'replicate_model', 'sdxl-inpaint')
|
||||
return ReplicateProvider(settings.replicate_api_key, default_model=default_model)
|
||||
|
||||
elif provider == "mock":
|
||||
return MockAIProvider()
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown AI provider: {provider}")
|
||||
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
from datetime import datetime
|
||||
from PIL import Image
|
||||
|
||||
from app.models.edit import Edit
|
||||
from app.models.project import Project
|
||||
from app.services.ai_provider import get_ai_provider
|
||||
from app.utils.image_processing import (
|
||||
bytes_to_image,
|
||||
image_to_bytes,
|
||||
crop_patch,
|
||||
blend_patch,
|
||||
insert_patch,
|
||||
create_mask_from_selection,
|
||||
resize_for_ai,
|
||||
scale_bbox
|
||||
)
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class EditService:
|
||||
"""Service for handling image edits"""
|
||||
|
||||
def __init__(self, data_dir: str = None):
|
||||
self.data_dir = data_dir or settings.data_dir
|
||||
self.ai_provider = get_ai_provider()
|
||||
|
||||
def get_project_dir(self, project_id: int) -> Path:
|
||||
"""Get project directory path"""
|
||||
return Path(self.data_dir) / "projects" / str(project_id)
|
||||
|
||||
def get_edit_dir(self, project_id: int, edit_id: int) -> Path:
|
||||
"""Get edit history directory path"""
|
||||
return self.get_project_dir(project_id) / "history" / str(edit_id)
|
||||
|
||||
def ensure_project_dir(self, project_id: int):
|
||||
"""Ensure project directory structure exists"""
|
||||
project_dir = self.get_project_dir(project_id)
|
||||
project_dir.mkdir(parents=True, exist_ok=True)
|
||||
(project_dir / "history").mkdir(exist_ok=True)
|
||||
|
||||
def get_current_image_path(self, project_id: int) -> Path:
|
||||
"""Get path to current image"""
|
||||
return self.get_project_dir(project_id) / "current.png"
|
||||
|
||||
def get_original_image_path(self, project_id: int) -> Path:
|
||||
"""Get path to original image"""
|
||||
return self.get_project_dir(project_id) / "original.png"
|
||||
|
||||
async def process_edit(
|
||||
self,
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
prompt: str,
|
||||
mode: str,
|
||||
selection_type: str,
|
||||
bbox: Dict[str, int],
|
||||
feather_px: int,
|
||||
selection_data: Optional[Dict] = None
|
||||
) -> str:
|
||||
"""
|
||||
Process an edit request
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
edit_id: Edit ID
|
||||
prompt: AI prompt
|
||||
mode: "A" or "B"
|
||||
selection_type: "rectangle", "ellipse", or "lasso"
|
||||
bbox: Bounding box {x, y, width, height}
|
||||
feather_px: Feather radius in pixels
|
||||
selection_data: Additional selection data (for lasso)
|
||||
|
||||
Returns:
|
||||
Path to the result image
|
||||
"""
|
||||
# Create edit directory
|
||||
edit_dir = self.get_edit_dir(project_id, edit_id)
|
||||
edit_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Load current image
|
||||
current_image_path = self.get_current_image_path(project_id)
|
||||
full_image = Image.open(current_image_path).convert('RGBA')
|
||||
|
||||
# Crop patch from current image
|
||||
original_patch = crop_patch(full_image, bbox)
|
||||
|
||||
# Save original patch
|
||||
original_patch.save(edit_dir / "patch_in.png")
|
||||
|
||||
# Create mask based on selection type
|
||||
mask = create_mask_from_selection(
|
||||
bbox['width'],
|
||||
bbox['height'],
|
||||
selection_type,
|
||||
selection_data or {}
|
||||
)
|
||||
|
||||
# Save mask
|
||||
mask.save(edit_dir / "mask.png")
|
||||
|
||||
# Resize patch and mask for AI if needed
|
||||
patch_for_ai, scale = resize_for_ai(original_patch)
|
||||
mask_for_ai = mask.resize(patch_for_ai.size, Image.Resampling.LANCZOS)
|
||||
|
||||
# Prepare full image for mode B
|
||||
full_image_bytes = None
|
||||
if mode == "B":
|
||||
full_image_for_ai, _ = resize_for_ai(full_image)
|
||||
full_image_bytes = image_to_bytes(full_image_for_ai)
|
||||
|
||||
# Call AI provider
|
||||
regenerated_patch_bytes = await self.ai_provider.edit_image(
|
||||
patch_image_bytes=image_to_bytes(patch_for_ai),
|
||||
mask_image_bytes=image_to_bytes(mask_for_ai),
|
||||
prompt=prompt,
|
||||
mode=mode,
|
||||
full_image_bytes=full_image_bytes
|
||||
)
|
||||
|
||||
# Convert regenerated patch back to PIL Image
|
||||
regenerated_patch = bytes_to_image(regenerated_patch_bytes)
|
||||
|
||||
# Resize back to original patch size if scaled
|
||||
if scale != 1.0:
|
||||
regenerated_patch = regenerated_patch.resize(
|
||||
original_patch.size,
|
||||
Image.Resampling.LANCZOS
|
||||
)
|
||||
|
||||
# Save regenerated patch
|
||||
regenerated_patch.save(edit_dir / "patch_out.png")
|
||||
|
||||
# Blend regenerated patch with original using mask
|
||||
blended_patch = blend_patch(
|
||||
original_patch,
|
||||
regenerated_patch,
|
||||
mask,
|
||||
feather_px
|
||||
)
|
||||
|
||||
# Insert blended patch back into full image
|
||||
result_image = insert_patch(full_image, blended_patch, bbox)
|
||||
|
||||
# Save result
|
||||
result_path = edit_dir / "result.png"
|
||||
result_image.save(result_path)
|
||||
|
||||
# Update current image
|
||||
result_image.save(current_image_path)
|
||||
|
||||
# Save metadata
|
||||
metadata = {
|
||||
'edit_id': edit_id,
|
||||
'project_id': project_id,
|
||||
'prompt': prompt,
|
||||
'mode': mode,
|
||||
'selection_type': selection_type,
|
||||
'bbox': bbox,
|
||||
'feather_px': feather_px,
|
||||
'selection_data': selection_data,
|
||||
'timestamp': datetime.utcnow().isoformat(),
|
||||
'ai_provider': settings.ai_provider
|
||||
}
|
||||
|
||||
with open(edit_dir / "meta.json", 'w') as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
return str(result_path)
|
||||
|
||||
def revert_to_edit(self, project_id: int, edit_id: int) -> str:
|
||||
"""
|
||||
Revert project to a specific edit
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
edit_id: Edit ID to revert to
|
||||
|
||||
Returns:
|
||||
Path to the reverted image
|
||||
"""
|
||||
edit_dir = self.get_edit_dir(project_id, edit_id)
|
||||
result_path = edit_dir / "result.png"
|
||||
|
||||
if not result_path.exists():
|
||||
raise FileNotFoundError(f"Edit {edit_id} result not found")
|
||||
|
||||
# Copy result to current (preserve alpha channel)
|
||||
current_path = self.get_current_image_path(project_id)
|
||||
img = Image.open(result_path)
|
||||
# Preserve original mode to maintain transparency
|
||||
img.save(current_path, format='PNG')
|
||||
|
||||
return str(current_path)
|
||||
|
||||
def reset_to_original(self, project_id: int) -> str:
|
||||
"""
|
||||
Reset project to original image
|
||||
|
||||
Args:
|
||||
project_id: Project ID
|
||||
|
||||
Returns:
|
||||
Path to the original image
|
||||
"""
|
||||
original_path = self.get_original_image_path(project_id)
|
||||
current_path = self.get_current_image_path(project_id)
|
||||
|
||||
if not original_path.exists():
|
||||
raise FileNotFoundError(f"Original image for project {project_id} not found")
|
||||
|
||||
# Copy original to current (preserve alpha channel)
|
||||
img = Image.open(original_path)
|
||||
# Preserve original mode to maintain transparency
|
||||
img.save(current_path, format='PNG')
|
||||
|
||||
return str(current_path)
|
||||
@@ -0,0 +1,391 @@
|
||||
"""
|
||||
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 ≥ 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
|
||||
otherwise → SD 1.5 + sequential CPU offload
|
||||
|
||||
Inpaint always uses SDXL/SD-family (no FLUX inpaint pipeline yet).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
# ── Model specification ───────────────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
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
|
||||
|
||||
|
||||
# ── 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 | sdxl_offload | sd2x | sd2x_low | 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)
|
||||
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}"
|
||||
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_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,
|
||||
recommended=_select_all_models(eff),
|
||||
warnings=warnings,
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
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_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,
|
||||
recommended=_select_all_models(eff / 2),
|
||||
warnings=["Apple MPS: using fp32 (fp16 less stable). Models load slower."],
|
||||
capabilities=_caps(tier),
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# CPU fallback
|
||||
return GpuCapabilities(
|
||||
backend="cpu",
|
||||
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",
|
||||
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"],
|
||||
)
|
||||
|
||||
|
||||
# ── 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)
|
||||
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)
|
||||
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 >= 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:
|
||||
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 >= 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"
|
||||
return "minimal"
|
||||
|
||||
|
||||
def _caps(tier: str) -> list[str]:
|
||||
base = ["txt2img", "inpaint", "img2img", "outpaint"]
|
||||
if tier in ("flux_full", "flux_offload", "sdxl", "sdxl_low", "sdxl_offload"):
|
||||
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":
|
||||
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 == "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). "
|
||||
"Using SD 1.5/2.x. Upgrade to ≥5.5 GB free for SDXL quality."
|
||||
)
|
||||
|
||||
if xf:
|
||||
w.append(
|
||||
"xformers detected — attention VRAM reduced ~20-30%. "
|
||||
"You may be able to run a higher-tier model than listed."
|
||||
)
|
||||
else:
|
||||
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."
|
||||
)
|
||||
|
||||
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 (≈ 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 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)
|
||||
|
||||
|
||||
# ── Singleton ─────────────────────────────────────────────────────────────────
|
||||
|
||||
_cached: Optional[GpuCapabilities] = None
|
||||
|
||||
|
||||
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()
|
||||
}
|
||||
@@ -0,0 +1,584 @@
|
||||
"""
|
||||
Local GPU diffusion provider — HuggingFace Diffusers backend.
|
||||
|
||||
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.
|
||||
|
||||
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.28.0,<0.29.0, transformers, accelerate, safetensors
|
||||
(all in requirements.gpu.txt — pinned <0.29.0 for PyTorch 2.1.x compatibility)
|
||||
"""
|
||||
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 (
|
||||
GpuCapabilities,
|
||||
ModelSpec,
|
||||
get_cached_gpu_info,
|
||||
infer_spec_from_model_id,
|
||||
)
|
||||
from app.services.remote_provider import RemoteAIProvider
|
||||
|
||||
# ── Model 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())
|
||||
|
||||
|
||||
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:
|
||||
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 = self._cache.popitem(last=False)
|
||||
_evict(evicted, evicted_key)
|
||||
self._cache[key] = pipe
|
||||
|
||||
|
||||
def _evict(pipe, key: str):
|
||||
try:
|
||||
import torch
|
||||
pipe.to("cpu")
|
||||
torch.cuda.empty_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):
|
||||
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()
|
||||
|
||||
@property
|
||||
def _info(self) -> GpuCapabilities:
|
||||
return get_cached_gpu_info()
|
||||
|
||||
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:
|
||||
info = self._info
|
||||
spec = _get_spec(pipe_type, info)
|
||||
|
||||
_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 {spec.model_id}…", error="")
|
||||
try:
|
||||
if spec.family == "flux":
|
||||
pipe = _load_flux_pipeline(pipe_type, spec, info)
|
||||
else:
|
||||
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
|
||||
|
||||
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:
|
||||
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 ──────────────────────────────────────────────────────
|
||||
|
||||
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
|
||||
pipe = await self._get_pipeline("inpaint")
|
||||
spec = _get_spec("inpaint", self._info)
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
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
|
||||
step_cb = _make_step_cb("inpaint", steps)
|
||||
|
||||
_set_state("inpaint", state="running", step=0, total_steps=steps, progress=0, message="Starting…")
|
||||
|
||||
def _run():
|
||||
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)
|
||||
|
||||
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")
|
||||
spec = _get_spec("txt2img", self._info)
|
||||
|
||||
max_dim = spec.native_res
|
||||
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
|
||||
|
||||
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]
|
||||
|
||||
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")
|
||||
spec = _get_spec("img2img", self._info)
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
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():
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
img = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
w, h = img.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)
|
||||
ImageDraw.Draw(mask).rectangle(mask_box, fill=255)
|
||||
|
||||
fill_prompt = prompt or "seamless natural continuation of the scene"
|
||||
return await self.inpaint(_to_png(expanded), _to_png(mask), fill_prompt, {})
|
||||
|
||||
async def health(self) -> bool:
|
||||
return True
|
||||
|
||||
def capabilities(self) -> list[str]:
|
||||
return self._info.capabilities
|
||||
|
||||
|
||||
# ── Image utilities ───────────────────────────────────────────────────────────
|
||||
|
||||
def _resize_pair(img: Image.Image, mask: Image.Image, target: int):
|
||||
w, h = img.size
|
||||
scale = target / max(w, h)
|
||||
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)
|
||||
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:
|
||||
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 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
|
||||
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()
|
||||
seen: set[str] = set()
|
||||
|
||||
for op, spec in info.recommended.items():
|
||||
if spec is None or spec.model_id in seen:
|
||||
continue
|
||||
seen.add(spec.model_id)
|
||||
|
||||
# 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
|
||||
|
||||
_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=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 — 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 {spec.model_id}: {exc}")
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
Local inpainting operations — LaMa, OpenCV, and background removal.
|
||||
All operations use GPU automatically if PyTorch detects one, CPU otherwise.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
# Lazy-loaded LaMa model (downloaded on first use, ~100MB)
|
||||
_lama = None
|
||||
|
||||
|
||||
def get_lama():
|
||||
global _lama
|
||||
if _lama is None:
|
||||
from simple_lama_inpainting import SimpleLama
|
||||
_lama = SimpleLama()
|
||||
return _lama
|
||||
|
||||
|
||||
def lama_available() -> bool:
|
||||
try:
|
||||
import simple_lama_inpainting # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def lama_inpaint(image_bytes: bytes, mask_bytes: bytes) -> bytes:
|
||||
"""LaMa structural inpainting — best for object removal and large fills."""
|
||||
lama = get_lama()
|
||||
image = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
if mask.size != image.size:
|
||||
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
|
||||
result = lama(image, mask)
|
||||
buf = BytesIO()
|
||||
result.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def opencv_inpaint(image_bytes: bytes, mask_bytes: bytes, method: str = "telea") -> bytes:
|
||||
"""OpenCV fast structural inpainting — CPU only, milliseconds."""
|
||||
image = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
mask = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
if mask.size != image.size:
|
||||
mask = mask.resize(image.size, Image.Resampling.LANCZOS)
|
||||
|
||||
img_np = np.array(image)
|
||||
mask_np = np.array(mask)
|
||||
_, mask_bin = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
|
||||
|
||||
flags = cv2.INPAINT_TELEA if method == "telea" else cv2.INPAINT_NS
|
||||
result = cv2.inpaint(img_np, mask_bin, inpaintRadius=3, flags=flags)
|
||||
|
||||
buf = BytesIO()
|
||||
Image.fromarray(result).save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def remove_background_rembg(image_bytes: bytes) -> bytes:
|
||||
"""Background removal using rembg."""
|
||||
from rembg import remove
|
||||
return remove(image_bytes)
|
||||
|
||||
|
||||
def rembg_available() -> bool:
|
||||
try:
|
||||
import rembg # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def gpu_available() -> bool:
|
||||
try:
|
||||
import torch
|
||||
return torch.cuda.is_available()
|
||||
except ImportError:
|
||||
return False
|
||||
@@ -0,0 +1,206 @@
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
from PIL import Image
|
||||
from datetime import datetime
|
||||
|
||||
from app.models.patch import Patch
|
||||
from app.config import settings
|
||||
|
||||
|
||||
class PatchLibraryService:
|
||||
"""Service for managing the patch library"""
|
||||
|
||||
def __init__(self, data_dir: str = None):
|
||||
self.data_dir = data_dir or settings.data_dir
|
||||
self.patch_library_dir = Path(self.data_dir) / "patch_library"
|
||||
self.patch_library_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_patch_path(self, patch_id: int) -> Path:
|
||||
"""Get path to patch file"""
|
||||
return self.patch_library_dir / f"{patch_id}.png"
|
||||
|
||||
def get_thumbnail_path(self, patch_id: int) -> Path:
|
||||
"""Get path to patch thumbnail"""
|
||||
return self.patch_library_dir / f"{patch_id}_thumb.png"
|
||||
|
||||
def create_thumbnail(self, image_path: Path, thumbnail_path: Path, size: tuple = (200, 200)):
|
||||
"""Create a thumbnail from an image"""
|
||||
img = Image.open(image_path)
|
||||
img.thumbnail(size, Image.Resampling.LANCZOS)
|
||||
img.save(thumbnail_path, 'PNG')
|
||||
|
||||
def save_patch_from_file(
|
||||
self,
|
||||
patch_id: int,
|
||||
image_path: str,
|
||||
create_thumb: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Save a patch from an existing file
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
image_path: Source image path
|
||||
create_thumb: Whether to create thumbnail
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
shutil.copy(image_path, patch_path)
|
||||
|
||||
if create_thumb:
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
self.create_thumbnail(patch_path, thumbnail_path)
|
||||
|
||||
return str(patch_path.relative_to(self.data_dir))
|
||||
|
||||
def save_patch_from_bytes(
|
||||
self,
|
||||
patch_id: int,
|
||||
image_bytes: bytes,
|
||||
create_thumb: bool = True
|
||||
) -> str:
|
||||
"""
|
||||
Save a patch from bytes
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
image_bytes: Image data as bytes
|
||||
create_thumb: Whether to create thumbnail
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
|
||||
# Save image
|
||||
with open(patch_path, 'wb') as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
if create_thumb:
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
self.create_thumbnail(patch_path, thumbnail_path)
|
||||
|
||||
return str(patch_path.relative_to(self.data_dir))
|
||||
|
||||
def save_ai_generated_patch(
|
||||
self,
|
||||
patch_id: int,
|
||||
edit_dir: Path
|
||||
) -> str:
|
||||
"""
|
||||
Save an AI-generated patch from an edit
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
edit_dir: Path to edit history directory
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
# Use the AI-generated output (patch_out.png)
|
||||
source_path = edit_dir / "patch_out.png"
|
||||
return self.save_patch_from_file(patch_id, str(source_path))
|
||||
|
||||
def save_manual_patch(
|
||||
self,
|
||||
patch_id: int,
|
||||
project_id: int,
|
||||
bbox: dict
|
||||
) -> str:
|
||||
"""
|
||||
Save a manually selected patch from current project image
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID
|
||||
project_id: Project ID
|
||||
bbox: Bounding box {x, y, width, height}
|
||||
|
||||
Returns:
|
||||
Relative path to saved patch
|
||||
"""
|
||||
from app.services.edit_service import EditService
|
||||
from app.utils.image_processing import crop_patch
|
||||
|
||||
edit_service = EditService(self.data_dir)
|
||||
current_image_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Load and crop current image
|
||||
img = Image.open(current_image_path)
|
||||
patch = crop_patch(img, bbox)
|
||||
|
||||
# Save patch
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
patch.save(patch_path, 'PNG')
|
||||
|
||||
# Create thumbnail
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
self.create_thumbnail(patch_path, thumbnail_path)
|
||||
|
||||
return str(patch_path.relative_to(self.data_dir))
|
||||
|
||||
def apply_patch_to_image(
|
||||
self,
|
||||
patch_id: int,
|
||||
target_image: Image.Image,
|
||||
bbox: dict,
|
||||
feather_px: int = 5
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Apply a saved patch to a target image
|
||||
|
||||
Args:
|
||||
patch_id: Patch ID to apply
|
||||
target_image: Target image to apply patch to
|
||||
bbox: Where to place the patch {x, y, width, height}
|
||||
feather_px: Feather radius for blending
|
||||
|
||||
Returns:
|
||||
Image with patch applied
|
||||
"""
|
||||
from app.utils.image_processing import insert_patch, create_feathered_mask
|
||||
from PIL import ImageOps
|
||||
|
||||
# Load patch
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
patch = Image.open(patch_path).convert('RGBA')
|
||||
|
||||
# Resize patch to match bbox if needed
|
||||
if patch.size != (bbox['width'], bbox['height']):
|
||||
patch = patch.resize((bbox['width'], bbox['height']), Image.Resampling.LANCZOS)
|
||||
|
||||
# Create a soft-edged mask for the patch
|
||||
mask = Image.new('L', patch.size, 255)
|
||||
if feather_px > 0:
|
||||
mask = create_feathered_mask(mask, feather_px)
|
||||
|
||||
# Apply mask to patch
|
||||
patch.putalpha(mask)
|
||||
|
||||
# Insert patch into target image
|
||||
result = insert_patch(target_image, patch, bbox)
|
||||
|
||||
return result
|
||||
|
||||
def delete_patch(self, patch_id: int):
|
||||
"""Delete a patch and its thumbnail"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
thumbnail_path = self.get_thumbnail_path(patch_id)
|
||||
|
||||
if patch_path.exists():
|
||||
patch_path.unlink()
|
||||
|
||||
if thumbnail_path.exists():
|
||||
thumbnail_path.unlink()
|
||||
|
||||
def get_patch_size(self, patch_id: int) -> tuple:
|
||||
"""Get patch dimensions"""
|
||||
patch_path = self.get_patch_path(patch_id)
|
||||
if not patch_path.exists():
|
||||
return (0, 0)
|
||||
|
||||
img = Image.open(patch_path)
|
||||
return img.size
|
||||
@@ -0,0 +1,474 @@
|
||||
"""
|
||||
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, AttributeError) as exc:
|
||||
print(f"[local_gpu] Cannot load diffusion provider: {exc}")
|
||||
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)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
SAM (Segment Anything Model) service.
|
||||
|
||||
Auto-downloads the ViT-B checkpoint (~375 MB) on first use.
|
||||
Caches the loaded model in memory; re-uses predictor across calls.
|
||||
|
||||
Prediction API:
|
||||
predict_points(image_bytes, points, labels) -> mask_bytes (PNG, white=selected)
|
||||
points: list of (x, y) in original image pixels
|
||||
labels: list of 1 (include) or 0 (exclude), same length as points
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# ── Model download ────────────────────────────────────────────────────────────
|
||||
|
||||
SAM_DIR = Path("/app/data/models/sam")
|
||||
SAM_FILENAME = "sam_vit_b_01ec64.pth"
|
||||
SAM_URL = f"https://dl.fbaipublicfiles.com/segment_anything/{SAM_FILENAME}"
|
||||
SAM_PATH = SAM_DIR / SAM_FILENAME
|
||||
|
||||
|
||||
class SamInstallState(str, Enum):
|
||||
idle = "idle"
|
||||
downloading = "downloading"
|
||||
done = "done"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamInstallStatus:
|
||||
state: SamInstallState = SamInstallState.idle
|
||||
progress: int = 0
|
||||
message: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
_install_status = SamInstallStatus()
|
||||
_install_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def get_install_status() -> dict:
|
||||
s = _install_status
|
||||
return {"state": s.state.value, "progress": s.progress,
|
||||
"message": s.message, "error": s.error}
|
||||
|
||||
|
||||
def sam_model_available() -> bool:
|
||||
return SAM_PATH.exists() and SAM_PATH.stat().st_size > 100_000_000
|
||||
|
||||
|
||||
async def ensure_sam_installed() -> bool:
|
||||
"""Download SAM ViT-B checkpoint if not present. Returns True on success."""
|
||||
global _install_status
|
||||
|
||||
if sam_model_available():
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
async with _install_lock:
|
||||
if sam_model_available():
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
if _install_status.state == SamInstallState.downloading:
|
||||
return False
|
||||
|
||||
try:
|
||||
SAM_DIR.mkdir(parents=True, exist_ok=True)
|
||||
_install_status = SamInstallStatus(
|
||||
state=SamInstallState.downloading, progress=0,
|
||||
message="Downloading SAM ViT-B model (~375 MB)…",
|
||||
)
|
||||
|
||||
def _download():
|
||||
def _progress(count, block, total):
|
||||
if total > 0:
|
||||
_install_status.progress = min(99, int(count * block * 99 / total))
|
||||
tmp = SAM_PATH.with_suffix(".tmp")
|
||||
urllib.request.urlretrieve(SAM_URL, tmp, _progress)
|
||||
tmp.rename(SAM_PATH)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, _download)
|
||||
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
except Exception as exc:
|
||||
_install_status = SamInstallStatus(
|
||||
state=SamInstallState.failed, error=str(exc),
|
||||
message="SAM download failed.",
|
||||
)
|
||||
print(f"[sam] Download failed: {exc}")
|
||||
return False
|
||||
|
||||
|
||||
# ── Model cache ───────────────────────────────────────────────────────────────
|
||||
|
||||
_predictor = None
|
||||
_predictor_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _load_predictor():
|
||||
"""Load SAM model and return a SamPredictor. Called in thread pool."""
|
||||
global _predictor
|
||||
if _predictor is not None:
|
||||
return _predictor
|
||||
|
||||
import torch
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
print(f"[sam] Loading SAM ViT-B on {device}…")
|
||||
sam = sam_model_registry["vit_b"](checkpoint=str(SAM_PATH))
|
||||
sam.to(device)
|
||||
_predictor = SamPredictor(sam)
|
||||
print("[sam] Model loaded.")
|
||||
return _predictor
|
||||
|
||||
|
||||
# ── Prediction ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _predict_sync(image_bytes: bytes,
|
||||
points: list[tuple[int, int]],
|
||||
labels: list[int]) -> bytes:
|
||||
"""
|
||||
Run SAM prediction synchronously (call via run_in_executor).
|
||||
Returns PNG bytes: white = selected, black = background.
|
||||
"""
|
||||
predictor = _load_predictor()
|
||||
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
img_array = np.array(image)
|
||||
|
||||
predictor.set_image(img_array)
|
||||
|
||||
pt_array = np.array(points, dtype=np.float32) # [[x, y], ...]
|
||||
lbl_array = np.array(labels, dtype=np.int32) # [1=fg, 0=bg, ...]
|
||||
|
||||
masks, scores, _ = predictor.predict(
|
||||
point_coords=pt_array,
|
||||
point_labels=lbl_array,
|
||||
multimask_output=True,
|
||||
)
|
||||
|
||||
# Pick the highest-confidence mask
|
||||
best = masks[int(np.argmax(scores))] # bool array H×W
|
||||
|
||||
mask_img = Image.fromarray((best * 255).astype(np.uint8), mode="L")
|
||||
buf = io.BytesIO()
|
||||
mask_img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
async def predict_points(image_bytes: bytes,
|
||||
points: list[tuple[int, int]],
|
||||
labels: list[int]) -> bytes:
|
||||
"""Async wrapper for SAM point prediction."""
|
||||
if not sam_model_available():
|
||||
ok = await ensure_sam_installed()
|
||||
if not ok:
|
||||
raise RuntimeError("SAM model not available.")
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, _predict_sync, image_bytes, points, labels)
|
||||
@@ -0,0 +1,494 @@
|
||||
"""
|
||||
Upscale service — auto-detects best available method and runs it.
|
||||
Auto-installs Real-ESRGAN NCNN Vulkan binary when Vulkan GPU is available.
|
||||
Skips NCNN on headless/CPU-only machines and uses PyTorch CPU or Lanczos instead.
|
||||
|
||||
Priority (auto mode):
|
||||
1. Real-ESRGAN PyTorch + CUDA GPU — fastest, best quality
|
||||
2. Real-ESRGAN PyTorch + Apple MPS — fast on Apple Silicon
|
||||
3. Real-ESRGAN NCNN Vulkan binary — fast on any Vulkan GPU
|
||||
4. Real-ESRGAN PyTorch CPU — AI quality, slow (~1-3 min)
|
||||
5. Lanczos — always available, instant
|
||||
|
||||
Capability probe is run once at first call and cached.
|
||||
NCNN binary is auto-downloaded only when Vulkan is detected.
|
||||
Set REALESRGAN_NCNN=force env var to override the Vulkan check.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import urllib.request
|
||||
import zipfile
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
# ── NCNN auto-install ─────────────────────────────────────────────────────────
|
||||
|
||||
NCNN_DEST_DIR = Path("/app/data/models/realesrgan")
|
||||
NCNN_VERSION = "v0.2.5.0"
|
||||
NCNN_BASE_URL = f"https://github.com/xinntao/Real-ESRGAN/releases/download/{NCNN_VERSION}"
|
||||
|
||||
_PLATFORM_ZIP = {
|
||||
"linux": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-ubuntu.zip",
|
||||
"darwin": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-macos.zip",
|
||||
"win32": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-windows.zip",
|
||||
"windows": f"realesrgan-ncnn-vulkan-{NCNN_VERSION}-windows.zip",
|
||||
}
|
||||
|
||||
|
||||
class InstallState(str, Enum):
|
||||
idle = "idle"
|
||||
skipped = "skipped" # headless / no Vulkan
|
||||
downloading = "downloading"
|
||||
extracting = "extracting"
|
||||
verifying = "verifying"
|
||||
done = "done"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class InstallStatus:
|
||||
state: InstallState = InstallState.idle
|
||||
progress: int = 0 # 0-100
|
||||
message: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
_install_status = InstallStatus()
|
||||
_install_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def get_install_status() -> dict:
|
||||
s = _install_status
|
||||
return {
|
||||
"state": s.state.value,
|
||||
"progress": s.progress,
|
||||
"message": s.message,
|
||||
"error": s.error,
|
||||
}
|
||||
|
||||
|
||||
def _ncnn_binary_name() -> str:
|
||||
return "realesrgan-ncnn-vulkan.exe" if "win" in sys.platform.lower() else "realesrgan-ncnn-vulkan"
|
||||
|
||||
|
||||
def _vulkan_available() -> bool:
|
||||
"""
|
||||
Check whether a Vulkan-capable GPU is accessible.
|
||||
Returns True if confident a GPU with Vulkan exists; False on headless/CPU-only.
|
||||
Set REALESRGAN_NCNN=force to bypass this check.
|
||||
"""
|
||||
if os.environ.get("REALESRGAN_NCNN", "").lower() == "force":
|
||||
return True
|
||||
|
||||
plat = sys.platform.lower()
|
||||
|
||||
if plat == "linux":
|
||||
# DRI render nodes exist when a GPU is present and drivers loaded
|
||||
dri = Path("/dev/dri")
|
||||
if dri.exists() and list(dri.glob("renderD*")):
|
||||
return True
|
||||
# Fallback: vulkaninfo (not always installed)
|
||||
if shutil.which("vulkaninfo"):
|
||||
r = subprocess.run(["vulkaninfo", "--summary"],
|
||||
capture_output=True, timeout=5)
|
||||
if r.returncode == 0 and b"GPU" in r.stdout:
|
||||
return True
|
||||
return False
|
||||
|
||||
if plat == "darwin":
|
||||
# macOS with Metal/MPS — Vulkan via MoltenVK always present on Apple Silicon/modern Intel
|
||||
return True
|
||||
|
||||
if "win" in plat:
|
||||
# Windows always has a display adapter; assume Vulkan available
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _test_ncnn_binary(binary_path: Path) -> bool:
|
||||
"""Run binary with --help to confirm it actually works (Vulkan loads ok)."""
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[str(binary_path), "--help"],
|
||||
capture_output=True, timeout=15,
|
||||
)
|
||||
# NCNN binary exits 255 for --help but prints usage; that's fine.
|
||||
# A Vulkan init failure produces "no vulkan device" on stderr.
|
||||
stderr = r.stderr.decode(errors="replace").lower()
|
||||
if "no vulkan" in stderr or "failed to create" in stderr:
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def ensure_ncnn_installed() -> Optional[Path]:
|
||||
"""
|
||||
Check for Vulkan, then download+install the NCNN binary if needed.
|
||||
Skips silently on headless/CPU-only machines.
|
||||
Returns binary Path on success, None otherwise.
|
||||
"""
|
||||
global _install_status
|
||||
|
||||
binary_path = NCNN_DEST_DIR / _ncnn_binary_name()
|
||||
|
||||
# Already installed — quick verify it still works
|
||||
if binary_path.exists() and os.access(binary_path, os.X_OK):
|
||||
loop = asyncio.get_event_loop()
|
||||
ok = await loop.run_in_executor(None, _test_ncnn_binary, binary_path)
|
||||
if ok:
|
||||
_install_status = InstallStatus(state=InstallState.done, progress=100,
|
||||
message="Already installed.")
|
||||
return binary_path
|
||||
else:
|
||||
# Binary exists but Vulkan broken — treat as headless
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.skipped,
|
||||
message="Vulkan unavailable — skipping NCNN (using PyTorch CPU or Lanczos).",
|
||||
)
|
||||
return None
|
||||
|
||||
async with _install_lock:
|
||||
# Re-check after lock
|
||||
if binary_path.exists() and os.access(binary_path, os.X_OK):
|
||||
_install_status = InstallStatus(state=InstallState.done, progress=100,
|
||||
message="Already installed.")
|
||||
return binary_path
|
||||
|
||||
if _install_status.state in (InstallState.downloading, InstallState.extracting,
|
||||
InstallState.verifying):
|
||||
return None # already running
|
||||
|
||||
# Check Vulkan before downloading anything
|
||||
loop = asyncio.get_event_loop()
|
||||
has_vulkan = await loop.run_in_executor(None, _vulkan_available)
|
||||
if not has_vulkan:
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.skipped,
|
||||
message="No Vulkan GPU detected — skipping NCNN install. "
|
||||
"AI upscaling via PyTorch CPU or set REALESRGAN_NCNN=force to override.",
|
||||
)
|
||||
print("[upscale] Headless/no-Vulkan detected — skipping NCNN download.")
|
||||
return None
|
||||
|
||||
plat = sys.platform.lower()
|
||||
zip_name = _PLATFORM_ZIP.get(plat)
|
||||
if not zip_name:
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.failed,
|
||||
error=f"Unsupported platform: {plat}",
|
||||
)
|
||||
return None
|
||||
|
||||
url = f"{NCNN_BASE_URL}/{zip_name}"
|
||||
|
||||
try:
|
||||
NCNN_DEST_DIR.mkdir(parents=True, exist_ok=True)
|
||||
zip_path = NCNN_DEST_DIR / zip_name
|
||||
|
||||
# Download
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.downloading, progress=0,
|
||||
message=f"Downloading Real-ESRGAN NCNN {NCNN_VERSION}…",
|
||||
)
|
||||
|
||||
def _do_download():
|
||||
def _progress(count, block, total):
|
||||
if total > 0:
|
||||
_install_status.progress = min(85, int(count * block * 85 / total))
|
||||
urllib.request.urlretrieve(url, zip_path, _progress)
|
||||
|
||||
await loop.run_in_executor(None, _do_download)
|
||||
|
||||
# Extract
|
||||
_install_status.state = InstallState.extracting
|
||||
_install_status.progress = 88
|
||||
_install_status.message = "Extracting…"
|
||||
|
||||
def _do_extract():
|
||||
with zipfile.ZipFile(zip_path, "r") as zf:
|
||||
zf.extractall(NCNN_DEST_DIR)
|
||||
found = list(NCNN_DEST_DIR.rglob(_ncnn_binary_name()))
|
||||
if not found:
|
||||
raise FileNotFoundError(f"Binary not found after extract: {_ncnn_binary_name()}")
|
||||
extracted = found[0]
|
||||
if extracted != binary_path:
|
||||
extracted.rename(binary_path)
|
||||
if "win" not in sys.platform.lower():
|
||||
binary_path.chmod(
|
||||
binary_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH
|
||||
)
|
||||
zip_path.unlink(missing_ok=True)
|
||||
|
||||
await loop.run_in_executor(None, _do_extract)
|
||||
|
||||
# Verify binary actually works
|
||||
_install_status.state = InstallState.verifying
|
||||
_install_status.progress = 95
|
||||
_install_status.message = "Verifying Vulkan…"
|
||||
|
||||
ok = await loop.run_in_executor(None, _test_ncnn_binary, binary_path)
|
||||
if not ok:
|
||||
binary_path.unlink(missing_ok=True)
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.skipped,
|
||||
message="Binary installed but Vulkan unavailable at runtime — "
|
||||
"falling back to PyTorch CPU / Lanczos.",
|
||||
)
|
||||
print("[upscale] NCNN binary installed but Vulkan check failed — skipping.")
|
||||
return None
|
||||
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.done, progress=100,
|
||||
message=f"Real-ESRGAN NCNN installed: {binary_path}",
|
||||
)
|
||||
invalidate_caps_cache()
|
||||
return binary_path
|
||||
|
||||
except Exception as exc:
|
||||
_install_status = InstallStatus(
|
||||
state=InstallState.failed,
|
||||
error=str(exc),
|
||||
message="Installation failed.",
|
||||
)
|
||||
print(f"[upscale] NCNN auto-install failed: {exc}")
|
||||
return None
|
||||
|
||||
|
||||
# ── Capability detection ──────────────────────────────────────────────────────
|
||||
|
||||
_caps: Optional[dict] = None
|
||||
|
||||
|
||||
def probe_upscale_capabilities() -> dict:
|
||||
"""Detect available upscaling methods. Cached after first call."""
|
||||
global _caps
|
||||
if _caps is not None:
|
||||
return _caps
|
||||
|
||||
caps = {
|
||||
"lanczos": True,
|
||||
"realesrgan_pytorch": False,
|
||||
"realesrgan_pytorch_device": None,
|
||||
"realesrgan_ncnn": False,
|
||||
"realesrgan_ncnn_path": None,
|
||||
"recommended": "lanczos",
|
||||
"recommended_label": "Lanczos (no AI upscaler found)",
|
||||
"methods": ["lanczos"],
|
||||
"ncnn_install_status": get_install_status(),
|
||||
}
|
||||
|
||||
# ── PyTorch path ──────────────────────────────────────────────────────────
|
||||
pytorch_device = None
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
pytorch_device = "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
pytorch_device = "mps"
|
||||
else:
|
||||
pytorch_device = "cpu"
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if pytorch_device:
|
||||
try:
|
||||
import realesrgan # noqa: F401
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet # noqa: F401
|
||||
caps["realesrgan_pytorch"] = True
|
||||
caps["realesrgan_pytorch_device"] = pytorch_device
|
||||
caps["methods"].append("realesrgan_pytorch")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# ── NCNN Vulkan binary ────────────────────────────────────────────────────
|
||||
ncnn_path = _find_ncnn_binary()
|
||||
if ncnn_path:
|
||||
caps["realesrgan_ncnn"] = True
|
||||
caps["realesrgan_ncnn_path"] = str(ncnn_path)
|
||||
caps["methods"].append("realesrgan_ncnn")
|
||||
|
||||
# ── Pick recommended ──────────────────────────────────────────────────────
|
||||
if caps["realesrgan_pytorch"] and pytorch_device in ("cuda", "mps"):
|
||||
device_label = "CUDA GPU" if pytorch_device == "cuda" else "Apple Silicon"
|
||||
caps["recommended"] = "realesrgan_pytorch"
|
||||
caps["recommended_label"] = f"Real-ESRGAN ({device_label})"
|
||||
elif caps["realesrgan_ncnn"]:
|
||||
caps["recommended"] = "realesrgan_ncnn"
|
||||
caps["recommended_label"] = "Real-ESRGAN NCNN (Vulkan)"
|
||||
elif caps["realesrgan_pytorch"] and pytorch_device == "cpu":
|
||||
caps["recommended"] = "realesrgan_pytorch"
|
||||
caps["recommended_label"] = "Real-ESRGAN (CPU — may be slow)"
|
||||
else:
|
||||
install_state = _install_status.state
|
||||
if install_state in (InstallState.downloading, InstallState.extracting, InstallState.verifying):
|
||||
caps["recommended_label"] = "Lanczos (AI upscaler installing…)"
|
||||
elif install_state == InstallState.skipped:
|
||||
caps["recommended_label"] = "Lanczos (headless — no Vulkan GPU)"
|
||||
else:
|
||||
caps["recommended_label"] = "Lanczos (no AI upscaler found)"
|
||||
|
||||
_caps = caps
|
||||
return caps
|
||||
|
||||
|
||||
def _find_ncnn_binary() -> Optional[Path]:
|
||||
found = shutil.which("realesrgan-ncnn-vulkan")
|
||||
if found:
|
||||
return Path(found)
|
||||
candidates = [
|
||||
NCNN_DEST_DIR / _ncnn_binary_name(),
|
||||
Path("/usr/local/bin/realesrgan-ncnn-vulkan"),
|
||||
Path.home() / ".local/bin/realesrgan-ncnn-vulkan",
|
||||
Path(r"C:/realesrgan-ncnn-vulkan/realesrgan-ncnn-vulkan.exe"),
|
||||
Path("/opt/homebrew/bin/realesrgan-ncnn-vulkan"),
|
||||
]
|
||||
for p in candidates:
|
||||
if p.exists() and os.access(p, os.X_OK):
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def invalidate_caps_cache():
|
||||
global _caps
|
||||
_caps = None
|
||||
|
||||
|
||||
# ── Upscale implementations ───────────────────────────────────────────────────
|
||||
|
||||
def _to_png_bytes(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def upscale_lanczos(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
||||
new_w = round(image.width * scale)
|
||||
new_h = round(image.height * scale)
|
||||
result = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
||||
return _to_png_bytes(result), "lanczos"
|
||||
|
||||
|
||||
def upscale_realesrgan_pytorch(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
caps = probe_upscale_capabilities()
|
||||
device = caps.get("realesrgan_pytorch_device", "cpu")
|
||||
|
||||
model_scale = 2 if scale <= 2.5 else 4
|
||||
model = RRDBNet(
|
||||
num_in_ch=3, num_out_ch=3, num_feat=64,
|
||||
num_block=23, num_grow_ch=32, scale=model_scale
|
||||
)
|
||||
|
||||
model_dir = Path("/app/data/models/realesrgan")
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
model_path = model_dir / f"RealESRGAN_x{model_scale}plus.pth"
|
||||
if not model_path.exists():
|
||||
model_path = None
|
||||
|
||||
upsampler = RealESRGANer(
|
||||
scale=model_scale,
|
||||
model_path=str(model_path) if model_path else None,
|
||||
model=model,
|
||||
tile=512,
|
||||
tile_pad=10,
|
||||
pre_pad=0,
|
||||
half=(device == "cuda"),
|
||||
device=torch.device(device),
|
||||
)
|
||||
|
||||
import numpy as np
|
||||
img_bgr = np.array(image)[:, :, ::-1].copy()
|
||||
enhanced, _ = upsampler.enhance(img_bgr, outscale=scale)
|
||||
result = Image.fromarray(enhanced[:, :, ::-1])
|
||||
return _to_png_bytes(result), f"realesrgan_pytorch_{device}"
|
||||
|
||||
|
||||
def upscale_realesrgan_ncnn(image: Image.Image, scale: float) -> tuple[bytes, str]:
|
||||
caps = probe_upscale_capabilities()
|
||||
binary = caps.get("realesrgan_ncnn_path")
|
||||
if not binary:
|
||||
raise RuntimeError("realesrgan-ncnn-vulkan binary not found")
|
||||
|
||||
model_scale = 4 if scale > 2.5 else 2
|
||||
target_w = round(image.width * scale)
|
||||
target_h = round(image.height * scale)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
in_path = Path(tmpdir) / "input.png"
|
||||
out_path = Path(tmpdir) / "output.png"
|
||||
image.save(in_path, format="PNG")
|
||||
|
||||
cmd = [
|
||||
binary,
|
||||
"-i", str(in_path), "-o", str(out_path),
|
||||
"-s", str(model_scale), "-n", f"realesrgan-x{model_scale}plus", "-f", "png",
|
||||
]
|
||||
r = subprocess.run(cmd, capture_output=True, timeout=300)
|
||||
if r.returncode != 0:
|
||||
raise RuntimeError(f"realesrgan-ncnn-vulkan failed: {r.stderr.decode()}")
|
||||
|
||||
result = Image.open(out_path).convert("RGB")
|
||||
if result.width != target_w or result.height != target_h:
|
||||
result = result.resize((target_w, target_h), Image.Resampling.LANCZOS)
|
||||
|
||||
return _to_png_bytes(result), "realesrgan_ncnn"
|
||||
|
||||
|
||||
# ── Public entry point ────────────────────────────────────────────────────────
|
||||
|
||||
def upscale_sync(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]:
|
||||
"""Upscale synchronously. Returns (png_bytes, method_label)."""
|
||||
caps = probe_upscale_capabilities()
|
||||
|
||||
if method == "auto":
|
||||
method = caps["recommended"]
|
||||
|
||||
if method == "realesrgan_pytorch":
|
||||
if caps["realesrgan_pytorch"]:
|
||||
try:
|
||||
return upscale_realesrgan_pytorch(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN PyTorch failed, falling back: {e}")
|
||||
if caps["realesrgan_ncnn"]:
|
||||
try:
|
||||
return upscale_realesrgan_ncnn(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN NCNN fallback failed: {e}")
|
||||
return upscale_lanczos(image, scale)
|
||||
|
||||
if method == "realesrgan_ncnn":
|
||||
if caps["realesrgan_ncnn"]:
|
||||
try:
|
||||
return upscale_realesrgan_ncnn(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN NCNN failed, falling back: {e}")
|
||||
if caps["realesrgan_pytorch"]:
|
||||
try:
|
||||
return upscale_realesrgan_pytorch(image, scale)
|
||||
except Exception as e:
|
||||
print(f"Real-ESRGAN PyTorch fallback failed: {e}")
|
||||
return upscale_lanczos(image, scale)
|
||||
|
||||
return upscale_lanczos(image, scale)
|
||||
|
||||
|
||||
async def upscale_image(image: Image.Image, scale: float, method: str = "auto") -> tuple[bytes, str]:
|
||||
"""Async wrapper — runs upscale in thread pool."""
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, upscale_sync, image, scale, method)
|
||||
Reference in New Issue
Block a user