image_server.py

"""
image_server.py — More AI Image Processing Server (port 7871)
=============================================================
Endpoints:
  GET  /            — health check
  GET  /gpu_info    — detect CUDA GPU and VRAM
  POST /remove_background — U2Net background removal (CPU, everyone)
  POST /upscale     — Real-ESRGAN upscale (2x or 4x)
  POST /adjust      — Pillow-based image adjustments
  POST /generate    — local text-to-image generation with cached PNG output
  POST /download_model — trigger model download with SSE progress
"""

import os, sys, io, base64, json, platform, time, hashlib, random, uuid, subprocess, urllib.error, urllib.parse, urllib.request, concurrent.futures
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import uvicorn

# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(title="More AI Image Server")

# CAT-5-002: Token auth middleware
import os as _os, sys as _sys
_SESSION_TOKEN = ""
for _i, _a in enumerate(_sys.argv):
    if _a == "--token" and _i + 1 < len(_sys.argv):
        _SESSION_TOKEN = _sys.argv[_i + 1]
        break

try:
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.responses import Response as _SR
    class _TokenAuth(BaseHTTPMiddleware):
        async def dispatch(self, request, call_next):
            if _SESSION_TOKEN and request.headers.get("x-more-ai-token","") != _SESSION_TOKEN:
                return _SR("Unauthorized", status_code=401)
            return await call_next(request)
    app.add_middleware(_TokenAuth)
except Exception as exc:
    raise RuntimeError("Sidecar auth middleware is required") from exc
MORE_AI_DIR = Path.home() / ".more_ai"
MODEL_DIR = MORE_AI_DIR / "image_models"
GENERATION_DIR = MORE_AI_DIR / "image_generations"
SERVER_SCRIPT_DIR = Path(__file__).resolve().parent
MODEL_DIR.mkdir(parents=True, exist_ok=True)
GENERATION_DIR.mkdir(parents=True, exist_ok=True)

