Spaces:
Sleeping
Sleeping
File size: 4,989 Bytes
871f869 b911f86 e7fd66f b911f86 727cb75 b911f86 727cb75 b911f86 727cb75 b911f86 a822dca b911f86 871f869 7a28b9f 727cb75 7a28b9f 727cb75 a822dca b911f86 727cb75 dd678c1 871f869 dd678c1 e7fd66f | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | from gradio_space.spaces_runtime import gpu_task
from inference.config import get_app_config, get_model_config
from inference.factory import get_backend, reset_backend
from inference.response_clean import strip_reasoning_output
_app_config = get_app_config()
_runtime_model_key: str | None = None
_current_model_key: str | None = None
_load_state: dict[str, bool] = {}
_load_errors: dict[str, str] = {}
def get_active_model_key() -> str:
return _runtime_model_key or _app_config.active_model
def set_runtime_model_key(key: str) -> str:
"""Pin the active preset for all tabs until process restart."""
global _runtime_model_key, _current_model_key
model = get_model_config(key)
previous = get_active_model_key()
if key != previous:
reset_backend()
_current_model_key = None
_load_state.pop(previous, None)
_load_errors.pop(previous, None)
_runtime_model_key = key
return model.label
def ensure_model_loaded(model_key: str) -> str | None:
global _current_model_key
if model_key != _current_model_key:
reset_backend()
_current_model_key = model_key
if _load_state.get(model_key):
return None
if model_key in _load_errors:
return _load_errors[model_key]
try:
get_backend(model_key).load()
_load_state[model_key] = True
return None
except Exception as exc: # noqa: BLE001 — surface model load failures in the UI
message = f"Failed to load model: {exc}"
_load_errors[model_key] = message
return message
def runtime_device_hint(model_key: str) -> str:
model = get_model_config(model_key)
if model.backend == "transformers":
try:
import torch
if torch.cuda.is_available():
return f"GPU ({torch.cuda.get_device_name(0)})"
except ImportError:
pass
return "CPU"
if model.n_gpu_layers > 0:
return f"llama.cpp GPU offload ({model.n_gpu_layers} layers)"
return "CPU"
def warmup(model_key: str | None = None) -> str:
key = model_key or get_active_model_key()
model = get_model_config(key)
if _load_state.get(key):
backend = get_backend(key)
device = (
backend.device_label
if hasattr(backend, "device_label")
else runtime_device_hint(key)
)
return f"Model ready: {model.label} on {device}"
if key in _load_errors:
return _load_errors[key]
device_hint = runtime_device_hint(key)
return (
f"Preset `{key}` selected ({model.backend}, {device_hint}). "
"Loading weights…"
)
@gpu_task(duration=120)
def reload_model(model_key: str) -> str:
"""Clear cached backend and reload weights for settings panel."""
global _current_model_key
key = model_key or get_active_model_key()
set_runtime_model_key(key)
reset_backend()
_current_model_key = None
_load_state.pop(key, None)
_load_errors.pop(key, None)
error = ensure_model_loaded(key)
if error:
return error
return warmup(key)
def select_and_reload_model(model_key: str) -> str:
"""Switch runtime preset and load weights (Settings dropdown)."""
return reload_model(model_key)
def preload_active_model() -> str:
"""Load the active preset at startup so the first request is fast."""
key = get_active_model_key()
print(f"[startup] Loading model preset `{key}`…", flush=True)
error = ensure_model_loaded(key)
if error:
print(f"[startup] {error}", flush=True)
return error
status = warmup(key)
print(f"[startup] {status}", flush=True)
return status
def model_status(model_key: str) -> str:
model = get_model_config(model_key)
notes = ""
if model.backend == "llama_cpp" and model.multimodal:
notes = (
"\n- Note: text-only on llama.cpp; use transformers preset for image/video input."
)
return (
f"**{model.label}**\n\n"
f"- Backend: `{model.backend}`\n"
f"- {warmup(model_key)}{notes}"
)
def _history_to_messages(history: list) -> list[dict[str, str]]:
messages: list[dict[str, str]] = []
for item in history:
if isinstance(item, dict):
messages.append({"role": item["role"], "content": item["content"]})
else:
user_msg, assistant_msg = item
messages.append({"role": "user", "content": user_msg})
if assistant_msg:
messages.append({"role": "assistant", "content": assistant_msg})
return messages
@gpu_task(duration=60)
def chat(message: str, history: list, model_key: str) -> str:
load_error = ensure_model_loaded(model_key)
if load_error:
return load_error
messages = _history_to_messages(history)
messages.append({"role": "user", "content": message})
reply = get_backend(model_key).chat(messages)
return strip_reasoning_output(reply)
|