141 lines
5.3 KiB
Python
141 lines
5.3 KiB
Python
"""Shared bge-small loader.
|
|
|
|
`bucket_search.py` and `modifier_search.py` originally each loaded
|
|
their own copy of `BAAI/bge-small-en-v1.5`. With `tag_search.py` added
|
|
that's three copies of the same ~30MB weights and three cold-start
|
|
loads on boot. This module owns one shared instance behind a lock.
|
|
|
|
Public API:
|
|
get() -> (model, tokenizer, device)
|
|
get_load_error() -> str | None # None means "not yet attempted or
|
|
# currently in transient retry state"
|
|
|
|
Behavior matches the pattern the call-sites already implemented locally:
|
|
- ModuleNotFoundError caches a hard failure (won't retry without pip install)
|
|
- ImportError / OSError leave load_error unset so the next call retries
|
|
(handles boot-time circular-import races between torch and joblib/loky)
|
|
- Any other exception caches as a hard failure
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger("promptchain.embed")
|
|
|
|
MODEL_ID = "BAAI/bge-small-en-v1.5"
|
|
EMBED_DIM = 384
|
|
|
|
_lock = threading.RLock()
|
|
_state: dict[str, Any] = {
|
|
"ready": False,
|
|
"model": None,
|
|
"tokenizer": None,
|
|
"device": None,
|
|
"load_error": None,
|
|
}
|
|
|
|
|
|
def get_load_error() -> str | None:
|
|
return _state["load_error"]
|
|
|
|
|
|
def get() -> tuple[Any, Any, str] | None:
|
|
"""Load (or return cached) shared encoder. Returns None on failure;
|
|
caller is expected to degrade gracefully (search returns empty list
|
|
rather than raising)."""
|
|
with _lock:
|
|
if _state["ready"]:
|
|
return _state["model"], _state["tokenizer"], _state["device"]
|
|
if _state["load_error"]:
|
|
return None
|
|
try:
|
|
import torch # noqa: F401 (only checking availability here)
|
|
from transformers import AutoTokenizer, AutoModel
|
|
except ModuleNotFoundError as e:
|
|
_state["load_error"] = f"missing dependency: {e}"
|
|
logger.warning(
|
|
"embed_model disabled: %s. Install with: pip install transformers",
|
|
e,
|
|
)
|
|
return None
|
|
except ImportError as e:
|
|
# Partial-init / circular import — torch and its joblib/loky
|
|
# transitive deps race against each other on cold ComfyUI
|
|
# boot. Don't cache; the next request after the import cycle
|
|
# resolves will succeed.
|
|
logger.warning(
|
|
"embed_model: deferred import (%s) — will retry on next request",
|
|
e,
|
|
)
|
|
return None
|
|
except ValueError as e:
|
|
# joblib double-registration: when a prior import partially
|
|
# initialized joblib (e.g. boot-time bucket_search.warmup
|
|
# failed mid-import), the next transformers import re-registers
|
|
# zlib and joblib raises ValueError. Treat as transient — the
|
|
# process state is already broken, but the search caller will
|
|
# degrade to alias-only matching and the request continues.
|
|
logger.warning(
|
|
"embed_model: deferred import (joblib state issue: %s) — "
|
|
"search will degrade to alias-only this request",
|
|
e,
|
|
)
|
|
return None
|
|
try:
|
|
# bge-small is ~30MB on GPU — negligible against image-gen VRAM
|
|
# budgets, and a 5090 rebuilds the 14k-row tag index in ~30s
|
|
# vs 5-10 min on CPU. Fall back to CPU when CUDA isn't there.
|
|
import torch
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModel.from_pretrained(MODEL_ID).to(device).eval()
|
|
_state["tokenizer"] = tokenizer
|
|
_state["model"] = model
|
|
_state["device"] = device
|
|
_state["ready"] = True
|
|
logger.info("embed_model: loaded %s on %s", MODEL_ID, device)
|
|
return model, tokenizer, device
|
|
except OSError as e:
|
|
# HF download / disk error — transient (network/lock). Retry.
|
|
logger.warning(
|
|
"embed_model: deferred fetch (%s) — will retry on next request",
|
|
e,
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
_state["load_error"] = str(e)
|
|
logger.exception("embed_model: load failed")
|
|
return None
|
|
|
|
|
|
def embed(texts: list[str], batch_size: int = 32):
|
|
"""L2-normalized CLS-pooled embeddings for `texts`. Returns torch
|
|
tensor [N, 384]. Returns None if the model isn't loaded.
|
|
|
|
Note: bge-small uses CLS pooling (per the model card). E5 family
|
|
requires mean pooling — if we ever switch models, callers may need
|
|
their own embed function. Today all three consumers (bucket, modifier,
|
|
tag) use bge-small, so this central embed is fine."""
|
|
loaded = get()
|
|
if loaded is None:
|
|
return None
|
|
model, tokenizer, device = loaded
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
out_chunks = []
|
|
for i in range(0, len(texts), batch_size):
|
|
chunk = texts[i:i + batch_size]
|
|
enc = tokenizer(
|
|
chunk, padding=True, truncation=True, max_length=512,
|
|
return_tensors="pt",
|
|
).to(device)
|
|
with torch.no_grad():
|
|
out = model(**enc)
|
|
cls = out.last_hidden_state[:, 0]
|
|
out_chunks.append(F.normalize(cls, p=2, dim=1))
|
|
return torch.cat(out_chunks, dim=0)
|