SERVER_VERSION = "1.1.3"
GENERATION_RUNTIME_PACKAGE = "stable-diffusion-cpp-python"
DEFAULT_LOCAL_GENERATION_MODEL = "sd15_q4"
IMAGE_MODEL_EXTS = {".gguf", ".safetensors", ".ckpt"}
COMFYUI_URL = os.environ.get("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
COMFYUI_COMMON_PORTS = ("8188", "8000", "8189", "8190")
_COMFYUI_URL_CACHE: dict[str, object] = {"checked_at": 0.0, "url": None, "pid": None, "process": None}
_COMFYUI_PROCESS_PORT_CACHE: dict[str, object] = {"checked_at": 0.0, "ports": []}
_LOCAL_LISTENER_CACHE: dict[str, object] = {"checked_at": 0.0, "listeners": []}
_PROCESS_INFO_CACHE: dict[int, dict[str, object]] = {}
_PROCESS_SNAPSHOT_CACHE: dict[str, object] = {"checked_at": 0.0, "processes": {}}
IMAGE_MODEL_HINTS = (
    "stable-diffusion",
    "stablediffusion",
    "diffusion",
    "sdxl",
    "sd-",
    "sd_",
    "flux",
    "dreamshaper",
    "juggernaut",
    "realvis",
    "realisticvision",
    "epicrealism",
    "deliberate",
    "animagine",
    "pony",
    "checkpoint",
    "checkpoints",
    "txt2img",
    "text-to-image",
    "fooocus",
    "invokeai",
    "comfyui",
    "automatic1111",
)
NON_IMAGE_MODEL_HINTS = (
    "ace_step",
    "acestep",
    "audio",
    "audioldm",
    "bark",
    "cogvideo",
    "cogvideox",
    "hunyuan3d",
    "hunyuan_3d",
    "hunyuan-video",
    "ltx-2",
    "ltx-video",
    "ltx_video",
    "ltxv",
    "music",
    "musicgen",
    "suno",
    "text-to-video",
    "t2v",
    "video",
    "wan2",
    "wan-ai",
)
NON_CHECKPOINT_PARTS = {
    "feature_extractor",
    "scheduler",
    "text_encoder",
    "text_encoder_2",
    "tokenizer",
    "tokenizer_2",
    "transformer",
    "unet",
    "vae",
}

def _normalise_base_url(value: str) -> Optional[str]:
    clean = str(value or "").strip().strip('"').strip("'").rstrip("/")
    if not clean:
        return None
    if clean.startswith("http://") or clean.startswith("https://"):
        return clean
    if clean.startswith("localhost") or clean.startswith("127.") or clean.startswith("[::1]"):
        return f"http://{clean}"
    return None

def _append_comfyui_url(urls: list[str], value: str):
    clean = _normalise_base_url(value)
    if clean and clean not in urls:
        urls.append(clean)

def _json_from_powershell(script: str, timeout: float = 3.0):
    try:
        proc = subprocess.run(
            ["powershell", "-NoProfile", "-Command", script],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except Exception:
        return None
    if proc.returncode != 0:
        return None
    raw = (proc.stdout or "").strip()
    if not raw:
        return None
    try:
        return json.loads(raw)
    except Exception:
        return None

def _netstat_local_listeners() -> list[dict[str, object]]:
    listeners: list[dict[str, object]] = []
    try:
        proc = subprocess.run(["netstat", "-ano", "-p", "tcp"], capture_output=True, text=True, timeout=3)
        for line in (proc.stdout or "").splitlines():
            parts = line.split()
            if len(parts) < 5 or parts[0].upper() != "TCP" or parts[3].upper() != "LISTENING":
                continue
            local = parts[1].strip()
            pid = int(parts[4])
            if local.startswith("["):
                end = local.rfind("]:")
                if end < 0:
                    continue
                address = local[1:end]
                port = int(local[end + 2:])
            else:
                address, raw_port = local.rsplit(":", 1)
                port = int(raw_port)
            if address in {"127.0.0.1", "::1", "0.0.0.0", "::"}:
                listeners.append({"address": address, "port": port, "pid": pid})
    except Exception:
        pass
    return listeners

def _local_listening_ports() -> list[dict[str, object]]:
    now = time.time()
    checked_at = float(_LOCAL_LISTENER_CACHE.get("checked_at") or 0.0)
    if now - checked_at < 15.0:
        return list(_LOCAL_LISTENER_CACHE.get("listeners") or [])

    listeners: list[dict[str, object]] = _netstat_local_listeners()
    if platform.system().lower() == "windows":
        if not listeners:
            data = _json_from_powershell(
                (
                    "$ErrorActionPreference='SilentlyContinue'; "
                    "Get-NetTCPConnection -State Listen | "
                    "Where-Object { $_.LocalAddress -in @('127.0.0.1','::1','0.0.0.0','::') } | "
                    "Select-Object "
                    "@{Name='address';Expression={$_.LocalAddress}},"
                    "@{Name='port';Expression={$_.LocalPort}},"
                    "@{Name='pid';Expression={$_.OwningProcess}} | "
                    "ConvertTo-Json -Compress"
                ),
                timeout=3.0,
            )
            rows = data if isinstance(data, list) else ([data] if isinstance(data, dict) else [])
            for row in rows:
                try:
                    port = int(row.get("port") or row.get("Port") or 0)
                    pid = int(row.get("pid") or row.get("PID") or 0)
                except Exception:
                    continue
                address = str(row.get("address") or row.get("Address") or "127.0.0.1")
                if port > 0:
                    listeners.append({"address": address, "port": port, "pid": pid})

    deduped: list[dict[str, object]] = []
    seen: set[tuple[int, int]] = set()
    for listener in listeners:
        key = (int(listener.get("port") or 0), int(listener.get("pid") or 0))
        if key not in seen:
            seen.add(key)
            deduped.append(listener)
    _LOCAL_LISTENER_CACHE.update({"checked_at": now, "listeners": deduped})
    return deduped

def _process_info(pid: int) -> dict[str, object]:
    if pid <= 0:
        return {}
    cached = _PROCESS_INFO_CACHE.get(pid)
    if cached and time.time() - float(cached.get("checked_at") or 0.0) < 60.0:
        return cached

    info: dict[str, object] = {"checked_at": time.time(), "pid": pid, "parent_pid": 0, "name": "", "command_line": ""}
    if platform.system().lower() == "windows":
        data = _json_from_powershell(
            (
                "$ErrorActionPreference='SilentlyContinue'; "
                f"Get-CimInstance Win32_Process -Filter \"ProcessId = {pid}\" | "
                "Select-Object "
                "@{Name='pid';Expression={$_.ProcessId}},"
                "@{Name='parent_pid';Expression={$_.ParentProcessId}},"
                "@{Name='name';Expression={$_.Name}},"
                "@{Name='command_line';Expression={$_.CommandLine}} | "
                "ConvertTo-Json -Compress"
            ),
            timeout=2.0,
        )
        if isinstance(data, dict):
            try:
                info.update({
                    "pid": int(data.get("pid") or pid),
                    "parent_pid": int(data.get("parent_pid") or 0),
                    "name": str(data.get("name") or ""),
                    "command_line": str(data.get("command_line") or ""),
                })
            except Exception:
                pass
    _PROCESS_INFO_CACHE[pid] = info
    return info

def _process_snapshot() -> dict[int, dict[str, object]]:
    now = time.time()
    checked_at = float(_PROCESS_SNAPSHOT_CACHE.get("checked_at") or 0.0)
    cached = _PROCESS_SNAPSHOT_CACHE.get("processes") or {}
    if now - checked_at < 30.0 and isinstance(cached, dict):
        return cached

    processes: dict[int, dict[str, object]] = {}
    if platform.system().lower() == "windows":
        data = _json_from_powershell(
            (
                "$ErrorActionPreference='SilentlyContinue'; "
                "Get-CimInstance Win32_Process | "
                "Select-Object "
                "@{Name='pid';Expression={$_.ProcessId}},"
                "@{Name='parent_pid';Expression={$_.ParentProcessId}},"
                "@{Name='name';Expression={$_.Name}},"
                "@{Name='command_line';Expression={$_.CommandLine}} | "
                "ConvertTo-Json -Compress"
            ),
            timeout=4.0,
        )
        rows = data if isinstance(data, list) else ([data] if isinstance(data, dict) else [])
        for row in rows:
            try:
                row_pid = int(row.get("pid") or 0)
                if row_pid <= 0:
                    continue
                info = {
                    "checked_at": now,
                    "pid": row_pid,
                    "parent_pid": int(row.get("parent_pid") or 0),
                    "name": str(row.get("name") or ""),
                    "command_line": str(row.get("command_line") or ""),
                }
                processes[row_pid] = info
                _PROCESS_INFO_CACHE[row_pid] = info
            except Exception:
                continue
    _PROCESS_SNAPSHOT_CACHE.update({"checked_at": now, "processes": processes})
    return processes

def _process_lineage_text(pid: int, max_depth: int = 5) -> str:
    parts: list[str] = []
    current = int(pid or 0)
    seen: set[int] = set()
    snapshot = _process_snapshot()
    for _ in range(max_depth):
        if current <= 0 or current in seen:
            break
        seen.add(current)
        info = snapshot.get(current) or _process_info(current)
        parts.append(f"{info.get('name','')} {info.get('command_line','')}")
        current = int(info.get("parent_pid") or 0)
    return " ".join(parts).lower()

def _lineage_looks_like_comfyui(pid: int) -> bool:
    text = _process_lineage_text(pid)
    return "comfyui" in text or "comfy desktop" in text

def _comfyui_ports_from_local_listeners(require_process_match: bool) -> list[str]:
    ports: list[str] = []
    for listener in _local_listening_ports():
        port = int(listener.get("port") or 0)
        pid = int(listener.get("pid") or 0)
        if port <= 0:
            continue
        if require_process_match and not _lineage_looks_like_comfyui(pid):
            continue
        ports.append(str(port))
    return [port for port in dict.fromkeys(ports) if port.isdigit()]

def _discover_comfyui_ports_from_processes() -> list[str]:
    if platform.system().lower() != "windows":
        return []
    now = time.time()
    checked_at = float(_COMFYUI_PROCESS_PORT_CACHE.get("checked_at") or 0.0)
    if now - checked_at < 30.0:
        return list(_COMFYUI_PROCESS_PORT_CACHE.get("ports") or [])
    try:
        proc = subprocess.run(
            [
                "powershell",
                "-NoProfile",
                "-Command",
                (
                    "Get-CimInstance Win32_Process | "
                    "Where-Object { $_.CommandLine -match 'ComfyUI|Comfy Desktop' } | "
                    "Select-Object -ExpandProperty CommandLine"
                ),
            ],
            capture_output=True,
            text=True,
            timeout=2,
        )
    except Exception:
        _COMFYUI_PROCESS_PORT_CACHE.update({"checked_at": now, "ports": []})
        return []
    ports: list[str] = []
    for line in (proc.stdout or "").splitlines():
        parts = line.replace("=", " ").split()
        for idx, part in enumerate(parts):
            if part == "--port" and idx + 1 < len(parts) and parts[idx + 1].isdigit():
                ports.append(parts[idx + 1])
            elif part.startswith("--port") and part.removeprefix("--port").strip().isdigit():
                ports.append(part.removeprefix("--port").strip())
    discovered = [port for port in dict.fromkeys(ports) if port.isdigit()]
    _COMFYUI_PROCESS_PORT_CACHE.update({"checked_at": now, "ports": discovered})
    return discovered

def _candidate_comfyui_urls(
    include_process_ports: bool = True,
    include_listener_ports: bool = True,
    process_matched_only: bool = False,
    include_common_ports: bool = True,
) -> tuple[str, ...]:
    urls: list[str] = []
    for value in [os.environ.get("COMFYUI_URL", "")]:
        _append_comfyui_url(urls, value)
    ports = [*COMFYUI_COMMON_PORTS] if include_common_ports else []
    if include_process_ports:
        ports.extend(_discover_comfyui_ports_from_processes())
    if include_listener_ports:
        ports.extend(_comfyui_ports_from_local_listeners(process_matched_only))
    for port in ports:
        for host in ("127.0.0.1", "localhost", "[::1]"):
            _append_comfyui_url(urls, f"http://{host}:{port}")
    return tuple(urls)

def _looks_like_comfyui_system_stats(payload: object) -> bool:
    if not isinstance(payload, dict):
        return False
    system = payload.get("system")
    if not isinstance(system, dict):
        return False
    if system.get("comfyui_version"):
        return True
    argv = " ".join(str(item) for item in system.get("argv") or []).lower()
    packages = " ".join(str(item.get("name", "")) for item in system.get("comfy_package_versions") or [] if isinstance(item, dict)).lower()
    return "comfyui" in argv or "comfy" in packages

def _listener_pid_for_url(base_url: Optional[str]) -> Optional[int]:
    if not base_url:
        return None
    try:
        port = urllib.parse.urlparse(base_url).port
    except Exception:
        port = None
    if not port:
        return None
    for listener in _local_listening_ports():
        try:
            if int(listener.get("port") or 0) == int(port):
                return int(listener.get("pid") or 0) or None
        except Exception:
            continue
    return None

def _process_label_for_pid(pid: Optional[int]) -> Optional[str]:
    if not pid:
        return None
    info = _process_info(int(pid))
    name = str(info.get("name") or "").strip()
    command = str(info.get("command_line") or "").strip()
    return f"{name} PID {pid}: {command[:260]}".strip()

def _add_model_root(roots: list[Path], candidate: Optional[Path | str]):
    if not candidate:
        return
    try:
        path = Path(candidate).expanduser()
    except Exception:
        return
    if str(path) in {"", "."}:
        return
    if path.exists() and path.is_dir():
        try:
            resolved = path.resolve()
        except Exception:
            resolved = path
        if resolved not in roots:
            roots.append(resolved)

def _candidate_generation_roots() -> list[Path]:
    roots: list[Path] = []
    home = Path.home()
    local_app = os.environ.get("LOCALAPPDATA")
    app_data = os.environ.get("APPDATA")

    for candidate in [
        MODEL_DIR,
        os.environ.get("HUGGINGFACE_HUB_CACHE"),
        Path(os.environ["HF_HOME"]) / "hub" if os.environ.get("HF_HOME") else None,
        home / ".cache" / "huggingface" / "hub",
        home / ".cache" / "lm-studio" / "models",
        Path(local_app) / "LM Studio" / "models" if local_app else None,
        Path(local_app) / "lm-studio" / "models" if local_app else None,
        os.environ.get("LM_STUDIO_MODELS"),
        os.environ.get("COMFYUI_MODEL_PATH"),
        home / "ComfyUI" / "models" / "checkpoints",
        home / "Documents" / "ComfyUI" / "models" / "checkpoints",
        Path(app_data) / "ComfyUI" / "models" / "checkpoints" if app_data else None,
        Path(local_app) / "ComfyUI" / "models" / "checkpoints" if local_app else None,
        os.environ.get("A1111_MODEL_PATH"),
        os.environ.get("STABLE_DIFFUSION_MODELS"),
        home / "stable-diffusion-webui" / "models" / "Stable-diffusion",
        home / "Documents" / "stable-diffusion-webui" / "models" / "Stable-diffusion",
        home / "Fooocus" / "models" / "checkpoints",
        home / "Documents" / "Fooocus" / "models" / "checkpoints",
        home / "InvokeAI" / "models",
        home / "Documents" / "InvokeAI" / "models",
    ]:
        _add_model_root(roots, candidate)
    return roots

def _candidate_more_ai_hf_roots() -> list[Path]:
    roots: list[Path] = []
    app_data = os.environ.get("APPDATA")
    local_app = os.environ.get("LOCALAPPDATA")
    for candidate in [
        SERVER_SCRIPT_DIR / "hf_models",
        Path(app_data) / "com.trieross.more-ai" / "hf_models" if app_data else None,
        Path(local_app) / "com.trieross.more-ai" / "hf_models" if local_app else None,
    ]:
        _add_model_root(roots, candidate)
    return roots

def _model_source_for_path(path: Path) -> str:
    lowered = str(path).lower()
    if "\\hf_models\\" in lowered or "/hf_models/" in lowered:
        return "more_ai_hf"
    if str(MODEL_DIR).lower() in lowered:
        return "more_ai"
    if "huggingface" in lowered or "models--" in lowered:
        return "huggingface"
    if "lm-studio" in lowered or "lm studio" in lowered:
        return "lm_studio"
    if "comfyui" in lowered:
        return "comfyui"
    if "stable-diffusion-webui" in lowered or "automatic1111" in lowered:
        return "automatic1111"
    if "fooocus" in lowered:
        return "fooocus"
    if "invokeai" in lowered:
        return "invokeai"
    return "local"

def _looks_like_generation_model(path: Path, trusted_root: bool = False) -> bool:
    if path.suffix.lower() not in IMAGE_MODEL_EXTS:
        return False
    if _blocked_non_image_model_path(path):
        return False
    if trusted_root:
        return True
    haystack = str(path).replace("\\", "/").lower()
    return any(hint in haystack for hint in IMAGE_MODEL_HINTS)

def _blocked_non_image_model_path(path: Path) -> bool:
    haystack = str(path).replace("\\", "/").lower()
    if any(hint in haystack for hint in NON_IMAGE_MODEL_HINTS):
        return True
    return any(part.lower() in NON_CHECKPOINT_PARTS for part in path.parts)

def _model_label_for_path(path: Path) -> str:
    for part in path.parts:
        if part.startswith("models--"):
            repo = part.removeprefix("models--").replace("--", "/")
            return f"{repo} - {path.name}"
    return path.stem

def _model_key_for_path(path: Path, existing: set[str]) -> str:
    key = "".join(ch if ch.isalnum() or ch in {"_", "-"} else "_" for ch in path.stem).strip("_").lower()
    key = key or f"local_{hashlib.sha1(str(path).encode('utf-8')).hexdigest()[:8]}"
    if key not in existing:
        existing.add(key)
        return key
    digest = hashlib.sha1(str(path).encode("utf-8")).hexdigest()[:8]
    unique = f"{key}_{digest}"
    existing.add(unique)
    return unique

def _scan_generation_models(max_models: int = 512) -> list[dict]:
    found: list[dict] = []
    seen: set[str] = set()
    for root in _candidate_generation_roots():
        trusted_root = False
        try:
            trusted_root = root.resolve() == MODEL_DIR.resolve()
        except Exception:
            trusted_root = root == MODEL_DIR
        visited = 0
        try:
            iterator = root.rglob("*")
            for path in iterator:
                visited += 1
                if visited > 20000:
                    break
                if not path.is_file() or not _looks_like_generation_model(path, trusted_root):
                    continue
                try:
                    resolved = path.resolve()
                    key = str(resolved).lower()
                except Exception:
                    resolved = path
                    key = str(path).lower()
                if key in seen:
                    continue
                seen.add(key)
                found.append({
                    "path": resolved,
                    "source": _model_source_for_path(resolved),
                    "label": _model_label_for_path(resolved),
                })
                if len(found) >= max_models:
                    return found
        except Exception:
            continue
    return sorted(found, key=lambda item: str(item["path"]).lower())

def _comfyui_request_json(
    path: str,
    payload: Optional[dict] = None,
    timeout: float = 3.0,
    base_url: Optional[str] = None,
) -> dict:
    url = f"{(base_url or COMFYUI_URL).rstrip('/')}{path}"
    data = json.dumps(payload).encode("utf-8") if payload is not None else None
    request = urllib.request.Request(
        url,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST" if payload is not None else "GET",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        raw = response.read().decode("utf-8")
    return json.loads(raw or "{}")

def _comfyui_request_bytes(
    path: str,
    query: dict,
    timeout: float = 30.0,
    base_url: Optional[str] = None,
) -> bytes:
    url = f"{(base_url or COMFYUI_URL).rstrip('/')}{path}?{urllib.parse.urlencode(query)}"
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return response.read()

def _verified_comfyui_url(base_url: str, timeout: float) -> Optional[str]:
    try:
        stats = _comfyui_request_json("/system_stats", timeout=timeout, base_url=base_url)
        return base_url if _looks_like_comfyui_system_stats(stats) else None
    except Exception:
        return None

def _first_verified_comfyui_url(candidates: tuple[str, ...], timeout: float, parallel: bool) -> Optional[str]:
    unique = tuple(dict.fromkeys(candidates))
    if not unique:
        return None
    if not parallel or len(unique) == 1:
        for base_url in unique:
            found = _verified_comfyui_url(base_url, timeout)
            if found:
                return found
        return None

    executor = concurrent.futures.ThreadPoolExecutor(max_workers=min(64, len(unique)))
    try:
        future_to_url = {
            executor.submit(_verified_comfyui_url, base_url, timeout): base_url
            for base_url in unique
        }
        for future in concurrent.futures.as_completed(future_to_url, timeout=max(1.0, timeout + 0.5)):
            try:
                found = future.result()
            except Exception:
                found = None
            if found:
                return found
    except concurrent.futures.TimeoutError:
        return None
    finally:
        executor.shutdown(wait=False, cancel_futures=True)
    return None

def _active_comfyui_url() -> Optional[str]:
    now = time.time()
    checked_at = float(_COMFYUI_URL_CACHE.get("checked_at") or 0.0)
    if now - checked_at < 2.0:
        cached_url = _COMFYUI_URL_CACHE.get("url")
        return str(cached_url) if cached_url else None

    search_phases = [
        (_candidate_comfyui_urls(include_process_ports=False, include_listener_ports=False, include_common_ports=False), 0.7, False),
        (_candidate_comfyui_urls(include_process_ports=False, include_listener_ports=True, process_matched_only=False, include_common_ports=False), 0.45, True),
        (_candidate_comfyui_urls(include_process_ports=True, include_listener_ports=False, include_common_ports=False), 0.7, False),
        (_candidate_comfyui_urls(include_process_ports=False, include_listener_ports=False, include_common_ports=True), 0.35, True),
    ]
    tried: set[str] = set()
    for candidates, timeout, parallel in search_phases:
        phase_candidates = tuple(url for url in candidates if url not in tried)
        tried.update(phase_candidates)
        base_url = _first_verified_comfyui_url(phase_candidates, timeout=timeout, parallel=parallel)
        if base_url:
            pid = _listener_pid_for_url(base_url)
            _COMFYUI_URL_CACHE.update({
                "checked_at": now,
                "url": base_url,
                "pid": pid,
                "process": None,
            })
            return base_url
    _COMFYUI_URL_CACHE.update({"checked_at": now, "url": None, "pid": None, "process": None})
    return None

def _comfyui_available() -> bool:
    return _active_comfyui_url() is not None

def _comfyui_checkpoint_names(base_url: Optional[str] = None) -> list[str]:
    active_url = base_url or _active_comfyui_url()
    if not active_url:
        return []
    try:
        info = _comfyui_request_json("/object_info/CheckpointLoaderSimple", timeout=3.0, base_url=active_url)
    except Exception:
        return []

    node_info = info.get("CheckpointLoaderSimple", info if isinstance(info, dict) else {})
    ckpt_spec = (
        node_info.get("input", {})
        .get("required", {})
        .get("ckpt_name")
    )
    if isinstance(ckpt_spec, list) and ckpt_spec:
        names = ckpt_spec[0] if isinstance(ckpt_spec[0], list) else ckpt_spec
    else:
        names = []
    return sorted(
        {
            str(name).strip()
            for name in names
            if str(name).strip() and not _blocked_non_image_model_path(Path(str(name)))
        },
        key=lambda name: name.lower(),
    )

def _comfyui_model_label(ckpt_name: str) -> str:
    return Path(ckpt_name).stem or ckpt_name

def _scan_comfyui_models(max_models: int = 256) -> list[dict]:
    active_url = _active_comfyui_url()
    if not active_url:
        return []
    models: list[dict] = []
    for ckpt_name in _comfyui_checkpoint_names(active_url):
        models.append({
            "path": ckpt_name,
            "source": "comfyui",
            "label": _comfyui_model_label(ckpt_name),
            "engine": "comfyui",
            "runnable": True,
            "requires_backend": "ComfyUI",
            "setup_hint": f"Uses the running ComfyUI server at {active_url}.",
        })
        if len(models) >= max_models:
            break
    return models

def _find_comfyui_model(model_name: Optional[str]) -> Optional[str]:
    names = _comfyui_checkpoint_names()
    if not names:
        return None
    if not model_name or model_name in {"texttoimage", "local_diffusion"}:
        return names[0]

    requested = str(model_name).strip()
    requested_path = Path(requested)
    candidates = {
        requested.lower(),
        requested_path.name.lower(),
        requested_path.stem.lower(),
    }
    for ckpt_name in names:
        ckpt_path = Path(ckpt_name)
        ckpt_candidates = {
            ckpt_name.lower(),
            ckpt_path.name.lower(),
            ckpt_path.stem.lower(),
        }
        if candidates & ckpt_candidates:
            return ckpt_name
    return None

def _read_json_file(path: Path) -> Optional[dict]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None

def _looks_like_hf_image_repo(repo_dir: Path) -> bool:
    manifest = _read_json_file(repo_dir / ".moreai-model.json") or {}
    category = str(manifest.get("category") or manifest.get("source_class") or "").lower()
    if "image" in category:
        return True
    if "video" in category or "music" in category or "audio" in category:
        return False

    model_index = _read_json_file(repo_dir / "model_index.json") or {}
    pipeline = str(model_index.get("_class_name") or model_index.get("pipeline_tag") or "").lower()
    if any(token in pipeline for token in ("video", "cogvideo", "wan")):
        return False
    if any(token in pipeline for token in (
        "stablediffusion",
        "stable_diffusion",
        "stablecascade",
        "flux",
        "pixart",
        "kandinsky",
        "ifpipeline",
    )):
        return True

    name = repo_dir.name.lower()
    if any(token in name for token in ("music", "audio", "cogvideo", "wan2", "t2v", "video")):
        return False
    return any(token in name for token in IMAGE_MODEL_HINTS)

def _hf_repo_label(repo_dir: Path) -> str:
    model_index = _read_json_file(repo_dir / "model_index.json") or {}
    name = str(model_index.get("_name_or_path") or "").strip()
    if name:
        return name
    return repo_dir.name.replace("__", "/")

def _scan_more_ai_hf_image_repos(max_models: int = 256) -> list[dict]:
    found: list[dict] = []
    seen: set[str] = set()
    for root in _candidate_more_ai_hf_roots():
        try:
            for repo_dir in sorted([p for p in root.iterdir() if p.is_dir()], key=lambda p: p.name.lower()):
                try:
                    resolved = repo_dir.resolve()
                except Exception:
                    resolved = repo_dir
                key = str(resolved).lower()
                if key in seen or not _looks_like_hf_image_repo(resolved):
                    continue
                seen.add(key)
                found.append({
                    "path": resolved,
                    "source": "more_ai_hf",
                    "label": _hf_repo_label(resolved),
                    "engine": "diffusers",
                    "runnable": False,
                    "requires_backend": "Diffusers / ComfyUI / Automatic1111",
                    "setup_hint": "Downloaded from Hugging Face. Paint can see it, but this repo-style Diffusers model needs a Diffusers, ComfyUI, or Automatic1111 backend before local generation can run it.",
                })
                if len(found) >= max_models:
                    return found
        except Exception:
            continue
    return found

# ── Health ────────────────────────────────────────────────────────────────────

@app.get("/")
def health():
    return {"status": "ok", "service": "image_server", "version": SERVER_VERSION}

# ── GPU detection ─────────────────────────────────────────────────────────────

@app.get("/gpu_info")
def gpu_info():
    info = {"gpu_name": "No GPU detected", "vram_gb": 0.0, "has_cuda": False}
    try:
        import torch
        if torch.cuda.is_available():
            idx = torch.cuda.current_device()
            name = torch.cuda.get_device_name(idx)
            vram = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 3)
            info = {"gpu_name": name, "vram_gb": round(vram, 1), "has_cuda": True}
    except Exception:
        # Try nvidia-smi as fallback
        try:
            import subprocess
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
                capture_output=True, text=True, timeout=5
            )
            if result.returncode == 0:
                parts = result.stdout.strip().split(",")
                if len(parts) >= 2:
                    name = parts[0].strip()
                    vram_str = parts[1].strip().replace(" MiB", "").replace(" MB", "")
                    try:
                        vram = float(vram_str) / 1024
                        info = {"gpu_name": name, "vram_gb": round(vram, 1), "has_cuda": True}
                    except ValueError:
                        pass
        except Exception:
            pass
    return info

@app.get("/generation_capabilities")
def generation_capabilities():
    gpu = gpu_info()
    generation_models = _scan_generation_models()
    comfyui_models = _scan_comfyui_models()
    active_comfyui_url = _active_comfyui_url()
    active_comfyui_pid = _listener_pid_for_url(active_comfyui_url)
    hf_image_repos = _scan_more_ai_hf_image_repos()
    stable_diffusion_cpp_available = False
    try:
        import stable_diffusion_cpp  # noqa: F401
        stable_diffusion_cpp_available = True
    except Exception:
        stable_diffusion_cpp_available = False

    vram = float(gpu.get("vram_gb") or 0.0)
    if gpu.get("has_cuda") and vram >= 10:
        recommended_max_resolution = 1024
    elif gpu.get("has_cuda") and vram >= 6:
        recommended_max_resolution = 768
    else:
        recommended_max_resolution = 512

    comfyui_ready = len(comfyui_models) > 0
    stable_cpp_ready = stable_diffusion_cpp_available and any(
        model.get("source") != "comfyui" for model in generation_models
    )
    backend = "comfyui" if comfyui_ready else ("stable_diffusion_cpp" if stable_cpp_ready else "unavailable")
    return {
        "backend": backend,
        "supports_gguf": stable_diffusion_cpp_available,
        "comfyui_available": active_comfyui_url is not None,
        "comfyui_url": active_comfyui_url or COMFYUI_URL,
        "comfyui_pid": active_comfyui_pid,
        "comfyui_process": _process_label_for_pid(active_comfyui_pid),
        "comfyui_match": "verified_system_stats" if active_comfyui_url else None,
        "comfyui_model_count": len(comfyui_models),
        "comfyui_models": [model["label"] for model in comfyui_models],
        "vram_total_gb": vram,
        "vram_free_gb": None,
        "recommended_max_resolution": recommended_max_resolution,
        "local_generation_model_count": len(generation_models),
        "local_generation_models": [model["label"] for model in generation_models],
        "downloaded_hf_image_model_count": len(hf_image_repos),
        "downloaded_hf_image_models": [model["label"] for model in hf_image_repos],
        "generation_runtime_available": comfyui_ready or stable_cpp_ready,
        "generation_runtime_package": GENERATION_RUNTIME_PACKAGE,
        "runtime_installable": True,
        "recommended_local_model": DEFAULT_LOCAL_GENERATION_MODEL,
        "hardware_local_recommended": bool(gpu.get("has_cuda") and vram >= 4),
    }

# ── Background removal ────────────────────────────────────────────────────────

class ImageRequest(BaseModel):
    image_b64: str
    format: Optional[str] = "png"

@app.post("/remove_background")
def remove_background(req: ImageRequest):
    """
    Remove image background using rembg (U2Net, CPU-compatible).
    Returns PNG with transparent background.
    """
    try:
        from rembg import remove
        img_bytes = base64.b64decode(req.image_b64)
        result = remove(img_bytes)
        b64 = base64.b64encode(result).decode("utf-8")
        return {"image_b64": b64, "format": "png", "success": True}
    except ImportError:
        return {"success": False, "error": "rembg not installed. Run: pip install rembg"}
    except Exception as e:
        return {"success": False, "error": str(e)}

# ── Upscale ───────────────────────────────────────────────────────────────────

class UpscaleRequest(BaseModel):
    image_b64: str
    scale: int = 2   # 2 or 4

@app.post("/upscale")
def upscale_image(req: UpscaleRequest):
    """
    Upscale image using Real-ESRGAN (if available) or Pillow LANCZOS fallback.
    scale=2 uses lightweight model (CPU), scale=4 requires more VRAM.
    """
    try:
        img_bytes = base64.b64decode(req.image_b64)
        from PIL import Image

        # Try Real-ESRGAN first
        try:
            from basicsr.archs.rrdbnet_arch import RRDBNet
            from realesrgan import RealESRGANer
            model_name = "RealESRGAN_x2plus.pth" if req.scale == 2 else "RealESRGAN_x4plus.pth"
            model_path = MODEL_DIR / model_name
            if not model_path.exists():
                raise FileNotFoundError(f"Model not downloaded: {model_name}")
            num_block = 23
            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=req.scale)
            upsampler = RealESRGANer(scale=req.scale, model_path=str(model_path), model=model, tile=0, tile_pad=10, pre_pad=0, half=False)
            import numpy as np, cv2
            img_arr = np.frombuffer(img_bytes, np.uint8)
            img_cv = cv2.imdecode(img_arr, cv2.IMREAD_UNCHANGED)
            output, _ = upsampler.enhance(img_cv, outscale=req.scale)
            _, buf = cv2.imencode(".png", output)
            b64 = base64.b64encode(buf.tobytes()).decode()
            return {"image_b64": b64, "format": "png", "success": True, "method": "realesrgan"}
        except Exception:
            pass

        # Pillow LANCZOS fallback
        img = Image.open(io.BytesIO(img_bytes))
        new_w = img.width * req.scale
        new_h = img.height * req.scale
        upscaled = img.resize((new_w, new_h), Image.LANCZOS)
        buf = io.BytesIO()
        upscaled.save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode()
        return {"image_b64": b64, "format": "png", "success": True, "method": "lanczos"}

    except Exception as e:
        return {"success": False, "error": str(e)}

