# Design: model menu + persistent daemon + docker compose

Date: 2026-07-08
Status: approved (brainstorm), pending implementation-plan

## Goal

Extend the local SDXL image generator with:

1. **Model menu** — choose between several SDXL-arch models (permissive set).
2. **Persistent daemon** — hold one model resident in XPU VRAM so prompt
   iteration on the same image skips the model-reload tax.
3. **Explicit lifecycle** — start the daemon while revising, stop it to free
   VRAM when done.
4. **docker compose** — replace the raw `docker run` with compose for the
   long-lived service.

The existing bounded scope (SDXL arch, single 1024x1024 image, no refiner, no
Flux, no batching) is unchanged. The daemon is in scope specifically as a
local burst-iteration convenience, not general server hosting.

## Non-goals

- No multiple models resident at once (3 * ~9GB > 12GB VRAM; won't fit).
- No refiner, Flux, batching, video (still RunPod territory).
- No arbitrary HF model ids from the CLI — only the registry keys.
- No auth / remote exposure — daemon binds localhost only.

## Architecture

Three roles (was two):

```
imggen (host bash wrapper)
   ├─ subcommands: start / stop / status / list  → docker compose + curl
   └─ default (prompt given)                      → curl daemon, or one-shot fallback
        │
        ▼  HTTP 127.0.0.1:PORT
server.py (in container, long-lived)   ← NEW
   ├─ holds ONE model resident in XPU VRAM
   ├─ POST /generate  → PNG bytes
   ├─ POST /model     → swap resident model (unload + load)
   ├─ GET  /status    → {model, ready}
   └─ imports generate.py core (build_pipe / generate / resolve_model)
        │
        ▼
generate.py   ← refactored: model registry + reusable funcs, keeps one-shot CLI
```

- `generate.py` stays runnable standalone (one-shot) AND importable by
  `server.py`. No logic duplicated.
- Daemon holds one model. Swap on `-m` mismatch (reload tax only on switch).
- compose defines the service. `imggen start` = `docker compose up -d`,
  `stop` = `down`.

File count grows from 3 to 5: `imggen`, `generate.py`, `server.py`,
`compose.yaml`, `.env.example` (+ `.env` gitignored, not committed).

## Model registry

Single dict in `generate.py`, source of truth for both CLI and daemon:

```python
MODELS = {
    "sdxl":    "stabilityai/stable-diffusion-xl-base-1.0",  # default, general
    "realvis": "SG161222/RealVisXL_V5.0",                   # photoreal, permissive
    "pony":    "AstraliteHeart/pony-diffusion-v6-xl",       # unrestricted, versatile
}
DEFAULT_MODEL = "sdxl"
VAE = "madebyollin/sdxl-vae-fp16-fix"   # shared, all three are SDXL arch
```

- All three are SDXL arch → same `StableDiffusionXLPipeline`, same fp16-fix
  VAE, same 12GB fit. One `build_pipe(key)` handles all.
- `-m/--model` accepts a **registry key only** (`realvis`), not raw HF ids.
  Unknown key → error listing valid keys, exit 2.
- First use of a model downloads ~7GB to the `/models` cache (HF_HOME), once.

### Load robustness (MUST handle — real risk)

`build_pipe` must be robust to fine-tunes that differ from SDXL base:

```
try:    from_pretrained(..., variant="fp16")     # SDXL base publishes fp16 variant
except: from_pretrained(..., torch_dtype=float16) # fine-tunes may lack an fp16 variant
```

- **Repo IDs must be verified to load before locking.** RealVis and Pony repo
  ids + diffusers-format layout + fp16-variant availability are unconfirmed
  (the verification WebFetch was interrupted during brainstorm). If a repo
  ships only a single-file `.safetensors` (no diffusers folder), `from_pretrained`
  fails and it needs `from_single_file` or a diffusers-format mirror instead.
  This is the one unknown that a real download test resolves.

## HTTP daemon API

`server.py`, stdlib `http.server` only (no new dependency). Binds
`127.0.0.1:8765` inside the container; port published to host localhost only.

| method | path       | body                                              | returns |
|--------|------------|---------------------------------------------------|---------|
| GET    | `/status`  | —                                                 | `{"model":"realvis","ready":true}` |
| POST   | `/generate`| `{prompt,negative,steps,guidance,seed,width,height}` | PNG bytes (`image/png`) |
| POST   | `/model`   | `{"name":"pony"}`                                 | `{"model":"pony","ready":true}` after swap |

- Holds one pipe resident. `/generate` = pure inference, no reload.
- `/model` with a different key → unload current (`del pipe;
  torch.xpu.empty_cache()`), load new. Same key → no-op.
- Loads `IMGGEN_MODEL` (default `sdxl`) at startup.
- **Single-threaded / serialized** — one GPU, concurrent generate = OOM risk.
  `http.server` default is serial, which is what we want. Requests queue.
- Image returned as **bytes**; the **host wrapper writes the PNG** to
  `$OUT_DIR`. Daemon stays stateless about filenames; output-path logic stays
  host-side as it is now.

## compose.yaml

```yaml
services:
  imggen:
    build: .
    image: imggen:local
    user: "${IMGGEN_UID:-1000}:${IMGGEN_GID:-100}"
    group_add:
      - video          # gid for /dev/dri/card* access as non-root
    devices:
      - /dev/dri
    volumes:
      - ${IMGGEN_OUT:-./out}:/out
      - ${IMGGEN_CACHE:-./models}:/models
    environment:
      - HF_HOME=/models
      - IMGGEN_MODEL=${IMGGEN_MODEL:-sdxl}
    ports:
      - "127.0.0.1:${IMGGEN_PORT:-8765}:8765"
    command: ["python3", "server.py"]
```

- `devices: /dev/dri` = compose equivalent of `--device /dev/dri`. Verify Arc
  XPU is visible under compose (same flag, expected fine).
- Port default 8765, overridable via `IMGGEN_PORT`.

## Configuration: .env (real values out of git)

- Real host paths live in `.env` (gitignored). Compose auto-reads it.
- `.env.example` committed with generic placeholders, documents the shape.
- Real defaults (from the user's actual `~/bin/imggen`):

```
# .env  (gitignored)
IMGGEN_OUT=/home/danix/Pictures/AI-gen
IMGGEN_CACHE=/data/LLM-models/imggen-models
# IMGGEN_PORT=8765
# IMGGEN_MODEL=sdxl
```

Note these differ from the paths documented in CLAUDE.md (`~/imggen-out`,
`~/.cache/imggen-models`) — CLAUDE.md is stale and must be updated (see Docs).

## Host wrapper (imggen)

| command                        | action |
|--------------------------------|--------|
| `imggen start [-m KEY]`        | `IMGGEN_MODEL=KEY docker compose up -d` (idempotent) |
| `imggen stop`                  | `docker compose down` (frees VRAM) |
| `imggen status`                | `curl -s localhost:PORT/status` else "daemon down" |
| `imggen list`                  | print registry keys + descriptions |
| `imggen [-m KEY] "prompt" ...` | daemon up → POST /model if key differs, then POST /generate, write PNG. daemon down → one-shot fallback |

### Interactive menu

`imggen "prompt"` with **no `-m`**:

- daemon up + TTY present → prompt `which model? 1)sdxl 2)realvis 3)pony`,
  then swap if chosen ≠ resident, then generate.
- non-TTY (piped/scripted) → fall back to `sdxl`.
- daemon down → one-shot fallback (`docker compose run --rm`), non-interactive
  → `sdxl` or the `-m` key.

## Error handling

| case                          | behavior |
|-------------------------------|----------|
| unknown `-m KEY`              | error + list valid keys, exit 2 |
| daemon down, prompt sent      | fall back to one-shot (no error) |
| model download fails (network)| HF raises → daemon 500 w/ msg; CLI prints + exits |
| generate OOM                  | catch, `empty_cache()`, 500 "OOM, try fewer steps"; daemon survives |
| bad JSON to daemon            | 400 |
| swap during generate          | serialized (single-thread), no race |
| XPU unavailable in daemon     | startup crash → container exits → `status` shows down (fail-loud preserved) |

## Testing

Hardware (Intel Arc B580) is on the same machine, so integration runs locally,
not hand-waved. No test framework added (ponytail): plain `assert`
self-checks / one `test_*.py` for pure logic; GPU path exercised manually.

**Pure-logic, no GPU (automatable):**
- `resolve_model` — valid keys resolve, invalid key raises listing valid keys.
- registry integrity (keys map to non-empty ids, DEFAULT_MODEL in MODELS).
- wrapper arg dispatch routes start/stop/status/list vs prompt correctly.
- daemon JSON request parsing (valid body, bad JSON → 400).

**Integration on real Arc (run by the agent on this machine):**
- one-shot SDXL end-to-end → valid PNG (weights cached, fast).
- daemon lifecycle: `start` → `status` ready → `generate` → PNG valid → stop
  frees VRAM.
- model swap: `-m realvis` while sdxl resident → `/model` swaps → generate.
- **Model download tests (realvis, pony) — approved to pull** to the real
  cache `/data/LLM-models/imggen-models`. These confirm diffusers-format load
  + variant fallback for each fine-tune (resolves the load-robustness unknown).

## Known wrinkles (flagged, not blockers)

- **Container runs as the host user** (`user: 1000:100` + `group_add: video`)
  so output PNGs + cache land `danix:users`, not `root:root`. GPU access holds
  because uid 1000 is in `video` and `group_add: video` gives the container
  that gid; `renderD*` is world-rw as a fallback. Prereqs: the HF cache under
  `/data/LLM-models/imggen-models` must be owned by `danix:users` so downloads
  as uid 1000 can write it (the user is handling this chown; older runs left
  `root:root` files there).
- VRAM ~9GB is **locked while the daemon lives**. Nothing else can use the GPU
  meanwhile. `imggen stop` frees it. This is the intended tradeoff, documented
  for the user in the README warning.

## Docs to update (deliverables)

- **README.md** — new subcommands, model menu, daemon lifecycle, VRAM-locked
  warning, `.env` setup. Keep the "Development Approach" section.
- **CLAUDE.md** — new architecture (3 → 5 files), daemon in scope as local
  burst-iteration, corrected real `.env` defaults, compose replaces raw
  `docker run`.