embeddings: multimodal (text+image) support + correct VRAM tracking/eviction

Wire up real image embeddings and make embedding models first-class in the
VRAM lifecycle.

- api/embeddings.py: detect CLIP/SigLIP dual encoders and drive them through
  transformers get_text_features/get_image_features so text and images share
  one projected space (ST path kept for repos shipping a native recipe).
  request.image is now actually read (URL/data-URI/path/base64), vectors are
  appended after the text ones, and text-only models return a clear 400.
  Handle the transformers 5.x pooled-output return shape.
- Wrap the loaded model in _EmbeddingModel (unpacks as (backend, model) but
  exposes cleanup()) and register it via add_model() + record_vram_delta() on
  the request path, so it is measured, LRU-tracked, and cleanly evicted like
  every other model type instead of leaking as a bare tuple.
- admin: fix the model-load button for embeddings (was routed to the diffusers
  loader) to use _load_embedding_model, matching the request path.
- admin: backfill used_vram_gb after a download completes, since the entry is
  saved before the weights exist on disk; factor the estimate into a shared
  _estimate_used_vram_gb helper. A freshly-downloaded CLIP/SigLIP entry now
  always carries a reasonable estimate so pre-load eviction sizes correctly.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EPLnsRpNBzWCHLgkXATqRz

codai/admin/routes.py

@@ -671,6 +671,14 @@ def _run_download_thread(session_id: str, model_id: str, file_pattern: str, pq):
                 if tail:
                     detail += f". Last output: {tail}"
                 push({"type": "error", "message": detail})
+        # On a successful download, ensure any config entry for this model carries
+        # a used_vram_gb estimate — the save that registered it ran before the
+        # weights existed on disk, so its own auto-estimate came back empty.
+        if terminal == "done" and session_id not in _download_cancelled:
+            try:
+                _backfill_vram_estimate(model_id, file_pattern)
+            except Exception as _bf_err:
+                print(f"[download] used_vram_gb backfill failed for '{model_id}': {_bf_err}")
     finally:
         _download_cancelled.discard(session_id)
 
@@ -984,6 +992,60 @@ def _basename_key(key: str) -> str:
     return _os.path.basename(key) if ("/" in key or _os.sep in key) else key
 
 
+def _estimate_used_vram_gb(path: str):
+    """Best-effort VRAM estimate (GB) for a model path/HF id, or None.
+
+    Local weight file → its size; HF repo id → the cached shard total. A small
+    multiplier covers runtime overhead (KV/activations/VAE spike). Deliberately
+    on the conservative (slightly high) side: over-estimating makes eviction free
+    a touch too much, whereas under-estimating risks an OOM on load.
+    """
+    import os
+    from codai.models.cache import is_huggingface_model_id
+    if os.path.isfile(path):
+        size_bytes = os.path.getsize(path)
+        multiplier = 1.1 if path.endswith(".gguf") else 1.2
+        return round(size_bytes / 1e9 * multiplier, 2)
+    if is_huggingface_model_id(path):
+        from codai.models.manager import MultiModelManager
+        size_bytes = MultiModelManager._hf_cached_model_size_bytes(path)
+        if size_bytes > 0:
+            return round(size_bytes / 1e9 * 1.2, 2)
+    return None
+
+
+def _backfill_vram_estimate(model_id: str, file_pattern: str = "") -> bool:
+    """After a download, fill in used_vram_gb on any matching config entry that
+    lacks one — the save that created the entry ran before the weights existed on
+    disk, so its estimate came back empty. Returns True if anything changed.
+
+    A CLIP/SigLIP (or any) entry thus always ends up with a reasonable estimate,
+    so the pre-load eviction can size how much VRAM to reclaim instead of freeing
+    only the 1 GB fallback headroom.
+    """
+    if config_manager is None:
+        return False
+    fname = _basename_key(model_id)
+    changed = False
+    for cat in _VALID_MODEL_CATS | {"gguf_models"}:
+        for m in config_manager.models_data.get(cat, []):
+            if not isinstance(m, dict):
+                continue
+            key = _entry_key(m)
+            if not (key == model_id or (fname and _basename_key(key) == fname)):
+                continue
+            if m.get("used_vram_gb") is not None:
+                continue
+            est = _estimate_used_vram_gb(key or model_id)
+            if est is not None:
+                m["used_vram_gb"] = est
+                changed = True
+                print(f"[download] backfilled used_vram_gb={est} GB for '{key or model_id}'")
+    if changed:
+        config_manager.save_models()
+    return changed
+
+
 def _is_model_configured(model_id: str) -> bool:
     """True if model_id is already a configured model (matched by id or basename)."""
     if config_manager is None:
@@ -2158,7 +2220,23 @@ async def api_model_load(request: Request, username: str = Depends(require_admin
             multi_model_manager.active_in_vram = model_key
             multi_model_manager.models_in_vram.add(model_key)
             multi_model_manager.record_vram_delta(model_key, _snap)
-        elif model_type in ("embedding", "spatial", "vision"):
+        elif model_type == "embedding":
+            # Embedding models are sentence-transformers / CLIP encoders, not
+            # diffusers pipelines — use the same loader the /v1/embeddings path
+            # uses so a preload from the interface produces a reusable model.
+            from codai.api.embeddings import _load_embedding_model, _derive_device as _emb_device
+            model_key = f"embedding:{path}"
+            _snap = multi_model_manager.vram_before_load()
+            cfg = (multi_model_manager.config.get(model_key)
+                   or multi_model_manager.config.get(path) or model_cfg or {})
+            emb_obj = await asyncio.to_thread(_load_embedding_model, path, _emb_device(), cfg)
+            if emb_obj is None:
+                raise RuntimeError("Embedding model failed to load")
+            multi_model_manager.add_model(model_key, emb_obj)
+            multi_model_manager.active_in_vram = model_key
+            multi_model_manager.models_in_vram.add(model_key)
+            multi_model_manager.record_vram_delta(model_key, _snap)
+        elif model_type in ("spatial", "vision"):
             from codai.api.images import _load_diffusers_pipeline
             from codai.api.state import get_global_args
             model_key = f"{model_type}:{path}"
@@ -2420,20 +2498,12 @@ async def api_model_configure(request: Request, username: str = Depends(require_
         lst = config_manager.models_data.get(cat, [])
         config_manager.models_data[cat] = [m for m in lst if not _should_remove(m)]
 
-    # Auto-estimate used_vram_gb from file size if not provided
+    # Auto-estimate used_vram_gb from file size if not provided. For a not-yet-
+    # downloaded HF model this returns None (nothing on disk to measure); the
+    # download-completion backfill fills it in once the weights land.
     used_vram_gb = data.get("used_vram_gb")
     if used_vram_gb is None:
-        import os
-        from codai.models.cache import is_huggingface_model_id
-        if os.path.isfile(path):
-            size_bytes = os.path.getsize(path)
-            multiplier = 1.1 if path.endswith(".gguf") else 1.2
-            used_vram_gb = round(size_bytes / 1e9 * multiplier, 2)
-        elif is_huggingface_model_id(path):
-            from codai.models.manager import MultiModelManager
-            size_bytes = MultiModelManager._hf_cached_model_size_bytes(path)
-            if size_bytes > 0:
-                used_vram_gb = round(size_bytes / 1e9 * 1.2, 2)
+        used_vram_gb = _estimate_used_vram_gb(path)
 
     # Build settings entry
     entry: dict = {"path": path, "model_type": model_types[0], "model_types": model_types, "config_id": config_id}

codai/api/embeddings.py

@@ -48,38 +48,208 @@ def _derive_device() -> str:
     return "cuda:0"
 
 
+class _EmbeddingModel:
+    """A loaded embedding model, tagged with its backend.
+
+    Iterates as ``(backend, model)`` so ``backend, model = obj`` unpacking keeps
+    working everywhere, but also exposes ``cleanup()`` — which is what the model
+    manager's eviction path (``_evict_one`` / ``ModelInstancePool.cleanup_all``)
+    calls to move weights off the GPU. Without it a bare tuple falls through
+    those branches and its VRAM is only reclaimed implicitly by gc.
+    """
+
+    __slots__ = ("backend", "model")
+
+    def __init__(self, backend, model):
+        self.backend = backend
+        self.model = model
+
+    def __iter__(self):
+        yield self.backend
+        yield self.model
+
+    def cleanup(self):
+        try:
+            if self.backend == 'sentence_transformers':
+                if hasattr(self.model, 'to'):
+                    self.model.to('cpu')
+            elif self.backend in ('clip', 'transformers'):
+                # model is (processor_or_tokenizer, hf_model, device)
+                hf_model = self.model[1]
+                if hf_model is not None and hasattr(hf_model, 'to'):
+                    hf_model.to('cpu')
+        except Exception:
+            pass
+        self.model = None
+
+
+def _trust_remote_code(model_config: dict = None) -> bool:
+    cfg = model_config or {}
+    raw = cfg.get('_raw_cfg') if isinstance(cfg.get('_raw_cfg'), dict) else {}
+    return bool(cfg.get('trust_remote_code') or raw.get('trust_remote_code'))
+
+
+# Vision+text dual encoders (CLIP/SigLIP family). These expose get_text_features()
+# and get_image_features(), whose projection heads put both modalities in ONE
+# shared space — the whole point of a multimodal embedding model.
+_DUAL_ENCODER_TYPES = {
+    'clip', 'clip_vision_model', 'siglip', 'siglip2', 'chinese_clip',
+    'altclip', 'blip', 'blip-2', 'blip_2', 'x_clip', 'metaclip_2',
+}
+
+
+def _is_dual_encoder(model_name: str, trust: bool) -> bool:
+    """True if the HF config describes a vision+text dual encoder."""
+    try:
+        from transformers import AutoConfig
+        cfg = AutoConfig.from_pretrained(model_name, trust_remote_code=trust)
+    except Exception:
+        return False
+    return (hasattr(cfg, 'vision_config')
+            or str(getattr(cfg, 'model_type', '')) in _DUAL_ENCODER_TYPES)
+
+
+def _has_st_modules(model_name: str) -> bool:
+    """True if the repo ships a sentence-transformers modules.json.
+
+    Matters for dual encoders: when ST has a native recipe (e.g.
+    sentence-transformers/clip-ViT-B-32, jina-clip-v2) it handles both
+    modalities itself. When it doesn't, wrapping a raw CLIPModel in ST gives
+    unprojected text hidden states — a *different* space from the image
+    features — so we must drive the model through transformers instead.
+    """
+    import os
+    try:
+        if os.path.isdir(model_name):
+            return os.path.isfile(os.path.join(model_name, 'modules.json'))
+        from huggingface_hub import file_exists
+        return bool(file_exists(model_name, 'modules.json'))
+    except Exception:
+        return False
+
+
 def _load_embedding_model(model_name: str, device: str, model_config: dict = None):
     from codai.models.hf_loading import build_from_pretrained_kwargs
-    try:
-        from sentence_transformers import SentenceTransformer
-        # sentence-transformers honours quantization via model_kwargs.
-        fp = build_from_pretrained_kwargs(model_config)
-        st_kwargs = {}
-        if 'quantization_config' in fp:
-            st_kwargs['model_kwargs'] = {'quantization_config': fp['quantization_config']}
-        model = SentenceTransformer(model_name, device=device, **st_kwargs)
-        return ('sentence_transformers', model)
-    except ImportError:
-        pass
+    trust = _trust_remote_code(model_config)
+
+    # A dual encoder without an ST recipe must go down the transformers path so
+    # text and images share one space; everything else prefers ST.
+    prefer_clip = _is_dual_encoder(model_name, trust) and not _has_st_modules(model_name)
+
+    if not prefer_clip:
+        try:
+            from sentence_transformers import SentenceTransformer
+            # sentence-transformers honours quantization via model_kwargs.
+            fp = build_from_pretrained_kwargs(model_config)
+            st_kwargs = {}
+            if 'quantization_config' in fp:
+                st_kwargs['model_kwargs'] = {'quantization_config': fp['quantization_config']}
+            if trust:
+                st_kwargs['trust_remote_code'] = True
+            model = SentenceTransformer(model_name, device=device, **st_kwargs)
+            return _EmbeddingModel('sentence_transformers', model)
+        except ImportError:
+            pass
 
     try:
         from transformers import AutoTokenizer, AutoModel
         import torch
         fp = build_from_pretrained_kwargs(model_config)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if trust:
+            fp['trust_remote_code'] = True
         model = AutoModel.from_pretrained(model_name, **fp)
         if 'quantization_config' not in fp and 'device_map' not in fp:
             model = model.to(device)
-        return ('transformers', (tokenizer, model, device))
+        if hasattr(model, 'get_text_features') and hasattr(model, 'get_image_features'):
+            from transformers import AutoProcessor
+            processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=trust)
+            return _EmbeddingModel('clip', (processor, model, device))
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust)
+        return _EmbeddingModel('transformers', (tokenizer, model, device))
     except Exception as e:
         raise RuntimeError(f"Cannot load embedding model '{model_name}': {e}")
 
 
+def _supports_images(model_obj) -> bool:
+    """True if this loaded model can embed images into the same space as text."""
+    backend, model = model_obj
+    if backend == 'clip':
+        return True
+    if backend != 'sentence_transformers':
+        return False
+    try:
+        first = model._first_module()
+    except Exception:
+        return False
+    # ST's CLIPModel module, or a custom (trust_remote_code) module that carries
+    # an image processor — both accept PIL images in encode().
+    return (type(first).__name__.lower().startswith('clip')
+            or hasattr(first, 'processor')
+            or hasattr(first, 'image_processor'))
+
+
+def _decode_image(src: str):
+    """Accept a data URI, http(s) URL, local file path or bare base64 blob."""
+    import io
+    import os
+    import re
+    from PIL import Image
+
+    if not isinstance(src, str) or not src.strip():
+        raise ValueError("empty image reference")
+    s = src.strip()
+
+    if s.startswith(('http://', 'https://')):
+        import requests
+        resp = requests.get(s, timeout=30)
+        resp.raise_for_status()
+        return Image.open(io.BytesIO(resp.content)).convert('RGB')
+    if s.startswith('data:'):
+        s = s.split(',', 1)[1] if ',' in s else ''
+    elif os.path.isfile(s):
+        return Image.open(s).convert('RGB')
+
+    try:
+        raw = base64.b64decode(re.sub(r'\s+', '', s), validate=True)
+    except Exception as e:
+        raise ValueError(
+            f"image is not a URL, data URI, file path or base64 blob: {e}")
+    return Image.open(io.BytesIO(raw)).convert('RGB')
+
+
+def _clip_feats(raw):
+    """Extract the projected shared-space vector from a get_*_features() result.
+
+    transformers <5 returns the projected tensor directly; transformers 5.x
+    returns a BaseModelOutputWithPooling whose `pooler_output` holds the
+    projected (shared-space) embedding.
+    """
+    import torch
+    if isinstance(raw, torch.Tensor):
+        return raw
+    for attr in ('text_embeds', 'image_embeds', 'pooler_output'):
+        val = getattr(raw, attr, None)
+        if val is not None:
+            return val
+    raise RuntimeError("CLIP model returned no usable feature tensor")
+
+
 def _embed_texts(model_obj, texts: List[str], dimensions=None) -> List[List[float]]:
     backend, model = model_obj
     if backend == 'sentence_transformers':
         vecs = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
         results = [v.tolist() for v in vecs]
+    elif backend == 'clip':
+        import torch
+        import torch.nn.functional as F
+        processor, hf_model, device = model
+        inputs = processor(text=texts, padding=True, truncation=True,
+                           return_tensors='pt')
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        with torch.no_grad():
+            feats = _clip_feats(hf_model.get_text_features(**inputs))
+        feats = F.normalize(feats, dim=-1)
+        results = [row.cpu().tolist() for row in feats]
     else:
         import torch
         tokenizer, hf_model, device = model
@@ -101,6 +271,33 @@ def _embed_texts(model_obj, texts: List[str], dimensions=None) -> List[List[floa
     return results
 
 
+def _embed_images(model_obj, images: List[str], dimensions=None) -> List[List[float]]:
+    """Embed images into the same vector space as _embed_texts()."""
+    backend, model = model_obj
+    pil_images = [_decode_image(src) for src in images]
+
+    if backend == 'sentence_transformers':
+        vecs = model.encode(pil_images, convert_to_numpy=True,
+                            normalize_embeddings=True)
+        results = [v.tolist() for v in vecs]
+    elif backend == 'clip':
+        import torch
+        import torch.nn.functional as F
+        processor, hf_model, device = model
+        inputs = processor(images=pil_images, return_tensors='pt')
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        with torch.no_grad():
+            feats = _clip_feats(hf_model.get_image_features(**inputs))
+        feats = F.normalize(feats, dim=-1)
+        results = [row.cpu().tolist() for row in feats]
+    else:
+        raise ValueError("model is text-only")
+
+    if dimensions:
+        results = [v[:dimensions] for v in results]
+    return results
+
+
 @router.post("/v1/embeddings", response_model=EmbeddingsResponse, summary="Create embeddings")
 async def create_embeddings(request: EmbeddingsRequest, http_request: Request = None):
     """
@@ -109,7 +306,8 @@ async def create_embeddings(request: EmbeddingsRequest, http_request: Request =
     # Register a task so embeddings appear in the unified task list, like every
     # other model type. Finished on success or error below.
     from codai.tasks import task_registry
-    _title = request.input if isinstance(request.input, str) else "embeddings"
+    _title = (request.input if isinstance(request.input, str)
+              else ("image embeddings" if request.image is not None else "embeddings"))
     _tid = task_registry.register(
         "embedding", title=str(_title)[:80], model=(request.model or "embedding"))
     task_registry.start(_tid)
@@ -143,20 +341,50 @@ async def _run_embeddings(request: EmbeddingsRequest, http_request: Request = No
     if model_obj is None:
         device = _derive_device()
         from codai.tasks import loading_task
+        # Snapshot VRAM around the load so the model's real footprint is measured
+        # and recorded — this is what lets a later request for another model size
+        # its eviction correctly (and lets this model be evicted to reclaim VRAM).
+        _snap = multi_model_manager.vram_before_load()
         try:
             with loading_task(model_name, model_type="embedding"):
                 model_obj = await asyncio.get_event_loop().run_in_executor(
                     None, _load_embedding_model, model_name, device, _emb_cfg)
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Failed to load embedding model: {e}")
-        multi_model_manager.models[model_key] = model_obj
+        # Register through add_model (pool + models_in_vram bookkeeping) rather than
+        # a bare dict assignment, so eviction/unload treat it like every other model.
+        multi_model_manager.add_model(model_key, model_obj)
         multi_model_manager.current_model_key = model_key
+        multi_model_manager.record_vram_delta(model_key, _snap)
+
+    texts: List[str] = []
+    if request.input is not None:
+        texts = [request.input] if isinstance(request.input, str) else list(request.input)
+    images: List[str] = []
+    if request.image is not None:
+        images = [request.image] if isinstance(request.image, str) else list(request.image)
+    if not texts and not images:
+        raise HTTPException(
+            status_code=400, detail="Provide 'input' (text) and/or 'image'.")
 
-    texts = [request.input] if isinstance(request.input, str) else request.input
+    if images and not _supports_images(model_obj):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Model '{model_name}' is text-only; image embedding needs a "
+                   "multimodal model (CLIP/SigLIP family, e.g. "
+                   "sentence-transformers/clip-ViT-B-32 or jinaai/jina-clip-v2).")
 
+    # Text vectors first, then image vectors — indices follow that order.
+    vectors: List[List[float]] = []
     try:
-        vectors = await asyncio.get_event_loop().run_in_executor(
-            None, _embed_texts, model_obj, texts, request.dimensions)
+        if texts:
+            vectors += await asyncio.get_event_loop().run_in_executor(
+                None, _embed_texts, model_obj, texts, request.dimensions)
+        if images:
+            vectors += await asyncio.get_event_loop().run_in_executor(
+                None, _embed_images, model_obj, images, request.dimensions)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Embedding failed: {e}")
 

codai/pydantic/embedrequest.py

@@ -22,8 +22,8 @@ from pydantic import BaseModel, ConfigDict, Field
 
 class EmbeddingsRequest(BaseModel):
     model: str = Field(..., description="Embedding model id to use.")
-    input: Union[str, List[str]] = Field(..., description="Text or list of texts to embed.")
-    image: Optional[Union[str, List[str]]] = Field(None, description="Base64/URL image(s) for multimodal embedding models.")
+    input: Optional[Union[str, List[str]]] = Field(None, description="Text or list of texts to embed. Optional only when 'image' is given.")
+    image: Optional[Union[str, List[str]]] = Field(None, description="Image(s) for multimodal embedding models: http(s) URL, data URI, local path or bare base64. Vectors are returned after the text ones.")
     encoding_format: Optional[str] = Field("float", description="Return embeddings as 'float' arrays or 'base64'.")
     dimensions: Optional[int] = Field(None, description="Truncate embeddings to N dimensions (if the model supports it).")
     quantization: Optional[str] = Field(None, description="Optional TurboQuant vector quantization: 'turbo' (8-bit), 'turbo8', 'turbo6', 'turbo4' or 'turbo2'. With encoding_format='float' the (lossy) reconstructed vectors are returned; with 'base64' the compact packed bytes are returned plus a 'quantization' metadata block describing how to decode them.")