Spaces:
Running
Running
File size: 3,119 Bytes
3865888 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """Concrete local loaders registered into ModelManager (V6).
Each loader reads from models/<task>/ (offline). Missing files raise
FileNotFoundError -> manager surfaces UNAVAILABLE, never fake handles.
"""
import os
from src.models.offline import apply_offline_env
apply_offline_env()
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
def _p(*parts):
return os.path.join(PROJECT_ROOT, "models", *[p.lower() for p in parts])
def load_text():
from llama_cpp import Llama
d = _p("TEXT_MODEL")
gguf = next((os.path.join(d, f) for f in sorted(os.listdir(d))
if f.endswith(".gguf")), None)
if gguf is None:
raise FileNotFoundError("no GGUF in models/text_model")
return Llama(model_path=gguf, n_ctx=2048, n_threads=6, verbose=False)
def load_embedding():
from sentence_transformers import SentenceTransformer
d = _p("EMBEDDING_MODEL")
# local directory first (offline-capable); hub ID only as online fallback
local_files = os.listdir(d) if os.path.isdir(d) else []
if any(f.endswith(".safetensors") or f.endswith(".bin") for f in local_files):
try:
return SentenceTransformer(d)
except Exception:
pass
return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2",
cache_folder=d)
def load_stt():
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
d = _p("STT_MODEL")
proc = WhisperProcessor.from_pretrained(d, local_files_only=True)
model = WhisperForConditionalGeneration.from_pretrained(
d, local_files_only=True, dtype=torch.float32).eval()
return {"processor": proc, "model": model}
def load_tts():
from kokoro import KModel, KPipeline
d = _p("TTS_MODEL")
ckpt = os.path.join(d, "kokoro-v1_0.pth")
if not os.path.exists(ckpt):
raise FileNotFoundError("no kokoro checkpoint in models/tts_model")
model = KModel(repo_id=None, config=os.path.join(d, "config.json"), model=ckpt)
return KPipeline(lang_code="a", model=model)
def load_image():
from src.models.image_adapter import LocalImageModel
m = LocalImageModel()
m.load()
return m
def load_vision():
import torch
from transformers import AutoProcessor, SmolVLMForConditionalGeneration
d = _p("VISION_MODEL")
proc = AutoProcessor.from_pretrained(d, local_files_only=True,
trust_remote_code=True)
model = SmolVLMForConditionalGeneration.from_pretrained(
d, local_files_only=True, trust_remote_code=True,
dtype=torch.float32).eval()
return {"processor": proc, "model": model}
LOADERS = {"TEXT_MODEL": load_text, "EMBEDDING_MODEL": load_embedding,
"STT_MODEL": load_stt, "TTS_MODEL": load_tts,
"IMAGE_MODEL": load_image, "VISION_MODEL": load_vision}
def default_manager(**kw):
from src.models.manager import ModelManager
mm = ModelManager(**kw)
for task, fn in LOADERS.items():
mm.register_loader(task, fn)
return mm
|