Update app.py

app.py

@@ -1,8 +1,11 @@
 """
-MiniCPM Forge — Unified Edge AI Showcase
-Five tiny models, one interface. llama.cpp-powered.
+MiniCPM Forge — Hybrid Multi-Model Showcase
+Two backends, one interface:
+  • Vision  (v46, v46t) → transformers + ZeroGPU  (same as the working single-model space)
+  • Text    (cpm5, cpm41) → llama-cpp-python        (CPU/GPU, works on free tier)
+  • Omni   (o45)          → API placeholder
 
-build-small-hackathon 2026 · Chris4K · ki-fusion-labs.de
+build-small-hackathon 2026 · Chris4K  
 """
 from __future__ import annotations
 
@@ -11,42 +14,35 @@ import base64
 import json
 import os
 import pathlib
+import re
 import threading
 import time
 import uuid
 from io import BytesIO
-from typing import Optional
+from typing import Generator, Optional
 
 from PIL import Image
 from fastapi import Request
 from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 from gradio import Server
-from huggingface_hub import hf_hub_download
 
 # ─────────────────────────────────────────────────────────────────────────────
-# ZeroGPU / HF Spaces GPU decorator (optional)
+# ZeroGPU (optional — works on HF Spaces GPU / ZeroGPU tiers)
 # ─────────────────────────────────────────────────────────────────────────────
 try:
     import spaces  # type: ignore
+    HAS_SPACES_GPU = True
 except ImportError:
-    class _NoSpaces:
+    HAS_SPACES_GPU = False
+    class _FakeSpaces:
         @staticmethod
         def GPU(duration: int = 120):
-            def _wrap(fn):
-                return fn
+            def _wrap(fn): return fn
             return _wrap