# ── Adjustments ───────────────────────────────────────────────────────────────

class AdjustRequest(BaseModel):
    image_b64: str
    brightness: float = 0.0   # -100 to +100
    contrast:   float = 0.0
    saturation: float = 0.0
    temperature: float = 0.0  # -100 (cool) to +100 (warm) → hue shift
    sharpness:  float = 0.0   # 0 to +100

@app.post("/adjust")
def adjust_image(req: AdjustRequest):
    """Apply Pillow-based image adjustments and return the result."""
    try:
        from PIL import Image, ImageEnhance, ImageFilter
        img_bytes = base64.b64decode(req.image_b64)
        img = Image.open(io.BytesIO(img_bytes)).convert("RGBA")

        # Work on RGB for adjustments
        rgb = img.convert("RGB")

        # Brightness
        if req.brightness != 0:
            factor = 1.0 + req.brightness / 100.0
            rgb = ImageEnhance.Brightness(rgb).enhance(max(0.0, factor))

        # Contrast
        if req.contrast != 0:
            factor = 1.0 + req.contrast / 100.0
            rgb = ImageEnhance.Contrast(rgb).enhance(max(0.0, factor))

        # Saturation
        if req.saturation != 0:
            factor = 1.0 + req.saturation / 100.0
            rgb = ImageEnhance.Color(rgb).enhance(max(0.0, factor))

        # Temperature (warm/cool — shift red/blue channels)
        if req.temperature != 0:
            import numpy as np
            arr = np.array(rgb, dtype=np.float32)
            shift = req.temperature * 0.3
            arr[:, :, 0] = np.clip(arr[:, :, 0] + shift, 0, 255)  # red
            arr[:, :, 2] = np.clip(arr[:, :, 2] - shift, 0, 255)  # blue
            rgb = Image.fromarray(arr.astype(np.uint8))

        # Sharpness
        if req.sharpness > 0:
            factor = 1.0 + req.sharpness / 50.0
            rgb = ImageEnhance.Sharpness(rgb).enhance(factor)

        # Restore alpha
        if img.mode == "RGBA":
            result = Image.new("RGBA", rgb.size)
            result.paste(rgb)
            r, g, b, _ = img.split()
            result.putalpha(img.split()[3])
        else:
            result = rgb

        buf = io.BytesIO()
        result.save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode()
        return {"image_b64": b64, "format": "png", "success": True}

    except Exception as e:
        return {"success": False, "error": str(e)}

