iopaint, ai-gpu: interactive model selection, auto-pull, fixes

iopaint: - Add model selection menu (10 choices) with CPU/GPU/SD tiers and size/use-case descriptions shown at install time - Fix volume mount: ./models:/root/.cache (was only /root/.cache/iopaint) — now persists both torch hub cache (LaMa) and HuggingFace cache (SD/PowerPaint) - Refactor compose to use ${MODEL} and ${DEVICE} env vars so switching models only requires editing .env + restart, no compose file edit needed - Add PowerPaint-V2-filling and SD 1.5 inpainting as explicit menu choices for text-guided object replacement - Update header and README to document all three use cases (erase, fill, replace) and note that IOPaint is local-only (cannot use a remote GPU) ai-gpu: - Add Ollama model selection menu (8 models, multi-select with sizes/descriptions) defaulting to llama3.2:3b + nomic-embed-text - Auto-pull selected Ollama models immediately after LLM stack starts - Add InvokeAI starter model selection (SD 1.5 / SDXL Turbo / SDXL Base / skip) - Queue InvokeAI model download via REST API (POST /api/v2/models/install) with fallback instructions if the API is unavailable - Add HuggingFace token prompt; stored as HUGGING_FACE_HUB_TOKEN in image-gen .env - Wire SearXNG into Open WebUI via ENABLE_RAG_WEB_SEARCH + SEARXNG_QUERY_URL in llm .env - Default start choice is now 2 (portal + LLM + Ollama pull) so the stack is ready to use immediately after install https://claude.ai/code/session_01JEu7LgCWXKhXo18MeYFRZp
2026-06-09 03:57:01 +00:00
parent 5c394815b4
commit f9013287d9
2 changed files with 391 additions and 156 deletions
@@ -206,19 +206,94 @@ install_ai_gpu() {
    if [ "$DRY_RUN" = true ]; then
        echo "[DRY-RUN] Would clone $REPO_URL to $REPO_DIR"
        echo "[DRY-RUN] Would create stacks: image-gen (InvokeAI:9090), llm (Ollama:11434 + OpenWebUI:3000), portal (8080)"
-        echo "[DRY-RUN] Would write .env files and patch portal volume paths to $ACTUAL_HOME/docker"
-        echo "[DRY-RUN] Would configure Caddy for portal (localai) and InvokeAI (images)"
+        echo "[DRY-RUN] Would prompt for Ollama and InvokeAI model selection"
+        echo "[DRY-RUN] Would auto-pull selected Ollama models after LLM stack starts"
+        echo "[DRY-RUN] Would queue InvokeAI starter model via REST API"
        return 0
    fi

-    # TZ — repo defaults to America/New_York, replace with user preference
+    # ── Timezone ──────────────────────────────────────────────────────────────
    local TZ_VAL="${SITE_TZ:-UTC}"
    prompt_text "Timezone (e.g. America/New_York) [$TZ_VAL]:" "$TZ_VAL" TZ_VAL
    TZ_VAL="${TZ_VAL:-UTC}"

-    mkdir -p "$AI_DIR"
+    # ── Ollama model selection ────────────────────────────────────────────────
+    echo ""
+    log_info "Ollama LLM models — select which to download (enter numbers separated by spaces):"
+    log_info "All models are quantized (Q4_K_M) and run comfortably on 6 GB VRAM."
+    echo ""
+    log_info "  1) llama3.2:3b       ~2.0 GB  Fast general chat. Great all-rounder.  ← Recommended"
+    log_info "  2) llama3.2:1b       ~1.3 GB  Ultra-fast. Light tasks, low latency."
+    log_info "  3) qwen2.5:7b        ~4.7 GB  Top code + math model. Strong reasoning."
+    log_info "  4) mistral:7b        ~4.1 GB  Solid all-rounder. Good at instruction follow."
+    log_info "  5) phi4-mini         ~2.5 GB  Microsoft Phi-4 mini. Excellent for coding."
+    log_info "  6) gemma3:4b         ~2.5 GB  Google Gemma 3. Well-rounded, multilingual."
+    log_info "  7) deepseek-r1:7b    ~4.7 GB  Strong reasoning and math. Think-step model."
+    log_info "  8) nomic-embed-text  ~274 MB  Embedding model — enables RAG/doc search."
+    log_info "                                 Recommended to add alongside a chat model."
+    echo ""
+    log_info "  Example: '1 8' pulls llama3.2:3b + nomic-embed-text"
+    log_info "  Enter '0' or leave blank to skip and pull models manually later."
+    echo ""
+
+    local OLLAMA_CHOICES=""
+    prompt_text "Models to download [1 8]:" "1 8" OLLAMA_CHOICES
+
+    declare -a OLLAMA_MODELS=()
+    for _n in $OLLAMA_CHOICES; do
+        case "$_n" in
+            1) OLLAMA_MODELS+=("llama3.2:3b") ;;
+            2) OLLAMA_MODELS+=("llama3.2:1b") ;;
+            3) OLLAMA_MODELS+=("qwen2.5:7b") ;;
+            4) OLLAMA_MODELS+=("mistral:7b") ;;
+            5) OLLAMA_MODELS+=("phi4-mini") ;;
+            6) OLLAMA_MODELS+=("gemma3:4b") ;;
+            7) OLLAMA_MODELS+=("deepseek-r1:7b") ;;
+            8) OLLAMA_MODELS+=("nomic-embed-text") ;;
+        esac
+    done
+
+    # ── InvokeAI model selection ──────────────────────────────────────────────
+    echo ""
+    log_info "InvokeAI image generation models (for 6 GB VRAM with partial GPU offload):"
+    log_info "InvokeAI uses VRAM=3 GB + 8 GB RAM cache, so all models below work on 6 GB."
+    echo ""
+    log_info "  1) stabilityai/stable-diffusion-v1-5   ~4 GB  SD 1.5 — fast, huge style/LoRA library."
+    log_info "                                                  Best starting model for most uses."
+    log_info "  2) stabilityai/sdxl-turbo              ~7 GB  SDXL Turbo — 4-step generation."
+    log_info "                                                  Fast, high quality, slightly slower on 6 GB."
+    log_info "  3) stabilityai/stable-diffusion-xl-base-1.0"
+    log_info "                                         ~7 GB  SDXL base — best quality at 1024px."
+    log_info "                                                  Slowest due to RAM offload on 6 GB."
+    log_info "  4) Skip — install models via the Model Manager at http://localhost:9090"
+    echo ""
+    log_info "  Tip: SD 1.5 (choice 1) is fastest and most compatible. Start here."
+    log_info "  HuggingFace token: required for some gated models (free at huggingface.co/settings/tokens)"
+    echo ""
+
+    local INVOKE_CHOICE=""
+    prompt_text "InvokeAI starter model [1]:" "1" INVOKE_CHOICE
+
+    local INVOKE_MODEL_SOURCE=""
+    local INVOKE_MODEL_NAME=""
+    case "$INVOKE_CHOICE" in
+        2) INVOKE_MODEL_SOURCE="stabilityai/sdxl-turbo"
+           INVOKE_MODEL_NAME="SDXL Turbo" ;;
+        3) INVOKE_MODEL_SOURCE="stabilityai/stable-diffusion-xl-base-1.0"
+           INVOKE_MODEL_NAME="SDXL Base" ;;
+        4) INVOKE_MODEL_SOURCE=""
+           INVOKE_MODEL_NAME="" ;;
+        *) INVOKE_MODEL_SOURCE="stabilityai/stable-diffusion-v1-5"
+           INVOKE_MODEL_NAME="SD 1.5" ;;
+    esac
+
+    local HF_TOKEN=""
+    if [ -n "$INVOKE_MODEL_SOURCE" ]; then
+        prompt_text "HuggingFace token (optional — needed for gated models, enter to skip):" "" HF_TOKEN
+    fi

    # ── Clone / update repo ───────────────────────────────────────────────────