-    spaces = _NoSpaces()  # type: ignore
+    spaces = _FakeSpaces()  # type: ignore
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Model registry  (filenames verified June 2026 against HF repo trees)
-# ─────────────────────────────────────────────────────────────────────────────
-# Naming quirk: openbmb GGUF repos replace dots with underscores in version
-# numbers:  4.6  →  4_6  in filenames, but the repo slug itself keeps dots.
-#
-# Confirmed sources:
-#   LM   : https://github.com/OpenBMB/MiniCPM-V-Apps (official mobile app)
-#   mmproj: same + HF llama-cpp-python auto-snippet
-#   CPM4.1-8B: openbmb/MiniCPM4.1-8B model card llama.cpp section
+# Model registry
 # ─────────────────────────────────────────────────────────────────────────────
 MODELS: dict[str, dict] = {
     "v46": {
@@ -54,286 +50,304 @@ MODELS: dict[str, dict] = {
         "name": "MiniCPM-V 4.6",
         "tag": "Vision · OCR",
         "color": "#06b6d4",
-        "repo": "openbmb/MiniCPM-V-4.6-gguf",
-        "file": "MiniCPM-V-4_6-Q4_K_M.gguf",   # underscore: 4_6, NOT 4.6
-        "mmproj": "mmproj-model-f16.gguf",        # short name, no model prefix
+        "backend": "transformers",
+        "hf_id": "openbmb/MiniCPM-V-4.6",
+        "thinking": False,
         "ctx": 8192,
         "vision": True,
-        "thinking": False,
-        "handler": "minicpmv",
     },
     "v46t": {
         "id": "v46t",
         "name": "MiniCPM-V 4.6-T",
         "tag": "Vision · Thinking",
         "color": "#a855f7",
-        "repo": "openbmb/MiniCPM-V-4.6-Thinking-gguf",
-        "file": "MiniCPM-V-4_6-Thinking-Q4_K_M.gguf",  # same underscore convention
-        "mmproj": "mmproj-model-f16.gguf",
+        "backend": "transformers",
+        "hf_id": "openbmb/MiniCPM-V-4.6-Thinking",
+        "thinking": True,
         "ctx": 8192,
         "vision": True,
-        "thinking": True,
-        "handler": "minicpmv",
     },
     "cpm5": {
         "id": "cpm5",
         "name": "MiniCPM5-1B",
         "tag": "⚡ Ultra-light",
         "color": "#22c55e",
+        "backend": "llama",
         "repo": "openbmb/MiniCPM5-1B-GGUF",
-        "file": "MiniCPM5-1B-Q4_K_M.gguf",  # ✓ confirmed working
+        "file": "MiniCPM5-1B-Q4_K_M.gguf",      # ✓ confirmed working
         "ctx": 4096,
         "vision": False,
         "thinking": False,
-        "handler": "text",
     },
     "cpm41": {
         "id": "cpm41",
         "name": "MiniCPM4.1-8B",
         "tag": "🧠 Reasoning",
         "color": "#f97316",
-        "repo": "openbmb/MiniCPM4.1-8B-GGUF",       # .1 in repo slug
-        "file": "MiniCPM4.1-8B-Q4_K_M.gguf",        # .1 in filename (dots kept here)
+        "backend": "llama",
+        "repo": "openbmb/MiniCPM4.1-8B-GGUF",
+        "file": "MiniCPM4.1-8B-Q4_K_M.gguf",
         "ctx": 16384,
         "vision": False,
         "thinking": True,
-        "handler": "text",
     },
     "o45": {
         "id": "o45",
         "name": "MiniCPM-o 4.5",
         "tag": "🌐 Omni",
         "color": "#ec4899",
-        "api": True,
+        "backend": "api",
         "ctx": 8192,
         "vision": True,
         "thinking": False,
-        "handler": "api",
-        "note": "Served via OpenBMB API — no local download required",
     },
 }
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Model loading — lazy, thread-safe, cached
+# Shared state
 # ─────────────────────────────────────────────────────────────────────────────
 CACHE_DIR = pathlib.Path(os.environ.get("HF_HOME", "/tmp/hf_cache")) / "forge"
-_loaded: dict[str, object] = {}
+_load_status: dict[str, str] = {}
 _load_lock = threading.Lock()
-_load_status: dict[str, str] = {}  # model_id → "idle"|"loading"|"ready"|"error:<msg>"
 
+# llama-cpp loaded models
+_llama_models: dict[str, object] = {}
 
-def _get_vision_handler(mmproj_path: str):
-    """Return the best available llama-cpp vision handler for MiniCPM-V 4.6.
+# transformers loaded models (processor + model per id)
+_tr_processors: dict[str, object] = {}
+_tr_models: dict[str, object] = {}
 
-    llama-cpp-python >= 0.3.9 ships MiniCPMv2_6ChatHandler which is also
-    compatible with 4.6 at the GGUF level (the new merger architecture is
-    abstracted away internally by llama.cpp >= b9049).
+# ─────────────────────────────────────────────────────────────────────────────
+# Text normalisation (from official openbmb demo)
+# ─────────────────────────────────────────────────────────────────────────────
+_NORM_PATTERN = re.compile(
+    r'(```[\s\S]*?```|`[^`]+`|\$\$[\s\S]*?\$\$|\$[^$]+\$'
+    r'|\\\([\s\S]*?\\\)|\\\[[\s\S]*?\\\])'
+    r'|(?<!\\)(?:\\r\\n|\\[nr])'
+)
+
+def normalize_response_text(text: str) -> str:
+    if not isinstance(text, str) or "\\" not in text:
+        return text
+    return _NORM_PATTERN.sub(lambda m: m.group(1) or '\n', text)
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Backend A: transformers (vision models, runs via ZeroGPU when available)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _load_transformers(model_id: str) -> None:
+    """Load processor + model into _tr_processors/_tr_models. Thread-safe."""
+    if model_id in _tr_models:
+        return
+    with _load_lock:
+        if model_id in _tr_models:
+            return
+        _load_status[model_id] = "loading"
+        cfg = MODELS[model_id]
+        hf_id = cfg["hf_id"]
+        print(f"[forge] Loading transformers model: {hf_id}")
+        try:
+            import torch
+            from transformers import AutoProcessor, AutoModelForImageTextToText  # type: ignore
+
+            processor = AutoProcessor.from_pretrained(hf_id, trust_remote_code=True)
+
+            if torch.cuda.is_available():
+                model = AutoModelForImageTextToText.from_pretrained(
+                    hf_id,
+                    torch_dtype=torch.bfloat16,
+                    attn_implementation="sdpa",
+                    trust_remote_code=True,
+                    device_map="cuda",
+                ).eval()
+            else:
+                # CPU fallback — slow but functional for demos
+                model = AutoModelForImageTextToText.from_pretrained(
+                    hf_id,
+                    torch_dtype=torch.float32,
+                    trust_remote_code=True,
+                    device_map="cpu",
+                ).eval()
+
+            _tr_processors[model_id] = processor
+            _tr_models[model_id] = model
+            _load_status[model_id] = "ready"
+            print(f"[forge] ✓ transformers {model_id} ready")
+        except Exception as exc:
+            _load_status[model_id] = f"error:{exc}"
+            print(f"[forge] ✗ transformers {model_id} failed: {exc}")
+            raise
+
+
+@spaces.GPU(duration=120)
+def _run_transformers(
+    model_id: str,
+    messages: list,
+    params: dict,
+) -> Generator[str, None, None]:
     """
-    try:
-        from llama_cpp.llama_chat_format import MiniCPMv2_6ChatHandler  # type: ignore
-        return MiniCPMv2_6ChatHandler(clip_model_path=mmproj_path)
-    except (ImportError, AttributeError):
-        pass
-    try:
-        from llama_cpp.llama_chat_format import Llava16ChatHandler  # type: ignore
-        return Llava16ChatHandler(clip_model_path=mmproj_path)
-    except (ImportError, AttributeError):
-        pass
-    raise RuntimeError(
-        "No compatible vision chat handler found. "
-        "Run:  pip install 'llama-cpp-python>=0.3.9'"
+    Sync generator — yields *delta* text chunks.
+    Decorated with @spaces.GPU so it runs on ZeroGPU when available;
+    falls back to CPU silently when spaces is not installed.
+    """
+    import torch
+    from transformers import TextIteratorStreamer  # type: ignore
+
+    processor = _tr_processors[model_id]
+    model     = _tr_models[model_id]
+    thinking  = params.get("thinking_mode", False)
+
+    is_video = any(
+        it.get("type") == "video"
+        for msg in messages
+        for it in (msg.get("content") or [])
     )
 
+    with torch.no_grad():
+        inputs = processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            enable_thinking=thinking,
+            processor_kwargs={
+                "downsample_mode": "16x",
+                "max_slice_nums": 1 if is_video else 9,
+                "use_image_id": not is_video,
+            },
+        ).to(model.device)
+
+    if model.device.type == "cuda":
+        import torch as _torch
+        for k, v in inputs.items():
+            if isinstance(v, _torch.Tensor) and _torch.is_floating_point(v):
+                inputs[k] = v.to(dtype=_torch.bfloat16)
+
+    streamer = TextIteratorStreamer(
+        processor.tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True,
+        timeout=60.0,
+    )
 
-def _hub_download_with_glob(
-    repo_id: str,
-    filename: str,
-    local_dir: str,
-) -> str:
-    """Download filename from repo_id.
+    gen_kw = {
+        **inputs,
+        "max_new_tokens": params.get("max_tokens", 1024),
+        "do_sample": True,
+        "temperature": params.get("temperature", 0.7),
+        "top_p": params.get("top_p", 0.8),
+        "top_k": int(params.get("top_k", 100)),
+        "streamer": streamer,
+        "downsample_mode": "16x",
+    }
 
-    If the exact filename 404s, falls back to glob-matching the repo file list
-    so minor naming inconsistencies across openbmb releases don't break things.
-    """
-    import fnmatch
-    from huggingface_hub import list_repo_files  # type: ignore
+    t = threading.Thread(target=model.generate, kwargs=gen_kw, daemon=True)
+    t.start()
 
-    local_dir_path = pathlib.Path(local_dir)
-    local_dir_path.mkdir(parents=True, exist_ok=True)
+    for chunk in streamer:
+        yield normalize_response_text(chunk)
 
-    # 1) Try exact name first
-    try:
-        return hf_hub_download(repo_id=repo_id, filename=filename,
-                                local_dir=local_dir)
-    except Exception as exact_err:
-        pass  # fall through to glob
-
-    # 2) Build a glob from the filename: replace version numbers with *
-    #    e.g. MiniCPM-V-4_6-Q4_K_M.gguf → *Q4_K_M*.gguf
-    import re
-    base = pathlib.Path(filename).name
-    # Keep quant type suffix as anchor
-    quant_match = re.search(r'(Q\d_K_[MS]|Q\d_\d|F16|BF16)', base)
-    glob_pattern = f"*{quant_match.group(1)}*.gguf" if quant_match else f"*{base}*"
-
-    all_files = list(list_repo_files(repo_id))
-    candidates = [f for f in all_files if fnmatch.fnmatch(f, glob_pattern)
-                  and not f.endswith(".md") and not f.endswith(".json")]
+    t.join(timeout=10)
 
-    if not candidates:
-        raise FileNotFoundError(
-            f"Exact file {filename!r} not found in {repo_id}, "
-            f"and glob {glob_pattern!r} matched nothing. "
-            f"Available GGUF files: {[f for f in all_files if f.endswith('.gguf')]}"
-        )
 
-    # Prefer Q4_K_M over Q4_K_S etc if multiple hits
-    best = next((f for f in candidates if "Q4_K_M" in f), candidates[0])
-    print(f"  → glob fallback: using {best!r} instead of {filename!r}")
-    return hf_hub_download(repo_id=repo_id, filename=best, local_dir=local_dir)
+async def _stream_transformers(
+    model_id: str,
+    messages: list,
+    params: dict,
+    loop: asyncio.AbstractEventLoop,
+):
+    """Async generator: bridges sync _run_transformers → async SSE."""
+    queue: asyncio.Queue = asyncio.Queue(maxsize=256)
+
+    def _worker():
+        try:
+            for chunk in _run_transformers(model_id, messages, params):
+                loop.call_soon_threadsafe(queue.put_nowait, chunk)
+        except Exception as exc:
+            loop.call_soon_threadsafe(queue.put_nowait, f"\n\n[⚠ {exc}]")
+        finally:
+            loop.call_soon_threadsafe(queue.put_nowait, None)
+
+    loop.run_in_executor(None, _worker)
+
+    while True:
+        token = await queue.get()
+        if token is None:
+            break
+        yield token
+
 
+# ─────────────────────────────────────────────────────────────────────────────
+# Backend B: llama-cpp (text models)
+# ─────────────────────────────────────────────────────────────────────────────
 
-def _hub_download_mmproj(repo_id: str, mmproj_hint: str, local_dir: str) -> str:
-    """Download mmproj file; glob-fallback to any *mmproj*f16*.gguf."""
+def _hub_download_robust(repo_id: str, filename: str, local_dir: str) -> str:
+    """hf_hub_download with glob fallback for filename drift."""
     import fnmatch
-    from huggingface_hub import list_repo_files  # type: ignore
+    from huggingface_hub import hf_hub_download, list_repo_files  # type: ignore
 
+    pathlib.Path(local_dir).mkdir(parents=True, exist_ok=True)
     try:
-        return hf_hub_download(repo_id=repo_id, filename=mmproj_hint,
-                                local_dir=local_dir)
+        return hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
     except Exception:
         pass
-
-    all_files = list(list_repo_files(repo_id))
-    candidates = [f for f in all_files
-                  if "mmproj" in f.lower() and f.endswith(".gguf")
-                  and "iOS" not in f and "ios" not in f]  # skip mobile-only variants
+    # Glob fallback
+    quant = re.search(r'(Q\d_K_[MS]|Q\d_\d|F16|BF16)', filename)
+    pat = f"*{quant.group(1)}*.gguf" if quant else f"*{pathlib.Path(filename).stem}*"
+    candidates = [f for f in list_repo_files(repo_id)
+                  if fnmatch.fnmatch(f, pat) and f.endswith(".gguf")]
     if not candidates:
         raise FileNotFoundError(
-            f"No mmproj file found in {repo_id}. "
-            f"Files: {[f for f in all_files if f.endswith('.gguf')]}"
+            f"No file matching {pat!r} in {repo_id}. "
+            f"GGUFs: {[f for f in list_repo_files(repo_id) if f.endswith('.gguf')]}"
         )
-    best = next((f for f in candidates if "f16" in f.lower()), candidates[0])
-    print(f"  → mmproj glob fallback: {best!r}")
+    best = next((f for f in candidates if "Q4_K_M" in f), candidates[0])
+    print(f"  → glob fallback: {best!r}")
     return hf_hub_download(repo_id=repo_id, filename=best, local_dir=local_dir)
 
 
-def load_model(model_id: str) -> Optional[object]:
-    """Download and load a GGUF model. Returns the Llama instance or None on error."""
-    if model_id in _loaded:
-        return _loaded[model_id]
-
-    cfg = MODELS.get(model_id)
-    if not cfg or cfg.get("api"):
-        return None
-
+def _load_llama(model_id: str) -> None:
+    """Load a text-only GGUF model via llama-cpp-python."""
+    if model_id in _llama_models:
+        return
     with _load_lock:
-        if model_id in _loaded:
-            return _loaded[model_id]
-
+        if model_id in _llama_models:
+            return
         _load_status[model_id] = "loading"
+        cfg = MODELS[model_id]
         local_dir = str(CACHE_DIR / model_id)
-
+        print(f"[forge] Downloading LM: {cfg['repo']} / {cfg['file']}")
         try:
             from llama_cpp import Llama  # type: ignore
 
-            CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
-            print(f"[forge] Downloading LM: {cfg['repo']} / {cfg['file']}")
-            model_path = _hub_download_with_glob(
-                repo_id=cfg["repo"],
-                filename=cfg["file"],
-                local_dir=local_dir,
+            model_path = _hub_download_robust(cfg["repo"], cfg["file"], local_dir)
+            n_gpu = int(os.environ.get("N_GPU_LAYERS", "-1"))
+            llm = Llama(
+                model_path=model_path,
+                n_ctx=cfg["ctx"],
+                n_gpu_layers=n_gpu,
+                verbose=False,
             )
-
-            n_gpu_layers = int(os.environ.get("N_GPU_LAYERS", "-1"))
-
-            if cfg.get("vision"):
-                print(f"[forge] Downloading mmproj: {cfg['mmproj']}")
-                mmproj_path = _hub_download_mmproj(
-                    repo_id=cfg["repo"],
-                    mmproj_hint=cfg["mmproj"],
-                    local_dir=local_dir,
-                )
-                chat_handler = _get_vision_handler(mmproj_path)
-                llm = Llama(
-                    model_path=model_path,
-                    chat_handler=chat_handler,
-                    n_ctx=cfg["ctx"],
-                    n_gpu_layers=n_gpu_layers,
-                    verbose=False,
-                )
-            else:
-                llm = Llama(
-                    model_path=model_path,
-                    n_ctx=cfg["ctx"],
-                    n_gpu_layers=n_gpu_layers,
-                    verbose=False,
-                )
-
-            _loaded[model_id] = llm
+            _llama_models[model_id] = llm
             _load_status[model_id] = "ready"
-            print(f"[forge] ✓ {model_id} ready")
-            return llm
-
+            print(f"[forge] ✓ llama {model_id} ready")
         except Exception as exc:
             _load_status[model_id] = f"error:{exc}"
-            print(f"[forge] ✗ {model_id} failed: {exc}")
-            return None
-
+            print(f"[forge] ✗ llama {model_id} failed: {exc}")
+            raise
 
-# ─────────────────────────────────────────────────────────────────────────────
-# Inference helpers
-# ─────────────────────────────────────────────────────────────────────────────
-
-def _image_to_data_uri(b64_string: str) -> str:
-    """Convert raw base64 image string to data: URI."""
-    # Detect mime type from magic bytes
-    raw = base64.b64decode(b64_string[:32])
-    if raw[:8] == b"\x89PNG\r\n\x1a\n":
-        mime = "image/png"
-    elif raw[:3] == b"\xff\xd8\xff":
-        mime = "image/jpeg"
-    else:
-        mime = "image/png"
-    return f"data:{mime};base64,{b64_string}"
-
-
-def _build_messages(
-    message: str,
-    history: list,
-    image_b64: Optional[str] = None,
-) -> list[dict]:
-    """Build OpenAI-style message list from chat history + current turn."""
-    messages: list[dict] = []
-
-    for turn in history or []:
-        user_text = turn.get("user", "")
-        asst_text = turn.get("assistant", "")
-        if user_text:
-            messages.append({"role": "user", "content": user_text})
-        if asst_text:
-            messages.append({"role": "assistant", "content": asst_text})
-
-    # Current user turn — inject image if provided
-    if image_b64:
-        data_uri = _image_to_data_uri(image_b64)
-        content: list[dict] = [
-            {"type": "image_url", "image_url": {"url": data_uri}},
-            {"type": "text", "text": message or "Describe this image."},
-        ]
-        messages.append({"role": "user", "content": content})
-    else:
-        messages.append({"role": "user", "content": message})
 
-    return messages
+async def _stream_llama(
+    llm,
+    messages: list,
+    params: dict,
+    loop: asyncio.AbstractEventLoop,
+):
+    """Async generator: bridges sync llama-cpp stream → async SSE."""
+    queue: asyncio.Queue = asyncio.Queue(maxsize=256)
 
-
-async def _stream_tokens(llm, messages: list, params: dict, loop: asyncio.AbstractEventLoop):
-    """Async generator: yields tokens from a sync llama.cpp inference run."""
-    queue: asyncio.Queue = asyncio.Queue(maxsize=128)
-
-    def _run() -> None:
+    def _worker():
         try:
             output = llm.create_chat_completion(
                 messages=messages,
@@ -341,7 +355,7 @@ async def _stream_tokens(llm, messages: list, params: dict, loop: asyncio.Abstra
                 max_tokens=params.get("max_tokens", 1024),
                 temperature=params.get("temperature", 0.7),
                 top_p=params.get("top_p", 0.8),
-                top_k=params.get("top_k", 100),
+                top_k=int(params.get("top_k", 100)),
                 repeat_penalty=params.get("repeat_penalty", 1.05),
             )
             for chunk in output:
@@ -349,11 +363,11 @@ async def _stream_tokens(llm, messages: list, params: dict, loop: asyncio.Abstra
                 if delta:
                     loop.call_soon_threadsafe(queue.put_nowait, delta)
         except Exception as exc:
-            loop.call_soon_threadsafe(queue.put_nowait, f"\n\n[⚠ Error: {exc}]")
+            loop.call_soon_threadsafe(queue.put_nowait, f"\n\n[⚠ {exc}]")
         finally:
-            loop.call_soon_threadsafe(queue.put_nowait, None)  # sentinel
+            loop.call_soon_threadsafe(queue.put_nowait, None)
 
-    loop.run_in_executor(None, _run)
+    loop.run_in_executor(None, _worker)
 
     while True:
         token = await queue.get()
@@ -362,6 +376,46 @@ async def _stream_tokens(llm, messages: list, params: dict, loop: asyncio.Abstra
         yield token
 
 
+# ─────────────────────────────────────────────────────────────────────────────
+# Shared message builder
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _build_messages(
+    message: str,
+    history: list,
+    image_b64: Optional[str],
+    backend: str,
+) -> list[dict]:
+    msgs: list[dict] = []
+
+    for turn in history or []:
+        if turn.get("user"):
+            msgs.append({"role": "user",
+                          "content": [{"type": "text", "text": turn["user"]}]})
+        if turn.get("assistant"):
+            msgs.append({"role": "assistant",
+                          "content": [{"type": "text", "text": turn["assistant"]}]})
+
+    content: list[dict] = []
+
+    if image_b64:
+        if backend == "transformers":
+            # transformers expects a PIL Image object
+            img_bytes = base64.b64decode(image_b64)
+            img = Image.open(BytesIO(img_bytes)).convert("RGB")
+            content.append({"type": "image", "image": img})
+        else:
+            # llama-cpp expects a data: URI
+            raw = base64.b64decode(image_b64[:32])
+            mime = "image/png" if raw[:8] == b"\x89PNG\r\n\x1a\n" else "image/jpeg"
+            content.append({"type": "image_url",
+                             "image_url": {"url": f"data:{mime};base64,{image_b64}"}})
+
+    content.append({"type": "text", "text": message or "Describe this image."})
+    msgs.append({"role": "user", "content": content})
+    return msgs
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # Gradio Server
 # ─────────────────────────────────────────────────────────────────────────────
@@ -370,114 +424,116 @@ demo = Server()
 
 @demo.get("/", response_class=HTMLResponse)
 async def homepage():
-    html_path = pathlib.Path(__file__).parent / "index.html"
-    with open(html_path, encoding="utf-8") as f:
-        return f.read()
+    html = pathlib.Path(__file__).parent / "index.html"
+    return html.read_text(encoding="utf-8")
 
 
 @demo.get("/api/models")
 async def api_models():
-    """Return model registry (without sensitive fields)."""
-    safe = {}
+    out = {}
     for mid, cfg in MODELS.items():
-        safe[mid] = {
-            k: cfg[k]
-            for k in ("id", "name", "tag", "color", "ctx", "vision", "thinking")
-            if k in cfg
-        }
-        safe[mid]["status"] = _load_status.get(mid, "idle")
-        safe[mid]["api"] = cfg.get("api", False)
-    return JSONResponse(safe)
+        out[mid] = {k: cfg[k] for k in ("id", "name", "tag", "color", "ctx", "vision", "thinking")
+                    if k in cfg}
+        out[mid]["status"] = _load_status.get(mid, "idle")
+        out[mid]["api"]    = cfg.get("backend") == "api"
+        out[mid]["backend"] = cfg.get("backend", "llama")
+    return JSONResponse(out)
 
 
 @demo.post("/api/load")
 async def api_load(request: Request):
-    """Background-load a model and return status."""
     data = await request.json()
-    model_id = data.get("model_id", "")
-    if model_id not in MODELS:
+    mid = data.get("model_id", "")
+    cfg = MODELS.get(mid)
+    if not cfg:
         return JSONResponse({"error": "unknown model"}, status_code=400)
-    if MODELS[model_id].get("api"):
-        return JSONResponse({"status": "api", "message": "No local download needed."})
+    if cfg["backend"] == "api":
+        return JSONResponse({"status": "api"})
 
-    # Trigger load in background thread
-    if _load_status.get(model_id) not in ("loading", "ready"):
+    # Kick off load in a background thread if not already running
+    current = _load_status.get(mid, "idle")
+    if current not in ("loading", "ready"):
         loop = asyncio.get_event_loop()
-        loop.run_in_executor(None, load_model, model_id)
+        loader = _load_transformers if cfg["backend"] == "transformers" else _load_llama
+        loop.run_in_executor(None, loader, mid)
 
-    return JSONResponse({"status": _load_status.get(model_id, "loading")})
+    return JSONResponse({"status": _load_status.get(mid, "loading")})
 
 
 @demo.post("/stream/chat")
 async def stream_chat(request: Request):
-    """SSE streaming inference endpoint."""
-    data = await request.json()
-    model_id = data.get("model_id", "v46")
+    data      = await request.json()
+    mid       = data.get("model_id", "cpm5")
     message   = data.get("message", "")
     history   = data.get("history", [])
     image_b64 = data.get("image_b64")
     params    = data.get("params", {})
 
-    cfg = MODELS.get(model_id)
+    cfg = MODELS.get(mid)
     if not cfg:
         return JSONResponse({"error": "unknown model"}, status_code=400)
 
-    # ── API-mode models (MiniCPM-o 4.5) ──────────────────────────────────
-    if cfg.get("api"):
+    backend = cfg.get("backend", "llama")
+
+    # ── API mode ──────────────────────────────────────────────────────────────
+    if backend == "api":
         async def _api_sse():
-            # Placeholder: replace with real openbmb API endpoint when available
-            yield f"data: {json.dumps({'token': '⚡ MiniCPM-o 4.5 (API mode) — '})}\n\n"
-            yield f"data: {json.dumps({'token': 'Connect your OpenBMB API key via the settings panel.'})}\n\n"
+            yield f"data: {json.dumps({'token': '🌐 MiniCPM-o 4.5 API mode — set OPENBMB_API_KEY in Space secrets.'})}\n\n"
             yield f"data: {json.dumps({'done': True})}\n\n"
-        return StreamingResponse(
-            _api_sse(),
-            media_type="text/event-stream",
-            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
-        )
-
-    # ── llama.cpp inference ───────────────────────────────────────────────
-    llm = load_model(model_id)
-    if llm is None:
-        status = _load_status.get(model_id, "unknown")
-        err_msg = f"Model '{model_id}' not ready (status: {status}). Load it first."
+        return StreamingResponse(_api_sse(), media_type="text/event-stream",
+                                  headers={"Cache-Control": "no-cache"})
 
+    # ── Check model is loaded ─────────────────────────────────────────────────
+    store = _tr_models if backend == "transformers" else _llama_models
+    if mid not in store:
+        msg = f"Model '{mid}' not loaded (status: {_load_status.get(mid, 'idle')}). Click ⬇ LOAD first."
         async def _err_sse():
-            yield f"data: {json.dumps({'token': f'⚠ {err_msg}'})}\n\n"
+            yield f"data: {json.dumps({'token': f'⚠ {msg}'})}\n\n"
             yield f"data: {json.dumps({'done': True})}\n\n"
         return StreamingResponse(_err_sse(), media_type="text/event-stream",
                                   headers={"Cache-Control": "no-cache"})
 
-    messages = _build_messages(message, history, image_b64)
-    loop = asyncio.get_event_loop()
-    start_time = time.monotonic()
-    token_count = [0]
+    messages = _build_messages(message, history, image_b64, backend)
+    loop     = asyncio.get_event_loop()
+    t0       = time.monotonic()
+    n_tok    = [0]
+
+    async def sse_gen():
+        if backend == "transformers":
+            gen = _stream_transformers(mid, messages, params, loop)
+        else:
+            gen = _stream_llama(_llama_models[mid], messages, params, loop)
+
+        async for token in gen:
+            n_tok[0] += 1
+            elapsed = time.monotonic() - t0
+            speed   = round(n_tok[0] / elapsed, 1) if elapsed > 0 else 0
+            yield f"data: {json.dumps({'token': token, 'speed': speed, 'n': n_tok[0]})}\n\n"
 
-    async def sse_generator():
-        async for token in _stream_tokens(llm, messages, params, loop):
-            token_count[0] += 1
-            elapsed = time.monotonic() - start_time
-            speed = round(token_count[0] / elapsed, 1) if elapsed > 0 else 0
-            payload = json.dumps({"token": token, "speed": speed, "n": token_count[0]})
-            yield f"data: {payload}\n\n"
-        yield f"data: {json.dumps({'done': True, 'total': token_count[0]})}\n\n"
+        yield f"data: {json.dumps({'done': True, 'total': n_tok[0]})}\n\n"
 
     return StreamingResponse(
-        sse_generator(),
+        sse_gen(),
         media_type="text/event-stream",
-        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"},
+        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
+                 "Connection": "keep-alive"},
     )
 
 
 @demo.get("/health")
 async def health():
-    loaded = [mid for mid, s in _load_status.items() if s == "ready"]
-    return JSONResponse({"status": "ok", "loaded_models": loaded})
+    return JSONResponse({
+        "status": "ok",
+        "backends": {
+            "transformers": list(_tr_models.keys()),
+            "llama": list(_llama_models.keys()),
+        },
+        "spaces_gpu": HAS_SPACES_GPU,
+    })
 
 
 # ─────────────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
-    # Pre-load the smallest model on startup (optional)
-    # threading.Thread(target=load_model, args=("v46",), daemon=True).start()
     demo.launch(
         server_name="0.0.0.0",
         server_port=int(os.environ.get("PORT", 7860)),