# ── Text-to-image generation ──────────────────────────────────────────────────

class GenerateRequest(BaseModel):
    prompt: str
    negative_prompt: Optional[str] = ""
    model_name: Optional[str] = "texttoimage"
    width: int = 512
    height: int = 512
    steps: int = 24
    cfg_scale: float = 7.0
    seed: Optional[int] = None
    source: Optional[str] = "typed"

def _clamp_generation_size(value: int) -> int:
    value = max(128, min(1536, int(value or 512)))
    return max(128, (value // 8) * 8)

def _slug(value: str) -> str:
    clean = "".join(ch.lower() if ch.isalnum() else "-" for ch in value).strip("-")
    while "--" in clean:
        clean = clean.replace("--", "-")
    return clean[:64] or "image"

def _generation_output_paths(prompt: str, generation_id: str):
    day_dir = GENERATION_DIR / datetime.now().strftime("%Y-%m-%d")
    day_dir.mkdir(parents=True, exist_ok=True)
    stem = f"{datetime.now().strftime('%H%M%S')}_{_slug(prompt)}_{generation_id[:8]}"
    return day_dir / f"{stem}.png", day_dir / f"{stem}.json"

def _find_generation_model(model_name: Optional[str]) -> Optional[Path]:
    if not model_name:
        return None
    raw = Path(model_name)
    if raw.exists():
        if _model_source_for_path(raw) == "comfyui":
            return None
        return raw
    direct = MODEL_DIR / model_name
    if direct.exists():
        return direct
    for ext in IMAGE_MODEL_EXTS:
        candidate = MODEL_DIR / f"{model_name}{ext}"
        if candidate.exists():
            return candidate
    for model in _scan_generation_models():
        if model["source"] == "comfyui":
            continue
        candidate = model["path"]
        if model_name in {"texttoimage", "local_diffusion"}:
            return candidate
        if candidate.stem.lower() == model_name.lower() or candidate.name.lower() == model_name.lower():
            return candidate
    return None

def _try_stable_diffusion_cpp(req: GenerateRequest, out_path: Path, model_path: Path, seed: int):
    from stable_diffusion_cpp import StableDiffusion

    kwargs = {
        "model_path": str(model_path),
        "n_threads": max(1, (os.cpu_count() or 4) - 1),
    }
    sd = StableDiffusion(**kwargs)
    common = {
        "prompt": req.prompt,
        "negative_prompt": req.negative_prompt or "",
        "width": _clamp_generation_size(req.width),
        "height": _clamp_generation_size(req.height),
        "sample_steps": max(1, min(80, int(req.steps or 24))),
        "cfg_scale": float(req.cfg_scale or 7.0),
        "seed": int(seed),
    }
    try:
        images = sd.txt_to_img(**common)
    except TypeError:
        common["steps"] = common.pop("sample_steps")
        images = sd.txt_to_img(**common)

    image = images[0] if isinstance(images, (list, tuple)) else images
    if hasattr(image, "save"):
        image.save(out_path)
        return
    if isinstance(image, bytes):
        out_path.write_bytes(image)
        return
    raise RuntimeError("stable_diffusion_cpp returned an unsupported image result")

def _try_comfyui(req: GenerateRequest, out_path: Path, ckpt_name: str, seed: int):
    active_url = _active_comfyui_url()
    if not active_url:
        raise RuntimeError("ComfyUI is not reachable. Start ComfyUI and refresh the model list.")

    workflow = {
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(seed),
                "steps": max(1, min(80, int(req.steps or 24))),
                "cfg": float(req.cfg_scale or 7.0),
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1.0,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": ckpt_name},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {
                "width": _clamp_generation_size(req.width),
                "height": _clamp_generation_size(req.height),
                "batch_size": 1,
            },
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": req.prompt, "clip": ["4", 1]},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": req.negative_prompt or "", "clip": ["4", 1]},
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "more_ai", "images": ["8", 0]},
        },
    }

    prompt_response = _comfyui_request_json(
        "/prompt",
        {"prompt": workflow, "client_id": str(uuid.uuid4())},
        timeout=10.0,
        base_url=active_url,
    )
    prompt_id = str(prompt_response.get("prompt_id") or "")
    if not prompt_id:
        raise RuntimeError(f"ComfyUI did not return a prompt id: {prompt_response}")

    deadline = time.time() + 300
    last_status = ""
    while time.time() < deadline:
        time.sleep(0.5)
        history = _comfyui_request_json(
            f"/history/{urllib.parse.quote(prompt_id)}",
            timeout=10.0,
            base_url=active_url,
        )
        item = history.get(prompt_id) or history.get(str(prompt_id)) or {}
        status = item.get("status") or {}
        last_status = str(status.get("status_str") or status.get("completed") or last_status)
        outputs = item.get("outputs") or {}
        for output in outputs.values():
            for image in output.get("images", []) or []:
                filename = image.get("filename")
                if not filename:
                    continue
                data = _comfyui_request_bytes(
                    "/view",
                    {
                        "filename": filename,
                        "subfolder": image.get("subfolder", ""),
                        "type": image.get("type", "output"),
                    },
                    base_url=active_url,
                )
                out_path.write_bytes(data)
                return
        if status.get("status_str") == "error":
            raise RuntimeError(f"ComfyUI generation failed: {status}")

    raise RuntimeError(f"ComfyUI generation timed out waiting for image output. Last status: {last_status}")

