Move ai-stack and paintplus vendored source under vendor/
Matches the existing vendor/easy-asterisk convention (used by services/asterisk.sh) instead of two one-off top-level directories that cluttered the repo root and didn't look like anything else next to setup.sh, lib/, services/, extras/. Only the two services' own SRC_DIR path resolution and header comments needed updating — nothing else in the repo referenced the old ./ai-stack / ./paintplus paths. Also documents vendor/ in README.md's Layout section.
This commit is contained in:
Claude
@@ -0,0 +1,23 @@
|
||||
# Generated by setup scripts — not tracked
|
||||
docker-compose.yml
|
||||
.env
|
||||
requirements.txt
|
||||
mcp_requirements.txt
|
||||
start.sh
|
||||
stop.sh
|
||||
status.sh
|
||||
pull-models.sh
|
||||
Caddyfile.example
|
||||
.ufw-done
|
||||
|
||||
# Data directories
|
||||
workspace/
|
||||
repos/
|
||||
papers/
|
||||
index/
|
||||
logs/
|
||||
kiwix/
|
||||
gitea/
|
||||
portainer-data/
|
||||
invokeai-data/
|
||||
invokeai-outputs/
|
||||
Vendored
+988
@@ -0,0 +1,988 @@
|
||||
# Local AI Stack
|
||||
|
||||
A fully offline, self-hosted AI environment for Ubuntu 24.04. Runs on any NVIDIA GPU (or CPU-only).
|
||||
|
||||
**Services:** Ollama · Open WebUI · RAG · MCP · ChromaDB · SearXNG · Kiwix · Gitea · InvokeAI · ComfyUI · Portainer
|
||||
|
||||
> **Vendored into [ubuntu-post-install](https://github.com/outis1one/ubuntu-post-install) as the `ai-stack` service.**
|
||||
> `services/ai-stack.sh` copies this source to `~/docker/ai-stack/` and hands off to `local-ai-setup.sh`
|
||||
> (the VRAM-aware installer below). The upstream `ubuntu-post-install.sh` and `ubuntu-post-install-main.zip`
|
||||
> — old standalone copies of that project, "independent of the AI stack" per this README — are omitted from
|
||||
> the vendored copy. Source: github.com/outis1one/local-ai
|
||||
|
||||
---
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
git clone <this-repo>
|
||||
cd local-ai
|
||||
./laptop_full_setup.sh
|
||||
```
|
||||
|
||||
That's it. The script installs Docker, NVIDIA drivers (if needed), generates all config, starts the stack, and optionally pulls models.
|
||||
|
||||
---
|
||||
|
||||
## Service URLs
|
||||
|
||||
After setup, all services are available on your LAN:
|
||||
|
||||
| Service | URL | Purpose |
|
||||
|------------|----------------------------|--------------------------------|
|
||||
| Open WebUI | `http://<ip>:3000` | Chat interface (Ollama + RAG) |
|
||||
| InvokeAI | `http://<ip>:9090` | Image generation (standalone) |
|
||||
| ComfyUI | `http://<ip>:8188` | Image generation (OWUI integration) |
|
||||
| SearXNG | `http://<ip>:8888` | Private web search |
|
||||
| Kiwix | `http://<ip>:8181` | Offline Wikipedia / docs |
|
||||
| Gitea | `http://<ip>:3001` | Self-hosted Git |
|
||||
| RAG Health | `http://<ip>:8001/health` | RAG server status |
|
||||
| MCP SSE | `http://<ip>:8002/sse` | MCP endpoint for Claude Code |
|
||||
| Portainer | `https://<ip>:9443` | Docker management UI |
|
||||
|
||||
---
|
||||
|
||||
## Day-to-Day Commands
|
||||
|
||||
All generated into `~/docker/ai-stack/` by the setup script:
|
||||
|
||||
```bash
|
||||
bash ~/docker/ai-stack/start.sh # pull latest images + docker compose up -d
|
||||
bash ~/docker/ai-stack/stop.sh # docker compose down
|
||||
bash ~/docker/ai-stack/status.sh # GPU / container / RAG health
|
||||
bash ~/docker/ai-stack/pull-models.sh # pull Ollama models (run once after first install)
|
||||
```
|
||||
|
||||
The stack also registers as a **systemd service** that starts on boot:
|
||||
|
||||
```bash
|
||||
sudo systemctl start local-ai
|
||||
sudo systemctl stop local-ai
|
||||
sudo systemctl status local-ai
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Script Reference
|
||||
|
||||
| Script | Lines | What it does |
|
||||
|---|---|---|
|
||||
| `laptop_full_setup.sh` | 620 | **Main setup.** Installs Docker + NVIDIA toolkit, creates `~/docker/ai-stack/`, writes `docker-compose.yml`, starts stack, registers systemd service. |
|
||||
| `local-ai-setup.sh` | 837 | Alternative setup script. Same as above but also **auto-detects VRAM** and selects models accordingly (14B for ≥14GB VRAM, 7B for CPU). Use this instead of `laptop_full_setup.sh` if you want VRAM-aware model selection. |
|
||||
| `ubuntu-post-install.sh` | 8,889 | Full Ubuntu 24.04 post-install (dev tools, fonts, apps, tweaks). Run once on a fresh OS install. Independent of the AI stack. |
|
||||
| `configure-storage.sh` | 239 | Storage/mount configuration helper. Run separately if you have a secondary drive for AI data. |
|
||||
| `kiwix_download.sh` | 198 | Downloads ZIM files (Wikipedia, Stack Overflow, etc.) for offline use. Run separately — files are large. |
|
||||
| `invokeai-import-lora.sh` | 85 | Copies a LoRA `.safetensors` file into InvokeAI's Docker model volume. |
|
||||
| `comfyui-import-lora.sh` | 97 | Copies a LoRA into ComfyUI and prints workflow setup instructions. |
|
||||
| `comfyui-install-ipadapter.sh` | 185 | Installs IP-Adapter nodes + models into ComfyUI for reference-image workflows (same face, different settings). |
|
||||
| `setup-image-models.sh` | 200 | **GPU-aware** image model installer. Detects VRAM, offers appropriate SD/SDXL/Flux models, installs into InvokeAI and/or ComfyUI. |
|
||||
|
||||
### Which setup script should I use?
|
||||
|
||||
- `laptop_full_setup.sh` — fixed model selection (`qwen2.5:14b` / `qwen2.5-coder:7b`), simpler
|
||||
- `local-ai-setup.sh` — detects your VRAM at runtime and picks appropriate models, also embeds `server.py` and `mcp_server.py` directly (doesn't need repo files copied separately)
|
||||
|
||||
Both scripts are **idempotent** — safe to re-run for updates. Config files are kept on re-run unless you pass `--force`.
|
||||
|
||||
---
|
||||
|
||||
## Generated File Layout
|
||||
|
||||
```
|
||||
~/docker/ai-stack/
|
||||
├── docker-compose.yml # generated by setup script
|
||||
├── .env # API tokens — edit this, never overwritten
|
||||
├── server.py # RAG server (copied from repo)
|
||||
├── mcp_server.py # MCP server (copied from repo)
|
||||
├── requirements.txt # RAG Python deps
|
||||
├── mcp_requirements.txt # MCP Python deps
|
||||
├── start.sh # start the stack
|
||||
├── stop.sh # stop the stack
|
||||
├── status.sh # GPU + container + RAG health
|
||||
├── pull-models.sh # pull Ollama models
|
||||
├── Caddyfile.example # reverse proxy config template
|
||||
├── papers/ # drop PDFs here for RAG indexing
|
||||
├── repos/ # git repos indexed by RAG
|
||||
├── workspace/ # MCP working directory
|
||||
├── index/ # ChromaDB vector store (persistent)
|
||||
├── kiwix/ # ZIM files for Kiwix
|
||||
├── gitea/ # Gitea data
|
||||
├── invokeai-outputs/ # InvokeAI generated images
|
||||
├── comfyui-output/ # ComfyUI generated images
|
||||
├── comfyui-data/ # ComfyUI custom nodes
|
||||
└── logs/
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## First Run Checklist
|
||||
|
||||
1. **Run setup:**
|
||||
```bash
|
||||
./laptop_full_setup.sh
|
||||
```
|
||||
2. **Pull models** (prompted at end of setup, or run manually):
|
||||
```bash
|
||||
bash ~/docker/ai-stack/pull-models.sh
|
||||
```
|
||||
Downloads ~15-30GB. Takes 10-40 min depending on connection.
|
||||
|
||||
3. **Add API tokens** (optional — for Gitea/GitHub MCP tools):
|
||||
```bash
|
||||
nano ~/docker/ai-stack/.env
|
||||
```
|
||||
|
||||
4. **Connect Claude Code to MCP:**
|
||||
```bash
|
||||
claude mcp add local http://<your-ip>:8002/sse
|
||||
```
|
||||
|
||||
5. **Download ZIMs** for offline docs (optional, large):
|
||||
```bash
|
||||
./kiwix_download.sh
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Querying Your Codebase from Open WebUI
|
||||
|
||||
The stack includes a **code-aware RAG server** that sits between Open WebUI and Ollama. When you chat in Open WebUI, the RAG server automatically retrieves relevant code from your indexed repos and injects it into the prompt — so the model answers with your actual code as context.
|
||||
|
||||
This is NOT the same as Open WebUI's built-in Knowledge Collections. This is a separate, always-on layer that understands code structure.
|
||||
|
||||
### How it works (architecture)
|
||||
|
||||
```
|
||||
You type a question in Open WebUI
|
||||
↓
|
||||
Open WebUI → RAG Server (port 8001) /v1/chat/completions
|
||||
↓
|
||||
RAG Server: searches ChromaDB for relevant code chunks
|
||||
(AST-parsed Python functions, regex-split JS/Go/Rust, etc.)
|
||||
↓
|
||||
RAG Server: prepends code snippets to your prompt:
|
||||
"### auth.py:authenticate
|
||||
def authenticate(user, password): ..."
|
||||
↓
|
||||
RAG Server → Ollama: generates response WITH your code as context
|
||||
↓
|
||||
Response sent back to Open WebUI
|
||||
```
|
||||
|
||||
Open WebUI is already configured to route through the RAG server via:
|
||||
```
|
||||
OPENAI_API_BASE_URL=http://rag-server:8001/v1
|
||||
```
|
||||
|
||||
### Step 1: Index your code
|
||||
|
||||
**Option A: Drop repos in the repos directory**
|
||||
```bash
|
||||
cd ~/docker/ai-stack/repos
|
||||
git clone http://localhost:3001/your-user/your-repo.git
|
||||
# Restart RAG server to trigger indexing:
|
||||
docker restart rag-server
|
||||
```
|
||||
|
||||
**Option B: Use the ingest API**
|
||||
```bash
|
||||
# From Gitea:
|
||||
curl -X POST http://localhost:8001/ingest/repo \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"url": "http://gitea:3000/user/repo", "name": "my-repo"}'
|
||||
|
||||
# From GitHub:
|
||||
curl -X POST http://localhost:8001/ingest/repo \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"url": "https://github.com/user/repo", "name": "my-repo", "branch": "main"}'
|
||||
```
|
||||
|
||||
**Option C: Auto-index on push (Gitea webhook)**
|
||||
1. In Gitea → your repo → **Settings** → **Webhooks** → **Add Webhook** → **Gitea**
|
||||
2. Target URL: `http://rag-server:8001/webhook/gitea`
|
||||
3. Trigger: Push events
|
||||
4. Now every `git push` to Gitea auto-reindexes that repo
|
||||
|
||||
### Step 2: Chat about your code
|
||||
|
||||
Just ask questions in Open WebUI. The RAG server automatically retrieves relevant code:
|
||||
|
||||
- "What does the authenticate function do?"
|
||||
- "How is the database connection configured?"
|
||||
- "Show me all the API endpoints"
|
||||
- "What tests exist for the user model?"
|
||||
|
||||
The model sees the actual code snippets and responds based on them — not hallucinating.
|
||||
|
||||
### What gets indexed
|
||||
|
||||
**Supported file types:**
|
||||
`.py`, `.js`, `.ts`, `.tsx`, `.jsx`, `.go`, `.rs`, `.java`, `.c`, `.cpp`, `.h`, `.hpp`, `.cs`, `.rb`, `.sh`, `.yaml`, `.yml`, `.toml`, `.sql`, `.md`
|
||||
|
||||
**Smart chunking:**
|
||||
- **Python:** AST-parsed — each function and class is its own searchable chunk
|
||||
- **Other languages:** Regex-split on `function`, `class`, `const`, `func`, `impl`, etc.
|
||||
- **Fallback:** Sliding window (1200 chars, 200 char overlap)
|
||||
|
||||
**Skipped directories:**
|
||||
`node_modules`, `.git`, `__pycache__`, `dist`, `build`, `.venv`, `venv`, `vendor`, `target`, `bin`, `obj`
|
||||
|
||||
### Checking RAG status
|
||||
|
||||
```bash
|
||||
curl http://localhost:8001/health
|
||||
```
|
||||
|
||||
Returns document counts per collection (code, papers) and overall status.
|
||||
|
||||
### RAG server vs Knowledge Collections vs Memories
|
||||
|
||||
These are **three separate systems** in the stack. Understanding the difference matters:
|
||||
|
||||
| | **RAG Server** | **Knowledge Collections** | **Memories** |
|
||||
|---|---|---|---|
|
||||
| **What** | Code-aware retrieval layer | Open WebUI's built-in document RAG | Short facts about the user |
|
||||
| **Content** | Git repos (auto-indexed) | PDFs, text files you upload | Extracted from conversations |
|
||||
| **Chunking** | AST/regex (code-aware) | Generic text chunking | Single sentences |
|
||||
| **Activation** | Always on (every chat) | Per-chat (`#` tag) | Global (every chat) |
|
||||
| **Best for** | "What does this function do?" | Project docs, research notes | "Remember I prefer Python" |
|
||||
| **Context cost** | ~1500-2500 tokens (top 6 chunks) | ~1500-2500 tokens (top K chunks) | ~200 tokens |
|
||||
|
||||
**For code questions:** The RAG server handles this automatically — no setup needed beyond indexing your repos.
|
||||
|
||||
**For project docs/notes:** Use Knowledge Collections — type `#` to scope per-chat.
|
||||
|
||||
**For personal preferences:** Use Memories (sparingly — they're global).
|
||||
|
||||
---
|
||||
|
||||
## Image Generation from Open WebUI
|
||||
|
||||
Open WebUI can generate images inline in chat conversations using ComfyUI as the backend. When configured, you can ask any model to "generate an image of..." and it will call ComfyUI to create the image.
|
||||
|
||||
### How it works
|
||||
|
||||
Open WebUI natively supports these image generation engines:
|
||||
- **ComfyUI** — Node-based, best Open WebUI integration, local
|
||||
- **AUTOMATIC1111** — Stable Diffusion WebUI, local
|
||||
- **OpenAI DALL-E** — Cloud API
|
||||
- **Gemini** — Cloud API
|
||||
|
||||
**InvokeAI** does NOT have a compatible API for Open WebUI integration. It works great as a standalone tool at `http://<ip>:9090` but cannot be called from within Open WebUI chats. For chat-integrated image generation, use **ComfyUI**.
|
||||
|
||||
### Setup: ComfyUI + Open WebUI (recommended)
|
||||
|
||||
If you selected ComfyUI during setup, the environment variables are already configured. You just need to install a model and set up a workflow.
|
||||
|
||||
#### Step 1: Install a Stable Diffusion model in ComfyUI
|
||||
|
||||
```bash
|
||||
# Open ComfyUI at http://<ip>:8188
|
||||
# Use the built-in Model Manager to download a model, or manually:
|
||||
docker exec comfyui bash -c "cd /opt/ComfyUI/models/checkpoints && \
|
||||
wget -q 'https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors'"
|
||||
```
|
||||
|
||||
Or download any `.safetensors` checkpoint and copy it in:
|
||||
```bash
|
||||
docker cp ~/Downloads/my-model.safetensors comfyui:/opt/ComfyUI/models/checkpoints/
|
||||
```
|
||||
|
||||
#### Step 2: Create and export a workflow
|
||||
|
||||
1. Open ComfyUI at `http://<ip>:8188`
|
||||
2. Build or load a workflow (the default text-to-image workflow works)
|
||||
3. Click the **gear icon** → enable **Dev Mode**
|
||||
4. Click **Save (API Format)** — this downloads `workflow_api.json`
|
||||
|
||||
#### Step 3: Configure Open WebUI
|
||||
|
||||
1. Open WebUI → **Admin** → **Settings** → **Images**
|
||||
2. Set **Engine** to `ComfyUI`
|
||||
3. Set **URL** to `http://comfyui:8188` (container networking, already set via env vars)
|
||||
4. Click **Import Workflow** and upload your `workflow_api.json`
|
||||
5. Map the **prompt node** (usually the KSampler or CLIPTextEncode node)
|
||||
6. Save settings
|
||||
|
||||
#### Step 4: Generate images in chat
|
||||
|
||||
In any Open WebUI chat, type something like:
|
||||
- "Generate an image of a mountain landscape at sunset"
|
||||
- "Create a photo of a cyberpunk city"
|
||||
|
||||
The model will detect the image generation request and pass it to ComfyUI.
|
||||
|
||||
### Using LoRAs with Open WebUI (workflow-per-style)
|
||||
|
||||
LoRAs let you apply trained styles (anime, photorealistic, specific characters, etc.) to image generation. The trick: **bake the LoRA into a ComfyUI workflow, export it, and import into Open WebUI.** After that, you iterate from chat without touching ComfyUI.
|
||||
|
||||
#### Step 1: Import your LoRA file
|
||||
|
||||
```bash
|
||||
./comfyui-import-lora.sh ~/Downloads/my-style-lora.safetensors "Anime Style"
|
||||
```
|
||||
|
||||
Or manually:
|
||||
```bash
|
||||
docker cp ~/Downloads/my-style-lora.safetensors comfyui:/opt/ComfyUI/models/loras/
|
||||
```
|
||||
|
||||
#### Step 2: Build a workflow with the LoRA (one-time)
|
||||
|
||||
1. Open ComfyUI at `http://<ip>:8188`
|
||||
2. Load the default text-to-image workflow
|
||||
3. Add a **Load LoRA** node: right-click → Add Node → loaders → Load LoRA
|
||||
4. Wire it between the checkpoint and the rest of the pipeline:
|
||||
```
|
||||
[Load Checkpoint] → MODEL → [Load LoRA] → MODEL → [KSampler]
|
||||
→ CLIP → → CLIP → [CLIP Text Encode]
|
||||
```
|
||||
5. Select your LoRA file, set `strength_model` and `strength_clip` (start 0.7–0.85)
|
||||
6. Test it — click **Queue Prompt** and verify it works
|
||||
7. Enable **Dev Mode** (gear icon) → click **Save (API Format)**
|
||||
|
||||
#### Step 3: Import into Open WebUI
|
||||
|
||||
1. Open WebUI → **Admin** → **Settings** → **Images**
|
||||
2. Click **Import Workflow** → upload the `workflow_api.json`
|
||||
3. Map the prompt node (usually CLIPTextEncode)
|
||||
4. Save
|
||||
|
||||
Now every "Generate an image of..." in chat uses your LoRA automatically.
|
||||
|
||||
#### Combining multiple LoRAs in one image
|
||||
|
||||
You can chain multiple Load LoRA nodes to blend styles:
|
||||
|
||||
```
|
||||
[Load Checkpoint] → [Load LoRA 1] → [Load LoRA 2] → [KSampler]
|
||||
(anime, 0.7) (lighting, 0.5)
|
||||
```
|
||||
|
||||
Each LoRA has its own `strength_model` / `strength_clip` sliders. Export the combined workflow as usual — both LoRAs are baked in and always applied.
|
||||
|
||||
#### Managing multiple LoRA styles
|
||||
|
||||
Each LoRA combo needs its own exported workflow. Practical approach:
|
||||
|
||||
| Workflow file | LoRAs | Use case |
|
||||
|---------------|-------|----------|
|
||||
| `anime-only.json` | Anime @ 0.8 | Anime style |
|
||||
| `anime-cinematic.json` | Anime @ 0.7 + Cinematic @ 0.5 | Anime + dramatic lighting |
|
||||
| `photorealistic.json` | Photorealistic @ 0.9 | Photo look |
|
||||
| `base-model.json` | None | No LoRA, base checkpoint only |
|
||||
|
||||
Switch between them in **Admin → Settings → Images → Import Workflow**. The active workflow applies to all image generation requests.
|
||||
|
||||
> **Important:** There are no keywords or trigger words to activate LoRAs from chat. LoRAs are hardwired in the workflow JSON — whatever prompt you type, the LoRA(s) always apply. To change styles, swap the workflow in OWUI settings.
|
||||
|
||||
#### LoRA compatibility
|
||||
|
||||
| LoRA trained on | Must use checkpoint |
|
||||
|----------------|-------------------|
|
||||
| SD 1.5 | Any SD 1.5 model (e.g. `v1-5-pruned-emaonly.safetensors`) |
|
||||
| SDXL | Any SDXL model (e.g. `sd_xl_base_1.0.safetensors`) |
|
||||
| Flux | Flux checkpoint |
|
||||
|
||||
Mismatched architectures will produce errors or garbage output.
|
||||
|
||||
### IP-Adapter: same face, different settings (reference image workflows)
|
||||
|
||||
IP-Adapter lets you upload a reference photo and generate new images that preserve the face/subject while changing everything else — age, emotion, setting, art style. This is different from LoRAs (which bake in a trained style). IP-Adapter works from a single image with no training needed.
|
||||
|
||||
#### What you can do
|
||||
|
||||
| Task | How | Example prompt |
|
||||
|------|-----|----------------|
|
||||
| Different setting | Reference face + new scene | "same person, sitting in a cafe in Paris" |
|
||||
| Age up/down | Reference face + age prompt | "same person as an elderly man" or "same person as a child" |
|
||||
| Emotions | Reference face + emotion | "same person, laughing" or "same person, crying" |
|
||||
| Art style | Reference face + style | "same person, oil painting, renaissance style" |
|
||||
| Pose change | Reference face + pose | "same person, looking over shoulder, dramatic lighting" |
|
||||
|
||||
#### Quick install
|
||||
|
||||
```bash
|
||||
./comfyui-install-ipadapter.sh # SDXL models (~2.5GB)
|
||||
./comfyui-install-ipadapter.sh --faceid # + FaceID for stronger face lock
|
||||
./comfyui-install-ipadapter.sh --all # SDXL + SD 1.5 models
|
||||
```
|
||||
|
||||
Then restart ComfyUI: `docker restart comfyui`
|
||||
|
||||
#### Basic workflow in ComfyUI
|
||||
|
||||
1. Open `http://<ip>:8188`
|
||||
2. Add nodes: right-click → Add Node → ipadapter
|
||||
3. Wire up:
|
||||
```
|
||||
[Load Checkpoint] → MODEL → [IPAdapter Apply] → [KSampler] → [VAE Decode] → [Save Image]
|
||||
→ CLIP → [CLIP Text Encode] → positive →
|
||||
[Load Image (your photo)] → [IPAdapter Unified Loader] → ipadapter →
|
||||
```
|
||||
4. **IPAdapter Unified Loader**: set preset to `PLUS FACE (portrait)`
|
||||
5. **IPAdapter Apply**: set weight 0.7–1.0 (higher = more faithful to reference)
|
||||
6. Text prompt: describe the new scene/emotion/age
|
||||
7. Click **Queue Prompt**
|
||||
|
||||
#### Adjusting face fidelity
|
||||
|
||||
| Weight | Effect |
|
||||
|--------|--------|
|
||||
| 0.5–0.6 | Loose reference — inspired by the face but not a match |
|
||||
| 0.7–0.8 | Good balance — recognizable face, creative freedom in scene |
|
||||
| 0.9–1.0 | Strong lock — very close to reference face |
|
||||
| FaceID preset | Strongest — uses face detection for identity lock |
|
||||
|
||||
#### Important: ComfyUI only (not from Open WebUI chat)
|
||||
|
||||
IP-Adapter workflows require uploading a reference image to ComfyUI. Open WebUI's image generation integration only sends text prompts — it can't attach reference images. For this use case, work directly in ComfyUI at `http://<ip>:8188`.
|
||||
|
||||
For text-only image generation from Open WebUI chat, use [LoRA workflows](#using-loras-with-open-webui-workflow-per-style) instead.
|
||||
|
||||
### Setup: AUTOMATIC1111 (alternative)
|
||||
|
||||
If you prefer AUTOMATIC1111 over ComfyUI:
|
||||
|
||||
1. Run AUTOMATIC1111 with the `--api` flag
|
||||
2. In Open WebUI → **Admin** → **Settings** → **Images**:
|
||||
- Engine: `Automatic1111`
|
||||
- URL: `http://host.docker.internal:7860` (or container name if in Docker)
|
||||
3. Environment variables (alternative to UI config):
|
||||
```
|
||||
ENABLE_IMAGE_GENERATION=true
|
||||
IMAGE_GENERATION_ENGINE=automatic1111
|
||||
AUTOMATIC1111_BASE_URL=http://host.docker.internal:7860
|
||||
```
|
||||
|
||||
### Environment variables reference
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `ENABLE_IMAGE_GENERATION` | `false` | Enable image generation feature |
|
||||
| `IMAGE_GENERATION_ENGINE` | — | `comfyui`, `automatic1111`, `openai`, or `gemini` |
|
||||
| `IMAGE_GENERATION_MODEL` | — | Model ID for generation |
|
||||
| `IMAGE_SIZE` | `512x512` | Default output size |
|
||||
| `COMFYUI_BASE_URL` | — | ComfyUI API URL (e.g. `http://comfyui:8188`) |
|
||||
| `COMFYUI_API_KEY` | — | ComfyUI API key (if auth enabled) |
|
||||
| `COMFYUI_WORKFLOW` | — | Custom workflow JSON (API format) |
|
||||
| `AUTOMATIC1111_BASE_URL` | — | AUTOMATIC1111 API URL |
|
||||
| `AUTOMATIC1111_API_AUTH` | — | Auth credentials (`user:pass`) |
|
||||
|
||||
### Installing Functions & Actions (without community signup)
|
||||
|
||||
The Open WebUI community hub (`openwebui.com`) requires a free account to download functions. If you get parse errors pasting URLs or don't want to sign up, you can install functions manually by pasting their source code directly.
|
||||
|
||||
#### Manual installation (no account needed)
|
||||
|
||||
1. Find the function's source on GitHub — most are in the [open-webui/functions](https://github.com/open-webui/functions) repo or linked from community pages
|
||||
2. Copy the **raw Python source code** (the entire `.py` file)
|
||||
3. In Open WebUI → **Workspace** → **Functions** → click **＋** (Create)
|
||||
4. Paste the code into the editor
|
||||
5. Give it a name and save
|
||||
6. Enable it under **Workspace** → **Functions** — toggle it on globally or per-model
|
||||
|
||||
For **Actions** (like Generate Image), the process is the same but go to **Workspace** → **Functions** and set the function type to **Action** in the metadata.
|
||||
|
||||
#### Recommended functions
|
||||
|
||||
| Function | Type | What it does |
|
||||
|----------|------|-------------|
|
||||
| **Auto Memory** | Filter | Automatically extracts and stores facts from conversations as persistent memories |
|
||||
| **Generate Image** | Action | Adds a "Generate Image" button to messages for quick image generation |
|
||||
|
||||
### Auto Memory setup
|
||||
|
||||
Auto Memory is a filter function that runs after each message exchange, uses an LLM call to extract noteworthy facts, and stores them as user memories in Open WebUI's built-in memory system.
|
||||
|
||||
#### Step 1: Install the function
|
||||
|
||||
Use the manual method above. The source code is at:
|
||||
`https://github.com/open-webui/functions` — search for `auto_memory` or `adaptive_memory`
|
||||
|
||||
#### Step 2: Configure the function settings
|
||||
|
||||
After installing, click the **gear icon** on the function to configure:
|
||||
|
||||
| Setting | Value for local stack | Notes |
|
||||
|---------|----------------------|-------|
|
||||
| `openai_api_url` | `http://ollama:11434/v1` | Ollama's OpenAI-compatible endpoint |
|
||||
| `model` | Your chat model (e.g. `qwen2.5:14b`) | Used for memory extraction LLM calls |
|
||||
| `api_key` | `ollama` | Any non-empty string — Ollama ignores auth |
|
||||
| `context_window_n` | `4` (default) | Number of related memories injected per message |
|
||||
| `similarity_score_filter` | Leave default | Cosine similarity threshold for memory retrieval |
|
||||
| `messages_to_consider` | Leave default | How many recent messages to analyze |
|
||||
| `allow_modify_user_memories` | `false` | **Keep off** — lets LLM delete/edit memories via API if enabled |
|
||||
|
||||
#### Step 3: Enable memories per user
|
||||
|
||||
Each user must enable the memory feature individually:
|
||||
1. Click your **profile icon** → **Settings** → **Personalization**
|
||||
2. Toggle **Memory** to **ON**
|
||||
|
||||
Without this, the function installs but does nothing.
|
||||
|
||||
### Generate Image action setup
|
||||
|
||||
1. Install via manual method above (search for `generate_image` in the functions repo)
|
||||
2. Configure the function settings:
|
||||
- Uses your existing Open WebUI image generation settings (ComfyUI/A1111)
|
||||
- No additional API configuration needed if image generation already works
|
||||
3. After installing, a **"Generate Image"** button appears on AI messages
|
||||
4. Clicking it sends the message content as an image generation prompt
|
||||
|
||||
### Memory system: how it works and context impact
|
||||
|
||||
Open WebUI has two separate systems for persistent knowledge. Understanding the difference is critical for managing your context window.
|
||||
|
||||
#### Built-in Memories vs Knowledge Collections
|
||||
|
||||
| | **Memories** | **Knowledge** |
|
||||
|---|---|---|
|
||||
| **What** | Short facts extracted from chats | Document collections (PDFs, text files) |
|
||||
| **How stored** | Text + vector embeddings in DB | Chunked documents + embeddings in ChromaDB |
|
||||
| **How injected** | Prepended to system prompt every message | Retrieved via RAG only when relevant chunks match |
|
||||
| **Scope** | **Global per user** — injected into ALL chats | **Per-chat** — you choose which collection with `#` |
|
||||
| **Context cost** | Always consumed, every message | Only consumed when you tag a collection |
|
||||
| **Control** | All or nothing (on/off globally) | Fine-grained (pick per conversation) |
|
||||
|
||||
#### Context window consumption by memories
|
||||
|
||||
Each memory is roughly 20–80 tokens (a sentence or two). With the default `context_window_n=4` setting in Auto Memory, 4 related memories are injected per message.
|
||||
|
||||
**Estimated context consumption per message:**
|
||||
|
||||
| Stored memories | Injected per msg | Tokens consumed | % of 8K context | % of 32K context |
|
||||
|----------------|-----------------|-----------------|-----------------|------------------|
|
||||
| 10 | ~4 | ~200 | 2.5% | 0.6% |
|
||||
| 50 | ~4 | ~200 | 2.5% | 0.6% |
|
||||
| 200 | ~4 | ~200 | 2.5% | 0.6% |
|
||||
|
||||
The key insight: **only the top N similar memories are injected** (default 4), not all of them. So having 200 memories doesn't consume more context than having 10 — the retrieval system picks the most relevant ones.
|
||||
|
||||
However, **the memories are injected on every single message** in every chat. This is the overhead cost.
|
||||
|
||||
#### Context consumed by chat history (the bigger problem)
|
||||
|
||||
The real context pressure comes from **conversation history**, not memories. Here's what actually fills your context window:
|
||||
|
||||
| Messages in chat | ~Tokens used | % of 8K | % of 32K |
|
||||
|-----------------|-------------|---------|----------|
|
||||
| 5 exchanges (10 msgs) | ~2,000–4,000 | 25–50% | 6–12% |
|
||||
| 20 exchanges (40 msgs) | ~8,000–16,000 | **100%+ (truncated)** | 25–50% |
|
||||
| 100 exchanges (200 msgs) | ~40,000–80,000 | **way over** | **100%+ (truncated)** |
|
||||
|
||||
With a small model running 8K context, you'll hit the limit after ~10-20 exchanges. The model starts dropping earlier messages. Memories add a small fixed overhead (~200 tokens) on top of this.
|
||||
|
||||
#### Across multiple separate chats
|
||||
|
||||
**Good news:** separate chats do NOT share context windows. Each chat starts fresh. The only cross-chat cost is the ~200 tokens of memories injected into each new chat's system prompt.
|
||||
|
||||
So 5, 20, or 100 separate chats don't accumulate — each one independently uses the context window. Memories are the only thing that carries over.
|
||||
|
||||
### Project-scoped memory (avoiding global memory pollution)
|
||||
|
||||
**Open WebUI does NOT have native project-scoped memory.** Memories are global per user — every fact extracted from any chat gets injected into every other chat.
|
||||
|
||||
This is a known limitation. Here are workarounds:
|
||||
|
||||
#### Option 1: Use Knowledge Collections instead of Memories (recommended)
|
||||
|
||||
Knowledge collections give you the scoping you want:
|
||||
|
||||
1. **Create a collection:** Workspace → Knowledge → Create Collection (e.g. "GPU Research Project")
|
||||
2. **Add documents:** Upload PDFs, text files, or paste notes into the collection
|
||||
3. **Use per-chat:** In any chat, type `#` and select your collection — only that chat gets the context
|
||||
4. **One-off chats stay clean:** Don't tag a collection, and no project context is injected
|
||||
|
||||
This is the closest thing to "projects" in Open WebUI. You can have separate knowledge collections for separate projects, and only pull them in when relevant.
|
||||
|
||||
#### Option 2: Disable Auto Memory, use manual memories
|
||||
|
||||
1. Turn off the Auto Memory function
|
||||
2. Manually add memories via **Profile → Settings → Personalization → Memories**
|
||||
3. Keep only universally useful facts (your name, preferences, etc.)
|
||||
4. Use Knowledge collections for project-specific context
|
||||
|
||||
#### Option 3: Periodically clear memories
|
||||
|
||||
**Profile → Settings → Personalization → Memories → Clear All** — nuclear option, but keeps things clean between projects.
|
||||
|
||||
### Working with Knowledge Collections (small context window survival guide)
|
||||
|
||||
With small local models (4B-9B, 4K-8K context), you'll hit the context window limit fast — often after 10-20 exchanges. Here's how to work effectively despite that.
|
||||
|
||||
#### How Knowledge Collections actually work
|
||||
|
||||
When you type `#` in chat and select a collection, Open WebUI does **RAG retrieval** — it finds the most relevant chunks from the collection, not the whole thing. This is efficient and doesn't blow up your context window.
|
||||
|
||||
- Default chunk size: ~500 tokens
|
||||
- Typically 3-5 relevant chunks are retrieved per query (~1500-2500 tokens)
|
||||
- The chunks are injected as context alongside your message
|
||||
|
||||
#### Creating Knowledge Collections from chat conversations
|
||||
|
||||
There is **no built-in "export chat to knowledge" button** — this is a requested but unimplemented feature. Here's the practical workflow:
|
||||
|
||||
**The handoff method (like the ChatGPT dark ages, but structured):**
|
||||
|
||||
1. When you're ~60-70% through your context window (you'll feel the model getting fuzzy), ask:
|
||||
```
|
||||
Summarize everything we've discussed and decided so far. Include:
|
||||
- Key decisions made
|
||||
- Current state of the work
|
||||
- What still needs to be done
|
||||
- Any important details or constraints
|
||||
Format as a structured document I can use to continue this conversation.
|
||||
```
|
||||
2. **Copy the summary output**
|
||||
3. Go to **Workspace → Knowledge → Create Collection** (e.g. "GPU Build - Session 1")
|
||||
4. Click **Add Content** → paste the summary as a text file (`.txt` or `.md`)
|
||||
5. Start a **new chat**, type `#` and select your collection, then continue where you left off
|
||||
|
||||
**Ongoing project workflow:**
|
||||
|
||||
```
|
||||
Chat 1: Research phase
|
||||
→ Ask for summary at end
|
||||
→ Save summary to "Project X" knowledge collection
|
||||
|
||||
Chat 2: Type # → select "Project X" → continue
|
||||
→ Model gets relevant context via RAG
|
||||
→ Ask for updated summary at end
|
||||
→ Add updated summary to collection (replace or append)
|
||||
|
||||
Chat 3: Type # → select "Project X" → continue
|
||||
→ Repeat...
|
||||
```
|
||||
|
||||
Each chat starts fresh with full context window available, but can pull in relevant history from previous sessions via RAG.
|
||||
|
||||
#### Useful community functions for context management
|
||||
|
||||
Install these manually (Workspace → Functions → Create, paste the code):
|
||||
|
||||
| Function | What it does |
|
||||
|----------|-------------|
|
||||
| **Checkpoint Summarization Filter** | Auto-summarizes conversation history when context gets long |
|
||||
| **Chat Context Clipper** | Keeps only the last N messages, preserves system prompt and first message |
|
||||
| **Context Length Filter** | Hard limits on turns (default 25) and tokens (default 10,000) |
|
||||
|
||||
The **Checkpoint Summarization Filter** is the closest to automatic handoff — it summarizes older messages so the model can keep going without losing track.
|
||||
|
||||
#### How Open WebUI handles context overflow
|
||||
|
||||
By default, Open WebUI **truncates** old messages (drops them silently) — it does NOT auto-summarize. The model just loses access to earlier conversation. This is why you suddenly feel like the model "forgot" what you were talking about.
|
||||
|
||||
You can control this with:
|
||||
- **Context Length Filter** function: set max turns and token limits explicitly
|
||||
- **Chat Context Clipper**: keeps latest N messages, always preserves system prompt + first message pair
|
||||
- Or do manual handoffs before you hit the limit
|
||||
|
||||
#### RAG tuning for better knowledge retrieval
|
||||
|
||||
If your knowledge collections aren't returning good results, tune these in **Admin → Settings → Documents**:
|
||||
|
||||
| Setting | Default | Recommendation |
|
||||
|---------|---------|---------------|
|
||||
| Chunk Size | 500 | 300 for factual docs, 800 for narrative/code |
|
||||
| Chunk Overlap | 100 | 50-100 (higher = better continuity, more tokens) |
|
||||
| Top K | 4 | 3-5 (more = more context consumed) |
|
||||
| Relevance Threshold | 0.0 | 0.3-0.5 (filters out low-quality matches) |
|
||||
|
||||
**Important:** Chunk size cannot exceed your embedding model's token limit. `nomic-embed-text` supports up to 8192 tokens, so you have plenty of headroom.
|
||||
|
||||
#### The honest comparison: Local AI vs Claude Code
|
||||
|
||||
| | **Local AI (Open WebUI + Ollama)** | **Claude Code** |
|
||||
|---|---|---|
|
||||
| **Context window** | 4K-32K (model dependent) | 200K |
|
||||
| **Cross-session memory** | Manual handoffs or Auto Memory | Automatic (CLAUDE.md, project memory) |
|
||||
| **Project scoping** | Knowledge collections (manual) | Built-in (each project has its own context) |
|
||||
| **Code awareness** | RAG server auto-retrieves relevant chunks from indexed repos | Reads your entire codebase directly |
|
||||
| **Continuation** | New chat + `#collection` handoff | "Let's finish X" just works |
|
||||
| **Code editing** | Can't edit files (chat only) | Reads, writes, runs code directly |
|
||||
| **Cost** | Free (your electricity) | API usage fees |
|
||||
| **Privacy** | 100% local | Cloud-based |
|
||||
| **Offline** | Works without internet | Requires internet |
|
||||
|
||||
Local AI requires more manual workflow management but has real code awareness via the RAG server. For code-heavy editing and multi-step tasks, Claude Code is dramatically better. For private code Q&A, document research, and learning — the local stack with RAG + knowledge collections is solid.
|
||||
|
||||
### Channels (beta feature)
|
||||
|
||||
Channels are persistent chat rooms (like Slack/Discord channels) with multi-model support. They do NOT scope memories differently — memories are still global per user. Channels are useful for team collaboration, not memory isolation.
|
||||
|
||||
### Realistic expectations by model size
|
||||
|
||||
Not all models can do all tasks. Here's what to actually expect:
|
||||
|
||||
| Task | 4B (qwen3.5:4b) | 9B (qwen3.5:9b) | 35B MoE (qwen3.5-35b-a3b) | 14B+ dense |
|
||||
|------|:---:|:---:|:---:|:---:|
|
||||
| Answer simple questions | OK | Good | Good | Good |
|
||||
| Explain existing code (with RAG) | OK | Good | Good | Good |
|
||||
| Fix a simple bug (typo, off-by-one) | Maybe | Usually | Usually | Yes |
|
||||
| Write a small utility function | Shaky | OK | Good | Good |
|
||||
| Fix logic error across 2-3 functions | No | Maybe | Usually | Usually |
|
||||
| Write a new feature (multiple files) | No | Shaky | Maybe | Maybe |
|
||||
| Refactor with style consistency | No | No | Sometimes | Sometimes |
|
||||
| Summarize a conversation for handoff | OK | Good | Good | Good |
|
||||
|
||||
**The 35B MoE model (`qwen3.5-35b-a3b`) is the sweet spot for small GPUs.** It was trained as a 35B model but only activates 3B parameters per token. This means it has the *knowledge* of a 35B model with the VRAM footprint closer to a 4B. On a 6GB card it may fit (VRAM usage varies with context length and KV cache settings).
|
||||
|
||||
**Bottom line for a 6GB GPU:**
|
||||
- Use `qwen3.5:4b` for quick chat, explanations, and summarization
|
||||
- Try `qwen3.5-35b-a3b` for code tasks — if it fits, it will be significantly better than 4B
|
||||
- Use Claude Code for anything that requires reading/writing multiple files or complex reasoning
|
||||
- The RAG server helps a lot — even a 4B model gives useful answers when it has the right code chunks in context
|
||||
|
||||
### Fixing models that output code instead of natural language
|
||||
|
||||
If your model (especially smaller ones like Qwen 3.5) responds with Python code blocks instead of plain English answers (as shown in the screenshot), this is a common behavior with code-optimized models.
|
||||
|
||||
#### Why it happens
|
||||
|
||||
- Small models like `qwen3.5:4b` are heavily optimized for coding tasks
|
||||
- They interpret ambiguous questions as "write me a program" instead of "answer my question"
|
||||
- The model's training data skews toward code generation at smaller parameter counts
|
||||
- Thinking/reasoning mode (if enabled) amplifies this tendency
|
||||
|
||||
#### Fix: Set a system prompt in Open WebUI
|
||||
|
||||
1. Go to **Admin** → **Settings** → **Interface** → **Default System Prompt** (for all chats)
|
||||
|
||||
Or per-model: **Workspace** → **Models** → select model → **System Prompt**
|
||||
|
||||
2. Use a system prompt like:
|
||||
|
||||
```
|
||||
You are a helpful assistant. Answer questions in clear, natural language.
|
||||
Explain concepts conversationally. Only include code if the user explicitly
|
||||
asks for code or a script. When discussing technical topics, use plain
|
||||
English explanations with examples, not programs.
|
||||
```
|
||||
|
||||
3. For the specific GPU pricing question in the screenshot, the model should have responded with a comparison table in text, not a Python script. A good system prompt prevents this.
|
||||
|
||||
#### Alternative: Use a chat-optimized model
|
||||
|
||||
Some models are better at conversational responses:
|
||||
|
||||
| Model | Size | Better for chat? |
|
||||
|-------|------|-----------------|
|
||||
| `qwen2.5:14b` | 14B | Yes — more balanced |
|
||||
| `qwen2.5:7b` | 7B | Yes — good general chat |
|
||||
| `llama3.1:8b` | 8B | Yes — conversational |
|
||||
| `qwen2.5-coder:7b` | 7B | No — code-focused |
|
||||
| `qwen3.5:4b` | 4B | No — too small, code-biased |
|
||||
|
||||
Avoid using `-coder` variants or very small models (≤4B) for general chat. Use them only when you actually want code.
|
||||
|
||||
### VRAM considerations
|
||||
|
||||
Image generation and LLM inference compete for GPU memory. With a single GPU:
|
||||
|
||||
| VRAM | Recommendation |
|
||||
|------|---------------|
|
||||
| ≥ 24GB | Run both LLM + image gen simultaneously |
|
||||
| 12–24GB | Use smaller LLM when generating images, or stop Ollama first |
|
||||
| < 12GB | Run one at a time — stop Ollama before generating images |
|
||||
|
||||
ComfyUI models typically need 4–8GB VRAM (SD 1.5: ~4GB, SDXL: ~7GB, Flux: ~12GB).
|
||||
|
||||
---
|
||||
|
||||
## InvokeAI vs ComfyUI — which to use when
|
||||
|
||||
Both are installed by the setup script. Here's when to use each:
|
||||
|
||||
| Task | InvokeAI (`:9090`) | ComfyUI (`:8188`) |
|
||||
|------|:---:|:---:|
|
||||
| Friendly UI, sliders, drag-and-drop | Yes | No (node editor) |
|
||||
| Use a LoRA you trained on RunPod | Yes — just import + select | Yes — but wire LoRA Loader node |
|
||||
| Same face, different settings (img2img) | Yes — Image to Image tab | Yes — IP-Adapter nodes |
|
||||
| Age/emotion/scene changes | Yes — img2img + prompt | Yes — IP-Adapter + prompt |
|
||||
| Chat-integrated image gen (from Open WebUI) | No (no OWUI API) | Yes (workflow export) |
|
||||
| Maximum flexibility / custom pipelines | No | Yes |
|
||||
| Learning curve | Low | High |
|
||||
|
||||
**Bottom line:** Use InvokeAI for the interactive "play with images" workflow. Use ComfyUI only when you need Open WebUI chat integration or advanced node pipelines.
|
||||
|
||||
## Using InvokeAI for image iteration
|
||||
|
||||
InvokeAI is the easier path for what you want — take a reference image and generate variations with different settings, ages, emotions, art styles.
|
||||
|
||||
### Step 1: Install a base model (required first)
|
||||
|
||||
Before LoRAs or img2img will work, InvokeAI needs a base Stable Diffusion model.
|
||||
|
||||
```bash
|
||||
# Option A: Download SD 1.5 (~4GB VRAM needed) via the container
|
||||
docker exec invokeai invokeai-model-install --add stabilityai/stable-diffusion-v1-5
|
||||
|
||||
# Option B: Download SDXL (~7GB VRAM needed)
|
||||
docker exec invokeai invokeai-model-install --add stabilityai/stable-diffusion-xl-base-1.0
|
||||
```
|
||||
|
||||
Or use the UI: open `http://<ip>:9090` → **Model Manager** (cube icon) → **Starter Models** tab → install SD 1.5 or SDXL.
|
||||
|
||||
> **6GB GPU note:** SD 1.5 fits comfortably. SDXL is tight — it may work with float16 (already set) but could be slow. Start with SD 1.5.
|
||||
|
||||
### Step 2: Import your RunPod LoRA
|
||||
|
||||
```bash
|
||||
# Copy the .safetensors file from wherever you downloaded it:
|
||||
./invokeai-import-lora.sh ~/Downloads/my-lora.safetensors "My Character"
|
||||
```
|
||||
|
||||
Then in the InvokeAI UI:
|
||||
1. **Model Manager** (cube icon) → **Scan for Models**
|
||||
2. Your LoRA should appear in the list
|
||||
|
||||
If it doesn't:
|
||||
- Click **Add Model** → **Scan Folder** → enter `/invokeai/models/lora`
|
||||
- Make sure the LoRA architecture matches your base model (SD 1.5 LoRA needs SD 1.5 base)
|
||||
|
||||
### Step 3: Generate with your LoRA (text-to-image)
|
||||
|
||||
1. Go to the **Text to Image** tab
|
||||
2. Select your **base model** (must match the LoRA's training base)
|
||||
3. In the left panel, expand the **LoRA** section (below model selector)
|
||||
4. Click **+** → select your LoRA → set weight to **0.7–0.85**
|
||||
5. Write a prompt: "portrait of [subject], smiling, studio lighting"
|
||||
6. Click **Invoke**
|
||||
|
||||
### Step 4: Iterate on an image (img2img)
|
||||
|
||||
This is where InvokeAI shines for your use case — take an image and riff on it:
|
||||
|
||||
1. Switch to the **Image to Image** tab
|
||||
2. Drag your reference photo onto the canvas (or click to upload)
|
||||
3. Keep your LoRA active (same as above)
|
||||
4. Set the **Denoising Strength** slider:
|
||||
|
||||
| Strength | Effect |
|
||||
|----------|--------|
|
||||
| 0.2–0.3 | Subtle tweaks — mostly keeps the original, minor style changes |
|
||||
| 0.4–0.5 | Moderate changes — recognizable but different mood/lighting |
|
||||
| 0.6–0.7 | Significant changes — same composition, new details/style |
|
||||
| 0.8–1.0 | Major rewrite — loosely inspired by original, mostly new |
|
||||
|
||||
5. Change the prompt to describe what you want different:
|
||||
- **Age up:** "same person, elderly, wrinkles, grey hair, wise expression"
|
||||
- **Age down:** "same person as a young child, bright eyes, playground"
|
||||
- **Emotion:** "same person, laughing joyfully" or "same person, crying, dramatic lighting"
|
||||
- **Setting:** "same person, sitting in a Parisian cafe, afternoon light"
|
||||
- **Art style:** "same person, oil painting, renaissance style, dramatic chiaroscuro"
|
||||
6. Click **Invoke** — iterate by adjusting strength and prompt
|
||||
|
||||
### Step 5: Use the Unified Canvas for inpainting
|
||||
|
||||
This is the "fix this specific thing" workflow — brush over a hand, arm, face, background, whatever, and regenerate just that area while keeping everything else untouched.
|
||||
|
||||
1. Switch to the **Unified Canvas** tab
|
||||
2. Upload or paste your image
|
||||
3. Select the **Mask** brush tool (not the paint brush)
|
||||
4. Brush over **only the area you want to change** — everything else stays locked
|
||||
5. Write a prompt describing what the masked area should become
|
||||
6. Set **Denoising Strength** to 0.6–0.8 (higher = more change)
|
||||
7. Click **Invoke** — only the masked pixels regenerate
|
||||
|
||||
**Common inpainting fixes:**
|
||||
|
||||
| Problem | Mask | Prompt |
|
||||
|---------|------|--------|
|
||||
| Hand in wrong position | Brush over the arm/hand | "natural hand resting at side, relaxed pose" |
|
||||
| Extra fingers | Brush over the hand | "normal human hand, five fingers, anatomically correct" |
|
||||
| Weird eyes | Brush over both eyes | "natural eyes, looking at camera, detailed iris" |
|
||||
| Bad background | Brush over background only | "clean studio backdrop" or "forest trail, golden hour" |
|
||||
| Wrong clothing | Brush over the clothing area | "wearing blue denim jacket, casual style" |
|
||||
| Face swap / aging | Brush over the face | "same person, elderly, wrinkles" or "same person as child" |
|
||||
|
||||
**Tips for better inpainting results:**
|
||||
- **Mask slightly larger** than the problem area — gives the model room to blend edges
|
||||
- **Use soft brush edges** (lower brush hardness) for more natural blending
|
||||
- If the result has visible seams, increase your mask area and try again
|
||||
- **Lower denoising (0.4–0.5)** for subtle fixes, **higher (0.7–0.9)** for major changes
|
||||
- Keep your LoRA active during inpainting — it maintains the trained style/face consistency
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
- **Greyed-out upload/LoRA buttons:** Install a base model first (Step 1). InvokeAI disables most features until a checkpoint is loaded.
|
||||
- **Model not synced:** After copying files via script, click **Scan for Models** in Model Manager.
|
||||
- **Architecture mismatch:** A LoRA trained on SD 1.5 only works with SD 1.5 base models — not SDXL. Check what your RunPod training used.
|
||||
- **Out of VRAM:** Try SD 1.5 instead of SDXL, or reduce image size to 512x512.
|
||||
- **Use the import script:** The greyed-out UI upload can be bypassed entirely by using `invokeai-import-lora.sh` to copy files directly into the model volume.
|
||||
|
||||
---
|
||||
|
||||
## Updating
|
||||
|
||||
Re-run the setup script — it detects an existing install and skips prereqs:
|
||||
|
||||
```bash
|
||||
./laptop_full_setup.sh
|
||||
# or
|
||||
./laptop_full_setup.sh --force # also overwrites config files
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## GPU / Model Tiers (`local-ai-setup.sh`)
|
||||
|
||||
| VRAM | Chat model | Code model | Context |
|
||||
|---------|-----------------|-----------------------|---------|
|
||||
| ≥ 14 GB | qwen2.5:14b | qwen2.5-coder:14b | 32k |
|
||||
| 8–14 GB | qwen2.5:14b | qwen2.5-coder:7b | 16k |
|
||||
| 4–8 GB | qwen2.5:7b | qwen2.5-coder:7b | 8k |
|
||||
| CPU | qwen2.5:7b | qwen2.5-coder:7b | 4k |
|
||||
|
||||
Embed model is always `nomic-embed-text` (required for RAG).
|
||||
|
||||
### VRAM reality check
|
||||
|
||||
Ollama will **always try to run** any model — it silently offloads layers to CPU when VRAM is insufficient. The model still works but gets significantly slower. The setup script's "fully in VRAM" label can be misleading.
|
||||
|
||||
**Actual VRAM needed for common models (Q4_K_M quantization):**
|
||||
|
||||
| Model | Download size | VRAM for inference | Fits in 6GB? | Fits in 8GB? |
|
||||
|-------|-------------|-------------------|-------------|-------------|
|
||||
| `qwen3.5:4b` | ~2.5 GB | ~3.5–4 GB | Yes | Yes |
|
||||
| `qwen2.5:7b` | ~4.4 GB | ~5.5–6 GB | Tight | Yes |
|
||||
| `qwen3.5:9b` | ~5.5 GB | ~6.5–7 GB | **No — partial CPU offload** | Tight |
|
||||
| `qwen2.5:14b` | ~8.7 GB | ~10–11 GB | No | No |
|
||||
| `qwen3.5-35b-a3b` (MoE) | ~20 GB | ~3.5 GB active | Yes (only 3B active) | Yes |
|
||||
|
||||
**Why the file size != VRAM needed:** Inference requires additional memory for KV cache, attention buffers, and CUDA overhead. Expect ~1-2 GB more than the model file size.
|
||||
|
||||
**Signs of CPU offload (model too big for your VRAM):**
|
||||
- Tokens per second drops from 20-40 to 2-8
|
||||
- `nvidia-smi` shows VRAM maxed out
|
||||
- CPU usage spikes during generation
|
||||
- First token takes much longer than usual
|
||||
|
||||
### Image Generation Model Tiers
|
||||
|
||||
The `setup-image-models.sh` script detects your GPU and offers appropriate models:
|
||||
|
||||
| VRAM | Available Models | Default | Notes |
|
||||
|------|-----------------|---------|-------|
|
||||
| ≥ 24GB | SD 1.5, SDXL, SDXL Turbo, Flux.1-schnell, Flux.1-dev | SDXL | All models, no constraints |
|
||||
| 12–23GB | SD 1.5, SDXL, SDXL Turbo, Flux.1-schnell | SDXL | Flux-dev too tight |
|
||||
| 8–11GB | SD 1.5, SDXL (tight), SDXL Turbo | SD 1.5 | SDXL works at 512px, may be slow |
|
||||
| 4–7GB | SD 1.5 (float16) | SD 1.5 | Only SD 1.5 fits |
|
||||
| < 4GB | none | — | CPU generation not recommended |
|
||||
|
||||
**GPU sharing:** Ollama and image generation share the GPU. Ollama auto-unloads models
|
||||
after its `KEEP_ALIVE` timeout (default 24h), so image gen gets full VRAM when the LLM
|
||||
is idle. For immediate unload: `docker exec ollama ollama stop <model-name>`
|
||||
|
||||
**Multi-GPU scaling:** With dual GPUs (e.g., 2× RTX 5000 = 32GB total), the VRAM
|
||||
is summed for tier selection. Both InvokeAI and ComfyUI will use all available GPUs.
|
||||
|
||||
```bash
|
||||
# Install image models (auto-detects GPU):
|
||||
./setup-image-models.sh
|
||||
|
||||
# Or auto-install the recommended default:
|
||||
./setup-image-models.sh --auto
|
||||
```
|
||||
+97
@@ -0,0 +1,97 @@
|
||||
#!/usr/bin/env bash
|
||||
# Import a LoRA (.safetensors) into ComfyUI's Docker volume
|
||||
# Usage: ./comfyui-import-lora.sh /path/to/my-lora.safetensors
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
|
||||
if [[ $# -lt 1 ]]; then
|
||||
echo -e "${BOLD}Usage:${NC} $0 <lora-file.safetensors> [display-name]"
|
||||
echo ""
|
||||
echo " Copies a LoRA file into ComfyUI's models/loras/ directory so it"
|
||||
echo " can be used in workflows (including via Open WebUI)."
|
||||
echo ""
|
||||
echo " Examples:"
|
||||
echo " $0 ~/Downloads/my-style-lora.safetensors"
|
||||
echo " $0 ~/Downloads/my-style-lora.safetensors \"Anime Style\""
|
||||
echo ""
|
||||
echo " After importing, create a workflow in ComfyUI that uses this LoRA,"
|
||||
echo " export it as API format, and import into Open WebUI for chat-based"
|
||||
echo " image generation with this LoRA applied automatically."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
LORA_FILE="$1"
|
||||
DISPLAY_NAME="${2:-$(basename "$LORA_FILE" .safetensors)}"
|
||||
|
||||
# Validate file exists
|
||||
if [[ ! -f "$LORA_FILE" ]]; then
|
||||
echo -e "${RED}Error:${NC} File not found: $LORA_FILE"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Validate file extension
|
||||
if [[ "$LORA_FILE" != *.safetensors && "$LORA_FILE" != *.ckpt && "$LORA_FILE" != *.pt ]]; then
|
||||
echo -e "${YELLOW}Warning:${NC} File doesn't have a typical LoRA extension (.safetensors, .ckpt, .pt)"
|
||||
read -rp "Continue anyway? [y/N] " yn
|
||||
[[ "$yn" != [yY]* ]] && exit 1
|
||||
fi
|
||||
|
||||
# Check if ComfyUI container exists
|
||||
if ! docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^comfyui$'; then
|
||||
echo -e "${RED}Error:${NC} ComfyUI container not found. Is the AI stack running?"
|
||||
echo " Try: docker compose up -d comfyui"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if container is running
|
||||
if ! docker ps --format '{{.Names}}' 2>/dev/null | grep -q '^comfyui$'; then
|
||||
echo -e "${YELLOW}ComfyUI container is stopped. Starting it...${NC}"
|
||||
docker start comfyui
|
||||
sleep 3
|
||||
fi
|
||||
|
||||
FILENAME=$(basename "$LORA_FILE")
|
||||
|
||||
echo -e "${CYAN}[..]${NC} Copying ${BOLD}$FILENAME${NC} into ComfyUI..."
|
||||
|
||||
# Copy the LoRA file into the container's loras directory
|
||||
docker exec comfyui mkdir -p /opt/ComfyUI/models/loras
|
||||
docker cp "$LORA_FILE" "comfyui:/opt/ComfyUI/models/loras/$FILENAME"
|
||||
|
||||
echo -e "${GREEN}[OK]${NC} LoRA '${DISPLAY_NAME}' copied to ComfyUI!"
|
||||
echo ""
|
||||
echo -e "${BOLD}Next: Create a workflow that uses this LoRA${NC}"
|
||||
echo ""
|
||||
echo " 1. Open ComfyUI at http://localhost:8188"
|
||||
echo " 2. Load or create a text-to-image workflow"
|
||||
echo " 3. Add a ${BOLD}LoRA Loader${NC} node:"
|
||||
echo " Right-click canvas → Add Node → loaders → Load LoRA"
|
||||
echo " 4. Wire it between the ${BOLD}checkpoint loader${NC} and the ${BOLD}CLIP/sampler${NC}:"
|
||||
echo ""
|
||||
echo " [Load Checkpoint] → MODEL → [Load LoRA] → MODEL → [KSampler]"
|
||||
echo " → CLIP → → CLIP → [CLIP Text Encode]"
|
||||
echo ""
|
||||
echo " 5. In the LoRA Loader node, select ${BOLD}${FILENAME}${NC}"
|
||||
echo " 6. Set ${BOLD}strength_model${NC} and ${BOLD}strength_clip${NC} (start with 0.7–0.85)"
|
||||
echo ""
|
||||
echo -e "${BOLD}To use from Open WebUI chat (no more ComfyUI interaction needed):${NC}"
|
||||
echo ""
|
||||
echo " 7. Click the ${BOLD}gear icon${NC} → enable ${BOLD}Dev Mode${NC}"
|
||||
echo " 8. Click ${BOLD}Save (API Format)${NC} → saves workflow_api.json"
|
||||
echo " 9. In Open WebUI → Admin → Settings → Images → ${BOLD}Import Workflow${NC}"
|
||||
echo " 10. Upload the workflow_api.json and map the prompt node"
|
||||
echo " 11. Now just chat: \"Generate an image of a forest in ${DISPLAY_NAME} style\""
|
||||
echo ""
|
||||
echo -e "${YELLOW}Tip:${NC} The LoRA must match your base model architecture."
|
||||
echo " SD 1.5 LoRA → SD 1.5 checkpoint. SDXL LoRA → SDXL checkpoint."
|
||||
echo ""
|
||||
echo -e "${YELLOW}Tip:${NC} You can ${BOLD}chain multiple LoRAs${NC} in one workflow:"
|
||||
echo " [Checkpoint] → [Load LoRA 1] → [Load LoRA 2] → [KSampler]"
|
||||
echo " Each has its own strength slider so you can blend styles."
|
||||
echo ""
|
||||
echo -e "${YELLOW}Tip:${NC} Export multiple workflows (one per LoRA combo) and switch"
|
||||
echo " between them in Open WebUI's image settings as needed."
|
||||
echo " LoRAs are baked into the workflow — there's no keyword to"
|
||||
echo " toggle them on/off from chat."
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
#!/usr/bin/env bash
|
||||
# Install IP-Adapter custom nodes + models into ComfyUI's Docker container
|
||||
# Enables: same face/different settings, age up/down, emotion changes, style transfer
|
||||
#
|
||||
# Usage:
|
||||
# ./comfyui-install-ipadapter.sh — install for SDXL (recommended)
|
||||
# ./comfyui-install-ipadapter.sh --sd15 — install for SD 1.5
|
||||
# ./comfyui-install-ipadapter.sh --all — install both SDXL + SD 1.5
|
||||
# ./comfyui-install-ipadapter.sh --faceid — also install FaceID (better face lock)
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
info() { echo -e "${CYAN}[..]${NC} $*"; }
|
||||
ok() { echo -e "${GREEN}[OK]${NC} $*"; }
|
||||
warn() { echo -e "${YELLOW}[!!]${NC} $*"; }
|
||||
die() { echo -e "${RED}[XX]${NC} $*" >&2; exit 1; }
|
||||
|
||||
# ── parse args ──────────────────────────────────────────────────────────────
|
||||
INSTALL_SD15=false
|
||||
INSTALL_SDXL=true
|
||||
INSTALL_FACEID=false
|
||||
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--sd15) INSTALL_SD15=true; INSTALL_SDXL=false ;;
|
||||
--all) INSTALL_SD15=true; INSTALL_SDXL=true ;;
|
||||
--faceid) INSTALL_FACEID=true ;;
|
||||
--help|-h)
|
||||
echo -e "${BOLD}Usage:${NC} $0 [--sd15] [--all] [--faceid]"
|
||||
echo ""
|
||||
echo " Installs IP-Adapter custom nodes and models into ComfyUI."
|
||||
echo " Enables reference-image workflows: same face in different"
|
||||
echo " settings, age changes, emotions, style transfer."
|
||||
echo ""
|
||||
echo " Options:"
|
||||
echo " --sd15 Install SD 1.5 models (instead of SDXL)"
|
||||
echo " --all Install both SDXL + SD 1.5 models"
|
||||
echo " --faceid Also install FaceID models (better face lock,"
|
||||
echo " requires insightface — adds ~1GB)"
|
||||
echo ""
|
||||
echo " Default: SDXL models only (~2.5GB download)"
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# ── check container ─────────────────────────────────────────────────────────
|
||||
if ! docker ps --format '{{.Names}}' 2>/dev/null | grep -q '^comfyui$'; then
|
||||
if docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^comfyui$'; then
|
||||
warn "ComfyUI container is stopped. Starting it..."
|
||||
docker start comfyui
|
||||
sleep 5
|
||||
else
|
||||
die "ComfyUI container not found. Is the AI stack running?"
|
||||
fi
|
||||
fi
|
||||
|
||||
# ── helper: download into container ─────────────────────────────────────────
|
||||
# Usage: dl_model <url> <container_path>
|
||||
dl_model() {
|
||||
local url="$1" dest="$2"
|
||||
local filename; filename=$(basename "$dest")
|
||||
if docker exec comfyui test -f "$dest" 2>/dev/null; then
|
||||
ok "Already exists: $filename"
|
||||
return 0
|
||||
fi
|
||||
local dir; dir=$(dirname "$dest")
|
||||
docker exec comfyui mkdir -p "$dir"
|
||||
info "Downloading $filename..."
|
||||
if ! docker exec comfyui wget -q --show-progress -O "$dest" "$url" 2>&1; then
|
||||
# wget --show-progress not always available
|
||||
docker exec comfyui wget -q -O "$dest" "$url"
|
||||
fi
|
||||
ok "Downloaded $filename"
|
||||
}
|
||||
|
||||
HF_IPA="https://huggingface.co/h94/IP-Adapter/resolve/main"
|
||||
MODELS="/opt/ComfyUI/models"
|
||||
|
||||
# ── Step 1: Install custom nodes ────────────────────────────────────────────
|
||||
echo ""
|
||||
echo -e "${BOLD}━━━ IP-Adapter Installation for ComfyUI ━━━${NC}"
|
||||
echo ""
|
||||
|
||||
info "Installing ComfyUI_IPAdapter_plus custom nodes..."
|
||||
if docker exec comfyui test -d /opt/ComfyUI/custom_nodes/ComfyUI_IPAdapter_plus 2>/dev/null; then
|
||||
info "Custom nodes already installed — pulling updates..."
|
||||
docker exec comfyui bash -c "cd /opt/ComfyUI/custom_nodes/ComfyUI_IPAdapter_plus && git pull -q"
|
||||
ok "Custom nodes updated"
|
||||
else
|
||||
docker exec comfyui bash -c "cd /opt/ComfyUI/custom_nodes && git clone --depth 1 https://github.com/cubiq/ComfyUI_IPAdapter_plus.git"
|
||||
ok "Custom nodes installed"
|
||||
fi
|
||||
|
||||
# Install Python dependencies if requirements.txt exists
|
||||
if docker exec comfyui test -f /opt/ComfyUI/custom_nodes/ComfyUI_IPAdapter_plus/requirements.txt 2>/dev/null; then
|
||||
info "Installing Python dependencies..."
|
||||
docker exec comfyui pip install -q -r /opt/ComfyUI/custom_nodes/ComfyUI_IPAdapter_plus/requirements.txt 2>/dev/null || true
|
||||
fi
|
||||
|
||||
# ── Step 2: CLIP Vision models (required by all variants) ───────────────────
|
||||
echo ""
|
||||
info "Downloading CLIP Vision encoders..."
|
||||
|
||||
dl_model \
|
||||
"https://huggingface.co/h94/IP-Adapter/resolve/main/models/image_encoder/model.safetensors" \
|
||||
"$MODELS/clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"
|
||||
|
||||
if $INSTALL_SDXL; then
|
||||
dl_model \
|
||||
"https://huggingface.co/h94/IP-Adapter/resolve/main/sdxl_models/image_encoder/model.safetensors" \
|
||||
"$MODELS/clip_vision/CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors"
|
||||
fi
|
||||
|
||||
# ── Step 3: IP-Adapter models ──────────────────────────────────────────────
|
||||
echo ""
|
||||
docker exec comfyui mkdir -p "$MODELS/ipadapter"
|
||||
|
||||
if $INSTALL_SDXL; then
|
||||
info "Downloading SDXL IP-Adapter models..."
|
||||
dl_model "$HF_IPA/sdxl_models/ip-adapter-plus_sdxl_vit-h.safetensors" \
|
||||
"$MODELS/ipadapter/ip-adapter-plus_sdxl_vit-h.safetensors"
|
||||
dl_model "$HF_IPA/sdxl_models/ip-adapter-plus-face_sdxl_vit-h.safetensors" \
|
||||
"$MODELS/ipadapter/ip-adapter-plus-face_sdxl_vit-h.safetensors"
|
||||
fi
|
||||
|
||||
if $INSTALL_SD15; then
|
||||
info "Downloading SD 1.5 IP-Adapter models..."
|
||||
dl_model "$HF_IPA/models/ip-adapter-plus_sd15.safetensors" \
|
||||
"$MODELS/ipadapter/ip-adapter-plus_sd15.safetensors"
|
||||
dl_model "$HF_IPA/models/ip-adapter-plus-face_sd15.safetensors" \
|
||||
"$MODELS/ipadapter/ip-adapter-plus-face_sd15.safetensors"
|
||||
fi
|
||||
|
||||
# ── Step 4 (optional): FaceID models ───────────────────────────────────────
|
||||
if $INSTALL_FACEID; then
|
||||
echo ""
|
||||
info "Installing FaceID dependencies (insightface, onnxruntime)..."
|
||||
docker exec comfyui pip install -q insightface onnxruntime 2>/dev/null \
|
||||
|| warn "Could not install insightface — FaceID nodes may not work"
|
||||
|
||||
HF_FACEID="https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main"
|
||||
info "Downloading FaceID models..."
|
||||
|
||||
if $INSTALL_SDXL; then
|
||||
dl_model "$HF_FACEID/ip-adapter-faceid-plusv2_sdxl.bin" \
|
||||
"$MODELS/ipadapter/ip-adapter-faceid-plusv2_sdxl.bin"
|
||||
fi
|
||||
if $INSTALL_SD15; then
|
||||
dl_model "$HF_FACEID/ip-adapter-faceid-plusv2_sd15.bin" \
|
||||
"$MODELS/ipadapter/ip-adapter-faceid-plusv2_sd15.bin"
|
||||
fi
|
||||
|
||||
# FaceID LoRAs (required for FaceID models)
|
||||
docker exec comfyui mkdir -p "$MODELS/loras"
|
||||
if $INSTALL_SDXL; then
|
||||
dl_model "$HF_FACEID/ip-adapter-faceid-plusv2_sdxl_lora.safetensors" \
|
||||
"$MODELS/loras/ip-adapter-faceid-plusv2_sdxl_lora.safetensors"
|
||||
fi
|
||||
if $INSTALL_SD15; then
|
||||
dl_model "$HF_FACEID/ip-adapter-faceid-plusv2_sd15_lora.safetensors" \
|
||||
"$MODELS/loras/ip-adapter-faceid-plusv2_sd15_lora.safetensors"
|
||||
fi
|
||||
fi
|
||||
|
||||
# ── Done ────────────────────────────────────────────────────────────────────
|
||||
echo ""
|
||||
echo -e "${GREEN}━━━ Installation complete! ━━━${NC}"
|
||||
echo ""
|
||||
echo -e "${BOLD}Restart ComfyUI to load new nodes:${NC}"
|
||||
echo " docker restart comfyui"
|
||||
echo ""
|
||||
echo -e "${BOLD}Then open ComfyUI at http://localhost:8188 and try it:${NC}"
|
||||
echo ""
|
||||
echo " Basic reference image (style/scene transfer):"
|
||||
echo " ┌──────────────────────────────────────────────────────────────────┐"
|
||||
echo " │ [Load Checkpoint]──→[Load Image]──→[IPAdapter Unified Loader] │"
|
||||
echo " │ │ │ │"
|
||||
echo " │ ├── MODEL ──────────────→ [IPAdapter Apply] → [KSampler]│"
|
||||
echo " │ └── CLIP → [CLIP Text Encode]──→ positive ──→ │"
|
||||
echo " │ \"a castle at sunset\" │"
|
||||
echo " └──────────────────────────────────────────────────────────────────┘"
|
||||
echo ""
|
||||
echo " 1. Add nodes: right-click → Add Node → ipadapter"
|
||||
echo " 2. ${BOLD}IPAdapter Unified Loader${NC}: set preset to '${BOLD}PLUS FACE (portrait)${NC}'"
|
||||
echo " 3. ${BOLD}Load Image${NC}: upload your reference photo"
|
||||
echo " 4. ${BOLD}IPAdapter Apply${NC}: connect model + image, set weight 0.7–1.0"
|
||||
echo " 5. Text prompt controls the new scene: \"elderly, sitting in cafe, smiling\""
|
||||
echo ""
|
||||
echo -e "${BOLD}What each preset does:${NC}"
|
||||
echo " PLUS — general style/scene transfer"
|
||||
echo " PLUS FACE — preserves face likeness (best for your use case)"
|
||||
if $INSTALL_FACEID; then
|
||||
echo " FACEID PLUSV2 — strongest face lock (uses insightface for detection)"
|
||||
fi
|
||||
echo ""
|
||||
echo -e "${BOLD}Example prompts with a reference face:${NC}"
|
||||
echo " • \"same person, elderly, wise expression, studio lighting\""
|
||||
echo " • \"same person as a child, playing in a park, happy\""
|
||||
echo " • \"same person, crying, dramatic lighting, black and white\""
|
||||
echo " • \"same person, oil painting style, renaissance setting\""
|
||||
echo ""
|
||||
echo -e "${YELLOW}Note:${NC} This works directly in ComfyUI. Open WebUI's ComfyUI integration"
|
||||
echo " only sends text prompts — it can't attach a reference image."
|
||||
echo " For chat-based image gen (without reference images), use LoRAs instead."
|
||||
+234
@@ -0,0 +1,234 @@
|
||||
#!/usr/bin/env bash
|
||||
# configure-searxng-safesearch.sh
|
||||
# Set SearXNG safe-search level; optionally disable categories or engines.
|
||||
#
|
||||
# Usage:
|
||||
# ./configure-searxng-safesearch.sh [strict|moderate|none] [OPTIONS]
|
||||
#
|
||||
# Options:
|
||||
# --disable-categories cat1,cat2 videos images news science social
|
||||
# --disable-engines eng1,eng2 e.g. duckduckgo,bing,yandex
|
||||
# --enable-engines eng1,eng2 re-enable engines auto-disabled by level
|
||||
#
|
||||
# Examples:
|
||||
# ./configure-searxng-safesearch.sh strict
|
||||
# ./configure-searxng-safesearch.sh moderate --disable-categories videos,images
|
||||
# ./configure-searxng-safesearch.sh none --disable-engines bing,duckduckgo
|
||||
# ./configure-searxng-safesearch.sh strict --disable-categories videos \
|
||||
# --disable-engines bing --enable-engines yandex
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
red() { printf '\e[31m%s\e[0m\n' "$*"; }
|
||||
grn() { printf '\e[32m%s\e[0m\n' "$*"; }
|
||||
blu() { printf '\e[34m%s\e[0m\n' "$*"; }
|
||||
yel() { printf '\e[33m%s\e[0m\n' "$*"; }
|
||||
die() { red "ERROR: $*"; exit 1; }
|
||||
ok() { grn " ✓ $*"; }
|
||||
info() { blu " → $*"; }
|
||||
warn() { yel " ! $*"; }
|
||||
|
||||
# ── Defaults ──────────────────────────────────────────────────────────────────
|
||||
LEVEL="moderate"
|
||||
DISABLE_CATS=""
|
||||
DISABLE_ENGINES_EXTRA=""
|
||||
ENABLE_ENGINES_EXTRA=""
|
||||
BASE="${BASE:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
|
||||
|
||||
# ── Parse arguments ───────────────────────────────────────────────────────────
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
strict|moderate|none) LEVEL="$1"; shift ;;
|
||||
--disable-categories) DISABLE_CATS="$2"; shift 2 ;;
|
||||
--disable-engines) DISABLE_ENGINES_EXTRA="$2"; shift 2 ;;
|
||||
--enable-engines) ENABLE_ENGINES_EXTRA="$2"; shift 2 ;;
|
||||
--base) BASE="$2"; shift 2 ;;
|
||||
-h|--help)
|
||||
sed -n '2,20p' "$0" | sed 's/^# \?//'
|
||||
exit 0
|
||||
;;
|
||||
*) die "Unknown argument: $1 (run with --help)" ;;
|
||||
esac
|
||||
done
|
||||
|
||||
SETTINGS="$BASE/searxng/settings.yml"
|
||||
COMPOSE="$BASE/docker-compose.yml"
|
||||
|
||||
# Ensure settings directory exists (fix ownership if Docker created it as root)
|
||||
mkdir -p "$(dirname "$SETTINGS")" 2>/dev/null || \
|
||||
sudo mkdir -p "$(dirname "$SETTINGS")"
|
||||
if [[ -e "$SETTINGS" ]] && [[ ! -w "$SETTINGS" ]]; then
|
||||
warn "settings.yml not writable — fixing ownership"
|
||||
sudo chown "$(id -u):$(id -g)" "$SETTINGS" || \
|
||||
die "Cannot write to $SETTINGS — run: sudo chown $USER $SETTINGS"
|
||||
fi
|
||||
|
||||
# ── Level → integer ───────────────────────────────────────────────────────────
|
||||
case "$LEVEL" in
|
||||
none) SAFE_INT=0 ;;
|
||||
moderate) SAFE_INT=1 ;;
|
||||
strict) SAFE_INT=2 ;;
|
||||
esac
|
||||
|
||||
info "Safe search: $LEVEL (${SAFE_INT})"
|
||||
|
||||
# ── Engines with no safe-search support ───────────────────────────────────────
|
||||
# Auto-disabled when level is moderate or strict.
|
||||
NO_SAFESEARCH_ENGINES=(
|
||||
# Torrent / P2P — no filtering possible
|
||||
"1337x" "piratebay" "nyaa" "torrentz" "kickass torrents"
|
||||
# Web engines without safe-search API
|
||||
"mojeek" "naver" "baidu"
|
||||
# Yandex: parameter exists but not reliably enforced for non-Russian queries
|
||||
"yandex" "yandex images"
|
||||
# Video frontends — no safe-search passthrough
|
||||
"invidious" "piped" "peertube" "sepiasearch"
|
||||
)
|
||||
|
||||
# ── Category → engine lists ───────────────────────────────────────────────────
|
||||
VIDEOS_ENGINES=(
|
||||
"youtube" "bing videos" "brave videos" "duckduckgo videos" "google videos" "qwant videos"
|
||||
"dailymotion" "media.ccc.de" "wikcommons.videos"
|
||||
"vimeo" "odysee" "rumble" "bitchute"
|
||||
"invidious" "piped" "peertube" "sepiasearch"
|
||||
)
|
||||
IMAGES_ENGINES=(
|
||||
"google images" "bing images" "duckduckgo images" "brave images" "qwant images"
|
||||
"startpage images" "mojeek images" "presearch images"
|
||||
"openverse" "unsplash" "pexels" "pixabay images" "pinterest" "flickr"
|
||||
"wikcommons.images" "artic" "yandex images"
|
||||
"imgur" "deviantart" "artstation" "adobe stock"
|
||||
)
|
||||
NEWS_ENGINES=(
|
||||
"google news" "bing news" "duckduckgo news" "brave news" "qwant news"
|
||||
"startpage news" "presearch news" "mojeek news"
|
||||
"reuters" "yahoo news" "wikinews" "yep news"
|
||||
)
|
||||
SCIENCE_ENGINES=(
|
||||
"arxiv" "semantic scholar" "pubmed" "crossref" "base"
|
||||
)
|
||||
SOCIAL_ENGINES=(
|
||||
"reddit" "lemmy" "mastodon"
|
||||
)
|
||||
|
||||
# ── Build disable/enable maps ─────────────────────────────────────────────────
|
||||
declare -A DISABLE_MAP # engine → 1
|
||||
declare -A ENABLE_MAP # engine → 1 (overrides everything)
|
||||
|
||||
# Parse --enable-engines
|
||||
if [[ -n "$ENABLE_ENGINES_EXTRA" ]]; then
|
||||
IFS=',' read -ra _engs <<< "$ENABLE_ENGINES_EXTRA"
|
||||
for e in "${_engs[@]}"; do
|
||||
e="${e#"${e%%[![:space:]]*}"}"; e="${e%"${e##*[![:space:]]}"}" # trim
|
||||
[[ -n "$e" ]] && ENABLE_MAP["$e"]=1
|
||||
done
|
||||
fi
|
||||
|
||||
# Helper: add to DISABLE_MAP unless explicitly re-enabled
|
||||
mark_disabled() {
|
||||
local eng="$1"
|
||||
# Use set +u to safely check array key existence on bash < 4.4 (set -u quirk)
|
||||
set +u; local _chk="${ENABLE_MAP[$eng]+x}"; set -u
|
||||
[[ -n "$_chk" ]] && return # user said keep it
|
||||
DISABLE_MAP["$eng"]=1
|
||||
}
|
||||
|
||||
# Auto-disable no-safesearch engines for moderate/strict
|
||||
if [[ "$LEVEL" != "none" ]]; then
|
||||
for eng in "${NO_SAFESEARCH_ENGINES[@]}"; do
|
||||
mark_disabled "$eng"
|
||||
done
|
||||
fi
|
||||
|
||||
# Category disables
|
||||
if [[ -n "$DISABLE_CATS" ]]; then
|
||||
IFS=',' read -ra _cats <<< "$DISABLE_CATS"
|
||||
for cat in "${_cats[@]}"; do
|
||||
cat="${cat#"${cat%%[![:space:]]*}"}"; cat="${cat%"${cat##*[![:space:]]}"}"
|
||||
cat="${cat,,}"
|
||||
case "$cat" in
|
||||
videos) for e in "${VIDEOS_ENGINES[@]}"; do mark_disabled "$e"; done ;;
|
||||
images) for e in "${IMAGES_ENGINES[@]}"; do mark_disabled "$e"; done ;;
|
||||
news) for e in "${NEWS_ENGINES[@]}"; do mark_disabled "$e"; done ;;
|
||||
science) for e in "${SCIENCE_ENGINES[@]}"; do mark_disabled "$e"; done ;;
|
||||
social) for e in "${SOCIAL_ENGINES[@]}"; do mark_disabled "$e"; done ;;
|
||||
"") ;;
|
||||
*) warn "Unknown category '$cat' — valid: videos images news science social" ;;
|
||||
esac
|
||||
done
|
||||
fi
|
||||
|
||||
# Extra engine disables
|
||||
if [[ -n "$DISABLE_ENGINES_EXTRA" ]]; then
|
||||
IFS=',' read -ra _engs <<< "$DISABLE_ENGINES_EXTRA"
|
||||
for e in "${_engs[@]}"; do
|
||||
e="${e#"${e%%[![:space:]]*}"}"; e="${e%"${e##*[![:space:]]}"}"
|
||||
[[ -n "$e" ]] && mark_disabled "$e"
|
||||
done
|
||||
fi
|
||||
|
||||
# ── Preserve existing secret key ─────────────────────────────────────────────
|
||||
SECRET_KEY=$(grep -oP '(?<=secret_key: ")[^"]+' "$SETTINGS" 2>/dev/null || true)
|
||||
[[ -z "$SECRET_KEY" ]] && SECRET_KEY=$(openssl rand -hex 32)
|
||||
|
||||
# ── Build engine override block ───────────────────────────────────────────────
|
||||
ENGINE_BLOCK=""
|
||||
# set +u: iterating empty associative arrays throws "unbound variable" on bash <4.4
|
||||
set +u
|
||||
for eng in "${!DISABLE_MAP[@]}"; do
|
||||
ENGINE_BLOCK+=" - name: ${eng}\n disabled: true\n"
|
||||
done
|
||||
for eng in "${!ENABLE_MAP[@]}"; do
|
||||
ENGINE_BLOCK+=" - name: ${eng}\n disabled: false\n"
|
||||
done
|
||||
set -u
|
||||
|
||||
# ── Write settings.yml ────────────────────────────────────────────────────────
|
||||
{
|
||||
printf 'use_default_settings: true\n'
|
||||
printf 'general:\n instance_name: "Local Search"\n'
|
||||
printf 'server:\n secret_key: "%s"\n limiter: false\n' "$SECRET_KEY"
|
||||
printf 'search:\n safe_search: %d\n default_lang: "en"\n formats: [html, json]\n' "$SAFE_INT"
|
||||
# Lock the safe-search preference so users cannot override it via the UI
|
||||
if [[ "$LEVEL" != "none" ]]; then
|
||||
printf 'preferences:\n lock:\n - safesearch\n'
|
||||
fi
|
||||
if [[ -n "$ENGINE_BLOCK" ]]; then
|
||||
printf 'engines:\n'
|
||||
printf '%b' "$ENGINE_BLOCK"
|
||||
fi
|
||||
} > "$SETTINGS"
|
||||
|
||||
ok "Updated settings.yml (safe_search: $SAFE_INT)"
|
||||
# bash < 4.4: ${#assoc[@]} on an empty declared array throws "unbound variable"
|
||||
# under set -u — disable nounset for the rest of the reporting section
|
||||
set +u
|
||||
if [[ ${#DISABLE_MAP[@]} -gt 0 ]]; then
|
||||
info "Disabled (${#DISABLE_MAP[@]}): $(printf '%s, ' "${!DISABLE_MAP[@]}" | sed 's/, $//')"
|
||||
fi
|
||||
if [[ ${#ENABLE_MAP[@]} -gt 0 ]]; then
|
||||
info "Re-enabled: $(printf '%s, ' "${!ENABLE_MAP[@]}" | sed 's/, $//')"
|
||||
fi
|
||||
|
||||
# ── Update &safesearch= in SEARXNG_QUERY_URL inside docker-compose.yml ────────
|
||||
if [[ -f "$COMPOSE" ]]; then
|
||||
sed -i -E \
|
||||
"s|(SEARXNG_QUERY_URL=http://searxng:[0-9]+/search\?[^&[:space:]]*)(&safesearch=[0-9])?|\1\&safesearch=${SAFE_INT}|g" \
|
||||
"$COMPOSE"
|
||||
ok "Updated SEARXNG_QUERY_URL (&safesearch=${SAFE_INT})"
|
||||
fi
|
||||
|
||||
# ── Restart SearXNG ───────────────────────────────────────────────────────────
|
||||
if docker ps --format '{{.Names}}' 2>/dev/null | grep -q '^searxng$'; then
|
||||
info "Restarting SearXNG..."
|
||||
docker restart searxng
|
||||
ok "SearXNG restarted"
|
||||
else
|
||||
info "SearXNG not running — changes take effect on next start"
|
||||
fi
|
||||
|
||||
echo
|
||||
grn "Done — safe search: $LEVEL"
|
||||
[[ "$LEVEL" != "none" ]] && \
|
||||
info "Engines skipped (can't enforce '$LEVEL'): ${#DISABLE_MAP[@]} total"
|
||||
+239
@@ -0,0 +1,239 @@
|
||||
#!/usr/bin/env bash
|
||||
# configure-storage.sh — interactively assign drives for ZIM and model storage
|
||||
# then patches docker-compose.yml in place
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
info() { echo -e "${CYAN}[..]${NC} $*"; }
|
||||
ok() { echo -e "${GREEN}[OK]${NC} $*"; }
|
||||
warn() { echo -e "${YELLOW}[!!]${NC} $*"; }
|
||||
die() { echo -e "${RED}[ERR]${NC} $*"; exit 1; }
|
||||
|
||||
BASE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
COMPOSE="$BASE/docker-compose.yml"
|
||||
|
||||
[[ -f "$COMPOSE" ]] || die "docker-compose.yml not found at $BASE — run local-ai-setup.sh first"
|
||||
|
||||
# ── collect drives ─────────────────────────────────────────────────────────────
|
||||
echo -e "\n${BOLD}━━━ Detected Drives ━━━${NC}\n"
|
||||
|
||||
# build array of mountpoints excluding tiny/system ones
|
||||
mapfile -t MOUNTS < <(
|
||||
df -h --output=target,size,avail,pcent 2>/dev/null \
|
||||
| tail -n +2 \
|
||||
| grep -v -E '^\s*(/$|/boot|/sys|/proc|/dev|/run|/snap|/var/lib/docker|/etc|tmpfs)' \
|
||||
| awk '$3 ~ /[0-9]/ {print}' \
|
||||
| sort -u
|
||||
)
|
||||
|
||||
if [[ ${#MOUNTS[@]} -eq 0 ]]; then
|
||||
warn "No additional mounted drives found beyond system disk."
|
||||
warn "Mount your drives first (e.g. /mnt/storage), then re-run this script."
|
||||
echo ""
|
||||
echo "Quick mount example:"
|
||||
echo " sudo mkdir -p /mnt/storage"
|
||||
echo " sudo mount /dev/sdb1 /mnt/storage"
|
||||
echo " # To make permanent, add to /etc/fstab"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# display numbered list
|
||||
echo -e " ${BOLD}# Mount Point Size Free Used${NC}"
|
||||
echo " ─────────────────────────────────────────────────────"
|
||||
IDX=0
|
||||
declare -a MOUNT_PATHS
|
||||
for line in "${MOUNTS[@]}"; do
|
||||
mp=$(echo "$line" | awk '{print $1}')
|
||||
sz=$(echo "$line" | awk '{print $2}')
|
||||
av=$(echo "$line" | awk '{print $3}')
|
||||
pc=$(echo "$line" | awk '{print $4}')
|
||||
printf " ${CYAN}%-3d${NC} %-30s %-7s %-7s %s\n" "$IDX" "$mp" "$sz" "$av" "$pc"
|
||||
MOUNT_PATHS[$IDX]="$mp"
|
||||
(( IDX++ ))
|
||||
done
|
||||
echo ""
|
||||
|
||||
pick_drive() {
|
||||
local purpose="$1"
|
||||
local varname="$2"
|
||||
local skip_ok="${3:-true}"
|
||||
|
||||
echo -e "${BOLD}$purpose storage${NC}"
|
||||
if $skip_ok; then
|
||||
read -rp " Enter drive number (or Enter to skip): " choice
|
||||
else
|
||||
read -rp " Enter drive number: " choice
|
||||
fi
|
||||
|
||||
if [[ -z "$choice" ]]; then
|
||||
eval "$varname=''"
|
||||
info "Skipping $purpose storage."
|
||||
elif [[ "$choice" =~ ^[0-9]+$ ]] && [[ "$choice" -lt "$IDX" ]]; then
|
||||
eval "$varname='${MOUNT_PATHS[$choice]}'"
|
||||
ok "Selected: ${MOUNT_PATHS[$choice]}"
|
||||
else
|
||||
warn "Invalid choice — skipping $purpose."
|
||||
eval "$varname=''"
|
||||
fi
|
||||
echo ""
|
||||
}
|
||||
|
||||
# ── ZIM storage ────────────────────────────────────────────────────────────────
|
||||
echo -e "${BOLD}━━━ Kiwix ZIM Storage ━━━${NC}\n"
|
||||
echo " ZIM files can be large (Wikipedia alone is ~100GB)."
|
||||
echo " Pick a drive with plenty of space."
|
||||
echo ""
|
||||
pick_drive "ZIM" ZIM_DRIVE
|
||||
|
||||
ZIM_PATH=""
|
||||
if [[ -n "$ZIM_DRIVE" ]]; then
|
||||
read -rp " Folder name on that drive [zims]: " ZIM_FOLDER
|
||||
ZIM_FOLDER="${ZIM_FOLDER:-zims}"
|
||||
ZIM_PATH="$ZIM_DRIVE/$ZIM_FOLDER"
|
||||
mkdir -p "$ZIM_PATH"
|
||||
ok "ZIM path: $ZIM_PATH"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# ── Model storage ──────────────────────────────────────────────────────────────
|
||||
echo -e "${BOLD}━━━ Ollama Model Storage ━━━${NC}\n"
|
||||
echo " Models range from 4GB (7B) to 30GB+ (70B)."
|
||||
echo " Pick a fast drive (SSD preferred) with 50–200GB free."
|
||||
echo ""
|
||||
pick_drive "Model" MODEL_DRIVE
|
||||
|
||||
MODEL_PATH=""
|
||||
if [[ -n "$MODEL_DRIVE" ]]; then
|
||||
read -rp " Folder name on that drive [ollama-models]: " MODEL_FOLDER
|
||||
MODEL_FOLDER="${MODEL_FOLDER:-ollama-models}"
|
||||
MODEL_PATH="$MODEL_DRIVE/$MODEL_FOLDER"
|
||||
mkdir -p "$MODEL_PATH"
|
||||
ok "Model path: $MODEL_PATH"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# ── bail if nothing selected ───────────────────────────────────────────────────
|
||||
if [[ -z "$ZIM_PATH" && -z "$MODEL_PATH" ]]; then
|
||||
info "Nothing selected — no changes made."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# ── summary + confirm ──────────────────────────────────────────────────────────
|
||||
echo -e "${BOLD}━━━ Planned Changes ━━━${NC}\n"
|
||||
[[ -n "$ZIM_PATH" ]] && echo " Kiwix ZIMs → $ZIM_PATH"
|
||||
[[ -n "$MODEL_PATH" ]] && echo " Ollama models → $MODEL_PATH"
|
||||
echo ""
|
||||
read -rp "Apply these changes to docker-compose.yml? [Y/n]: " CONFIRM
|
||||
[[ "${CONFIRM,,}" == "n" ]] && { info "Aborted — no changes made."; exit 0; }
|
||||
|
||||
# ── backup ────────────────────────────────────────────────────────────────────
|
||||
cp "$COMPOSE" "$COMPOSE.bak"
|
||||
info "Backed up to docker-compose.yml.bak"
|
||||
|
||||
# ── patch kiwix volumes ────────────────────────────────────────────────────────
|
||||
if [[ -n "$ZIM_PATH" ]]; then
|
||||
# Replace the kiwix volumes line to add extra ZIM mounts
|
||||
# Current: volumes: [$BASE/kiwix:/data]
|
||||
# New: volumes:
|
||||
# - $BASE/kiwix:/data
|
||||
# - /mnt/drive/zims:/data/external
|
||||
|
||||
python3 - "$COMPOSE" "$BASE" "$ZIM_PATH" << 'PY'
|
||||
import sys, re
|
||||
|
||||
compose_file = sys.argv[1]
|
||||
base = sys.argv[2]
|
||||
zim_path = sys.argv[3]
|
||||
|
||||
text = open(compose_file).read()
|
||||
|
||||
# Find the kiwix service volumes line (single-line array style)
|
||||
old = f" volumes: [{base}/kiwix:/data]"
|
||||
new = (
|
||||
f" volumes:\n"
|
||||
f" - {base}/kiwix:/data\n"
|
||||
f" - {zim_path}:/data/external"
|
||||
)
|
||||
|
||||
if old in text:
|
||||
text = text.replace(old, new)
|
||||
open(compose_file, "w").write(text)
|
||||
print(f"[OK] Kiwix volumes updated")
|
||||
else:
|
||||
# Already multi-line — append the new mount if not already there
|
||||
marker = "container_name: kiwix"
|
||||
if marker in text and zim_path not in text:
|
||||
# find the volumes block under kiwix and append
|
||||
lines = text.splitlines()
|
||||
in_kiwix = False
|
||||
in_vols = False
|
||||
insert_after = -1
|
||||
for i, line in enumerate(lines):
|
||||
if marker in line:
|
||||
in_kiwix = True
|
||||
if in_kiwix and "volumes:" in line:
|
||||
in_vols = True
|
||||
if in_vols and line.strip().startswith("- ") and ":/data" in line:
|
||||
insert_after = i
|
||||
if in_vols and insert_after > 0 and not line.strip().startswith("- "):
|
||||
break
|
||||
if insert_after > 0:
|
||||
lines.insert(insert_after + 1, f" - {zim_path}:/data/external")
|
||||
open(compose_file, "w").write("\n".join(lines) + "\n")
|
||||
print(f"[OK] Kiwix extra ZIM volume appended")
|
||||
else:
|
||||
print(f"[!!] Could not patch kiwix volumes — edit manually: add - {zim_path}:/data/external")
|
||||
else:
|
||||
print(f"[!!] Kiwix volumes line not found in expected format — edit manually")
|
||||
PY
|
||||
fi
|
||||
|
||||
# ── patch ollama volumes ───────────────────────────────────────────────────────
|
||||
if [[ -n "$MODEL_PATH" ]]; then
|
||||
python3 - "$COMPOSE" "$MODEL_PATH" << 'PY'
|
||||
import sys
|
||||
|
||||
compose_file = sys.argv[1]
|
||||
model_path = sys.argv[2]
|
||||
|
||||
text = open(compose_file).read()
|
||||
|
||||
# Replace Docker-managed volume with host path
|
||||
old = " volumes: [ollama-models:/root/.ollama]"
|
||||
new = f" volumes: [{model_path}:/root/.ollama]"
|
||||
|
||||
if old in text:
|
||||
text = text.replace(old, new)
|
||||
# Remove ollama-models from the top-level volumes section if present
|
||||
text = text.replace(" ollama-models:\n", "")
|
||||
open(compose_file, "w").write(text)
|
||||
print(f"[OK] Ollama volume → {model_path}")
|
||||
else:
|
||||
print(f"[!!] ollama volumes line not in expected format — edit manually:")
|
||||
print(f" change ollama-models:/root/.ollama to {model_path}:/root/.ollama")
|
||||
PY
|
||||
fi
|
||||
|
||||
# ── restart if running ─────────────────────────────────────────────────────────
|
||||
echo ""
|
||||
if docker compose -f "$COMPOSE" ps --quiet 2>/dev/null | grep -q .; then
|
||||
read -rp "Stack is running — restart now to apply changes? [Y/n]: " RESTART
|
||||
if [[ "${RESTART,,}" != "n" ]]; then
|
||||
cd "$BASE"
|
||||
docker compose down
|
||||
docker compose up -d
|
||||
ok "Stack restarted"
|
||||
else
|
||||
warn "Remember to restart: cd $BASE && docker compose down && docker compose up -d"
|
||||
fi
|
||||
else
|
||||
info "Stack not running — changes will apply on next start."
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo -e "${GREEN}${BOLD}Done!${NC}"
|
||||
[[ -n "$ZIM_PATH" ]] && echo " Drop .zim files into: $ZIM_PATH"
|
||||
[[ -n "$MODEL_PATH" ]] && echo " Ollama models stored in: $MODEL_PATH"
|
||||
echo ""
|
||||
echo " Undo: cp $COMPOSE.bak $COMPOSE"
|
||||
+868
@@ -0,0 +1,868 @@
|
||||
# GPU Setup Research: Rack Server AI Workloads
|
||||
|
||||
*Last updated: March 22, 2026*
|
||||
|
||||
## Goal
|
||||
Cost-efficient rack-mountable GPU setup for:
|
||||
1. **LLM coding inference** — Run 32B+ parameter coding models with maximum context windows
|
||||
2. **Image generation** — ComfyUI / InvokeAI with Stable Diffusion SDXL / Flux
|
||||
|
||||
Target servers: Dell R720/R730 or HP DL380 equivalent (2U rack)
|
||||
|
||||
## Why 48GB VRAM is the Right Target
|
||||
|
||||
### The Problem with 24GB
|
||||
32B coding models at Q4_K_M quantization use ~20GB of weights, leaving only ~4GB for KV cache
|
||||
on a 24GB card. This severely limits context window size — the key ingredient for complex
|
||||
coding sessions where the model needs to understand your entire codebase.
|
||||
|
||||
### What 48GB Unlocks
|
||||
- **32B models at higher quantization** (Q6_K/Q8_0) = better output quality
|
||||
- **28GB+ free for KV cache** = massive context windows (32K+ tokens)
|
||||
- **70B models** in aggressive quantization (~12 t/s but functional)
|
||||
- **Simultaneous model loading** — coding model + image gen model at once
|
||||
- Room for future larger models without hardware changes
|
||||
|
||||
## Best Local Coding Models (2026)
|
||||
|
||||
### Qwen 3.5 Family (February 2026 — Gated Delta Networks)
|
||||
|
||||
Architecture breakthrough: 3 of every 4 layers use **linear attention** (O(n) scaling),
|
||||
drastically reducing KV cache memory. These models need far less VRAM for long contexts
|
||||
than traditional transformers.
|
||||
|
||||
| Model | Type | Active Params | Size at Q4_K_M | Max Context | Quality | Notes |
|
||||
|-------|------|---------------|---------------|-------------|---------|-------|
|
||||
| **Qwen3.5-35B-A3B** | **MoE** | **3B** | **~12GB** | **262K** | **B+ to A-** | 35B total but only 3B active — quality tracks active params |
|
||||
| **Qwen3.5-27B** | Dense | 27B | ~17GB | 262K | **A-** | 72.4% SWE-bench, ties GPT-5 mini. The real A- option. |
|
||||
| **Qwen3.5-122B-A10B** | MoE | 10B | ~76GB | 262K | A | Matches GPT-5 mini across the board |
|
||||
| **Qwen3.5-9B** | Dense | 9B | ~6GB | 262K | B+ | Fits on any modern GPU |
|
||||
| **Qwen3.5-4B** | Dense | 4B | ~3GB | 262K | B | Tiny but capable |
|
||||
|
||||
**Quality reality check:** MoE models route tokens through only a subset of parameters.
|
||||
The 35B-A3B activates **3B params per token** — think of it as a smart 7B model, not a 35B.
|
||||
Quality is closer to B+ for complex coding. The 27B dense model is genuinely A- but needs
|
||||
17GB weights (leaving less room for context on 32GB). At Q4 quantization there's a further
|
||||
small quality loss. And 262K is a VRAM ceiling, not a quality guarantee — models degrade
|
||||
at the edges of their context window. Practical high-quality context is more like 64-128K.
|
||||
|
||||
**No local model approaches Claude Opus on hard problems.** The strategy isn't to replace
|
||||
Opus — it's to offload the 80% of routine work so your Pro plan limits stop being an issue.
|
||||
|
||||
### Previous Generation (Still Relevant)
|
||||
|
||||
| Model | Size at Q4_K_M | Quality | Notes |
|
||||
|-------|---------------|---------|-------|
|
||||
| **Qwen2.5-Coder 32B** | ~20GB | 73.7 Aider (≈ GPT-4o) | FIM king, 92.7% HumanEval |
|
||||
| **Qwen3-Coder 30B-A3B** (MoE) | ~18GB | #1 SWE-rebench (64.6%) | Only 3.3B active, very fast |
|
||||
| **Qwen3-Coder-Next 80B** (MoE) | needs 64GB+ RAM offload | Beats Claude Opus 4.6 on SWE-rebench | Hybrid attention, 256K context |
|
||||
|
||||
### Honest Assessment: Local vs Claude Code
|
||||
|
||||
Nothing local approaches Claude Opus 4.6 quality for complex multi-file agentic coding.
|
||||
These 32B models are competitive with **GPT-4o** — a tier below Claude Sonnet, two tiers below Opus.
|
||||
|
||||
**The real strategy: Drop Max ($100/mo), keep Pro ($20/mo), offload bulk work to local.**
|
||||
|
||||
The problem with Pro for large projects: rate limits. A 10,000-line codebase needs the model
|
||||
to read, understand, and hold context across many files. On Pro you'll hit usage caps mid-session
|
||||
on complex multi-file work. Max ($100/mo) removes those limits — but that's $80/mo extra.
|
||||
|
||||
Local AI eliminates this problem differently:
|
||||
- **Local model (262K context)**: Reads your entire 10K-line project at once. No rate limits,
|
||||
no usage caps, runs 24/7. Handles the bulk work — understanding codebase structure, routine
|
||||
bug fixes, simple refactors, code explanation, test writing, boilerplate generation.
|
||||
- **Claude Pro ($20/mo)**: Reserved for the hard problems — complex multi-file architectural
|
||||
changes, subtle bugs that need Opus-level reasoning, code review on critical paths.
|
||||
Pro limits are fine when you're only sending Claude the *hard* 20% instead of everything.
|
||||
|
||||
This is the unlock: local doesn't replace Claude, it **reduces your Claude usage enough
|
||||
that Pro limits stop being a problem.** The 80% of routine work that was burning through
|
||||
your Max quota now runs locally with zero limits.
|
||||
|
||||
| Plan | Monthly | What You Get | Limit Problem |
|
||||
|------|---------|--------------|---------------|
|
||||
| Max only | $100 | Opus unlimited | Paying $80/mo for unlimited when you don't need it |
|
||||
| Pro only | $20 | Opus with rate limits | **Hits caps on 10K-line projects** |
|
||||
| **Pro + Local GPU** | **$29** | Opus for hard stuff + unlimited local | **No caps — bulk work is local** |
|
||||
| Local only (no Claude) | $9 | A- quality only | Stuck on hard problems with no escape hatch |
|
||||
|
||||
## 48GB GPU Market (March 22, 2026 — Real Prices)
|
||||
|
||||
| GPU | Arch | Used Price | TDP | Cooling | Tensor Cores | Mem BW |
|
||||
|-----|------|------------|-----|---------|--------------|--------|
|
||||
| **Quadro RTX 8000** | Turing (2018) | **$2,000–2,900** | 260W | Passive variant | Yes (576) | 672 GB/s |
|
||||
| **A40** | Ampere (2020) | **~$5,050+** | 300W | Passive | Yes (336 3rd-gen) | 696 GB/s |
|
||||
| **RTX A6000** | Ampere (2020) | **~$5,400+** | 300W | Active (blower) | Yes (336 3rd-gen) | 768 GB/s |
|
||||
| **L40** | Ada (2022) | **~$6,500+** | 300W | Passive | Yes (568 4th-gen) | 864 GB/s |
|
||||
| **RTX 6000 Ada** | Ada (2022) | **~$6,500+** | 300W | Active | Yes (568 4th-gen) | 960 GB/s |
|
||||
|
||||
Sources: eBay active/sold listings, GPUPoet price tracking, Pangoly, CamelCamelCamel (all March 2026)
|
||||
Note: One outlier RTX 8000 listing at ~$750 exists but is not representative of the market.
|
||||
|
||||
### Cheapest 48GB Option: Quadro RTX 8000 Passive ($2,000–2,900)
|
||||
|
||||
The RTX 8000 is still the cheapest 48GB card — roughly half the price of an A40 and a
|
||||
third of an A6000. The passive variant is purpose-built for rack servers — no fan, relies
|
||||
on chassis airflow, designed for 24/7 operation in 2U/4U systems.
|
||||
|
||||
Key advantages over the P40:
|
||||
- **48GB vs 24GB** — room for models + massive context
|
||||
- **Has Tensor Cores** (576 Turing) — native FP16, no `--force-fp32` hacks for image gen
|
||||
- **NVLink support** — pair two for 96GB combined (100 GB/s bidirectional)
|
||||
- 10W idle power draw
|
||||
|
||||
### Cost Reality Check
|
||||
At $2,000–2,900 the RTX 8000 is a significant investment. The key question: is unified
|
||||
48GB VRAM worth 4–6x the cost of dual P40s ($400–500)?
|
||||
|
||||
**Yes, if** you need large context windows (32K+) for complex coding — KV cache can't
|
||||
be split across two GPUs without NVLink (which P40s don't have).
|
||||
|
||||
**No, if** you're mostly doing short-prompt coding tasks and image gen — dual P40s give
|
||||
you 48GB total (split) at a fraction of the cost, and each card can handle its own workload.
|
||||
|
||||
## RTX 8000 Performance Benchmarks
|
||||
|
||||
### LLM Inference (Exllama, 5.0 bpw quantization)
|
||||
|
||||
| Model | Context | Prompt Processing | Generation |
|
||||
|-------|---------|-------------------|------------|
|
||||
| Qwen3 30B-A3B (MoE) | 8K | 950 t/s | **34 t/s** |
|
||||
| Qwen3 30B-A3B (MoE) | 16K | 673 t/s | **21 t/s** |
|
||||
| Qwen3 30B-A3B (MoE) | 32K | 345 t/s | **11 t/s** |
|
||||
| Llama 3.3 70B | short | 36 t/s | **13 t/s** |
|
||||
| Llama 3.1 8B | — | — | **72 t/s** |
|
||||
|
||||
### Compared to P40 (24GB)
|
||||
|
||||
| Metric | P40 (24GB) | RTX 8000 (48GB) |
|
||||
|--------|-----------|-----------------|
|
||||
| **Used price** | **$150–320** | **$2,000–2,900** |
|
||||
| 32B model fit | Barely (~2GB free) | Comfortable (~28GB free) |
|
||||
| 32B generation speed | ~5-12 t/s (est.) | ~20-34 t/s |
|
||||
| Max practical context | ~4K tokens | **32K+ tokens** |
|
||||
| Image gen (SDXL) | ~49s (`--force-fp32`) | Faster (native FP16) |
|
||||
| Rack server ready | Yes (passive) | Yes (passive variant) |
|
||||
|
||||
### Image Generation
|
||||
The RTX 8000 has Turing Tensor Cores with native FP16 support. Unlike the P40, it does NOT
|
||||
need `--force-fp32` workarounds. Image gen performance is significantly better than the P40,
|
||||
though still behind Ampere/Ada cards.
|
||||
|
||||
## Budget Build: 2x Quadro RTX 5000 + NVLink ($850 Total)
|
||||
|
||||
*The best price-to-capability ratio for local AI coding in 2026.*
|
||||
|
||||
### Why This Works Now
|
||||
|
||||
Qwen 3.5 (February 2026) introduced **Gated Delta Networks** — 3 out of 4 layers use linear
|
||||
attention (O(n) scaling) instead of quadratic. KV cache memory usage is dramatically lower
|
||||
than traditional transformers. A 35B MoE model with 262K context now fits in ~25GB VRAM.
|
||||
|
||||
### Hardware
|
||||
|
||||
#### GPU: NVIDIA Quadro RTX 5000 (Turing, TU104)
|
||||
|
||||
| Spec | Value |
|
||||
|------|-------|
|
||||
| VRAM | 16GB GDDR6 |
|
||||
| CUDA Cores | 3072 |
|
||||
| Tensor Cores | 384 (Gen 2, FP16) |
|
||||
| TDP | ~230W |
|
||||
| NVLink | **Yes — 50 GB/s bidirectional** |
|
||||
| Form Factor | Dual-slot, blower cooler (rack-friendly) |
|
||||
| PCIe | 3.0 x16 |
|
||||
| Used Price | **~$400** |
|
||||
| Part Number | VCQRTX5000-PB |
|
||||
|
||||
#### NVLink Bridge (CRITICAL: RTX 5000 uses a unique smaller connector)
|
||||
|
||||
The Quadro RTX 5000 has a **shorter NVLink connector** than all other Quadro RTX cards.
|
||||
Bridges from the RTX 6000/8000 will NOT physically fit. You must buy the RTX 5000-specific bridge.
|
||||
|
||||
| Detail | Value |
|
||||
|--------|-------|
|
||||
| Product | NVIDIA Quadro RTX 5000 NVLink HB Bridge 2-Slot |
|
||||
| SKU | NVLINKX8-2SLOT-PB |
|
||||
| Part Numbers | 1JF3K, 699-54934-0500-000, 900-54934-0100-000, P4934, 6FY12AA, L55997-001 |
|
||||
| Price | **~$30-80** (eBay, Amazon) |
|
||||
| Bandwidth | 50 GB/s total (25 GB/s per direction) |
|
||||
| Sizing | 2-slot (cards adjacent) or 3-slot (one slot gap — better thermals) |
|
||||
|
||||
**Where to buy:**
|
||||
- eBay: search "Quadro RTX 5000 NVLink" or part numbers P4934 / L55997-001 / 1JF3K
|
||||
- Amazon: search part number 6FY12AA or 1JF3K
|
||||
|
||||
**WARNING:** The 3-slot bridge is recommended over 2-slot. With a 2-slot bridge the cards
|
||||
sit directly adjacent — the top card's blower intake gets blocked by the bottom card.
|
||||
A 3-slot bridge leaves an air gap for proper cooling.
|
||||
|
||||
#### Motherboard Requirements
|
||||
|
||||
| Requirement | Details |
|
||||
|-------------|---------|
|
||||
| PCIe slots | Two x16 slots (x8 electrical is fine — LLM inference is VRAM-bound, not PCIe-bound) |
|
||||
| Slot spacing | Must match your NVLink bridge size (2-slot or 3-slot gap) |
|
||||
| Power supply | 650W+ minimum (80 PLUS Gold recommended), 850W+ for headroom |
|
||||
| Power connectors | 2x 8-pin PCIe power (one per card). Do NOT daisy-chain — use separate cables |
|
||||
| CPU platform | Any modern platform works. Threadripper/Xeon not required |
|
||||
|
||||
**Recommended motherboards (workstation/server):**
|
||||
- Any board with 2x PCIe x16 slots spaced 2-3 slots apart
|
||||
- Server: Dell R730/R740 with GPU riser (but verify 3-slot bridge clearance in 2U)
|
||||
- Workstation: MSI X399 Creation, ASUS WS series, Supermicro X11/X12 boards
|
||||
- Desktop: Most ATX boards with 2 full-length x16 slots work
|
||||
|
||||
**Rack server note:** The Quadro RTX 5000's blower cooler exhausts out the bracket —
|
||||
this works well in rack airflow. If using a 2U server, measure clearance for the NVLink
|
||||
bridge sitting on top of the cards. A 4U chassis gives the most room.
|
||||
|
||||
### What Runs on 16GB (Single RTX 5000 — Start Here)
|
||||
|
||||
| Model | Quant | Context | Quality | Notes |
|
||||
|-------|-------|---------|---------|-------|
|
||||
| **Qwen3.5-35B-A3B** | Q4_K_L | ~64-128K | **B+** | MoE, 3B active. Good but VRAM is tight — context may be lower |
|
||||
| Qwen3.5-9B | Q8 | 128K+ | B+ | Fits comfortably, high quant |
|
||||
| Qwen3.5-4B | Q8 | 262K | B | Tiny model, long context |
|
||||
| Qwen2.5-Coder-7B | Q8 | 128K | B | Solid for simple tasks |
|
||||
| Qwen2.5-Coder-14B | Q4_K_M | 16-32K | B+ | Tight fit, limited context |
|
||||
|
||||
A single card is a solid start — B+ coding with decent context. But 16GB is the ceiling.
|
||||
You can't run bigger dense models, can't use higher quantization, and context is squeezed.
|
||||
|
||||
### What the Second Card + NVLink Unlocks (32GB)
|
||||
|
||||
The second card doesn't just double context — it opens models that **don't fit on 16GB at all:**
|
||||
|
||||
| Model | Arch | Quant | Weights | Context | Total VRAM | Quality | **Why it needs 32GB** |
|
||||
|-------|------|-------|---------|---------|------------|---------|----------------------|
|
||||
| **Qwen3.5-27B** | **Dense** | **Q4_K_M** | **~17GB** | **128K+** | **~25GB** | **A-** | **17GB weights won't fit on 16GB** |
|
||||
| **Qwen2.5-Coder-32B** | **Dense** | **Q4_K_M** | **~20GB** | **16-24K** | **~28GB** | **A-** | **20GB weights won't fit on 16GB** |
|
||||
| Qwen2.5-Coder-14B | Dense | **Q8** | ~16GB | 64K | ~28GB | A- | Q8 quant = better output, needs 16GB for weights alone |
|
||||
| Qwen3-Coder-Next (80B) | MoE | Q4 | ~20GB | 128K | ~28GB | A | 20GB weights won't fit on 16GB |
|
||||
| Qwen3.5-35B-A3B | MoE (3B active) | Q4_K_M | ~12GB | 262K | ~25GB | B+ | Fits on 1 card at reduced context, but 32GB = full 262K + headroom |
|
||||
|
||||
**The real upgrade isn't 262K context — it's access to dense 27B/32B models that are
|
||||
genuinely A- quality.** The 35B-A3B MoE runs on both setups, but its 3B active params
|
||||
limit quality. The Qwen3.5-27B dense model uses all 27B params on every token — that's
|
||||
the quality jump. And its 17GB of weights physically can't fit on a single 16GB card.
|
||||
|
||||
Think of it this way:
|
||||
- **1 card**: B+ coding (MoE or small dense models, squeezed context)
|
||||
- **2 cards**: **A- coding** (full dense 27B/32B models, comfortable context, higher quant options)
|
||||
|
||||
### Practical Context Windows (Usability, Not Ceilings)
|
||||
|
||||
Context window "support" is a ceiling, not what you actually get. VRAM must hold both the
|
||||
model weights AND the KV cache. What's left after weights determines your real context.
|
||||
Quality also degrades toward the edges of a model's context window.
|
||||
|
||||
**Reference: A 10,000-line codebase ≈ 100-150K tokens** (varies by language/comments).
|
||||
This Claude Opus session uses a **1 million token** context window for comparison.
|
||||
|
||||
#### 1 Card (16GB) — Practical
|
||||
|
||||
| Model | Weights | Free for KV | **Usable context** | 10K-line project? |
|
||||
|-------|---------|-------------|-------------------|-------------------|
|
||||
| Qwen3.5-35B-A3B (MoE) | ~12GB | ~3GB | **32-50K tokens** | **No — ~1/3 of it** |
|
||||
| Qwen3.5-9B (dense) | ~6GB | ~9GB | **80-100K tokens** | **Mostly — but B+ quality** |
|
||||
| Qwen2.5-Coder-14B | ~10GB | ~5GB | **16-24K tokens** | **No — a few files at a time** |
|
||||
|
||||
**Workflow on 1 card:** You're feeding files in chunks. Good for "fix this function" or
|
||||
"explain this file." Not for "read my whole project and refactor the auth system."
|
||||
|
||||
#### 2 Cards (32GB via NVLink) — Practical
|
||||
|
||||
| Model | Weights | Free for KV | **Usable context** | 10K-line project? |
|
||||
|-------|---------|-------------|-------------------|-------------------|
|
||||
| **Qwen3.5-27B (dense)** | ~17GB | ~14GB | **80-128K tokens** | **Yes — most/all of it at A-** |
|
||||
| Qwen3.5-35B-A3B (MoE) | ~12GB | ~19GB | **128-180K tokens** | **Yes with room to spare (B+)** |
|
||||
| Qwen2.5-Coder-32B | ~20GB | ~11GB | **32-48K tokens** | **Partial — but strong A- on what it sees** |
|
||||
|
||||
**Workflow on 2 cards:** You can dump most/all of a 10K-line project in one shot with the
|
||||
27B dense model. That's the real workflow change — "here's my whole project, find the bug"
|
||||
becomes possible locally.
|
||||
|
||||
#### vs This Claude Session
|
||||
|
||||
| Setup | Usable context | vs Opus 1M | Whole-project workflow? |
|
||||
|-------|---------------|------------|----------------------|
|
||||
| 1x RTX 5000 (best) | ~50-100K | 5-10% | No — file by file |
|
||||
| **2x RTX 5000 (best)** | **~128-180K** | **13-18%** | **Yes — for 10K-line projects** |
|
||||
| Claude Opus (this session) | 1,000K | 100% | Yes — for anything |
|
||||
|
||||
**Neither setup replaces this session** for complex multi-file work across a 50K+ line
|
||||
codebase. That's why you keep Pro. But 2 cards handles the daily "read my project and
|
||||
help me code" workflow locally with no rate limits — and that's 80% of the work.
|
||||
|
||||
### Squeezing Every Byte: Single-Card Optimization (16GB)
|
||||
|
||||
Before buying a second card, stack these techniques. They're cumulative — use all of them
|
||||
together. The gains compound because they all free VRAM from the same bottleneck: KV cache.
|
||||
|
||||
#### 1. Quantize the KV Cache (Biggest Single Win)
|
||||
|
||||
By default, llama.cpp stores the KV cache in FP16. That's 2 bytes per value. You can
|
||||
compress it with zero code changes — just flags:
|
||||
|
||||
| Cache Type | Bytes/value | vs FP16 | Quality Impact | Verdict |
|
||||
|-----------|-------------|---------|----------------|---------|
|
||||
| FP16 (default) | 2.0 | baseline | none | wasteful on 16GB |
|
||||
| **Q8_0** | **1.0** | **50% smaller** | **~0.002-0.05 perplexity** | **Always use this** |
|
||||
| Q4_0 | 0.5 | 75% smaller | ~0.2 perplexity (noticeable) | Use if desperate |
|
||||
| **Asymmetric: K=Q8_0, V=Q4_0** | **0.75 avg** | **62% smaller** | **Better than uniform Q4** | **Best bang/buck** |
|
||||
|
||||
The K cache is more sensitive to quantization than V. Asymmetric (Q8 keys, Q4 values) gives
|
||||
you ~62% savings with quality closer to Q8 than Q4.
|
||||
|
||||
**Concrete example — Qwen3.5-35B-A3B on 1 card (16GB):**
|
||||
- Weights: ~12GB → 4GB free for KV cache
|
||||
- FP16 KV cache: 4GB → **~50K context**
|
||||
- Q8_0 KV cache: 4GB buys 2x → **~100K context**
|
||||
- K=Q8/V=Q4 KV cache: 4GB buys 2.6x → **~130K context**
|
||||
|
||||
That's the difference between "a few files" and "a meaningful chunk of a project."
|
||||
|
||||
```bash
|
||||
# llama.cpp — always use these three flags together
|
||||
llama-server \
|
||||
--cache-type-k q8_0 \
|
||||
--cache-type-v q4_0 \
|
||||
--flash-attn \
|
||||
-m model.gguf -ngl 99 -c 131072
|
||||
|
||||
# Ollama — set environment variable before starting
|
||||
export OLLAMA_KV_CACHE_TYPE=q8_0 # or q4_0 for aggressive
|
||||
export OLLAMA_FLASH_ATTENTION=1
|
||||
ollama serve
|
||||
```
|
||||
|
||||
#### 2. Flash Attention (Free Speed + VRAM)
|
||||
|
||||
Flash attention restructures how attention is computed — instead of materializing the full
|
||||
attention matrix in VRAM, it computes it in tiles. Result: less VRAM used during inference,
|
||||
slightly faster, **zero quality loss**.
|
||||
|
||||
Always enable it. There's no downside on Turing GPUs with quantized KV cache.
|
||||
|
||||
```bash
|
||||
# llama.cpp
|
||||
--flash-attn
|
||||
|
||||
# Ollama
|
||||
export OLLAMA_FLASH_ATTENTION=1
|
||||
```
|
||||
|
||||
#### 3. Host-Memory Prompt Caching (`--cram`) — System RAM as L2 Cache
|
||||
|
||||
This is the smart use of system RAM. The `--cram` flag in llama-server stores pre-computed
|
||||
prompt representations in host memory (system RAM). When you send the same system prompt
|
||||
or reuse a conversation prefix, it skips reprocessing — hot-swaps the cached computation
|
||||
back onto the GPU.
|
||||
|
||||
This doesn't increase context window size, but it **dramatically reduces time-to-first-token**
|
||||
for repeated workflows (which is most coding — same system prompt, same project context).
|
||||
|
||||
```bash
|
||||
# llama-server with 16GB RAM cache for prompts
|
||||
llama-server \
|
||||
--cram 16384 \
|
||||
--cache-type-k q8_0 --cache-type-v q4_0 --flash-attn \
|
||||
-m model.gguf -ngl 99 -c 131072
|
||||
```
|
||||
|
||||
Your R720/R730 has 128-384GB of DDR3/DDR4 RAM. Use it. `--cram 65536` (64GB) is reasonable
|
||||
for a dedicated inference server — it costs nothing, and repeat prompts become near-instant.
|
||||
|
||||
#### 4. KV Cache to System RAM (`-nkvo`) — Last Resort for Context
|
||||
|
||||
The `-nkvo` (no KV offload) flag moves the entire KV cache to system RAM, freeing all 16GB
|
||||
of VRAM for model weights. This sounds great but comes with a brutal speed penalty:
|
||||
|
||||
| Scenario | Speed Impact |
|
||||
|----------|-------------|
|
||||
| Full VRAM (normal) | Baseline (25-35 tok/s) |
|
||||
| KV in system RAM via PCIe | **5-20x slower** (~2-7 tok/s) |
|
||||
| KV on NVMe via mmap | **30x+ slower** (~1 tok/s) |
|
||||
|
||||
**When it makes sense:** Loading a model that barely doesn't fit (e.g., Qwen3.5-27B dense
|
||||
at 17GB weights on a 16GB card). You'd get ~2-5 tok/s with KV in RAM — painfully slow, but
|
||||
it's the difference between "runs slowly" and "doesn't run at all." Fine for a batch job
|
||||
where you walk away and come back. Not viable for interactive coding.
|
||||
|
||||
**Don't do this routinely.** Quantized KV cache (technique #1) is 50-100x better because
|
||||
the cache stays on the GPU. Only use `-nkvo` for models that literally can't fit otherwise.
|
||||
|
||||
#### 5. Pick the Right Architecture (GQA + MoE = VRAM Efficient)
|
||||
|
||||
Not all models consume KV cache equally. Modern architectures with **Grouped Query Attention
|
||||
(GQA)** use far less KV cache than older Multi-Head Attention (MHA):
|
||||
|
||||
| Architecture | KV cache at 64K context | Examples |
|
||||
|-------------|------------------------|---------|
|
||||
| MHA (old) | ~8-12GB | LLaMA-1, GPT-J |
|
||||
| **GQA (modern)** | **~1-3GB** | **Qwen3.5 series, LLaMA-3** |
|
||||
| **GQA + MoE** | **~1.2GB** | **Qwen3.5-35B-A3B** |
|
||||
|
||||
The Qwen3.5-35B-A3B is almost purpose-built for your situation: 3B active params (fast on
|
||||
Turing), MoE architecture (small memory footprint during inference), and GQA (tiny KV cache).
|
||||
With quantized KV on top of that, 130K+ context on a single 16GB card is realistic.
|
||||
|
||||
#### 6. NVMe as mmap Backing Store
|
||||
|
||||
Your fast NVMe matters for **model loading**, not inference. llama.cpp uses mmap by default
|
||||
to stream model weights from disk, so a fast NVMe means:
|
||||
- Near-instant cold starts (weights stream in as needed)
|
||||
- Graceful degradation if model slightly exceeds RAM (OS pages out unused layers)
|
||||
|
||||
But NVMe is **not** a viable substitute for VRAM during inference. The bandwidth gap is too
|
||||
large: VRAM runs at ~400 GB/s (RTX 5000), system RAM at ~50-100 GB/s (DDR4 quad-channel),
|
||||
NVMe at ~3-7 GB/s. Three orders of magnitude difference from VRAM.
|
||||
|
||||
**Practical use:** Keep all your GGUF model files on NVMe. Enable mmap (default). That's it.
|
||||
Don't try to use NVMe as overflow for the KV cache — the latency kills interactive use.
|
||||
|
||||
#### Stacking Everything: Revised Single-Card Numbers
|
||||
|
||||
| Model | Optimization | Usable Context | Speed | Quality |
|
||||
|-------|-------------|---------------|-------|---------|
|
||||
| Qwen3.5-35B-A3B Q4 | None (defaults) | ~32-50K | 25-35 tok/s | B+ |
|
||||
| Qwen3.5-35B-A3B Q4 | **KV Q8 + flash** | **~80-100K** | **25-35 tok/s** | **B+** |
|
||||
| Qwen3.5-35B-A3B Q4 | **KV asym + flash** | **~100-130K** | **25-35 tok/s** | **B+ (tiny quality dip)** |
|
||||
| Qwen3.5-9B Q8 | KV Q8 + flash | ~120-160K | 30-45 tok/s | B |
|
||||
| Qwen2.5-Coder-14B Q4 | KV Q8 + flash | ~40-64K | 20-30 tok/s | B+ |
|
||||
|
||||
Add `--cram` on top for instant repeated prompts. That's your real single-card ceiling.
|
||||
|
||||
**The honest answer:** With all optimizations stacked, a single card goes from "a few files
|
||||
at a time" to "maybe half a 10K-line project." That's a meaningful upgrade from the
|
||||
unoptimized baseline, but it still doesn't match what 2 cards with a dense 27B model gives
|
||||
you. The second card isn't about optimization tricks — it's about physics (more VRAM = more
|
||||
data on the fast bus).
|
||||
|
||||
### Estimated Inference Speed
|
||||
|
||||
| Model | 1x RTX 5000 | 2x RTX 5000 (NVLink) |
|
||||
|-------|-------------|---------------------|
|
||||
| Qwen3.5-35B-A3B Q4 (short ctx) | ~25-35 tok/s | ~25-35 tok/s |
|
||||
| Qwen3.5-35B-A3B Q4 (128K ctx) | ~10-18 tok/s | ~15-25 tok/s |
|
||||
| Qwen3.5-35B-A3B Q4 (262K ctx) | Won't fit | ~10-18 tok/s |
|
||||
| Qwen2.5-Coder-14B Q4 | ~20-30 tok/s | ~25-35 tok/s |
|
||||
|
||||
NVLink matters most at large context windows where KV cache spans both cards.
|
||||
At short contexts that fit on one card, the second GPU adds less benefit.
|
||||
|
||||
**Speed reality check:** NVLink doesn't make it faster — it prevents the slowdown you'd get
|
||||
from PCIe when the model spans both cards. The base speed is still Turing (2018 silicon).
|
||||
10-35 tok/s is fast enough for coding (you read slower than that), but it's not instant.
|
||||
The MoE architecture (only 3B active params at inference) is what makes it viable on older
|
||||
hardware — NVLink just removes the inter-GPU bottleneck for 262K context.
|
||||
|
||||
### Image Generation (Included — No Extra Cost)
|
||||
|
||||
The RTX 5000 has **384 Tensor Cores with native FP16** — full SDXL/Flux support, no hacks.
|
||||
|
||||
| Workload | VRAM Needed | Where It Runs |
|
||||
|----------|-------------|---------------|
|
||||
| SDXL (1024x1024) | ~8-10GB | Either card alone |
|
||||
| Flux Dev | ~12-14GB | Single card (16GB) |
|
||||
| Flux Dev (high-res / batched) | ~18-24GB | Both cards via NVLink (32GB) |
|
||||
| ComfyUI / InvokeAI | Works natively | No `--force-fp32` needed |
|
||||
|
||||
**Important: 262K context requires both cards unified.** You can't split one off for image
|
||||
gen and keep 262K. It's one task at a time:
|
||||
|
||||
```bash
|
||||
# CODING SESSION: Both cards unified → 32GB → 262K context
|
||||
ollama run qwen3.5:35b-a3b-q4_K_M # Uses both GPUs via NVLink
|
||||
|
||||
# IMAGE GEN SESSION: Stop LLM, run image gen on one card (16GB is plenty)
|
||||
ollama stop # Frees VRAM
|
||||
comfyui --listen 0.0.0.0 # SDXL/Flux fits easily in 16GB
|
||||
|
||||
# Swap takes a few seconds, not simultaneous but not painful
|
||||
```
|
||||
|
||||
If you want simultaneous coding + image gen, you'd run a smaller model at shorter context
|
||||
on one card (e.g., Qwen3.5-35B-A3B at ~64K on 16GB) and image gen on the other. But for
|
||||
full 262K context, both cards must be dedicated to the LLM.
|
||||
|
||||
### Hardware Longevity: 3-5 Years Realistic
|
||||
|
||||
- **2026-2027**: Sweet spot. MoE + linear attention models are getting smaller active params.
|
||||
32GB unified handles the best coding models at full context. Peak value.
|
||||
- **2028-2029**: Still useful. The trend is more efficient models, not bigger ones.
|
||||
32GB likely still runs the best ~35-70B MoE coding models of that era.
|
||||
- **2030+**: Questionable. New architectures may need FP8, newer tensor core ops that
|
||||
Turing lacks. But VRAM is VRAM — something useful will always run on 32GB.
|
||||
- **The cards themselves won't die** — Quadro-grade, designed for 24/7 data center use.
|
||||
They'll be outclassed before they fail.
|
||||
|
||||
### Power Consumption & Cost
|
||||
|
||||
| Config | Idle | Load | Monthly (8hr/day @ $0.09/kWh) | Annual |
|
||||
|--------|------|------|-------------------------------|--------|
|
||||
| 1x Quadro RTX 5000 | ~15W | ~210W | **~$4.50** | ~$54 |
|
||||
| 2x Quadro RTX 5000 | ~30W | ~420W | **~$9.00** | ~$108 |
|
||||
|
||||
### Software Setup
|
||||
|
||||
#### llama.cpp (Recommended — Best Multi-GPU Support)
|
||||
|
||||
```bash
|
||||
# Build with CUDA support
|
||||
git clone https://github.com/ggerganov/llama.cpp
|
||||
cd llama.cpp
|
||||
cmake -B build -DGGML_CUDA=ON
|
||||
cmake --build build --config Release -j$(nproc)
|
||||
|
||||
# Download Qwen3.5-35B-A3B GGUF (Q4_K_M)
|
||||
# Get from: https://huggingface.co/unsloth/Qwen3.5-35B-A3B-GGUF
|
||||
|
||||
# Run on dual GPU with NVLink (all optimizations on)
|
||||
./build/bin/llama-server \
|
||||
-m Qwen3.5-35B-A3B-Q4_K_M.gguf \
|
||||
-ngl 999 \
|
||||
-c 262144 \
|
||||
--cache-type-k q8_0 \
|
||||
--cache-type-v q4_0 \
|
||||
--flash-attn \
|
||||
--cram 65536 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8080
|
||||
|
||||
# --cache-type-k q8_0 / --cache-type-v q4_0 = asymmetric KV quantization (62% smaller cache)
|
||||
# --flash-attn = tiled attention (less VRAM, no quality loss)
|
||||
# --cram 65536 = 64GB host RAM prompt cache (instant repeat prompts)
|
||||
# llama.cpp auto-detects NVLink and splits layers across both GPUs
|
||||
# Use -ts 1,1 to manually set equal split if needed
|
||||
|
||||
# Single card variant (no NVLink) — same flags, smaller context
|
||||
./build/bin/llama-server \
|
||||
-m Qwen3.5-35B-A3B-Q4_K_M.gguf \
|
||||
-ngl 999 \
|
||||
-c 131072 \
|
||||
--cache-type-k q8_0 \
|
||||
--cache-type-v q4_0 \
|
||||
--flash-attn \
|
||||
--cram 65536 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8080
|
||||
```
|
||||
|
||||
#### Ollama
|
||||
|
||||
```bash
|
||||
# Requires Ollama v0.17+ for Qwen3.5 support
|
||||
# NOTE: As of March 2026, some Qwen3.5 GGUFs have compatibility issues
|
||||
# with Ollama due to mmproj vision files. llama.cpp may be more reliable.
|
||||
|
||||
# Environment variables for multi-GPU
|
||||
export OLLAMA_GPU_SPLIT=16,16 # Equal split across both 16GB cards
|
||||
export OLLAMA_KV_CACHE_TYPE=q8_0 # Halves KV cache VRAM with minimal quality loss
|
||||
export OLLAMA_KEEP_ALIVE=24h # Keep model loaded in VRAM
|
||||
export OLLAMA_FLASH_ATTENTION=1 # Enable flash attention for VRAM savings
|
||||
|
||||
# Pull and run
|
||||
ollama pull qwen3.5:35b-a3b-q4_K_M
|
||||
ollama run qwen3.5:35b-a3b-q4_K_M
|
||||
```
|
||||
|
||||
#### Verify NVLink Is Working
|
||||
|
||||
```bash
|
||||
# Check NVLink status
|
||||
nvidia-smi nvlink --status
|
||||
|
||||
# Check NVLink bandwidth
|
||||
nvidia-smi nvlink -gt d
|
||||
|
||||
# Monitor both GPUs during inference
|
||||
watch -n 0.5 nvidia-smi
|
||||
```
|
||||
|
||||
### Total Cost Summary
|
||||
|
||||
| Item | Cost |
|
||||
|------|------|
|
||||
| 2x Quadro RTX 5000 | ~$800 |
|
||||
| NVLink HB Bridge 2-slot (P4934) | ~$50 |
|
||||
| Dell cables/riser (R720/R730) | ~$60 |
|
||||
| Dell 1100W PSUs (if needed) | ~$60-100 |
|
||||
| **Hardware total** | **~$960** |
|
||||
| Monthly power (2 cards, 8hr/day) | $9/mo |
|
||||
| Claude Pro subscription (keep) | $20/mo |
|
||||
| Claude Max subscription (drop) | -$100/mo saved |
|
||||
| **Net monthly cost** | **$29/mo (was $100/mo)** |
|
||||
|
||||
### The Math: Drop Max, Keep Pro, Add Local
|
||||
|
||||
| | Year 1 | Year 2 | Year 3 | **3-Year Total** |
|
||||
|---|--------|--------|--------|-----------------|
|
||||
| **Claude Max (current)** | $1,200 | $1,200 | $1,200 | **$3,600** |
|
||||
| **Pro + Local GPU** | $960 + $348 | $348 | $348 | **$2,004** |
|
||||
| **Savings** | | | | **$1,596** |
|
||||
|
||||
You save ~$71/mo after hardware payoff. The GPU pays for itself in **13 months**.
|
||||
After that, you're saving $80/mo vs Max with no usage limits on bulk work.
|
||||
|
||||
### Comparison: This Build vs Alternatives
|
||||
|
||||
| Setup | Monthly | 3yr Total | Limits? | Quality |
|
||||
|-------|---------|-----------|---------|---------|
|
||||
| **Pro + 2x RTX 5000** | **$29** | **$2,004** | **Unlimited local, Pro limits for Opus** | **A- local, A+ cloud** |
|
||||
| Pro + 1x RTX 3090 | $25 | $1,620 | Unlimited local (128K ctx), Pro limits | A- local, A+ cloud |
|
||||
| Pro + RTX 8000 (48GB) | $29 | $3,040+ | Unlimited local, Pro limits | A- local, A+ cloud |
|
||||
| **Claude Max (no GPU)** | **$100** | **$3,600** | **Unlimited Opus** | **A+ cloud only** |
|
||||
| Claude Pro only (no GPU) | $20 | $720 | **Hits caps on large projects** | A+ cloud, limited |
|
||||
| API-only (Opus heavy use) | $500+ | $18,000+ | Pay per token | A+ cloud only |
|
||||
|
||||
### Phased Build Plan
|
||||
|
||||
**Phase 1 — Start with one card ($400)**
|
||||
1. Buy Quadro RTX 5000 (VCQRTX5000-PB) — ~$400 on eBay
|
||||
2. Install in any PCIe x16 slot
|
||||
3. Install llama.cpp or Ollama v0.17+
|
||||
4. Run Qwen3.5-35B-A3B at Q4_K_L with 64-128K context
|
||||
5. Already A- quality for coding — test if local inference fits your workflow
|
||||
|
||||
**Phase 2 — Add second card + NVLink ($450)**
|
||||
1. Buy matching Quadro RTX 5000 — ~$400
|
||||
2. Buy NVLink HB Bridge 2-slot (part: P4934 / 1JF3K / 6FY12AA) — ~$50
|
||||
3. Install second card in adjacent/nearby x16 slot
|
||||
4. Connect NVLink bridge
|
||||
5. Verify with `nvidia-smi nvlink --status`
|
||||
6. Now running 32GB unified — Qwen3.5-35B-A3B at Q4_K_M with full 262K context
|
||||
|
||||
### Dell R720/R730 Installation Guide
|
||||
|
||||
#### Prerequisites (MUST HAVE before buying GPUs)
|
||||
|
||||
| Requirement | R720 | R730 | Why |
|
||||
|-------------|-------|-------|-----|
|
||||
| **Dual CPUs** | Required | Required | GPU riser slots are wired to CPU2 — dead without it |
|
||||
| **2x 1100W PSUs** | Required | Required | 2x 230W GPUs + system = ~600W+ under load |
|
||||
| **GPU Riser 3** | Required for 2nd GPU | Required for GPUs | Provides the PCIe x16 slot + 8-pin power |
|
||||
| **GPU Power Cable** | Required | Required | Riser-to-GPU power, not included by default |
|
||||
| **Low-Profile Heatsinks** | Must swap (part of enablement kit) | Usually pre-installed | Standard heatsinks block GPU riser clearance |
|
||||
| **Max ambient temp** | 30°C (not the usual 35°C) | 30°C | High GPU TDP restricts cooling headroom |
|
||||
|
||||
#### Shopping List: Dell-Specific Parts
|
||||
|
||||
| Part | Dell P/N | What It Is | Price | Where |
|
||||
|------|----------|------------|-------|-------|
|
||||
| **GPU Power Cable** | **9H6FV** (09H6FV) | 8-pin EPS (riser) → 6-pin + 6+2-pin PCIe. One cable powers one GPU | ~$10-15 | Amazon, eBay |
|
||||
| **GPU Power Cable (alt)** | **N08NH** (0N08NH) | Same function, alternate Dell part number | ~$10-15 | Amazon, eBay |
|
||||
| **GPU Riser 3** (R720) | Check eBay for "R720 riser 3" or "R720 GPU riser" | Second riser card that provides GPU-capable x16 slot | ~$15-30 | eBay |
|
||||
| **GPU Riser 3** (R730) | Check eBay for "R730 riser 3" or "R730 GPU riser" | R730 version — NOT interchangeable with R720 | ~$15-30 | eBay |
|
||||
| **Low-Profile Heatsinks** (R720 only) | Part of original GPU enablement kit | Shorter heatsinks that clear the GPU riser. Search "R720 low profile heatsink" | ~$10-20/pair | eBay |
|
||||
|
||||
**You need 2x power cables** (one per GPU). Search Amazon for "Dell R720 R730 GPU power cable 9H6FV" — multiple sellers (COMeap, ZAHARA, BestParts) stock them for ~$10-15 each.
|
||||
|
||||
#### How It Fits
|
||||
|
||||
```
|
||||
Dell R720/R730 Riser Layout (rear view):
|
||||
┌─────────────────────────────────────┐
|
||||
│ Riser 1 Riser 2 Riser 3│
|
||||
│ (network/ (GPU 1) (GPU 2)│
|
||||
│ storage) PCIe x16 PCIe x16│
|
||||
│ Gen2(720) Gen2(720)│
|
||||
│ Gen3(730) Gen3(730)│
|
||||
└─────────────────────────────────────┘
|
||||
↑ RTX 5000 ↑ ↑ RTX 5000 ↑
|
||||
└── NVLink Bridge ──┘
|
||||
```
|
||||
|
||||
- Both GPUs sit on **adjacent risers** (Riser 2 + Riser 3) — this is 2-slot spacing
|
||||
- The **2-slot NVLink bridge** (P4934) is the correct size for R720/R730
|
||||
- The cards mount **vertically** via risers, parallel to each other
|
||||
- NVLink bridge connects across the top of both cards
|
||||
|
||||
#### R720 vs R730
|
||||
|
||||
| Feature | R720 | R730 |
|
||||
|---------|------|------|
|
||||
| **PCIe** | Gen2 x16 | **Gen3 x16** |
|
||||
| **Impact on LLM** | None — VRAM-bound | None — VRAM-bound |
|
||||
| **Impact on NVLink** | None — NVLink bypasses PCIe | None — NVLink bypasses PCIe |
|
||||
| **GPU power delivery** | Same 8-pin from riser | Same 8-pin from riser |
|
||||
| **Heatsink swap** | Usually required | Usually already low-profile |
|
||||
| **Used price** | ~$100-150 cheaper | Preferred if budget allows |
|
||||
| **Recommendation** | Fine if you already have one | **Buy this one** if shopping new |
|
||||
|
||||
#### Potential Issues
|
||||
|
||||
1. **NVLink bridge clearance in 2U** — The bridge sits on top of both GPUs. In a 2U chassis
|
||||
this is tight. The R720/R730 riser design mounts cards vertically which actually helps —
|
||||
the bridge faces the chassis side panel, not the lid. Should fit, but measure before buying.
|
||||
|
||||
2. **Blower fan noise** — The RTX 5000 has an active blower (unlike passive Tesla cards).
|
||||
The server's own fans may spin higher to compensate. The blower exhausts out the bracket
|
||||
which is correct for rack airflow.
|
||||
|
||||
3. **"Unsupported" GPU warning** — Dell officially supports Tesla/Quadro cards from their era.
|
||||
The Quadro RTX 5000 is a later generation than R720/R730 was designed for, but community
|
||||
reports confirm Quadro RTX and even consumer RTX cards work fine. You won't get Dell support
|
||||
if something goes wrong, but electrically it's standard PCIe.
|
||||
|
||||
4. **CPU TDP limit** — Dell requires CPUs of 115W or less when GPUs are installed (R720).
|
||||
Check your CPU model. Most common Xeon E5-2600 v1/v2 (R720) and E5-2600 v3/v4 (R730)
|
||||
processors are within this range, but some high-core-count variants exceed it.
|
||||
|
||||
5. **PSU mode** — With dual 300W GPUs, set PSU configuration to **non-redundant mode**
|
||||
to use combined wattage from both PSUs. In redundant mode, you're limited to one PSU's
|
||||
capacity (1100W) which may not be enough under full GPU + CPU load.
|
||||
|
||||
#### Complete R720/R730 Shopping List
|
||||
|
||||
```
|
||||
GPUS + NVLINK
|
||||
2x Quadro RTX 5000 ~$800
|
||||
1x NVLink Bridge 2-slot (P4934 / L55997-001) ~$50
|
||||
|
||||
DELL-SPECIFIC PARTS
|
||||
2x GPU Power Cable (9H6FV or N08NH) ~$25
|
||||
1x GPU Riser 3 (match your server model!) ~$20
|
||||
2x Low-Profile Heatsinks (R720 only) ~$15
|
||||
|
||||
POWER (if not already installed)
|
||||
2x Dell 1100W PSU ~$30-50 ea
|
||||
|
||||
TOTAL (assuming you have the server + dual CPUs) ~$940-960
|
||||
```
|
||||
|
||||
## 24GB GPU Options (Previous Research — Still Valid for Tighter Budgets)
|
||||
|
||||
| GPU | VRAM | Price Range | Best Deals | Notes |
|
||||
|-----|------|-------------|------------|-------|
|
||||
| **Tesla P40** | 24GB | $150-320 | Newegg refurb $219-270; eBay used $150-200 | Best VRAM/$ at 24GB |
|
||||
| **RTX A2000 12GB** | 12GB | $250-535 | eBay used ~$250-350; one listing at $490 | Can't run 32B models |
|
||||
| **Tesla T4** | 16GB | $150-350 | eBay used $150-250 | Great power efficiency |
|
||||
| **RTX A4000** | 16GB | $700-750+ | eBay used ~$700; new $720+ | Too expensive for 16GB |
|
||||
|
||||
## Rack Server Compatibility
|
||||
|
||||
### Quadro RTX 8000 Passive in R720/R730
|
||||
- **Physical fit**: Full-length, dual-slot — fits in GPU riser slots
|
||||
- **Power**: 260W, requires 8-pin aux power + GPU enablement kit
|
||||
- **Cooling**: Passive — relies on server chassis fans (same as P40)
|
||||
- **Requirement**: Dual CPUs, redundant 1100W PSUs recommended
|
||||
- **NVLink**: Can pair two RTX 8000s for 96GB combined VRAM
|
||||
- Very similar physical/power requirements to the Tesla P40
|
||||
|
||||
### RTX A2000 in R720/R730
|
||||
- **Physical fit**: Yes. Dual-slot, low-profile, 167mm length
|
||||
- **Power**: 70W bus-powered, no aux cable needed. Must use 75W slots (slots 4-7 on R720)
|
||||
- **Cooling**: Blower-style fan exhausts out bracket — ideal for rack airflow
|
||||
- **Requirement**: Dual CPUs needed for GPU PCIe slots
|
||||
- **Confirmed working** in Dell R740XD (similar architecture)
|
||||
|
||||
### Tesla P40 in R720/R730
|
||||
- **Physical fit**: Yes. Full-length, single-slot, designed for rack servers
|
||||
- **Power**: 250W, requires 8-pin aux power. Needs GPU enablement kit
|
||||
- **Cooling**: Passive — relies on server chassis fans
|
||||
- **Requirement**: Dual CPUs, redundant 1100W PSUs recommended
|
||||
- **Natively supported** in these servers
|
||||
|
||||
### R720 vs R730
|
||||
- R720: PCIe Gen2 (not a bottleneck for LLM inference, which is VRAM-bound)
|
||||
- R730: PCIe Gen3, generally preferred
|
||||
- Both support up to 2x double-wide or 4x single-wide GPUs
|
||||
|
||||
## Recommended Setups
|
||||
|
||||
### If budget allows ($2,000–2,900): RTX 8000 Passive
|
||||
Single card handles both coding and image gen. 48GB VRAM fits 32B models with massive
|
||||
context windows (32K+). Passive cooling is rack-native. Tensor cores handle FP16 image gen
|
||||
properly. One card, one slot, simple setup. The premium buys you unified VRAM = big context.
|
||||
|
||||
### If budget allows + dedicated image gen ($2,300–3,250): RTX 8000 + A2000
|
||||
RTX 8000 for coding with full 48GB dedicated to LLM context.
|
||||
A2000 for image gen (3x faster than Turing, 70W, bus-powered, blower cooled).
|
||||
Best separation of concerns — no model swapping needed.
|
||||
|
||||
### Best value ($400–500): Dual P40
|
||||
Two P40s for 48GB total, but split across cards (can't combine for one model without
|
||||
NVLink, which P40s lack). One for 32B coding (tight fit, ~4K context), one for image gen
|
||||
(slow, needs --force-fp32). **5x cheaper than RTX 8000** but with significant context limitations.
|
||||
|
||||
### Cheapest entry ($200–300): Single P40
|
||||
Run 32B coding model with very limited context (~4K tokens). Swap to image gen when needed.
|
||||
Good for testing whether local LLM coding works for your workflow before investing more.
|
||||
|
||||
## Configuration Notes for local-ai stack
|
||||
|
||||
### For 48GB RTX 8000
|
||||
```bash
|
||||
# Ollama — take advantage of the full 48GB
|
||||
OLLAMA_NUM_GPU=999
|
||||
OLLAMA_NUM_CTX=32768 # Large context window — 48GB can handle it
|
||||
OLLAMA_KEEP_ALIVE=24h
|
||||
|
||||
# Pull best coding models
|
||||
ollama pull qwen2.5-coder:32b-instruct-q4_K_M # ~20GB, leaves 28GB for context
|
||||
ollama pull qwen3.5:27b # ~16GB at Q4, even more context room
|
||||
ollama pull qwen3-coder:30b # MoE, very fast inference
|
||||
|
||||
# Higher quantization for better quality (48GB allows this)
|
||||
# Look for Q6_K or Q8_0 variants on Ollama for better output quality
|
||||
```
|
||||
|
||||
### For 32B models on P40 (24GB — tight fit)
|
||||
```bash
|
||||
OLLAMA_NUM_GPU=999
|
||||
OLLAMA_NUM_CTX=4096 # Keep context small to fit in remaining VRAM
|
||||
OLLAMA_KEEP_ALIVE=24h
|
||||
|
||||
ollama pull qwen2.5-coder:32b-instruct-q4_K_M
|
||||
```
|
||||
|
||||
### For dual-GPU setup (RTX 8000 + A2000 or P40 + anything)
|
||||
```bash
|
||||
# Assign GPU 0 to Ollama (coding), GPU 1 to InvokeAI (image gen)
|
||||
# In docker-compose.yml for Ollama:
|
||||
CUDA_VISIBLE_DEVICES=0
|
||||
|
||||
# In docker-compose.yml for InvokeAI:
|
||||
CUDA_VISIBLE_DEVICES=1
|
||||
```
|
||||
|
||||
### For image gen on P40 (no tensor cores)
|
||||
```bash
|
||||
# InvokeAI
|
||||
INVOKEAI_PRECISION=float32
|
||||
|
||||
# ComfyUI launch args
|
||||
--force-fp32
|
||||
```
|
||||
|
||||
### For image gen on RTX 8000 / A2000 / T4 (has tensor cores)
|
||||
```bash
|
||||
# InvokeAI — native FP16 works fine
|
||||
INVOKEAI_PRECISION=float16
|
||||
|
||||
# ComfyUI — no special flags needed
|
||||
```
|
||||
|
||||
## Sources
|
||||
- [Quadro RTX 8000 for Local LLMs — Hardware Corner](https://www.hardware-corner.net/guides/quadro-rtx-8000-for-llm/)
|
||||
- [RTX 8000 Passive — Network Outlet](https://networkoutlet.com/blogs/articles/nvidia-quadro-rtx-8000-48gb-passive-cooling-powering-ai-rendering-server-workloads)
|
||||
- [LLM Benchmarks on Turing/Ampere GPUs — Stefandroid](https://blog.stefandroid.com/2025/06/02/benchmark-llm-performance-nvidia-gpus.html)
|
||||
- [NVIDIA A40 Price Tracking — GPUPoet](https://gpupoet.com/gpu/learn/card/nvidia-a40)
|
||||
- [NVIDIA L40 Price Tracking — GPUPoet](https://gpupoet.com/gpu/learn/card/nvidia-l40)
|
||||
- [RTX A6000 Price History — CamelCamelCamel](https://camelcamelcamel.com/product/B09BDH8VZV)
|
||||
- [RTX A6000 Price History — Pangoly](https://pangoly.com/en/price-history/pny-nvidia-quadro-rtx-a6000)
|
||||
- [NVIDIA RTX A2000 Datasheet](https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/rtx-a2000/nvidia-rtx-a2000-datasheet-1987439-r5.pdf)
|
||||
- [Dell R730 Owner's Manual — Expansion Cards](https://www.dell.com/support/manuals/en-us/poweredge-r730/r730_ompublication/expansion-card-installation-guidelines)
|
||||
- [Dell R720 Owner's Manual — Expansion Cards](https://www.dell.com/support/manuals/en-us/poweredge-r720/720720xdom/expansion-card-installation-guidelines)
|
||||
- [ComfyUI GPU Benchmarks Discussion](https://github.com/Comfy-Org/ComfyUI/discussions/2970)
|
||||
- [ComfyUI P40 FP32 Issue](https://github.com/Comfy-Org/ComfyUI/issues/4363)
|
||||
- [Best Local LLMs for 24GB VRAM 2026](https://localllm.in/blog/best-local-llms-24gb-vram)
|
||||
- [Best Coding Models 2026](https://localvram.com/en/guides/best-coding-models/)
|
||||
- [Ollama VRAM Requirements Guide](https://localllm.in/blog/ollama-vram-requirements-for-local-llms)
|
||||
- [Local LLMs That Can Replace Claude Code](https://agentnativedev.medium.com/local-llms-that-can-replace-claude-code-6f5b6cac93bf)
|
||||
- [7 Local LLM Families to Replace Claude/Codex](https://agentnativedev.medium.com/7-local-llm-families-to-replace-claude-codex-for-everyday-tasks-25ba74c3635d)
|
||||
- [Qwen2.5-Coder 32B on Ollama](https://ollama.com/library/qwen2.5-coder:32b-instruct-q4_K_M)
|
||||
- [Qwen3-Coder — How to Run Locally](https://unsloth.ai/docs/models/qwen3-coder-how-to-run-locally)
|
||||
+464
@@ -0,0 +1,464 @@
|
||||
#!/usr/bin/env bash
|
||||
# =============================================================================
|
||||
# Gitea ↔ GitHub Mirror Sync
|
||||
#
|
||||
# Mirrors repos between your local Gitea and GitHub in both directions:
|
||||
# GitHub → Gitea: Pulls repos you own on GitHub into Gitea (backup/offline use)
|
||||
# Gitea → GitHub: Pushes Gitea repos to GitHub (remote backup)
|
||||
#
|
||||
# Usage:
|
||||
# ./gitea-github-sync.sh — sync all configured repos
|
||||
# ./gitea-github-sync.sh --pull-only — GitHub → Gitea only
|
||||
# ./gitea-github-sync.sh --push-only — Gitea → GitHub only
|
||||
# ./gitea-github-sync.sh --repo owner/name — sync one specific repo
|
||||
# ./gitea-github-sync.sh --list — list what would sync (dry run)
|
||||
# ./gitea-github-sync.sh --init — interactive first-time setup
|
||||
#
|
||||
# Config: ~/.config/gitea-github-sync/config
|
||||
# Tokens: reads from .env in the same directory as this script (or $SYNC_ENV)
|
||||
#
|
||||
# Schedule: install the systemd timer with --install-timer
|
||||
# ./gitea-github-sync.sh --install-timer — every 6 hours (default)
|
||||
# ./gitea-github-sync.sh --install-timer 1h — custom interval
|
||||
# ./gitea-github-sync.sh --remove-timer — remove the timer
|
||||
# =============================================================================
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
info() { echo -e "${CYAN}[sync]${NC} $*"; }
|
||||
ok() { echo -e "${GREEN}[ ok ]${NC} $*"; }
|
||||
warn() { echo -e "${YELLOW}[warn]${NC} $*"; }
|
||||
err() { echo -e "${RED}[err ]${NC} $*" >&2; }
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
CONFIG_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/gitea-github-sync"
|
||||
CONFIG_FILE="$CONFIG_DIR/config"
|
||||
WORK_DIR="$CONFIG_DIR/repos"
|
||||
LOG_FILE="$CONFIG_DIR/sync.log"
|
||||
|
||||
# ── load tokens from .env ───────────────────────────────────────────────────
|
||||
ENV_FILE="${SYNC_ENV:-$SCRIPT_DIR/.env}"
|
||||
if [[ -f "$ENV_FILE" ]]; then
|
||||
# shellcheck disable=SC1090
|
||||
set -a; source <(grep -E '^(GITEA_TOKEN|GITHUB_TOKEN|GITEA_URL)=' "$ENV_FILE" | sed 's/ *#.*//'); set +a
|
||||
fi
|
||||
|
||||
GITEA_URL="${GITEA_URL:-http://localhost:3001}"
|
||||
GITEA_TOKEN="${GITEA_TOKEN:-}"
|
||||
GITHUB_TOKEN="${GITHUB_TOKEN:-}"
|
||||
|
||||
# ── parse args ──────────────────────────────────────────────────────────────
|
||||
MODE="all" # all | pull | push | list | init | install-timer | remove-timer
|
||||
SINGLE_REPO=""
|
||||
TIMER_INTERVAL="6h"
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--pull-only) MODE="pull"; shift ;;
|
||||
--push-only) MODE="push"; shift ;;
|
||||
--list) MODE="list"; shift ;;
|
||||
--init) MODE="init"; shift ;;
|
||||
--install-timer) MODE="install-timer"; shift; [[ "${1:-}" =~ ^[0-9]+[smhd]$ ]] && { TIMER_INTERVAL="$1"; shift; } ;;
|
||||
--remove-timer) MODE="remove-timer"; shift ;;
|
||||
--repo) shift; SINGLE_REPO="${1:-}"; shift ;;
|
||||
-h|--help)
|
||||
sed -n '2,/^# =====/{ /^# =====/d; s/^# \?//p; }' "$0"; exit 0 ;;
|
||||
*) err "Unknown arg: $1"; exit 1 ;;
|
||||
esac
|
||||
done
|
||||
|
||||
# ── helpers ─────────────────────────────────────────────────────────────────
|
||||
_gitea_api() {
|
||||
local method="$1" path="$2"; shift 2
|
||||
curl -sfL -X "$method" \
|
||||
-H "Authorization: token $GITEA_TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
"$GITEA_URL/api/v1$path" "$@"
|
||||
}
|
||||
|
||||
_github_api() {
|
||||
local method="$1" path="$2"; shift 2
|
||||
curl -sfL -X "$method" \
|
||||
-H "Authorization: Bearer $GITHUB_TOKEN" \
|
||||
-H "Accept: application/vnd.github+json" \
|
||||
"https://api.github.com$path" "$@"
|
||||
}
|
||||
|
||||
_log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" >> "$LOG_FILE"; }
|
||||
|
||||
# ── config management ──────────────────────────────────────────────────────
|
||||
load_config() {
|
||||
mkdir -p "$CONFIG_DIR" "$WORK_DIR"
|
||||
GITHUB_USER=""
|
||||
GITEA_USER=""
|
||||
SYNC_REPOS=() # explicit list (empty = auto-discover)
|
||||
EXCLUDE_REPOS=() # repos to skip
|
||||
PUSH_PRIVATE=false # push private Gitea repos to GitHub?
|
||||
PULL_PRIVATE=true # pull private GitHub repos to Gitea?
|
||||
PULL_FORKS=false # pull forked repos from GitHub?
|
||||
|
||||
if [[ -f "$CONFIG_FILE" ]]; then
|
||||
# shellcheck disable=SC1090
|
||||
source "$CONFIG_FILE"
|
||||
fi
|
||||
}
|
||||
|
||||
save_config() {
|
||||
mkdir -p "$CONFIG_DIR"
|
||||
cat > "$CONFIG_FILE" << EOF
|
||||
# Gitea-GitHub Sync — configuration
|
||||
# Generated $(date '+%Y-%m-%d %H:%M:%S')
|
||||
|
||||
# GitHub username (for discovering repos to pull)
|
||||
GITHUB_USER="$GITHUB_USER"
|
||||
|
||||
# Gitea username (for discovering repos to push)
|
||||
GITEA_USER="$GITEA_USER"
|
||||
|
||||
# Explicit repo list — if set, only these sync. Format: owner/repo
|
||||
# Leave empty () to auto-discover from both platforms.
|
||||
SYNC_REPOS=($(printf '"%s" ' "${SYNC_REPOS[@]}"))
|
||||
|
||||
# Repos to skip (pattern matched against owner/repo)
|
||||
EXCLUDE_REPOS=($(printf '"%s" ' "${EXCLUDE_REPOS[@]}"))
|
||||
|
||||
# Push private Gitea repos to GitHub as private repos?
|
||||
PUSH_PRIVATE=$PUSH_PRIVATE
|
||||
|
||||
# Pull private GitHub repos to Gitea?
|
||||
PULL_PRIVATE=$PULL_PRIVATE
|
||||
|
||||
# Pull forked repos from GitHub?
|
||||
PULL_FORKS=$PULL_FORKS
|
||||
EOF
|
||||
ok "Config saved: $CONFIG_FILE"
|
||||
}
|
||||
|
||||
# ── init (first-time setup) ────────────────────────────────────────────────
|
||||
do_init() {
|
||||
echo -e "\n${BOLD}Gitea ↔ GitHub Sync — First-Time Setup${NC}\n"
|
||||
|
||||
# Check tokens
|
||||
if [[ -z "$GITEA_TOKEN" || "$GITEA_TOKEN" == "your-gitea-token-here" ]]; then
|
||||
err "GITEA_TOKEN not set. Add it to $ENV_FILE first."
|
||||
echo " Generate at: $GITEA_URL/user/settings/applications"
|
||||
exit 1
|
||||
fi
|
||||
if [[ -z "$GITHUB_TOKEN" || "$GITHUB_TOKEN" == "your-github-token-here" ]]; then
|
||||
err "GITHUB_TOKEN not set. Add it to $ENV_FILE first."
|
||||
echo " Generate at: https://github.com/settings/tokens"
|
||||
echo " Scopes needed: repo (full control)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Discover usernames
|
||||
info "Detecting GitHub user..."
|
||||
GITHUB_USER=$(_github_api GET /user | python3 -c "import sys,json; print(json.load(sys.stdin)['login'])" 2>/dev/null) \
|
||||
|| { err "Failed to reach GitHub API. Check GITHUB_TOKEN."; exit 1; }
|
||||
ok "GitHub user: $GITHUB_USER"
|
||||
|
||||
info "Detecting Gitea user..."
|
||||
GITEA_USER=$(_gitea_api GET /user | python3 -c "import sys,json; print(json.load(sys.stdin)['login'])" 2>/dev/null) \
|
||||
|| { err "Failed to reach Gitea API. Check GITEA_TOKEN and GITEA_URL ($GITEA_URL)."; exit 1; }
|
||||
ok "Gitea user: $GITEA_USER"
|
||||
|
||||
# Ask about sync scope
|
||||
echo ""
|
||||
read -rp "Pull private GitHub repos to Gitea? [Y/n] " ans
|
||||
PULL_PRIVATE=true; [[ "${ans,,}" == "n" ]] && PULL_PRIVATE=false
|
||||
|
||||
read -rp "Pull forked repos from GitHub? [y/N] " ans
|
||||
PULL_FORKS=false; [[ "${ans,,}" == "y" ]] && PULL_FORKS=true
|
||||
|
||||
read -rp "Push private Gitea repos to GitHub? [y/N] " ans
|
||||
PUSH_PRIVATE=false; [[ "${ans,,}" == "y" ]] && PUSH_PRIVATE=true
|
||||
|
||||
save_config
|
||||
|
||||
echo ""
|
||||
info "Run '$(basename "$0") --list' to preview what would sync."
|
||||
info "Run '$(basename "$0")' to sync now."
|
||||
info "Run '$(basename "$0") --install-timer' to sync automatically."
|
||||
}
|
||||
|
||||
# ── discover repos ─────────────────────────────────────────────────────────
|
||||
get_github_repos() {
|
||||
local page=1 repos=()
|
||||
while true; do
|
||||
local batch
|
||||
batch=$(_github_api GET "/user/repos?per_page=100&page=$page&affiliation=owner" \
|
||||
| python3 -c "
|
||||
import sys, json
|
||||
for r in json.load(sys.stdin):
|
||||
if r.get('fork') and not $($PULL_FORKS && echo True || echo False):
|
||||
continue
|
||||
if r.get('private') and not $($PULL_PRIVATE && echo True || echo False):
|
||||
continue
|
||||
print(r['full_name'] + '|' + r['clone_url'] + '|' + str(r.get('private',False)).lower())
|
||||
" 2>/dev/null) || break
|
||||
[[ -z "$batch" ]] && break
|
||||
while IFS= read -r line; do repos+=("$line"); done <<< "$batch"
|
||||
((page++))
|
||||
done
|
||||
printf '%s\n' "${repos[@]}"
|
||||
}
|
||||
|
||||
get_gitea_repos() {
|
||||
local page=1 repos=()
|
||||
while true; do
|
||||
local batch
|
||||
batch=$(_gitea_api GET "/repos/search?limit=50&page=$page" \
|
||||
| python3 -c "
|
||||
import sys, json
|
||||
for r in json.load(sys.stdin).get('data', []):
|
||||
if r.get('private') and not $($PUSH_PRIVATE && echo True || echo False):
|
||||
continue
|
||||
print(r['full_name'] + '|' + r['clone_url'] + '|' + str(r.get('private',False)).lower())
|
||||
" 2>/dev/null) || break
|
||||
[[ -z "$batch" ]] && break
|
||||
while IFS= read -r line; do repos+=("$line"); done <<< "$batch"
|
||||
((page++))
|
||||
done
|
||||
printf '%s\n' "${repos[@]}"
|
||||
}
|
||||
|
||||
is_excluded() {
|
||||
local repo="$1"
|
||||
for pat in "${EXCLUDE_REPOS[@]}"; do
|
||||
[[ "$repo" == $pat ]] && return 0
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
# ── sync: GitHub → Gitea (pull) ───────────────────────────────────────────
|
||||
sync_github_to_gitea() {
|
||||
local full_name="$1" clone_url="$2" is_private="$3"
|
||||
local repo_name="${full_name#*/}"
|
||||
local local_path="$WORK_DIR/$full_name"
|
||||
|
||||
# Clone or fetch from GitHub
|
||||
if [[ -d "$local_path" ]]; then
|
||||
info "Fetching $full_name from GitHub..."
|
||||
git -C "$local_path" fetch --all --prune --quiet 2>/dev/null || {
|
||||
err "Failed to fetch $full_name"; return 1; }
|
||||
else
|
||||
info "Cloning $full_name from GitHub..."
|
||||
mkdir -p "$(dirname "$local_path")"
|
||||
local auth_url="${clone_url/https:\/\//https:\/\/$GITHUB_TOKEN@}"
|
||||
git clone --bare --quiet "$auth_url" "$local_path" 2>/dev/null || {
|
||||
err "Failed to clone $full_name"; return 1; }
|
||||
fi
|
||||
|
||||
# Ensure repo exists on Gitea
|
||||
local gitea_check
|
||||
gitea_check=$(_gitea_api GET "/repos/$GITEA_USER/$repo_name" 2>/dev/null) || true
|
||||
if ! echo "$gitea_check" | python3 -c "import sys,json; json.load(sys.stdin)['id']" &>/dev/null; then
|
||||
info "Creating $repo_name on Gitea..."
|
||||
_gitea_api POST "/user/repos" \
|
||||
-d "{\"name\":\"$repo_name\",\"private\":$is_private,\"description\":\"Mirror of $full_name from GitHub\"}" \
|
||||
>/dev/null || { err "Failed to create $repo_name on Gitea"; return 1; }
|
||||
fi
|
||||
|
||||
# Push to Gitea
|
||||
local gitea_push_url="${GITEA_URL/https:\/\//https:\/\/$GITEA_USER:$GITEA_TOKEN@}"
|
||||
gitea_push_url="${gitea_push_url/http:\/\//http:\/\/$GITEA_USER:$GITEA_TOKEN@}"
|
||||
gitea_push_url="$gitea_push_url/$GITEA_USER/$repo_name.git"
|
||||
|
||||
git -C "$local_path" push --mirror "$gitea_push_url" --quiet 2>/dev/null || {
|
||||
err "Failed to push $full_name to Gitea"; return 1; }
|
||||
ok "GitHub → Gitea: $full_name"
|
||||
_log "PULL $full_name OK"
|
||||
}
|
||||
|
||||
# ── sync: Gitea → GitHub (push) ───────────────────────────────────────────
|
||||
sync_gitea_to_github() {
|
||||
local full_name="$1" clone_url="$2" is_private="$3"
|
||||
local repo_name="${full_name#*/}"
|
||||
local local_path="$WORK_DIR/gitea/$full_name"
|
||||
|
||||
# Clone or fetch from Gitea
|
||||
local gitea_auth_url="${clone_url/https:\/\//https:\/\/$GITEA_USER:$GITEA_TOKEN@}"
|
||||
gitea_auth_url="${gitea_auth_url/http:\/\//http:\/\/$GITEA_USER:$GITEA_TOKEN@}"
|
||||
|
||||
if [[ -d "$local_path" ]]; then
|
||||
info "Fetching $full_name from Gitea..."
|
||||
git -C "$local_path" fetch --all --prune --quiet 2>/dev/null || {
|
||||
err "Failed to fetch $full_name from Gitea"; return 1; }
|
||||
else
|
||||
info "Cloning $full_name from Gitea..."
|
||||
mkdir -p "$(dirname "$local_path")"
|
||||
git clone --bare --quiet "$gitea_auth_url" "$local_path" 2>/dev/null || {
|
||||
err "Failed to clone $full_name from Gitea"; return 1; }
|
||||
fi
|
||||
|
||||
# Ensure repo exists on GitHub
|
||||
local gh_check
|
||||
gh_check=$(_github_api GET "/repos/$GITHUB_USER/$repo_name" 2>/dev/null) || true
|
||||
if ! echo "$gh_check" | python3 -c "import sys,json; json.load(sys.stdin)['id']" &>/dev/null; then
|
||||
info "Creating $repo_name on GitHub..."
|
||||
_github_api POST "/user/repos" \
|
||||
-d "{\"name\":\"$repo_name\",\"private\":$is_private,\"description\":\"Mirror from Gitea\"}" \
|
||||
>/dev/null || { err "Failed to create $repo_name on GitHub"; return 1; }
|
||||
fi
|
||||
|
||||
# Push to GitHub
|
||||
local github_push_url="https://$GITHUB_TOKEN@github.com/$GITHUB_USER/$repo_name.git"
|
||||
git -C "$local_path" push --mirror "$github_push_url" --quiet 2>/dev/null || {
|
||||
err "Failed to push $full_name to GitHub"; return 1; }
|
||||
ok "Gitea → GitHub: $full_name"
|
||||
_log "PUSH $full_name OK"
|
||||
}
|
||||
|
||||
# ── list (dry run) ─────────────────────────────────────────────────────────
|
||||
do_list() {
|
||||
echo -e "\n${BOLD}Repos that would sync:${NC}\n"
|
||||
|
||||
if [[ ${#SYNC_REPOS[@]} -gt 0 ]]; then
|
||||
echo -e "${CYAN}Explicit list:${NC}"
|
||||
printf ' %s\n' "${SYNC_REPOS[@]}"
|
||||
else
|
||||
if [[ "$MODE" != "push" ]]; then
|
||||
echo -e "${CYAN}GitHub → Gitea (pull):${NC}"
|
||||
get_github_repos | while IFS='|' read -r name url priv; do
|
||||
is_excluded "$name" && echo " $name (excluded)" && continue
|
||||
echo " $name $([ "$priv" = "true" ] && echo "[private]")"
|
||||
done
|
||||
fi
|
||||
echo ""
|
||||
if [[ "$MODE" != "pull" ]]; then
|
||||
echo -e "${CYAN}Gitea → GitHub (push):${NC}"
|
||||
get_gitea_repos | while IFS='|' read -r name url priv; do
|
||||
is_excluded "$name" && echo " $name (excluded)" && continue
|
||||
echo " $name $([ "$priv" = "true" ] && echo "[private]")"
|
||||
done
|
||||
fi
|
||||
fi
|
||||
echo ""
|
||||
}
|
||||
|
||||
# ── main sync ──────────────────────────────────────────────────────────────
|
||||
do_sync() {
|
||||
local pull_count=0 push_count=0 fail_count=0
|
||||
|
||||
_log "=== Sync started (mode=$MODE) ==="
|
||||
|
||||
# GitHub → Gitea
|
||||
if [[ "$MODE" == "all" || "$MODE" == "pull" ]]; then
|
||||
info "Discovering GitHub repos..."
|
||||
while IFS='|' read -r name url priv; do
|
||||
[[ -z "$name" ]] && continue
|
||||
[[ -n "$SINGLE_REPO" && "$name" != "$SINGLE_REPO" ]] && continue
|
||||
is_excluded "$name" && continue
|
||||
if sync_github_to_gitea "$name" "$url" "$priv"; then
|
||||
((pull_count++))
|
||||
else
|
||||
((fail_count++))
|
||||
fi
|
||||
done < <(get_github_repos)
|
||||
fi
|
||||
|
||||
# Gitea → GitHub
|
||||
if [[ "$MODE" == "all" || "$MODE" == "push" ]]; then
|
||||
info "Discovering Gitea repos..."
|
||||
while IFS='|' read -r name url priv; do
|
||||
[[ -z "$name" ]] && continue
|
||||
[[ -n "$SINGLE_REPO" && "${name#*/}" != "${SINGLE_REPO#*/}" ]] && continue
|
||||
is_excluded "$name" && continue
|
||||
# Skip repos that came from GitHub (already mirrored)
|
||||
local repo_name="${name#*/}"
|
||||
if [[ -d "$WORK_DIR/$GITHUB_USER/$repo_name" ]]; then
|
||||
info "Skipping $name (already a GitHub mirror)"
|
||||
continue
|
||||
fi
|
||||
if sync_gitea_to_github "$name" "$url" "$priv"; then
|
||||
((push_count++))
|
||||
else
|
||||
((fail_count++))
|
||||
fi
|
||||
done < <(get_gitea_repos)
|
||||
fi
|
||||
|
||||
echo ""
|
||||
ok "Sync complete: ${pull_count} pulled, ${push_count} pushed, ${fail_count} failed"
|
||||
_log "=== Sync complete: pull=$pull_count push=$push_count fail=$fail_count ==="
|
||||
}
|
||||
|
||||
# ── systemd timer ──────────────────────────────────────────────────────────
|
||||
install_timer() {
|
||||
local service_file="/etc/systemd/system/gitea-github-sync.service"
|
||||
local timer_file="/etc/systemd/system/gitea-github-sync.timer"
|
||||
local script_path
|
||||
script_path="$(readlink -f "$0")"
|
||||
|
||||
info "Installing systemd timer (interval: $TIMER_INTERVAL)..."
|
||||
|
||||
sudo tee "$service_file" > /dev/null << EOF
|
||||
[Unit]
|
||||
Description=Gitea-GitHub Mirror Sync
|
||||
After=network-online.target docker.service
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
User=$USER
|
||||
ExecStart=$script_path
|
||||
Environment=HOME=$HOME
|
||||
StandardOutput=append:$LOG_FILE
|
||||
StandardError=append:$LOG_FILE
|
||||
EOF
|
||||
|
||||
sudo tee "$timer_file" > /dev/null << EOF
|
||||
[Unit]
|
||||
Description=Gitea-GitHub Sync Timer
|
||||
|
||||
[Timer]
|
||||
OnBootSec=5min
|
||||
OnUnitActiveSec=$TIMER_INTERVAL
|
||||
Persistent=true
|
||||
|
||||
[Install]
|
||||
WantedBy=timers.target
|
||||
EOF
|
||||
|
||||
sudo systemctl daemon-reload
|
||||
sudo systemctl enable --now gitea-github-sync.timer
|
||||
ok "Timer installed: every $TIMER_INTERVAL"
|
||||
ok "Check status: systemctl status gitea-github-sync.timer"
|
||||
ok "Run now: sudo systemctl start gitea-github-sync.service"
|
||||
ok "Logs: $LOG_FILE"
|
||||
}
|
||||
|
||||
remove_timer() {
|
||||
info "Removing systemd timer..."
|
||||
sudo systemctl disable --now gitea-github-sync.timer 2>/dev/null || true
|
||||
sudo rm -f /etc/systemd/system/gitea-github-sync.{service,timer}
|
||||
sudo systemctl daemon-reload
|
||||
ok "Timer removed"
|
||||
}
|
||||
|
||||
# ── preflight checks ──────────────────────────────────────────────────────
|
||||
preflight() {
|
||||
local ok=true
|
||||
if [[ -z "$GITEA_TOKEN" || "$GITEA_TOKEN" == "your-gitea-token-here" ]]; then
|
||||
err "GITEA_TOKEN not set. Edit $ENV_FILE"; ok=false
|
||||
fi
|
||||
if [[ -z "$GITHUB_TOKEN" || "$GITHUB_TOKEN" == "your-github-token-here" ]]; then
|
||||
err "GITHUB_TOKEN not set. Edit $ENV_FILE"; ok=false
|
||||
fi
|
||||
if [[ -z "$GITEA_USER" || -z "$GITHUB_USER" ]]; then
|
||||
err "Run --init first to configure usernames"; ok=false
|
||||
fi
|
||||
$ok || exit 1
|
||||
}
|
||||
|
||||
# ── main ───────────────────────────────────────────────────────────────────
|
||||
load_config
|
||||
|
||||
case "$MODE" in
|
||||
init) do_init ;;
|
||||
install-timer) install_timer ;;
|
||||
remove-timer) remove_timer ;;
|
||||
list) preflight; do_list ;;
|
||||
*) preflight; do_sync ;;
|
||||
esac
|
||||
+81
@@ -0,0 +1,81 @@
|
||||
#!/usr/bin/env bash
|
||||
# Import a LoRA (.safetensors) file into InvokeAI's Docker volume
|
||||
# Usage: ./invokeai-import-lora.sh /path/to/my-lora.safetensors
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
|
||||
if [[ $# -lt 1 ]]; then
|
||||
echo -e "${BOLD}Usage:${NC} $0 <lora-file.safetensors> [display-name]"
|
||||
echo ""
|
||||
echo " Copies a LoRA file into InvokeAI's model volume so it appears"
|
||||
echo " in the Model Manager automatically."
|
||||
echo ""
|
||||
echo " Examples:"
|
||||
echo " $0 ~/Downloads/my-character-lora.safetensors"
|
||||
echo " $0 ~/Downloads/my-character-lora.safetensors \"My Character\""
|
||||
exit 1
|
||||
fi
|
||||
|
||||
LORA_FILE="$1"
|
||||
DISPLAY_NAME="${2:-$(basename "$LORA_FILE" .safetensors)}"
|
||||
|
||||
# Validate file exists
|
||||
if [[ ! -f "$LORA_FILE" ]]; then
|
||||
echo -e "${RED}Error:${NC} File not found: $LORA_FILE"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Validate file extension
|
||||
if [[ "$LORA_FILE" != *.safetensors && "$LORA_FILE" != *.ckpt && "$LORA_FILE" != *.pt ]]; then
|
||||
echo -e "${YELLOW}Warning:${NC} File doesn't have a typical LoRA extension (.safetensors, .ckpt, .pt)"
|
||||
read -rp "Continue anyway? [y/N] " yn
|
||||
[[ "$yn" != [yY]* ]] && exit 1
|
||||
fi
|
||||
|
||||
# Check if InvokeAI container exists
|
||||
if ! docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^invokeai$'; then
|
||||
echo -e "${RED}Error:${NC} InvokeAI container not found. Is the AI stack running?"
|
||||
echo " Try: bash ~/docker/ai-stack/start.sh"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check if container is running
|
||||
if ! docker ps --format '{{.Names}}' 2>/dev/null | grep -q '^invokeai$'; then
|
||||
echo -e "${YELLOW}InvokeAI container is stopped. Starting it...${NC}"
|
||||
docker start invokeai
|
||||
sleep 3
|
||||
fi
|
||||
|
||||
FILENAME=$(basename "$LORA_FILE")
|
||||
|
||||
echo -e "${CYAN}[..]${NC} Copying ${BOLD}$FILENAME${NC} into InvokeAI..."
|
||||
|
||||
# Copy the LoRA file into the container's models directory
|
||||
# InvokeAI looks for LoRA files in /invokeai/models/lora/
|
||||
docker exec invokeai mkdir -p /invokeai/models/lora
|
||||
docker cp "$LORA_FILE" "invokeai:/invokeai/models/lora/$FILENAME"
|
||||
|
||||
echo -e "${GREEN}[OK]${NC} LoRA file copied successfully!"
|
||||
echo ""
|
||||
echo -e "${BOLD}Next steps in InvokeAI (http://localhost:9090):${NC}"
|
||||
echo ""
|
||||
echo " 1. Open the ${BOLD}Model Manager${NC} (cube icon in the left sidebar)"
|
||||
echo " 2. Click ${BOLD}\"Scan for Models\"${NC} or ${BOLD}\"Sync Models\"${NC} button"
|
||||
echo " - Your LoRA '${DISPLAY_NAME}' should appear in the list"
|
||||
echo " 3. If it doesn't auto-detect, click ${BOLD}\"Add Model\" > \"Scan Folder\"${NC}"
|
||||
echo " and enter: ${BOLD}/invokeai/models/lora${NC}"
|
||||
echo ""
|
||||
echo -e "${BOLD}To use the LoRA when generating images:${NC}"
|
||||
echo ""
|
||||
echo " 1. Go to the ${BOLD}Text to Image${NC} or ${BOLD}Image to Image${NC} tab"
|
||||
echo " 2. In the left panel, find the ${BOLD}\"LoRA\"${NC} section"
|
||||
echo " (expand it if collapsed — it's below the main model selector)"
|
||||
echo " 3. Click ${BOLD}\"+\"${NC} to add your LoRA from the dropdown"
|
||||
echo " 4. Adjust the ${BOLD}weight${NC} slider (start with 0.7–0.85)"
|
||||
echo " 5. Make sure your ${BOLD}base model${NC} matches what the LoRA was trained on"
|
||||
echo " (e.g., if trained on SD 1.5, select a SD 1.5 checkpoint)"
|
||||
echo ""
|
||||
echo -e "${YELLOW}Tip:${NC} If the LoRA was trained on SD 1.5, you MUST use an SD 1.5"
|
||||
echo " base model — it won't work with SDXL or other architectures."
|
||||
Vendored
+218
@@ -0,0 +1,218 @@
|
||||
#!/bin/bash
|
||||
# =============================================================================
|
||||
# Kiwix ZIM Download Script
|
||||
# Downloads all ZIM files to ~/docker/ai-stack/kiwix/
|
||||
# Finds latest version of each file automatically
|
||||
# Usage: bash kiwix-download.sh
|
||||
# =============================================================================
|
||||
|
||||
KIWIX_DIR="$HOME/docker/ai-stack/kiwix"
|
||||
MIRROR="https://ftp.fau.de/kiwix/zim"
|
||||
MIRROR2="https://download.kiwix.org/zim" # fallback for files missing on fau.de
|
||||
LOG="$HOME/docker/ai-stack/logs/kiwix-download.log"
|
||||
mkdir -p "$KIWIX_DIR" "$(dirname "$LOG")"
|
||||
|
||||
G='\033[0;32m'; Y='\033[1;33m'; C='\033[0;36m'; N='\033[0m'
|
||||
ok() { echo -e "${G}[OK]${N} $1" | tee -a "$LOG"; }
|
||||
inf() { echo -e "${C}[..]${N} $1" | tee -a "$LOG"; }
|
||||
wrn() { echo -e "${Y}[!!]${N} $1" | tee -a "$LOG"; }
|
||||
|
||||
# Find latest version of a ZIM file from mirror
|
||||
# Usage: latest_zim "base_url" "pattern"
|
||||
latest_zim() {
|
||||
local base_url="$1"
|
||||
local pattern="$2"
|
||||
curl -s "$base_url/" | \
|
||||
grep -oP "${pattern}_[\d-]+\.zim" | \
|
||||
sort -u | tail -1
|
||||
}
|
||||
|
||||
# Download a ZIM file if not already present
|
||||
# Usage: download_zim "category" "filename" "description" "size" ["mirror"]
|
||||
download_zim() {
|
||||
local category="$1"
|
||||
local filename="$2"
|
||||
local description="$3"
|
||||
local size="$4"
|
||||
local mirror="${5:-$MIRROR}"
|
||||
local url="$mirror/$category/$filename"
|
||||
local dest="$KIWIX_DIR/$filename"
|
||||
|
||||
if [[ -f "$dest" ]]; then
|
||||
ok "$description already downloaded — skipping"
|
||||
return
|
||||
fi
|
||||
|
||||
inf "Queuing: $description (~$size)"
|
||||
inf " URL: $url"
|
||||
|
||||
nohup wget -c "$url" -O "$dest" \
|
||||
>> "$LOG" 2>&1 &
|
||||
echo $! >> "$KIWIX_DIR/.download_pids"
|
||||
ok "Started download PID $! — $description"
|
||||
}
|
||||
|
||||
echo "" | tee -a "$LOG"
|
||||
echo "=== Kiwix Download Started: $(date) ===" | tee -a "$LOG"
|
||||
echo "" | tee -a "$LOG"
|
||||
echo -e "${C}Finding latest versions...${N}"
|
||||
echo ""
|
||||
|
||||
# ── Find latest filenames ──────────────────────────────────────────────────
|
||||
inf "Checking Wikipedia..."
|
||||
WIKI=$(latest_zim "$MIRROR/wikipedia" "wikipedia_en_all_nopic")
|
||||
|
||||
inf "Checking Wiktionary..."
|
||||
WIKT=$(latest_zim "$MIRROR/wiktionary" "wiktionary_en_all_nopic")
|
||||
|
||||
inf "Checking Wikiquote..."
|
||||
WIKQ=$(latest_zim "$MIRROR/wikiquote" "wikiquote_en_all_nopic")
|
||||
|
||||
inf "Checking Wikisource..."
|
||||
WIKS=$(latest_zim "$MIRROR/wikisource" "wikisource_en_all_nopic")
|
||||
|
||||
inf "Checking Wikibooks..."
|
||||
WIKB=$(latest_zim "$MIRROR/wikibooks" "wikibooks_en_all_nopic")
|
||||
|
||||
inf "Checking Wikivoyage..."
|
||||
WIKV=$(latest_zim "$MIRROR/wikivoyage" "wikivoyage_en_all_nopic")
|
||||
|
||||
inf "Checking Wikiversity..."
|
||||
WIKUNI=$(latest_zim "$MIRROR/wikiversity" "wikiversity_en_all_nopic")
|
||||
|
||||
inf "Checking WikiNews..."
|
||||
WIKNEWS=$(latest_zim "$MIRROR/wikinews" "wikinews_en_all_nopic")
|
||||
|
||||
inf "Checking Vikidia (kids K-8)..."
|
||||
VIKIDIA=$(latest_zim "$MIRROR/vikidia" "vikidia_en_all_nopic")
|
||||
|
||||
# Stack Exchange sites: domain-style filenames in stack_exchange/
|
||||
inf "Checking Stack Overflow..."
|
||||
SO=$(latest_zim "$MIRROR/stack_exchange" "stackoverflow.com_en_all")
|
||||
|
||||
inf "Checking Ask Ubuntu..."
|
||||
ASKUBUNTU=$(latest_zim "$MIRROR/stack_exchange" "askubuntu.com_en_all")
|
||||
|
||||
inf "Checking Super User..."
|
||||
SUPERUSER=$(latest_zim "$MIRROR/stack_exchange" "superuser.com_en_all")
|
||||
|
||||
inf "Checking Unix & Linux SE..."
|
||||
UNIX=$(latest_zim "$MIRROR/stack_exchange" "unix.stackexchange.com_en_all")
|
||||
|
||||
inf "Checking Server Fault..."
|
||||
SERVERFAULT=$(latest_zim "$MIRROR/stack_exchange" "serverfault.com_en_all")
|
||||
|
||||
# Arch Wiki: _maxi is part of the base name, not the date stamp
|
||||
inf "Checking Arch Wiki..."
|
||||
ARCH=$(latest_zim "$MIRROR/other" "archlinux_en_all_maxi")
|
||||
|
||||
# TED: lives in ted/ folder, mul_youth variant
|
||||
inf "Checking TED Talks..."
|
||||
TED=$(latest_zim "$MIRROR/ted" "ted_mul_youth")
|
||||
|
||||
# PhET: own folder
|
||||
inf "Checking PhET Simulations..."
|
||||
PHET=$(latest_zim "$MIRROR/phet" "phet_en_all")
|
||||
|
||||
# DevDocs: variant is zig not all
|
||||
inf "Checking DevDocs..."
|
||||
DEVDOCS=$(latest_zim "$MIRROR/devdocs" "devdocs_en_zig")
|
||||
|
||||
inf "Checking FreeCodeCamp..."
|
||||
FCC=$(latest_zim "$MIRROR/freecodecamp" "freecodecamp_en_all")
|
||||
|
||||
inf "Checking iFixit..."
|
||||
IFIX=$(latest_zim "$MIRROR/ifixit" "ifixit_en_all")
|
||||
|
||||
# LibreTexts: own folder, workforce variant (no _all on this mirror)
|
||||
inf "Checking LibreTexts..."
|
||||
LIBRE=$(latest_zim "$MIRROR/libretexts" "libretexts.org_en_workforce")
|
||||
|
||||
# Gutenberg: en_all exists on download.kiwix.org, not ftp.fau.de
|
||||
inf "Checking Project Gutenberg..."
|
||||
GUT=$(latest_zim "$MIRROR2/gutenberg" "gutenberg_en_all")
|
||||
|
||||
echo ""
|
||||
echo -e "${C}═══════════════════════════════════════════════════${N}"
|
||||
echo -e "${C} Download Queue${N}"
|
||||
echo -e "${C}═══════════════════════════════════════════════════${N}"
|
||||
echo ""
|
||||
echo " Wikipedia (no images) ${WIKI:-NOT FOUND} ~46GB"
|
||||
echo " Wiktionary ${WIKT:-NOT FOUND} ~2GB"
|
||||
echo " Wikiquote ${WIKQ:-NOT FOUND} ~300MB"
|
||||
echo " Wikisource ${WIKS:-NOT FOUND} ~4GB"
|
||||
echo " Wikibooks ${WIKB:-NOT FOUND} ~500MB"
|
||||
echo " Wikivoyage ${WIKV:-NOT FOUND} ~200MB"
|
||||
echo " Wikiversity ${WIKUNI:-NOT FOUND} ~500MB"
|
||||
echo " WikiNews ${WIKNEWS:-NOT FOUND} ~300MB"
|
||||
echo " Vikidia (kids K-8) ${VIKIDIA:-NOT FOUND} ~66MB"
|
||||
echo " Stack Overflow ${SO:-NOT FOUND} ~3GB"
|
||||
echo " Ask Ubuntu ${ASKUBUNTU:-NOT FOUND} ~1GB"
|
||||
echo " Super User ${SUPERUSER:-NOT FOUND} ~1GB"
|
||||
echo " Unix & Linux SE ${UNIX:-NOT FOUND} ~500MB"
|
||||
echo " Server Fault ${SERVERFAULT:-NOT FOUND} ~500MB"
|
||||
echo " Arch Linux Wiki ${ARCH:-NOT FOUND} ~30MB"
|
||||
echo " TED Talks ${TED:-NOT FOUND} ~5GB"
|
||||
echo " PhET Simulations ${PHET:-NOT FOUND} ~500MB"
|
||||
echo " DevDocs ${DEVDOCS:-NOT FOUND} ~1GB"
|
||||
echo " FreeCodeCamp ${FCC:-NOT FOUND} ~small"
|
||||
echo " iFixit (repair guides) ${IFIX:-NOT FOUND} ~2GB"
|
||||
echo " LibreTexts (textbooks) ${LIBRE:-NOT FOUND} ~varies"
|
||||
echo " Project Gutenberg ${GUT:-NOT FOUND} ~60GB"
|
||||
echo ""
|
||||
echo -e "${Y} Total estimate: ~133GB — make sure you have space!${N}"
|
||||
echo ""
|
||||
|
||||
df -h "$KIWIX_DIR" | tail -1 | awk '{print " Available disk space: " $4}'
|
||||
echo ""
|
||||
|
||||
read -rp " Proceed with all downloads? (Y/n): " CONFIRM
|
||||
[[ "${CONFIRM,,}" == "n" ]] && exit 0
|
||||
|
||||
# ── Start downloads ────────────────────────────────────────────────────────
|
||||
echo ""
|
||||
rm -f "$KIWIX_DIR/.download_pids"
|
||||
|
||||
[[ -n "$WIKI" ]] && download_zim "wikipedia" "$WIKI" "Wikipedia (no images)" "46GB"
|
||||
[[ -n "$WIKT" ]] && download_zim "wiktionary" "$WIKT" "Wiktionary" "2GB"
|
||||
[[ -n "$WIKQ" ]] && download_zim "wikiquote" "$WIKQ" "Wikiquote" "300MB"
|
||||
[[ -n "$WIKS" ]] && download_zim "wikisource" "$WIKS" "Wikisource" "4GB"
|
||||
[[ -n "$WIKB" ]] && download_zim "wikibooks" "$WIKB" "Wikibooks" "500MB"
|
||||
[[ -n "$WIKV" ]] && download_zim "wikivoyage" "$WIKV" "Wikivoyage" "200MB"
|
||||
[[ -n "$WIKUNI" ]] && download_zim "wikiversity" "$WIKUNI" "Wikiversity" "500MB"
|
||||
[[ -n "$WIKNEWS" ]] && download_zim "wikinews" "$WIKNEWS" "WikiNews" "300MB"
|
||||
[[ -n "$VIKIDIA" ]] && download_zim "vikidia" "$VIKIDIA" "Vikidia (kids K-8)" "66MB"
|
||||
[[ -n "$SO" ]] && download_zim "stack_exchange" "$SO" "Stack Overflow" "3GB"
|
||||
[[ -n "$ASKUBUNTU" ]] && download_zim "stack_exchange" "$ASKUBUNTU" "Ask Ubuntu" "1GB"
|
||||
[[ -n "$SUPERUSER" ]] && download_zim "stack_exchange" "$SUPERUSER" "Super User" "1GB"
|
||||
[[ -n "$UNIX" ]] && download_zim "stack_exchange" "$UNIX" "Unix & Linux SE" "500MB"
|
||||
[[ -n "$SERVERFAULT" ]] && download_zim "stack_exchange" "$SERVERFAULT" "Server Fault" "500MB"
|
||||
[[ -n "$ARCH" ]] && download_zim "other" "$ARCH" "Arch Linux Wiki" "30MB"
|
||||
[[ -n "$TED" ]] && download_zim "ted" "$TED" "TED Talks" "5GB"
|
||||
[[ -n "$PHET" ]] && download_zim "phet" "$PHET" "PhET Simulations" "500MB"
|
||||
[[ -n "$DEVDOCS" ]] && download_zim "devdocs" "$DEVDOCS" "DevDocs" "1GB"
|
||||
[[ -n "$FCC" ]] && download_zim "freecodecamp" "$FCC" "FreeCodeCamp" "small"
|
||||
[[ -n "$IFIX" ]] && download_zim "ifixit" "$IFIX" "iFixit (repair guides)" "2GB"
|
||||
[[ -n "$LIBRE" ]] && download_zim "libretexts" "$LIBRE" "LibreTexts (textbooks)" "varies"
|
||||
[[ -n "$GUT" ]] && download_zim "gutenberg" "$GUT" "Project Gutenberg" "60GB" "$MIRROR2"
|
||||
|
||||
echo ""
|
||||
echo -e "${G}════════════════════════════════════════════════════${N}"
|
||||
echo -e "${G} All downloads started in background${N}"
|
||||
echo -e "${G}════════════════════════════════════════════════════${N}"
|
||||
echo ""
|
||||
echo " Monitor progress:"
|
||||
echo " tail -f $LOG"
|
||||
echo ""
|
||||
echo " Check file sizes growing:"
|
||||
echo " watch -n 60 'ls -lh $KIWIX_DIR/'"
|
||||
echo ""
|
||||
echo " Check active downloads:"
|
||||
echo " jobs -l"
|
||||
echo " ps aux | grep wget"
|
||||
echo ""
|
||||
echo " Once all downloads complete, restart Kiwix:"
|
||||
echo " cd ~/docker/ai-stack"
|
||||
echo " docker compose up -d kiwix"
|
||||
echo ""
|
||||
echo -e "${Y} NOTE: Downloads resume automatically if interrupted (-c flag)${N}"
|
||||
+130
@@ -0,0 +1,130 @@
|
||||
#!/bin/bash
|
||||
# =============================================================================
|
||||
# Kiwix ZIM Auto-Updater
|
||||
# Checks each ZIM for a newer version on the mirror, downloads it, removes old.
|
||||
# Safe to run as a cron job — only acts when a newer version exists.
|
||||
# Usage: bash kiwix_update.sh
|
||||
# Cron example (monthly): 0 3 1 * * /bin/bash /path/to/kiwix_update.sh
|
||||
# =============================================================================
|
||||
|
||||
KIWIX_DIR="$HOME/docker/ai-stack/kiwix"
|
||||
MIRROR="https://ftp.fau.de/kiwix/zim"
|
||||
MIRROR2="https://download.kiwix.org/zim"
|
||||
LOG="$HOME/docker/ai-stack/logs/kiwix-update.log"
|
||||
mkdir -p "$KIWIX_DIR" "$(dirname "$LOG")"
|
||||
|
||||
G='\033[0;32m'; Y='\033[1;33m'; C='\033[0;36m'; R='\033[0;31m'; N='\033[0m'
|
||||
ok() { echo -e "${G}[OK]${N} $1" | tee -a "$LOG"; }
|
||||
inf() { echo -e "${C}[..]${N} $1" | tee -a "$LOG"; }
|
||||
wrn() { echo -e "${Y}[!!]${N} $1" | tee -a "$LOG"; }
|
||||
err() { echo -e "${R}[XX]${N} $1" | tee -a "$LOG"; }
|
||||
|
||||
echo "" | tee -a "$LOG"
|
||||
echo "=== Kiwix Update Check: $(date) ===" | tee -a "$LOG"
|
||||
echo "" | tee -a "$LOG"
|
||||
|
||||
# Find latest filename on mirror
|
||||
latest_zim() {
|
||||
local base_url="$1"
|
||||
local pattern="$2"
|
||||
curl -s "$base_url/" | \
|
||||
grep -oP "${pattern}_[\d-]+\.zim" | \
|
||||
sort -u | tail -1
|
||||
}
|
||||
|
||||
# Find what's currently on disk matching a pattern
|
||||
current_zim() {
|
||||
local pattern="$1"
|
||||
ls "$KIWIX_DIR"/${pattern}_*.zim 2>/dev/null | sort | tail -1 | xargs basename 2>/dev/null
|
||||
}
|
||||
|
||||
# Check and update one ZIM
|
||||
# Usage: check_update "category" "pattern" "description" ["mirror"]
|
||||
UPDATED=0
|
||||
check_update() {
|
||||
local category="$1"
|
||||
local pattern="$2"
|
||||
local description="$3"
|
||||
local mirror="${4:-$MIRROR}"
|
||||
|
||||
inf "Checking $description..."
|
||||
local latest
|
||||
latest=$(latest_zim "$mirror/$category" "$pattern")
|
||||
|
||||
if [[ -z "$latest" ]]; then
|
||||
wrn "$description: could not find latest on mirror — skipping"
|
||||
return
|
||||
fi
|
||||
|
||||
local current
|
||||
current=$(current_zim "$pattern")
|
||||
|
||||
if [[ -z "$current" ]]; then
|
||||
inf "$description: not downloaded yet — skipping (run kiwix_download.sh first)"
|
||||
return
|
||||
fi
|
||||
|
||||
if [[ "$latest" == "$current" ]]; then
|
||||
ok "$description: up to date ($current)"
|
||||
return
|
||||
fi
|
||||
|
||||
echo ""
|
||||
inf "$description: update available"
|
||||
inf " Current: $current"
|
||||
inf " Latest: $latest"
|
||||
|
||||
local url="$mirror/$category/$latest"
|
||||
local dest="$KIWIX_DIR/$latest"
|
||||
local old="$KIWIX_DIR/$current"
|
||||
|
||||
inf " Downloading $latest..."
|
||||
if wget -q --show-progress -c "$url" -O "$dest" >> "$LOG" 2>&1; then
|
||||
ok " Download complete — removing old file"
|
||||
rm -f "$old"
|
||||
UPDATED=$((UPDATED + 1))
|
||||
else
|
||||
err " Download failed — keeping old file"
|
||||
rm -f "$dest"
|
||||
fi
|
||||
echo ""
|
||||
}
|
||||
|
||||
# ── Check each ZIM ─────────────────────────────────────────────────────────
|
||||
check_update "wikipedia" "wikipedia_en_all_nopic" "Wikipedia"
|
||||
check_update "wiktionary" "wiktionary_en_all_nopic" "Wiktionary"
|
||||
check_update "wikiquote" "wikiquote_en_all_nopic" "Wikiquote"
|
||||
check_update "wikisource" "wikisource_en_all_nopic" "Wikisource"
|
||||
check_update "wikibooks" "wikibooks_en_all_nopic" "Wikibooks"
|
||||
check_update "wikivoyage" "wikivoyage_en_all_nopic" "Wikivoyage"
|
||||
check_update "wikiversity" "wikiversity_en_all_nopic" "Wikiversity"
|
||||
check_update "wikinews" "wikinews_en_all_nopic" "WikiNews"
|
||||
check_update "vikidia" "vikidia_en_all_nopic" "Vikidia"
|
||||
check_update "stack_exchange" "stackoverflow.com_en_all" "Stack Overflow"
|
||||
check_update "stack_exchange" "askubuntu.com_en_all" "Ask Ubuntu"
|
||||
check_update "stack_exchange" "superuser.com_en_all" "Super User"
|
||||
check_update "stack_exchange" "unix.stackexchange.com_en_all" "Unix & Linux SE"
|
||||
check_update "stack_exchange" "serverfault.com_en_all" "Server Fault"
|
||||
check_update "other" "archlinux_en_all_maxi" "Arch Wiki"
|
||||
check_update "ted" "ted_mul_youth" "TED Talks"
|
||||
check_update "phet" "phet_en_all" "PhET Simulations"
|
||||
check_update "devdocs" "devdocs_en_zig" "DevDocs"
|
||||
check_update "freecodecamp" "freecodecamp_en_all" "FreeCodeCamp"
|
||||
check_update "ifixit" "ifixit_en_all" "iFixit"
|
||||
check_update "libretexts" "libretexts.org_en_workforce" "LibreTexts"
|
||||
check_update "gutenberg" "gutenberg_en_all" "Project Gutenberg" "$MIRROR2"
|
||||
|
||||
# ── Restart Kiwix if anything changed ─────────────────────────────────────
|
||||
echo ""
|
||||
if [[ $UPDATED -gt 0 ]]; then
|
||||
ok "$UPDATED ZIM(s) updated — restarting Kiwix container"
|
||||
docker compose -f "$HOME/docker/ai-stack/docker-compose.yml" restart kiwix >> "$LOG" 2>&1 \
|
||||
&& ok "Kiwix restarted" \
|
||||
|| wrn "Could not restart Kiwix — do it manually: docker compose restart kiwix"
|
||||
else
|
||||
ok "All ZIMs are up to date — no restart needed"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "=== Update check complete: $(date) ===" | tee -a "$LOG"
|
||||
echo ""
|
||||
+1689
File diff suppressed because it is too large
Load Diff
+986
@@ -0,0 +1,986 @@
|
||||
#!/usr/bin/env bash
|
||||
# Local AI Stack — single script, new install or update
|
||||
# Usage: ./local-ai-setup.sh [--force] [--no-pull]
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
info() { echo -e "${CYAN}[..]${NC} $*"; }
|
||||
ok() { echo -e "${GREEN}[OK]${NC} $*"; }
|
||||
warn() { echo -e "${YELLOW}[!!]${NC} $*"; }
|
||||
section() { echo -e "\n${BOLD}━━━ $* ━━━${NC}"; }
|
||||
|
||||
FORCE=false; NO_PULL=false
|
||||
for a in "$@"; do [[ "$a" == "--force" ]] && FORCE=true; [[ "$a" == "--no-pull" ]] && NO_PULL=true; done
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
BASE="$SCRIPT_DIR"
|
||||
LOCAL_IP=$(ip route get 1.1.1.1 2>/dev/null | grep -oP 'src \K\S+' || hostname -I | awk '{print $1}')
|
||||
[[ -z "$LOCAL_IP" ]] && read -rp "Enter LAN IP: " LOCAL_IP
|
||||
|
||||
IS_UPDATE=false; [[ -f "$BASE/docker-compose.yml" ]] && IS_UPDATE=true
|
||||
|
||||
# ── detect VRAM and set models accordingly ────────────────────────────────────
|
||||
VRAM_GB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null \
|
||||
| head -1 | awk '{printf "%d", $1/1024}' 2>/dev/null || echo "0")
|
||||
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | wc -l || echo "0")
|
||||
TOTAL_VRAM=$((VRAM_GB * GPU_COUNT))
|
||||
|
||||
# Ollama optimization flags (stacked — see docs/gpu-setup-research.md)
|
||||
OLLAMA_KV_CACHE="q8_0" # halves KV cache VRAM (q4_0 for aggressive)
|
||||
OLLAMA_FLASH="1" # flash attention: less VRAM, no quality loss
|
||||
|
||||
if [[ "$TOTAL_VRAM" -ge 40 ]]; then
|
||||
CHAT_MODEL="qwen3.5:27b"; CODE_MODEL="qwen3.5:27b"
|
||||
CTX=131072; TIER="${TOTAL_VRAM}GB — 27B dense, 128K context"
|
||||
elif [[ "$TOTAL_VRAM" -ge 28 ]]; then
|
||||
CHAT_MODEL="qwen3.5-35b-a3b"; CODE_MODEL="qwen3.5-35b-a3b"
|
||||
CTX=131072; TIER="${TOTAL_VRAM}GB — 35B MoE, 128K context"
|
||||
elif [[ "$TOTAL_VRAM" -ge 14 ]]; then
|
||||
CHAT_MODEL="qwen3.5-35b-a3b"; CODE_MODEL="qwen3.5-35b-a3b"
|
||||
CTX=65536; TIER="${TOTAL_VRAM}GB — 35B MoE + KV quant, 64K context"
|
||||
elif [[ "$TOTAL_VRAM" -ge 8 ]]; then
|
||||
CHAT_MODEL="qwen3.5:9b"; CODE_MODEL="qwen3.5:9b"
|
||||
CTX=32768; TIER="${TOTAL_VRAM}GB — 9B dense, 32K context"
|
||||
elif [[ "$TOTAL_VRAM" -ge 6 ]]; then
|
||||
CHAT_MODEL="qwen3.5:4b"; CODE_MODEL="qwen3.5:4b"
|
||||
CTX=32768; TIER="${TOTAL_VRAM}GB — 4B + KV quant, 32K context"
|
||||
elif [[ "$TOTAL_VRAM" -ge 4 ]]; then
|
||||
CHAT_MODEL="qwen3.5:4b"; CODE_MODEL="qwen3.5:4b"
|
||||
CTX=16384; TIER="${TOTAL_VRAM}GB — 4B models, 16K context"
|
||||
else
|
||||
CHAT_MODEL="qwen3.5:4b"; CODE_MODEL="qwen3.5:4b"
|
||||
CTX=4096; OLLAMA_KV_CACHE="q4_0"; TIER="CPU-only — 4B models, 4K context"
|
||||
fi
|
||||
EMBED_MODEL="nomic-embed-text"
|
||||
|
||||
# ── Image generation model tiers (VRAM-aware) ────────────────────────────────
|
||||
# These vars are used by setup-image-models.sh and printed in status output.
|
||||
# Image gen shares GPU with Ollama — Ollama unloads after KEEP_ALIVE timeout,
|
||||
# so image gen gets full VRAM when Ollama is idle.
|
||||
if [[ "$TOTAL_VRAM" -ge 24 ]]; then
|
||||
IMG_MODELS="SD 1.5, SDXL, SDXL Turbo, Flux.1-dev, Flux.1-schnell"
|
||||
IMG_TIER="all models including Flux"
|
||||
IMG_DEFAULT="SDXL"
|
||||
elif [[ "$TOTAL_VRAM" -ge 12 ]]; then
|
||||
IMG_MODELS="SD 1.5, SDXL, SDXL Turbo, Flux.1-schnell (tight)"
|
||||
IMG_TIER="SDXL + Flux-schnell"
|
||||
IMG_DEFAULT="SDXL"
|
||||
elif [[ "$TOTAL_VRAM" -ge 8 ]]; then
|
||||
IMG_MODELS="SD 1.5, SDXL (tight at 512px), SDXL Turbo"
|
||||
IMG_TIER="SD 1.5 comfortable, SDXL possible"
|
||||
IMG_DEFAULT="SD 1.5"
|
||||
elif [[ "$TOTAL_VRAM" -ge 4 ]]; then
|
||||
IMG_MODELS="SD 1.5 (float16)"
|
||||
IMG_TIER="SD 1.5 only"
|
||||
IMG_DEFAULT="SD 1.5"
|
||||
else
|
||||
IMG_MODELS="none (CPU generation extremely slow)"
|
||||
IMG_TIER="CPU only — not recommended"
|
||||
IMG_DEFAULT=""
|
||||
fi
|
||||
|
||||
# ── InvokeAI precision (GPU-aware) ───────────────────────────────────────────
|
||||
GPU_COMPUTE=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \
|
||||
| head -1 | tr -d '.' || echo "0")
|
||||
if [[ "$GPU_COMPUTE" -ge 80 ]]; then
|
||||
INVOKEAI_PRECISION="bfloat16" # Ampere+ (RTX 30xx/40xx/50xx, A-series)
|
||||
elif [[ "$GPU_COMPUTE" -ge 60 ]]; then
|
||||
INVOKEAI_PRECISION="float16" # Pascal+ (GTX 10xx, RTX 20xx, Tesla P40/V100)
|
||||
else
|
||||
INVOKEAI_PRECISION="auto" # Let InvokeAI decide
|
||||
fi
|
||||
|
||||
section "Local AI Stack — $($IS_UPDATE && echo UPDATE || echo NEW INSTALL)"
|
||||
info "Base : $BASE"
|
||||
info "IP : $LOCAL_IP"
|
||||
info "GPU : ${VRAM_GB}GB VRAM → $TIER"
|
||||
info "Image : $IMG_TIER ($IMG_MODELS)"
|
||||
|
||||
|
||||
write_if_new() {
|
||||
local dest="$1"; local body; body=$(cat)
|
||||
if [[ ! -f "$dest" ]] || $FORCE; then
|
||||
printf '%s\n' "$body" > "$dest"; ok "Wrote $(basename "$dest")"
|
||||
else
|
||||
info "Kept $(basename "$dest") (--force to overwrite)"
|
||||
fi
|
||||
}
|
||||
|
||||
# ── prereqs (new install only) ────────────────────────────────────────────────
|
||||
if ! $IS_UPDATE; then
|
||||
section "Prerequisites"
|
||||
if ! command -v docker &>/dev/null; then
|
||||
info "Installing Docker..."
|
||||
curl -fsSL https://get.docker.com | sh
|
||||
sudo usermod -aG docker "$USER"
|
||||
warn "Run: newgrp docker (or log out/in)"
|
||||
else
|
||||
ok "Docker: $(docker --version | cut -d' ' -f3)"
|
||||
fi
|
||||
|
||||
if command -v nvidia-smi &>/dev/null && ! dpkg -l 2>/dev/null | grep -q nvidia-container-toolkit; then
|
||||
info "Installing NVIDIA Container Toolkit..."
|
||||
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
|
||||
| sudo gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
|
||||
curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
|
||||
| sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
|
||||
| sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
|
||||
sudo apt-get update -qq && sudo apt-get install -y nvidia-container-toolkit
|
||||
sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker
|
||||
ok "NVIDIA Container Toolkit installed"
|
||||
fi
|
||||
command -v rg &>/dev/null || sudo apt-get install -y ripgrep
|
||||
fi
|
||||
|
||||
# ── directories ───────────────────────────────────────────────────────────────
|
||||
section "Directories"
|
||||
for d in papers repos workspace index invokeai-data invokeai-outputs comfyui-data comfyui-output kiwix gitea portainer-data logs; do
|
||||
mkdir -p "$BASE/$d"
|
||||
done
|
||||
ok "Ready under $BASE"
|
||||
|
||||
# =============================================================================
|
||||
section "Writing server.py (RAG)"
|
||||
# =============================================================================
|
||||
write_if_new "$BASE/server.py" << 'PY'
|
||||
import ast, fnmatch, hashlib, json, logging, os, re, subprocess, threading, time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
import chromadb, httpx
|
||||
from chromadb.utils.embedding_functions import OllamaEmbeddingFunction
|
||||
from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("rag")
|
||||
|
||||
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://ollama:11434")
|
||||
CHROMA_URL = os.getenv("CHROMA_URL", "http://chromadb:8000")
|
||||
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
|
||||
CHAT_MODEL = os.getenv("CHAT_MODEL", "qwen2.5:14b")
|
||||
PAPERS_DIR = Path(os.getenv("PAPERS_DIR", "/papers"))
|
||||
REPOS_DIR = Path(os.getenv("REPOS_DIR", "/repos"))
|
||||
TOP_K = int(os.getenv("TOP_K", "6"))
|
||||
|
||||
CODE_EXTS = {".py",".js",".ts",".tsx",".jsx",".go",".rs",".java",".c",".cpp",
|
||||
".h",".cs",".rb",".sh",".yaml",".yml",".toml",".sql",".md"}
|
||||
SKIP_DIRS = {"node_modules",".git","__pycache__","dist","build",".venv","venv","target"}
|
||||
MAX_BYTES = 400_000
|
||||
|
||||
app = FastAPI(title="RAG Server")
|
||||
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
|
||||
|
||||
def _embed_fn():
|
||||
return OllamaEmbeddingFunction(url=f"{OLLAMA_URL}/api/embeddings", model_name=EMBED_MODEL)
|
||||
|
||||
def _chroma():
|
||||
host, port = CHROMA_URL.replace("http://","").split(":")
|
||||
return chromadb.HttpClient(host=host, port=int(port))
|
||||
|
||||
def get_col(name):
|
||||
return _chroma().get_or_create_collection(name, embedding_function=_embed_fn())
|
||||
|
||||
def _doc_id(text, key):
|
||||
return hashlib.md5(f"{key}|{text[:200]}".encode()).hexdigest()
|
||||
|
||||
def _sliding(text, size=1000, overlap=150):
|
||||
chunks, i = [], 0
|
||||
while i < len(text):
|
||||
chunks.append(text[i:i+size]); i += size - overlap
|
||||
return [c for c in chunks if c.strip()]
|
||||
|
||||
def _chunk_python(src):
|
||||
try: tree = ast.parse(src)
|
||||
except SyntaxError: return []
|
||||
lines = src.splitlines(); out = []
|
||||
for node in ast.iter_child_nodes(tree):
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
out.append((node.name, "\n".join(lines[node.lineno-1:node.end_lineno])[:4000]))
|
||||
return out
|
||||
|
||||
def _chunk_file(path, src):
|
||||
if path.suffix == ".py":
|
||||
pairs = _chunk_python(src)
|
||||
if pairs: return pairs
|
||||
pat = re.compile(r'(?:^|\n)(?=(?:export\s+)?(?:async\s+)?(?:function|class)|^func |^type |^impl |^pub fn |^fn )', re.M)
|
||||
parts = [p.strip() for p in pat.split(src) if p.strip()]
|
||||
if len(parts) > 1: return [(f"s{i}", p[:4000]) for i,p in enumerate(parts)]
|
||||
return [(f"c{i}", c) for i,c in enumerate(_sliding(src, 1200, 200))]
|
||||
|
||||
def ingest_file(col, fpath, repo=""):
|
||||
if fpath.stat().st_size > MAX_BYTES: return
|
||||
if any(fnmatch.fnmatch(fpath.name, p) for p in ("*.min.js","*.map","package-lock.json","yarn.lock")): return
|
||||
try: src = fpath.read_text(encoding="utf-8", errors="ignore")
|
||||
except: return
|
||||
if not src.strip(): return
|
||||
rel = str(fpath)
|
||||
pairs = _chunk_file(fpath, src) if fpath.suffix in CODE_EXTS else \
|
||||
[(f"c{i}",c) for i,c in enumerate(_sliding(src))]
|
||||
ids,docs,metas = [],[],[]
|
||||
for label,chunk in pairs:
|
||||
if not chunk.strip(): continue
|
||||
ids.append(_doc_id(chunk, rel+label)); docs.append(chunk)
|
||||
metas.append({"source":rel,"label":label,"repo":repo,"lang":fpath.suffix.lstrip(".")})
|
||||
if ids: col.upsert(ids=ids, documents=docs, metadatas=metas)
|
||||
|
||||
def ingest_dir(col, directory, repo=""):
|
||||
count = 0
|
||||
for f in directory.rglob("*"):
|
||||
if not f.is_file() or any(p in f.parts for p in SKIP_DIRS): continue
|
||||
ingest_file(col, f, repo); count += 1
|
||||
log.info("Indexed %d files from %s", count, directory); return count
|
||||
|
||||
def ingest_pdfs(col):
|
||||
try: import pypdf
|
||||
except ImportError: return 0
|
||||
n = 0
|
||||
for pdf in PAPERS_DIR.glob("*.pdf"):
|
||||
try:
|
||||
text = "\n".join(p.extract_text() or "" for p in pypdf.PdfReader(str(pdf)).pages)
|
||||
for i,chunk in enumerate(_sliding(text)):
|
||||
col.upsert(ids=[_doc_id(chunk,str(pdf)+str(i))], documents=[chunk],
|
||||
metadatas=[{"source":str(pdf),"label":f"p{i}","repo":"","lang":"pdf"}])
|
||||
n += 1
|
||||
except Exception as e: log.warning("PDF %s: %s", pdf.name, e)
|
||||
return n
|
||||
|
||||
def _startup():
|
||||
for _ in range(40):
|
||||
try:
|
||||
r = httpx.get(f"{OLLAMA_URL}/api/tags", timeout=5)
|
||||
if any(EMBED_MODEL in m["name"] for m in r.json().get("models",[])): break
|
||||
except: pass
|
||||
log.info("Waiting for embed model..."); time.sleep(5)
|
||||
code_col = get_col("code"); papers_col = get_col("papers")
|
||||
for d in REPOS_DIR.iterdir():
|
||||
if d.is_dir(): ingest_dir(code_col, d, d.name)
|
||||
ingest_pdfs(papers_col)
|
||||
for f in PAPERS_DIR.glob("*.txt"): ingest_file(papers_col, f)
|
||||
log.info("Startup index done")
|
||||
|
||||
@app.on_event("startup")
|
||||
async def on_startup(): threading.Thread(target=_startup, daemon=True).start()
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
try:
|
||||
cc = _chroma()
|
||||
return {"status":"ok",
|
||||
"code": cc.get_collection("code", embedding_function=_embed_fn()).count(),
|
||||
"papers": cc.get_collection("papers", embedding_function=_embed_fn()).count()}
|
||||
except Exception as e: return {"status":"error","detail":str(e)}
|
||||
|
||||
class RepoReq(BaseModel):
|
||||
url: str; name: str = ""; branch: str = "main"
|
||||
|
||||
@app.post("/ingest/repo")
|
||||
async def ingest_repo(req: RepoReq):
|
||||
name = req.name or req.url.rstrip("/").split("/")[-1].removesuffix(".git")
|
||||
dest = REPOS_DIR / name
|
||||
try:
|
||||
if dest.exists(): subprocess.run(["git","pull"], cwd=dest, check=True, timeout=120)
|
||||
else: subprocess.run(["git","clone","--depth=1","-b",req.branch,req.url,str(dest)], check=True, timeout=300)
|
||||
except subprocess.CalledProcessError as e: raise HTTPException(400, str(e))
|
||||
return {"status":"ok","repo":name,"files":ingest_dir(get_col("code"),dest,name)}
|
||||
|
||||
@app.post("/ingest/papers")
|
||||
async def trigger_papers(bg: BackgroundTasks):
|
||||
bg.add_task(ingest_pdfs, get_col("papers")); return {"status":"queued"}
|
||||
|
||||
async def _webhook(payload):
|
||||
repo = payload.get("repository") or {}
|
||||
url = repo.get("clone_url") or repo.get("html_url",""); name = repo.get("name","unknown")
|
||||
if not url: return {"status":"ignored"}
|
||||
dest = REPOS_DIR / name
|
||||
if dest.exists(): subprocess.run(["git","pull"], cwd=dest, timeout=120)
|
||||
else: subprocess.run(["git","clone","--depth=1",url,str(dest)], timeout=300)
|
||||
return {"status":"ok","repo":name,"files":ingest_dir(get_col("code"),dest,name)}
|
||||
|
||||
@app.post("/webhook/gitea")
|
||||
async def wh_gitea(r: Request): return await _webhook(await r.json())
|
||||
@app.post("/webhook/github")
|
||||
async def wh_github(r: Request): return await _webhook(await r.json())
|
||||
|
||||
def _ctx(query, cols):
|
||||
parts = []
|
||||
for cn in cols:
|
||||
try:
|
||||
col = get_col(cn)
|
||||
if col.count() == 0: continue
|
||||
res = col.query(query_texts=[query], n_results=min(TOP_K, col.count()))
|
||||
for doc,meta in zip(res["documents"][0], res["metadatas"][0]):
|
||||
parts.append(f"### {meta.get('source','')}:{meta.get('label','')}\n```{meta.get('lang','')}\n{doc}\n```")
|
||||
except Exception as e: log.warning("col %s: %s", cn, e)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
class ChatReq(BaseModel):
|
||||
model: str = CHAT_MODEL; messages: list[dict[str,Any]]
|
||||
stream: bool = False; collections: list[str] = ["code","papers"]
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat(req: ChatReq):
|
||||
query = next((m["content"] for m in reversed(req.messages) if m.get("role")=="user"), "")
|
||||
ctx = _ctx(query, req.collections)
|
||||
msgs = [{"role":"system","content":f"You are a coding assistant. Use context below.\n\n## Context\n{ctx}"}] + req.messages
|
||||
payload = {"model":req.model,"messages":msgs,"stream":req.stream}
|
||||
if req.stream:
|
||||
async def gen():
|
||||
async with httpx.AsyncClient(timeout=300) as c:
|
||||
async with c.stream("POST",f"{OLLAMA_URL}/v1/chat/completions",json=payload) as r:
|
||||
async for chunk in r.aiter_bytes(): yield chunk
|
||||
return StreamingResponse(gen(), media_type="text/event-stream")
|
||||
async with httpx.AsyncClient(timeout=300) as c:
|
||||
r = await c.post(f"{OLLAMA_URL}/v1/chat/completions", json=payload)
|
||||
return r.json()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn; uvicorn.run("server:app", host="0.0.0.0", port=8001, reload=False)
|
||||
PY
|
||||
|
||||
# =============================================================================
|
||||
section "Writing mcp_server.py"
|
||||
# =============================================================================
|
||||
write_if_new "$BASE/mcp_server.py" << 'PY'
|
||||
import json, os, re, subprocess
|
||||
from pathlib import Path
|
||||
import httpx
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
WORKSPACE = Path(os.getenv("WORKSPACE_DIR", "/workspace"))
|
||||
REPOS_DIR = Path(os.getenv("REPOS_DIR", "/repos"))
|
||||
GITEA_URL = os.getenv("GITEA_URL", "http://gitea:3000")
|
||||
GITEA_TOKEN = os.getenv("GITEA_TOKEN", "")
|
||||
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN","")
|
||||
RAG_URL = os.getenv("RAG_URL", "http://rag-server:8001")
|
||||
|
||||
mcp = FastMCP("local-dev-tools")
|
||||
|
||||
@mcp.tool()
|
||||
def bash(command: str, cwd: str = "") -> str:
|
||||
"""Run a shell command. Default cwd is /workspace."""
|
||||
work = Path(cwd) if cwd else WORKSPACE
|
||||
work.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
r = subprocess.run(command, shell=True, cwd=work, timeout=120, capture_output=True, text=True)
|
||||
out = r.stdout + (f"\n[stderr]\n{r.stderr}" if r.stderr else "")
|
||||
if r.returncode != 0: out += f"\n[exit {r.returncode}]"
|
||||
return out or "(no output)"
|
||||
except subprocess.TimeoutExpired: return "[timeout]"
|
||||
except Exception as e: return f"[error] {e}"
|
||||
|
||||
@mcp.tool()
|
||||
def read_file(path: str) -> str:
|
||||
"""Read a file. Absolute or relative to /workspace."""
|
||||
p = Path(path) if Path(path).is_absolute() else WORKSPACE / path
|
||||
if not p.exists(): return f"[not found] {p}"
|
||||
if p.stat().st_size > 500_000: return f"[too large: {p.stat().st_size//1024}KB]"
|
||||
return p.read_text(encoding="utf-8", errors="replace")
|
||||
|
||||
@mcp.tool()
|
||||
def write_file(path: str, content: str) -> str:
|
||||
"""Write content to a file. Relative to /workspace."""
|
||||
p = Path(path) if Path(path).is_absolute() else WORKSPACE / path
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(content, encoding="utf-8")
|
||||
return f"Wrote {len(content)} chars to {p}"
|
||||
|
||||
@mcp.tool()
|
||||
def list_files(path: str = "", pattern: str = "**/*") -> str:
|
||||
"""List files matching a glob pattern."""
|
||||
base = Path(path) if path else WORKSPACE
|
||||
if not base.exists(): return f"[not found] {base}"
|
||||
files = sorted(str(f.relative_to(base)) for f in base.glob(pattern) if f.is_file())
|
||||
return "\n".join(files[:500]) or "(empty)"
|
||||
|
||||
@mcp.tool()
|
||||
def search_code(query: str, path: str = "", glob: str = "", case_sensitive: bool = False) -> str:
|
||||
"""Search file contents with ripgrep."""
|
||||
base = path or str(WORKSPACE)
|
||||
cmd = ["rg", "--line-number", "--no-heading"]
|
||||
if not case_sensitive: cmd.append("-i")
|
||||
if glob: cmd += ["-g", glob]
|
||||
cmd += [query, base]
|
||||
try:
|
||||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||||
lines = r.stdout.strip().splitlines()
|
||||
if len(lines) > 200: lines = lines[:200] + [f"...({len(r.stdout.splitlines())-200} more)"]
|
||||
return "\n".join(lines) or "(no matches)"
|
||||
except FileNotFoundError:
|
||||
r = subprocess.run(["grep","-rn",query,base], capture_output=True, text=True, timeout=30)
|
||||
return r.stdout[:8000] or "(no matches)"
|
||||
|
||||
@mcp.tool()
|
||||
def fetch_url(url: str, extract_text: bool = True) -> str:
|
||||
"""Fetch the content of a URL."""
|
||||
try:
|
||||
r = httpx.get(url, timeout=30, follow_redirects=True, headers={"User-Agent":"Mozilla/5.0"})
|
||||
content = r.text
|
||||
if extract_text:
|
||||
content = re.sub(r'<script[^>]*>.*?</script>', '', content, flags=re.DOTALL)
|
||||
content = re.sub(r'<style[^>]*>.*?</style>', '', content, flags=re.DOTALL)
|
||||
content = re.sub(r'<[^>]+>', '', content)
|
||||
content = re.sub(r'\n{3,}', '\n\n', content).strip()
|
||||
return content[:20000]
|
||||
except Exception as e: return f"[error] {e}"
|
||||
|
||||
def _git(args, repo=""):
|
||||
cwd = Path(repo) if repo else WORKSPACE
|
||||
r = subprocess.run(["git"]+args, cwd=cwd, capture_output=True, text=True, timeout=60)
|
||||
return (r.stdout + r.stderr).strip() or "(no output)"
|
||||
|
||||
@mcp.tool()
|
||||
def git_status(repo: str = "") -> str:
|
||||
"""Show git status."""
|
||||
return _git(["status","--short"], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_diff(repo: str = "", cached: bool = False) -> str:
|
||||
"""Show git diff."""
|
||||
flag = ["--cached"] if cached else []
|
||||
return _git(["diff","--stat"]+flag, repo) + "\n\n" + _git(["diff"]+flag, repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_log(repo: str = "", n: int = 10) -> str:
|
||||
"""Show last n commits."""
|
||||
return _git(["log",f"-{n}","--oneline","--decorate"], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_commit(message: str, repo: str = "", add_all: bool = True) -> str:
|
||||
"""Stage all and commit."""
|
||||
if add_all: _git(["add","-A"], repo)
|
||||
return _git(["commit","-m",message], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_checkout(branch: str, repo: str = "", create: bool = False) -> str:
|
||||
"""Checkout or create a branch."""
|
||||
return _git(["checkout","-b",branch] if create else ["checkout",branch], repo)
|
||||
|
||||
def _gitea(method, path, body=None):
|
||||
if not GITEA_TOKEN: return {"error":"GITEA_TOKEN not set in .env"}
|
||||
r = httpx.request(method, f"{GITEA_URL}/api/v1{path}",
|
||||
json=body, headers={"Authorization":f"token {GITEA_TOKEN}"}, timeout=30)
|
||||
try: return r.json()
|
||||
except: return {"status":r.status_code,"text":r.text}
|
||||
|
||||
@mcp.tool()
|
||||
def gitea_list_repos() -> str:
|
||||
"""List your Gitea repos."""
|
||||
d = _gitea("GET", "/repos/search?limit=50")
|
||||
if "error" in d: return d["error"]
|
||||
return "\n".join(f"{r['full_name']} — {r.get('description','')}" for r in d.get("data",[]))
|
||||
|
||||
@mcp.tool()
|
||||
def gitea_create_repo(name: str, private: bool = True, description: str = "") -> str:
|
||||
"""Create a Gitea repo."""
|
||||
r = _gitea("POST","/user/repos",{"name":name,"private":private,"description":description,"auto_init":True,"default_branch":"main"})
|
||||
return r.get("html_url") or str(r)
|
||||
|
||||
@mcp.tool()
|
||||
def gitea_create_issue(repo: str, title: str, body: str = "") -> str:
|
||||
"""Create a Gitea issue (owner/repo)."""
|
||||
r = _gitea("POST",f"/repos/{repo}/issues",{"title":title,"body":body})
|
||||
return r.get("html_url") or str(r)
|
||||
|
||||
@mcp.tool()
|
||||
def github_api(method: str, endpoint: str, body: str = "") -> str:
|
||||
"""Call GitHub REST API. endpoint e.g. /repos/owner/repo/issues"""
|
||||
if not GITHUB_TOKEN: return "GITHUB_TOKEN not set in .env"
|
||||
r = httpx.request(method.upper(), f"https://api.github.com{endpoint}",
|
||||
json=json.loads(body) if body else None,
|
||||
headers={"Authorization":f"Bearer {GITHUB_TOKEN}","Accept":"application/vnd.github+json"},
|
||||
timeout=30)
|
||||
try: return json.dumps(r.json(), indent=2)
|
||||
except: return r.text
|
||||
|
||||
@mcp.tool()
|
||||
def ingest_repo(url: str, name: str = "", branch: str = "main") -> str:
|
||||
"""Clone a repo and index it in RAG."""
|
||||
r = httpx.post(f"{RAG_URL}/ingest/repo", json={"url":url,"name":name,"branch":branch}, timeout=300)
|
||||
return r.text
|
||||
|
||||
@mcp.tool()
|
||||
def rag_health() -> str:
|
||||
"""Check RAG server status."""
|
||||
try: return httpx.get(f"{RAG_URL}/health", timeout=10).text
|
||||
except Exception as e: return f"RAG unreachable: {e}"
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(mcp.sse_app(), host="0.0.0.0", port=8002)
|
||||
PY
|
||||
|
||||
# =============================================================================
|
||||
section "Requirements"
|
||||
# =============================================================================
|
||||
cat > "$BASE/requirements.txt" << 'REQ'
|
||||
fastapi
|
||||
uvicorn[standard]
|
||||
httpx
|
||||
pydantic
|
||||
chromadb
|
||||
pypdf
|
||||
python-multipart
|
||||
REQ
|
||||
|
||||
cat > "$BASE/mcp_requirements.txt" << 'REQ'
|
||||
mcp[cli]
|
||||
fastapi
|
||||
uvicorn[standard]
|
||||
httpx
|
||||
duckduckgo-search
|
||||
REQ
|
||||
ok "requirements.txt + mcp_requirements.txt"
|
||||
|
||||
# =============================================================================
|
||||
section ".env (tokens — never overwritten)"
|
||||
# =============================================================================
|
||||
if [[ ! -f "$BASE/.env" ]]; then
|
||||
cat > "$BASE/.env" << ENV
|
||||
# Local AI Stack — edit to add your API tokens
|
||||
GITEA_TOKEN=your-gitea-token-here
|
||||
GITHUB_TOKEN=your-github-token-here
|
||||
GITEA_URL=http://$LOCAL_IP:3001
|
||||
ENV
|
||||
ok "Created .env"
|
||||
else
|
||||
info "Kept .env"
|
||||
fi
|
||||
|
||||
# =============================================================================
|
||||
section "Docker Compose"
|
||||
# =============================================================================
|
||||
cat > "$BASE/docker-compose.yml" << COMPOSE
|
||||
# Local AI Stack — generated $(date '+%Y-%m-%d')
|
||||
# GPU: OLLAMA_NUM_GPU=999 uses all available VRAM automatically (V100/RTX/any)
|
||||
# Context: OLLAMA_NUM_CTX= set by detected VRAM (GB)
|
||||
#
|
||||
# ── Common commands (run from this folder) ─────────────────────────────────────
|
||||
# Start everything: docker compose up -d
|
||||
# Stop everything: docker compose down
|
||||
# Restart one service: docker compose restart <service>
|
||||
# Stop one service: docker compose stop <service>
|
||||
# Start one service: docker compose up -d <service>
|
||||
# Follow all logs: docker compose logs -f
|
||||
# Follow one service logs: docker compose logs -f <service>
|
||||
# Pull latest images: docker compose pull && docker compose up -d
|
||||
# Show status: docker compose ps
|
||||
#
|
||||
# Services: ollama open-webui chromadb rag-server mcp-server aider
|
||||
# kiwix gitea invokeai portainer
|
||||
# ───────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
services:
|
||||
|
||||
ollama:
|
||||
image: ollama/ollama:latest
|
||||
container_name: ollama
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:11434:11434"]
|
||||
volumes: [ollama-models:/root/.ollama]
|
||||
environment:
|
||||
- OLLAMA_NUM_GPU=999
|
||||
- OLLAMA_NUM_CTX=$CTX
|
||||
- OLLAMA_KEEP_ALIVE=24h
|
||||
- OLLAMA_MAX_LOADED_MODELS=1
|
||||
- OLLAMA_KV_CACHE_TYPE=$OLLAMA_KV_CACHE
|
||||
- OLLAMA_FLASH_ATTENTION=$OLLAMA_FLASH
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
healthcheck:
|
||||
test: ["CMD","ollama","list"]
|
||||
interval: 30s; timeout: 10s; retries: 5
|
||||
|
||||
open-webui:
|
||||
image: ghcr.io/open-webui/open-webui:main
|
||||
container_name: open-webui
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:3000:8080"]
|
||||
volumes: [open-webui-data:/app/backend/data]
|
||||
environment:
|
||||
- OLLAMA_BASE_URL=http://ollama:11434
|
||||
- OPENAI_API_BASE_URL=http://rag-server:8001/v1
|
||||
- OPENAI_API_KEY=local-rag
|
||||
- ENABLE_OPENAI_API=true
|
||||
- ENABLE_TOOL_SERVERS=true
|
||||
- WEBUI_AUTH=true
|
||||
- ENABLE_RAG_WEB_SEARCH=true
|
||||
- RAG_WEB_SEARCH_ENGINE=duckduckgo
|
||||
- ENABLE_IMAGE_GENERATION=true
|
||||
- IMAGE_GENERATION_ENGINE=comfyui
|
||||
- COMFYUI_BASE_URL=http://comfyui:8188
|
||||
depends_on:
|
||||
ollama: {condition: service_healthy}
|
||||
|
||||
chromadb:
|
||||
image: chromadb/chroma:latest
|
||||
container_name: chromadb
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:8000:8000"]
|
||||
volumes: [$BASE/index:/chroma/chroma]
|
||||
environment:
|
||||
- IS_PERSISTENT=TRUE
|
||||
- ANONYMIZED_TELEMETRY=FALSE
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL","wget -qO- http://localhost:8000/api/v2/heartbeat || exit 1"]
|
||||
interval: 15s; timeout: 5s; retries: 5
|
||||
|
||||
rag-server:
|
||||
image: python:3.11-slim
|
||||
container_name: rag-server
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:8001:8001"]
|
||||
volumes:
|
||||
- $BASE/papers:/papers
|
||||
- $BASE/repos:/repos
|
||||
- $BASE/index:/index
|
||||
- $BASE/server.py:/app/server.py
|
||||
- $BASE/requirements.txt:/app/requirements.txt
|
||||
working_dir: /app
|
||||
environment:
|
||||
- OLLAMA_URL=http://ollama:11434
|
||||
- CHROMA_URL=http://chromadb:8000
|
||||
- EMBED_MODEL=nomic-embed-text
|
||||
- CHAT_MODEL=$CHAT_MODEL
|
||||
command: >
|
||||
bash -c "apt-get update -qq && apt-get install -y --no-install-recommends git &&
|
||||
pip install --no-cache-dir -r requirements.txt &&
|
||||
uvicorn server:app --host 0.0.0.0 --port 8001"
|
||||
depends_on:
|
||||
chromadb: {condition: service_healthy}
|
||||
ollama: {condition: service_healthy}
|
||||
|
||||
mcp-server:
|
||||
image: python:3.11-slim
|
||||
container_name: mcp-server
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:8002:8002"]
|
||||
volumes:
|
||||
- $BASE/workspace:/workspace
|
||||
- $BASE/repos:/repos
|
||||
- $BASE/mcp_server.py:/app/mcp_server.py
|
||||
- $BASE/mcp_requirements.txt:/app/mcp_requirements.txt
|
||||
- $SCRIPT_DIR/gitea-github-sync.sh:/app/gitea-github-sync.sh:ro
|
||||
working_dir: /app
|
||||
env_file: $BASE/.env
|
||||
environment:
|
||||
- WORKSPACE_DIR=/workspace
|
||||
- REPOS_DIR=/repos
|
||||
- GITEA_URL=http://gitea:3000
|
||||
- RAG_URL=http://rag-server:8001
|
||||
- KIWIX_URL=http://kiwix:80
|
||||
command: >
|
||||
bash -c "apt-get update -qq && apt-get install -y --no-install-recommends git ripgrep curl &&
|
||||
pip install --no-cache-dir -r mcp_requirements.txt &&
|
||||
python mcp_server.py"
|
||||
depends_on: [rag-server, kiwix]
|
||||
|
||||
aider:
|
||||
image: paulgauthier/aider:latest
|
||||
container_name: aider
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:8080:8501"]
|
||||
volumes:
|
||||
- $BASE/workspace:/workspace
|
||||
- $BASE/repos:/repos
|
||||
working_dir: /workspace
|
||||
environment:
|
||||
- OLLAMA_API_BASE=http://ollama:11434
|
||||
- GIT_AUTHOR_NAME=aider
|
||||
- GIT_AUTHOR_EMAIL=aider@local
|
||||
- GIT_COMMITTER_NAME=aider
|
||||
- GIT_COMMITTER_EMAIL=aider@local
|
||||
command: >
|
||||
--gui --no-auto-commits --no-check-update
|
||||
--model ollama/$CODE_MODEL
|
||||
depends_on:
|
||||
ollama: {condition: service_healthy}
|
||||
|
||||
kiwix:
|
||||
image: ghcr.io/kiwix/kiwix-serve:latest
|
||||
container_name: kiwix
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:8181:80"]
|
||||
volumes: [$BASE/kiwix:/data]
|
||||
entrypoint: ["sh", "-c"]
|
||||
command: ["ls /data/*.zim >/dev/null 2>&1 && exec kiwix-serve /data/*.zim || { echo 'No ZIM files in /data yet - sleeping. Add .zim files and restart kiwix.'; exec sleep infinity; }"]
|
||||
|
||||
gitea:
|
||||
image: gitea/gitea:latest
|
||||
container_name: gitea
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:3001:3000","0.0.0.0:2222:22"]
|
||||
volumes:
|
||||
- $BASE/gitea:/data
|
||||
- /etc/timezone:/etc/timezone:ro
|
||||
- /etc/localtime:/etc/localtime:ro
|
||||
environment:
|
||||
- USER_UID=1000
|
||||
- USER_GID=1000
|
||||
- GITEA__database__DB_TYPE=sqlite3
|
||||
- GITEA__database__PATH=/data/gitea/gitea.db
|
||||
- GITEA__webhook__ALLOWED_HOST_LIST=rag-server,mcp-server
|
||||
|
||||
invokeai:
|
||||
image: ghcr.io/invoke-ai/invokeai:latest
|
||||
container_name: invokeai
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:9090:9090"]
|
||||
volumes:
|
||||
- invokeai-models:/invokeai/models
|
||||
- $BASE/invokeai-outputs:/invokeai/outputs
|
||||
- $BASE/invokeai-data:/invokeai/databases
|
||||
environment:
|
||||
- INVOKEAI_HOST=0.0.0.0
|
||||
- INVOKEAI_PORT=9090
|
||||
- INVOKEAI_PRECISION=$INVOKEAI_PRECISION
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
|
||||
comfyui:
|
||||
image: ghcr.io/ai-dock/comfyui:latest
|
||||
container_name: comfyui
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:8188:8188"]
|
||||
volumes:
|
||||
- comfyui-models:/opt/ComfyUI/models
|
||||
- $BASE/comfyui-output:/opt/ComfyUI/output
|
||||
- $BASE/comfyui-data:/opt/ComfyUI/custom_nodes
|
||||
environment:
|
||||
- CLI_ARGS=--listen 0.0.0.0
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
|
||||
portainer:
|
||||
image: portainer/portainer-ce:latest
|
||||
container_name: portainer
|
||||
restart: unless-stopped
|
||||
ports: ["0.0.0.0:9000:9000","0.0.0.0:9443:9443"]
|
||||
volumes:
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
- $BASE/portainer-data:/data
|
||||
|
||||
volumes:
|
||||
ollama-models:
|
||||
open-webui-data:
|
||||
invokeai-models:
|
||||
comfyui-models:
|
||||
COMPOSE
|
||||
ok "docker-compose.yml"
|
||||
|
||||
# =============================================================================
|
||||
section "Firewall"
|
||||
# =============================================================================
|
||||
if command -v ufw &>/dev/null && [[ ! -f "$BASE/.ufw-done" ]] || $FORCE; then
|
||||
read -rp " LAN subnet [192.168.1.0/24]: " LAN; LAN="${LAN:-192.168.1.0/24}"
|
||||
[[ "$LAN" =~ /[0-9]+$ ]] || LAN="${LAN}/24"
|
||||
for pc in "3000:Open WebUI" "11434:Ollama" "8001:RAG" "8002:MCP" \
|
||||
"8000:ChromaDB" "8080:Aider" "8181:Kiwix" \
|
||||
"3001:Gitea" "2222:Gitea SSH" "9090:InvokeAI" "8188:ComfyUI" \
|
||||
"9000:Portainer" "9443:Portainer S"; do
|
||||
sudo ufw allow from "$LAN" to any port "${pc%%:*}" proto tcp comment "${pc##*:}" >/dev/null
|
||||
done
|
||||
sudo ufw reload >/dev/null; ok "UFW rules set for $LAN"; touch "$BASE/.ufw-done"
|
||||
fi
|
||||
|
||||
# =============================================================================
|
||||
section "Helper Scripts"
|
||||
# =============================================================================
|
||||
cat > "$BASE/start.sh" << STARTSH
|
||||
#!/bin/bash
|
||||
cd "$BASE"
|
||||
docker compose pull --quiet 2>/dev/null
|
||||
docker compose up -d
|
||||
echo ""
|
||||
echo " Open WebUI → http://$LOCAL_IP:3000"
|
||||
echo " Aider UI → http://$LOCAL_IP:8080 (Claude Code-like editor)"
|
||||
echo " InvokeAI → http://$LOCAL_IP:9090 (inpainting, img2img)"
|
||||
echo " ComfyUI → http://$LOCAL_IP:8188 (chat-integrated image gen)"
|
||||
echo " Kiwix → http://$LOCAL_IP:8181 (run kiwix_download.sh first)"
|
||||
echo " Gitea → http://$LOCAL_IP:3001"
|
||||
echo " RAG → http://$LOCAL_IP:8001/health"
|
||||
echo " MCP SSE → http://$LOCAL_IP:8002/sse"
|
||||
echo " Portainer → https://$LOCAL_IP:9443"
|
||||
echo ""
|
||||
echo " claude mcp add local http://$LOCAL_IP:8002/sse"
|
||||
STARTSH
|
||||
chmod +x "$BASE/start.sh"
|
||||
|
||||
cat > "$BASE/stop.sh" << STOPSH
|
||||
#!/bin/bash
|
||||
cd "$BASE" && docker compose down
|
||||
STOPSH
|
||||
chmod +x "$BASE/stop.sh"
|
||||
|
||||
cat > "$BASE/status.sh" << 'STATUSSH'
|
||||
#!/bin/bash
|
||||
echo "=== GPU ===" && nvidia-smi --query-gpu=name,memory.used,memory.total \
|
||||
--format=csv,noheader 2>/dev/null || echo "(no GPU)"
|
||||
echo "" && echo "=== Containers ===" && docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"
|
||||
echo "" && echo "=== Ollama ===" && docker exec ollama ollama ps 2>/dev/null || echo "(not running)"
|
||||
echo "" && echo "=== RAG ===" && curl -s http://localhost:8001/health | python3 -m json.tool 2>/dev/null
|
||||
STATUSSH
|
||||
chmod +x "$BASE/status.sh"
|
||||
|
||||
cat > "$BASE/pull-models.sh" << PULLSH
|
||||
#!/bin/bash
|
||||
echo "Waiting for Ollama..."
|
||||
until docker exec ollama ollama list &>/dev/null; do sleep 3; done
|
||||
|
||||
echo "Embed model (RAG — required)..."
|
||||
docker exec ollama ollama pull $EMBED_MODEL
|
||||
|
||||
echo "Fast chat model..."
|
||||
docker exec ollama ollama pull qwen2.5:7b
|
||||
|
||||
echo "Smart chat model..."
|
||||
docker exec ollama ollama pull $CHAT_MODEL
|
||||
|
||||
echo "Code model..."
|
||||
docker exec ollama ollama pull $CODE_MODEL
|
||||
|
||||
echo "Reasoning model (DeepSeek-R1 14B — optional)..."
|
||||
read -rp "Pull DeepSeek-R1:14b for planning/reasoning? [y/N]: " DR
|
||||
[[ "\${DR,,}" == "y" ]] && docker exec ollama ollama pull deepseek-r1:14b
|
||||
|
||||
echo "" && docker exec ollama ollama list
|
||||
PULLSH
|
||||
chmod +x "$BASE/pull-models.sh"
|
||||
ok "start.sh stop.sh status.sh pull-models.sh"
|
||||
|
||||
cat > "$BASE/aider.sh" << AIDSH
|
||||
#!/bin/bash
|
||||
# Aider CLI — Claude Code-like terminal experience against any local git repo.
|
||||
# Usage:
|
||||
# ./aider.sh # interactive, files from stdin
|
||||
# ./aider.sh src/main.py # open specific files
|
||||
# ./aider.sh --model ollama/qwen2.5-coder:7b src/foo.py # override model
|
||||
#
|
||||
# Tip: clone Gitea repos into $BASE/repos/, then:
|
||||
# cd $BASE/repos/my-project && $BASE/aider.sh <files>
|
||||
|
||||
MODEL="\${AIDER_MODEL:-ollama/$CODE_MODEL}"
|
||||
REPO_ROOT="\$(git rev-parse --show-toplevel 2>/dev/null || pwd)"
|
||||
|
||||
docker run --rm -it \\
|
||||
--network local-ai_default \\
|
||||
-v "\$REPO_ROOT:\$REPO_ROOT" \\
|
||||
-v "$BASE/repos:/repos" \\
|
||||
-w "\$REPO_ROOT" \\
|
||||
-e OLLAMA_API_BASE=http://ollama:11434 \\
|
||||
-e GIT_AUTHOR_NAME=aider \\
|
||||
-e GIT_AUTHOR_EMAIL=aider@local \\
|
||||
-e GIT_COMMITTER_NAME=aider \\
|
||||
-e GIT_COMMITTER_EMAIL=aider@local \\
|
||||
paulgauthier/aider:latest \\
|
||||
--no-auto-commits \\
|
||||
--no-check-update \\
|
||||
--model "\$MODEL" \\
|
||||
"\$@"
|
||||
AIDSH
|
||||
chmod +x "$BASE/aider.sh"
|
||||
ok "aider.sh"
|
||||
|
||||
# =============================================================================
|
||||
section "Systemd"
|
||||
# =============================================================================
|
||||
sudo tee /etc/systemd/system/local-ai.service >/dev/null << SYSD
|
||||
[Unit]
|
||||
Description=Local AI Stack
|
||||
After=docker.service network-online.target
|
||||
Requires=docker.service
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
RemainAfterExit=yes
|
||||
User=$USER
|
||||
WorkingDirectory=$BASE
|
||||
ExecStart=/bin/bash $BASE/start.sh
|
||||
ExecStop=/bin/bash $BASE/stop.sh
|
||||
TimeoutStartSec=300
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
SYSD
|
||||
sudo systemctl daemon-reload && sudo systemctl enable local-ai.service
|
||||
ok "Systemd: local-ai.service enabled"
|
||||
|
||||
# =============================================================================
|
||||
section "Starting Stack"
|
||||
# =============================================================================
|
||||
cd "$BASE"
|
||||
info "Pulling images..."
|
||||
docker compose pull --quiet
|
||||
docker compose up -d
|
||||
ok "Stack running"
|
||||
|
||||
if ! $IS_UPDATE && ! $NO_PULL; then
|
||||
echo ""
|
||||
read -rp "Pull Ollama models now? (~15-30 min) [Y/n]: " DO_PULL
|
||||
[[ "${DO_PULL,,}" != "n" ]] && bash "$BASE/pull-models.sh"
|
||||
|
||||
# Install image generation base model (GPU-aware)
|
||||
if [[ -n "$IMG_DEFAULT" ]] && [[ -x "$SCRIPT_DIR/setup-image-models.sh" ]]; then
|
||||
echo ""
|
||||
section "Image Generation Models"
|
||||
info "GPU: ${TOTAL_VRAM}GB VRAM → $IMG_TIER"
|
||||
info "Auto-installing recommended model: $IMG_DEFAULT"
|
||||
bash "$SCRIPT_DIR/setup-image-models.sh" --auto
|
||||
fi
|
||||
fi
|
||||
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo -e "${GREEN}${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo -e "${GREEN}${BOLD} Done! GPU: ${VRAM_GB}GB → $TIER${NC}"
|
||||
echo -e "${GREEN}${BOLD} Image gen: $IMG_TIER${NC}"
|
||||
echo -e "${GREEN}${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo ""
|
||||
echo -e " ${CYAN}Open WebUI${NC} → http://$LOCAL_IP:3000"
|
||||
echo -e " ${CYAN}Aider UI${NC} → http://$LOCAL_IP:8080 (browser coding assistant)"
|
||||
echo -e " ${CYAN}InvokeAI${NC} → http://$LOCAL_IP:9090 (inpainting, img2img)"
|
||||
echo -e " ${CYAN}ComfyUI${NC} → http://$LOCAL_IP:8188 (chat-integrated image gen)"
|
||||
echo -e " ${CYAN}Kiwix${NC} → http://$LOCAL_IP:8181"
|
||||
echo -e " ${CYAN}Gitea${NC} → http://$LOCAL_IP:3001"
|
||||
echo -e " ${CYAN}RAG${NC} → http://$LOCAL_IP:8001/health"
|
||||
echo -e " ${CYAN}MCP SSE${NC} → http://$LOCAL_IP:8002/sse"
|
||||
echo -e " ${CYAN}Portainer${NC} → https://$LOCAL_IP:9443"
|
||||
echo ""
|
||||
echo -e " ${YELLOW}Add MCP to Claude Code:${NC}"
|
||||
echo " claude mcp add local http://$LOCAL_IP:8002/sse"
|
||||
echo ""
|
||||
echo -e " ${YELLOW}Tokens:${NC} $BASE/.env"
|
||||
echo -e " ${YELLOW}PDFs:${NC} $BASE/papers/"
|
||||
echo -e " ${YELLOW}Workspace:${NC} $BASE/workspace/"
|
||||
echo -e " ${YELLOW}ZIMs:${NC} ./kiwix_download.sh"
|
||||
echo ""
|
||||
echo -e " ${YELLOW}Aider CLI:${NC} cd your-repo && $BASE/aider.sh <files>"
|
||||
echo " (or open browser UI above — both use your local code model)"
|
||||
echo ""
|
||||
echo -e " ${YELLOW}Recommended Open WebUI Functions${NC} (install from Admin → Functions → ＋):"
|
||||
echo " Context tracker: https://openwebui.com/f/centrisic/context_tracker"
|
||||
echo " → Shows tokens used vs available, progress bar, context % remaining"
|
||||
echo " Context compaction: https://openwebui.com/f/projectmoon/checkpoint_summarization_filter"
|
||||
echo " → Auto-summarizes old messages when context fills up (like Claude)"
|
||||
echo ""
|
||||
echo -e " ${YELLOW}Save Claude usage:${NC} use local models for boilerplate, docs,"
|
||||
echo " simple fixes. Use Claude for hard bugs,"
|
||||
echo " multi-file refactoring, architecture decisions."
|
||||
echo ""
|
||||
Vendored
+366
@@ -0,0 +1,366 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP Server — Claude Code-equivalent tools for Open WebUI / Claude Code CLI.
|
||||
Tools: bash, file read/write/list, code search, git ops, Gitea API, repo ingest,
|
||||
offline doc search (Kiwix), web search (DuckDuckGo), Gitea↔GitHub sync.
|
||||
Connects via SSE on port 8002 — add to Open WebUI Tools or ~/.claude/mcp.json
|
||||
"""
|
||||
import os, subprocess, textwrap
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
WORKSPACE = Path(os.getenv("WORKSPACE_DIR", "/workspace"))
|
||||
REPOS_DIR = Path(os.getenv("REPOS_DIR", "/repos"))
|
||||
GITEA_URL = os.getenv("GITEA_URL", "http://gitea:3000")
|
||||
GITEA_TOKEN = os.getenv("GITEA_TOKEN", "")
|
||||
GITHUB_TOKEN= os.getenv("GITHUB_TOKEN","")
|
||||
RAG_URL = os.getenv("RAG_URL", "http://rag-server:8001")
|
||||
KIWIX_URL = os.getenv("KIWIX_URL", "http://kiwix:80")
|
||||
|
||||
mcp = FastMCP("local-dev-tools")
|
||||
|
||||
# ── bash ──────────────────────────────────────────────────────────────────────
|
||||
@mcp.tool()
|
||||
def bash(command: str, cwd: str = "") -> str:
|
||||
"""Run a shell command. Default cwd is /workspace."""
|
||||
work = Path(cwd) if cwd else WORKSPACE
|
||||
work.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
r = subprocess.run(command, shell=True, cwd=work, timeout=120,
|
||||
capture_output=True, text=True)
|
||||
out = r.stdout + (f"\n[stderr]\n{r.stderr}" if r.stderr else "")
|
||||
if r.returncode != 0:
|
||||
out += f"\n[exit {r.returncode}]"
|
||||
return out or "(no output)"
|
||||
except subprocess.TimeoutExpired:
|
||||
return "[timeout after 120s]"
|
||||
except Exception as e:
|
||||
return f"[error] {e}"
|
||||
|
||||
# ── file ops ──────────────────────────────────────────────────────────────────
|
||||
@mcp.tool()
|
||||
def read_file(path: str) -> str:
|
||||
"""Read a file. Use absolute path or relative to /workspace."""
|
||||
p = Path(path) if Path(path).is_absolute() else WORKSPACE / path
|
||||
if not p.exists():
|
||||
return f"[not found] {p}"
|
||||
if p.stat().st_size > 500_000:
|
||||
return f"[too large — {p.stat().st_size//1024}KB]"
|
||||
return p.read_text(encoding="utf-8", errors="replace")
|
||||
|
||||
@mcp.tool()
|
||||
def write_file(path: str, content: str) -> str:
|
||||
"""Write content to a file (creates parent dirs). Relative to /workspace."""
|
||||
p = Path(path) if Path(path).is_absolute() else WORKSPACE / path
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(content, encoding="utf-8")
|
||||
return f"Wrote {len(content)} chars to {p}"
|
||||
|
||||
@mcp.tool()
|
||||
def list_files(path: str = "", pattern: str = "**/*") -> str:
|
||||
"""List files matching a glob pattern."""
|
||||
base = Path(path) if path else WORKSPACE
|
||||
if not base.exists():
|
||||
return f"[not found] {base}"
|
||||
files = sorted(str(f.relative_to(base)) for f in base.glob(pattern) if f.is_file())
|
||||
return "\n".join(files[:500]) or "(empty)"
|
||||
|
||||
@mcp.tool()
|
||||
def search_code(query: str, path: str = "", glob: str = "",
|
||||
case_sensitive: bool = False) -> str:
|
||||
"""Search file contents with ripgrep. Returns file:line matches."""
|
||||
base = path or str(WORKSPACE)
|
||||
cmd = ["rg", "--line-number", "--no-heading"]
|
||||
if not case_sensitive:
|
||||
cmd.append("-i")
|
||||
if glob:
|
||||
cmd += ["-g", glob]
|
||||
cmd += [query, base]
|
||||
try:
|
||||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||||
lines = r.stdout.strip().splitlines()
|
||||
if len(lines) > 200:
|
||||
lines = lines[:200] + [f"… ({len(r.stdout.splitlines())-200} more)"]
|
||||
return "\n".join(lines) or "(no matches)"
|
||||
except FileNotFoundError:
|
||||
# ripgrep not installed, fall back to grep
|
||||
r = subprocess.run(["grep", "-rn", query, base],
|
||||
capture_output=True, text=True, timeout=30)
|
||||
return r.stdout[:8000] or "(no matches)"
|
||||
except Exception as e:
|
||||
return f"[error] {e}"
|
||||
|
||||
# ── git ───────────────────────────────────────────────────────────────────────
|
||||
def _git(args: list[str], repo: str = "") -> str:
|
||||
cwd = Path(repo) if repo else WORKSPACE
|
||||
r = subprocess.run(["git"] + args, cwd=cwd,
|
||||
capture_output=True, text=True, timeout=60)
|
||||
return (r.stdout + r.stderr).strip() or "(no output)"
|
||||
|
||||
@mcp.tool()
|
||||
def git_status(repo: str = "") -> str:
|
||||
"""Show git status of a repo (default: /workspace)."""
|
||||
return _git(["status", "--short"], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_diff(repo: str = "", cached: bool = False) -> str:
|
||||
"""Show git diff (staged if cached=True)."""
|
||||
args = ["diff", "--stat", "--cached"] if cached else ["diff", "--stat"]
|
||||
return _git(args, repo) + "\n\n" + _git(
|
||||
["diff", "--cached"] if cached else ["diff"], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_log(repo: str = "", n: int = 10) -> str:
|
||||
"""Show last n git commits."""
|
||||
return _git(["log", f"-{n}", "--oneline", "--decorate"], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_commit(message: str, repo: str = "", add_all: bool = True) -> str:
|
||||
"""Stage all changes and create a commit."""
|
||||
if add_all:
|
||||
_git(["add", "-A"], repo)
|
||||
return _git(["commit", "-m", message], repo)
|
||||
|
||||
@mcp.tool()
|
||||
def git_checkout(branch: str, repo: str = "", create: bool = False) -> str:
|
||||
"""Checkout a branch, optionally creating it."""
|
||||
args = ["checkout", "-b", branch] if create else ["checkout", branch]
|
||||
return _git(args, repo)
|
||||
|
||||
# ── Search: unified (Kiwix offline + DuckDuckGo live) ───────────────────────
|
||||
# Kiwix ZIMs have complete, high-quality articles but may be months old.
|
||||
# DDG has live results but lower signal-to-noise. The unified search tool
|
||||
# checks both and lets the model see freshness info to judge which to trust.
|
||||
#
|
||||
# Heuristic: topics that change fast (releases, CVEs, "latest X") get flagged
|
||||
# as potentially stale in offline results. Timeless topics (algorithms, language
|
||||
# docs, math) are fine from Kiwix and skip the web hit entirely.
|
||||
|
||||
import re as _re
|
||||
from datetime import datetime as _dt
|
||||
|
||||
# Words that suggest the query needs fresh data
|
||||
_FRESH_KEYWORDS = _re.compile(
|
||||
r'\b(latest|newest|recent|2025|2026|update|release|version|changelog|CVE|vulnerability|'
|
||||
r'breaking change|deprecat|current|today|this year|this month|announce|just released)\b',
|
||||
_re.IGNORECASE
|
||||
)
|
||||
|
||||
def _kiwix_search(query: str, limit: int = 5) -> list[dict]:
|
||||
"""Search Kiwix, return list of {title, snippet, path, source}."""
|
||||
try:
|
||||
r = httpx.get(f"{KIWIX_URL}/search",
|
||||
params={"pattern": query, "pageLength": limit},
|
||||
timeout=15, follow_redirects=True)
|
||||
if r.status_code != 200:
|
||||
return []
|
||||
html = r.text
|
||||
results = []
|
||||
# Try structured parse first
|
||||
articles = _re.findall(
|
||||
r'<a[^>]+href="(/[^"]+)"[^>]*>\s*<span[^>]*>([^<]*)</span>.*?'
|
||||
r'(?:<cite[^>]*>([^<]*)</cite>)?.*?'
|
||||
r'(?:<p[^>]*>(.*?)</p>)?',
|
||||
html, _re.DOTALL
|
||||
)
|
||||
if articles:
|
||||
for path, title, cite, snippet in articles[:limit]:
|
||||
snippet_clean = _re.sub(r'<[^>]+>', '', snippet or '').strip()[:300]
|
||||
results.append({"title": title.strip(), "snippet": snippet_clean,
|
||||
"path": path, "source": cite.strip() if cite else "kiwix"})
|
||||
else:
|
||||
# Fallback: grab any links
|
||||
for path, title in _re.findall(r'<a[^>]+href="(/[^"]+)"[^>]*>([^<]+)</a>', html)[:limit]:
|
||||
results.append({"title": title.strip(), "snippet": "",
|
||||
"path": path, "source": "kiwix"})
|
||||
return results
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def _ddg_search(query: str, limit: int = 5) -> list[dict]:
|
||||
"""Search DuckDuckGo, return list of {title, snippet, url}."""
|
||||
try:
|
||||
from duckduckgo_search import DDGS
|
||||
results = []
|
||||
with DDGS() as ddgs:
|
||||
for r in ddgs.text(query, max_results=limit):
|
||||
results.append({"title": r["title"], "snippet": r["body"], "url": r["href"]})
|
||||
return results
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
@mcp.tool()
|
||||
def search(query: str, limit: int = 5) -> str:
|
||||
"""Unified search: checks offline docs (Kiwix) AND live web (DuckDuckGo).
|
||||
Returns results from both with freshness guidance.
|
||||
For timeless topics (algorithms, docs): offline results are sufficient.
|
||||
For time-sensitive topics (releases, CVEs): live results are flagged as preferred."""
|
||||
needs_fresh = bool(_FRESH_KEYWORDS.search(query))
|
||||
output_parts = []
|
||||
|
||||
# Always search Kiwix (fast, local)
|
||||
kiwix_results = _kiwix_search(query, limit)
|
||||
if kiwix_results:
|
||||
header = "## Offline Docs (Kiwix)"
|
||||
if needs_fresh:
|
||||
header += " ⚠️ POSSIBLY STALE — query looks time-sensitive, prefer live results below"
|
||||
output_parts.append(header)
|
||||
for i, r in enumerate(kiwix_results, 1):
|
||||
entry = f"{i}. **{r['title']}**"
|
||||
if r["source"] and r["source"] != "kiwix":
|
||||
entry += f" ({r['source']})"
|
||||
if r["snippet"]:
|
||||
entry += f"\n {r['snippet']}"
|
||||
entry += f"\n → read_doc('{r['path']}')"
|
||||
output_parts.append(entry)
|
||||
|
||||
# Search DDG if: query needs fresh data, OR Kiwix returned nothing, OR always (to compare)
|
||||
do_web = needs_fresh or not kiwix_results
|
||||
ddg_results = []
|
||||
if do_web:
|
||||
ddg_results = _ddg_search(query, limit)
|
||||
|
||||
if ddg_results:
|
||||
header = "## Live Web (DuckDuckGo)"
|
||||
if needs_fresh:
|
||||
header += " ✓ PREFER THESE for this query"
|
||||
output_parts.append(header)
|
||||
for i, r in enumerate(ddg_results, 1):
|
||||
output_parts.append(f"{i}. **{r['title']}**\n {r['snippet']}\n {r['url']}")
|
||||
elif do_web:
|
||||
output_parts.append("## Live Web (DuckDuckGo)\n(no results or DDG unreachable)")
|
||||
|
||||
if not kiwix_results and not ddg_results:
|
||||
return f"No results for '{query}' from either offline docs or web search."
|
||||
|
||||
# Freshness note
|
||||
if kiwix_results and not needs_fresh and not ddg_results:
|
||||
output_parts.append("\n_Offline results look sufficient for this topic. "
|
||||
"Use web_search() if you need to verify currency._")
|
||||
|
||||
return "\n\n".join(output_parts)
|
||||
|
||||
@mcp.tool()
|
||||
def read_doc(path: str) -> str:
|
||||
"""Read a full article from Kiwix by its path (from search results).
|
||||
Example: read_doc('/wikipedia_en_all/A/Python_(programming_language)')"""
|
||||
try:
|
||||
r = httpx.get(f"{KIWIX_URL}{path}", timeout=15, follow_redirects=True)
|
||||
if r.status_code != 200:
|
||||
return f"Not found: {path} (HTTP {r.status_code})"
|
||||
# Strip HTML tags, keep text content
|
||||
text = _re.sub(r'<script[^>]*>.*?</script>', '', r.text, flags=_re.DOTALL)
|
||||
text = _re.sub(r'<style[^>]*>.*?</style>', '', text, flags=_re.DOTALL)
|
||||
text = _re.sub(r'<[^>]+>', ' ', text)
|
||||
text = _re.sub(r'\s+', ' ', text).strip()
|
||||
if len(text) > 8000:
|
||||
text = text[:8000] + "\n\n[... truncated — article continues ...]"
|
||||
return text
|
||||
except Exception as e:
|
||||
return f"Error reading doc: {e}"
|
||||
|
||||
@mcp.tool()
|
||||
def web_search(query: str, num_results: int = 5) -> str:
|
||||
"""Search ONLY the live web via DuckDuckGo. Use search() instead for most queries —
|
||||
it checks both offline and live. Use this directly only when you specifically need
|
||||
live-only results (e.g., verifying if offline info is current)."""
|
||||
results = _ddg_search(query, num_results)
|
||||
if not results:
|
||||
return f"No web results for '{query}'"
|
||||
return "\n\n".join(f"**{r['title']}**\n {r['snippet']}\n {r['url']}" for r in results)
|
||||
|
||||
# ── Gitea API ─────────────────────────────────────────────────────────────────
|
||||
def _gitea(method: str, path: str, body: dict = {}) -> dict:
|
||||
if not GITEA_TOKEN:
|
||||
return {"error": "GITEA_TOKEN not set in .env"}
|
||||
url = f"{GITEA_URL}/api/v1{path}"
|
||||
headers = {"Authorization": f"token {GITEA_TOKEN}",
|
||||
"Content-Type": "application/json"}
|
||||
r = httpx.request(method, url, json=body or None, headers=headers, timeout=30)
|
||||
try:
|
||||
return r.json()
|
||||
except Exception:
|
||||
return {"status": r.status_code, "text": r.text}
|
||||
|
||||
@mcp.tool()
|
||||
def gitea_list_repos() -> str:
|
||||
"""List your Gitea repos."""
|
||||
repos = _gitea("GET", "/repos/search?limit=50")
|
||||
if "error" in repos:
|
||||
return repos["error"]
|
||||
return "\n".join(f"{r['full_name']} — {r.get('description','')}"
|
||||
for r in repos.get("data", []))
|
||||
|
||||
@mcp.tool()
|
||||
def gitea_create_repo(name: str, private: bool = True, description: str = "") -> str:
|
||||
"""Create a new Gitea repository."""
|
||||
r = _gitea("POST", "/user/repos",
|
||||
{"name": name, "private": private, "description": description,
|
||||
"auto_init": True, "default_branch": "main"})
|
||||
return r.get("html_url") or str(r)
|
||||
|
||||
@mcp.tool()
|
||||
def gitea_create_issue(repo: str, title: str, body: str = "") -> str:
|
||||
"""Create an issue on a Gitea repo (format: owner/repo)."""
|
||||
r = _gitea("POST", f"/repos/{repo}/issues", {"title": title, "body": body})
|
||||
return r.get("html_url") or str(r)
|
||||
|
||||
# ── GitHub API ────────────────────────────────────────────────────────────────
|
||||
@mcp.tool()
|
||||
def github_api(method: str, endpoint: str, body: str = "") -> str:
|
||||
"""Call the GitHub REST API. endpoint e.g. /repos/owner/repo/issues"""
|
||||
if not GITHUB_TOKEN:
|
||||
return "GITHUB_TOKEN not set in .env"
|
||||
import json as _json
|
||||
headers = {"Authorization": f"Bearer {GITHUB_TOKEN}",
|
||||
"Accept": "application/vnd.github+json"}
|
||||
r = httpx.request(method.upper(), f"https://api.github.com{endpoint}",
|
||||
json=_json.loads(body) if body else None,
|
||||
headers=headers, timeout=30)
|
||||
try:
|
||||
return _json.dumps(r.json(), indent=2)
|
||||
except Exception:
|
||||
return r.text
|
||||
|
||||
# ── Gitea ↔ GitHub sync ──────────────────────────────────────────────────────
|
||||
@mcp.tool()
|
||||
def gitea_github_sync(mode: str = "all", repo: str = "") -> str:
|
||||
"""Run Gitea↔GitHub mirror sync. mode: all|pull|push|list. repo: optional owner/name."""
|
||||
cmd = ["/app/gitea-github-sync.sh"]
|
||||
if mode == "pull": cmd.append("--pull-only")
|
||||
elif mode == "push": cmd.append("--push-only")
|
||||
elif mode == "list": cmd.append("--list")
|
||||
if repo:
|
||||
cmd.extend(["--repo", repo])
|
||||
try:
|
||||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=600,
|
||||
env={**os.environ, "SYNC_ENV": "/app/.env"})
|
||||
return (r.stdout + r.stderr).strip() or "Sync completed (no output)"
|
||||
except subprocess.TimeoutExpired:
|
||||
return "Sync timed out after 10 minutes"
|
||||
except Exception as e:
|
||||
return f"Sync failed: {e}"
|
||||
|
||||
# ── RAG ingest ────────────────────────────────────────────────────────────────
|
||||
@mcp.tool()
|
||||
def ingest_repo(url: str, name: str = "", branch: str = "main") -> str:
|
||||
"""Clone a git repo and index it in the RAG code collection."""
|
||||
r = httpx.post(f"{RAG_URL}/ingest/repo",
|
||||
json={"url": url, "name": name, "branch": branch}, timeout=300)
|
||||
return r.text
|
||||
|
||||
@mcp.tool()
|
||||
def rag_health() -> str:
|
||||
"""Check RAG server status and indexed document counts."""
|
||||
try:
|
||||
r = httpx.get(f"{RAG_URL}/health", timeout=10)
|
||||
return r.text
|
||||
except Exception as e:
|
||||
return f"RAG server unreachable: {e}"
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
app = mcp.sse_app()
|
||||
uvicorn.run(app, host="0.0.0.0", port=8002)
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
#!/bin/bash
|
||||
# =============================================================================
|
||||
# Ollama Model Auto-Updater
|
||||
# Pulls latest version of every installed model. Ollama compares digests
|
||||
# server-side — no download happens if the model is already current.
|
||||
# Safe to run as a systemd timer.
|
||||
# =============================================================================
|
||||
|
||||
OLLAMA_URL="${OLLAMA_HOST:-http://localhost:11434}"
|
||||
LOG="$HOME/docker/ai-stack/logs/ollama-update.log"
|
||||
mkdir -p "$(dirname "$LOG")"
|
||||
|
||||
G='\033[0;32m'; Y='\033[1;33m'; C='\033[0;36m'; R='\033[0;31m'; N='\033[0m'
|
||||
ok() { echo -e "${G}[OK]${N} $1" | tee -a "$LOG"; }
|
||||
inf() { echo -e "${C}[..]${N} $1" | tee -a "$LOG"; }
|
||||
wrn() { echo -e "${Y}[!!]${N} $1" | tee -a "$LOG"; }
|
||||
err() { echo -e "${R}[XX]${N} $1" | tee -a "$LOG"; }
|
||||
|
||||
echo "" | tee -a "$LOG"
|
||||
echo "=== Ollama Model Update: $(date) ===" | tee -a "$LOG"
|
||||
echo "" | tee -a "$LOG"
|
||||
|
||||
# Check Ollama is reachable
|
||||
if ! curl -sf "$OLLAMA_URL/api/tags" > /dev/null; then
|
||||
err "Ollama not reachable at $OLLAMA_URL — is the container running?"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Get installed model names from the API
|
||||
mapfile -t MODELS < <(
|
||||
curl -sf "$OLLAMA_URL/api/tags" | \
|
||||
grep -oP '"name"\s*:\s*"\K[^"]+' | \
|
||||
sort -u
|
||||
)
|
||||
|
||||
if [[ ${#MODELS[@]} -eq 0 ]]; then
|
||||
wrn "No models found — nothing to update"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
inf "Found ${#MODELS[@]} model(s): ${MODELS[*]}"
|
||||
echo "" | tee -a "$LOG"
|
||||
|
||||
UPDATED=0
|
||||
FAILED=0
|
||||
|
||||
for model in "${MODELS[@]}"; do
|
||||
inf "Pulling $model..."
|
||||
output=$(ollama pull "$model" 2>&1)
|
||||
exit_code=$?
|
||||
|
||||
if [[ $exit_code -ne 0 ]]; then
|
||||
err "$model: pull failed"
|
||||
echo "$output" >> "$LOG"
|
||||
FAILED=$((FAILED + 1))
|
||||
elif echo "$output" | grep -q "up to date"; then
|
||||
ok "$model: already up to date"
|
||||
else
|
||||
ok "$model: updated"
|
||||
UPDATED=$((UPDATED + 1))
|
||||
fi
|
||||
done
|
||||
|
||||
echo "" | tee -a "$LOG"
|
||||
inf "Done — $UPDATED updated, $FAILED failed"
|
||||
echo "=== Complete: $(date) ===" | tee -a "$LOG"
|
||||
echo "" | tee -a "$LOG"
|
||||
Vendored
+281
@@ -0,0 +1,281 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
RAG Server — code-aware chunking, multi-collection, repo ingest, webhooks.
|
||||
Collections: papers (PDFs/text), code (source files, AST-split for Python)
|
||||
"""
|
||||
import ast, fnmatch, hashlib, json, logging, os, re, subprocess, threading, time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
import chromadb
|
||||
import httpx
|
||||
from chromadb.utils.embedding_functions import OllamaEmbeddingFunction
|
||||
from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger("rag")
|
||||
|
||||
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://ollama:11434")
|
||||
CHROMA_URL = os.getenv("CHROMA_URL", "http://chromadb:8000")
|
||||
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
|
||||
CHAT_MODEL = os.getenv("CHAT_MODEL", "qwen2.5:14b")
|
||||
PAPERS_DIR = Path(os.getenv("PAPERS_DIR", "/papers"))
|
||||
REPOS_DIR = Path(os.getenv("REPOS_DIR", "/repos"))
|
||||
TOP_K = int(os.getenv("TOP_K", "6"))
|
||||
|
||||
CODE_EXTS = {".py",".js",".ts",".tsx",".jsx",".go",".rs",".java",".c",".cpp",
|
||||
".h",".hpp",".cs",".rb",".sh",".yaml",".yml",".toml",".sql",".md"}
|
||||
SKIP_DIRS = {"node_modules",".git","__pycache__","dist","build",".venv",
|
||||
"venv","env",".next","vendor","target","bin","obj"}
|
||||
SKIP_FILES = {"package-lock.json","yarn.lock","pnpm-lock.yaml","Cargo.lock"}
|
||||
MAX_BYTES = 400_000
|
||||
|
||||
app = FastAPI(title="RAG Server")
|
||||
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
|
||||
|
||||
# ── ChromaDB ──────────────────────────────────────────────────────────────────
|
||||
def _embed_fn():
|
||||
return OllamaEmbeddingFunction(
|
||||
url=f"{OLLAMA_URL}/api/embeddings", model_name=EMBED_MODEL)
|
||||
|
||||
def _chroma():
|
||||
host, port = CHROMA_URL.replace("http://","").split(":")
|
||||
return chromadb.HttpClient(host=host, port=int(port))
|
||||
|
||||
def get_col(name: str):
|
||||
return _chroma().get_or_create_collection(name, embedding_function=_embed_fn())
|
||||
|
||||
# ── chunkers ─────────────────────────────────────────────────────────────────
|
||||
def _doc_id(text: str, key: str) -> str:
|
||||
return hashlib.md5(f"{key}|{text[:200]}".encode()).hexdigest()
|
||||
|
||||
def _sliding(text: str, size=1000, overlap=150) -> list[str]:
|
||||
chunks, i = [], 0
|
||||
while i < len(text):
|
||||
chunks.append(text[i:i+size])
|
||||
i += size - overlap
|
||||
return [c for c in chunks if c.strip()]
|
||||
|
||||
def _chunk_python(src: str) -> list[tuple[str,str]]:
|
||||
try:
|
||||
tree = ast.parse(src)
|
||||
except SyntaxError:
|
||||
return []
|
||||
lines = src.splitlines()
|
||||
out = []
|
||||
for node in ast.iter_child_nodes(tree):
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
chunk = "\n".join(lines[node.lineno-1:node.end_lineno])
|
||||
out.append((node.name, chunk[:4000]))
|
||||
return out
|
||||
|
||||
def _chunk_file(path: Path, src: str) -> list[tuple[str,str]]:
|
||||
if path.suffix == ".py":
|
||||
pairs = _chunk_python(src)
|
||||
if pairs:
|
||||
return pairs
|
||||
# function/class boundary split for JS/TS/Go/Rust etc.
|
||||
pat = re.compile(
|
||||
r'(?:^|\n)(?=(?:export\s+)?(?:async\s+)?(?:function|class|const\s+\w+\s*=\s*(?:async\s+)?\()'
|
||||
r'|^func |^type |^impl |^pub fn |^fn )',
|
||||
re.MULTILINE)
|
||||
parts = [p.strip() for p in pat.split(src) if p.strip()]
|
||||
if len(parts) > 1:
|
||||
return [(f"s{i}", p[:4000]) for i, p in enumerate(parts)]
|
||||
return [(f"c{i}", c) for i, c in enumerate(_sliding(src, 1200, 200))]
|
||||
|
||||
# ── ingest helpers ────────────────────────────────────────────────────────────
|
||||
def ingest_file(col, fpath: Path, repo: str = ""):
|
||||
if fpath.stat().st_size > MAX_BYTES or fpath.name in SKIP_FILES:
|
||||
return
|
||||
if any(fnmatch.fnmatch(fpath.name, p) for p in ("*.min.js","*.min.css","*.map")):
|
||||
return
|
||||
try:
|
||||
src = fpath.read_text(encoding="utf-8", errors="ignore")
|
||||
except Exception:
|
||||
return
|
||||
if not src.strip():
|
||||
return
|
||||
rel = str(fpath)
|
||||
pairs = _chunk_file(fpath, src) if fpath.suffix in CODE_EXTS else \
|
||||
[(f"c{i}", c) for i, c in enumerate(_sliding(src))]
|
||||
ids, docs, metas = [], [], []
|
||||
for label, chunk in pairs:
|
||||
if not chunk.strip():
|
||||
continue
|
||||
ids.append(_doc_id(chunk, rel+label))
|
||||
docs.append(chunk)
|
||||
metas.append({"source": rel, "label": label, "repo": repo,
|
||||
"lang": fpath.suffix.lstrip(".")})
|
||||
if ids:
|
||||
col.upsert(ids=ids, documents=docs, metadatas=metas)
|
||||
|
||||
def ingest_dir(col, directory: Path, repo: str = "") -> int:
|
||||
count = 0
|
||||
for f in directory.rglob("*"):
|
||||
if not f.is_file():
|
||||
continue
|
||||
if any(p in f.parts for p in SKIP_DIRS):
|
||||
continue
|
||||
ingest_file(col, f, repo)
|
||||
count += 1
|
||||
log.info("Indexed %d files from %s", count, directory)
|
||||
return count
|
||||
|
||||
def ingest_pdfs(col) -> int:
|
||||
try:
|
||||
import pypdf
|
||||
except ImportError:
|
||||
log.warning("pypdf not installed — skipping PDFs")
|
||||
return 0
|
||||
n = 0
|
||||
for pdf in PAPERS_DIR.glob("*.pdf"):
|
||||
try:
|
||||
text = "\n".join(p.extract_text() or ""
|
||||
for p in pypdf.PdfReader(str(pdf)).pages)
|
||||
for i, chunk in enumerate(_sliding(text)):
|
||||
col.upsert(ids=[_doc_id(chunk, str(pdf)+str(i))],
|
||||
documents=[chunk],
|
||||
metadatas=[{"source": str(pdf), "label": f"p{i}",
|
||||
"repo": "", "lang": "pdf"}])
|
||||
n += 1
|
||||
except Exception as e:
|
||||
log.warning("PDF %s: %s", pdf.name, e)
|
||||
return n
|
||||
|
||||
# ── startup ───────────────────────────────────────────────────────────────────
|
||||
def _startup_index():
|
||||
# wait for embed model
|
||||
for _ in range(40):
|
||||
try:
|
||||
r = httpx.get(f"{OLLAMA_URL}/api/tags", timeout=5)
|
||||
if any(EMBED_MODEL in m["name"] for m in r.json().get("models", [])):
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
log.info("Waiting for embed model %s…", EMBED_MODEL)
|
||||
time.sleep(5)
|
||||
|
||||
code_col = get_col("code")
|
||||
papers_col = get_col("papers")
|
||||
for d in REPOS_DIR.iterdir():
|
||||
if d.is_dir():
|
||||
ingest_dir(code_col, d, d.name)
|
||||
ingest_pdfs(papers_col)
|
||||
for f in PAPERS_DIR.glob("*.txt"):
|
||||
ingest_file(papers_col, f)
|
||||
log.info("Startup index complete")
|
||||
|
||||
@app.on_event("startup")
|
||||
async def on_startup():
|
||||
threading.Thread(target=_startup_index, daemon=True).start()
|
||||
|
||||
# ── endpoints ─────────────────────────────────────────────────────────────────
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
try:
|
||||
cc = _chroma()
|
||||
return {"status": "ok",
|
||||
"code": cc.get_collection("code", embedding_function=_embed_fn()).count(),
|
||||
"papers": cc.get_collection("papers", embedding_function=_embed_fn()).count(),
|
||||
"embed": EMBED_MODEL, "chat": CHAT_MODEL}
|
||||
except Exception as e:
|
||||
return {"status": "error", "detail": str(e)}
|
||||
|
||||
class RepoRequest(BaseModel):
|
||||
url: str
|
||||
name: str = ""
|
||||
branch: str = "main"
|
||||
|
||||
@app.post("/ingest/repo")
|
||||
async def ingest_repo(req: RepoRequest):
|
||||
name = req.name or req.url.rstrip("/").split("/")[-1].removesuffix(".git")
|
||||
dest = REPOS_DIR / name
|
||||
try:
|
||||
if dest.exists():
|
||||
subprocess.run(["git","pull"], cwd=dest, check=True, timeout=120)
|
||||
else:
|
||||
subprocess.run(["git","clone","--depth=1","-b",req.branch,
|
||||
req.url, str(dest)], check=True, timeout=300)
|
||||
except subprocess.CalledProcessError as e:
|
||||
raise HTTPException(400, str(e))
|
||||
n = ingest_dir(get_col("code"), dest, name)
|
||||
return {"status": "ok", "repo": name, "files": n}
|
||||
|
||||
@app.post("/ingest/papers")
|
||||
async def trigger_papers(bg: BackgroundTasks):
|
||||
bg.add_task(ingest_pdfs, get_col("papers"))
|
||||
return {"status": "queued"}
|
||||
|
||||
async def _webhook(payload: dict):
|
||||
repo = payload.get("repository") or {}
|
||||
url = repo.get("clone_url") or repo.get("html_url","")
|
||||
name = repo.get("name","unknown")
|
||||
if not url:
|
||||
return {"status": "ignored"}
|
||||
dest = REPOS_DIR / name
|
||||
if dest.exists():
|
||||
subprocess.run(["git","pull"], cwd=dest, timeout=120)
|
||||
else:
|
||||
subprocess.run(["git","clone","--depth=1",url,str(dest)], timeout=300)
|
||||
n = ingest_dir(get_col("code"), dest, name)
|
||||
return {"status": "ok", "repo": name, "files": n}
|
||||
|
||||
@app.post("/webhook/gitea")
|
||||
async def webhook_gitea(r: Request): return await _webhook(await r.json())
|
||||
|
||||
@app.post("/webhook/github")
|
||||
async def webhook_github(r: Request): return await _webhook(await r.json())
|
||||
|
||||
# ── RAG chat ──────────────────────────────────────────────────────────────────
|
||||
class ChatRequest(BaseModel):
|
||||
model: str = CHAT_MODEL
|
||||
messages: list[dict[str,Any]]
|
||||
stream: bool = False
|
||||
collections: list[str] = ["code","papers"]
|
||||
|
||||
def _context(query: str, cols: list[str]) -> str:
|
||||
parts = []
|
||||
for cname in cols:
|
||||
try:
|
||||
col = get_col(cname)
|
||||
if col.count() == 0:
|
||||
continue
|
||||
res = col.query(query_texts=[query], n_results=min(TOP_K, col.count()))
|
||||
for doc, meta in zip(res["documents"][0], res["metadatas"][0]):
|
||||
parts.append(f"### {meta.get('source','')}:{meta.get('label','')}\n"
|
||||
f"```{meta.get('lang','')}\n{doc}\n```")
|
||||
except Exception as e:
|
||||
log.warning("col %s: %s", cname, e)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def chat(req: ChatRequest):
|
||||
query = next((m["content"] for m in reversed(req.messages)
|
||||
if m.get("role")=="user"), "")
|
||||
context = _context(query, req.collections)
|
||||
msgs = [{"role":"system","content":
|
||||
"You are a helpful coding assistant. Use the retrieved context below.\n\n"
|
||||
f"## Context\n{context}"}] + req.messages
|
||||
|
||||
payload = {"model": req.model, "messages": msgs, "stream": req.stream}
|
||||
|
||||
if req.stream:
|
||||
async def gen():
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
async with client.stream("POST",
|
||||
f"{OLLAMA_URL}/v1/chat/completions", json=payload) as r:
|
||||
async for chunk in r.aiter_bytes():
|
||||
yield chunk
|
||||
return StreamingResponse(gen(), media_type="text/event-stream")
|
||||
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
r = await client.post(f"{OLLAMA_URL}/v1/chat/completions", json=payload)
|
||||
return r.json()
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run("server:app", host="0.0.0.0", port=8001, reload=False)
|
||||
+275
@@ -0,0 +1,275 @@
|
||||
#!/usr/bin/env bash
|
||||
# Detect GPU VRAM and install appropriate Stable Diffusion models for
|
||||
# InvokeAI and/or ComfyUI. Run anytime — safe to re-run.
|
||||
#
|
||||
# Usage: ./setup-image-models.sh [--auto]
|
||||
# --auto Skip prompts, install the recommended default for your GPU
|
||||
set -euo pipefail
|
||||
|
||||
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
|
||||
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
|
||||
info() { echo -e "${CYAN}[..]${NC} $*"; }
|
||||
ok() { echo -e "${GREEN}[OK]${NC} $*"; }
|
||||
warn() { echo -e "${YELLOW}[!!]${NC} $*"; }
|
||||
|
||||
AUTO=false
|
||||
[[ "${1:-}" == "--auto" ]] && AUTO=true
|
||||
|
||||
# ── Detect GPU ────────────────────────────────────────────────────────────────
|
||||
VRAM_GB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null \
|
||||
| head -1 | awk '{printf "%d", $1/1024}' 2>/dev/null || echo "0")
|
||||
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | wc -l || echo "0")
|
||||
GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 || echo "None")
|
||||
TOTAL_VRAM=$((VRAM_GB * GPU_COUNT))
|
||||
|
||||
echo ""
|
||||
echo -e "${BOLD}━━━ Image Generation Model Setup ━━━${NC}"
|
||||
echo ""
|
||||
info "GPU : $GPU_NAME"
|
||||
[[ "$GPU_COUNT" -gt 1 ]] && info "GPU count : $GPU_COUNT"
|
||||
info "VRAM/card : ${VRAM_GB}GB"
|
||||
info "Total VRAM: ${TOTAL_VRAM}GB"
|
||||
echo ""
|
||||
|
||||
# ── Determine which containers are available ──────────────────────────────────
|
||||
HAS_INVOKEAI=false
|
||||
HAS_COMFYUI=false
|
||||
docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^invokeai$' && HAS_INVOKEAI=true
|
||||
docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^comfyui$' && HAS_COMFYUI=true
|
||||
|
||||
if ! $HAS_INVOKEAI && ! $HAS_COMFYUI; then
|
||||
echo -e "${RED}Error:${NC} Neither InvokeAI nor ComfyUI containers found."
|
||||
echo " Run the setup script first to deploy the AI stack."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
$HAS_INVOKEAI && info "InvokeAI : found"
|
||||
$HAS_COMFYUI && info "ComfyUI : found"
|
||||
echo ""
|
||||
|
||||
# ── Build model menu based on VRAM ────────────────────────────────────────────
|
||||
# Model VRAM requirements (generation, not just loading):
|
||||
# SD 1.5 ~4GB 512x512 native
|
||||
# SDXL ~7GB 1024x1024 native
|
||||
# SDXL Turbo ~7GB 512x512 (4-step)
|
||||
# Flux.1-schnell ~12GB fast, high quality
|
||||
# Flux.1-dev ~20GB best quality, slow
|
||||
|
||||
declare -a MODEL_IDS=()
|
||||
declare -a MODEL_NAMES=()
|
||||
declare -a MODEL_VRAM=()
|
||||
declare -a MODEL_NOTES=()
|
||||
|
||||
add_model() {
|
||||
MODEL_IDS+=("$1"); MODEL_NAMES+=("$2"); MODEL_VRAM+=("$3"); MODEL_NOTES+=("$4")
|
||||
}
|
||||
|
||||
# Always offer SD 1.5 if any GPU exists
|
||||
if [[ "$TOTAL_VRAM" -ge 4 ]]; then
|
||||
add_model "sd15" "Stable Diffusion 1.5" "4" "512px native, most LoRA compatible, fast"
|
||||
fi
|
||||
|
||||
if [[ "$TOTAL_VRAM" -ge 8 ]]; then
|
||||
add_model "sdxl" "Stable Diffusion XL" "7" "1024px native, better quality, more detail"
|
||||
add_model "sdxl-turbo" "SDXL Turbo" "7" "4-step generation, very fast, good quality"
|
||||
fi
|
||||
|
||||
if [[ "$TOTAL_VRAM" -ge 12 ]]; then
|
||||
add_model "flux-schnell" "Flux.1-schnell" "12" "fast Flux variant, excellent quality"
|
||||
fi
|
||||
|
||||
if [[ "$TOTAL_VRAM" -ge 20 ]]; then
|
||||
add_model "flux-dev" "Flux.1-dev" "20" "best quality, slower, needs lots of VRAM"
|
||||
fi
|
||||
|
||||
if [[ ${#MODEL_IDS[@]} -eq 0 ]]; then
|
||||
warn "No GPU with sufficient VRAM detected (need at least 4GB)."
|
||||
warn "CPU-only image generation is extremely slow and not recommended."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# ── Determine default recommendation ─────────────────────────────────────────
|
||||
if [[ "$TOTAL_VRAM" -ge 20 ]]; then DEFAULT_ID="flux-dev"
|
||||
elif [[ "$TOTAL_VRAM" -ge 12 ]]; then DEFAULT_ID="sdxl"
|
||||
elif [[ "$TOTAL_VRAM" -ge 8 ]]; then DEFAULT_ID="sdxl"
|
||||
elif [[ "$TOTAL_VRAM" -ge 4 ]]; then DEFAULT_ID="sd15"
|
||||
else DEFAULT_ID="sd15"
|
||||
fi
|
||||
|
||||
echo -e "${BOLD}Available models for your ${TOTAL_VRAM}GB GPU:${NC}"
|
||||
echo ""
|
||||
for i in "${!MODEL_IDS[@]}"; do
|
||||
DEFAULT_TAG=""
|
||||
[[ "${MODEL_IDS[$i]}" == "$DEFAULT_ID" ]] && DEFAULT_TAG=" ${GREEN}← recommended${NC}"
|
||||
printf " ${BOLD}%d)${NC} %-25s ~%sGB VRAM %s%b\n" \
|
||||
$((i+1)) "${MODEL_NAMES[$i]}" "${MODEL_VRAM[$i]}" "${MODEL_NOTES[$i]}" "$DEFAULT_TAG"
|
||||
done
|
||||
echo ""
|
||||
|
||||
if $AUTO; then
|
||||
SELECTED="$DEFAULT_ID"
|
||||
info "Auto mode: installing $SELECTED"
|
||||
else
|
||||
echo -e " Enter number(s) separated by spaces, or press Enter for recommended."
|
||||
echo -e " Example: ${BOLD}1 2${NC} to install both SD 1.5 and SDXL"
|
||||
echo ""
|
||||
read -rp " Selection [recommended]: " CHOICE
|
||||
|
||||
if [[ -z "$CHOICE" ]]; then
|
||||
SELECTED="$DEFAULT_ID"
|
||||
else
|
||||
SELECTED=""
|
||||
for num in $CHOICE; do
|
||||
idx=$((num - 1))
|
||||
if [[ $idx -ge 0 && $idx -lt ${#MODEL_IDS[@]} ]]; then
|
||||
SELECTED+=" ${MODEL_IDS[$idx]}"
|
||||
else
|
||||
warn "Invalid selection: $num (skipping)"
|
||||
fi
|
||||
done
|
||||
SELECTED="${SELECTED# }"
|
||||
fi
|
||||
fi
|
||||
|
||||
[[ -z "$SELECTED" ]] && { warn "No models selected."; exit 1; }
|
||||
|
||||
echo ""
|
||||
info "Will install: $SELECTED"
|
||||
echo ""
|
||||
|
||||
# ── HuggingFace model identifiers ────────────────────────────────────────────
|
||||
declare -A HF_MODELS=(
|
||||
[sd15]="stabilityai/stable-diffusion-v1-5"
|
||||
[sdxl]="stabilityai/stable-diffusion-xl-base-1.0"
|
||||
[sdxl-turbo]="stabilityai/sdxl-turbo"
|
||||
[flux-schnell]="black-forest-labs/FLUX.1-schnell"
|
||||
[flux-dev]="black-forest-labs/FLUX.1-dev"
|
||||
)
|
||||
|
||||
declare -A MODEL_SIZES=(
|
||||
[sd15]="~4GB"
|
||||
[sdxl]="~7GB"
|
||||
[sdxl-turbo]="~7GB"
|
||||
[flux-schnell]="~12GB"
|
||||
[flux-dev]="~24GB"
|
||||
)
|
||||
|
||||
# ── Install into InvokeAI ────────────────────────────────────────────────────
|
||||
if $HAS_INVOKEAI; then
|
||||
echo -e "${BOLD}━━━ Installing into InvokeAI ━━━${NC}"
|
||||
|
||||
# Make sure container is running
|
||||
if ! docker ps --format '{{.Names}}' | grep -q '^invokeai$'; then
|
||||
info "Starting InvokeAI container..."
|
||||
docker start invokeai
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
for model_id in $SELECTED; do
|
||||
hf_id="${HF_MODELS[$model_id]:-}"
|
||||
[[ -z "$hf_id" ]] && { warn "Unknown model: $model_id"; continue; }
|
||||
info "Installing $model_id (${MODEL_SIZES[$model_id]}) → ${hf_id}..."
|
||||
info " This may take a while depending on your connection."
|
||||
|
||||
if docker exec invokeai invokeai-model-install --add "$hf_id" 2>&1; then
|
||||
ok "$model_id installed in InvokeAI"
|
||||
else
|
||||
warn "$model_id install failed in InvokeAI — try manually via Model Manager at :9090"
|
||||
fi
|
||||
echo ""
|
||||
done
|
||||
fi
|
||||
|
||||
# ── Install into ComfyUI ─────────────────────────────────────────────────────
|
||||
if $HAS_COMFYUI; then
|
||||
echo -e "${BOLD}━━━ Installing into ComfyUI ━━━${NC}"
|
||||
info "ComfyUI downloads models on first use via its UI."
|
||||
info "To pre-download, use the ComfyUI Manager at http://localhost:8188"
|
||||
echo ""
|
||||
|
||||
# Make sure container is running
|
||||
if ! docker ps --format '{{.Names}}' | grep -q '^comfyui$'; then
|
||||
info "Starting ComfyUI container..."
|
||||
docker start comfyui
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# For ComfyUI, download checkpoints into the models volume
|
||||
for model_id in $SELECTED; do
|
||||
hf_id="${HF_MODELS[$model_id]:-}"
|
||||
[[ -z "$hf_id" ]] && continue
|
||||
|
||||
# Check if model already exists
|
||||
CKPT_DIR="/opt/ComfyUI/models/checkpoints"
|
||||
if docker exec comfyui ls "$CKPT_DIR" 2>/dev/null | grep -qi "${model_id//-/_}"; then
|
||||
ok "$model_id already present in ComfyUI"
|
||||
continue
|
||||
fi
|
||||
|
||||
info "Downloading $model_id for ComfyUI (${MODEL_SIZES[$model_id]})..."
|
||||
info " Downloading from HuggingFace: $hf_id"
|
||||
|
||||
# Use ComfyUI's built-in download mechanism via python
|
||||
case "$model_id" in
|
||||
sd15)
|
||||
docker exec comfyui bash -c \
|
||||
"cd /opt/ComfyUI && python -c \"
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download('$hf_id', 'v1-5-pruned-emaonly.safetensors', local_dir='models/checkpoints')
|
||||
\" 2>&1" && ok "$model_id downloaded for ComfyUI" \
|
||||
|| warn "$model_id download failed — install via ComfyUI Manager UI"
|
||||
;;
|
||||
sdxl)
|
||||
docker exec comfyui bash -c \
|
||||
"cd /opt/ComfyUI && python -c \"
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download('$hf_id', 'sd_xl_base_1.0.safetensors', local_dir='models/checkpoints')
|
||||
\" 2>&1" && ok "$model_id downloaded for ComfyUI" \
|
||||
|| warn "$model_id download failed — install via ComfyUI Manager UI"
|
||||
;;
|
||||
*)
|
||||
info "$model_id: use ComfyUI Manager to install (complex model structure)"
|
||||
;;
|
||||
esac
|
||||
echo ""
|
||||
done
|
||||
fi
|
||||
|
||||
# ── Summary ───────────────────────────────────────────────────────────────────
|
||||
echo ""
|
||||
echo -e "${GREEN}${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo -e "${GREEN}${BOLD} Image model setup complete!${NC}"
|
||||
echo -e "${GREEN}${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
|
||||
echo ""
|
||||
echo -e " ${BOLD}Your GPU:${NC} $GPU_NAME — ${TOTAL_VRAM}GB VRAM"
|
||||
echo -e " ${BOLD}Installed:${NC} $SELECTED"
|
||||
echo ""
|
||||
$HAS_INVOKEAI && echo -e " ${CYAN}InvokeAI${NC} → http://localhost:9090"
|
||||
$HAS_COMFYUI && echo -e " ${CYAN}ComfyUI${NC} → http://localhost:8188"
|
||||
echo ""
|
||||
|
||||
if $HAS_INVOKEAI; then
|
||||
echo -e " ${YELLOW}InvokeAI quick start:${NC}"
|
||||
echo " 1. Open Model Manager → verify your model appears"
|
||||
echo " 2. Go to Text to Image → select the model"
|
||||
echo " 3. For inpainting: use the Unified Canvas tab"
|
||||
echo " - Upload image → brush over the area to change"
|
||||
echo " - Write what you want in that area → Invoke"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
if $HAS_COMFYUI; then
|
||||
echo -e " ${YELLOW}ComfyUI quick start:${NC}"
|
||||
echo " 1. Open ComfyUI → load a basic txt2img workflow"
|
||||
echo " 2. Select your checkpoint in the Load Checkpoint node"
|
||||
echo " 3. For Open WebUI integration: enable Dev Mode → export API workflow"
|
||||
echo ""
|
||||
fi
|
||||
|
||||
echo -e " ${YELLOW}GPU sharing:${NC} Ollama and image gen share the GPU."
|
||||
echo " Ollama auto-unloads models after 24h idle (KEEP_ALIVE=24h)."
|
||||
echo " For immediate unload before heavy image gen:"
|
||||
echo " docker exec ollama ollama stop <model-name>"
|
||||
echo ""
|
||||
echo -e " ${YELLOW}Import LoRAs:${NC} ./invokeai-import-lora.sh <file.safetensors>"
|
||||
echo ""
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
#!/bin/bash
|
||||
# Install all local-ai systemd user units and enable timers
|
||||
set -e
|
||||
|
||||
UNIT_DIR="$HOME/.config/systemd/user"
|
||||
SRC="$(cd "$(dirname "$0")" && pwd)"
|
||||
|
||||
mkdir -p "$UNIT_DIR"
|
||||
|
||||
for unit in "$SRC"/*.service "$SRC"/*.timer; do
|
||||
name=$(basename "$unit")
|
||||
ln -sf "$unit" "$UNIT_DIR/$name"
|
||||
echo "Linked $name"
|
||||
done
|
||||
|
||||
systemctl --user daemon-reload
|
||||
|
||||
for timer in "$SRC"/*.timer; do
|
||||
name=$(basename "$timer")
|
||||
systemctl --user enable --now "$name"
|
||||
echo "Enabled $name"
|
||||
done
|
||||
|
||||
echo ""
|
||||
systemctl --user list-timers kiwix-update.timer ollama-update.timer
|
||||
echo ""
|
||||
echo "Run manually:"
|
||||
echo " systemctl --user start kiwix-update.service"
|
||||
echo " systemctl --user start ollama-update.service"
|
||||
echo ""
|
||||
echo "Logs:"
|
||||
echo " journalctl --user -u kiwix-update.service -f"
|
||||
echo " journalctl --user -u ollama-update.service -f"
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
[Unit]
|
||||
Description=Kiwix ZIM auto-updater
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
ExecStart=/bin/bash %h/local-ai/kiwix_update.sh
|
||||
StandardOutput=journal
|
||||
StandardError=journal
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
[Unit]
|
||||
Description=Monthly Kiwix ZIM update check
|
||||
Requires=kiwix-update.service
|
||||
|
||||
[Timer]
|
||||
OnCalendar=monthly
|
||||
RandomizedDelaySec=1h
|
||||
Persistent=true
|
||||
|
||||
[Install]
|
||||
WantedBy=timers.target
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
[Unit]
|
||||
Description=Ollama model auto-updater
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
ExecStart=/bin/bash %h/local-ai/ollama_update.sh
|
||||
StandardOutput=journal
|
||||
StandardError=journal
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
[Unit]
|
||||
Description=Monthly Ollama model update check
|
||||
Requires=ollama-update.service
|
||||
|
||||
[Timer]
|
||||
OnCalendar=monthly
|
||||
RandomizedDelaySec=2h
|
||||
Persistent=true
|
||||
|
||||
[Install]
|
||||
WantedBy=timers.target
|
||||
Vendored
+222
@@ -0,0 +1,222 @@
|
||||
# =============================================================================
|
||||
# AI Photo Edit - Environment Configuration
|
||||
# =============================================================================
|
||||
#
|
||||
# SETUP INSTRUCTIONS:
|
||||
# 1. Copy this file to .env: cp .env.example .env
|
||||
# 2. Get API key from Replicate (see below)
|
||||
# 3. Paste your key in the REPLICATE_API_KEY line
|
||||
# 4. Rebuild: docker-compose up -d --build
|
||||
#
|
||||
# =============================================================================
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# STEP 1: Choose AI Provider
|
||||
# =============================================================================
|
||||
# Options: local_gpu, mock, openai, stability, replicate, invokeai, comfyui
|
||||
#
|
||||
# local_gpu = FREE, runs on YOUR GPU — best option if you have an NVIDIA card
|
||||
# (use docker-compose.gpu.yml — models auto-download on first use)
|
||||
# mock = Free, returns original image unchanged (UI testing only)
|
||||
# openai = gpt-image-2 generation + edits (~$0.006-0.21/image, see below)
|
||||
# stability = NOT IMPLEMENTED YET — config field exists but no driver in
|
||||
# remote_provider.py; setting this breaks every AI call
|
||||
# replicate = NOT IMPLEMENTED YET — same as above (Replicate IS used for
|
||||
# Smart Select's SAM fallback, but that's a separate code path)
|
||||
# invokeai = Self-hosted InvokeAI running on another machine
|
||||
# comfyui = Self-hosted ComfyUI running on another machine
|
||||
#
|
||||
# GPU QUICK-START:
|
||||
# docker compose -f docker-compose.gpu.yml up --build
|
||||
# (AI_PROVIDER defaults to local_gpu in that compose file — but only if this
|
||||
# var is unset/blank; since AI_PROVIDER=local_gpu is set explicitly below,
|
||||
# that's what you get either way)
|
||||
# =============================================================================
|
||||
|
||||
AI_PROVIDER=local_gpu
|
||||
|
||||
# ── Local GPU settings (only relevant when AI_PROVIDER=local_gpu) ────────────
|
||||
# Auto-download HuggingFace models on first request (true/false)
|
||||
AUTO_DOWNLOAD_MODELS=true
|
||||
# Max diffusion pipelines to keep loaded in GPU memory (each is 2–7 GB)
|
||||
LOCAL_GPU_MAX_PIPELINES=2
|
||||
# HuggingFace token — only needed for gated/private models
|
||||
#HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
# Override auto-selected model for any operation (leave blank = auto by VRAM tier)
|
||||
# On 4-6GB cards (Quadro P2200, GTX 1060/1660, etc.) the auto-tier picks either
|
||||
# slow SDXL+CPU-offload or a generic (not hand/face-tuned) SD checkpoint.
|
||||
# Lykon/dreamshaper-8-inpainting is SD1.5-based (1.7GB, fast, no offload needed)
|
||||
# and noticeably better on hands/faces — worth forcing on small cards:
|
||||
#HF_MODEL_INPAINT=Lykon/dreamshaper-8-inpainting
|
||||
#HF_MODEL_TXT2IMG=your-org/your-txt2img-model
|
||||
#HF_MODEL_IMG2IMG=your-org/your-img2img-model
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Per-operation provider overrides (optional — blank means use AI_PROVIDER above)
|
||||
# Example: use OpenAI for text-to-image (best quality) but InvokeAI for everything else
|
||||
#AI_PROVIDER_TXT2IMG=openai
|
||||
#AI_PROVIDER_INPAINT=invokeai
|
||||
#AI_PROVIDER_IMG2IMG=invokeai
|
||||
#AI_PROVIDER_OUTPAINT=invokeai
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# STEP 2: Get Your API Key
|
||||
# =============================================================================
|
||||
#
|
||||
# ╔═══════════════════════════════════════════════════════════════════════════╗
|
||||
# ║ REPLICATE (RECOMMENDED) ║
|
||||
# ╠═══════════════════════════════════════════════════════════════════════════╣
|
||||
# ║ ║
|
||||
# ║ 1. Go to: https://replicate.com ║
|
||||
# ║ 2. Click "Sign in" (use GitHub, Google, or email) ║
|
||||
# ║ 3. Go to: https://replicate.com/account/api-tokens ║
|
||||
# ║ 4. Click "Create token" ║
|
||||
# ║ 5. Copy the token (starts with "r8_") ║
|
||||
# ║ 6. Paste it below after REPLICATE_API_KEY= ║
|
||||
# ║ ║
|
||||
# ║ FREE TIER: New accounts get some free credits to try models! ║
|
||||
# ║ PRICING: ~$0.002-0.03 per image depending on model ║
|
||||
# ║ ║
|
||||
# ╚═══════════════════════════════════════════════════════════════════════════╝
|
||||
|
||||
REPLICATE_API_KEY=r8_PASTE_YOUR_KEY_HERE
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# OPENAI (cloud, gpt-image-2)
|
||||
# Get key at: https://platform.openai.com/api-keys
|
||||
# AI_PROVIDER=openai
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# dall-e-2 and dall-e-3 were retired May 12, 2026; gpt-image-1 deprecates
|
||||
# Oct 23, 2026. gpt-image-2 is the current model and handles both
|
||||
# generation and masked edits (inpaint/img2img/outpaint) — no org
|
||||
# verification step needed, unlike gpt-image-1.
|
||||
# Pricing (1024x1024): ~$0.006 low / $0.053 medium / $0.211 high quality.
|
||||
#OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
#OPENAI_MODEL=gpt-image-2
|
||||
# Model for inpaint/img2img/outpaint (the /v1/images/edits endpoint).
|
||||
#OPENAI_EDIT_MODEL=gpt-image-2
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# INVOKEAI (self-hosted, best for Flux/SDXL)
|
||||
# Run InvokeAI on your local machine or NAS, point URL here.
|
||||
# AI_PROVIDER=invokeai
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#INVOKEAI_URL=http://192.168.1.x:9090
|
||||
#INVOKEAI_DEFAULT_MODEL=flux-dev
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# COMFYUI (self-hosted, workflow JSON API)
|
||||
# AI_PROVIDER=comfyui
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#COMFYUI_URL=http://192.168.1.x:8188
|
||||
#COMFYUI_DEFAULT_MODEL=v1-5-pruned-emaonly.ckpt
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# STABILITY AI (Alternative)
|
||||
# Get key at: https://platform.stability.ai/account/keys
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#STABILITY_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# STEP 3: Model Selection (OPTIONAL - for advanced users)
|
||||
# =============================================================================
|
||||
#
|
||||
# By default, the system AUTO-SELECTS the best model based on your prompt:
|
||||
# - Prompt contains "remove/erase/delete" → Uses LaMa (fast removal)
|
||||
# - Prompt contains "face/hands/person" → Uses Realistic Vision
|
||||
# - Everything else → Uses SDXL Inpaint
|
||||
#
|
||||
# To FORCE a specific model, uncomment ONE line below:
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# REPLICATE_MODEL=sdxl-inpaint # General purpose, good quality (~$0.01)
|
||||
# REPLICATE_MODEL=lama # Object removal ONLY (~$0.002, fastest)
|
||||
# REPLICATE_MODEL=realistic-vision # Faces, hands, skin (~$0.02)
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# IMPORTANT: About Flux and other text-to-image models
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#
|
||||
# Models like "black-forest-labs/flux-kontext-pro" are TEXT-TO-IMAGE models.
|
||||
# They generate NEW images from text, they DON'T edit existing images.
|
||||
#
|
||||
# For EDITING (inpainting), you need models that accept:
|
||||
# - An existing image
|
||||
# - A mask showing what to change
|
||||
# - A prompt describing the change
|
||||
#
|
||||
# WORKS for editing: DOESN'T work for editing:
|
||||
# ✓ sdxl-inpaint ✗ flux-kontext-pro (text-to-image)
|
||||
# ✓ lama ✗ flux-dev (text-to-image)
|
||||
# ✓ realistic-vision ✗ ideogram (text-to-image)
|
||||
#
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Stability AI model selection (if using AI_PROVIDER=stability)
|
||||
#STABILITY_MODEL=sdxl # Options: sdxl, sd15, sd21
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SECURITY (Change this in production!)
|
||||
# =============================================================================
|
||||
|
||||
SECRET_KEY=change-this-to-a-long-random-string-in-production
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ADVANCED SETTINGS (Usually don't need to change)
|
||||
# =============================================================================
|
||||
|
||||
# CORS origins (comma-separated)
|
||||
CORS_ORIGINS=http://localhost:5173,http://localhost:3000,http://localhost:3080,http://localhost
|
||||
|
||||
# Database path
|
||||
DATABASE_URL=sqlite:///./data/ai_photo_edit.db
|
||||
|
||||
# Auto-download SAM model on startup (true/false)
|
||||
# When true (default): Downloads SAM model (~375MB) on first startup for offline Smart Select
|
||||
# When false: Skips download, Smart Select uses Replicate API (requires REPLICATE_API_KEY)
|
||||
AUTO_DOWNLOAD_SAM=true
|
||||
|
||||
# Auto-download U2Net model on startup (true/false)
|
||||
# When true (default): Downloads U2Net model (~176MB) on first startup for offline Remove Background
|
||||
# When false: Skips download, Remove Background falls back to rembg (if installed)
|
||||
AUTO_DOWNLOAD_U2NET=true
|
||||
|
||||
# Background removal model (Remove Background tool) — used when request.model="auto"
|
||||
# Options: ben2 (default — best for clean cutouts, hair/edges), birefnet-hr
|
||||
# (best for high-res/print work, slower), u2net (lightweight, always-on fallback)
|
||||
# ben2 and birefnet-hr download weights from HuggingFace on first use (GPU image only).
|
||||
BG_REMOVAL_MODEL=ben2
|
||||
|
||||
# Allow users to select model per-edit
|
||||
ALLOW_MODEL_OVERRIDE=true
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TROUBLESHOOTING
|
||||
# =============================================================================
|
||||
#
|
||||
# PROBLEM: "405 Method Not Allowed" errors
|
||||
# FIX: Rebuild container: docker-compose build --no-cache && docker-compose up -d
|
||||
#
|
||||
# PROBLEM: "REPLICATE_API_KEY not configured"
|
||||
# FIX: 1. Make sure .env file exists (not just .env.example)
|
||||
# 2. Make sure REPLICATE_API_KEY has your actual key
|
||||
# 3. Restart: docker-compose down && docker-compose up -d
|
||||
#
|
||||
# PROBLEM: Edits don't change the image
|
||||
# FIX: Check AI_PROVIDER isn't set to "mock"
|
||||
#
|
||||
# PROBLEM: "rembg not installed"
|
||||
# FIX: Rebuild: docker-compose build --no-cache backend
|
||||
#
|
||||
# PROBLEM: Smart Select uses flood-fill instead of AI
|
||||
# FIX: Smart Select needs REPLICATE_API_KEY for SAM model
|
||||
#
|
||||
# CHECK LOGS: docker-compose logs -f backend
|
||||
#
|
||||
# =============================================================================
|
||||
@@ -0,0 +1,70 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
venv/
|
||||
env/
|
||||
ENV/
|
||||
|
||||
# Node
|
||||
node_modules/
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
.pnp
|
||||
.pnp.js
|
||||
coverage/
|
||||
build/
|
||||
.DS_Store
|
||||
.env.local
|
||||
.env.development.local
|
||||
.env.test.local
|
||||
.env.production.local
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Environment
|
||||
.env
|
||||
|
||||
# Data
|
||||
data/projects/*/
|
||||
!data/projects/.gitkeep
|
||||
*.db
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
|
||||
# Docker
|
||||
*.log
|
||||
docker-compose.override.yml
|
||||
|
||||
# Testing
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
Vendored
+86
@@ -0,0 +1,86 @@
|
||||
# Contributing to AI Photo Edit
|
||||
|
||||
Thank you for your interest in contributing to AI Photo Edit!
|
||||
|
||||
## Development Setup
|
||||
|
||||
1. Fork the repository
|
||||
2. Clone your fork
|
||||
3. Create a feature branch
|
||||
4. Make your changes
|
||||
5. Test your changes
|
||||
6. Submit a pull request
|
||||
|
||||
## Development Environment
|
||||
|
||||
### Using Docker (Recommended)
|
||||
|
||||
```bash
|
||||
# Start dev environment with hot-reload
|
||||
docker-compose -f docker-compose.dev.yml up
|
||||
```
|
||||
|
||||
### Local Development
|
||||
|
||||
**Backend**
|
||||
```bash
|
||||
cd backend
|
||||
python -m venv venv
|
||||
source venv/bin/activate # On Windows: venv\Scripts\activate
|
||||
pip install -r requirements.txt
|
||||
uvicorn app.main:app --reload
|
||||
```
|
||||
|
||||
**Frontend**
|
||||
```bash
|
||||
cd frontend
|
||||
npm install
|
||||
npm run dev
|
||||
```
|
||||
|
||||
## Code Style
|
||||
|
||||
### Python (Backend)
|
||||
- Follow PEP 8
|
||||
- Use type hints where appropriate
|
||||
- Add docstrings to functions and classes
|
||||
|
||||
### JavaScript/React (Frontend)
|
||||
- Use functional components with hooks
|
||||
- Follow React best practices
|
||||
- Use meaningful variable names
|
||||
|
||||
## Pull Request Process
|
||||
|
||||
1. Update the README.md with details of changes if needed
|
||||
2. Ensure all tests pass
|
||||
3. Update documentation as needed
|
||||
4. Get approval from maintainers
|
||||
5. Squash commits if requested
|
||||
|
||||
## Reporting Bugs
|
||||
|
||||
When reporting bugs, please include:
|
||||
- Description of the issue
|
||||
- Steps to reproduce
|
||||
- Expected behavior
|
||||
- Actual behavior
|
||||
- Screenshots if applicable
|
||||
- Environment details (OS, Docker version, etc.)
|
||||
|
||||
## Feature Requests
|
||||
|
||||
We welcome feature requests! Please:
|
||||
- Check if the feature already exists
|
||||
- Explain the use case
|
||||
- Describe the expected behavior
|
||||
- Consider if it aligns with project goals
|
||||
|
||||
## Code of Conduct
|
||||
|
||||
- Be respectful and inclusive
|
||||
- Welcome newcomers
|
||||
- Focus on constructive feedback
|
||||
- Respect differing opinions
|
||||
|
||||
Thank you for contributing!
|
||||
Vendored
+122
@@ -0,0 +1,122 @@
|
||||
# Caddy 2 Configuration for EditmaskwithAI
|
||||
# ==========================================
|
||||
#
|
||||
# SETUP INSTRUCTIONS:
|
||||
# 1. Replace 'your-subdomain.yourdomain.com' with your actual domain
|
||||
# 2. Make sure DNS CNAME record points to your server
|
||||
# 3. Ensure ports 80 and 443 are open (Caddy handles SSL automatically)
|
||||
# 4. The frontend runs on port 3080 by default (docker-compose)
|
||||
#
|
||||
# Common Issues:
|
||||
# - "Connection refused": Check if the frontend container is running
|
||||
# - "Bad gateway": Check if localhost:3080 is accessible
|
||||
# - "SSL error": Make sure ports 80/443 are open for Let's Encrypt
|
||||
|
||||
# ============================================
|
||||
# OPTION 1: Domain with automatic HTTPS (recommended)
|
||||
# ============================================
|
||||
# Replace with your actual domain
|
||||
your-subdomain.yourdomain.com {
|
||||
# Reverse proxy to frontend (nginx serves both frontend and proxies API)
|
||||
reverse_proxy localhost:3080 {
|
||||
# Health checks
|
||||
health_uri /health
|
||||
health_interval 30s
|
||||
health_timeout 10s
|
||||
|
||||
# Headers for proper proxying
|
||||
header_up Host {upstream_hostport}
|
||||
header_up X-Real-IP {remote_host}
|
||||
header_up X-Forwarded-For {remote_host}
|
||||
header_up X-Forwarded-Proto {scheme}
|
||||
}
|
||||
|
||||
# Enable compression
|
||||
encode gzip zstd
|
||||
|
||||
# Logging (optional - uncomment for debugging)
|
||||
# log {
|
||||
# output file /var/log/caddy/access.log
|
||||
# format json
|
||||
# }
|
||||
}
|
||||
|
||||
# ============================================
|
||||
# OPTION 2: IP address or localhost (no HTTPS)
|
||||
# ============================================
|
||||
# Uncomment this block and comment out Option 1 if you don't have a domain
|
||||
# or want to test locally
|
||||
|
||||
# :8080 {
|
||||
# reverse_proxy localhost:3080 {
|
||||
# header_up Host {upstream_hostport}
|
||||
# header_up X-Real-IP {remote_host}
|
||||
# header_up X-Forwarded-For {remote_host}
|
||||
# }
|
||||
# encode gzip zstd
|
||||
# }
|
||||
|
||||
# ============================================
|
||||
# OPTION 3: Multiple subdomains
|
||||
# ============================================
|
||||
# If you want both www and non-www versions
|
||||
|
||||
# yourdomain.com, www.yourdomain.com {
|
||||
# reverse_proxy localhost:3080 {
|
||||
# header_up Host {upstream_hostport}
|
||||
# header_up X-Real-IP {remote_host}
|
||||
# header_up X-Forwarded-For {remote_host}
|
||||
# header_up X-Forwarded-Proto {scheme}
|
||||
# }
|
||||
# encode gzip zstd
|
||||
# }
|
||||
|
||||
# ============================================
|
||||
# OPTION 4: Behind another reverse proxy (Cloudflare, etc.)
|
||||
# ============================================
|
||||
# Use this if Caddy is behind Cloudflare or another proxy
|
||||
|
||||
# your-subdomain.yourdomain.com {
|
||||
# # Trust proxy headers from upstream
|
||||
# servers {
|
||||
# trusted_proxies static 173.245.48.0/20 103.21.244.0/22 103.22.200.0/22 103.31.4.0/22 141.101.64.0/18 108.162.192.0/18 190.93.240.0/20 188.114.96.0/20 197.234.240.0/22 198.41.128.0/17 162.158.0.0/15 104.16.0.0/13 104.24.0.0/14 172.64.0.0/13 131.0.72.0/22
|
||||
# }
|
||||
#
|
||||
# reverse_proxy localhost:3080 {
|
||||
# header_up Host {upstream_hostport}
|
||||
# header_up X-Real-IP {http.request.header.CF-Connecting-IP}
|
||||
# header_up X-Forwarded-For {http.request.header.CF-Connecting-IP}
|
||||
# header_up X-Forwarded-Proto {scheme}
|
||||
# }
|
||||
# encode gzip zstd
|
||||
# }
|
||||
|
||||
# ============================================
|
||||
# TROUBLESHOOTING
|
||||
# ============================================
|
||||
#
|
||||
# 1. Check Caddy logs:
|
||||
# docker logs caddy
|
||||
# OR: journalctl -u caddy -f
|
||||
#
|
||||
# 2. Test backend connectivity:
|
||||
# curl -I http://localhost:3080
|
||||
#
|
||||
# 3. Check DNS resolution:
|
||||
# dig your-subdomain.yourdomain.com
|
||||
# nslookup your-subdomain.yourdomain.com
|
||||
#
|
||||
# 4. Verify ports are open:
|
||||
# sudo netstat -tlnp | grep -E ':(80|443|3080)'
|
||||
#
|
||||
# 5. Check firewall:
|
||||
# sudo ufw status
|
||||
# sudo iptables -L -n
|
||||
#
|
||||
# 6. For Let's Encrypt issues:
|
||||
# - Ensure ports 80 and 443 are accessible from internet
|
||||
# - Check if domain resolves to your server's IP
|
||||
# - Try: caddy validate --config /path/to/Caddyfile
|
||||
#
|
||||
# 7. Force reload Caddy config:
|
||||
# caddy reload --config /path/to/Caddyfile
|
||||
Vendored
+68
@@ -0,0 +1,68 @@
|
||||
# ==============================================================================
|
||||
# AI Photo Edit - Unified Container
|
||||
# Builds miniPaint frontend and serves it alongside FastAPI backend
|
||||
# ==============================================================================
|
||||
|
||||
# Stage 1: Build miniPaint frontend
|
||||
FROM node:20-alpine AS frontend-build
|
||||
|
||||
WORKDIR /frontend
|
||||
|
||||
# Copy package files and install dependencies
|
||||
COPY frontend/package.json frontend/package-lock.json* ./
|
||||
RUN npm install
|
||||
|
||||
# Copy frontend source and build
|
||||
COPY frontend/ ./
|
||||
RUN npm run build
|
||||
|
||||
# Stage 2: Python backend with frontend static files
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies for OpenCV, rembg, SAM, and image processing
|
||||
# Note: onnxruntime 1.17+ fixed executable stack issues, no longer need execstack
|
||||
RUN apt-get update && apt-get install -y \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
libgomp1 \
|
||||
wget \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy requirements and install Python dependencies
|
||||
COPY backend/requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Verify rembg loads correctly (model downloads on first use)
|
||||
# rembg now supports BiRefNet models which are state-of-the-art for background removal
|
||||
RUN python -c "from rembg import remove; print('rembg ready')" || echo "WARNING: rembg not available - Remove Background will be disabled"
|
||||
|
||||
# Copy backend application
|
||||
COPY backend/ .
|
||||
|
||||
# Copy entrypoint script
|
||||
COPY backend/entrypoint.sh /entrypoint.sh
|
||||
RUN chmod +x /entrypoint.sh
|
||||
|
||||
# Copy scripts
|
||||
COPY scripts/ /scripts/
|
||||
|
||||
# Copy miniPaint frontend files from stage 1
|
||||
COPY --from=frontend-build /frontend/index.html /app/static/
|
||||
COPY --from=frontend-build /frontend/dist /app/static/dist
|
||||
COPY --from=frontend-build /frontend/images /app/static/images
|
||||
COPY --from=frontend-build /frontend/src/css /app/static/src/css
|
||||
|
||||
# Create data directories
|
||||
RUN mkdir -p /app/data/projects /app/data/patches /app/data/models
|
||||
|
||||
# Expose port
|
||||
EXPOSE 8000
|
||||
|
||||
# Use entrypoint script
|
||||
ENTRYPOINT ["/entrypoint.sh"]
|
||||
Vendored
+82
@@ -0,0 +1,82 @@
|
||||
# =============================================================================
|
||||
# EditmaskwithAI — GPU Container (NVIDIA CUDA)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker-compose.gpu.yml up --build
|
||||
#
|
||||
# Requirements on host:
|
||||
# - NVIDIA driver ≥ 525 (for CUDA 12.x)
|
||||
# - nvidia-container-toolkit installed and configured
|
||||
# - docker compose v2 (or docker-compose with GPU device support)
|
||||
#
|
||||
# AMD ROCm users: replace the pytorch base image with a ROCm variant, e.g.
|
||||
# rocm/pytorch:latest (and remove the nvidia-smi check below)
|
||||
# =============================================================================
|
||||
|
||||
# ── Stage 1: Build miniPaint frontend ────────────────────────────────────────
|
||||
FROM node:20-alpine AS frontend-build
|
||||
|
||||
WORKDIR /frontend
|
||||
COPY frontend/package.json frontend/package-lock.json* ./
|
||||
RUN npm install
|
||||
COPY frontend/ ./
|
||||
RUN npm run build
|
||||
|
||||
# ── Stage 2: PyTorch CUDA runtime ────────────────────────────────────────────
|
||||
# pytorch/pytorch already includes torch + torchvision built for CUDA 12.1.
|
||||
# Using the runtime (not devel) image keeps the layer lean.
|
||||
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# System dependencies
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
libgomp1 \
|
||||
wget \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Python dependencies — base + GPU extras
|
||||
# BUILDID forces pip layers to re-run when you need fresh packages without a full --no-cache:
|
||||
# BUILDID=$(date +%s) docker compose -f docker-compose.gpu.yml up --build
|
||||
ARG BUILDID=1
|
||||
COPY backend/requirements.txt .
|
||||
COPY backend/requirements.gpu.txt .
|
||||
RUN echo "BUILDID=$BUILDID" && pip install --no-cache-dir -r requirements.txt
|
||||
RUN echo "BUILDID=$BUILDID" && pip install --no-cache-dir -r requirements.gpu.txt
|
||||
|
||||
# Smoke-test rembg (model downloads on first use)
|
||||
RUN python -c "from rembg import remove; print('rembg OK')" \
|
||||
|| echo "WARNING: rembg unavailable — Remove Background disabled"
|
||||
|
||||
# Smoke-test ben2 (weights download from HuggingFace on first use)
|
||||
RUN python -c "import ben2; print('ben2 OK')" \
|
||||
|| echo "WARNING: ben2 unavailable — Remove Background falls back to U2Net/rembg"
|
||||
|
||||
# Copy backend application
|
||||
COPY backend/ .
|
||||
|
||||
# Entrypoint
|
||||
COPY backend/entrypoint.sh /entrypoint.sh
|
||||
RUN chmod +x /entrypoint.sh
|
||||
|
||||
# Scripts (SAM download, DB init, GPU setup, etc.)
|
||||
COPY scripts/ /scripts/
|
||||
RUN chmod +x /scripts/*.py 2>/dev/null || true
|
||||
|
||||
# Copy built frontend from Stage 1
|
||||
COPY --from=frontend-build /frontend/index.html /app/static/
|
||||
COPY --from=frontend-build /frontend/dist /app/static/dist
|
||||
COPY --from=frontend-build /frontend/images /app/static/images
|
||||
COPY --from=frontend-build /frontend/src/css /app/static/src/css
|
||||
|
||||
# Persistent data directories
|
||||
RUN mkdir -p /app/data/projects /app/data/patches /app/data/models
|
||||
|
||||
EXPOSE 8000
|
||||
ENTRYPOINT ["/entrypoint.sh"]
|
||||
Vendored
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 AI Photo Edit Contributors
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
Vendored
+35
@@ -0,0 +1,35 @@
|
||||
.PHONY: help up down build logs clean dev test
|
||||
|
||||
help: ## Show this help message
|
||||
@echo 'Usage: make [target]'
|
||||
@echo ''
|
||||
@echo 'Available targets:'
|
||||
@awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z_-]+:.*?## / {printf " %-15s %s\n", $$1, $$2}' $(MAKEFILE_LIST)
|
||||
|
||||
up: ## Start the application (production)
|
||||
docker-compose up -d
|
||||
|
||||
down: ## Stop the application
|
||||
docker-compose down
|
||||
|
||||
build: ## Build all containers
|
||||
docker-compose build
|
||||
|
||||
logs: ## Show logs
|
||||
docker-compose logs -f
|
||||
|
||||
clean: ## Remove all containers, volumes, and data
|
||||
docker-compose down -v
|
||||
rm -rf data/
|
||||
|
||||
dev: ## Start the application (development mode)
|
||||
docker-compose -f docker-compose.dev.yml up
|
||||
|
||||
test: ## Run tests
|
||||
@echo "Tests not yet implemented"
|
||||
|
||||
restart: ## Restart the application
|
||||
docker-compose restart
|
||||
|
||||
ps: ## Show running containers
|
||||
docker-compose ps
|
||||
Vendored
+292
@@ -0,0 +1,292 @@
|
||||
# PaintPlus
|
||||
|
||||
_Vendored into ubuntu-post-install as the `paintplus` service. Based on EditmaskwithAI (github.com/outis1one/EditmaskwithAI)._
|
||||
|
||||
A self-hosted, web-based AI photo editor. Paint over any object, describe what you want, and the AI replaces just that region — every pixel outside your selection stays untouched.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### GPU machine (recommended — free inference, best quality)
|
||||
|
||||
```bash
|
||||
git clone https://github.com/outis1one/editmaskwithai
|
||||
cd editmaskwithai
|
||||
|
||||
# One-time setup: installs nvidia-container-toolkit, configures Docker,
|
||||
# sets up a permanent DNS fix, and prefetches all AI models on the host.
|
||||
chmod +x install-local-gpu.sh
|
||||
./install-local-gpu.sh
|
||||
|
||||
# Start the app (run this each time):
|
||||
chmod +x bring-up-local-gpu.sh
|
||||
./bring-up-local-gpu.sh
|
||||
```
|
||||
|
||||
Open **http://localhost:3080**
|
||||
|
||||
**Models (~13 GB total, one time) download automatically on the host**, outside Docker — both scripts call `./prefetch-models.sh` for you, since in-container DNS is unreliable on some hosts. They're cached in `./data/hf_cache/` and `./data/models/`, and survive rebuilds.
|
||||
|
||||
---
|
||||
|
||||
### Cloud API (no GPU required)
|
||||
|
||||
```bash
|
||||
git clone https://github.com/outis1one/editmaskwithai
|
||||
cd editmaskwithai
|
||||
cp .env.example .env
|
||||
# Edit .env: set AI_PROVIDER and your API key (see .env.example for options)
|
||||
docker compose up -d --build
|
||||
```
|
||||
|
||||
Open **http://localhost:3080**
|
||||
|
||||
---
|
||||
|
||||
### Updates (any machine)
|
||||
|
||||
```bash
|
||||
git pull
|
||||
# GPU:
|
||||
./bring-up-local-gpu.sh
|
||||
# or cloud (no GPU):
|
||||
docker compose up -d --build
|
||||
```
|
||||
|
||||
If pip packages seem stale after a pull (e.g., wrong diffusers version), force a pip layer rebuild without re-downloading the entire PyTorch base image:
|
||||
|
||||
```bash
|
||||
BUILDID=$(date +%s) ./bring-up-local-gpu.sh
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## AI Providers
|
||||
|
||||
| Provider | Setup | Cost | Quality |
|
||||
|---|---|---|---|
|
||||
| `local_gpu` | GPU machine + nvidia-container-toolkit | Free | Best (SDXL/FLUX auto-selected by VRAM) |
|
||||
| `openai` | `OPENAI_API_KEY=sk-...` | ~$0.006–0.21/image | gpt-image-2 |
|
||||
| `replicate` | `REPLICATE_API_KEY=r8_...` | ~$0.002–0.03/image | Multiple models |
|
||||
| `invokeai` | InvokeAI running on another machine | Self-hosted | FLUX/SDXL |
|
||||
| `comfyui` | ComfyUI running on another machine | Self-hosted | Any model |
|
||||
|
||||
You can also mix: set a default provider in `.env` and override per-operation in the **Image → AI Provider Settings** dialog inside the app.
|
||||
|
||||
---
|
||||
|
||||
## GPU Tier Auto-Selection
|
||||
|
||||
The app detects your GPU at startup and picks the best model it can run:
|
||||
|
||||
| Effective VRAM | Model selected | Notes |
|
||||
|---|---|---|
|
||||
| ≥ 24 GB | FLUX.1-schnell | Best quality, 4-step generation |
|
||||
| 12–24 GB | SDXL | Excellent quality |
|
||||
| 8–12 GB | SDXL + xformers | Good quality |
|
||||
| 6–8 GB | SDXL + attention slicing | Good quality, slightly slower |
|
||||
| 4–6 GB | SDXL + CPU offload | Good quality, slower (GTX 1060 6GB range) |
|
||||
| 2–4 GB | SD 1.5 | Fast, lower detail |
|
||||
| < 2 GB | SD 1.5 + CPU offload | Very slow — consider a cloud provider |
|
||||
|
||||
Override the auto-selected model with `HF_MODEL_TXT2IMG`, `HF_MODEL_INPAINT` in `.env`.
|
||||
|
||||
---
|
||||
|
||||
## What it can do
|
||||
|
||||
### Selection
|
||||
- **Smart Select (SAM brush)** — paint over an object, AI detects its exact boundaries
|
||||
- **Smart Select (click)** — click any object, SAM selects it
|
||||
- **Rectangle / Ellipse / Lasso** — classic selection tools
|
||||
|
||||
### After selecting
|
||||
- **AI Edit** — describe what to change ("add a scar", "make it look aged")
|
||||
- **Make less symmetrical** — AI adds natural organic variation
|
||||
- **Replace with clipboard** — paste any image into the selection shape
|
||||
- **Scale by %** — make the selected object bigger/smaller, AI fills the gap
|
||||
- **Copy / Cut to layer** — non-destructive layer workflow
|
||||
- **Erase** — remove the selected region with AI fill
|
||||
|
||||
### Image tools
|
||||
- **Text → Image** — generate from a text description (GPU or cloud)
|
||||
- **Upscale** — Real-ESRGAN AI upscaling (genuinely adds detail, not just resize)
|
||||
- **Prepare for Print** — one-click: AI upscale to target DPI + fit to frame
|
||||
- **Fit to Frame** — resize/crop/AI-extend to standard print sizes
|
||||
- **Expand Canvas (Outpaint)** — AI extends the image in any direction
|
||||
- **Remove Background** — one-click background removal (BEN2 by default, BiRefNet-HR or U2Net selectable)
|
||||
|
||||
### Print presets
|
||||
Frame sizes: 4×6, 5×7, 8×10, 11×14, 16×20, 18×24, 20×24, 24×36 (portrait + landscape)
|
||||
DPI options: 72, 150, 200, 300 — 200 DPI is fine for 18×24" and larger (viewed from distance)
|
||||
|
||||
---
|
||||
|
||||
## Progress bars
|
||||
|
||||
All AI operations show a real-time progress overlay. For local GPU inference, the bar advances step-by-step as the model denoises (e.g. "Step 14 / 30"). For cloud providers and upscale operations, it animates to indicate activity.
|
||||
|
||||
---
|
||||
|
||||
## Logs
|
||||
|
||||
```bash
|
||||
# GPU container:
|
||||
docker compose -f docker-compose.gpu.yml logs -f
|
||||
|
||||
# Standard container:
|
||||
docker compose logs -f
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## File structure
|
||||
|
||||
```
|
||||
EditmaskwithAI/
|
||||
├── backend/
|
||||
│ ├── app/
|
||||
│ │ ├── routers/ # API endpoints (ai_tools, print_tools, …)
|
||||
│ │ ├── services/ # gpu_detect, local_diffusion, upscale, …
|
||||
│ │ └── config.py
|
||||
│ ├── requirements.txt
|
||||
│ └── requirements.gpu.txt
|
||||
├── frontend/
|
||||
│ └── src/js/
|
||||
│ ├── tools/ # brush_select (SAM paint), smart_select, …
|
||||
│ ├── modules/
|
||||
│ │ ├── generate/ # text_to_image, outpaint
|
||||
│ │ └── image/ # upscale, frame_fit, print_prepare, …
|
||||
│ └── libs/
|
||||
│ └── progress_overlay.js
|
||||
├── docker-compose.yml # Cloud / no-GPU
|
||||
├── docker-compose.gpu.yml # NVIDIA GPU (recommended)
|
||||
├── docker-compose.dev.yml # Dev with hot reload
|
||||
├── Dockerfile
|
||||
├── Dockerfile.gpu
|
||||
├── install-local-gpu.sh # One-time GPU host setup
|
||||
├── bring-up-local-gpu.sh # Start/stop the GPU container
|
||||
├── prefetch-models.sh # Download AI models on the host (called automatically; also runnable standalone)
|
||||
└── .env.example
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**GPU not detected in Docker**
|
||||
```bash
|
||||
# Check toolkit is installed and Docker restarted:
|
||||
docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi
|
||||
# If that fails, re-run: sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker
|
||||
```
|
||||
|
||||
**Model download stalls or fails**
|
||||
```bash
|
||||
# Check logs for HuggingFace errors:
|
||||
docker compose -f docker-compose.gpu.yml logs -f | grep -E "local_gpu|Error|Failed"
|
||||
# If a private/gated model: add HF_TOKEN=hf_... to .env
|
||||
```
|
||||
|
||||
**SAM model fails to download (DNS error / firewall blocking port 53)**
|
||||
|
||||
If the container can't reach `dl.fbaipublicfiles.com` (you'll see `Errno -3 Name or service not known` in the logs), download SAM directly on the host and let the bind mount make it visible to the container — no rebuild needed:
|
||||
|
||||
```bash
|
||||
./prefetch-models.sh
|
||||
# or manually:
|
||||
mkdir -p ./data/models
|
||||
# sudo needed if ./data/ was created by Docker (root-owned):
|
||||
sudo curl -L -o ./data/models/sam_vit_b_01ec64.pth \
|
||||
https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
|
||||
```
|
||||
|
||||
The file is ~375 MB. Once it exists at `./data/models/sam_vit_b_01ec64.pth`, the container picks it up on the next startup (no rebuild required). Verify with:
|
||||
```bash
|
||||
docker compose -f docker-compose.gpu.yml logs | grep -i sam
|
||||
# Should show: "SAM model loaded on cuda" (or cpu)
|
||||
```
|
||||
|
||||
If Docker created `./data/` as root and you can't write there without `sudo`, you can also use root's curl as above — the container reads the file regardless of owner.
|
||||
|
||||
**Remove Background fails ("Install ben2, u2net, or rembg")**
|
||||
|
||||
Remove Background tries, in order: the model set by `BG_REMOVAL_MODEL` (default `ben2`), then the other local models, then `rembg` as a last resort. You'll see this error only if all of them fail.
|
||||
|
||||
- **ben2 / birefnet-hr** (GPU image only) download their weights from HuggingFace on first use, cached under `./data/hf_cache`. If that download fails (DNS/firewall, see above), run `./prefetch-models.sh` to fetch both directly on the host, or check the logs for the specific error:
|
||||
```bash
|
||||
docker compose -f docker-compose.gpu.yml logs -f | grep -iE "ben2|birefnet"
|
||||
```
|
||||
- **u2net** auto-downloads (~176MB) from GitHub on first use, same as SAM. If that fails too, download it directly on the host:
|
||||
```bash
|
||||
./prefetch-models.sh
|
||||
# or manually:
|
||||
mkdir -p ./data/models
|
||||
sudo curl -L -o ./data/models/u2net.onnx \
|
||||
https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
|
||||
```
|
||||
The file is ~176 MB. Once it exists at `./data/models/u2net.onnx`, the next "Remove Background" click picks it up — no rebuild or restart needed. Verify with:
|
||||
```bash
|
||||
docker compose logs -f | grep -i u2net
|
||||
# Should show: "U2Net model loaded successfully with OpenCV DNN"
|
||||
```
|
||||
|
||||
You can also pick a specific model per-edit from the Remove Background dialog's model dropdown, overriding `BG_REMOVAL_MODEL` for that one call.
|
||||
|
||||
**AI models not downloading (container DNS blocked)**
|
||||
|
||||
`install-local-gpu.sh` and `bring-up-local-gpu.sh` already run this for you automatically on every start, so you normally don't need to think about it. If a model still didn't download (no network at the time, etc.), re-run it manually — it lands in `./data/`, which is already bind-mounted into the container, so it's picked up with no rebuild:
|
||||
|
||||
```bash
|
||||
./prefetch-models.sh # SAM + U2Net + BEN2 + BiRefNet-HR (~1.5GB)
|
||||
./prefetch-models.sh --sdxl # also Text→Image / AI Edit models (~13GB)
|
||||
./bring-up-local-gpu.sh
|
||||
```
|
||||
|
||||
If that also fails to reach the network, the problem is host-level (firewall/DNS), not Docker-specific — see your network/firewall configuration.
|
||||
|
||||
Alternatively, if you ran `./install-local-gpu.sh`, container DNS is already permanently fixed via a systemd-managed iptables rule. If you skipped that script, apply the same fix manually (does **not** affect container isolation):
|
||||
|
||||
```bash
|
||||
sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT
|
||||
./bring-up-local-gpu.sh
|
||||
```
|
||||
|
||||
The container will now resolve hostnames and download models automatically (~13 GB on first run, then cached). Watch progress:
|
||||
```bash
|
||||
docker compose -f docker-compose.gpu.yml logs -f | grep -E "local_gpu|Cached|failed"
|
||||
```
|
||||
|
||||
**Alternative: download with a Docker helper container** (no host Python needed):
|
||||
|
||||
```bash
|
||||
# Inpainting model (~6.5 GB) — needed for AI Edit, Make less symmetrical, etc.
|
||||
docker run --rm \
|
||||
-v "$(pwd)/data/hf_cache:/root/.cache/huggingface" \
|
||||
python:3.11-slim \
|
||||
bash -c "pip install -q huggingface-hub && \
|
||||
huggingface-cli download diffusers/stable-diffusion-xl-1.0-inpainting-0.1 \
|
||||
--exclude '*.msgpack' 'flax_*' 'tf_*'"
|
||||
|
||||
# Text-to-image model (~6.5 GB) — needed for Text → Image
|
||||
docker run --rm \
|
||||
-v "$(pwd)/data/hf_cache:/root/.cache/huggingface" \
|
||||
python:3.11-slim \
|
||||
bash -c "pip install -q huggingface-hub && \
|
||||
huggingface-cli download stabilityai/stable-diffusion-xl-base-1.0 \
|
||||
--exclude '*.msgpack' 'flax_*' 'tf_*'"
|
||||
```
|
||||
|
||||
Then restart: `docker compose -f docker-compose.gpu.yml restart`
|
||||
|
||||
**Out of VRAM during generation**
|
||||
- Reduce `LOCAL_GPU_MAX_PIPELINES=1` in `.env` (default 2)
|
||||
- Or override to a smaller model: `HF_MODEL_TXT2IMG=runwayml/stable-diffusion-v1-5`
|
||||
|
||||
**Settings saved locally only**
|
||||
- The in-app AI Provider Settings dialog saves to localStorage for the session
|
||||
- To make settings permanent: edit `.env` and rebuild
|
||||
|
||||
**Check API docs**
|
||||
```
|
||||
http://localhost:3080/api/docs
|
||||
```
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
# Database
|
||||
DATABASE_URL=sqlite:///./data/ai_photo_edit.db
|
||||
|
||||
# Security
|
||||
SECRET_KEY=your-secret-key-here-change-in-production
|
||||
ALGORITHM=HS256
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES=30
|
||||
|
||||
# AI Provider (configure based on your provider)
|
||||
AI_PROVIDER=openai
|
||||
OPENAI_API_KEY=your-openai-api-key-here
|
||||
# Alternative providers (uncomment as needed)
|
||||
# AI_PROVIDER=stability
|
||||
# STABILITY_API_KEY=your-stability-api-key-here
|
||||
|
||||
# File Storage
|
||||
DATA_DIR=/app/data
|
||||
MAX_UPLOAD_SIZE_MB=50
|
||||
|
||||
# CORS
|
||||
CORS_ORIGINS=http://localhost:3000,http://localhost:5173
|
||||
Vendored
+37
@@ -0,0 +1,37 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies for OpenCV, SAM, and image processing
|
||||
RUN apt-get update && apt-get install -y \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
libgomp1 \
|
||||
wget \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy requirements
|
||||
COPY requirements.txt .
|
||||
|
||||
# Install Python dependencies
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy application
|
||||
COPY . .
|
||||
|
||||
# Copy entrypoint script and make it executable
|
||||
COPY entrypoint.sh /entrypoint.sh
|
||||
RUN chmod +x /entrypoint.sh
|
||||
|
||||
# Create data directories
|
||||
RUN mkdir -p /app/data/projects /app/data/patches /app/data/models
|
||||
|
||||
# Expose port
|
||||
EXPOSE 8000
|
||||
|
||||
# Use entrypoint script (auto-populates eyes on first run, then starts server)
|
||||
ENTRYPOINT ["/entrypoint.sh"]
|
||||
+80
@@ -0,0 +1,80 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
from typing import List
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
# Database
|
||||
database_url: str = "sqlite:///./data/ai_photo_edit.db"
|
||||
|
||||
# Security
|
||||
secret_key: str = "your-secret-key-change-in-production"
|
||||
algorithm: str = "HS256"
|
||||
access_token_expire_minutes: int = 30
|
||||
|
||||
# AI Provider
|
||||
# Local: blank or "mock" — always available, no config needed
|
||||
# Remote default (used for any operation without a specific override):
|
||||
# openai | invokeai | comfyui | replicate | stability
|
||||
ai_provider: str = "mock"
|
||||
|
||||
# Per-operation provider overrides — blank means use ai_provider default.
|
||||
# Operations: inpaint, txt2img, img2img, outpaint
|
||||
# Example: AI_PROVIDER_TXT2IMG=openai (use OpenAI for text-to-image only)
|
||||
ai_provider_inpaint: str = "" # remote inpaint / replace selection
|
||||
ai_provider_txt2img: str = "" # text-to-image
|
||||
ai_provider_img2img: str = "" # image-to-image
|
||||
ai_provider_outpaint: str = "" # expand canvas
|
||||
|
||||
# Provider API Keys
|
||||
openai_api_key: str = ""
|
||||
openai_model: str = "gpt-image-2" # text-to-image (generations endpoint)
|
||||
openai_edit_model: str = "gpt-image-2" # inpaint/img2img/outpaint (edits endpoint)
|
||||
stability_api_key: str = ""
|
||||
replicate_api_key: str = ""
|
||||
|
||||
# InvokeAI (self-hosted)
|
||||
invokeai_url: str = ""
|
||||
invokeai_default_model: str = "flux-dev"
|
||||
|
||||
# ComfyUI (self-hosted)
|
||||
comfyui_url: str = ""
|
||||
comfyui_default_model: str = "v1-5-pruned-emaonly.ckpt"
|
||||
|
||||
# Model Selection (optional, provider-specific)
|
||||
stability_model: str = "sdxl" # Options: sdxl, sd15, sd21
|
||||
replicate_model: str = "sdxl-inpaint" # Options: sdxl-inpaint, lama, realistic-vision
|
||||
|
||||
# Allow per-edit model override
|
||||
allow_model_override: bool = True
|
||||
|
||||
# Remove Background — preferred local model when request.model="auto"
|
||||
# Options: ben2 (default, best for clean cutouts/hair), birefnet-hr (best
|
||||
# for high-res/print work), u2net (lightweight, smallest download)
|
||||
bg_removal_model: str = "ben2"
|
||||
|
||||
# Local GPU diffusion (AI_PROVIDER=local_gpu)
|
||||
auto_download_models: bool = True # download HF models on first use
|
||||
local_gpu_max_pipelines: int = 2 # max diffusion pipelines kept in GPU memory
|
||||
hf_token: str = "" # HuggingFace token (only needed for gated models)
|
||||
# Override auto-selected models per operation (leave blank = auto-pick by VRAM tier)
|
||||
hf_model_inpaint: str = ""
|
||||
hf_model_txt2img: str = ""
|
||||
hf_model_img2img: str = ""
|
||||
|
||||
# File Storage
|
||||
data_dir: str = "./data"
|
||||
max_upload_size_mb: int = 50
|
||||
|
||||
# CORS
|
||||
cors_origins: str = "http://localhost:3000,http://localhost:5173"
|
||||
|
||||
@property
|
||||
def cors_origins_list(self) -> List[str]:
|
||||
return [origin.strip() for origin in self.cors_origins.split(",")]
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
settings = Settings()
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from app.config import settings
|
||||
import os
|
||||
|
||||
# Ensure data directory exists
|
||||
os.makedirs("./data", exist_ok=True)
|
||||
|
||||
engine = create_engine(
|
||||
settings.database_url,
|
||||
connect_args={"check_same_thread": False} # Needed for SQLite
|
||||
)
|
||||
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
def get_db():
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def init_db():
|
||||
"""Initialize database tables"""
|
||||
Base.metadata.create_all(bind=engine)
|
||||
Vendored
+162
@@ -0,0 +1,162 @@
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.responses import FileResponse, HTMLResponse
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from app.config import settings
|
||||
from app.database import init_db
|
||||
from app.routers import projects, edits, images, patches, generate, tools, ai_tools, print_tools
|
||||
from app.routers import gpu_status
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Initialize database on startup; auto-install Real-ESRGAN NCNN in background."""
|
||||
init_db()
|
||||
# Kick off NCNN install in background if no AI upscaler detected
|
||||
from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed
|
||||
caps = probe_upscale_capabilities()
|
||||
if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]:
|
||||
asyncio.create_task(ensure_ncnn_installed())
|
||||
# Pre-download SAM model in background so first click is fast
|
||||
from app.services.sam_service import ensure_sam_installed
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
|
||||
# If local GPU provider is active, log GPU info at startup
|
||||
if settings.ai_provider.lower() == "local_gpu" or any(
|
||||
v.lower() == "local_gpu"
|
||||
for v in [
|
||||
settings.ai_provider_inpaint,
|
||||
settings.ai_provider_txt2img,
|
||||
settings.ai_provider_img2img,
|
||||
settings.ai_provider_outpaint,
|
||||
]
|
||||
if v
|
||||
):
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
info = get_cached_gpu_info()
|
||||
cc_str = f" | CC={info.compute_capability}" if info.compute_capability else ""
|
||||
print(
|
||||
f"[gpu] {info.device_name} | {info.vram_total_gb:.1f} GB{cc_str} | "
|
||||
f"tier={info.tier} | fp16={info.fp16}"
|
||||
)
|
||||
for w in info.warnings:
|
||||
print(f"[gpu] ⚠ {w}")
|
||||
if settings.auto_download_models:
|
||||
# Download model weight files to disk cache in background so first
|
||||
# user request loads from local disk instead of the internet.
|
||||
from app.services.local_diffusion import prefetch_model_files
|
||||
asyncio.create_task(prefetch_model_files())
|
||||
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="AI Photo Edit API",
|
||||
description="API for AI-powered photo editing with mask-scoped regeneration",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
# Configure CORS
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.cors_origins_list,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Include routers
|
||||
app.include_router(projects.router)
|
||||
app.include_router(edits.router)
|
||||
app.include_router(images.router)
|
||||
app.include_router(patches.router)
|
||||
app.include_router(generate.router)
|
||||
app.include_router(tools.router)
|
||||
app.include_router(ai_tools.router)
|
||||
app.include_router(print_tools.router)
|
||||
app.include_router(gpu_status.router)
|
||||
|
||||
|
||||
@app.get("/api")
|
||||
def api_root():
|
||||
"""API info endpoint"""
|
||||
return {
|
||||
"name": "AI Photo Edit API",
|
||||
"version": "1.0.0",
|
||||
"status": "running"
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
"""Health check endpoint"""
|
||||
return {"status": "healthy"}
|
||||
|
||||
|
||||
# Static files directory
|
||||
STATIC_DIR = Path("/app/static")
|
||||
|
||||
|
||||
class NoCacheStaticFiles(StaticFiles):
|
||||
"""webpack outputs a fixed 'bundle.js' filename (no content hash), so
|
||||
browsers can keep serving a stale cached copy after a rebuild unless
|
||||
forced to revalidate on every request."""
|
||||
|
||||
def file_response(self, *args, **kwargs):
|
||||
response = super().file_response(*args, **kwargs)
|
||||
response.headers["Cache-Control"] = "no-cache"
|
||||
return response
|
||||
|
||||
|
||||
# Serve static assets - mount subdirectories if they exist
|
||||
if STATIC_DIR.exists():
|
||||
# React-style assets folder
|
||||
if (STATIC_DIR / "assets").exists():
|
||||
app.mount("/assets", StaticFiles(directory=STATIC_DIR / "assets"), name="assets")
|
||||
# miniPaint dist folder (webpack bundle) - no-cache so code updates are picked up immediately
|
||||
if (STATIC_DIR / "dist").exists():
|
||||
app.mount("/dist", NoCacheStaticFiles(directory=STATIC_DIR / "dist"), name="dist")
|
||||
# miniPaint images folder
|
||||
if (STATIC_DIR / "images").exists():
|
||||
app.mount("/images", StaticFiles(directory=STATIC_DIR / "images"), name="images")
|
||||
# miniPaint CSS folder - no-cache, same reasoning as /dist
|
||||
if (STATIC_DIR / "src").exists():
|
||||
app.mount("/src", NoCacheStaticFiles(directory=STATIC_DIR / "src"), name="src")
|
||||
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
async def serve_spa():
|
||||
"""Serve miniPaint index.html"""
|
||||
index_path = STATIC_DIR / "index.html"
|
||||
if index_path.exists():
|
||||
return FileResponse(index_path, headers={"Cache-Control": "no-cache"})
|
||||
return HTMLResponse("<h1>Frontend not built. Run npm build in frontend/</h1>")
|
||||
|
||||
|
||||
@app.get("/{full_path:path}")
|
||||
async def serve_spa_routes(request: Request, full_path: str):
|
||||
"""
|
||||
Catch-all route for serving static files.
|
||||
Serves static files if they exist, otherwise returns index.html.
|
||||
"""
|
||||
# Don't catch API routes
|
||||
if full_path.startswith(("projects", "edits", "patches", "tools", "generate", "health", "docs", "openapi.json", "api")):
|
||||
return {"detail": "Not Found"}
|
||||
|
||||
# Check if it's a static file
|
||||
static_file = STATIC_DIR / full_path
|
||||
if static_file.exists() and static_file.is_file():
|
||||
return FileResponse(static_file)
|
||||
|
||||
# Otherwise serve index.html
|
||||
index_path = STATIC_DIR / "index.html"
|
||||
if index_path.exists():
|
||||
return FileResponse(index_path, headers={"Cache-Control": "no-cache"})
|
||||
|
||||
return HTMLResponse("<h1>Frontend not built</h1>", status_code=404)
|
||||
@@ -0,0 +1,6 @@
|
||||
from app.models.user import User
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.models.patch import Patch
|
||||
|
||||
__all__ = ["User", "Project", "Edit", "Patch"]
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class Edit(Base):
|
||||
__tablename__ = "edits"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
project_id = Column(Integer, ForeignKey("projects.id"), nullable=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
mode = Column(String, nullable=False) # "A" or "B"
|
||||
prompt = Column(Text, nullable=False)
|
||||
selection_type = Column(String, nullable=False) # "rectangle", "ellipse", "lasso"
|
||||
bbox_json = Column(Text, nullable=False) # JSON string of {x, y, width, height}
|
||||
feather_px = Column(Integer, default=0)
|
||||
ai_provider = Column(String, nullable=False)
|
||||
status = Column(String, nullable=False) # "pending", "processing", "completed", "failed"
|
||||
error_message = Column(Text, nullable=True)
|
||||
|
||||
# Relationships
|
||||
project = relationship("Project", back_populates="edits")
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class Patch(Base):
|
||||
__tablename__ = "patches"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=True)
|
||||
name = Column(String, nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
# Source information
|
||||
source_type = Column(String, nullable=False) # "ai_generated", "manual_selection", "imported"
|
||||
source_project_id = Column(Integer, ForeignKey("projects.id"), nullable=True)
|
||||
source_edit_id = Column(Integer, ForeignKey("edits.id"), nullable=True)
|
||||
|
||||
# Patch metadata
|
||||
width = Column(Integer, nullable=False)
|
||||
height = Column(Integer, nullable=False)
|
||||
tags = Column(Text, nullable=True) # Comma-separated tags
|
||||
category = Column(String, nullable=True) # "hand", "face", "body", "object", "texture", etc.
|
||||
|
||||
# Is this patch shared/public?
|
||||
is_public = Column(Boolean, default=False)
|
||||
|
||||
# File path (relative to data dir)
|
||||
file_path = Column(String, nullable=False)
|
||||
thumbnail_path = Column(String, nullable=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="patches")
|
||||
source_project = relationship("Project")
|
||||
source_edit = relationship("Edit")
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class Project(Base):
|
||||
__tablename__ = "projects"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
name = Column(String, nullable=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="projects")
|
||||
edits = relationship("Edit", back_populates="project", cascade="all, delete-orphan")
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
email = Column(String, unique=True, index=True, nullable=False)
|
||||
password_hash = Column(String, nullable=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
projects = relationship("Project", back_populates="user", cascade="all, delete-orphan")
|
||||
patches = relationship("Patch", back_populates="user", cascade="all, delete-orphan")
|
||||
+989
@@ -0,0 +1,989 @@
|
||||
"""
|
||||
AI tools router — LaMa inpaint, background removal, remote generation, config.
|
||||
All endpoints are under /api prefix.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
import base64
|
||||
import asyncio
|
||||
import json
|
||||
from io import BytesIO
|
||||
|
||||
from app.services.local_inpaint import (
|
||||
lama_inpaint, opencv_inpaint, lama_available, gpu_available, rembg_available,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["ai-tools"])
|
||||
|
||||
|
||||
# ─── Request / response models ───────────────────────────────────────────────
|
||||
|
||||
class EraseRequest(BaseModel):
|
||||
image: str # base64
|
||||
mask: str # base64
|
||||
|
||||
|
||||
class InpaintRemoteRequest(BaseModel):
|
||||
image: str
|
||||
mask: str
|
||||
prompt: str
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class Txt2ImgRequest(BaseModel):
|
||||
prompt: str
|
||||
width: Optional[int] = 1024
|
||||
height: Optional[int] = 1024
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
seed: Optional[int] = 0
|
||||
|
||||
|
||||
class Img2ImgRequest(BaseModel):
|
||||
image: str
|
||||
prompt: str
|
||||
strength: Optional[float] = 0.75
|
||||
negative_prompt: Optional[str] = ""
|
||||
steps: Optional[int] = 30
|
||||
cfg_scale: Optional[float] = 7.5
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class OutpaintRequest(BaseModel):
|
||||
image: str
|
||||
direction: str # left | right | top | bottom
|
||||
size: Optional[int] = 256
|
||||
prompt: Optional[str] = ""
|
||||
|
||||
|
||||
class BgRemoveRequest(BaseModel):
|
||||
image: str
|
||||
|
||||
|
||||
# ─── Helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
def _decode(b64: str) -> bytes:
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _encode(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _require_remote(operation: str = None):
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
from app.config import settings
|
||||
provider = get_remote_provider(operation)
|
||||
if provider is None:
|
||||
if (settings.ai_provider or "").lower() == "local_gpu":
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"local_gpu provider failed to load — diffusers may be incompatible with "
|
||||
"the installed PyTorch version. Check container logs for details. "
|
||||
"If you see 'torch has no attribute xpu', rebuild the container from the "
|
||||
"correct branch so the pinned diffusers<0.29.0 is installed."
|
||||
)
|
||||
)
|
||||
op_hint = f"AI_PROVIDER_{operation.upper()} or " if operation else ""
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=f"No remote AI provider configured for '{operation or 'default'}'. "
|
||||
f"Set {op_hint}AI_PROVIDER in .env (openai / invokeai / comfyui)."
|
||||
)
|
||||
return provider
|
||||
|
||||
|
||||
# ─── Local inpaint endpoints ─────────────────────────────────────────────────
|
||||
|
||||
@router.post("/erase")
|
||||
async def erase(req: EraseRequest):
|
||||
"""
|
||||
Magic eraser: remove object / fill region using LaMa (local, no API key needed).
|
||||
Falls back to OpenCV if LaMa not installed.
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
mask_bytes = _decode(req.mask)
|
||||
|
||||
if lama_available():
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, image_bytes, mask_bytes
|
||||
)
|
||||
method = "lama"
|
||||
else:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, image_bytes, mask_bytes
|
||||
)
|
||||
method = "opencv"
|
||||
|
||||
return {"result": _encode(result), "method": method}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/inpaint/lama")
|
||||
async def inpaint_lama(req: EraseRequest):
|
||||
"""LaMa structural inpainting."""
|
||||
if not lama_available():
|
||||
raise HTTPException(status_code=503, detail="simple-lama-inpainting not installed.")
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, _decode(req.image), _decode(req.mask)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/inpaint/fast")
|
||||
async def inpaint_fast(req: EraseRequest):
|
||||
"""OpenCV fast inpainting (CPU, milliseconds)."""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, _decode(req.image), _decode(req.mask)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/background/remove")
|
||||
async def background_remove(req: BgRemoveRequest):
|
||||
"""Remove background — rembg if available, else U2Net."""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
|
||||
# Try rembg first
|
||||
if rembg_available():
|
||||
from app.services.local_inpaint import remove_background_rembg
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, remove_background_rembg, image_bytes
|
||||
)
|
||||
return {"result": _encode(result), "method": "rembg"}
|
||||
|
||||
# Fall back to U2Net (existing implementation)
|
||||
from PIL import Image
|
||||
from io import BytesIO as _BytesIO
|
||||
img = Image.open(_BytesIO(image_bytes)).convert("RGB")
|
||||
from app.routers.tools import _remove_background_u2net
|
||||
result = await _remove_background_u2net(img)
|
||||
return {"result": _encode(result), "method": "u2net"}
|
||||
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Remote provider endpoints ───────────────────────────────────────────────
|
||||
|
||||
@router.post("/inpaint/remote")
|
||||
async def inpaint_remote(req: InpaintRemoteRequest):
|
||||
"""Inpaint via configured remote provider (InvokeAI / ComfyUI / OpenAI)."""
|
||||
provider = _require_remote("inpaint")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.inpaint(_decode(req.image), _decode(req.mask), req.prompt, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/generate/progress")
|
||||
async def generation_progress_stream():
|
||||
"""
|
||||
SSE stream of local GPU pipeline inference progress.
|
||||
Events are JSON arrays of pipeline state objects, emitted every 200 ms.
|
||||
Each object: {pipeline, state, step, total_steps, progress, message, model_id, …}
|
||||
Clients open this with EventSource before firing a generation POST,
|
||||
then close it when the POST resolves.
|
||||
"""
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
|
||||
async def event_gen():
|
||||
try:
|
||||
while True:
|
||||
states = get_all_model_states()
|
||||
yield f"data: {json.dumps(states)}\n\n"
|
||||
await asyncio.sleep(0.2)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
return StreamingResponse(
|
||||
event_gen(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/generate/txt2img")
|
||||
async def txt2img(req: Txt2ImgRequest):
|
||||
"""Text-to-image via configured remote provider."""
|
||||
provider = _require_remote("txt2img")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
"seed": req.seed or 0,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.txt2img(req.prompt, req.width, req.height, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/generate/img2img")
|
||||
async def img2img(req: Img2ImgRequest):
|
||||
"""Image-to-image via configured remote provider."""
|
||||
provider = _require_remote("img2img")
|
||||
try:
|
||||
params = {
|
||||
"negative_prompt": req.negative_prompt or "",
|
||||
"steps": req.steps,
|
||||
"cfg_scale": req.cfg_scale,
|
||||
}
|
||||
if req.model:
|
||||
params["model"] = req.model
|
||||
result = await provider.img2img(_decode(req.image), req.prompt, req.strength, params)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/generate/outpaint")
|
||||
async def outpaint(req: OutpaintRequest):
|
||||
"""Expand canvas in given direction via remote provider."""
|
||||
provider = _require_remote("outpaint")
|
||||
if req.direction not in ("left", "right", "top", "bottom"):
|
||||
raise HTTPException(status_code=400, detail="direction must be left/right/top/bottom")
|
||||
try:
|
||||
result = await provider.outpaint(_decode(req.image), req.direction, req.size, req.prompt or "")
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Config / capabilities ────────────────────────────────────────────────────
|
||||
|
||||
class ConfigUpdateRequest(BaseModel):
|
||||
ai_provider: Optional[str] = None
|
||||
# Per-operation overrides (blank = use default)
|
||||
ai_provider_inpaint: Optional[str] = None
|
||||
ai_provider_txt2img: Optional[str] = None
|
||||
ai_provider_img2img: Optional[str] = None
|
||||
ai_provider_outpaint: Optional[str] = None
|
||||
# Credentials / URLs
|
||||
openai_api_key: Optional[str] = None
|
||||
openai_model: Optional[str] = None
|
||||
invokeai_url: Optional[str] = None
|
||||
invokeai_default_model: Optional[str] = None
|
||||
comfyui_url: Optional[str] = None
|
||||
comfyui_default_model: Optional[str] = None
|
||||
replicate_api_key: Optional[str] = None
|
||||
stability_api_key: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/config")
|
||||
async def update_config(req: ConfigUpdateRequest):
|
||||
"""
|
||||
Apply runtime provider settings (no restart needed).
|
||||
Values are applied to the live settings object in-process.
|
||||
They do NOT persist across restarts — set them in .env for permanence.
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
_str_fields = [
|
||||
"ai_provider", "ai_provider_inpaint", "ai_provider_txt2img",
|
||||
"ai_provider_img2img", "ai_provider_outpaint",
|
||||
"openai_api_key", "openai_model",
|
||||
"invokeai_url", "invokeai_default_model",
|
||||
"comfyui_url", "comfyui_default_model",
|
||||
"replicate_api_key", "stability_api_key",
|
||||
]
|
||||
for field in _str_fields:
|
||||
val = getattr(req, field, None)
|
||||
if val is not None:
|
||||
setattr(settings, field, val)
|
||||
|
||||
return {
|
||||
"status": "ok",
|
||||
"ai_provider": settings.ai_provider,
|
||||
"overrides": {
|
||||
"inpaint": settings.ai_provider_inpaint or None,
|
||||
"txt2img": settings.ai_provider_txt2img or None,
|
||||
"img2img": settings.ai_provider_img2img or None,
|
||||
"outpaint": settings.ai_provider_outpaint or None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
async def _check_provider(operation: str) -> dict:
|
||||
"""Health-check the provider for a specific operation."""
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
try:
|
||||
p = get_remote_provider(operation)
|
||||
if p is None:
|
||||
return {"provider": None, "healthy": False}
|
||||
healthy = await asyncio.wait_for(p.health(), timeout=5.0)
|
||||
return {"provider": p.__class__.__name__.replace("Provider", "").lower(), "healthy": healthy}
|
||||
except Exception:
|
||||
return {"provider": None, "healthy": False}
|
||||
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config():
|
||||
"""
|
||||
Return capability flags so the frontend can show/hide tools.
|
||||
Includes per-operation provider assignments and health status.
|
||||
"""
|
||||
from app.config import settings
|
||||
|
||||
# Run health checks for each operation concurrently
|
||||
ops = ["inpaint", "txt2img", "img2img", "outpaint"]
|
||||
results = await asyncio.gather(*[_check_provider(op) for op in ops])
|
||||
op_status = dict(zip(ops, results))
|
||||
|
||||
# Default provider for display (used when no per-op override)
|
||||
default_name = (settings.ai_provider or "").lower() or None
|
||||
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
gpu_info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
"local": {
|
||||
"lama": lama_available(),
|
||||
"rembg": rembg_available(),
|
||||
"opencv": True,
|
||||
"gpu_detected": gpu_available(),
|
||||
"gpu_backend": gpu_info.backend,
|
||||
"gpu_device": gpu_info.device_name,
|
||||
"gpu_vram_total": gpu_info.vram_total_gb,
|
||||
"gpu_vram_free": gpu_info.vram_free_gb,
|
||||
"gpu_cc": gpu_info.compute_capability,
|
||||
"gpu_fp16": gpu_info.fp16,
|
||||
"gpu_bf16": gpu_info.bf16,
|
||||
"gpu_fp8": gpu_info.fp8,
|
||||
"gpu_tensor_cores": gpu_info.tensor_cores,
|
||||
"gpu_tier": gpu_info.tier,
|
||||
"gpu_eff_vram": gpu_info.effective_vram_gb,
|
||||
"local_gpu_available": gpu_info.backend in ("cuda", "mps"),
|
||||
"local_gpu_capabilities": gpu_info.capabilities,
|
||||
"local_gpu_warnings": gpu_info.warnings,
|
||||
},
|
||||
"remote": {
|
||||
"default_provider": default_name,
|
||||
# Legacy field kept for backwards compat with badge/capabilities checks
|
||||
"provider": default_name,
|
||||
"healthy": any(v["healthy"] for v in op_status.values()),
|
||||
"operations": op_status,
|
||||
"overrides": {
|
||||
"inpaint": settings.ai_provider_inpaint or None,
|
||||
"txt2img": settings.ai_provider_txt2img or None,
|
||||
"img2img": settings.ai_provider_img2img or None,
|
||||
"outpaint": settings.ai_provider_outpaint or None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# ─── Selection image operations ─────────────────────────────────────────────
|
||||
|
||||
class ScaleSelectionRequest(BaseModel):
|
||||
image: str # base64 full canvas
|
||||
mask: str # base64 selection mask (white = object)
|
||||
scale_pct: float = 103.0 # 103 = 3% bigger, 95 = 5% smaller
|
||||
|
||||
|
||||
class AiEditRegionRequest(BaseModel):
|
||||
image: str
|
||||
mask: str
|
||||
instruction: str
|
||||
negative_prompt: str = ""
|
||||
steps: int = 30
|
||||
cfg_scale: float = 7.5
|
||||
|
||||
|
||||
class PasteIntoSelectionRequest(BaseModel):
|
||||
image: str # base64 target canvas
|
||||
mask: str # base64 selection mask
|
||||
paste_image: str # base64 image to paste
|
||||
|
||||
|
||||
@router.post("/image/scale-selection")
|
||||
async def scale_selection(req: ScaleSelectionRequest):
|
||||
"""
|
||||
Scale the object selected by mask by scale_pct%, AI-fill the exposed gap.
|
||||
Works purely with local tools (LaMa/OpenCV) — no remote provider needed.
|
||||
"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||||
|
||||
img = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||||
if img.size != mask.size:
|
||||
mask = mask.resize(img.size, Image.LANCZOS)
|
||||
|
||||
mask_arr = np.array(mask)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
if len(xs) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty mask — nothing to scale")
|
||||
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
cx, cy = (minx + maxx) / 2.0, (miny + maxy) / 2.0
|
||||
obj_w, obj_h = maxx - minx + 1, maxy - miny + 1
|
||||
|
||||
scale = req.scale_pct / 100.0
|
||||
new_w = max(1, round(obj_w * scale))
|
||||
new_h = max(1, round(obj_h * scale))
|
||||
|
||||
# Extract masked object crop (RGBA with mask as alpha)
|
||||
img_rgba = img.convert("RGBA")
|
||||
obj_crop = img_rgba.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
r, g, b, _ = obj_crop.split()
|
||||
obj_masked = Image.merge("RGBA", (r, g, b, mask_crop))
|
||||
scaled_obj = obj_masked.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# AI-fill the original mask area (gap) with LaMa/OpenCV
|
||||
gap_mask = mask.filter(ImageFilter.MaxFilter(9)) # expand ~4px for clean seam
|
||||
gap_bytes = BytesIO()
|
||||
img.save(gap_bytes, format="PNG")
|
||||
gap_mask_bytes = BytesIO()
|
||||
gap_mask.save(gap_mask_bytes, format="PNG")
|
||||
|
||||
try:
|
||||
if lama_available():
|
||||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||||
None, lama_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||||
)
|
||||
else:
|
||||
filled_bytes = await asyncio.get_event_loop().run_in_executor(
|
||||
None, opencv_inpaint, gap_bytes.getvalue(), gap_mask_bytes.getvalue()
|
||||
)
|
||||
filled = Image.open(BytesIO(filled_bytes)).convert("RGBA")
|
||||
except Exception as exc:
|
||||
print(f"[scale-selection] fill fallback: {exc}")
|
||||
filled = img.convert("RGBA")
|
||||
|
||||
# Paste scaled object centered on original centroid
|
||||
px = round(cx - new_w / 2)
|
||||
py = round(cy - new_h / 2)
|
||||
result = filled.copy()
|
||||
result.paste(scaled_obj, (px, py), scaled_obj.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return {"result": _encode(out.getvalue())}
|
||||
|
||||
|
||||
@router.post("/image/ai-edit-region")
|
||||
async def ai_edit_region(req: AiEditRegionRequest):
|
||||
"""
|
||||
AI-edit the selected region using the configured inpaint provider.
|
||||
Works with local_gpu, InvokeAI, ComfyUI, or OpenAI.
|
||||
"""
|
||||
provider = _require_remote("inpaint")
|
||||
try:
|
||||
result_bytes = await provider.inpaint(
|
||||
_decode(req.image),
|
||||
_decode(req.mask),
|
||||
req.instruction,
|
||||
{"negative_prompt": req.negative_prompt, "steps": req.steps, "cfg_scale": req.cfg_scale},
|
||||
)
|
||||
except Exception as exc:
|
||||
import traceback; traceback.print_exc()
|
||||
msg = str(exc)
|
||||
if "Errno -3" in msg or "Name or service not known" in msg or "ConnectError" in msg:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=(
|
||||
"AI model files not yet downloaded — container DNS appears to be blocked. "
|
||||
"Fix: sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT on the host, "
|
||||
"or pre-download the model: pip install huggingface-hub && "
|
||||
"huggingface-cli download diffusers/stable-diffusion-xl-1.0-inpainting-0.1 "
|
||||
"--cache-dir ./data/hf_cache"
|
||||
)
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=msg)
|
||||
return {"result": _encode(result_bytes)}
|
||||
|
||||
|
||||
@router.post("/image/paste-into-selection")
|
||||
async def paste_into_selection(req: PasteIntoSelectionRequest):
|
||||
"""
|
||||
Scale a clipboard image to the selection bounding box, mask it to the
|
||||
selection shape, and composite it over the original canvas.
|
||||
"""
|
||||
try:
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
raise HTTPException(status_code=500, detail="PIL/numpy not available")
|
||||
|
||||
img = Image.open(BytesIO(_decode(req.image))).convert("RGBA")
|
||||
mask = Image.open(BytesIO(_decode(req.mask))).convert("L")
|
||||
paste_img = Image.open(BytesIO(_decode(req.paste_image))).convert("RGBA")
|
||||
|
||||
if img.size != mask.size:
|
||||
mask = mask.resize(img.size, Image.LANCZOS)
|
||||
|
||||
mask_arr = np.array(mask)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
if len(xs) == 0:
|
||||
raise HTTPException(status_code=400, detail="Empty mask")
|
||||
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale clipboard image to fit the selection bounding box
|
||||
paste_scaled = paste_img.resize((target_w, target_h), Image.LANCZOS)
|
||||
|
||||
# Clip paste to selection shape using mask
|
||||
mask_crop = mask.crop((minx, miny, maxx + 1, maxy + 1))
|
||||
r, g, b, a = paste_scaled.split()
|
||||
mask_np = np.array(mask_crop)
|
||||
alpha_np = np.array(a)
|
||||
combined = (alpha_np.astype(np.uint16) * mask_np.astype(np.uint16) // 255).astype(np.uint8)
|
||||
paste_final = Image.merge("RGBA", (r, g, b, Image.fromarray(combined)))
|
||||
|
||||
result = img.copy()
|
||||
result.paste(paste_final, (minx, miny), paste_final.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return {"result": _encode(out.getvalue())}
|
||||
|
||||
|
||||
# ─── SAM (Segment Anything) ──────────────────────────────────────────────────
|
||||
|
||||
class SegmentPointRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
points: list[list[int]] # [[x, y], ...] original image coords
|
||||
labels: list[int] # 1=include, 0=exclude — same length as points
|
||||
|
||||
|
||||
@router.post("/segment/point")
|
||||
async def segment_point(req: SegmentPointRequest):
|
||||
"""
|
||||
Run SAM point-prompt segmentation.
|
||||
Returns a binary mask PNG (white = selected area).
|
||||
Auto-downloads the SAM ViT-B model (~375 MB) on first call.
|
||||
"""
|
||||
if not req.points:
|
||||
raise HTTPException(status_code=400, detail="At least one point required.")
|
||||
if len(req.points) != len(req.labels):
|
||||
raise HTTPException(status_code=400, detail="points and labels must have the same length.")
|
||||
|
||||
try:
|
||||
image_bytes = base64.b64decode(req.image)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
from app.services.sam_service import predict_points, get_install_status
|
||||
try:
|
||||
mask_bytes = await predict_points(
|
||||
image_bytes,
|
||||
[tuple(p) for p in req.points],
|
||||
req.labels,
|
||||
)
|
||||
return {
|
||||
"mask": base64.b64encode(mask_bytes).decode(),
|
||||
"sam_install": get_install_status(),
|
||||
}
|
||||
except RuntimeError as e:
|
||||
raise HTTPException(status_code=503, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/segment/install-status")
|
||||
def segment_install_status():
|
||||
"""Poll SAM model download progress."""
|
||||
from app.services.sam_service import get_install_status, sam_model_available
|
||||
status = get_install_status()
|
||||
status["model_ready"] = sam_model_available()
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/segment/install")
|
||||
async def segment_install():
|
||||
"""Trigger SAM model download explicitly (also auto-triggered on first /segment/point call)."""
|
||||
from app.services.sam_service import ensure_sam_installed, get_install_status
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
return get_install_status()
|
||||
|
||||
|
||||
# ─── Enhance ─────────────────────────────────────────────────────────────────
|
||||
|
||||
import io as _io
|
||||
import numpy as _np
|
||||
import cv2 as _cv2
|
||||
from PIL import Image as _Image
|
||||
|
||||
class EnhanceRequest(BaseModel):
|
||||
image: str # base64
|
||||
strength: float = 1.0
|
||||
|
||||
|
||||
def _enhance_image(image_bytes: bytes, strength: float) -> bytes:
|
||||
"""
|
||||
Apply a chain of non-AI image enhancements, each blended with `strength` (0–1).
|
||||
|
||||
Steps:
|
||||
1. Auto white balance (gray-world)
|
||||
2. CLAHE on L channel of LAB colorspace
|
||||
3. Auto saturation boost in HSV (×1.15, clamped)
|
||||
4. Mild unsharp mask (gaussian sigma=1.0, delta weight=0.3)
|
||||
"""
|
||||
strength = max(0.0, min(1.0, float(strength)))
|
||||
|
||||
# Decode to RGB numpy array
|
||||
pil = _Image.open(_io.BytesIO(image_bytes)).convert("RGB")
|
||||
orig = _np.array(pil, dtype=_np.float32) # H×W×3, float [0,255]
|
||||
|
||||
img = orig.copy()
|
||||
|
||||
# ── Step 1: Auto white balance (gray-world) ──────────────────────────────
|
||||
mean_r = img[:, :, 0].mean()
|
||||
mean_g = img[:, :, 1].mean()
|
||||
mean_b = img[:, :, 2].mean()
|
||||
overall_mean = (mean_r + mean_g + mean_b) / 3.0
|
||||
|
||||
def _scale(channel, channel_mean):
|
||||
if channel_mean == 0:
|
||||
return channel
|
||||
return channel * (overall_mean / channel_mean)
|
||||
|
||||
wb = img.copy()
|
||||
wb[:, :, 0] = _np.clip(_scale(img[:, :, 0], mean_r), 0, 255)
|
||||
wb[:, :, 1] = _np.clip(_scale(img[:, :, 1], mean_g), 0, 255)
|
||||
wb[:, :, 2] = _np.clip(_scale(img[:, :, 2], mean_b), 0, 255)
|
||||
|
||||
img = (orig + strength * (wb - orig)).clip(0, 255)
|
||||
|
||||
# ── Step 2: CLAHE on L channel (LAB) ────────────────────────────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
lab = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2LAB)
|
||||
clahe = _cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
l_orig = lab[:, :, 0].copy()
|
||||
lab[:, :, 0] = clahe.apply(l_orig)
|
||||
# Blend L channel back using strength
|
||||
lab_blended = lab.copy()
|
||||
lab_blended[:, :, 0] = (l_orig + strength * (lab[:, :, 0].astype(_np.float32) - l_orig.astype(_np.float32))).clip(0, 255).astype(_np.uint8)
|
||||
img = _cv2.cvtColor(lab_blended, _cv2.COLOR_LAB2RGB).astype(_np.float32)
|
||||
|
||||
# ── Step 3: Auto saturation boost (HSV, ×1.15) ──────────────────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
hsv = _cv2.cvtColor(img_u8, _cv2.COLOR_RGB2HSV).astype(_np.float32)
|
||||
s_orig = hsv[:, :, 1].copy()
|
||||
s_boosted = _np.clip(s_orig * 1.15, 0, 255)
|
||||
hsv[:, :, 1] = s_orig + strength * (s_boosted - s_orig)
|
||||
hsv = hsv.clip(0, 255).astype(_np.uint8)
|
||||
img = _cv2.cvtColor(hsv, _cv2.COLOR_HSV2RGB).astype(_np.float32)
|
||||
|
||||
# ── Step 4: Mild unsharp mask (sigma=1.0, delta weight=0.3) ─────────────
|
||||
img_u8 = img.astype(_np.uint8)
|
||||
blurred = _cv2.GaussianBlur(img_u8, (0, 0), sigmaX=1.0)
|
||||
sharpness_delta = img_u8.astype(_np.float32) - blurred.astype(_np.float32)
|
||||
sharpened = img_u8.astype(_np.float32) + 0.3 * sharpness_delta * strength
|
||||
img = sharpened.clip(0, 255)
|
||||
|
||||
# Encode result as PNG
|
||||
result_pil = _Image.fromarray(img.astype(_np.uint8), mode="RGB")
|
||||
buf = _io.BytesIO()
|
||||
result_pil.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
@router.post("/enhance")
|
||||
async def enhance(req: EnhanceRequest):
|
||||
"""
|
||||
Non-AI image enhancement: auto white balance, CLAHE, saturation boost,
|
||||
and unsharp mask. Each step is blended proportionally to `strength` (0–1).
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _enhance_image, image_bytes, req.strength
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Subject replace ─────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractSubjectRequest(BaseModel):
|
||||
image: str # base64
|
||||
|
||||
|
||||
class ReplaceSubjectRequest(BaseModel):
|
||||
background_image: str # base64 — image whose background we keep
|
||||
subject_image: str # base64 — image whose subject we extract
|
||||
mask: Optional[str] = None # base64 — white = where the subject should land
|
||||
match_colors: bool = True # blend subject color stats toward background
|
||||
|
||||
|
||||
def _extract_subject_bytes(image_bytes: bytes) -> bytes:
|
||||
"""Remove background from image using rembg; return RGBA PNG bytes."""
|
||||
if rembg_available():
|
||||
return remove_background_rembg(image_bytes)
|
||||
raise RuntimeError(
|
||||
"rembg is not installed. Run: pip install rembg (or add it to requirements.txt)"
|
||||
)
|
||||
|
||||
|
||||
def _color_transfer_lab(subj_rgba: "Image", bg_rgb: "Image", blend: float = 0.45) -> "Image":
|
||||
"""
|
||||
Partial LAB color transfer: nudge subject color statistics 'blend' fraction
|
||||
toward the background's statistics so it looks like it belongs in the scene.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
src_arr = np.array(subj_rgba.convert("RGB"), dtype=np.float32)
|
||||
tgt_arr = np.array(bg_rgb.convert("RGB"), dtype=np.float32)
|
||||
|
||||
alpha = np.array(subj_rgba.split()[3])
|
||||
subject_mask = alpha > 10
|
||||
|
||||
if not subject_mask.any():
|
||||
return subj_rgba
|
||||
|
||||
src_lab = cv2.cvtColor(src_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
tgt_lab = cv2.cvtColor(tgt_arr.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
|
||||
|
||||
for ch in range(3):
|
||||
src_ch = src_lab[:, :, ch]
|
||||
src_pixels = src_ch[subject_mask]
|
||||
tgt_pixels = tgt_lab[:, :, ch].flatten()
|
||||
|
||||
src_mean, src_std = float(src_pixels.mean()), float(src_pixels.std()) + 1e-6
|
||||
tgt_mean, tgt_std = float(tgt_pixels.mean()), float(tgt_pixels.std()) + 1e-6
|
||||
|
||||
adjusted_std = src_std + blend * (tgt_std - src_std)
|
||||
adjusted = (src_ch - src_mean) * (adjusted_std / src_std) + src_mean + blend * (tgt_mean - src_mean)
|
||||
src_lab[:, :, ch] = np.clip(adjusted, 0, 255)
|
||||
|
||||
result_rgb = cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2RGB)
|
||||
r, g, b = result_rgb[:, :, 0], result_rgb[:, :, 1], result_rgb[:, :, 2]
|
||||
return Image.merge("RGBA", [
|
||||
Image.fromarray(r), Image.fromarray(g),
|
||||
Image.fromarray(b), Image.fromarray(alpha),
|
||||
])
|
||||
|
||||
|
||||
def _do_replace_subject(
|
||||
bg_bytes: bytes,
|
||||
subj_bytes: bytes,
|
||||
mask_bytes: Optional[bytes],
|
||||
match_colors: bool,
|
||||
) -> bytes:
|
||||
"""Core compositing: extract subject → scale → color-match → paste onto background."""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
bg_img = Image.open(BytesIO(bg_bytes)).convert("RGBA")
|
||||
|
||||
subj_rgba = Image.open(BytesIO(_extract_subject_bytes(subj_bytes))).convert("RGBA")
|
||||
|
||||
# Determine target placement bounding box from mask or full canvas
|
||||
if mask_bytes:
|
||||
mask_img = Image.open(BytesIO(mask_bytes)).convert("L")
|
||||
if mask_img.size != bg_img.size:
|
||||
mask_img = mask_img.resize(bg_img.size, Image.LANCZOS)
|
||||
mask_arr = np.array(mask_img)
|
||||
ys, xs = np.where(mask_arr > 128)
|
||||
else:
|
||||
mask_img = None
|
||||
mask_arr = None
|
||||
ys, xs = np.array([]), np.array([])
|
||||
|
||||
if len(xs) > 0:
|
||||
minx, maxx = int(xs.min()), int(xs.max())
|
||||
miny, maxy = int(ys.min()), int(ys.max())
|
||||
else:
|
||||
minx, miny = 0, 0
|
||||
maxx, maxy = bg_img.width - 1, bg_img.height - 1
|
||||
|
||||
target_w = maxx - minx + 1
|
||||
target_h = maxy - miny + 1
|
||||
|
||||
# Scale subject to fit target area, preserving aspect ratio
|
||||
sw, sh = subj_rgba.size
|
||||
scale = min(target_w / sw, target_h / sh)
|
||||
new_w = max(1, round(sw * scale))
|
||||
new_h = max(1, round(sh * scale))
|
||||
subj_scaled = subj_rgba.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
# Optional color transfer to blend lighting/tone
|
||||
if match_colors:
|
||||
subj_scaled = _color_transfer_lab(subj_scaled, bg_img.convert("RGB"))
|
||||
|
||||
# Center in target area
|
||||
px = minx + (target_w - new_w) // 2
|
||||
py = miny + (target_h - new_h) // 2
|
||||
|
||||
result = bg_img.copy()
|
||||
|
||||
if mask_img is not None and len(xs) > 0:
|
||||
# Build a full-canvas RGBA layer for the subject
|
||||
subj_canvas = Image.new("RGBA", bg_img.size, (0, 0, 0, 0))
|
||||
subj_canvas.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
# Clip subject's alpha to the selection mask
|
||||
sc_arr = np.array(subj_canvas)
|
||||
sc_arr[:, :, 3] = np.minimum(sc_arr[:, :, 3], mask_arr).astype(np.uint8)
|
||||
subj_canvas = Image.fromarray(sc_arr)
|
||||
result.paste(subj_canvas, (0, 0), subj_canvas.split()[3])
|
||||
else:
|
||||
result.paste(subj_scaled, (px, py), subj_scaled.split()[3])
|
||||
|
||||
out = BytesIO()
|
||||
result.convert("RGB").save(out, format="PNG")
|
||||
return out.getvalue()
|
||||
|
||||
|
||||
@router.post("/image/extract-subject")
|
||||
async def extract_subject(req: ExtractSubjectRequest):
|
||||
"""
|
||||
Remove background from an image and return the subject with transparency (RGBA PNG).
|
||||
Uses rembg (AI-powered) when available.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_subject_bytes, _decode(req.image)
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/image/replace-subject")
|
||||
async def replace_subject(req: ReplaceSubjectRequest):
|
||||
"""
|
||||
Extract the primary subject from `subject_image` (via rembg background removal),
|
||||
scale it to fit the `mask` selection on `background_image`, apply optional LAB
|
||||
color transfer for lighting consistency, and composite the result.
|
||||
|
||||
Returns the composited image as base64 PNG.
|
||||
"""
|
||||
try:
|
||||
result = await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
_do_replace_subject,
|
||||
_decode(req.background_image),
|
||||
_decode(req.subject_image),
|
||||
_decode(req.mask) if req.mask else None,
|
||||
req.match_colors,
|
||||
)
|
||||
return {"result": _encode(result)}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# ─── Extract colors ───────────────────────────────────────────────────────────
|
||||
|
||||
class ExtractColorsRequest(BaseModel):
|
||||
image: str # base64
|
||||
count: int = 6
|
||||
|
||||
|
||||
def _extract_colors(image_bytes: bytes, count: int) -> list[str]:
|
||||
"""
|
||||
Resize image to 150×150, k-means cluster pixels into `count` groups
|
||||
using pure numpy (no sklearn dependency), return hex strings by frequency.
|
||||
"""
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
|
||||
count = max(1, min(count, 32))
|
||||
|
||||
pil = Image.open(BytesIO(image_bytes)).convert("RGB").resize((150, 150))
|
||||
pixels = np.array(pil, dtype=np.float32).reshape(-1, 3) # (22500, 3)
|
||||
n = len(pixels)
|
||||
|
||||
# Initialise centers with k-means++ seeding
|
||||
rng = np.random.default_rng(42)
|
||||
centers = [pixels[rng.integers(n)]]
|
||||
for _ in range(count - 1):
|
||||
dists = np.min([np.sum((pixels - c) ** 2, axis=1) for c in centers], axis=0)
|
||||
probs = dists / dists.sum()
|
||||
centers.append(pixels[rng.choice(n, p=probs)])
|
||||
centers = np.array(centers)
|
||||
|
||||
labels = np.zeros(n, dtype=np.int32)
|
||||
for _ in range(20): # max 20 iterations
|
||||
# Assign each pixel to nearest center
|
||||
dists = np.sum((pixels[:, None] - centers[None]) ** 2, axis=2) # (n, k)
|
||||
new_labels = np.argmin(dists, axis=1)
|
||||
if np.all(new_labels == labels):
|
||||
break
|
||||
labels = new_labels
|
||||
# Recompute centers
|
||||
for k in range(count):
|
||||
mask = labels == k
|
||||
if mask.any():
|
||||
centers[k] = pixels[mask].mean(axis=0)
|
||||
|
||||
counts = np.bincount(labels, minlength=count)
|
||||
order = np.argsort(-counts)
|
||||
|
||||
return [
|
||||
"#{:02x}{:02x}{:02x}".format(*centers[i].astype(int).clip(0, 255))
|
||||
for i in order
|
||||
]
|
||||
|
||||
|
||||
@router.post("/extract-colors")
|
||||
async def extract_colors(req: ExtractColorsRequest):
|
||||
"""
|
||||
Extract dominant colors from an image using k-means clustering.
|
||||
Returns hex color strings sorted by frequency (most dominant first).
|
||||
"""
|
||||
try:
|
||||
image_bytes = _decode(req.image)
|
||||
colors = await asyncio.get_event_loop().run_in_executor(
|
||||
None, _extract_colors, image_bytes, req.count
|
||||
)
|
||||
return {"colors": colors}
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
+171
@@ -0,0 +1,171 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
|
||||
from sqlalchemy.orm import Session
|
||||
import json
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import EditRequest, EditResponse, StatusResponse
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/edits", tags=["edits"])
|
||||
|
||||
|
||||
async def process_edit_background(
|
||||
edit_id: int,
|
||||
project_id: int,
|
||||
request: EditRequest,
|
||||
db: Session
|
||||
):
|
||||
"""Background task to process edit"""
|
||||
edit_service = EditService()
|
||||
|
||||
try:
|
||||
# Process the edit
|
||||
result_path = await edit_service.process_edit(
|
||||
project_id=project_id,
|
||||
edit_id=edit_id,
|
||||
prompt=request.prompt,
|
||||
mode=request.mode,
|
||||
selection_type=request.selection_type,
|
||||
bbox=request.bbox,
|
||||
feather_px=request.feather_px,
|
||||
selection_data=request.selection_data
|
||||
)
|
||||
|
||||
# Update edit status
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if edit:
|
||||
edit.status = "completed"
|
||||
db.commit()
|
||||
|
||||
except Exception as e:
|
||||
# Update edit with error
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if edit:
|
||||
edit.status = "failed"
|
||||
edit.error_message = str(e)
|
||||
db.commit()
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/fix", response_model=EditResponse)
|
||||
async def create_edit(
|
||||
project_id: int,
|
||||
request: EditRequest,
|
||||
background_tasks: BackgroundTasks,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new edit request (Fix button)
|
||||
|
||||
This endpoint accepts the selection data and prompt,
|
||||
then processes the edit in the background.
|
||||
"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Validate mode
|
||||
if request.mode not in ["A", "B"]:
|
||||
raise HTTPException(status_code=400, detail="Mode must be 'A' or 'B'")
|
||||
|
||||
# Validate selection type
|
||||
if request.selection_type not in ["rectangle", "ellipse", "lasso"]:
|
||||
raise HTTPException(status_code=400, detail="Invalid selection type")
|
||||
|
||||
# Create edit record
|
||||
edit = Edit(
|
||||
project_id=project_id,
|
||||
mode=request.mode,
|
||||
prompt=request.prompt,
|
||||
selection_type=request.selection_type,
|
||||
bbox_json=json.dumps(request.bbox),
|
||||
feather_px=request.feather_px,
|
||||
ai_provider=settings.ai_provider,
|
||||
status="pending"
|
||||
)
|
||||
db.add(edit)
|
||||
db.commit()
|
||||
db.refresh(edit)
|
||||
|
||||
# Process edit in background
|
||||
background_tasks.add_task(
|
||||
process_edit_background,
|
||||
edit.id,
|
||||
project_id,
|
||||
request,
|
||||
db
|
||||
)
|
||||
|
||||
return edit
|
||||
|
||||
|
||||
@router.get("/{edit_id}", response_model=EditResponse)
|
||||
def get_edit(
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get edit details and status"""
|
||||
edit = db.query(Edit).filter(Edit.id == edit_id).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
return edit
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/revert/{edit_id}", response_model=StatusResponse)
|
||||
def revert_to_edit(
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Revert project to a specific edit"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Verify edit exists and belongs to project
|
||||
edit = db.query(Edit).filter(
|
||||
Edit.id == edit_id,
|
||||
Edit.project_id == project_id
|
||||
).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
|
||||
# Revert
|
||||
edit_service = EditService()
|
||||
try:
|
||||
result_path = edit_service.revert_to_edit(project_id, edit_id)
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Reverted to edit {edit_id}",
|
||||
data={"image_url": f"/projects/{project_id}/current"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/projects/{project_id}/reset", response_model=StatusResponse)
|
||||
def reset_to_original(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Reset project to original image"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Reset
|
||||
edit_service = EditService()
|
||||
try:
|
||||
result_path = edit_service.reset_to_original(project_id)
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message="Reset to original image",
|
||||
data={"image_url": f"/projects/{project_id}/current"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Form
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import Optional
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import os
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.schemas import TextToImageRequest, TextToImageResponse
|
||||
from app.services.ai_provider import get_ai_provider
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/generate", tags=["generate"])
|
||||
|
||||
|
||||
@router.post("/text-to-image", response_model=TextToImageResponse)
|
||||
async def text_to_image(
|
||||
prompt: str = Form(...),
|
||||
width: int = Form(1024),
|
||||
height: int = Form(1024),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
ai_provider: Optional[str] = Form(None),
|
||||
ai_model: Optional[str] = Form(None),
|
||||
create_project: bool = Form(True),
|
||||
project_name: Optional[str] = Form(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate an image from text prompt
|
||||
|
||||
Args:
|
||||
prompt: Text description of desired image
|
||||
width: Image width (default 1024)
|
||||
height: Image height (default 1024)
|
||||
negative_prompt: What to avoid in generation
|
||||
ai_provider: Override default AI provider
|
||||
ai_model: Specific model to use
|
||||
create_project: Whether to create a new project with the result
|
||||
project_name: Name for the new project (if create_project=True)
|
||||
|
||||
Returns:
|
||||
Generated image info and optionally project details
|
||||
"""
|
||||
|
||||
# Validate dimensions
|
||||
if width < 256 or width > 2048 or height < 256 or height > 2048:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Width and height must be between 256 and 2048"
|
||||
)
|
||||
|
||||
# Get AI provider
|
||||
provider = get_ai_provider(ai_provider, ai_model)
|
||||
|
||||
try:
|
||||
# Generate image
|
||||
image_bytes = await provider.text_to_image(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
model=ai_model,
|
||||
negative_prompt=negative_prompt
|
||||
)
|
||||
|
||||
project_id = None
|
||||
image_url = None
|
||||
|
||||
if create_project:
|
||||
# Create a new project
|
||||
project = Project(
|
||||
name=project_name or f"Generated: {prompt[:50]}",
|
||||
user_id=None # TODO: Add authentication
|
||||
)
|
||||
db.add(project)
|
||||
db.commit()
|
||||
db.refresh(project)
|
||||
|
||||
project_id = project.id
|
||||
|
||||
# Save image as both original and current
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(project_id)
|
||||
|
||||
original_path = edit_service.get_original_image_path(project_id)
|
||||
current_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Save image
|
||||
img = Image.open(BytesIO(image_bytes))
|
||||
img.save(original_path, 'PNG')
|
||||
img.save(current_path, 'PNG')
|
||||
|
||||
image_url = f"/projects/{project_id}/current"
|
||||
|
||||
return TextToImageResponse(
|
||||
status="success",
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
project_id=project_id,
|
||||
image_url=image_url,
|
||||
ai_provider=ai_provider or settings.ai_provider,
|
||||
ai_model=ai_model
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/layer/text-to-image", response_model=TextToImageResponse)
|
||||
async def text_to_image_layer(
|
||||
project_id: int = Form(...),
|
||||
prompt: str = Form(...),
|
||||
width: int = Form(512),
|
||||
height: int = Form(512),
|
||||
x: int = Form(0),
|
||||
y: int = Form(0),
|
||||
negative_prompt: Optional[str] = Form(None),
|
||||
ai_provider: Optional[str] = Form(None),
|
||||
ai_model: Optional[str] = Form(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate an image as a new layer in an existing project
|
||||
|
||||
This generates a smaller image that can be placed as a layer
|
||||
on top of the current project canvas.
|
||||
"""
|
||||
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Get AI provider
|
||||
provider = get_ai_provider(ai_provider, ai_model)
|
||||
|
||||
try:
|
||||
# Generate image
|
||||
image_bytes = await provider.text_to_image(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
model=ai_model,
|
||||
negative_prompt=negative_prompt
|
||||
)
|
||||
|
||||
# Save as temporary layer file
|
||||
edit_service = EditService()
|
||||
layers_dir = edit_service.get_project_dir(project_id) / "layers"
|
||||
layers_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Generate unique layer filename
|
||||
import time
|
||||
layer_filename = f"generated_{int(time.time())}.png"
|
||||
layer_path = layers_dir / layer_filename
|
||||
|
||||
# Save layer image
|
||||
with open(layer_path, 'wb') as f:
|
||||
f.write(image_bytes)
|
||||
|
||||
return TextToImageResponse(
|
||||
status="success",
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
project_id=project_id,
|
||||
image_url=f"/projects/{project_id}/layers/{layer_filename}",
|
||||
layer_position={"x": x, "y": y},
|
||||
ai_provider=ai_provider or settings.ai_provider,
|
||||
ai_model=ai_model
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
GPU status and model management endpoints.
|
||||
All under /api/gpu prefix.
|
||||
"""
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List
|
||||
import asyncio
|
||||
|
||||
router = APIRouter(prefix="/api/gpu", tags=["gpu"])
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
async def gpu_status():
|
||||
"""
|
||||
Full GPU capability report: hardware, feature flags, VRAM budget,
|
||||
and which model was selected for each operation.
|
||||
Frontend polls this to show GPU badge and tool availability.
|
||||
"""
|
||||
from app.services.gpu_detect import get_cached_gpu_info
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
|
||||
info = get_cached_gpu_info()
|
||||
|
||||
return {
|
||||
# Hardware
|
||||
"backend": info.backend,
|
||||
"device_name": info.device_name,
|
||||
"vram_total_gb": info.vram_total_gb,
|
||||
"vram_free_gb": info.vram_free_gb,
|
||||
"compute_capability": info.compute_capability,
|
||||
# Feature flags
|
||||
"fp16": info.fp16,
|
||||
"bf16": info.bf16,
|
||||
"fp8": info.fp8,
|
||||
"int8": info.int8,
|
||||
"tensor_cores": info.tensor_cores,
|
||||
"xformers": info.xformers,
|
||||
# Derived
|
||||
"effective_vram_gb": info.effective_vram_gb,
|
||||
"tier": info.tier,
|
||||
# Selected models per operation
|
||||
"recommended": {
|
||||
op: (
|
||||
{
|
||||
"model_id": spec.model_id,
|
||||
"family": spec.family,
|
||||
"memory_opt": spec.memory_opt,
|
||||
"native_res": spec.native_res,
|
||||
"vram_fp16_gb": spec.vram_fp16_gb,
|
||||
}
|
||||
if spec else None
|
||||
)
|
||||
for op, spec in info.recommended.items()
|
||||
},
|
||||
"pipeline_states": get_all_model_states(),
|
||||
"warnings": info.warnings,
|
||||
"capabilities": info.capabilities,
|
||||
}
|
||||
|
||||
|
||||
class PrefetchRequest(BaseModel):
|
||||
operations: Optional[List[str]] = None
|
||||
|
||||
|
||||
@router.post("/prefetch")
|
||||
async def prefetch_models(req: PrefetchRequest = PrefetchRequest()):
|
||||
"""
|
||||
Eagerly load pipelines into GPU memory for the requested operations.
|
||||
Returns immediately; poll /api/gpu/prefetch-status for progress.
|
||||
Default: inpaint, txt2img, img2img.
|
||||
"""
|
||||
ops = req.operations or ["inpaint", "txt2img", "img2img"]
|
||||
valid = {"inpaint", "txt2img", "img2img", "outpaint", "upscale"}
|
||||
ops = [op for op in ops if op in valid]
|
||||
|
||||
from app.services.local_diffusion import get_local_diffusion_provider
|
||||
provider = get_local_diffusion_provider()
|
||||
|
||||
async def _prefetch():
|
||||
for op in ops:
|
||||
try:
|
||||
await provider._get_pipeline(op)
|
||||
print(f"[gpu] Prefetch complete: {op}")
|
||||
except Exception as exc:
|
||||
print(f"[gpu] Prefetch failed for {op}: {exc}")
|
||||
|
||||
asyncio.create_task(_prefetch())
|
||||
return {"status": "prefetch_started", "operations": ops}
|
||||
|
||||
|
||||
@router.get("/prefetch-status")
|
||||
async def prefetch_status():
|
||||
"""Poll model download / load progress."""
|
||||
from app.services.local_diffusion import get_all_model_states
|
||||
return {"models": get_all_model_states()}
|
||||
+81
@@ -0,0 +1,81 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends
|
||||
from fastapi.responses import FileResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from pathlib import Path
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.services.edit_service import EditService
|
||||
|
||||
router = APIRouter(prefix="/projects", tags=["images"])
|
||||
|
||||
|
||||
@router.get("/{project_id}/original")
|
||||
def get_original_image(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the original uploaded image"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
image_path = edit_service.get_original_image_path(project_id)
|
||||
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Original image not found")
|
||||
|
||||
return FileResponse(
|
||||
image_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "public, max-age=3600"}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/current")
|
||||
def get_current_image(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the current edited image"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
image_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
if not image_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Current image not found")
|
||||
|
||||
return FileResponse(
|
||||
image_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "no-cache"}
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/history/{edit_id}/result")
|
||||
def get_edit_result(
|
||||
project_id: int,
|
||||
edit_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get the result image from a specific edit"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edit_service = EditService()
|
||||
edit_dir = edit_service.get_edit_dir(project_id, edit_id)
|
||||
result_path = edit_dir / "result.png"
|
||||
|
||||
if not result_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Edit result not found")
|
||||
|
||||
return FileResponse(
|
||||
result_path,
|
||||
media_type="image/png",
|
||||
headers={"Cache-Control": "public, max-age=3600"}
|
||||
)
|
||||
+308
@@ -0,0 +1,308 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form
|
||||
from fastapi.responses import FileResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List, Optional
|
||||
import json
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.patch import Patch
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import PatchCreate, PatchResponse, PatchApply, StatusResponse
|
||||
from app.services.patch_library import PatchLibraryService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/patches", tags=["patches"])
|
||||
|
||||
|
||||
@router.post("/", response_model=PatchResponse)
|
||||
async def create_patch(
|
||||
name: str = Form(...),
|
||||
description: Optional[str] = Form(None),
|
||||
source_type: str = Form(...),
|
||||
category: Optional[str] = Form(None),
|
||||
tags: Optional[str] = Form(None),
|
||||
source_project_id: Optional[int] = Form(None),
|
||||
source_edit_id: Optional[int] = Form(None),
|
||||
bbox: Optional[str] = Form(None),
|
||||
file: Optional[UploadFile] = File(None),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new patch in the library
|
||||
|
||||
Source types:
|
||||
- ai_generated: From an edit (requires source_edit_id)
|
||||
- manual_selection: Selected from current image (requires source_project_id and bbox)
|
||||
- imported: Uploaded file (requires file)
|
||||
"""
|
||||
|
||||
# Validate source_type
|
||||
if source_type not in ["ai_generated", "manual_selection", "imported"]:
|
||||
raise HTTPException(status_code=400, detail="Invalid source_type")
|
||||
|
||||
# Create patch record
|
||||
patch = Patch(
|
||||
name=name,
|
||||
description=description,
|
||||
source_type=source_type,
|
||||
source_project_id=source_project_id,
|
||||
source_edit_id=source_edit_id,
|
||||
tags=tags,
|
||||
category=category,
|
||||
file_path="", # Will be set after saving
|
||||
user_id=None # TODO: Add authentication
|
||||
)
|
||||
|
||||
db.add(patch)
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
# Save patch file based on source type
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
try:
|
||||
if source_type == "ai_generated":
|
||||
# Get edit directory and save AI-generated patch
|
||||
if not source_edit_id:
|
||||
raise HTTPException(status_code=400, detail="source_edit_id required for ai_generated")
|
||||
|
||||
edit = db.query(Edit).filter(Edit.id == source_edit_id).first()
|
||||
if not edit:
|
||||
raise HTTPException(status_code=404, detail="Edit not found")
|
||||
|
||||
from app.services.edit_service import EditService
|
||||
edit_service = EditService()
|
||||
edit_dir = edit_service.get_edit_dir(edit.project_id, edit.id)
|
||||
|
||||
file_path = patch_service.save_ai_generated_patch(patch.id, edit_dir)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
|
||||
elif source_type == "manual_selection":
|
||||
# Save manually selected patch from current image
|
||||
if not source_project_id or not bbox:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="source_project_id and bbox required for manual_selection"
|
||||
)
|
||||
|
||||
project = db.query(Project).filter(Project.id == source_project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
bbox_dict = json.loads(bbox) if isinstance(bbox, str) else bbox
|
||||
file_path = patch_service.save_manual_patch(patch.id, source_project_id, bbox_dict)
|
||||
|
||||
patch.width = bbox_dict['width']
|
||||
patch.height = bbox_dict['height']
|
||||
|
||||
elif source_type == "imported":
|
||||
# Save uploaded file
|
||||
if not file:
|
||||
raise HTTPException(status_code=400, detail="file required for imported")
|
||||
|
||||
image_bytes = await file.read()
|
||||
file_path = patch_service.save_patch_from_bytes(patch.id, image_bytes)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
|
||||
# Update patch with file path
|
||||
patch.file_path = file_path
|
||||
patch.thumbnail_path = str(patch_service.get_thumbnail_path(patch.id))
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
return patch
|
||||
|
||||
except Exception as e:
|
||||
# Cleanup on error
|
||||
patch_service.delete_patch(patch.id)
|
||||
db.delete(patch)
|
||||
db.commit()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/", response_model=List[PatchResponse])
|
||||
def list_patches(
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[str] = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List patches in the library with optional filtering"""
|
||||
|
||||
query = db.query(Patch)
|
||||
|
||||
if category:
|
||||
query = query.filter(Patch.category == category)
|
||||
|
||||
if tags:
|
||||
# Simple tag search (could be improved with full-text search)
|
||||
query = query.filter(Patch.tags.like(f"%{tags}%"))
|
||||
|
||||
patches = query.offset(offset).limit(limit).all()
|
||||
return patches
|
||||
|
||||
|
||||
@router.get("/{patch_id}", response_model=PatchResponse)
|
||||
def get_patch(
|
||||
patch_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get patch details"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
return patch
|
||||
|
||||
|
||||
@router.get("/{patch_id}/image")
|
||||
def get_patch_image(
|
||||
patch_id: int,
|
||||
thumbnail: bool = False,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get patch image file"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
if thumbnail:
|
||||
file_path = patch_service.get_thumbnail_path(patch_id)
|
||||
else:
|
||||
file_path = patch_service.get_patch_path(patch_id)
|
||||
|
||||
if not file_path.exists():
|
||||
raise HTTPException(status_code=404, detail="Patch image not found")
|
||||
|
||||
return FileResponse(file_path, media_type="image/png")
|
||||
|
||||
|
||||
@router.post("/apply", response_model=StatusResponse)
|
||||
async def apply_patch(
|
||||
project_id: int = Form(...),
|
||||
patch_id: int = Form(...),
|
||||
bbox: str = Form(...),
|
||||
feather_px: int = Form(5),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Apply a saved patch to a project image
|
||||
|
||||
This creates a new edit in the project history.
|
||||
"""
|
||||
# Verify project exists
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Verify patch exists
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
# Parse bbox
|
||||
bbox_dict = json.loads(bbox) if isinstance(bbox, str) else bbox
|
||||
|
||||
# Load current image
|
||||
from app.services.edit_service import EditService
|
||||
from PIL import Image
|
||||
|
||||
edit_service = EditService()
|
||||
current_image_path = edit_service.get_current_image_path(project_id)
|
||||
current_image = Image.open(current_image_path).convert('RGBA')
|
||||
|
||||
# Apply patch
|
||||
patch_service = PatchLibraryService()
|
||||
result_image = patch_service.apply_patch_to_image(
|
||||
patch_id,
|
||||
current_image,
|
||||
bbox_dict,
|
||||
feather_px
|
||||
)
|
||||
|
||||
# Save result as current image
|
||||
result_image.save(current_image_path)
|
||||
|
||||
# Create edit record
|
||||
edit = Edit(
|
||||
project_id=project_id,
|
||||
mode="patch_library",
|
||||
prompt=f"Applied saved patch: {patch.name}",
|
||||
selection_type="rectangle",
|
||||
bbox_json=json.dumps(bbox_dict),
|
||||
feather_px=feather_px,
|
||||
ai_provider="patch_library",
|
||||
status="completed"
|
||||
)
|
||||
db.add(edit)
|
||||
db.commit()
|
||||
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Applied patch '{patch.name}' to project",
|
||||
data={"edit_id": edit.id}
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/{patch_id}", response_model=StatusResponse)
|
||||
def delete_patch(
|
||||
patch_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Delete a patch from the library"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
# Delete files
|
||||
patch_service = PatchLibraryService()
|
||||
patch_service.delete_patch(patch_id)
|
||||
|
||||
# Delete record
|
||||
db.delete(patch)
|
||||
db.commit()
|
||||
|
||||
return StatusResponse(
|
||||
status="success",
|
||||
message=f"Deleted patch '{patch.name}'"
|
||||
)
|
||||
|
||||
|
||||
@router.put("/{patch_id}", response_model=PatchResponse)
|
||||
def update_patch(
|
||||
patch_id: int,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[str] = None,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Update patch metadata"""
|
||||
patch = db.query(Patch).filter(Patch.id == patch_id).first()
|
||||
if not patch:
|
||||
raise HTTPException(status_code=404, detail="Patch not found")
|
||||
|
||||
if name:
|
||||
patch.name = name
|
||||
if description is not None:
|
||||
patch.description = description
|
||||
if category:
|
||||
patch.category = category
|
||||
if tags is not None:
|
||||
patch.tags = tags
|
||||
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
return patch
|
||||
+487
@@ -0,0 +1,487 @@
|
||||
"""
|
||||
Print / frame tools — frame fit and upscale.
|
||||
All endpoints under /api/print prefix.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Literal
|
||||
import base64
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
router = APIRouter(prefix="/api/print", tags=["print-tools"])
|
||||
|
||||
# ── Frame size catalogue (inches) ──────────────────────────────────────────
|
||||
FRAME_SIZES = {
|
||||
"4x6": (4, 6),
|
||||
"5x7": (5, 7),
|
||||
"8x10": (8, 10),
|
||||
"11x14": (11, 14),
|
||||
"16x20": (16, 20),
|
||||
"18x24": (18, 24),
|
||||
"20x24": (20, 24),
|
||||
"24x36": (24, 36),
|
||||
# Square
|
||||
"4x4": (4, 4),
|
||||
"8x8": (8, 8),
|
||||
"12x12": (12, 12),
|
||||
}
|
||||
|
||||
|
||||
def _encode(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _decode(b64: str) -> bytes:
|
||||
return base64.b64decode(b64)
|
||||
|
||||
|
||||
def _to_png(img: Image.Image) -> bytes:
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
# ── Request models ─────────────────────────────────────────────────────────
|
||||
|
||||
class FrameFitRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
frame: str # e.g. "8x10"
|
||||
orientation: Literal["auto", "portrait", "landscape"] = "auto"
|
||||
mode: Literal["crop", "extend", "smart"] = "smart"
|
||||
dpi: int = 300
|
||||
# For extend mode: prompt passed to outpaint
|
||||
prompt: Optional[str] = ""
|
||||
# Smart mode threshold: extend if gap fraction < this, else crop
|
||||
smart_threshold: float = 0.15
|
||||
|
||||
|
||||
class UpscaleRequest(BaseModel):
|
||||
image: str # base64
|
||||
scale: float = 2.0 # 1.5, 2, 3, 4
|
||||
# auto = pick best available; lanczos = always works; realesrgan_pytorch / realesrgan_ncnn = explicit
|
||||
method: str = "auto"
|
||||
|
||||
|
||||
class PrepareRequest(BaseModel):
|
||||
image: str # base64
|
||||
frame: str # e.g. "8x10"
|
||||
orientation: Literal["auto", "portrait", "landscape"] = "auto"
|
||||
target_dpi: int = 300
|
||||
upscale_method: str = "auto" # auto / realesrgan_pytorch / realesrgan_ncnn / lanczos
|
||||
mode: Literal["crop", "extend", "smart"] = "smart"
|
||||
prompt: Optional[str] = ""
|
||||
|
||||
|
||||
# ── Frame sizes endpoint ───────────────────────────────────────────────────
|
||||
|
||||
@router.get("/frame-sizes")
|
||||
def list_frame_sizes():
|
||||
"""Return the catalogue of supported frame sizes."""
|
||||
return {
|
||||
"sizes": list(FRAME_SIZES.keys()),
|
||||
"catalogue": {k: {"inches": v, "pixels_300dpi": (v[0]*300, v[1]*300)}
|
||||
for k, v in FRAME_SIZES.items()},
|
||||
}
|
||||
|
||||
|
||||
# ── Frame fit ──────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/frame-fit")
|
||||
async def frame_fit(req: FrameFitRequest):
|
||||
"""
|
||||
Fit an image to a print frame size.
|
||||
|
||||
Modes:
|
||||
crop — center-crop to frame aspect ratio, then scale to print resolution.
|
||||
extend — scale to fill one dimension, outpaint the gap with AI.
|
||||
smart — extend if gap < smart_threshold of frame dimension, else crop.
|
||||
|
||||
Returns the fitted image plus a summary of what was done.
|
||||
"""
|
||||
if req.frame not in FRAME_SIZES:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}")
|
||||
if not (72 <= req.dpi <= 600):
|
||||
raise HTTPException(status_code=400, detail="dpi must be 72–600")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
fw, fh = FRAME_SIZES[req.frame] # frame inches (w, h in portrait)
|
||||
|
||||
# Resolve orientation
|
||||
img_w, img_h = image.size
|
||||
img_landscape = img_w >= img_h
|
||||
frame_landscape = fw >= fh
|
||||
|
||||
if req.orientation == "landscape":
|
||||
fw, fh = max(fw, fh), min(fw, fh)
|
||||
elif req.orientation == "portrait":
|
||||
fw, fh = min(fw, fh), max(fw, fh)
|
||||
else: # auto — match image orientation
|
||||
if img_landscape and not frame_landscape:
|
||||
fw, fh = fh, fw # rotate frame to landscape
|
||||
elif not img_landscape and frame_landscape:
|
||||
fw, fh = fh, fw # rotate frame to portrait
|
||||
|
||||
target_w = fw * req.dpi
|
||||
target_h = fh * req.dpi
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
# Determine actual mode
|
||||
mode = req.mode
|
||||
if mode == "smart":
|
||||
# Scale image to fill the frame — compute gap fraction
|
||||
if img_ratio > target_ratio:
|
||||
# Image wider → fits on height, gap on width
|
||||
scaled_h = target_h
|
||||
scaled_w = round(target_h * img_ratio)
|
||||
gap_frac = (scaled_w - target_w) / target_w # positive = overflow (crop)
|
||||
else:
|
||||
scaled_w = target_w
|
||||
scaled_h = round(target_w / img_ratio)
|
||||
gap_frac = (scaled_h - target_h) / target_h
|
||||
|
||||
# gap_frac > 0 means we'd need to crop; < 0 means we'd need to extend
|
||||
if gap_frac < 0:
|
||||
# Need to extend — use extend if gap is small enough
|
||||
mode = "extend" if abs(gap_frac) <= req.smart_threshold else "crop"
|
||||
else:
|
||||
mode = "crop"
|
||||
|
||||
if mode == "crop":
|
||||
result, summary = _crop_fit(image, target_w, target_h)
|
||||
else: # extend
|
||||
result, summary = await _extend_fit(image, target_w, target_h, req.prompt or "")
|
||||
|
||||
return {
|
||||
"result": _encode(_to_png(result)),
|
||||
"mode_used": mode,
|
||||
"frame": req.frame,
|
||||
"orientation": "landscape" if fw > fh else "portrait",
|
||||
"output_pixels": {"width": result.width, "height": result.height},
|
||||
"output_inches": {"width": fw, "height": fh},
|
||||
"dpi": req.dpi,
|
||||
"summary": summary,
|
||||
}
|
||||
|
||||
|
||||
def _crop_fit(image: Image.Image, target_w: int, target_h: int):
|
||||
"""Center-crop image to target aspect ratio, then Lanczos scale to target size."""
|
||||
img_w, img_h = image.size
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
if img_ratio > target_ratio:
|
||||
# Wider than target — crop sides
|
||||
new_w = round(img_h * target_ratio)
|
||||
x0 = (img_w - new_w) // 2
|
||||
cropped = image.crop((x0, 0, x0 + new_w, img_h))
|
||||
else:
|
||||
# Taller than target — crop top/bottom
|
||||
new_h = round(img_w / target_ratio)
|
||||
y0 = (img_h - new_h) // 2
|
||||
cropped = image.crop((0, y0, img_w, y0 + new_h))
|
||||
|
||||
result = cropped.resize((target_w, target_h), Image.Resampling.LANCZOS)
|
||||
summary = (
|
||||
f"Cropped from {img_w}×{img_h} to {cropped.width}×{cropped.height}, "
|
||||
f"scaled to {target_w}×{target_h}"
|
||||
)
|
||||
return result, summary
|
||||
|
||||
|
||||
async def _extend_fit(image: Image.Image, target_w: int, target_h: int, prompt: str):
|
||||
"""
|
||||
Scale image to fill one dimension exactly, then outpaint the gap with AI.
|
||||
Falls back to content-aware mirror fill if no remote provider configured.
|
||||
"""
|
||||
from app.services.remote_provider import get_remote_provider
|
||||
|
||||
img_w, img_h = image.size
|
||||
target_ratio = target_w / target_h
|
||||
img_ratio = img_w / img_h
|
||||
|
||||
if img_ratio > target_ratio:
|
||||
# Image wider — scale to target width, extend height
|
||||
scale = target_w / img_w
|
||||
scaled_w = target_w
|
||||
scaled_h = round(img_h * scale)
|
||||
gap_dir = "height"
|
||||
gap_top = (target_h - scaled_h) // 2
|
||||
gap_bottom = target_h - scaled_h - gap_top
|
||||
else:
|
||||
# Image taller — scale to target height, extend width
|
||||
scale = target_h / img_h
|
||||
scaled_h = target_h
|
||||
scaled_w = round(img_w * scale)
|
||||
gap_dir = "width"
|
||||
gap_left = (target_w - scaled_w) // 2
|
||||
gap_right = target_w - scaled_w - gap_left
|
||||
|
||||
scaled = image.resize((scaled_w, scaled_h), Image.Resampling.LANCZOS)
|
||||
|
||||
# Place scaled image on canvas
|
||||
canvas = Image.new("RGB", (target_w, target_h), (128, 128, 128))
|
||||
if gap_dir == "height":
|
||||
canvas.paste(scaled, (0, gap_top))
|
||||
# Build mask: top and bottom strips are white (to inpaint)
|
||||
mask = Image.new("L", (target_w, target_h), 0)
|
||||
if gap_top > 0:
|
||||
mask.paste(Image.new("L", (target_w, gap_top), 255), (0, 0))
|
||||
if gap_bottom > 0:
|
||||
mask.paste(Image.new("L", (target_w, gap_bottom), 255), (0, target_h - gap_bottom))
|
||||
else:
|
||||
canvas.paste(scaled, (gap_left, 0))
|
||||
mask = Image.new("L", (target_w, target_h), 0)
|
||||
if gap_left > 0:
|
||||
mask.paste(Image.new("L", (gap_left, target_h), 255), (0, 0))
|
||||
if gap_right > 0:
|
||||
mask.paste(Image.new("L", (gap_right, target_h), 255), (target_w - gap_right, 0))
|
||||
|
||||
# Try AI inpaint
|
||||
provider = get_remote_provider("inpaint")
|
||||
if provider:
|
||||
try:
|
||||
canvas_bytes = _to_png(canvas)
|
||||
mask_bytes = _to_png(mask)
|
||||
fill_prompt = prompt or "seamlessly continue the image, natural extension"
|
||||
result_bytes = await provider.inpaint(canvas_bytes, mask_bytes, fill_prompt, {})
|
||||
result = Image.open(BytesIO(result_bytes)).convert("RGB")
|
||||
summary = (
|
||||
f"Scaled {img_w}×{img_h} → {scaled_w}×{scaled_h}, "
|
||||
f"AI-extended {gap_dir} to {target_w}×{target_h}"
|
||||
)
|
||||
return result, summary
|
||||
except Exception as e:
|
||||
print(f"AI extend failed, using mirror fill: {e}")
|
||||
|
||||
# Fallback: mirror-fill the gap (looks decent for backgrounds/landscapes)
|
||||
result = _mirror_fill(canvas, mask, scaled, gap_dir,
|
||||
gap_top if gap_dir == "height" else gap_left,
|
||||
gap_bottom if gap_dir == "height" else gap_right,
|
||||
target_w, target_h)
|
||||
summary = (
|
||||
f"Scaled {img_w}×{img_h} → {scaled_w}×{scaled_h}, "
|
||||
f"mirror-filled {gap_dir} to {target_w}×{target_h} (no AI provider)"
|
||||
)
|
||||
return result, summary
|
||||
|
||||
|
||||
def _mirror_fill(canvas, mask, scaled, gap_dir, gap_a, gap_b, target_w, target_h):
|
||||
"""Fill gaps by reflecting the nearest edge strip."""
|
||||
result = canvas.copy()
|
||||
if gap_dir == "height":
|
||||
if gap_a > 0:
|
||||
strip = scaled.crop((0, 0, scaled.width, min(gap_a * 2, scaled.height)))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
strip = strip.resize((target_w, gap_a), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, 0))
|
||||
if gap_b > 0:
|
||||
strip = scaled.crop((0, max(0, scaled.height - gap_b * 2), scaled.width, scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
strip = strip.resize((target_w, gap_b), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, target_h - gap_b))
|
||||
else:
|
||||
if gap_a > 0:
|
||||
strip = scaled.crop((0, 0, min(gap_a * 2, scaled.width), scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
strip = strip.resize((gap_a, target_h), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (0, 0))
|
||||
if gap_b > 0:
|
||||
strip = scaled.crop((max(0, scaled.width - gap_b * 2), 0, scaled.width, scaled.height))
|
||||
strip = strip.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
strip = strip.resize((gap_b, target_h), Image.Resampling.LANCZOS)
|
||||
result.paste(strip, (target_w - gap_b, 0))
|
||||
return result
|
||||
|
||||
|
||||
# ── Upscale ────────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/upscale/refresh-caps")
|
||||
def upscale_refresh_caps():
|
||||
"""Bust the capability cache (call after installing Real-ESRGAN without restarting)."""
|
||||
from app.services.upscale import invalidate_caps_cache, probe_upscale_capabilities
|
||||
invalidate_caps_cache()
|
||||
return probe_upscale_capabilities()
|
||||
|
||||
|
||||
@router.get("/upscale/available")
|
||||
async def upscale_available():
|
||||
"""
|
||||
Return capability probe: which upscale methods are available,
|
||||
which device will be used, and which method is recommended.
|
||||
If no AI upscaler is found, triggers background NCNN auto-install.
|
||||
Frontend uses this to populate the method selector.
|
||||
"""
|
||||
from app.services.upscale import probe_upscale_capabilities, ensure_ncnn_installed, get_install_status
|
||||
caps = probe_upscale_capabilities()
|
||||
# Auto-install NCNN if no AI upscaler is available yet
|
||||
if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]:
|
||||
asyncio.create_task(ensure_ncnn_installed())
|
||||
caps["ncnn_install_status"] = get_install_status()
|
||||
return caps
|
||||
|
||||
|
||||
@router.get("/upscale/install-status")
|
||||
def upscale_install_status():
|
||||
"""Poll for Real-ESRGAN NCNN auto-install progress."""
|
||||
from app.services.upscale import get_install_status, probe_upscale_capabilities, _find_ncnn_binary
|
||||
status = get_install_status()
|
||||
# If install just finished, refresh caps
|
||||
if status["state"] == "done":
|
||||
from app.services.upscale import invalidate_caps_cache
|
||||
invalidate_caps_cache()
|
||||
caps = probe_upscale_capabilities()
|
||||
status["ncnn_available"] = caps["realesrgan_ncnn"]
|
||||
else:
|
||||
status["ncnn_available"] = False
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/prepare")
|
||||
async def prepare_for_print(req: PrepareRequest):
|
||||
"""
|
||||
One-shot Prepare for Print: AI upscale to reach target DPI, then fit to frame.
|
||||
|
||||
Steps:
|
||||
1. Resolve target pixel dimensions (frame × target_dpi, orientation-adjusted)
|
||||
2. Calculate needed upscale factor so the image meets the target resolution
|
||||
3. Run Real-ESRGAN if scale > 1.05 (else skip — already large enough)
|
||||
4. Run frame-fit (crop / extend / smart) to exact target dimensions
|
||||
5. Return the print-ready image and a quality report
|
||||
"""
|
||||
if req.frame not in FRAME_SIZES:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"Unknown frame '{req.frame}'. Valid: {list(FRAME_SIZES.keys())}")
|
||||
if not (72 <= req.target_dpi <= 600):
|
||||
raise HTTPException(status_code=400, detail="target_dpi must be 72–600")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
fw, fh = FRAME_SIZES[req.frame]
|
||||
img_w, img_h = image.size
|
||||
|
||||
# Resolve orientation (same logic as frame_fit)
|
||||
img_landscape = img_w >= img_h
|
||||
frame_landscape = fw >= fh
|
||||
if req.orientation == "landscape":
|
||||
fw, fh = max(fw, fh), min(fw, fh)
|
||||
elif req.orientation == "portrait":
|
||||
fw, fh = min(fw, fh), max(fw, fh)
|
||||
else:
|
||||
if img_landscape and not frame_landscape:
|
||||
fw, fh = fh, fw
|
||||
elif not img_landscape and frame_landscape:
|
||||
fw, fh = fh, fw
|
||||
|
||||
target_w = fw * req.target_dpi
|
||||
target_h = fh * req.target_dpi
|
||||
|
||||
# Scale factor needed so the shorter dimension fills the frame
|
||||
scale_w = target_w / img_w
|
||||
scale_h = target_h / img_h
|
||||
needed_scale = min(scale_w, scale_h) # fill-to-fit (extend) baseline
|
||||
# For crop mode we need max; use the larger to be safe and let frame-fit crop
|
||||
needed_scale_crop = max(scale_w, scale_h)
|
||||
|
||||
# Use the smaller (extend) scale as the upscale target; frame-fit handles the rest
|
||||
upscale_factor = max(1.0, needed_scale)
|
||||
upscale_applied = False
|
||||
method_used = "none"
|
||||
|
||||
upscaled = image
|
||||
if upscale_factor > 1.05:
|
||||
# Cap per-pass at 4× (Real-ESRGAN works best at 2–4×)
|
||||
remaining = upscale_factor
|
||||
while remaining > 1.05:
|
||||
pass_scale = min(remaining, 4.0)
|
||||
# Round to one decimal to keep scale in 1.1–8.0 range accepted by upscale service
|
||||
pass_scale = round(pass_scale, 1)
|
||||
if pass_scale < 1.1:
|
||||
break
|
||||
from app.services.upscale import upscale_image
|
||||
result_bytes, method_used = await upscale_image(upscaled, pass_scale, req.upscale_method)
|
||||
upscaled = Image.open(BytesIO(result_bytes)).convert("RGB")
|
||||
remaining /= pass_scale
|
||||
upscale_applied = True
|
||||
|
||||
# Encode upscaled image and run frame-fit
|
||||
upscaled_b64 = _encode(_to_png(upscaled))
|
||||
|
||||
fit_req = FrameFitRequest(
|
||||
image=upscaled_b64,
|
||||
frame=req.frame,
|
||||
orientation=req.orientation,
|
||||
mode=req.mode,
|
||||
dpi=req.target_dpi,
|
||||
prompt=req.prompt or "",
|
||||
)
|
||||
# Re-use the existing frame_fit logic inline
|
||||
fit_response = await frame_fit(fit_req)
|
||||
|
||||
return {
|
||||
"result": fit_response["result"],
|
||||
"frame": req.frame,
|
||||
"orientation": fit_response["orientation"],
|
||||
"output_pixels": fit_response["output_pixels"],
|
||||
"output_inches": fit_response["output_inches"],
|
||||
"dpi": req.target_dpi,
|
||||
"mode_used": fit_response["mode_used"],
|
||||
"upscale_applied": upscale_applied,
|
||||
"upscale_factor": round(upscale_factor, 2),
|
||||
"upscale_method": method_used,
|
||||
"summary": fit_response["summary"],
|
||||
}
|
||||
|
||||
|
||||
@router.post("/upscale")
|
||||
async def upscale(req: UpscaleRequest):
|
||||
"""
|
||||
Upscale image. method values:
|
||||
auto — pick best available (recommended)
|
||||
realesrgan_pytorch — Real-ESRGAN via PyTorch (CUDA/MPS/CPU)
|
||||
realesrgan_ncnn — Real-ESRGAN NCNN Vulkan binary
|
||||
lanczos — always available, instant
|
||||
Any AI method falls back to the next best if unavailable.
|
||||
"""
|
||||
if not (1.1 <= req.scale <= 8.0):
|
||||
raise HTTPException(status_code=400, detail="scale must be 1.1–8.0")
|
||||
|
||||
valid_methods = {"auto", "realesrgan_pytorch", "realesrgan_ncnn", "lanczos"}
|
||||
if req.method not in valid_methods:
|
||||
raise HTTPException(status_code=400,
|
||||
detail=f"method must be one of {sorted(valid_methods)}")
|
||||
|
||||
try:
|
||||
image = Image.open(BytesIO(_decode(req.image))).convert("RGB")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
orig_w, orig_h = image.size
|
||||
|
||||
try:
|
||||
from app.services.upscale import upscale_image
|
||||
result_bytes, method_used = await upscale_image(image, req.scale, req.method)
|
||||
result = Image.open(BytesIO(result_bytes))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
return {
|
||||
"result": _encode(result_bytes),
|
||||
"method": method_used,
|
||||
"original": {"width": orig_w, "height": orig_h},
|
||||
"output": {"width": result.width, "height": result.height},
|
||||
"scale": req.scale,
|
||||
}
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, status
|
||||
from sqlalchemy.orm import Session
|
||||
from typing import List
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
from app.database import get_db
|
||||
from app.models.project import Project
|
||||
from app.models.edit import Edit
|
||||
from app.schemas import ProjectCreate, ProjectResponse, EditResponse, UploadResponse
|
||||
from app.services.edit_service import EditService
|
||||
from app.config import settings
|
||||
|
||||
router = APIRouter(prefix="/projects", tags=["projects"])
|
||||
|
||||
|
||||
@router.post("/", response_model=ProjectResponse)
|
||||
def create_project(
|
||||
project: ProjectCreate,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Create a new project"""
|
||||
# For MVP, we'll use a default user_id of 1
|
||||
# In production, this would come from authentication
|
||||
user_id = 1
|
||||
|
||||
db_project = Project(
|
||||
user_id=user_id,
|
||||
name=project.name
|
||||
)
|
||||
db.add(db_project)
|
||||
db.commit()
|
||||
db.refresh(db_project)
|
||||
|
||||
# Create project directory
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(db_project.id)
|
||||
|
||||
return db_project
|
||||
|
||||
|
||||
@router.get("/", response_model=List[ProjectResponse])
|
||||
def list_projects(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List all projects"""
|
||||
projects = db.query(Project).offset(skip).limit(limit).all()
|
||||
return projects
|
||||
|
||||
|
||||
@router.get("/{project_id}", response_model=ProjectResponse)
|
||||
def get_project(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get a specific project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
return project
|
||||
|
||||
|
||||
@router.delete("/{project_id}")
|
||||
def delete_project(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Delete a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Delete project directory
|
||||
edit_service = EditService()
|
||||
project_dir = edit_service.get_project_dir(project_id)
|
||||
if project_dir.exists():
|
||||
shutil.rmtree(project_dir)
|
||||
|
||||
db.delete(project)
|
||||
db.commit()
|
||||
|
||||
return {"status": "success", "message": f"Project {project_id} deleted"}
|
||||
|
||||
|
||||
@router.post("/{project_id}/upload", response_model=UploadResponse)
|
||||
async def upload_image(
|
||||
project_id: int,
|
||||
file: UploadFile = File(...),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Upload an image to a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
# Validate file type
|
||||
if not file.content_type.startswith('image/'):
|
||||
raise HTTPException(status_code=400, detail="File must be an image")
|
||||
|
||||
# Create project directory
|
||||
edit_service = EditService()
|
||||
edit_service.ensure_project_dir(project_id)
|
||||
|
||||
# Save original and current images
|
||||
original_path = edit_service.get_original_image_path(project_id)
|
||||
current_path = edit_service.get_current_image_path(project_id)
|
||||
|
||||
# Read and validate image
|
||||
contents = await file.read()
|
||||
try:
|
||||
image = Image.open(BytesIO(contents))
|
||||
image = image.convert('RGBA')
|
||||
|
||||
# Save images
|
||||
image.save(original_path, 'PNG')
|
||||
image.save(current_path, 'PNG')
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid image file: {str(e)}")
|
||||
|
||||
return UploadResponse(
|
||||
project_id=project_id,
|
||||
original_url=f"/projects/{project_id}/original",
|
||||
current_url=f"/projects/{project_id}/current"
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{project_id}/edits", response_model=List[EditResponse])
|
||||
def list_edits(
|
||||
project_id: int,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List all edits for a project"""
|
||||
project = db.query(Project).filter(Project.id == project_id).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
edits = db.query(Edit).filter(Edit.project_id == project_id).order_by(Edit.created_at.desc()).all()
|
||||
return edits
|
||||
|
||||
|
||||
from io import BytesIO
|
||||
+1040
File diff suppressed because it is too large
Load Diff
+138
@@ -0,0 +1,138 @@
|
||||
from pydantic import BaseModel, EmailStr
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
# User schemas
|
||||
class UserCreate(BaseModel):
|
||||
email: EmailStr
|
||||
password: str
|
||||
|
||||
|
||||
class UserResponse(BaseModel):
|
||||
id: int
|
||||
email: str
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
# Project schemas
|
||||
class ProjectCreate(BaseModel):
|
||||
name: str
|
||||
|
||||
|
||||
class ProjectResponse(BaseModel):
|
||||
id: int
|
||||
user_id: int
|
||||
name: str
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
# Edit schemas
|
||||
class EditRequest(BaseModel):
|
||||
prompt: str
|
||||
mode: str = "A" # "A" or "B"
|
||||
selection_type: str # "rectangle", "ellipse", "lasso"
|
||||
bbox: Dict[str, int] # {x, y, width, height}
|
||||
feather_px: int = 0
|
||||
selection_data: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class EditResponse(BaseModel):
|
||||
id: int
|
||||
project_id: int
|
||||
created_at: datetime
|
||||
mode: str
|
||||
prompt: str
|
||||
selection_type: str
|
||||
bbox_json: str
|
||||
feather_px: int
|
||||
ai_provider: str
|
||||
status: str
|
||||
error_message: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
# Image upload
|
||||
class UploadResponse(BaseModel):
|
||||
project_id: int
|
||||
original_url: str
|
||||
current_url: str
|
||||
|
||||
|
||||
# Patch Library schemas
|
||||
class PatchCreate(BaseModel):
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
source_type: str # "ai_generated", "manual_selection", "imported"
|
||||
source_project_id: Optional[int] = None
|
||||
source_edit_id: Optional[int] = None
|
||||
category: Optional[str] = None
|
||||
tags: Optional[str] = None
|
||||
bbox: Optional[Dict[str, int]] = None
|
||||
|
||||
|
||||
class PatchResponse(BaseModel):
|
||||
id: int
|
||||
name: str
|
||||
description: Optional[str]
|
||||
created_at: datetime
|
||||
source_type: str
|
||||
source_project_id: Optional[int]
|
||||
source_edit_id: Optional[int]
|
||||
width: int
|
||||
height: int
|
||||
tags: Optional[str]
|
||||
category: Optional[str]
|
||||
is_public: bool
|
||||
file_path: str
|
||||
thumbnail_path: Optional[str]
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class PatchApply(BaseModel):
|
||||
project_id: int
|
||||
patch_id: int
|
||||
bbox: Dict[str, int]
|
||||
feather_px: int = 5
|
||||
|
||||
|
||||
# Text-to-Image schemas
|
||||
class TextToImageRequest(BaseModel):
|
||||
prompt: str
|
||||
width: int = 1024
|
||||
height: int = 1024
|
||||
negative_prompt: Optional[str] = None
|
||||
ai_provider: Optional[str] = None
|
||||
ai_model: Optional[str] = None
|
||||
create_project: bool = True
|
||||
project_name: Optional[str] = None
|
||||
|
||||
|
||||
class TextToImageResponse(BaseModel):
|
||||
status: str
|
||||
prompt: str
|
||||
width: int
|
||||
height: int
|
||||
project_id: Optional[int] = None
|
||||
image_url: Optional[str] = None
|
||||
layer_position: Optional[Dict[str, int]] = None
|
||||
ai_provider: str
|
||||
ai_model: Optional[str] = None
|
||||
|
||||
|
||||
# Generic responses
|
||||
class StatusResponse(BaseModel):
|
||||
status: str
|
||||
message: Optional[str] = None
|
||||
data: Optional[Any] = None
|
||||
@@ -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 image API — gpt-image-2 (generation + edits, same model/endpoint family)."""
|
||||
|
||||
def __init__(self, api_key: str, model: str = "gpt-image-2", edit_model: str = "gpt-image-2"):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.edit_model = edit_model
|
||||
self.base_url = "https://api.openai.com/v1"
|
||||
|
||||
async def _fetch_result(self, client: httpx.AsyncClient, item: dict) -> bytes:
|
||||
# gpt-image-1/2 only ever return b64_json; dall-e-2/dall-e-3 default to a url.
|
||||
if item.get('b64_json'):
|
||||
return base64.b64decode(item['b64_json'])
|
||||
image_response = await client.get(item['url'])
|
||||
image_response.raise_for_status()
|
||||
return image_response.content
|
||||
|
||||
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's images/edits endpoint"""
|
||||
|
||||
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 = {
|
||||
'model': model or self.edit_model,
|
||||
'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()
|
||||
return await self._fetch_result(client, result['data'][0])
|
||||
|
||||
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's images/generations endpoint"""
|
||||
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
data = {
|
||||
'model': model or self.model,
|
||||
'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()
|
||||
return await self._fetch_result(client, result['data'][0])
|
||||
|
||||
|
||||
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, settings.openai_model, settings.openai_edit_model)
|
||||
|
||||
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)
|
||||
+391
@@ -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,479 @@
|
||||
"""
|
||||
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-2 (generation + edits, same model/endpoint family)."""
|
||||
|
||||
def __init__(self, api_key: str, model: str = "gpt-image-2", edit_model: str = "gpt-image-2"):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.edit_model = edit_model
|
||||
self.base_url = "https://api.openai.com/v1"
|
||||
|
||||
def _headers(self):
|
||||
return {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
async def _fetch_result(self, client: httpx.AsyncClient, item: dict) -> bytes:
|
||||
# gpt-image-1/2 only ever return b64_json; dall-e-2/dall-e-3 default to a url.
|
||||
if item.get("b64_json"):
|
||||
return base64.b64decode(item["b64_json"])
|
||||
img_r = await client.get(item["url"])
|
||||
img_r.raise_for_status()
|
||||
return img_r.content
|
||||
|
||||
async def inpaint(self, image_bytes: bytes, mask_bytes: bytes, prompt: str, params: dict) -> bytes:
|
||||
model = (params or {}).get("model") or self.edit_model
|
||||
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 = {"model": model, "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()
|
||||
return await self._fetch_result(client, r.json()["data"][0])
|
||||
|
||||
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"
|
||||
model = (params or {}).get("model") or self.model
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
data = {"model": 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()
|
||||
return await self._fetch_result(client, r.json()["data"][0])
|
||||
|
||||
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, settings.openai_edit_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)
|
||||
+494
@@ -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)
|
||||
@@ -0,0 +1,228 @@
|
||||
from PIL import Image, ImageFilter, ImageDraw
|
||||
import numpy as np
|
||||
from io import BytesIO
|
||||
from typing import Tuple, Dict
|
||||
import cv2
|
||||
|
||||
|
||||
def bytes_to_image(image_bytes: bytes) -> Image.Image:
|
||||
"""Convert bytes to PIL Image"""
|
||||
return Image.open(BytesIO(image_bytes)).convert('RGBA')
|
||||
|
||||
|
||||
def image_to_bytes(image: Image.Image, format: str = 'PNG') -> bytes:
|
||||
"""Convert PIL Image to bytes"""
|
||||
buffer = BytesIO()
|
||||
image.save(buffer, format=format)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def crop_patch(image: Image.Image, bbox: Dict[str, int]) -> Image.Image:
|
||||
"""
|
||||
Crop a patch from the image using bounding box
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
bbox: Dictionary with x, y, width, height
|
||||
|
||||
Returns:
|
||||
Cropped patch as PIL Image
|
||||
"""
|
||||
x, y, width, height = bbox['x'], bbox['y'], bbox['width'], bbox['height']
|
||||
return image.crop((x, y, x + width, y + height))
|
||||
|
||||
|
||||
def create_feathered_mask(mask: Image.Image, feather_px: int) -> Image.Image:
|
||||
"""
|
||||
Apply feathering (Gaussian blur) to mask edges
|
||||
|
||||
Args:
|
||||
mask: Binary mask image (grayscale)
|
||||
feather_px: Feather radius in pixels
|
||||
|
||||
Returns:
|
||||
Feathered mask
|
||||
"""
|
||||
if feather_px <= 0:
|
||||
return mask
|
||||
|
||||
# Apply Gaussian blur for feathering
|
||||
feathered = mask.filter(ImageFilter.GaussianBlur(radius=feather_px))
|
||||
return feathered
|
||||
|
||||
|
||||
def blend_patch(
|
||||
original_patch: Image.Image,
|
||||
regenerated_patch: Image.Image,
|
||||
mask: Image.Image,
|
||||
feather_px: int = 0,
|
||||
preserve_alpha: bool = True
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Blend regenerated patch with original using mask.
|
||||
Preserves original alpha channel for semi-transparent areas (veils, glass, etc).
|
||||
|
||||
Args:
|
||||
original_patch: Original cropped patch
|
||||
regenerated_patch: AI-regenerated patch
|
||||
mask: Binary mask (same size as patches)
|
||||
feather_px: Feather radius for smooth blending
|
||||
preserve_alpha: If True, preserves original alpha channel
|
||||
|
||||
Returns:
|
||||
Blended patch with preserved transparency
|
||||
"""
|
||||
# Ensure all images are the same size
|
||||
if regenerated_patch.size != original_patch.size:
|
||||
regenerated_patch = regenerated_patch.resize(original_patch.size, Image.Resampling.LANCZOS)
|
||||
|
||||
if mask.size != original_patch.size:
|
||||
mask = mask.resize(original_patch.size, Image.Resampling.LANCZOS)
|
||||
|
||||
# Convert mask to grayscale if needed
|
||||
if mask.mode != 'L':
|
||||
mask = mask.convert('L')
|
||||
|
||||
# Apply feathering to mask
|
||||
feathered_mask = create_feathered_mask(mask, feather_px)
|
||||
|
||||
# Convert images to RGBA, storing original alpha
|
||||
original_rgba = original_patch.convert('RGBA')
|
||||
original_alpha = original_rgba.split()[3] # Store original alpha channel
|
||||
|
||||
regenerated_rgba = regenerated_patch.convert('RGBA')
|
||||
|
||||
# Blend using the feathered mask
|
||||
blended = Image.composite(regenerated_rgba, original_rgba, feathered_mask)
|
||||
|
||||
# Restore original alpha channel to preserve transparency
|
||||
# This keeps semi-transparent areas (veils, glass, smoke) intact
|
||||
if preserve_alpha:
|
||||
r, g, b, _ = blended.split()
|
||||
blended = Image.merge('RGBA', (r, g, b, original_alpha))
|
||||
|
||||
return blended
|
||||
|
||||
|
||||
def insert_patch(
|
||||
full_image: Image.Image,
|
||||
patch: Image.Image,
|
||||
bbox: Dict[str, int]
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Insert a patch back into the full image at the specified bbox
|
||||
|
||||
Args:
|
||||
full_image: Full original image
|
||||
patch: Patch to insert
|
||||
bbox: Bounding box {x, y, width, height}
|
||||
|
||||
Returns:
|
||||
Full image with patch inserted
|
||||
"""
|
||||
result = full_image.copy()
|
||||
x, y = bbox['x'], bbox['y']
|
||||
|
||||
# Ensure patch is the correct size
|
||||
if patch.size != (bbox['width'], bbox['height']):
|
||||
patch = patch.resize((bbox['width'], bbox['height']), Image.Resampling.LANCZOS)
|
||||
|
||||
# Paste the patch
|
||||
result.paste(patch, (x, y), patch if patch.mode == 'RGBA' else None)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def create_mask_from_selection(
|
||||
width: int,
|
||||
height: int,
|
||||
selection_type: str,
|
||||
selection_data: Dict
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Create a binary mask from selection data
|
||||
|
||||
Args:
|
||||
width: Mask width
|
||||
height: Mask height
|
||||
selection_type: "rectangle", "ellipse", or "lasso"
|
||||
selection_data: Selection-specific data
|
||||
|
||||
Returns:
|
||||
Binary mask (white = selected, black = not selected)
|
||||
"""
|
||||
mask = Image.new('L', (width, height), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
|
||||
if selection_type == "rectangle":
|
||||
# Fill entire rectangle
|
||||
draw.rectangle([0, 0, width, height], fill=255)
|
||||
|
||||
elif selection_type == "ellipse":
|
||||
# Fill entire ellipse
|
||||
draw.ellipse([0, 0, width, height], fill=255)
|
||||
|
||||
elif selection_type == "lasso":
|
||||
# Draw polygon from points
|
||||
points = selection_data.get('points', [])
|
||||
if points:
|
||||
# Convert points to relative coordinates within bbox
|
||||
draw.polygon(points, fill=255)
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def ensure_even_dimensions(image: Image.Image) -> Image.Image:
|
||||
"""
|
||||
Ensure image dimensions are even numbers (required by some AI providers)
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
|
||||
Returns:
|
||||
Image with even dimensions
|
||||
"""
|
||||
width, height = image.size
|
||||
new_width = width if width % 2 == 0 else width + 1
|
||||
new_height = height if height % 2 == 0 else height + 1
|
||||
|
||||
if (new_width, new_height) != (width, height):
|
||||
new_image = Image.new(image.mode, (new_width, new_height), (0, 0, 0, 0))
|
||||
new_image.paste(image, (0, 0))
|
||||
return new_image
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def resize_for_ai(image: Image.Image, max_size: int = 1024) -> Tuple[Image.Image, float]:
|
||||
"""
|
||||
Resize image if needed for AI processing (max dimension)
|
||||
|
||||
Args:
|
||||
image: PIL Image
|
||||
max_size: Maximum dimension size
|
||||
|
||||
Returns:
|
||||
Tuple of (resized image, scale factor)
|
||||
"""
|
||||
width, height = image.size
|
||||
max_dim = max(width, height)
|
||||
|
||||
if max_dim > max_size:
|
||||
scale = max_size / max_dim
|
||||
new_width = int(width * scale)
|
||||
new_height = int(height * scale)
|
||||
resized = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
return ensure_even_dimensions(resized), scale
|
||||
|
||||
return ensure_even_dimensions(image), 1.0
|
||||
|
||||
|
||||
def scale_bbox(bbox: Dict[str, int], scale: float) -> Dict[str, int]:
|
||||
"""Scale bounding box coordinates"""
|
||||
return {
|
||||
'x': int(bbox['x'] * scale),
|
||||
'y': int(bbox['y'] * scale),
|
||||
'width': int(bbox['width'] * scale),
|
||||
'height': int(bbox['height'] * scale)
|
||||
}
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
#!/bin/bash
|
||||
# =============================================================================
|
||||
# AI Photo Edit - Container Startup Script
|
||||
# =============================================================================
|
||||
# This script runs when the container starts. It:
|
||||
# 1. Initializes the database
|
||||
# 2. Downloads SAM model automatically (can be disabled with AUTO_DOWNLOAD_SAM=false)
|
||||
# 3. Downloads U2Net model automatically (can be disabled with AUTO_DOWNLOAD_U2NET=false)
|
||||
# 4. Downloads sample eye images if the catalog is empty
|
||||
# 5. Starts the FastAPI server
|
||||
# =============================================================================
|
||||
|
||||
set -e
|
||||
|
||||
echo "=========================================="
|
||||
echo "AI Photo Edit - Starting Up"
|
||||
echo "=========================================="
|
||||
|
||||
# Ensure data directories exist
|
||||
mkdir -p /app/data/projects
|
||||
mkdir -p /app/data/patches
|
||||
mkdir -p /app/data/models
|
||||
mkdir -p /app/data/patch_library
|
||||
|
||||
# Initialize database FIRST (before eye import)
|
||||
echo ""
|
||||
echo "Initializing database..."
|
||||
echo "------------------------------------------"
|
||||
cd /app && python /scripts/init_database.py || echo "Warning: Database init failed (non-fatal)"
|
||||
|
||||
# Check and download SAM model automatically
|
||||
echo ""
|
||||
echo "Checking SAM model (Smart Select)..."
|
||||
echo "------------------------------------------"
|
||||
if [ -f "/app/data/models/sam_model.pth" ] || \
|
||||
[ -f "/app/data/models/sam_vit_b_01ec64.pth" ] || \
|
||||
[ -f "/app/data/models/sam_vit_l_0b3195.pth" ] || \
|
||||
[ -f "/app/data/models/sam_vit_h_4b8939.pth" ]; then
|
||||
echo "✓ SAM model found - Smart Select will use local AI (free, offline)"
|
||||
else
|
||||
# Auto-download SAM unless explicitly disabled
|
||||
AUTO_DOWNLOAD_SAM="${AUTO_DOWNLOAD_SAM:-true}"
|
||||
if [ "$AUTO_DOWNLOAD_SAM" = "true" ]; then
|
||||
echo "SAM model not found. Downloading automatically..."
|
||||
echo "(This is a one-time ~375MB download that persists across rebuilds)"
|
||||
echo ""
|
||||
python /scripts/download_sam_model.py vit_b || {
|
||||
echo ""
|
||||
echo "⚠ SAM download failed (non-fatal)"
|
||||
echo " Smart Select will fall back to Replicate API (requires REPLICATE_API_KEY)"
|
||||
echo " To retry later: docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py"
|
||||
}
|
||||
else
|
||||
echo ""
|
||||
echo "⚠ SAM model not found (AUTO_DOWNLOAD_SAM=false)"
|
||||
echo ""
|
||||
echo " Smart Select will use Replicate API (requires REPLICATE_API_KEY)"
|
||||
echo ""
|
||||
echo " To enable FREE offline Smart Select, run:"
|
||||
echo " docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py"
|
||||
echo ""
|
||||
fi
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Checking U2Net model (Remove Background)..."
|
||||
echo "------------------------------------------"
|
||||
if [ -f "/app/data/models/u2net.onnx" ] || [ -f "/app/data/models/u2netp.onnx" ]; then
|
||||
echo "✓ U2Net model found - Remove Background will use local AI (free, offline)"
|
||||
else
|
||||
# Auto-download U2Net unless explicitly disabled
|
||||
AUTO_DOWNLOAD_U2NET="${AUTO_DOWNLOAD_U2NET:-true}"
|
||||
if [ "$AUTO_DOWNLOAD_U2NET" = "true" ]; then
|
||||
echo "U2Net model not found. Downloading automatically..."
|
||||
echo "(This is a one-time ~176MB download that persists across rebuilds)"
|
||||
echo ""
|
||||
python /scripts/download_u2net_model.py u2net || {
|
||||
echo ""
|
||||
echo "⚠ U2Net download failed (non-fatal)"
|
||||
echo " Remove Background will fall back to rembg (if installed)"
|
||||
echo " To retry later: docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py u2net"
|
||||
}
|
||||
else
|
||||
echo ""
|
||||
echo "⚠ U2Net model not found (AUTO_DOWNLOAD_U2NET=false)"
|
||||
echo ""
|
||||
echo " Remove Background will fall back to rembg (if installed)"
|
||||
echo ""
|
||||
echo " To enable FREE offline Remove Background, run:"
|
||||
echo " docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py u2net"
|
||||
echo ""
|
||||
fi
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Checking GPU capabilities..."
|
||||
echo "------------------------------------------"
|
||||
python /scripts/gpu_setup.py || echo "Warning: GPU detection failed (non-fatal)"
|
||||
|
||||
echo ""
|
||||
echo "=========================================="
|
||||
echo "Starting FastAPI server..."
|
||||
echo "=========================================="
|
||||
|
||||
# Start the server
|
||||
exec uvicorn app.main:app --host 0.0.0.0 --port 8000
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
# =============================================================================
|
||||
# GPU / Local Diffusion dependencies
|
||||
# Install alongside requirements.txt when running with AI_PROVIDER=local_gpu
|
||||
#
|
||||
# Usage:
|
||||
# pip install -r requirements.txt -r requirements.gpu.txt
|
||||
#
|
||||
# These are pre-installed in Dockerfile.gpu; optional in the standard image.
|
||||
# =============================================================================
|
||||
|
||||
# HuggingFace Diffusers ecosystem
|
||||
# Pinned <0.29.0: diffusers 0.29.0 added torch.xpu (Intel GPU) which fails on
|
||||
# PyTorch 2.1.x with "AttributeError: module 'torch' has no attribute 'xpu'".
|
||||
# Upgrade the base image in Dockerfile.gpu to pytorch 2.4+ before lifting this pin.
|
||||
# (FLUX support requires diffusers>=0.29 + PyTorch>=2.4; SDXL/SD works fine here.)
|
||||
diffusers>=0.28.0,<0.29.0
|
||||
transformers>=4.36.0,<4.40.0
|
||||
accelerate>=0.27.0
|
||||
huggingface-hub>=0.23.0
|
||||
safetensors>=0.4.0
|
||||
|
||||
# Required by SDXL pipelines
|
||||
invisible-watermark>=0.2.0
|
||||
omegaconf>=2.3.0
|
||||
|
||||
# Required by FLUX (T5 text encoder tokenizer)
|
||||
sentencepiece>=0.2.0
|
||||
|
||||
# xformers — reduces attention VRAM ~20-30%, often unlocks the next model tier
|
||||
# Must match your PyTorch+CUDA version; leave out if unsure.
|
||||
# Install post-container-start if needed:
|
||||
# pip install xformers --index-url https://download.pytorch.org/whl/cu121
|
||||
# xformers
|
||||
|
||||
# Background removal — BEN2 (default, clean cutouts/hair) + BiRefNet-HR
|
||||
# (high-res/print alternate). Both MIT-licensed. Verified against upstream
|
||||
# source: neither requires torch>=2.5 despite the BiRefNet repo's own
|
||||
# requirements.txt floor — that pin is for its training/eval scripts, not
|
||||
# the inference path used here. Weights download from HuggingFace on first
|
||||
# use (cached via the hf_cache bind mount, same as the diffusion models).
|
||||
ben2 @ git+https://github.com/PramaLLC/BEN2.git
|
||||
timm>=1.0.10
|
||||
einops>=0.6.0
|
||||
kornia>=0.7.0
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
fastapi==0.109.0
|
||||
uvicorn[standard]==0.27.0
|
||||
python-multipart==0.0.6
|
||||
Pillow>=10.0.0,<11.0.0
|
||||
numpy<2.0.0
|
||||
sqlalchemy==2.0.25
|
||||
python-jose[cryptography]==3.3.0
|
||||
passlib[bcrypt]==1.7.4
|
||||
python-dotenv==1.0.0
|
||||
aiofiles==23.2.1
|
||||
httpx==0.26.0
|
||||
pydantic==2.5.3
|
||||
pydantic-settings==2.1.0
|
||||
email-validator==2.1.0
|
||||
opencv-python-headless>=4.10.0
|
||||
# SAM (Segment Anything) for smart object selection - runs locally, no API needed
|
||||
torch==2.1.2
|
||||
torchvision==0.16.2
|
||||
segment-anything @ git+https://github.com/facebookresearch/segment-anything.git
|
||||
|
||||
# Local AI inpainting — LaMa model (auto GPU/CPU, no API key needed)
|
||||
simple-lama-inpainting
|
||||
|
||||
# Background removal — rembg enabled now that opencv 4.10+ supports numpy 2.x
|
||||
rembg[gpu]
|
||||
+225
@@ -0,0 +1,225 @@
|
||||
# Eye Catalog Import Scripts
|
||||
|
||||
Tools for populating your carved eye catalog with public domain examples.
|
||||
|
||||
---
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Download Classical Eyes
|
||||
|
||||
Follow the guide: `/docs/PUBLIC_DOMAIN_EYE_SOURCES.md`
|
||||
|
||||
**Best sources:**
|
||||
- Metropolitan Museum (CC0)
|
||||
- Smithsonian Open Access
|
||||
- Getty Museum Open Content
|
||||
|
||||
**Download 10-20 high-res images** of carved eyes from classical sculptures.
|
||||
|
||||
---
|
||||
|
||||
### 2. Crop the Eyes
|
||||
|
||||
Use any image editor (Photoshop, GIMP, Preview, etc.):
|
||||
|
||||
1. Open statue photo
|
||||
2. Zoom in on eye
|
||||
3. Crop just the eye (include eyelids, socket, tear duct)
|
||||
4. Save as PNG with descriptive name:
|
||||
- `greek_serene_left.png`
|
||||
- `roman_fierce_right.png`
|
||||
- `egyptian_wise_left.png`
|
||||
|
||||
---
|
||||
|
||||
### 3. Import to Catalog
|
||||
|
||||
**Easy way (one at a time):**
|
||||
```bash
|
||||
cd backend/scripts
|
||||
|
||||
# Import a Greek serene eye
|
||||
python import_eyes.py greek_serene_left.png \
|
||||
--emotion serene \
|
||||
--side left \
|
||||
--style greek
|
||||
|
||||
# Import a Roman fierce eye
|
||||
python import_eyes.py roman_fierce_right.png \
|
||||
--emotion fierce \
|
||||
--side right \
|
||||
--style roman
|
||||
```
|
||||
|
||||
**Batch import:**
|
||||
```bash
|
||||
# Import all Greek eyes at once
|
||||
python import_eyes.py greek_*.png \
|
||||
--emotion serene \
|
||||
--side both \
|
||||
--style greek
|
||||
|
||||
# Import all Roman eyes
|
||||
python import_eyes.py roman_*.png \
|
||||
--emotion fierce \
|
||||
--side both \
|
||||
--style roman
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Available Options
|
||||
|
||||
### Emotions
|
||||
- `serene` - Peaceful, calm (most classical Greek)
|
||||
- `fierce` - Intense, powerful (Hellenistic, Alexander)
|
||||
- `wise` - Aged, experienced (Roman senators)
|
||||
- `peaceful` - Gentle, kind (Archaic Greek)
|
||||
- `joyful` - Happy, smiling (rare in classical)
|
||||
- `sorrowful` - Sad, mourning (some Hellenistic)
|
||||
- `neutral` - Default, no strong emotion
|
||||
|
||||
### Styles
|
||||
- `greek` - Classical Greek (450-400 BCE)
|
||||
- `roman` - Roman Republican/Imperial
|
||||
- `egyptian` - Ancient Egyptian carved eyes
|
||||
- `renaissance` - Renaissance sculpture
|
||||
- `baroque` - Baroque period
|
||||
- `modern` - Contemporary carving
|
||||
- `custom` - Your own style
|
||||
|
||||
### Sides
|
||||
- `left` - Left eye
|
||||
- `right` - Right eye
|
||||
- `both` - Can be used for either (symmetric)
|
||||
|
||||
---
|
||||
|
||||
## Example Workflow
|
||||
|
||||
### Build a Complete Catalog
|
||||
|
||||
```bash
|
||||
# 1. Download eyes from Met Museum
|
||||
# (See PUBLIC_DOMAIN_EYE_SOURCES.md)
|
||||
|
||||
# 2. Crop and save them:
|
||||
# greek_serene_left.png
|
||||
# greek_serene_right.png
|
||||
# roman_fierce_left.png
|
||||
# roman_fierce_right.png
|
||||
# etc.
|
||||
|
||||
# 3. Import them all:
|
||||
|
||||
python import_eyes.py greek_serene_left.png --emotion serene --side left --style greek
|
||||
python import_eyes.py greek_serene_right.png --emotion serene --side right --style greek
|
||||
python import_eyes.py roman_fierce_left.png --emotion fierce --side left --style roman
|
||||
python import_eyes.py roman_fierce_right.png --emotion fierce --side right --style roman
|
||||
|
||||
# Or batch:
|
||||
python import_eyes.py greek_*.png --emotion serene --side both --style greek
|
||||
python import_eyes.py roman_*.png --emotion fierce --side both --style roman
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Recommended Starting Collection
|
||||
|
||||
### 10 Essential Eyes
|
||||
|
||||
1. **Greek Serene Left** (peaceful carvings)
|
||||
2. **Greek Serene Right**
|
||||
3. **Greek Archaic Left** (stylized, simple)
|
||||
4. **Greek Archaic Right**
|
||||
5. **Roman Fierce Left** (powerful portraits)
|
||||
6. **Roman Fierce Right**
|
||||
7. **Roman Wise Left** (aged, realistic)
|
||||
8. **Roman Wise Right**
|
||||
9. **Egyptian Stylized Left** (distinctive style)
|
||||
10. **Egyptian Stylized Right**
|
||||
|
||||
This gives you 5 styles/emotions to start!
|
||||
|
||||
---
|
||||
|
||||
## Check Your Catalog
|
||||
|
||||
### Via API
|
||||
|
||||
```bash
|
||||
# List all eyes in catalog
|
||||
curl http://localhost:8101/patches/?category=carved_eye
|
||||
|
||||
# Filter by emotion
|
||||
curl http://localhost:8101/patches/?category=carved_eye&tags=serene
|
||||
|
||||
# Filter by style
|
||||
curl http://localhost:8101/patches/?tags=greek
|
||||
```
|
||||
|
||||
### Via Web Interface
|
||||
|
||||
Go to: `http://your-server:3080`
|
||||
|
||||
Navigate to patch library to browse your eyes visually.
|
||||
|
||||
---
|
||||
|
||||
## Advanced: Seed Script
|
||||
|
||||
For batch importing from the `seed_data/eyes/` directory:
|
||||
|
||||
```bash
|
||||
# 1. Place all cropped eyes in:
|
||||
mkdir -p seed_data/eyes/
|
||||
# Copy your eye images there
|
||||
|
||||
# 2. Edit seed_eye_catalog.py to add metadata
|
||||
|
||||
# 3. Run:
|
||||
python seed_eye_catalog.py
|
||||
```
|
||||
|
||||
This auto-imports all eyes in `seed_data/eyes/` with pre-configured metadata.
|
||||
|
||||
---
|
||||
|
||||
## Tips
|
||||
|
||||
1. **High resolution:** Use images 1000px+ for best results
|
||||
2. **Clean crops:** Include some surrounding area, not just the eyeball
|
||||
3. **Consistent naming:** Use descriptive filenames
|
||||
4. **Test first:** Import 2-3 eyes to test the workflow
|
||||
5. **Build gradually:** Start with 10 eyes, expand as needed
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**"File not found":**
|
||||
- Make sure you're in `backend/scripts/` directory
|
||||
- Use full path or relative path to image
|
||||
|
||||
**"Database connection error":**
|
||||
- Make sure backend is running: `docker compose up backend`
|
||||
- Check database exists: `ls -la ../../data/`
|
||||
|
||||
**"Import failed":**
|
||||
- Check image format (PNG, JPG supported)
|
||||
- Verify file isn't corrupted
|
||||
- Check file permissions
|
||||
|
||||
---
|
||||
|
||||
## Next Steps
|
||||
|
||||
After importing eyes:
|
||||
|
||||
1. **Test them:** Apply to a colored photo via API
|
||||
2. **Refine:** Add more variations as needed
|
||||
3. **Build library:** Aim for 20-30 eyes covering all emotions
|
||||
4. **Share:** Your best eyes can be exported and shared
|
||||
|
||||
**Your catalog of master sculptor's eyes is ready to use!**
|
||||
+210
@@ -0,0 +1,210 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Quick eye import script
|
||||
|
||||
Usage:
|
||||
python import_eyes.py greek_serene_left.png --emotion serene --side left --style greek
|
||||
python import_eyes.py *.png --emotion fierce --style roman
|
||||
|
||||
This will add eyes to the patch library with proper metadata.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
# Add parent directory to path to import app modules
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from app.database import SessionLocal
|
||||
from app.models.patch import Patch
|
||||
from app.services.patch_library import PatchLibraryService
|
||||
|
||||
|
||||
def import_eye(
|
||||
image_path: Path,
|
||||
emotion: str,
|
||||
side: str,
|
||||
style: str,
|
||||
description: str = None
|
||||
):
|
||||
"""
|
||||
Import a single eye image into the catalog
|
||||
|
||||
Args:
|
||||
image_path: Path to eye image file
|
||||
emotion: serene, fierce, wise, peaceful, joyful, sorrowful
|
||||
side: left, right, both
|
||||
style: greek, roman, egyptian, renaissance, custom
|
||||
description: Optional custom description
|
||||
"""
|
||||
|
||||
db = SessionLocal()
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
# Generate name from filename if not provided
|
||||
name = image_path.stem.replace('_', ' ').title()
|
||||
|
||||
# Auto-generate description if not provided
|
||||
if not description:
|
||||
description = f"{style.title()} carved eye, {emotion} expression, {side} side. Suitable for CNC wood carving."
|
||||
|
||||
# Generate tags
|
||||
tags = f"{style}, {emotion}, {side}, carved, cnc-ready, wood-carving"
|
||||
|
||||
print(f"\n📸 Importing: {name}")
|
||||
print(f" File: {image_path.name}")
|
||||
print(f" Style: {style}")
|
||||
print(f" Emotion: {emotion}")
|
||||
print(f" Side: {side}")
|
||||
|
||||
try:
|
||||
# Create patch record
|
||||
patch = Patch(
|
||||
name=name,
|
||||
description=description,
|
||||
source_type="imported",
|
||||
category="carved_eye",
|
||||
tags=tags,
|
||||
width=0,
|
||||
height=0,
|
||||
user_id=None,
|
||||
file_path=""
|
||||
)
|
||||
|
||||
db.add(patch)
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
# Save image file
|
||||
file_path = patch_service.save_patch_from_file(
|
||||
patch.id,
|
||||
str(image_path),
|
||||
create_thumb=True
|
||||
)
|
||||
|
||||
# Get and update dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
patch.file_path = file_path
|
||||
patch.thumbnail_path = str(patch_service.get_thumbnail_path(patch.id))
|
||||
|
||||
db.commit()
|
||||
|
||||
print(f"✅ Successfully imported!")
|
||||
print(f" ID: {patch.id}")
|
||||
print(f" Size: {width}x{height}px")
|
||||
|
||||
return patch.id
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error importing: {e}")
|
||||
db.rollback()
|
||||
return None
|
||||
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Import carved eye images into the patch library",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Import a single eye
|
||||
python import_eyes.py greek_serene_left.png --emotion serene --side left --style greek
|
||||
|
||||
# Import multiple eyes with same metadata
|
||||
python import_eyes.py roman_*.png --emotion fierce --side both --style roman
|
||||
|
||||
# With custom description
|
||||
python import_eyes.py statue_eye.png --emotion wise --side right --style roman --description "Emperor Augustus portrait eye"
|
||||
|
||||
Emotions: serene, fierce, wise, peaceful, joyful, sorrowful, neutral
|
||||
Sides: left, right, both
|
||||
Styles: greek, roman, egyptian, renaissance, baroque, modern, custom
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'images',
|
||||
nargs='+',
|
||||
help='Image file(s) to import (supports wildcards)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--emotion',
|
||||
required=True,
|
||||
choices=['serene', 'fierce', 'wise', 'peaceful', 'joyful', 'sorrowful', 'neutral'],
|
||||
help='Emotional expression of the eye'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--side',
|
||||
required=True,
|
||||
choices=['left', 'right', 'both'],
|
||||
help='Which eye (left or right)'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--style',
|
||||
required=True,
|
||||
choices=['greek', 'roman', 'egyptian', 'renaissance', 'baroque', 'modern', 'custom'],
|
||||
help='Carving style/period'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--description',
|
||||
help='Custom description (optional)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Resolve wildcards and get all image files
|
||||
image_files = []
|
||||
for pattern in args.images:
|
||||
path = Path(pattern)
|
||||
if '*' in pattern:
|
||||
# Wildcard - expand it
|
||||
parent = path.parent if path.parent.exists() else Path('.')
|
||||
image_files.extend(parent.glob(path.name))
|
||||
else:
|
||||
# Single file
|
||||
if path.exists():
|
||||
image_files.append(path)
|
||||
else:
|
||||
print(f"⚠️ File not found: {pattern}")
|
||||
|
||||
if not image_files:
|
||||
print("❌ No image files found!")
|
||||
return
|
||||
|
||||
print(f"\n🎨 Importing {len(image_files)} eye image(s) into catalog...")
|
||||
print(f" Style: {args.style}")
|
||||
print(f" Emotion: {args.emotion}")
|
||||
print(f" Side: {args.side}")
|
||||
print("="*60)
|
||||
|
||||
imported_count = 0
|
||||
for image_path in image_files:
|
||||
patch_id = import_eye(
|
||||
image_path,
|
||||
args.emotion,
|
||||
args.side,
|
||||
args.style,
|
||||
args.description
|
||||
)
|
||||
if patch_id:
|
||||
imported_count += 1
|
||||
|
||||
print("="*60)
|
||||
print(f"\n✅ Import complete! {imported_count}/{len(image_files)} eyes added to catalog")
|
||||
print(f"\n💡 Access your eye catalog at: http://your-server:3080")
|
||||
print(f" Or via API: GET /patches/?category=carved_eye")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,176 @@
|
||||
"""
|
||||
Seed the patch library with classical carved eyes from public domain sources
|
||||
|
||||
This script helps pre-populate the eye catalog with examples from:
|
||||
- Greek statues (Metropolitan Museum, Louvre)
|
||||
- Roman sculptures (Smithsonian, British Museum)
|
||||
- Renaissance carvings
|
||||
- Ancient Egyptian carved eyes
|
||||
|
||||
All images should be public domain (CC0, Public Domain Mark)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import httpx
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
|
||||
# Public domain eye examples to seed the catalog
|
||||
# These are examples - you would add actual URLs from museum APIs
|
||||
CLASSICAL_EYES = [
|
||||
{
|
||||
"name": "Greek Statue - Serene Left Eye",
|
||||
"description": "Classical Greek marble carving, convex eyeball, defined upper lid, deep socket. Perfect for serene expressions.",
|
||||
"category": "carved_eye",
|
||||
"tags": "greek, serene, left, marble, classical, convex, deep-socket",
|
||||
"style": "greek_classical",
|
||||
"emotion": "serene",
|
||||
"side": "left",
|
||||
"source_url": "https://images.metmuseum.org/...", # Example
|
||||
"source": "Metropolitan Museum of Art - Public Domain"
|
||||
},
|
||||
{
|
||||
"name": "Roman Sculpture - Fierce Right Eye",
|
||||
"description": "Roman marble, prominent brow ridge, intense gaze, sharp eyelid definition.",
|
||||
"category": "carved_eye",
|
||||
"tags": "roman, fierce, right, marble, intense, sharp-detail",
|
||||
"style": "roman_classical",
|
||||
"emotion": "fierce",
|
||||
"side": "right",
|
||||
"source_url": "https://...",
|
||||
"source": "Smithsonian - CC0"
|
||||
},
|
||||
{
|
||||
"name": "Greek Kouros - Peaceful Left Eye",
|
||||
"description": "Archaic Greek style, almond-shaped, subtle carving, peaceful expression.",
|
||||
"category": "carved_eye",
|
||||
"tags": "greek, peaceful, left, archaic, almond-shaped, subtle",
|
||||
"style": "greek_archaic",
|
||||
"emotion": "peaceful",
|
||||
"side": "left",
|
||||
"source_url": "https://...",
|
||||
"source": "Getty Museum - Public Domain"
|
||||
},
|
||||
{
|
||||
"name": "Roman Portrait - Wise Right Eye",
|
||||
"description": "Late Roman period, detailed eyelids, slight downward gaze, wisdom and age.",
|
||||
"category": "carved_eye",
|
||||
"tags": "roman, wise, right, portrait, detailed, aged",
|
||||
"style": "roman_portrait",
|
||||
"emotion": "wise",
|
||||
"side": "right",
|
||||
"source_url": "https://...",
|
||||
"source": "British Museum - CC0"
|
||||
},
|
||||
]
|
||||
|
||||
async def download_image(url: str) -> bytes:
|
||||
"""Download image from URL"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
async def crop_eye_from_statue(image_bytes: bytes, crop_box: tuple) -> bytes:
|
||||
"""
|
||||
Crop just the eye from a full statue photo
|
||||
|
||||
Args:
|
||||
image_bytes: Full statue image
|
||||
crop_box: (left, top, right, bottom) coordinates
|
||||
|
||||
Returns:
|
||||
Cropped eye image bytes
|
||||
"""
|
||||
img = Image.open(BytesIO(image_bytes))
|
||||
eye = img.crop(crop_box)
|
||||
|
||||
# Save as PNG
|
||||
buffer = BytesIO()
|
||||
eye.save(buffer, format='PNG')
|
||||
return buffer.getvalue()
|
||||
|
||||
async def seed_eye_catalog():
|
||||
"""
|
||||
Seed the patch library with classical carved eyes
|
||||
|
||||
NOTE: This is a template. You need to:
|
||||
1. Get actual public domain image URLs
|
||||
2. Manually crop the eyes (or provide crop coordinates)
|
||||
3. Run this to populate the catalog
|
||||
"""
|
||||
|
||||
from app.database import SessionLocal
|
||||
from app.models.patch import Patch
|
||||
from app.services.patch_library import PatchLibraryService
|
||||
|
||||
db = SessionLocal()
|
||||
patch_service = PatchLibraryService()
|
||||
|
||||
print("Seeding eye catalog with classical carved eyes...")
|
||||
|
||||
for eye_data in CLASSICAL_EYES:
|
||||
print(f"\nAdding: {eye_data['name']}")
|
||||
|
||||
# Create patch record
|
||||
patch = Patch(
|
||||
name=eye_data['name'],
|
||||
description=eye_data['description'],
|
||||
source_type="imported",
|
||||
category=eye_data['category'],
|
||||
tags=eye_data['tags'],
|
||||
width=0, # Will be set after image save
|
||||
height=0,
|
||||
user_id=None,
|
||||
file_path=""
|
||||
)
|
||||
|
||||
db.add(patch)
|
||||
db.commit()
|
||||
db.refresh(patch)
|
||||
|
||||
# Download and save image
|
||||
# NOTE: You need to manually download/crop these first
|
||||
# This is just the structure
|
||||
|
||||
try:
|
||||
# image_bytes = await download_image(eye_data['source_url'])
|
||||
# cropped_eye = await crop_eye_from_statue(image_bytes, crop_box)
|
||||
|
||||
# For now, you would manually place images in:
|
||||
# ./seed_data/eyes/greek_serene_left.png
|
||||
# ./seed_data/eyes/roman_fierce_right.png
|
||||
# etc.
|
||||
|
||||
seed_image_path = Path(__file__).parent / "seed_data" / "eyes" / f"{eye_data['style']}_{eye_data['emotion']}_{eye_data['side']}.png"
|
||||
|
||||
if seed_image_path.exists():
|
||||
file_path = patch_service.save_patch_from_file(
|
||||
patch.id,
|
||||
str(seed_image_path),
|
||||
create_thumb=True
|
||||
)
|
||||
|
||||
# Get dimensions
|
||||
width, height = patch_service.get_patch_size(patch.id)
|
||||
patch.width = width
|
||||
patch.height = height
|
||||
patch.file_path = file_path
|
||||
patch.thumbnail_path = str(patch_service.get_thumbnail_path(patch.id))
|
||||
|
||||
db.commit()
|
||||
print(f"✅ Added {eye_data['name']}")
|
||||
else:
|
||||
print(f"⚠️ Image not found: {seed_image_path}")
|
||||
print(f" Please download and crop eye from: {eye_data['source']}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error adding {eye_data['name']}: {e}")
|
||||
db.rollback()
|
||||
|
||||
db.close()
|
||||
print("\n✅ Eye catalog seeding complete!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(seed_eye_catalog())
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
#!/usr/bin/env bash
|
||||
# bring-up-local-gpu.sh — start the GPU container.
|
||||
#
|
||||
# Run this each time you want to start the app.
|
||||
# Run ./install-local-gpu.sh once first on a new machine.
|
||||
#
|
||||
# Before starting, this fetches any missing models on the host (outside
|
||||
# Docker) via ./prefetch-models.sh — in-container DNS/network is unreliable
|
||||
# on some hosts, so this is the default now, not a manual troubleshooting
|
||||
# step. It never blocks startup: if it fails (no network, no python3, etc.)
|
||||
# the container still starts and falls back to its own in-container download.
|
||||
#
|
||||
# Usage:
|
||||
# ./bring-up-local-gpu.sh # start (detached, rebuild if needed)
|
||||
# ./bring-up-local-gpu.sh --no-build # start without rebuilding
|
||||
# ./bring-up-local-gpu.sh down # stop and remove container
|
||||
# ./bring-up-local-gpu.sh logs -f # tail logs
|
||||
#
|
||||
# Force pip layer rebuild (e.g. after requirements change):
|
||||
# BUILDID=$(date +%s) ./bring-up-local-gpu.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
# Pre-create ./data as the current (non-root) user. Otherwise, on a fresh
|
||||
# checkout, Docker's daemon (root) auto-creates these bind-mount sources on
|
||||
# the first 'up' — leaving them root-owned and blocking this same user from
|
||||
# later writing to them without sudo (e.g. ./prefetch-models.sh). No-op if
|
||||
# they already exist, regardless of current ownership.
|
||||
mkdir -p data/models data/hf_cache data/projects data/patches
|
||||
|
||||
if [ $# -eq 0 ]; then
|
||||
# Best-effort: fetch any missing models on the host first (see header).
|
||||
./prefetch-models.sh --sdxl \
|
||||
|| echo "⚠ Model prefetch had failures (see above) — continuing anyway, the container will retry in-container."
|
||||
exec docker compose -f docker-compose.gpu.yml up -d --build
|
||||
else
|
||||
exec docker compose -f docker-compose.gpu.yml "$@"
|
||||
fi
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,11 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert PPM files to PNG using PIL."""
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
ppm_dir = Path(__file__).parent
|
||||
for ppm_file in ppm_dir.glob("*.ppm"):
|
||||
png_file = ppm_file.with_suffix('.png')
|
||||
img = Image.open(ppm_file)
|
||||
img.save(png_file, 'PNG')
|
||||
print(f"Converted: {ppm_file.name} -> {png_file.name}")
|
||||
@@ -0,0 +1,164 @@
|
||||
[
|
||||
{
|
||||
"filename": "classic_realistic_blue.png",
|
||||
"ppm_filename": "classic_realistic_blue.ppm",
|
||||
"name": "Classic Realistic Eye - Blue",
|
||||
"description": "A realistic style eye with blue iris color",
|
||||
"tags": "eye,classic,realistic,blue,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_realistic_green.png",
|
||||
"ppm_filename": "classic_realistic_green.ppm",
|
||||
"name": "Classic Realistic Eye - Green",
|
||||
"description": "A realistic style eye with green iris color",
|
||||
"tags": "eye,classic,realistic,green,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_realistic_brown.png",
|
||||
"ppm_filename": "classic_realistic_brown.ppm",
|
||||
"name": "Classic Realistic Eye - Brown",
|
||||
"description": "A realistic style eye with brown iris color",
|
||||
"tags": "eye,classic,realistic,brown,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_realistic_hazel.png",
|
||||
"ppm_filename": "classic_realistic_hazel.ppm",
|
||||
"name": "Classic Realistic Eye - Hazel",
|
||||
"description": "A realistic style eye with hazel iris color",
|
||||
"tags": "eye,classic,realistic,hazel,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_realistic_grey.png",
|
||||
"ppm_filename": "classic_realistic_grey.ppm",
|
||||
"name": "Classic Realistic Eye - Grey",
|
||||
"description": "A realistic style eye with grey iris color",
|
||||
"tags": "eye,classic,realistic,grey,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_realistic_amber.png",
|
||||
"ppm_filename": "classic_realistic_amber.ppm",
|
||||
"name": "Classic Realistic Eye - Amber",
|
||||
"description": "A realistic style eye with amber iris color",
|
||||
"tags": "eye,classic,realistic,amber,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_anime_blue.png",
|
||||
"ppm_filename": "classic_anime_blue.ppm",
|
||||
"name": "Classic Anime Eye - Blue",
|
||||
"description": "A anime style eye with blue iris color",
|
||||
"tags": "eye,classic,anime,blue,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_anime_green.png",
|
||||
"ppm_filename": "classic_anime_green.ppm",
|
||||
"name": "Classic Anime Eye - Green",
|
||||
"description": "A anime style eye with green iris color",
|
||||
"tags": "eye,classic,anime,green,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_anime_brown.png",
|
||||
"ppm_filename": "classic_anime_brown.ppm",
|
||||
"name": "Classic Anime Eye - Brown",
|
||||
"description": "A anime style eye with brown iris color",
|
||||
"tags": "eye,classic,anime,brown,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_anime_hazel.png",
|
||||
"ppm_filename": "classic_anime_hazel.ppm",
|
||||
"name": "Classic Anime Eye - Hazel",
|
||||
"description": "A anime style eye with hazel iris color",
|
||||
"tags": "eye,classic,anime,hazel,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_anime_grey.png",
|
||||
"ppm_filename": "classic_anime_grey.ppm",
|
||||
"name": "Classic Anime Eye - Grey",
|
||||
"description": "A anime style eye with grey iris color",
|
||||
"tags": "eye,classic,anime,grey,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_anime_amber.png",
|
||||
"ppm_filename": "classic_anime_amber.ppm",
|
||||
"name": "Classic Anime Eye - Amber",
|
||||
"description": "A anime style eye with amber iris color",
|
||||
"tags": "eye,classic,anime,amber,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_cartoon_blue.png",
|
||||
"ppm_filename": "classic_cartoon_blue.ppm",
|
||||
"name": "Classic Cartoon Eye - Blue",
|
||||
"description": "A cartoon style eye with blue iris color",
|
||||
"tags": "eye,classic,cartoon,blue,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_cartoon_green.png",
|
||||
"ppm_filename": "classic_cartoon_green.ppm",
|
||||
"name": "Classic Cartoon Eye - Green",
|
||||
"description": "A cartoon style eye with green iris color",
|
||||
"tags": "eye,classic,cartoon,green,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_cartoon_brown.png",
|
||||
"ppm_filename": "classic_cartoon_brown.ppm",
|
||||
"name": "Classic Cartoon Eye - Brown",
|
||||
"description": "A cartoon style eye with brown iris color",
|
||||
"tags": "eye,classic,cartoon,brown,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_cartoon_hazel.png",
|
||||
"ppm_filename": "classic_cartoon_hazel.ppm",
|
||||
"name": "Classic Cartoon Eye - Hazel",
|
||||
"description": "A cartoon style eye with hazel iris color",
|
||||
"tags": "eye,classic,cartoon,hazel,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_cartoon_grey.png",
|
||||
"ppm_filename": "classic_cartoon_grey.ppm",
|
||||
"name": "Classic Cartoon Eye - Grey",
|
||||
"description": "A cartoon style eye with grey iris color",
|
||||
"tags": "eye,classic,cartoon,grey,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
},
|
||||
{
|
||||
"filename": "classic_cartoon_amber.png",
|
||||
"ppm_filename": "classic_cartoon_amber.ppm",
|
||||
"name": "Classic Cartoon Eye - Amber",
|
||||
"description": "A cartoon style eye with amber iris color",
|
||||
"tags": "eye,classic,cartoon,amber,medium",
|
||||
"width": 200,
|
||||
"height": 200
|
||||
}
|
||||
]
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
backend:
|
||||
build:
|
||||
context: ./backend
|
||||
dockerfile: Dockerfile
|
||||
container_name: paintplus-backend-dev
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
- ./data:/app/data
|
||||
- ./backend:/app
|
||||
environment:
|
||||
- DATABASE_URL=sqlite:///./data/ai_photo_edit.db
|
||||
- SECRET_KEY=dev-secret-key-change-in-production
|
||||
- AI_PROVIDER=${AI_PROVIDER:-mock}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- STABILITY_API_KEY=${STABILITY_API_KEY:-}
|
||||
- CORS_ORIGINS=http://localhost:5173,http://localhost:3000,http://localhost
|
||||
command: uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- paintplus-network
|
||||
|
||||
frontend:
|
||||
build:
|
||||
context: ./frontend
|
||||
dockerfile: Dockerfile.dev
|
||||
container_name: paintplus-frontend-dev
|
||||
ports:
|
||||
- "5173:5173"
|
||||
volumes:
|
||||
- ./frontend:/app
|
||||
- /app/node_modules
|
||||
environment:
|
||||
- VITE_API_BASE_URL=http://localhost:8000
|
||||
depends_on:
|
||||
- backend
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- paintplus-network
|
||||
|
||||
networks:
|
||||
paintplus-network:
|
||||
driver: bridge
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
# =============================================================================
|
||||
# EditmaskwithAI — GPU Docker Compose (NVIDIA CUDA)
|
||||
#
|
||||
# ── PREREQUISITES ─────────────────────────────────────────────────────────────
|
||||
#
|
||||
# 1. NVIDIA driver ≥ 525 installed on the host
|
||||
# Check: nvidia-smi
|
||||
#
|
||||
# 2. nvidia-container-toolkit installed and configured:
|
||||
# (Ubuntu/Debian)
|
||||
# curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
|
||||
# | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-ctk.gpg
|
||||
# curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
|
||||
# | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-ctk.gpg] https://#g' \
|
||||
# | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
|
||||
# sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
|
||||
# sudo nvidia-ctk runtime configure --runtime=docker
|
||||
# sudo systemctl restart docker
|
||||
#
|
||||
# (RHEL/Fedora/Rocky)
|
||||
# sudo dnf install -y nvidia-container-toolkit
|
||||
# sudo nvidia-ctk runtime configure --runtime=docker
|
||||
# sudo systemctl restart docker
|
||||
#
|
||||
# 3. Verify GPU access in Docker:
|
||||
# docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi
|
||||
#
|
||||
# ── QUICK START ───────────────────────────────────────────────────────────────
|
||||
#
|
||||
# docker compose -f docker-compose.gpu.yml up --build
|
||||
# Then open: http://localhost:3080
|
||||
#
|
||||
# ── OLDER DOCKER SETUPS (docker-compose v1 / nvidia-docker2) ─────────────────
|
||||
#
|
||||
# If you installed nvidia-docker2 (older approach) instead of nvidia-container-toolkit,
|
||||
# replace the 'deploy:' block below with:
|
||||
#
|
||||
# runtime: nvidia
|
||||
# environment:
|
||||
# - NVIDIA_VISIBLE_DEVICES=all
|
||||
# - NVIDIA_DRIVER_CAPABILITIES=compute,utility
|
||||
#
|
||||
# ── GPU TIER AUTO-SELECTION ───────────────────────────────────────────────────
|
||||
#
|
||||
# ≥16 GB VRAM → SDXL (best quality)
|
||||
# 8–16 GB → SDXL
|
||||
# 4–8 GB → Stable Diffusion 2.x
|
||||
# 2–4 GB → Stable Diffusion 1.5 (older GPUs: GTX 970/1060/RX 580)
|
||||
# <2 GB → SD 1.5 + CPU offload (very slow — consider a remote provider)
|
||||
#
|
||||
# ── AMD ROCm ──────────────────────────────────────────────────────────────────
|
||||
#
|
||||
# Swap the base image in Dockerfile.gpu:
|
||||
# FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
|
||||
# → FROM rocm/pytorch:rocm6.0_ubuntu22.04_py3.9_pytorch_2.1.0
|
||||
# Remove the 'driver: nvidia' line and add: device_ids: ['0']
|
||||
#
|
||||
# =============================================================================
|
||||
|
||||
services:
|
||||
app:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.gpu
|
||||
args:
|
||||
# Increment BUILDID to force pip layers to re-run without full --no-cache:
|
||||
# BUILDID=$(date +%s) docker compose -f docker-compose.gpu.yml up --build
|
||||
BUILDID: ${BUILDID:-1}
|
||||
container_name: paintplus
|
||||
ports:
|
||||
- "${PORT:-3080}:8000"
|
||||
volumes:
|
||||
# Persistent project data
|
||||
- ./data:/app/data
|
||||
# HuggingFace model cache — bind mount so models can be pre-downloaded on the host.
|
||||
# If container DNS is blocked, run ./prefetch-models.sh on the host first —
|
||||
# it downloads BEN2/BiRefNet-HR (and optionally SDXL with --sdxl) straight
|
||||
# into this directory, and the container picks them up on next start.
|
||||
# To free disk space: rm -rf ./data/hf_cache
|
||||
- ./data/hf_cache:/root/.cache/huggingface
|
||||
# Scripts (for exec access)
|
||||
- ./scripts:/scripts
|
||||
environment:
|
||||
# ── Local GPU (default for this compose) ────────────────────────────────
|
||||
- AI_PROVIDER=${AI_PROVIDER:-local_gpu}
|
||||
- AUTO_DOWNLOAD_MODELS=${AUTO_DOWNLOAD_MODELS:-true}
|
||||
|
||||
# ── Per-operation overrides (optional) ──────────────────────────────────
|
||||
# Leave blank to use AI_PROVIDER for all operations.
|
||||
# Example: use InvokeAI for inpaint, local GPU for everything else:
|
||||
# AI_PROVIDER_INPAINT=invokeai
|
||||
- AI_PROVIDER_INPAINT=${AI_PROVIDER_INPAINT:-}
|
||||
- AI_PROVIDER_TXT2IMG=${AI_PROVIDER_TXT2IMG:-}
|
||||
- AI_PROVIDER_IMG2IMG=${AI_PROVIDER_IMG2IMG:-}
|
||||
- AI_PROVIDER_OUTPAINT=${AI_PROVIDER_OUTPAINT:-}
|
||||
|
||||
# ── Remote/cloud providers (all optional) ────────────────────────────────
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- OPENAI_MODEL=${OPENAI_MODEL:-gpt-image-2}
|
||||
- OPENAI_EDIT_MODEL=${OPENAI_EDIT_MODEL:-gpt-image-2}
|
||||
- REPLICATE_API_KEY=${REPLICATE_API_KEY:-}
|
||||
- STABILITY_API_KEY=${STABILITY_API_KEY:-}
|
||||
|
||||
# ── InvokeAI / ComfyUI (running on another machine or container) ────────
|
||||
- INVOKEAI_URL=${INVOKEAI_URL:-}
|
||||
- INVOKEAI_DEFAULT_MODEL=${INVOKEAI_DEFAULT_MODEL:-flux-dev}
|
||||
- COMFYUI_URL=${COMFYUI_URL:-}
|
||||
- COMFYUI_DEFAULT_MODEL=${COMFYUI_DEFAULT_MODEL:-v1-5-pruned-emaonly.ckpt}
|
||||
|
||||
# ── HuggingFace model overrides (optional) ───────────────────────────────
|
||||
# Override the auto-selected model for any operation:
|
||||
# HF_MODEL_INPAINT=your-org/your-model
|
||||
- HF_MODEL_INPAINT=${HF_MODEL_INPAINT:-}
|
||||
- HF_MODEL_TXT2IMG=${HF_MODEL_TXT2IMG:-}
|
||||
- HF_MODEL_IMG2IMG=${HF_MODEL_IMG2IMG:-}
|
||||
- HF_TOKEN=${HF_TOKEN:-}
|
||||
|
||||
# ── App settings ─────────────────────────────────────────────────────────
|
||||
- DATABASE_URL=sqlite:///./data/ai_photo_edit.db
|
||||
- SECRET_KEY=${SECRET_KEY:-change-this-secret-key-in-production}
|
||||
- CORS_ORIGINS=*
|
||||
- AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true}
|
||||
- AUTO_DOWNLOAD_U2NET=${AUTO_DOWNLOAD_U2NET:-true}
|
||||
- BG_REMOVAL_MODEL=${BG_REMOVAL_MODEL:-ben2}
|
||||
|
||||
# ── NVIDIA GPU passthrough ────────────────────────────────────────────────
|
||||
# Requires nvidia-container-toolkit; see prerequisites at top of this file.
|
||||
# For older nvidia-docker2 setups, replace this block with:
|
||||
# runtime: nvidia
|
||||
# environment:
|
||||
# - NVIDIA_VISIBLE_DEVICES=all
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: 1
|
||||
capabilities: [gpu]
|
||||
|
||||
# DNS: try host resolver first (works on most networks including corporate/VPN),
|
||||
# fall back to Cloudflare then Google public resolvers.
|
||||
# If all three fail (Errno -3), your firewall is blocking port 53 UDP from Docker.
|
||||
# Fix on the host: sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT
|
||||
dns:
|
||||
- 1.1.1.1
|
||||
- 8.8.8.8
|
||||
- 8.8.4.4
|
||||
|
||||
restart: unless-stopped
|
||||
Vendored
+33
@@ -0,0 +1,33 @@
|
||||
services:
|
||||
app:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: paintplus
|
||||
ports:
|
||||
- "3080:8000"
|
||||
volumes:
|
||||
- ./data:/app/data
|
||||
- ./scripts:/scripts
|
||||
environment:
|
||||
- DATABASE_URL=sqlite:///./data/ai_photo_edit.db
|
||||
- SECRET_KEY=${SECRET_KEY:-change-this-secret-key-in-production}
|
||||
- AI_PROVIDER=${AI_PROVIDER:-mock}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- OPENAI_MODEL=${OPENAI_MODEL:-gpt-image-2}
|
||||
- OPENAI_EDIT_MODEL=${OPENAI_EDIT_MODEL:-gpt-image-2}
|
||||
- STABILITY_API_KEY=${STABILITY_API_KEY:-}
|
||||
- REPLICATE_API_KEY=${REPLICATE_API_KEY:-}
|
||||
- INVOKEAI_URL=${INVOKEAI_URL:-}
|
||||
- INVOKEAI_DEFAULT_MODEL=${INVOKEAI_DEFAULT_MODEL:-flux-dev}
|
||||
- COMFYUI_URL=${COMFYUI_URL:-}
|
||||
- COMFYUI_DEFAULT_MODEL=${COMFYUI_DEFAULT_MODEL:-v1-5-pruned-emaonly.ckpt}
|
||||
- CORS_ORIGINS=*
|
||||
# DNS servers for reliable external API access (Replicate, etc.)
|
||||
dns:
|
||||
- 8.8.8.8
|
||||
- 8.8.4.4
|
||||
restart: unless-stopped
|
||||
|
||||
volumes:
|
||||
data:
|
||||
+261
@@ -0,0 +1,261 @@
|
||||
# AI Provider Cost & Quality Comparison
|
||||
|
||||
## Provider Options for Inpainting/Image Editing
|
||||
|
||||
### 1. OpenAI DALL-E 2 ❌ (Not Recommended)
|
||||
**Current implementation uses this when `AI_PROVIDER=openai`**
|
||||
|
||||
**Pricing:**
|
||||
- $0.020 per image (1024x1024)
|
||||
- $0.018 per image (512x512)
|
||||
|
||||
**Quality:** ⭐⭐ (2/5)
|
||||
- Old model (2022)
|
||||
- Significantly lower quality than DALL-E 3
|
||||
- Cannot match ChatGPT web interface
|
||||
- Often produces artifacts
|
||||
|
||||
**Pros:**
|
||||
- Simple API
|
||||
- Fast responses
|
||||
|
||||
**Cons:**
|
||||
- Poor quality by modern standards
|
||||
- Limited to 1024x1024 max
|
||||
- No access to DALL-E 3 inpainting
|
||||
|
||||
**Verdict:** ❌ Don't use unless you need the cheapest option and quality doesn't matter
|
||||
|
||||
---
|
||||
|
||||
### 2. Stability AI (Stable Diffusion XL) ✅ (Good Choice)
|
||||
**Direct API to Stability AI**
|
||||
|
||||
**Pricing:**
|
||||
- Credits-based system
|
||||
- ~$0.010 per image (512x512)
|
||||
- ~$0.040 per image (1024x1024)
|
||||
- Must buy credit packs ($10 minimum = 1000 credits)
|
||||
|
||||
**Quality:** ⭐⭐⭐⭐ (4/5)
|
||||
- Excellent inpainting quality
|
||||
- Good at following prompts
|
||||
- Natural-looking results
|
||||
- Well-suited for photo editing
|
||||
|
||||
**Pros:**
|
||||
- Built specifically for inpainting
|
||||
- Good quality-to-cost ratio
|
||||
- Reliable API
|
||||
- Fast generation (15-30 seconds)
|
||||
|
||||
**Cons:**
|
||||
- Requires credit purchase upfront
|
||||
- Limited to SDXL models
|
||||
- Less flexible than Replicate
|
||||
|
||||
**Verdict:** ✅ Best balance of quality and cost for direct API
|
||||
|
||||
---
|
||||
|
||||
### 3. Replicate ⭐ (Most Flexible)
|
||||
**API marketplace with multiple models**
|
||||
|
||||
**Pricing:** Pay-per-second of GPU time
|
||||
- SDXL Inpainting: ~$0.0023/sec (~$0.01-0.03 per image)
|
||||
- Kandinsky 2.2: ~$0.0023/sec (~$0.01-0.02 per image)
|
||||
- LaMa (removal): ~$0.0005/sec (~$0.002 per image)
|
||||
- Varies by model and parameters
|
||||
|
||||
**Quality:** ⭐⭐⭐⭐⭐ (5/5 - depends on model choice)
|
||||
- Access to multiple models
|
||||
- Can choose best model for each use case
|
||||
- Community models available
|
||||
- Often better than Stability direct
|
||||
|
||||
**Pros:**
|
||||
- Multiple models to choose from
|
||||
- Pay only for what you use (no minimums)
|
||||
- Can use free models
|
||||
- New models added regularly
|
||||
- Fine-tuned models available
|
||||
|
||||
**Cons:**
|
||||
- More complex to implement
|
||||
- Pricing varies by model
|
||||
- Need to understand different models
|
||||
|
||||
**Best Models on Replicate:**
|
||||
- **SDXL Inpainting**: General purpose, excellent quality
|
||||
- **LaMa**: Best for object removal
|
||||
- **Kandinsky 2.2**: Good alternative to SDXL
|
||||
- **ControlNet Inpainting**: More control over results
|
||||
|
||||
**Verdict:** ⭐ Most flexible, best value if you implement multiple models
|
||||
|
||||
---
|
||||
|
||||
### 4. Local Models (Self-Hosted) 💰 (Best Quality, No Per-Use Cost)
|
||||
|
||||
**Pricing:**
|
||||
- $0 per image after setup
|
||||
- Requires GPU (RTX 3060 12GB minimum, RTX 4090 ideal)
|
||||
- Cloud GPU: $0.30-$1.00/hour (RunPod, Vast.ai)
|
||||
|
||||
**Quality:** ⭐⭐⭐⭐⭐ (5/5)
|
||||
- Best possible quality
|
||||
- Full control over model selection
|
||||
- Can use latest open-source models
|
||||
- No API limitations
|
||||
|
||||
**Setup Costs:**
|
||||
- GPU hardware: $300-$2000
|
||||
- OR Cloud GPU rental: $0.30-$1.00/hour
|
||||
|
||||
**Pros:**
|
||||
- Unlimited usage once set up
|
||||
- Best quality available
|
||||
- Complete privacy
|
||||
- No API rate limits
|
||||
- Can fine-tune models
|
||||
|
||||
**Cons:**
|
||||
- Requires GPU or cloud rental
|
||||
- More complex setup
|
||||
- Slower than cloud APIs (if CPU only)
|
||||
|
||||
**Verdict:** 💰 Best long-term if you have GPU or high volume
|
||||
|
||||
---
|
||||
|
||||
## Stability AI vs Replicate: What's the Difference?
|
||||
|
||||
### Stability AI (stability.ai)
|
||||
**What it is:**
|
||||
- The company that created Stable Diffusion
|
||||
- Direct API to their hosted models
|
||||
- Official source
|
||||
|
||||
**Business Model:**
|
||||
- Buy credits upfront
|
||||
- Credits expire after 3 months
|
||||
- Official support
|
||||
- Guaranteed uptime SLA
|
||||
|
||||
**Models Available:**
|
||||
- Stable Diffusion XL
|
||||
- Stable Diffusion 1.5
|
||||
- Their official models only
|
||||
|
||||
---
|
||||
|
||||
### Replicate (replicate.com)
|
||||
**What it is:**
|
||||
- Marketplace/platform for running ML models
|
||||
- Hosts models from many sources
|
||||
- Pay-per-use GPU time
|
||||
|
||||
**Business Model:**
|
||||
- Pay only for GPU seconds used
|
||||
- No upfront purchase
|
||||
- No credits that expire
|
||||
- $0.01 minimum charge per prediction
|
||||
|
||||
**Models Available:**
|
||||
- Stability AI's models (SDXL, SD 1.5)
|
||||
- Community models
|
||||
- Fine-tuned variants
|
||||
- Specialized models (LaMa, ControlNet, etc.)
|
||||
- 100+ image generation models
|
||||
|
||||
**Think of it like:**
|
||||
- **Stability AI** = Buying directly from Apple
|
||||
- **Replicate** = App Store with many developers
|
||||
|
||||
---
|
||||
|
||||
## Cost Comparison Examples
|
||||
|
||||
### Scenario: 100 edits per month
|
||||
|
||||
| Provider | Cost per Image | Monthly Cost | Quality |
|
||||
|----------|---------------|--------------|---------|
|
||||
| DALL-E 2 | $0.020 | $2.00 | ⭐⭐ Poor |
|
||||
| Stability AI | $0.040 | $4.00 | ⭐⭐⭐⭐ Good |
|
||||
| Replicate (SDXL) | $0.025 | $2.50 | ⭐⭐⭐⭐⭐ Excellent |
|
||||
| Replicate (LaMa) | $0.002 | $0.20 | ⭐⭐⭐⭐ Good for removal |
|
||||
| Local GPU | $0.00 | $0.00* | ⭐⭐⭐⭐⭐ Best |
|
||||
|
||||
*Requires $500+ GPU or $0.30-1.00/hr cloud GPU
|
||||
|
||||
### Scenario: 1000 edits per month (Heavy use)
|
||||
|
||||
| Provider | Monthly Cost | Notes |
|
||||
|----------|--------------|-------|
|
||||
| DALL-E 2 | $20.00 | Not worth it |
|
||||
| Stability AI | $40.00 | Need $10-40 credit refills |
|
||||
| Replicate (SDXL) | $25.00 | Pay as you go |
|
||||
| Local GPU | $0.00 | GPU pays for itself after ~50K images |
|
||||
| Cloud GPU (RunPod) | $20-60 | Depends on uptime needed |
|
||||
|
||||
---
|
||||
|
||||
## Quality Rankings for Inpainting
|
||||
|
||||
**Best to Worst:**
|
||||
|
||||
1. **Local SDXL Inpainting** ⭐⭐⭐⭐⭐ (self-hosted)
|
||||
2. **Replicate SDXL Inpainting** ⭐⭐⭐⭐⭐
|
||||
3. **Stability AI SDXL** ⭐⭐⭐⭐
|
||||
4. **Replicate LaMa** ⭐⭐⭐⭐ (for removal only)
|
||||
5. **DALL-E 2** ⭐⭐ (outdated)
|
||||
|
||||
---
|
||||
|
||||
## Recommendation by Use Case
|
||||
|
||||
### Best for Testing/Development: Mock Provider
|
||||
- Cost: $0
|
||||
- Quality: N/A (returns original)
|
||||
- Use when: Building/testing UI
|
||||
|
||||
### Best for Low Volume (< 100/month): Replicate
|
||||
- Cost: ~$2.50/month
|
||||
- Quality: ⭐⭐⭐⭐⭐
|
||||
- No minimum purchase
|
||||
- Multiple model options
|
||||
|
||||
### Best for Medium Volume (100-1000/month): Replicate or Stability AI
|
||||
- Replicate: ~$25/month, more flexibility
|
||||
- Stability AI: ~$40/month, simpler API
|
||||
|
||||
### Best for High Volume (1000+/month): Local GPU or Cloud GPU
|
||||
- Unlimited usage
|
||||
- Best quality
|
||||
- Full control
|
||||
|
||||
### Best Overall Value: Replicate
|
||||
- No minimum purchase
|
||||
- Pay only for what you use
|
||||
- Best model selection
|
||||
- Easy to try multiple models
|
||||
|
||||
---
|
||||
|
||||
## My Recommendation
|
||||
|
||||
Start with **Replicate** because:
|
||||
|
||||
1. ✅ No upfront cost (vs Stability's $10 minimum)
|
||||
2. ✅ Better quality than DALL-E 2
|
||||
3. ✅ Can try multiple models to find what works
|
||||
4. ✅ Cheapest per-image for low-medium volume
|
||||
5. ✅ Can switch to Stability AI later if needed
|
||||
|
||||
**Next Steps:**
|
||||
- I can add Replicate support (30 min of work)
|
||||
- Test with SDXL Inpainting first
|
||||
- Try LaMa for object removal
|
||||
- Fall back to Stability if needed
|
||||
|
||||
Would you like me to add Replicate support?
|
||||
+369
@@ -0,0 +1,369 @@
|
||||
# Model Selection Guide for Body Parts and Editing Tasks
|
||||
|
||||
## Quick Reference: Best Models by Use Case
|
||||
|
||||
### Human Features (Faces, Hands, Bodies)
|
||||
|
||||
**Best Choice: `realistic-vision` (Replicate)**
|
||||
|
||||
```env
|
||||
AI_PROVIDER=replicate
|
||||
REPLICATE_API_KEY=your-key
|
||||
REPLICATE_MODEL=realistic-vision
|
||||
```
|
||||
|
||||
**Why:** Trained specifically on human anatomy and realistic photos. Handles difficult features like:
|
||||
- ✅ Hands (notoriously hard for AI)
|
||||
- ✅ Faces and facial features
|
||||
- ✅ Skin textures and tones
|
||||
- ✅ Body proportions
|
||||
- ✅ Portraits
|
||||
|
||||
**Examples:**
|
||||
- "Fix the hand position"
|
||||
- "Remove red eye"
|
||||
- "Smooth skin blemishes"
|
||||
- "Adjust facial expression"
|
||||
- "Fix fingers"
|
||||
|
||||
**Cost:** ~$0.020/image
|
||||
**Quality:** ⭐⭐⭐⭐⭐
|
||||
|
||||
---
|
||||
|
||||
### Object Removal
|
||||
|
||||
**Best Choice: `lama` (Replicate)**
|
||||
|
||||
```env
|
||||
AI_PROVIDER=replicate
|
||||
REPLICATE_MODEL=lama
|
||||
```
|
||||
|
||||
**Why:** Specifically designed for inpainting and object removal. Excellent at:
|
||||
- ✅ Removing objects cleanly
|
||||
- ✅ Filling in backgrounds naturally
|
||||
- ✅ Maintaining surrounding context
|
||||
- ✅ Fast and cheap
|
||||
|
||||
**Examples:**
|
||||
- "Remove the person"
|
||||
- "Delete the watermark"
|
||||
- "Erase the object"
|
||||
- "Clean up the background"
|
||||
|
||||
**Cost:** ~$0.002/image (cheapest!)
|
||||
**Quality:** ⭐⭐⭐⭐
|
||||
|
||||
---
|
||||
|
||||
### General Purpose Editing
|
||||
|
||||
**Best Choice: `sdxl-inpaint` (Replicate or Stability AI)**
|
||||
|
||||
```env
|
||||
# Option 1: Replicate
|
||||
AI_PROVIDER=replicate
|
||||
REPLICATE_MODEL=sdxl-inpaint
|
||||
|
||||
# Option 2: Stability AI Direct
|
||||
AI_PROVIDER=stability
|
||||
STABILITY_MODEL=sdxl
|
||||
```
|
||||
|
||||
**Why:** SDXL (Stable Diffusion XL) is the best all-around model for:
|
||||
- ✅ Landscapes and scenery
|
||||
- ✅ Objects and textures
|
||||
- ✅ Creative edits
|
||||
- ✅ Style changes
|
||||
- ✅ Adding elements
|
||||
|
||||
**Examples:**
|
||||
- "Change sky to sunset"
|
||||
- "Add flowers"
|
||||
- "Make it autumn"
|
||||
- "Replace with grass"
|
||||
|
||||
**Cost:**
|
||||
- Replicate: ~$0.025/image
|
||||
- Stability AI: ~$0.040/image
|
||||
|
||||
**Quality:** ⭐⭐⭐⭐⭐
|
||||
|
||||
---
|
||||
|
||||
## Detailed Comparison by Body Part
|
||||
|
||||
### Hands ✋
|
||||
|
||||
**Challenge:** Hands are the hardest thing for AI to generate correctly. Common issues:
|
||||
- Wrong number of fingers
|
||||
- Unnatural finger positions
|
||||
- Distorted proportions
|
||||
- Weird joints
|
||||
|
||||
**Best Models (in order):**
|
||||
|
||||
1. **Realistic Vision** (Replicate) - ⭐⭐⭐⭐⭐
|
||||
- Best overall for hands
|
||||
- Understands hand anatomy
|
||||
- Cost: ~$0.020/image
|
||||
|
||||
2. **SDXL Inpainting** (Replicate/Stability) - ⭐⭐⭐
|
||||
- Decent but less consistent
|
||||
- Cost: ~$0.025-0.040/image
|
||||
|
||||
3. **DALL-E 2** (OpenAI) - ⭐⭐
|
||||
- Often struggles with hands
|
||||
- Not recommended
|
||||
|
||||
**Tips for Better Hand Edits:**
|
||||
- Use detailed prompts: "realistic human hand with five fingers"
|
||||
- Add negative prompts if provider supports: "deformed, extra fingers, missing fingers"
|
||||
- Use Mode B (full image context) for better results
|
||||
- Consider editing in multiple passes if needed
|
||||
|
||||
---
|
||||
|
||||
### Faces 😊
|
||||
|
||||
**Challenge:** Faces need to look natural and maintain proper proportions
|
||||
|
||||
**Best Models:**
|
||||
|
||||
1. **Realistic Vision** (Replicate) - ⭐⭐⭐⭐⭐
|
||||
- Excellent for facial features
|
||||
- Natural skin textures
|
||||
- Good expression handling
|
||||
|
||||
2. **SDXL Inpainting** - ⭐⭐⭐⭐
|
||||
- Good for general facial edits
|
||||
- Better for style than realism
|
||||
|
||||
**Use Cases:**
|
||||
- Remove blemishes
|
||||
- Fix red eye
|
||||
- Adjust expressions
|
||||
- Change hair
|
||||
- Smooth wrinkles
|
||||
|
||||
---
|
||||
|
||||
### Full Body / Torso 🧍
|
||||
|
||||
**Best Model:** Realistic Vision
|
||||
|
||||
**Why:** Maintains body proportions and realistic anatomy
|
||||
|
||||
**Examples:**
|
||||
- "Fix the clothing wrinkles"
|
||||
- "Change shirt color to blue"
|
||||
- "Remove the stain"
|
||||
|
||||
---
|
||||
|
||||
### Hearts ♥️ (Decorative Elements)
|
||||
|
||||
**Best Model:** SDXL Inpainting
|
||||
|
||||
**Why:** Great for creative and decorative elements
|
||||
|
||||
**Examples:**
|
||||
- "Add heart shape"
|
||||
- "Draw a heart pattern"
|
||||
- "Replace with hearts"
|
||||
|
||||
---
|
||||
|
||||
## Auto-Selection Feature
|
||||
|
||||
The system automatically selects the best model based on your prompt:
|
||||
|
||||
### Keywords that trigger `realistic-vision`:
|
||||
- hand, hands, finger, fingers
|
||||
- face, facial, portrait, eyes, nose, mouth
|
||||
- body, person, human, skin, people
|
||||
- realistic, photo, photograph
|
||||
|
||||
### Keywords that trigger `lama` (removal):
|
||||
- remove, delete, erase, cleanup
|
||||
- disappear, hide, clear
|
||||
|
||||
### Default: `sdxl-inpaint`
|
||||
- Everything else uses SDXL for best general quality
|
||||
|
||||
**Example Auto-Selection:**
|
||||
```python
|
||||
# User prompt: "Fix the hand" → auto-selects realistic-vision
|
||||
# User prompt: "Remove the person" → auto-selects lama
|
||||
# User prompt: "Change to sunset" → auto-selects sdxl-inpaint
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Manual Model Override
|
||||
|
||||
### Via Environment Variable
|
||||
Set default model in `.env`:
|
||||
```env
|
||||
REPLICATE_MODEL=realistic-vision
|
||||
```
|
||||
|
||||
### Via API Request
|
||||
Override per-edit in the API:
|
||||
```json
|
||||
{
|
||||
"prompt": "Fix the hand",
|
||||
"ai_provider": "replicate",
|
||||
"ai_model": "realistic-vision",
|
||||
"mode": "A",
|
||||
...
|
||||
}
|
||||
```
|
||||
|
||||
### Via Frontend (Future Feature)
|
||||
Model selector dropdown in the UI.
|
||||
|
||||
---
|
||||
|
||||
## Cost Optimization Strategies
|
||||
|
||||
### For Low-Volume Users (< 100 edits/month)
|
||||
**Recommendation:** Use Replicate with auto-selection
|
||||
|
||||
**Why:**
|
||||
- No minimum purchase
|
||||
- Pay only for what you use
|
||||
- Auto-selects cheapest appropriate model
|
||||
|
||||
**Estimated Cost:** $1-3/month
|
||||
|
||||
---
|
||||
|
||||
### For Medium-Volume Users (100-1000 edits/month)
|
||||
**Recommendation:** Replicate or Stability AI
|
||||
|
||||
**Strategy:**
|
||||
- Use `lama` for removals ($0.002/image)
|
||||
- Use `realistic-vision` for humans ($0.020/image)
|
||||
- Use `sdxl-inpaint` for general ($0.025/image)
|
||||
|
||||
**Estimated Cost:** $10-30/month
|
||||
|
||||
---
|
||||
|
||||
### For High-Volume Users (1000+ edits/month)
|
||||
**Recommendation:** Consider local GPU or cloud GPU
|
||||
|
||||
**Why:**
|
||||
- No per-image cost
|
||||
- Best quality control
|
||||
- Privacy
|
||||
|
||||
**Setup:**
|
||||
- Local: RTX 3060+ GPU ($300-2000 one-time)
|
||||
- Cloud: RunPod/Vast.ai ($0.30-1.00/hour)
|
||||
|
||||
---
|
||||
|
||||
## Quality Comparison Table
|
||||
|
||||
| Use Case | DALL-E 2 | Stability SDXL | Replicate SDXL | Replicate Realistic | Replicate LaMa |
|
||||
|----------|----------|----------------|----------------|---------------------|----------------|
|
||||
| Hands | ⭐⭐ | ⭐⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐ |
|
||||
| Faces | ⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐ |
|
||||
| Bodies | ⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐ |
|
||||
| Objects | ⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | N/A |
|
||||
| Landscapes | ⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | N/A |
|
||||
| Removal | ⭐⭐ | ⭐⭐⭐ | ⭐⭐⭐⭐ | ⭐⭐⭐ | ⭐⭐⭐⭐⭐ |
|
||||
| Creative | ⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | ⭐ |
|
||||
|
||||
---
|
||||
|
||||
## Advanced Tips
|
||||
|
||||
### For Difficult Hands
|
||||
1. **Use Mode B** - Provides full image context
|
||||
2. **Be specific** - "realistic five-fingered hand in natural pose"
|
||||
3. **Multiple passes** - Fix gross errors first, then refine
|
||||
4. **Reference images** - Mode B helps AI understand the pose
|
||||
|
||||
### For Facial Features
|
||||
1. **High feather value** - 10-15px for smooth blending
|
||||
2. **Small selections** - Target specific features
|
||||
3. **Natural lighting** - Mention lighting in prompt
|
||||
|
||||
### For Body Parts
|
||||
1. **Maintain proportions** - Use Mode B for body context
|
||||
2. **Clothing context** - Include clothing description in prompt
|
||||
3. **Skin tone consistency** - Mention skin tone if needed
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting Common Issues
|
||||
|
||||
### "Hands have too many fingers"
|
||||
- **Solution:** Switch to `realistic-vision` model
|
||||
- **Prompt:** "realistic human hand with exactly five fingers"
|
||||
- **Try:** Multiple generations, pick best result
|
||||
|
||||
### "Face looks unnatural"
|
||||
- **Solution:** Use `realistic-vision` model
|
||||
- **Increase:** Feather value to 15-20px
|
||||
- **Try:** Mode B for better context
|
||||
|
||||
### "Removal leaves artifacts"
|
||||
- **Solution:** Use `lama` model (designed for removal)
|
||||
- **Alternative:** SDXL with prompt "clean background"
|
||||
|
||||
### "Colors don't match"
|
||||
- **Increase:** Feather value to 20-30px
|
||||
- **Try:** Mode B for better color context
|
||||
- **Prompt:** Include color description
|
||||
|
||||
---
|
||||
|
||||
## Quick Start Examples
|
||||
|
||||
### Example 1: Fix a Hand
|
||||
```json
|
||||
{
|
||||
"prompt": "realistic human hand with five fingers, natural pose",
|
||||
"ai_provider": "replicate",
|
||||
"ai_model": "realistic-vision",
|
||||
"mode": "B",
|
||||
"feather_px": 10
|
||||
}
|
||||
```
|
||||
|
||||
### Example 2: Remove an Object
|
||||
```json
|
||||
{
|
||||
"prompt": "remove the object, clean background",
|
||||
"ai_provider": "replicate",
|
||||
"ai_model": "lama",
|
||||
"mode": "A",
|
||||
"feather_px": 5
|
||||
}
|
||||
```
|
||||
|
||||
### Example 3: Change Sky
|
||||
```json
|
||||
{
|
||||
"prompt": "sunset sky with orange and pink clouds",
|
||||
"ai_provider": "replicate",
|
||||
"ai_model": "sdxl-inpaint",
|
||||
"mode": "A",
|
||||
"feather_px": 15
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
**For Body Parts:** Use `realistic-vision` (Replicate)
|
||||
**For Removal:** Use `lama` (Replicate)
|
||||
**For Everything Else:** Use `sdxl-inpaint` (Replicate or Stability)
|
||||
|
||||
**Let the auto-selection do its job** - it's optimized for these use cases!
|
||||
+300
@@ -0,0 +1,300 @@
|
||||
# Public Domain Carved Eye Sources
|
||||
|
||||
Where to find high-quality images of carved eyes from classical sculptures (all public domain).
|
||||
|
||||
---
|
||||
|
||||
## Best Museums with Public Domain Images
|
||||
|
||||
### 1. Metropolitan Museum of Art (CC0 Public Domain)
|
||||
|
||||
**Website:** https://www.metmuseum.org/art/collection
|
||||
|
||||
**Search tips:**
|
||||
- Search: "greek statue marble head"
|
||||
- Search: "roman portrait bust"
|
||||
- Filter: "Public Domain" only
|
||||
- Download: Click "Download" for high-resolution
|
||||
|
||||
**Great examples:**
|
||||
- Greek Kouros heads (Archaic period)
|
||||
- Roman portrait busts
|
||||
- Hellenistic marble sculptures
|
||||
|
||||
**Direct collections:**
|
||||
- Greek & Roman Art: https://www.metmuseum.org/art/collection/search#!?department=13
|
||||
- Filter by "Images" → "Public Domain"
|
||||
|
||||
---
|
||||
|
||||
### 2. Smithsonian Open Access (CC0)
|
||||
|
||||
**Website:** https://www.si.edu/openaccess
|
||||
|
||||
**Features:**
|
||||
- 3 million+ images
|
||||
- All CC0 (no copyright restrictions)
|
||||
- High-resolution downloads
|
||||
|
||||
**Search:**
|
||||
- "roman marble head"
|
||||
- "greek sculpture eyes"
|
||||
- "classical portrait bust"
|
||||
|
||||
**API available:** https://api.si.edu/openaccess/api/v1.0/
|
||||
|
||||
---
|
||||
|
||||
### 3. Getty Museum (Open Content)
|
||||
|
||||
**Website:** https://www.getty.edu/art/collection/
|
||||
|
||||
**Search tips:**
|
||||
- Filter: "Open Content Program"
|
||||
- Greek and Roman antiquities
|
||||
- High-resolution IIIF images
|
||||
|
||||
**Great for:**
|
||||
- Archaic Greek sculptures
|
||||
- Classical period heads
|
||||
- Detailed close-ups
|
||||
|
||||
---
|
||||
|
||||
### 4. Rijksmuseum (Public Domain)
|
||||
|
||||
**Website:** https://www.rijksmuseum.nl/en/rijksstudio
|
||||
|
||||
**Features:**
|
||||
- Rijksstudio (free download tool)
|
||||
- High-resolution images
|
||||
- Classical sculpture collection
|
||||
|
||||
---
|
||||
|
||||
### 5. British Museum (CC BY-NC-SA 4.0)
|
||||
|
||||
**Website:** https://www.britishmuseum.org/collection
|
||||
|
||||
**Note:** Some restrictions, but many images free for non-commercial use
|
||||
|
||||
**Great for:**
|
||||
- Egyptian carved eyes
|
||||
- Greek marble heads
|
||||
- Roman portraits
|
||||
|
||||
---
|
||||
|
||||
### 6. Louvre Collections
|
||||
|
||||
**Website:** https://collections.louvre.fr/en/
|
||||
|
||||
**Search:** "sculpture greek head" or "sculpture roman portrait"
|
||||
|
||||
**Note:** Check individual image licenses
|
||||
|
||||
---
|
||||
|
||||
## How to Find the Perfect Eyes
|
||||
|
||||
### Search Strategy
|
||||
|
||||
1. **Search for heads/busts, not full statues:**
|
||||
- "greek marble head"
|
||||
- "roman portrait bust"
|
||||
- "classical sculpture face"
|
||||
|
||||
2. **Specific periods:**
|
||||
- "archaic greek kouros" (serene, stylized)
|
||||
- "classical greek sculpture" (idealized, peaceful)
|
||||
- "hellenistic sculpture" (emotional, dramatic)
|
||||
- "roman portrait" (realistic, wise)
|
||||
|
||||
3. **Look for close-ups:**
|
||||
- Museums often provide detail shots
|
||||
- Check "zoom" or "IIIF viewer" options
|
||||
|
||||
---
|
||||
|
||||
## Recommended Starting Collection
|
||||
|
||||
### Serene/Peaceful Eyes
|
||||
|
||||
**Greek Classical Period (450-400 BCE):**
|
||||
- Doryphoros (Spear Bearer) type
|
||||
- Athena heads
|
||||
- Apollo statues
|
||||
- Smooth, idealized features
|
||||
- Almond-shaped eyes
|
||||
- Minimal lid detail
|
||||
|
||||
**Best sources:** Met Museum, Getty
|
||||
|
||||
---
|
||||
|
||||
### Fierce/Intense Eyes
|
||||
|
||||
**Hellenistic Period (323-31 BCE):**
|
||||
- Alexander the Great portraits
|
||||
- Dying Gaul
|
||||
- Laocoon group
|
||||
- Dramatic expressions
|
||||
- Deep-set eyes
|
||||
- Strong brow ridges
|
||||
|
||||
**Best sources:** Smithsonian, British Museum
|
||||
|
||||
---
|
||||
|
||||
### Wise/Aged Eyes
|
||||
|
||||
**Roman Republican Period:**
|
||||
- Senator portraits
|
||||
- Veristic portraits
|
||||
- Realistic aging details
|
||||
- Detailed wrinkles
|
||||
- Saggy eyelids
|
||||
- Life-like features
|
||||
|
||||
**Best sources:** Met Museum, Getty
|
||||
|
||||
---
|
||||
|
||||
### Stylized/Archaic Eyes
|
||||
|
||||
**Greek Archaic Period (700-480 BCE):**
|
||||
- Kouros statues
|
||||
- Kore statues
|
||||
- Almond-shaped
|
||||
- Simplified forms
|
||||
- "Archaic smile"
|
||||
- Clean, simple carving
|
||||
|
||||
**Best sources:** Getty, Met Museum
|
||||
|
||||
---
|
||||
|
||||
## How to Download and Crop
|
||||
|
||||
### Step 1: Find the Statue
|
||||
|
||||
Example: Met Museum
|
||||
1. Go to https://www.metmuseum.org/art/collection
|
||||
2. Search: "roman portrait marble"
|
||||
3. Filter: Public Domain only
|
||||
4. Click on a good example
|
||||
|
||||
### Step 2: Download High-Res
|
||||
|
||||
1. Click "Download" button
|
||||
2. Choose largest size (usually 4000px+)
|
||||
3. Save to your computer
|
||||
|
||||
### Step 3: Crop the Eyes
|
||||
|
||||
Use any image editor (Photoshop, GIMP, etc.):
|
||||
|
||||
1. Open the full statue image
|
||||
2. Zoom in on one eye
|
||||
3. Crop just the eye area:
|
||||
- Include: eyeball, eyelids, tear duct, socket
|
||||
- Leave some surrounding area for context
|
||||
- Square or slightly rectangular crop
|
||||
|
||||
4. Save as PNG:
|
||||
- `greek_serene_left.png`
|
||||
- `roman_fierce_right.png`
|
||||
- etc.
|
||||
|
||||
5. Repeat for other eye (if different)
|
||||
|
||||
### Step 4: Organize
|
||||
|
||||
Place cropped eyes in:
|
||||
```
|
||||
./backend/scripts/seed_data/eyes/
|
||||
├── greek_serene_left.png
|
||||
├── greek_serene_right.png
|
||||
├── roman_fierce_left.png
|
||||
├── roman_fierce_right.png
|
||||
├── greek_peaceful_left.png
|
||||
└── ...
|
||||
```
|
||||
|
||||
### Step 5: Run Seed Script
|
||||
|
||||
```bash
|
||||
cd backend
|
||||
python scripts/seed_eye_catalog.py
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Recommended Starting Collection (10 Eyes)
|
||||
|
||||
To start, get these 10 eyes:
|
||||
|
||||
### Greek Classical (Serene)
|
||||
1. Left eye - Greek marble head
|
||||
2. Right eye - Greek marble head
|
||||
|
||||
### Greek Archaic (Stylized/Peaceful)
|
||||
3. Left eye - Kouros statue
|
||||
4. Right eye - Kouros statue
|
||||
|
||||
### Hellenistic (Fierce/Dramatic)
|
||||
5. Left eye - Alexander portrait
|
||||
6. Right eye - Alexander portrait
|
||||
|
||||
### Roman Republican (Wise/Aged)
|
||||
7. Left eye - Roman senator bust
|
||||
8. Right eye - Roman senator bust
|
||||
|
||||
### Roman Imperial (Powerful)
|
||||
9. Left eye - Emperor portrait
|
||||
10. Right eye - Emperor portrait
|
||||
|
||||
This gives you 5 emotional ranges × 2 eyes = 10 eyes to start!
|
||||
|
||||
---
|
||||
|
||||
## Quick Links
|
||||
|
||||
- **Met Museum Collection:** https://www.metmuseum.org/art/collection/search#!?department=13&showOnly=openAccess
|
||||
- **Smithsonian Open Access:** https://www.si.edu/openaccess
|
||||
- **Getty Open Content:** https://www.getty.edu/about/whatwedo/opencontent.html
|
||||
- **Rijksmuseum API:** https://data.rijksmuseum.nl/object-metadata/api/
|
||||
|
||||
---
|
||||
|
||||
## Legal Notes
|
||||
|
||||
- **CC0/Public Domain:** Use freely for any purpose
|
||||
- **CC BY:** Must credit the source
|
||||
- **CC BY-NC:** Non-commercial use only
|
||||
- **Always check** individual image licenses
|
||||
|
||||
For commercial carving business, stick to **CC0** or **Public Domain** images.
|
||||
|
||||
---
|
||||
|
||||
## Tips for Best Results
|
||||
|
||||
1. **High resolution:** Download largest size available (2000px+ minimum)
|
||||
2. **Good lighting:** Look for evenly lit photographs
|
||||
3. **Straight-on angle:** Avoid extreme angles
|
||||
4. **Clear detail:** Can you see the eyelid lines clearly?
|
||||
5. **Minimal damage:** Choose well-preserved sculptures
|
||||
|
||||
---
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. Browse the museums above
|
||||
2. Download 10-20 good eye examples
|
||||
3. Crop them in an image editor
|
||||
4. Place in `backend/scripts/seed_data/eyes/`
|
||||
5. Run the seed script
|
||||
6. Your catalog is ready!
|
||||
|
||||
**You'll have a library of proven carved eyes from master sculptors spanning 2000+ years!**
|
||||
Vendored
+350
@@ -0,0 +1,350 @@
|
||||
# Quick Start Guide
|
||||
|
||||
## How to Choose the Right AI Model
|
||||
|
||||
### For Body Parts (Hands, Faces, Bodies)
|
||||
|
||||
Use **Replicate with `realistic-vision`** model:
|
||||
|
||||
```env
|
||||
AI_PROVIDER=replicate
|
||||
REPLICATE_API_KEY=your-key-here
|
||||
REPLICATE_MODEL=realistic-vision
|
||||
```
|
||||
|
||||
**Why:** This model is specifically trained on human anatomy and handles difficult features like:
|
||||
- ✅ Hands (even complex finger positions)
|
||||
- ✅ Faces and expressions
|
||||
- ✅ Skin textures
|
||||
- ✅ Body proportions
|
||||
|
||||
**Cost:** ~$0.020/image
|
||||
|
||||
### For Removing Objects
|
||||
|
||||
Use **Replicate with `lama`** model:
|
||||
|
||||
```env
|
||||
AI_PROVIDER=replicate
|
||||
REPLICATE_MODEL=lama
|
||||
```
|
||||
|
||||
**Why:** Designed specifically for inpainting and removal
|
||||
**Cost:** ~$0.002/image (cheapest!)
|
||||
|
||||
### For General Edits (Landscapes, Objects, Creative)
|
||||
|
||||
Use **Replicate with `sdxl-inpaint`** model (default):
|
||||
|
||||
```env
|
||||
AI_PROVIDER=replicate
|
||||
REPLICATE_MODEL=sdxl-inpaint
|
||||
```
|
||||
|
||||
**Cost:** ~$0.025/image
|
||||
|
||||
---
|
||||
|
||||
## Auto-Model Selection
|
||||
|
||||
The system automatically picks the best model based on your prompt:
|
||||
|
||||
| Your Prompt | Auto-Selected Model | Why |
|
||||
|-------------|-------------------|-----|
|
||||
| "Fix the hand" | realistic-vision | Detects "hand" keyword |
|
||||
| "Remove person" | lama | Detects "remove" keyword |
|
||||
| "Change sky to sunset" | sdxl-inpaint | General purpose default |
|
||||
|
||||
**You don't need to manually specify models** - the auto-selection is optimized for quality and cost!
|
||||
|
||||
---
|
||||
|
||||
## Patch Library: Save and Reuse Parts
|
||||
|
||||
### What is the Patch Library?
|
||||
|
||||
A library where you can save image patches (regions) and reuse them across different images.
|
||||
|
||||
**Use Cases:**
|
||||
- Save a well-generated hand to reuse later
|
||||
- Save a perfect face for multiple photos
|
||||
- Build a collection of good body parts
|
||||
- Save textures, objects, or backgrounds
|
||||
- Reuse AI-generated elements that came out great
|
||||
|
||||
### How to Save a Patch
|
||||
|
||||
#### Option 1: Save AI-Generated Result
|
||||
|
||||
After an AI edit completes:
|
||||
|
||||
```bash
|
||||
POST /patches/
|
||||
{
|
||||
"name": "Perfect Hand",
|
||||
"description": "Well-formed left hand, palm up",
|
||||
"source_type": "ai_generated",
|
||||
"source_edit_id": 123,
|
||||
"category": "hand",
|
||||
"tags": "left, palm, realistic"
|
||||
}
|
||||
```
|
||||
|
||||
This saves the AI-generated output (`patch_out.png`) to your library.
|
||||
|
||||
#### Option 2: Save Manual Selection
|
||||
|
||||
Select any region from your current image:
|
||||
|
||||
```bash
|
||||
POST /patches/
|
||||
{
|
||||
"name": "Good Face",
|
||||
"description": "Frontal face with good lighting",
|
||||
"source_type": "manual_selection",
|
||||
"source_project_id": 456,
|
||||
"bbox": {"x": 100, "y": 100, "width": 200, "height": 200},
|
||||
"category": "face",
|
||||
"tags": "front, smile, female"
|
||||
}
|
||||
```
|
||||
|
||||
This saves whatever is currently in that region of your image.
|
||||
|
||||
#### Option 3: Import from File
|
||||
|
||||
Upload an external image:
|
||||
|
||||
```bash
|
||||
POST /patches/
|
||||
FormData:
|
||||
name: "Downloaded Hand"
|
||||
source_type: "imported"
|
||||
file: [uploaded PNG file]
|
||||
category: "hand"
|
||||
```
|
||||
|
||||
### How to Apply a Saved Patch
|
||||
|
||||
```bash
|
||||
POST /patches/apply
|
||||
{
|
||||
"project_id": 789,
|
||||
"patch_id": 123,
|
||||
"bbox": {"x": 300, "y": 400, "width": 200, "height": 200},
|
||||
"feather_px": 10
|
||||
}
|
||||
```
|
||||
|
||||
This places the saved patch at the specified location in your image.
|
||||
|
||||
### Browse Your Patch Library
|
||||
|
||||
```bash
|
||||
# List all patches
|
||||
GET /patches/
|
||||
|
||||
# Filter by category
|
||||
GET /patches/?category=hand
|
||||
|
||||
# Filter by tags
|
||||
GET /patches/?tags=realistic
|
||||
|
||||
# Get specific patch
|
||||
GET /patches/123
|
||||
|
||||
# Get patch image
|
||||
GET /patches/123/image
|
||||
|
||||
# Get patch thumbnail
|
||||
GET /patches/123/image?thumbnail=true
|
||||
```
|
||||
|
||||
### Organize Your Patches
|
||||
|
||||
**Categories:**
|
||||
- `hand` - Hand images
|
||||
- `face` - Facial features
|
||||
- `body` - Body parts
|
||||
- `object` - Objects and items
|
||||
- `texture` - Textures and patterns
|
||||
- `background` - Backgrounds and scenery
|
||||
|
||||
**Tags:** Comma-separated keywords for searching
|
||||
- "left, palm, realistic"
|
||||
- "front, smile, female"
|
||||
- "five fingers, open hand"
|
||||
|
||||
---
|
||||
|
||||
## Complete Workflow Example
|
||||
|
||||
### Scenario: Fix hands in a portrait photo
|
||||
|
||||
**Step 1: Create project and upload image**
|
||||
```bash
|
||||
POST /projects/ {"name": "Portrait Edit"}
|
||||
POST /projects/1/upload [upload photo]
|
||||
```
|
||||
|
||||
**Step 2: Try to fix the hand with AI**
|
||||
```bash
|
||||
POST /edits/projects/1/fix
|
||||
{
|
||||
"prompt": "realistic human hand with five fingers, natural pose",
|
||||
"mode": "B", # Use full image for context
|
||||
"selection_type": "rectangle",
|
||||
"bbox": {"x": 200, "y": 300, "width": 150, "height": 200},
|
||||
"feather_px": 10
|
||||
}
|
||||
```
|
||||
|
||||
The system auto-selects `realistic-vision` model because prompt mentions "hand".
|
||||
|
||||
**Step 3: If result is good, save it for later**
|
||||
```bash
|
||||
POST /patches/
|
||||
{
|
||||
"name": "Good Left Hand",
|
||||
"source_type": "ai_generated",
|
||||
"source_edit_id": 1,
|
||||
"category": "hand",
|
||||
"tags": "left, natural, realistic, five fingers"
|
||||
}
|
||||
```
|
||||
|
||||
**Step 4: Use saved hand on another photo**
|
||||
```bash
|
||||
# On a different project
|
||||
POST /patches/apply
|
||||
{
|
||||
"project_id": 2,
|
||||
"patch_id": 1,
|
||||
"bbox": {"x": 150, "y": 250, "width": 150, "height": 200},
|
||||
"feather_px": 15
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cost Comparison
|
||||
|
||||
### Example: Fixing 10 hands in different photos
|
||||
|
||||
**Option A: Generate each hand with AI**
|
||||
- 10 edits × $0.020 = **$0.20**
|
||||
|
||||
**Option B: Generate one good hand, save it, reuse it**
|
||||
- 1 AI generation: $0.020
|
||||
- 9 patch applications: $0.00 (no AI cost)
|
||||
- **Total: $0.020** (90% savings!)
|
||||
|
||||
### When to Use Saved Patches vs AI
|
||||
|
||||
**Use Saved Patches When:**
|
||||
- You have a perfect result you want to reuse
|
||||
- Same angle/lighting/style needed
|
||||
- Want to maintain consistency across images
|
||||
- Want to avoid AI generation costs
|
||||
|
||||
**Use AI Generation When:**
|
||||
- Need unique/different result each time
|
||||
- Different angle or perspective needed
|
||||
- Want variation and creativity
|
||||
- Patch doesn't fit the context
|
||||
|
||||
---
|
||||
|
||||
## Pro Tips
|
||||
|
||||
### Building a Good Patch Library
|
||||
|
||||
1. **Save your best AI results** - When AI generates something great, save it immediately
|
||||
2. **Organize with categories** - Use consistent categories for easy finding
|
||||
3. **Tag descriptively** - Include orientation (left/right), pose, lighting, etc.
|
||||
4. **Create variations** - Save multiple versions of common needs (left hand, right hand, etc.)
|
||||
5. **Build gradually** - Your library becomes more valuable over time
|
||||
|
||||
### Maximizing Quality
|
||||
|
||||
1. **For hands:** Always use `realistic-vision` model or save good results
|
||||
2. **For faces:** Use Mode B (full image context) for better matching
|
||||
3. **Use high feather values** (15-20px) when applying saved patches
|
||||
4. **Test positioning** before finalizing - patches work best when lighting/angle matches
|
||||
|
||||
### Saving Money
|
||||
|
||||
1. **Build a patch library** of common needs
|
||||
2. **Use `lama` for removals** instead of expensive models
|
||||
3. **Let auto-selection work** - it picks the cheapest appropriate model
|
||||
4. **Reuse successful patches** instead of regenerating
|
||||
|
||||
---
|
||||
|
||||
## API Quick Reference
|
||||
|
||||
```bash
|
||||
# List available patches
|
||||
GET /patches/
|
||||
|
||||
# Get patch details
|
||||
GET /patches/{id}
|
||||
|
||||
# Get patch image
|
||||
GET /patches/{id}/image
|
||||
GET /patches/{id}/image?thumbnail=true
|
||||
|
||||
# Create patch from AI edit
|
||||
POST /patches/
|
||||
{
|
||||
"name": "My Patch",
|
||||
"source_type": "ai_generated",
|
||||
"source_edit_id": 123,
|
||||
"category": "hand"
|
||||
}
|
||||
|
||||
# Create patch from manual selection
|
||||
POST /patches/
|
||||
{
|
||||
"name": "My Patch",
|
||||
"source_type": "manual_selection",
|
||||
"source_project_id": 456,
|
||||
"bbox": {"x": 100, "y": 100, "width": 200, "height": 200}
|
||||
}
|
||||
|
||||
# Apply saved patch
|
||||
POST /patches/apply
|
||||
{
|
||||
"project_id": 789,
|
||||
"patch_id": 123,
|
||||
"bbox": {"x": 300, "y": 400, "width": 200, "height": 200},
|
||||
"feather_px": 10
|
||||
}
|
||||
|
||||
# Delete patch
|
||||
DELETE /patches/{id}
|
||||
|
||||
# Update patch metadata
|
||||
PUT /patches/{id}
|
||||
{
|
||||
"name": "Updated Name",
|
||||
"tags": "new, tags",
|
||||
"category": "hand"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
✅ **For hands/faces/bodies:** Use `realistic-vision` model
|
||||
✅ **For removal:** Use `lama` model
|
||||
✅ **For general edits:** Use `sdxl-inpaint` (default)
|
||||
✅ **Auto-selection works great** - just write natural prompts
|
||||
✅ **Save good AI results** to patch library for reuse
|
||||
✅ **Save manual selections** from any image
|
||||
✅ **Reuse patches across images** to save money and maintain consistency
|
||||
|
||||
**You now have the best of both worlds:**
|
||||
- AI generation when you need something new
|
||||
- Saved patches when you need consistency or want to save money
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user