+    mkdir -p "$AI_DIR"
    if [ -d "$REPO_DIR/.git" ]; then
        log_info "Updating ai-6gb-gpu repo..."
        git -C "$REPO_DIR" pull --ff-only 2>/dev/null \
@@ -235,7 +310,6 @@ install_ai_gpu() {
    mkdir -p "$IMAGE_GEN_DIR"
    if [ -d "$REPO_DIR/ai-image-gen" ]; then
        cp -rn "$REPO_DIR/ai-image-gen/." "$IMAGE_GEN_DIR/" 2>/dev/null || true
-        # Replace any hardcoded timezone
        find "$IMAGE_GEN_DIR" -name "docker-compose.yml" -exec \
            sed -i "s|America/New_York|$TZ_VAL|g" {} \;
    fi
@@ -243,10 +317,11 @@ install_ai_gpu() {
    cat > "$IMAGE_GEN_DIR/.env" << IMGENV
 # InvokeAI — image generation
 TZ=${TZ_VAL}
-# VRAM cap: 3 GB leaves headroom on a 6 GB card
+# VRAM cap: 3 GB leaves headroom on a 6 GB card; remaining model layers go to RAM
 INVOKEAI_vram=3
-# RAM cache for model layers
+# RAM cache size for model layer offload
 INVOKEAI_ram=8
+${HF_TOKEN:+HUGGING_FACE_HUB_TOKEN=${HF_TOKEN}}
 CADDY_NET=${SITE_CADDY_NET}
 IMGENV
    chmod 600 "$IMAGE_GEN_DIR/.env"
@@ -266,8 +341,11 @@ IMGENV
    cat > "$LLM_DIR/.env" << LLMENV
 # Ollama + Open WebUI + SearXNG
 TZ=${TZ_VAL}
-# Open WebUI session secret
 WEBUI_SECRET_KEY=${WEBUI_SECRET}
+# Open WebUI: enable SearXNG for web search in chats
+ENABLE_RAG_WEB_SEARCH=true
+RAG_WEB_SEARCH_ENGINE=searxng
+SEARXNG_QUERY_URL=http://searxng:8080/search?q=<query>&format=json
 CADDY_NET=${SITE_CADDY_NET}
 LLMENV
    chmod 600 "$LLM_DIR/.env"
@@ -277,7 +355,6 @@ LLMENV
    mkdir -p "$PORTAL_DIR"
    if [ -d "$REPO_DIR/ai-portal" ]; then
        cp -rn "$REPO_DIR/ai-portal/." "$PORTAL_DIR/" 2>/dev/null || true
-        # Fix hardcoded home path in docker-compose.yml volume mounts
        if [ -f "$PORTAL_DIR/docker-compose.yml" ]; then
            sed -i \
                "s|/home/[^/]*/docker:|${ACTUAL_HOME}/docker:|g" \
@@ -289,25 +366,20 @@ LLMENV
    cat > "$PORTAL_DIR/.env" << PORTALENV
 # AI Portal — GPU stack swap controller
 TZ=${TZ_VAL}
-# Paths inside the container (Docker socket mount maps ACTUAL_HOME/docker → /docker)
+# Paths inside the container (/docker maps to ${ACTUAL_HOME}/docker via volume mount)
 IMAGE_STACK=/docker/ai-gpu/image-gen
 LLM_STACK=/docker/ai-gpu/llm
 CADDY_NET=${SITE_CADDY_NET}
 PORTALENV
    chmod 600 "$PORTAL_DIR/.env"

-    # Set ownership across everything
    ensure_docker_dir_ownership "$AI_DIR"

    echo ""
    log_success "AI GPU stacks configured under $AI_DIR"
-    log_info "Stack layout:"
-    log_info "  image-gen/  — InvokeAI  (port 9090,  nvidia GPU)"
-    log_info "  llm/        — Ollama (11434) + Open WebUI (3000) + SearXNG (internal)"
-    log_info "  portal/     — GPU swap portal (port 8080)"
-    echo ""
    log_warning "Only ONE GPU stack can run at a time on a 6 GB card."
-    log_warning "Use the portal to switch, or manually: docker compose -f <stack>/docker-compose.yml down/up."
+    log_info "Use the portal (port 8080) to hot-swap between image-gen and llm."
+    echo ""

    # ── Caddy ─────────────────────────────────────────────────────────────────
    configure_caddy_for_service "AI Portal" "ai-portal:8080" "localai"
@@ -324,125 +396,177 @@ Source: https://github.com/outis1one/ai-6gb-gpu

 | Stack | Service | Port | Notes |
 |-------|---------|------|-------|
-| \`image-gen/\` | InvokeAI | 9090 | Image generation (SDXL, Flux, etc.) |
+| \`image-gen/\` | InvokeAI | 9090 | Image generation (SD 1.5, SDXL, Flux…) |
 | \`llm/\` | Ollama | 11434 | LLM inference engine |
-| \`llm/\` | Open WebUI | 3000 | Chat UI for Ollama |
-| \`llm/\` | SearXNG | — | Internal web search for RAG |
-| \`portal/\` | AI Portal | 8080 | GPU swap controller UI |
+| \`llm/\` | Open WebUI | 3000 | Chat UI — models, RAG, web search |
+| \`llm/\` | SearXNG | internal | Web search backend for RAG in OpenWebUI |
+| \`portal/\` | AI Portal | 8080 | GPU swap controller — start/stop stacks |

-## Important: GPU time-sharing
+## GPU time-sharing (important)

-A 6 GB GPU can only run one AI stack at a time. Use the portal at
-http://localhost:8080 to swap between image-gen and llm — it stops the
-active stack before starting the requested one.
+A 6 GB GPU can only run one AI stack at a time.
+Use the portal at http://localhost:8080 to swap.

 Manual swap:
 \`\`\`bash
-# Stop image-gen, start llm
-docker compose -f $AI_DIR/image-gen/docker-compose.yml down
-docker compose -f $AI_DIR/llm/docker-compose.yml up -d
-
-# Stop llm, start image-gen
 docker compose -f $AI_DIR/llm/docker-compose.yml down
 docker compose -f $AI_DIR/image-gen/docker-compose.yml up -d
 \`\`\`

-## First-run setup
+## InvokeAI — add more models

-### InvokeAI
-1. Open http://localhost:9090
-2. Install models via the Model Manager (HuggingFace token may be needed)
-3. Recommended for 6 GB: SDXL-Turbo, Flux-Schnell-quantised
+Models installed at setup are in the Model Manager. To add more:
+1. Open http://localhost:9090 → Model Manager → Add Model
+2. Paste a HuggingFace repo ID (e.g. \`stabilityai/stable-diffusion-2-1\`)
+3. Or import a local .safetensors file
+
+For gated models, add \`HUGGING_FACE_HUB_TOKEN=xxx\` to \`image-gen/.env\`.
+
+Recommended models for 6 GB (with partial GPU offload):
+
+| Model | Source | Notes |
+|-------|--------|-------|
+| SD 1.5 | \`stabilityai/stable-diffusion-v1-5\` | Fast, huge LoRA library |
+| SDXL Turbo | \`stabilityai/sdxl-turbo\` | 4-step, good quality |
+| SDXL Base | \`stabilityai/stable-diffusion-xl-base-1.0\` | Best quality, slower |
+| SD 2.1 | \`stabilityai/stable-diffusion-2-1\` | Good mid-size choice |
+
+## Ollama — add more models

-### Ollama
 \`\`\`bash
-# Pull a model (while llm stack is running)
-docker exec ollama ollama pull llama3.2
-docker exec ollama ollama pull nomic-embed-text  # for RAG embeddings
+# Pull any model while llm stack is running
+docker exec ollama ollama pull llama3.2:3b
+docker exec ollama ollama pull nomic-embed-text    # RAG embeddings
+docker exec ollama ollama list                      # see installed models
 \`\`\`

-### Open WebUI
-Open http://localhost:3000 — create admin account on first visit.
+Browse models at: https://ollama.com/library
+For 6 GB cards, stick to 7B or smaller with Q4_K_M quantisation (~4.5 GB).
+
+## Open WebUI — first login
+
+Open http://localhost:3000 and create your admin account on first visit.
+Models pulled into Ollama appear automatically in the model dropdown.
+Enable web search: Settings → Admin → Web Search (SearXNG is pre-configured).

 ## Manage individual stacks
 \`\`\`bash
-# Image generation
-cd $AI_DIR/image-gen
-docker compose up -d
-docker compose down
-docker compose logs -f invokeai
-
-# LLM + chat
-cd $AI_DIR/llm
-docker compose up -d
-docker compose down
-docker compose logs -f openwebui
-
-# Portal
-cd $AI_DIR/portal
-docker compose up -d
-docker compose down
+cd $AI_DIR/image-gen && docker compose up -d    # start InvokeAI
+cd $AI_DIR/llm       && docker compose up -d    # start Ollama + OpenWebUI
+cd $AI_DIR/portal    && docker compose up -d    # start portal
+docker compose -f $AI_DIR/image-gen/docker-compose.yml logs -f invokeai
+docker compose -f $AI_DIR/llm/docker-compose.yml       logs -f openwebui
 \`\`\`

 ## Update
 \`\`\`bash
-# Pull latest repo changes and rebuild
 cd $REPO_DIR && git pull
 cd $AI_DIR/image-gen && docker compose pull && docker compose up -d
 cd $AI_DIR/llm       && docker compose pull && docker compose up -d
 cd $AI_DIR/portal    && docker compose build --pull && docker compose up -d
 \`\`\`
-
-## Files
- image-gen/.env  — InvokeAI VRAM/RAM limits and TZ
- llm/.env        — Open WebUI secret key and TZ
- portal/.env     — stack paths and TZ
 MD

-    # ── Start prompt ──────────────────────────────────────────────────────────
+    # ── Start stacks + pull models ────────────────────────────────────────────
+    echo ""
+    log_info "What would you like to start now?"
+    log_info "  1) Portal only           — start the swap controller, configure the rest later"
+    log_info "  2) Portal + LLM stack    — start Ollama/OpenWebUI and pull selected models"
+    log_info "  3) Portal + image-gen    — start InvokeAI and queue the starter model download"
+    log_info "  4) None                  — start manually later"
    echo ""
-    log_info "Which stack do you want to start now?"
-    log_info "  1) Portal only (recommended first — lets you manage the others)"
-    log_info "  2) Portal + image-gen (InvokeAI)"
-    log_info "  3) Portal + llm (Ollama/OpenWebUI)"
-    log_info "  4) None — start manually later"

    local START_CHOICE=""
-    prompt_text "Choice [1]:" "1" START_CHOICE
+    prompt_text "Choice [2]:" "2" START_CHOICE

-    case "$START_CHOICE" in
-        1)
-            docker compose -f "$PORTAL_DIR/docker-compose.yml" up -d \
-                && log_success "Portal started — http://localhost:8080" \
-                || log_warning "Portal start failed — check: docker compose -f $PORTAL_DIR/docker-compose.yml logs"
-            ;;
-        2)
-            docker compose -f "$PORTAL_DIR/docker-compose.yml" up -d \
-                && log_success "Portal started" \
-                || log_warning "Portal start failed"
-            docker compose -f "$IMAGE_GEN_DIR/docker-compose.yml" up -d \
-                && log_success "InvokeAI started — http://localhost:9090" \
-                || log_warning "InvokeAI start failed — check: docker compose -f $IMAGE_GEN_DIR/docker-compose.yml logs"
-            ;;
-        3)
-            docker compose -f "$PORTAL_DIR/docker-compose.yml" up -d \
-                && log_success "Portal started" \
-                || log_warning "Portal start failed"
-            docker compose -f "$LLM_DIR/docker-compose.yml" up -d \
-                && log_success "LLM stack started — Open WebUI: http://localhost:3000" \
-                || log_warning "LLM start failed — check: docker compose -f $LLM_DIR/docker-compose.yml logs"
-            ;;
-        *)
-            log_info "Skipped. Start when ready:"
-            log_info "  docker compose -f $PORTAL_DIR/docker-compose.yml up -d"
-            ;;
-    esac
+    # Always start portal if any stack is starting
+    if [[ "$START_CHOICE" =~ ^[123]$ ]]; then
+        docker compose -f "$PORTAL_DIR/docker-compose.yml" up -d \
+            && log_success "Portal started — http://localhost:8080" \
+            || log_warning "Portal start failed — check: docker compose -f $PORTAL_DIR/docker-compose.yml logs"
+    fi
+
+    if [[ "$START_CHOICE" == "2" ]]; then
+        # Start LLM stack
+        docker compose -f "$LLM_DIR/docker-compose.yml" up -d \
+            && log_success "LLM stack started" \
+            || { log_warning "LLM stack start failed"; START_CHOICE="0"; }
+
+        # Pull Ollama models if any were selected
+        if [ ${#OLLAMA_MODELS[@]} -gt 0 ] && [[ "$START_CHOICE" == "2" ]]; then
+            log_info "Waiting for Ollama to be ready..."
+            local _w=0
+            while ! curl -sf "http://localhost:11434/api/version" &>/dev/null; do
+                sleep 3; _w=$((_w+3))
+                [[ $_w -ge 90 ]] && { log_warning "Ollama not responding after 90s — pull models manually later"; break; }
+            done
+
+            if curl -sf "http://localhost:11434/api/version" &>/dev/null; then
+                for _m in "${OLLAMA_MODELS[@]}"; do
+                    log_info "Pulling $_m (this may take a while)..."
+                    docker exec ollama ollama pull "$_m" \
+                        && log_success "$_m ready" \
+                        || log_warning "Pull failed for $_m — retry: docker exec ollama ollama pull $_m"
+                done
+                log_success "Open WebUI ready at: http://localhost:3000"
+                log_info "Create your admin account on the first visit."
+            fi
+        fi
+    fi
+
+    if [[ "$START_CHOICE" == "3" ]]; then
+        # Start image-gen stack
+        docker compose -f "$IMAGE_GEN_DIR/docker-compose.yml" up -d \
+            && log_success "InvokeAI started" \
+            || { log_warning "InvokeAI start failed — check: docker compose -f $IMAGE_GEN_DIR/docker-compose.yml logs"; START_CHOICE="0"; }
+
+        # Queue starter model via InvokeAI REST API
+        if [ -n "$INVOKE_MODEL_SOURCE" ] && [[ "$START_CHOICE" == "3" ]]; then
+            log_info "Waiting for InvokeAI to be ready (model database initialises on first start)..."
+            local _w=0
+            while ! curl -sf "http://localhost:9090/api/v1/app/version" &>/dev/null; do
+                sleep 5; _w=$((_w+5))
+                [[ $_w -ge 180 ]] && { log_warning "InvokeAI not responding after 3 min"; break; }
+            done
+
+            if curl -sf "http://localhost:9090/api/v1/app/version" &>/dev/null; then
+                log_info "Queuing $INVOKE_MODEL_NAME download..."
+                local _resp
+                _resp=$(curl -s -X POST "http://localhost:9090/api/v2/models/install" \
+                    -H "Content-Type: application/json" \
+                    -d "{\"source\": \"${INVOKE_MODEL_SOURCE}\"}" 2>/dev/null)
+                if echo "$_resp" | grep -q '"id"'; then
+                    log_success "$INVOKE_MODEL_NAME queued — downloading in background"
+                    log_info "Track progress: http://localhost:9090 → Model Manager → In Progress"
+                else
+                    log_warning "Could not queue via API. Install manually:"
+                    log_info "  Open http://localhost:9090 → Model Manager → Add Model"
+                    log_info "  Source: $INVOKE_MODEL_SOURCE"
+                fi
+            fi
+        fi
+    fi
+
+    if [[ "$START_CHOICE" == "4" ]] || [[ "$START_CHOICE" == "0" ]]; then
+        log_info "Start when ready:"
+        log_info "  docker compose -f $PORTAL_DIR/docker-compose.yml up -d"
+        log_info "  docker compose -f $LLM_DIR/docker-compose.yml up -d"
+        log_info "  docker compose -f $IMAGE_GEN_DIR/docker-compose.yml up -d"
+    fi

    echo ""
-    echo "  Portal:    http://localhost:8080"
-    echo "  InvokeAI:  http://localhost:9090  (image-gen stack)"
-    echo "  OpenWebUI: http://localhost:3000  (llm stack)"
-    echo "  Source:    $REPO_DIR"
+    echo "  Portal:       http://localhost:8080  (GPU swap controller)"
+    echo "  InvokeAI:     http://localhost:9090  (image-gen stack)"
+    echo "  Open WebUI:   http://localhost:3000  (llm stack)"
+    echo "  Ollama API:   http://localhost:11434 (llm stack)"
+    echo "  Source repo:  $REPO_DIR"
+    echo ""
+    if [ ${#OLLAMA_MODELS[@]} -gt 0 ]; then
+        echo "  Ollama models queued: ${OLLAMA_MODELS[*]}"
+    fi
+    if [ -n "$INVOKE_MODEL_NAME" ]; then
+        echo "  InvokeAI starter:     $INVOKE_MODEL_NAME ($INVOKE_MODEL_SOURCE)"
+    fi
    echo ""
 }

@@ -1,13 +1,18 @@
 #!/bin/bash
-# services/iopaint.sh — AI image inpainting: object removal, fill, restore (IOPaint + LaMa).
+# services/iopaint.sh — AI image editing: erase objects, fill regions, replace with text prompt.
 # Part of the modular post-install system (sourced by setup.sh).
 #
 # Can also be run standalone on any machine:
 #   sudo bash iopaint.sh
 # (Docker must already be installed when run standalone)
 #
-# Runs the LaMa model (erase/fill objects) by default on CPU.
-# For GPU inference set DEVICE=cuda in .env and install nvidia-container-toolkit.
+# Three use cases:
+#   Erase/remove  — mask an object, AI fills the gap (LaMa, CPU-safe)
+#   Inpaint/fill  — restore damaged areas, remove watermarks (multiple models)
+#   Replace       — mask + text prompt → AI draws new content (PowerPaint, GPU required)
+#
+# IOPaint is local-only: it cannot call a remote GPU or InvokeAI on another machine.
+# For text-guided replacement the GPU must be on this same machine.
 # IOPaint has no built-in auth — protect with Authelia via Caddy.

 # ── Standalone bootstrap ──────────────────────────────────────────────────────
@@ -192,8 +197,9 @@ install_iopaint() {

    if [ "$DRY_RUN" = true ]; then
        echo "[DRY-RUN] Would create $IOPAINT_DIR"
+        echo "[DRY-RUN] Would prompt for model and GPU (CUDA) support"
        echo "[DRY-RUN] Would write docker-compose.yml and .env"
-        echo "[DRY-RUN] Would offer GPU (CUDA) support and Authelia SSO"
+        echo "[DRY-RUN] Would offer Authelia SSO"
        return 0
    fi

@@ -201,10 +207,67 @@ install_iopaint() {
    ensure_docker_dir_ownership "$IOPAINT_DIR"
    cd "$IOPAINT_DIR" || return 1

-    # Ask about GPU
+    # ── GPU ──────────────────────────────────────────────────────────────────
+    log_info "IOPaint can:"
+    log_info "  • Erase / remove  — mask an object, AI fills the gap (works CPU-only)"
+    log_info "  • Inpaint / fill  — restore damaged areas, remove watermarks"
+    log_info "  • Replace         — mask + type what goes there → AI draws it (needs GPU)"
+    log_info "  Note: inference is always local — IOPaint cannot use a GPU on another machine."
+    echo ""
    local USE_GPU=""
    prompt_yn "Enable CUDA GPU support? Requires nvidia-container-toolkit (y/n):" "n" USE_GPU

+    # ── Model selection ───────────────────────────────────────────────────────
+    echo ""
+    log_info "Select default model (you can switch models in the UI without restarting):"
+    echo ""
+    log_info "  ── CPU-safe — work on any machine ─────────────────────────────────────────"
+    log_info "  1) lama                    Erase / removal. Intelligent gap fill. ~200 MB  ← Recommended"
+    log_info "  2) cv2                     OpenCV fill. No download, instant. Rough quality."
+    log_info "  3) zits                    Portrait & face restoration. ~200 MB."
+    log_info "  4) manga                   Comic/manga text bubble removal. ~100 MB."
+    echo ""
+    if [[ "$USE_GPU" =~ ^[Yy]$ ]]; then
+        log_info "  ── GPU-accelerated fill ───────────────────────────────────────────────────"
+        log_info "  5) migan                   MiGAN: fast GPU inpainting. ~50 MB."
+        log_info "  6) fcf                     FcF: high-quality contextual fill. ~600 MB."
+        log_info "  7) mat                     MAT: large missing region fill. ~300 MB."
+        log_info "  8) ldm                     LDM: latent diffusion texture fill. ~1.2 GB."
+        echo ""
+        log_info "  ── Text-guided REPLACEMENT (GPU + Stable Diffusion) ───────────────────────"
+        log_info "  9) Sanster/PowerPaint-V2-filling"
+        log_info "                             Mask + type 'a red barn' → AI draws it. ~4 GB."
+        log_info "                             Modes: text-guided, shape-guided, erase, outpaint."
+        log_info " 10) runwayml/stable-diffusion-inpainting"
+        log_info "                             Classic SD 1.5 inpaint. Huge LoRA/style library. ~4 GB."
+        echo ""
+    fi
+
+    local MODEL_NUM=""
+    prompt_text "Model choice [1=lama]:" "1" MODEL_NUM
+
+    local IOPAINT_MODEL="lama"
+    case "$MODEL_NUM" in
+        2)  IOPAINT_MODEL="cv2" ;;
+        3)  IOPAINT_MODEL="zits" ;;
+        4)  IOPAINT_MODEL="manga" ;;
+        5)  IOPAINT_MODEL="migan" ;;
+        6)  IOPAINT_MODEL="fcf" ;;
+        7)  IOPAINT_MODEL="mat" ;;
+        8)  IOPAINT_MODEL="ldm" ;;
+        9)  IOPAINT_MODEL="Sanster/PowerPaint-V2-filling" ;;
+        10) IOPAINT_MODEL="runwayml/stable-diffusion-inpainting" ;;
+        *)  IOPAINT_MODEL="lama" ;;
+    esac
+
+    local DEVICE_VAL="cpu"
+    [[ "$USE_GPU" =~ ^[Yy]$ ]] && DEVICE_VAL="cuda"
+
+    # ── docker-compose.yml ────────────────────────────────────────────────────
+    # MODEL and DEVICE come from .env — change them there and restart to switch.
+    # Volume ./models:/root/.cache persists ALL model caches:
+    #   /root/.cache/torch/hub/checkpoints/  (LaMa, CV2, ZITS, etc.)
+    #   /root/.cache/huggingface/             (SD, PowerPaint, LDM, etc.)
    if [[ "$USE_GPU" =~ ^[Yy]$ ]]; then
        cat > docker-compose.yml << 'IOPAINT_GPU'
 name: iopaint
@@ -215,15 +278,19 @@ services:
    container_name: iopaint
    hostname: iopaint
    restart: unless-stopped
-    command: iopaint start --model=lama --device=cuda --port=8080 --host=0.0.0.0
+    command: >-
+      iopaint start
+      --model=${MODEL:-lama}
+      --device=${DEVICE:-cuda}
+      --port=8080
+      --host=0.0.0.0
    ports:
      - "8100:8080"
+    env_file: .env
    volumes:
-      - ./models:/root/.cache/iopaint
+      - ./models:/root/.cache
      - ./input:/app/input
      - ./output:/app/output
-    environment:
-      - DEVICE=cuda
    deploy:
      resources:
        reservations:
@@ -239,7 +306,6 @@ networks:
    external: true
    name: ${CADDY_NET:-caddy_net}
 IOPAINT_GPU
-        log_info "CUDA GPU mode enabled."
    else
        cat > docker-compose.yml << 'IOPAINT_CPU'
 name: iopaint
@@ -250,11 +316,17 @@ services:
    container_name: iopaint
    hostname: iopaint
    restart: unless-stopped
-    command: iopaint start --model=lama --device=cpu --port=8080 --host=0.0.0.0
+    command: >-
+      iopaint start
+      --model=${MODEL:-lama}
+      --device=${DEVICE:-cpu}
+      --port=8080
+      --host=0.0.0.0
    ports:
      - "8100:8080"
+    env_file: .env
    volumes:
-      - ./models:/root/.cache/iopaint
+      - ./models:/root/.cache
      - ./input:/app/input
      - ./output:/app/output
    networks:
@@ -265,18 +337,18 @@ networks:
    external: true
    name: ${CADDY_NET:-caddy_net}
 IOPAINT_CPU
-        log_info "CPU mode (default). Change DEVICE to cuda in .env and update the command to use GPU later."
    fi

+    # ── .env ─────────────────────────────────────────────────────────────────
    cat > .env << IOPAINT_ENV
-# IOPaint configuration
+# IOPaint — change MODEL and restart to switch (no need to edit docker-compose.yml)

-# Model to use for inpainting (lama recommended for object removal/erase)
-# Other models: ldm, zits, mat, fcf, manga, cv2, migan
-MODEL=lama
+# Current model (set during install — see README for full model list)
+MODEL=${IOPAINT_MODEL}

-# Device: cpu or cuda (cuda requires nvidia-container-toolkit + GPU deploy block)
-DEVICE=$([ "${USE_GPU,,}" = "y" ] && echo "cuda" || echo "cpu")
+# Device: cpu or cuda
+# For GPU: also requires the deploy: block in docker-compose.yml
+DEVICE=${DEVICE_VAL}

 # Caddy network
 CADDY_NET=${SITE_CADDY_NET}
@@ -287,9 +359,14 @@ IOPAINT_ENV
    chown -R "$ACTUAL_USER:$ACTUAL_USER" "$IOPAINT_DIR"

    echo ""
-    log_success "IOPaint configured at $IOPAINT_DIR"
-    log_info "LaMa model downloads on first start (~170 MB). Upload images via the web UI."
-    log_info "To switch models later: edit the --model= argument in docker-compose.yml and restart."
+    log_success "IOPaint configured — model: $IOPAINT_MODEL | device: $DEVICE_VAL"
+    if [[ "$IOPAINT_MODEL" == *"PowerPaint"* ]] || [[ "$IOPAINT_MODEL" == *"stable-diffusion"* ]]; then
+        log_info "SD-based model selected (~4 GB). It downloads from HuggingFace on first start."
+        log_info "If the download fails, try: HF_TOKEN=your_token docker compose up -d"
+    else
+        log_info "Model downloads automatically on first start (LaMa ~200 MB)."
+    fi
+    log_info "Switch models any time by editing MODEL= in .env and restarting."

    # No built-in auth — offer Authelia SSO protection
    local EXTRA_BLOCK=""
@@ -301,63 +378,97 @@ IOPAINT_ENV

    configure_caddy_for_service "IOPaint" "iopaint:8080" "inpaint" "$EXTRA_BLOCK"

-    write_readme "$IOPAINT_DIR" << MD
+    write_readme "$IOPAINT_DIR" << 'MD'
 # IOPaint

-AI-powered image inpainting — erase objects, fill regions, restore photos.
-Uses the LaMa (Large Mask) model for high-quality object removal by default.
+AI-powered image editing:
+- **Erase / remove** — mask an object, AI fills the background (LaMa, works CPU-only)
+- **Inpaint / restore** — fix damaged areas, remove watermarks
+- **Replace with AI** — mask something + type what goes there → AI draws it (PowerPaint, GPU)
+
+Note: IOPaint is **local only** — all inference runs on this machine.
+For text-guided replacement a CUDA GPU on this machine is required.

 ## Access
 - URL: http://localhost:8100
- No built-in login — protect via Authelia SSO if needed
+- No built-in login — protect via Authelia SSO if exposed

-## Models
-The \`--model\` argument in docker-compose.yml selects the AI model:
-| Model  | Best for |
-|--------|----------|
-| \`lama\` | Object removal, erase (default) |
-| \`ldm\`  | Texture-aware fill |
-| \`zits\` | Face/portrait restoration |
-| \`mat\`  | Large missing region fill |
-| \`manga\`| Manga/comic text removal |
+## Switching models
+Edit `MODEL=` in `.env` and restart — no need to touch `docker-compose.yml`:
+```bash
+cd ~/docker/iopaint
+nano .env          # change MODEL= line
+docker compose restart
+```

-Models download automatically on first use. Cached in \`./models/\`.
+## Model reference
+
+| # | Model | Type | Size | Best for |
+|---|-------|------|------|---------|
+| 1 | `lama` | CPU-safe | ~200 MB | **Object erase/removal** (default) |
+| 2 | `cv2` | CPU-safe | built-in | Basic fill, no download |
+| 3 | `zits` | CPU-safe | ~200 MB | Portrait & face restoration |
+| 4 | `manga` | CPU-safe | ~100 MB | Comic/manga text bubble removal |
+| 5 | `migan` | GPU | ~50 MB | Fast GPU inpainting |
+| 6 | `fcf` | GPU | ~600 MB | High-quality contextual fill |
+| 7 | `mat` | GPU | ~300 MB | Large missing region fill |
+| 8 | `ldm` | GPU | ~1.2 GB | Latent diffusion texture fill |
+| 9 | `Sanster/PowerPaint-V2-filling` | GPU+SD | ~4 GB | **Text-guided replacement** |
+| 10 | `runwayml/stable-diffusion-inpainting` | GPU+SD | ~4 GB | SD 1.5 inpaint, large LoRA library |
+
+SD-based models (9, 10) download from HuggingFace on first start.
+If a gated model needs a token: add `HF_TOKEN=xxx` to `.env`.
+
+## How to use
+1. Open http://localhost:8100
+2. Upload an image (or drag & drop)
+3. Paint a mask over the area to change
+4. For erase models: click Run → gap fills automatically
+5. For SD models (PowerPaint): type a text prompt → AI draws it into the masked area

 ## GPU acceleration
-Requires \`nvidia-container-toolkit\`. To enable:
-1. Edit docker-compose.yml: change \`--device=cpu\` → \`--device=cuda\`
-2. Uncomment the \`deploy:\` block (or re-run this installer with GPU=y)
-3. Restart: \`docker compose down && docker compose up -d\`
+Requires `nvidia-container-toolkit`. The GPU compose adds a `deploy:` block.
+Re-run the installer with GPU=y to regenerate docker-compose.yml, or manually add:
+```yaml
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+```
+Then change `DEVICE=cuda` in `.env` and restart.

 ## Manage
-\`\`\`bash
-cd $IOPAINT_DIR
-docker compose up -d                                      # start
-docker compose down                                       # stop
-docker compose logs -f                                    # logs
-docker compose pull && docker compose down && docker compose up -d  # update
-\`\`\`
+```bash
+cd ~/docker/iopaint
+docker compose up -d
+docker compose down
+docker compose logs -f
+docker compose pull && docker compose down && docker compose up -d
+```

 ## Files
- docker-compose.yml — stack definition (edit for model/device changes)
- .env               — runtime config
- models/            — cached AI model weights
- input/             — optional: place images here
- output/            — processed images written here
+- docker-compose.yml — stack (MODEL and DEVICE come from .env)
+- .env               — model and device config
+- models/            — all cached model weights (torch + HuggingFace)
+- input/, output/    — optional file staging
 MD

    local START_IO=""
    prompt_yn "Start IOPaint now? (y/n):" "y" START_IO
    if [ "$START_IO" = "y" ] || [ "$START_IO" = "Y" ]; then
        docker compose up -d \
-            && log_success "IOPaint started — LaMa model will download on first use" \
+            && log_success "IOPaint started — model downloads on first use" \
            || log_warning "Start failed — check: docker compose logs"
    fi

    echo ""
    echo "  URL:      http://localhost:8100"
-    echo "  Model:    LaMa (erase / object removal)"
-    echo "  Device:   $([ "${USE_GPU,,}" = "y" ] && echo "CUDA GPU" || echo "CPU")"
+    echo "  Model:    $IOPAINT_MODEL"
+    echo "  Device:   $DEVICE_VAL"
+    echo "  Switch:   edit MODEL= in $IOPAINT_DIR/.env and restart"
    echo ""
 }