@app.post("/generate")
def generate_image(req: GenerateRequest):
    started = time.time()
    prompt = req.prompt.strip()
    if not prompt:
        return {"success": False, "error": "Prompt is required."}

    req.width = _clamp_generation_size(req.width)
    req.height = _clamp_generation_size(req.height)
    req.steps = max(1, min(80, int(req.steps or 24)))
    req.cfg_scale = max(1.0, min(20.0, float(req.cfg_scale or 7.0)))

    seed = int(req.seed if req.seed is not None else random.randint(1, 2_147_483_647))
    generation_id = str(uuid.uuid4())
    out_path, meta_path = _generation_output_paths(prompt, generation_id)
    comfyui_model = _find_comfyui_model(req.model_name) if _comfyui_available() else None
    model_path: Optional[Path] = None
    model_identity = comfyui_model or req.model_name or "texttoimage"
    engine = "comfyui" if comfyui_model else "stable_diffusion_cpp"
    warning = None

    try:
        if comfyui_model:
            try:
                _try_comfyui(req, out_path, comfyui_model, seed)
            except Exception as e:
                return {
                    "success": False,
                    "error": f"ComfyUI local image generation failed before producing an image: {e}. Confirm ComfyUI is running and the selected checkpoint can be loaded.",
                    "code": "comfyui_generation_failed",
                }
        else:
            model_path = _find_generation_model(req.model_name)
            if not model_path:
                return {
                    "success": False,
                    "error": "Local image generation is not ready: no compatible local diffusion model is installed. Use an API image model or download a local diffusion model.",
                    "code": "local_image_model_missing",
                }

            try:
                _try_stable_diffusion_cpp(req, out_path, model_path, seed)
            except Exception as e:
                return {
                    "success": False,
                    "error": f"Local image generation failed before producing an image: {e}. Use an API image model or fix the local diffusion runtime.",
                    "code": "local_image_runtime_unavailable",
                }

        duration_ms = int((time.time() - started) * 1000)
        fingerprint = hashlib.sha256(json.dumps({
            "prompt": prompt,
            "negative_prompt": req.negative_prompt or "",
            "model_name": model_identity,
            "model_path": str(model_path) if model_path else comfyui_model,
            "seed": seed,
            "width": req.width,
            "height": req.height,
            "steps": req.steps,
            "cfg_scale": req.cfg_scale,
            "generation_type": "txt2img",
        }, sort_keys=True).encode("utf-8")).hexdigest()
        metadata = {
            "success": True,
            "generation_id": generation_id,
            "generation_type": "txt2img",
            "fingerprint": fingerprint,
            "prompt": prompt,
            "negative_prompt": req.negative_prompt or "",
            "model_name": model_identity,
            "model_path": str(model_path) if model_path else comfyui_model,
            "engine": engine,
            "is_real_model": True,
            "seed": seed,
            "width": req.width,
            "height": req.height,
            "steps": req.steps,
            "cfg_scale": req.cfg_scale,
            "source": req.source or "typed",
            "image_path": str(out_path),
            "metadata_path": str(meta_path),
            "duration_ms": duration_ms,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "warning": warning,
        }
        meta_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
        return metadata
    except Exception as e:
        return {"success": False, "error": str(e)}

# ── Local generation setup ────────────────────────────────────────────────────

@app.post("/install_generation_runtime")
def install_generation_runtime():
    """Install the local diffusion Python runtime when the user opts in."""
    try:
        import stable_diffusion_cpp  # noqa: F401
        return {
            "success": True,
            "already_installed": True,
            "package": GENERATION_RUNTIME_PACKAGE,
            "message": "Local image generation runtime is already installed.",
        }
    except Exception:
        pass

    try:
        proc = subprocess.run(
            [sys.executable, "-m", "pip", "install", "--upgrade", GENERATION_RUNTIME_PACKAGE],
            capture_output=True,
            text=True,
            timeout=1800,
        )
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "error": "Runtime install timed out before completion. Check your internet connection and try again.",
            "package": GENERATION_RUNTIME_PACKAGE,
        }
    except Exception as e:
        return {"success": False, "error": str(e), "package": GENERATION_RUNTIME_PACKAGE}

    if proc.returncode != 0:
        details = (proc.stderr or proc.stdout or "").strip()
        return {
            "success": False,
            "error": f"Runtime install failed: {details[:1200]}",
            "package": GENERATION_RUNTIME_PACKAGE,
        }

    try:
        import stable_diffusion_cpp  # noqa: F401
        return {
            "success": True,
            "already_installed": False,
            "package": GENERATION_RUNTIME_PACKAGE,
            "message": "Local image generation runtime installed.",
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Runtime installed but could not be imported yet: {e}",
            "package": GENERATION_RUNTIME_PACKAGE,
        }

# ── Model download with SSE progress ──────────────────────────────────────────

class DownloadRequest(BaseModel):
    model_name: str  # "u2net" | "realesrgan_x2" | "realesrgan_x4" | "sd15_q4"

MODEL_URLS = {
    "u2net":         ("https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx", "u2net.onnx"),
    "realesrgan_x2": ("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x2plus.pth", "RealESRGAN_x2plus.pth"),
    "realesrgan_x4": ("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", "RealESRGAN_x4plus.pth"),
    "sd15_q4":       ("https://huggingface.co/second-state/stable-diffusion-v1-5-GGUF/resolve/main/stable-diffusion-v1-5-pruned-emaonly-Q4_0.gguf", "stable-diffusion-v1-5-pruned-emaonly-Q4_0.gguf"),
}

@app.post("/download_model")
def download_model(req: DownloadRequest):
    """Download a model and return SSE progress events."""
    import urllib.request

    if req.model_name not in MODEL_URLS:
        return {"success": False, "error": f"Unknown model: {req.model_name}"}

    url, filename = MODEL_URLS[req.model_name]
    dest = MODEL_DIR / filename

    if dest.exists():
        return {"success": True, "message": f"{filename} already downloaded."}

    def generate():
        try:
            with urllib.request.urlopen(url) as resp:
                total = int(resp.headers.get("Content-Length", 0))
                done = 0
                chunk_size = 65536
                with open(dest, "wb") as f:
                    while True:
                        chunk = resp.read(chunk_size)
                        if not chunk:
                            break
                        f.write(chunk)
                        done += len(chunk)
                        pct = int(done * 100 / total) if total else 0
                        yield f"data: {json.dumps({'percent': pct, 'done': done, 'total': total})}\n\n"
            yield f"data: {json.dumps({'percent': 100, 'complete': True, 'path': str(dest)})}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

@app.get("/models")
def list_models():
    """Return which models are already downloaded."""
    models = {}
    used_keys: set[str] = set()
    for name, (url, filename) in MODEL_URLS.items():
        path = MODEL_DIR / filename
        used_keys.add(name)
        models[name] = {"downloaded": path.exists(), "path": str(path) if path.exists() else None}
        if name == DEFAULT_LOCAL_GENERATION_MODEL:
            models[name].update({
                "purpose": "text_to_image",
                "engine": "stable_diffusion_cpp",
                "label": "Stable Diffusion 1.5 Q4 (local)",
                "size_mb": 1570,
                "license": "CreativeML OpenRAIL-M",
                "source_url": url,
            })
    generation_models = _scan_generation_models()
    comfyui_models = _scan_comfyui_models()
    stable_generation_models = [model for model in generation_models if model["source"] != "comfyui"]
    first_model = comfyui_models[0] if comfyui_models else (stable_generation_models[0] if stable_generation_models else None)
    models["texttoimage"] = {
        "downloaded": first_model is not None,
        "path": str(first_model["path"]) if first_model else None,
        "purpose": "text_to_image",
        "engine": first_model.get("engine", "stable_diffusion_cpp") if first_model else "stable_diffusion_cpp",
        "label": first_model["label"] if first_model else "Local Diffusion Model",
        "source": first_model["source"] if first_model else "local",
        "detected": first_model is not None,
        "runnable": bool(first_model.get("runnable", True)) if first_model else False,
        "alias": True,
        "alias_of": str(first_model["path"]) if first_model else None,
    }
    used_keys.add("texttoimage")
    comfyui_names = {str(model["path"]).lower() for model in comfyui_models}
    for model in comfyui_models:
        key = _model_key_for_path(Path(str(model["path"])), used_keys)
        models[key] = {
            "downloaded": True,
            "path": str(model["path"]),
            "purpose": "text_to_image",
            "engine": "comfyui",
            "label": model["label"],
            "source": "comfyui",
            "detected": True,
            "runnable": True,
            "requires_backend": model["requires_backend"],
            "setup_hint": model["setup_hint"],
        }
    for model in generation_models:
        path = model["path"]
        key = _model_key_for_path(path, used_keys)
        if model["source"] == "comfyui":
            if str(path.name).lower() in comfyui_names or str(path).lower() in comfyui_names:
                continue
            models[key] = {
                "downloaded": True,
                "path": str(path),
                "purpose": "text_to_image",
                "engine": "comfyui",
                "label": model["label"],
                "source": "comfyui",
                "detected": True,
                "runnable": False,
                "requires_backend": "ComfyUI running",
                "setup_hint": f"Start ComfyUI, then More AI can use this checkpoint through {COMFYUI_URL}.",
            }
            continue
        models[key] = {
            "downloaded": True,
            "path": str(path),
            "purpose": "text_to_image",
            "engine": "stable_diffusion_cpp",
            "label": model["label"],
            "source": model["source"],
            "detected": True,
            "runnable": True,
        }
    for model in _scan_more_ai_hf_image_repos():
        path = model["path"]
        key = _model_key_for_path(path, used_keys)
        models[key] = {
            "downloaded": True,
            "path": str(path),
            "purpose": "text_to_image",
            "engine": model["engine"],
            "label": model["label"],
            "source": model["source"],
            "detected": True,
            "runnable": model["runnable"],
            "requires_backend": model["requires_backend"],
            "setup_hint": model["setup_hint"],
        }
    return models

# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=7871, log_level="warning")