From a40ebb1aabc7911cb5556a88351f72cffccaea53 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Tue, 24 Feb 2026 17:40:05 +0400 Subject: [PATCH] Add GGUF model inference via llama-server backend --- .gitignore | 3 + setup.sh | 62 ++- studio/backend/core/inference/__init__.py | 2 + studio/backend/core/inference/llama_cpp.py | 414 ++++++++++++++++++ studio/backend/models/inference.py | 2 + studio/backend/models/models.py | 1 + studio/backend/routes/inference.py | 268 ++++++++++-- studio/backend/routes/models.py | 18 + studio/backend/utils/models/model_config.py | 50 ++- .../chat/hooks/use-chat-model-runtime.ts | 6 +- .../frontend/src/features/chat/types/api.ts | 3 + .../src/features/chat/types/runtime.ts | 1 + .../frontend/src/hooks/use-hf-model-search.ts | 1 - 13 files changed, 791 insertions(+), 40 deletions(-) create mode 100644 studio/backend/core/inference/llama_cpp.py diff --git a/.gitignore b/.gitignore index 3901c1c699..e24c38c2b0 100755 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,9 @@ unsloth_training_checkpoints/ *.gguf *.safetensors +# Built binaries (llama-server etc.) +bin/ + # IDE / Editors .vscode/ .idea/ diff --git a/setup.sh b/setup.sh index 1f55e141fd..be17370df4 100755 --- a/setup.sh +++ b/setup.sh @@ -206,7 +206,67 @@ else fi fi -# ── 8. Add shell alias (skip in Colab) ── +# ── 8. Build llama-server for GGUF inference ── +# Builds in an isolated temp directory to avoid conflicts with unsloth-zoo's +# own llama.cpp management (used for GGUF export). Only the llama-server +# binary is extracted to $REPO/bin/. +LLAMA_SERVER_BIN="$SCRIPT_DIR/bin/llama-server" +if [ -f "$LLAMA_SERVER_BIN" ]; then + echo "" + echo "✅ llama-server already exists at $LLAMA_SERVER_BIN" +else + # Check prerequisites + if ! command -v cmake &>/dev/null; then + echo "" + echo "⚠️ cmake not found — skipping llama-server build (GGUF inference won't be available)" + echo " Install cmake and re-run setup.sh to enable GGUF inference." + elif ! command -v git &>/dev/null; then + echo "" + echo "⚠️ git not found — skipping llama-server build (GGUF inference won't be available)" + else + echo "" + echo "Building llama-server for GGUF inference..." + LLAMA_BUILD_TMP=$(mktemp -d) + + BUILD_OK=true + run_quiet "clone llama.cpp" git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$LLAMA_BUILD_TMP/llama.cpp" || BUILD_OK=false + + if [ "$BUILD_OK" = true ]; then + CMAKE_ARGS="" + if command -v nvcc &>/dev/null; then + echo " Building with CUDA support..." + CMAKE_ARGS="-DGGML_CUDA=ON" + else + echo " Building CPU-only (no CUDA detected)..." + fi + + NCPU=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4) + + run_quiet "cmake llama.cpp" cmake -S "$LLAMA_BUILD_TMP/llama.cpp" -B "$LLAMA_BUILD_TMP/llama.cpp/build" $CMAKE_ARGS || BUILD_OK=false + fi + + if [ "$BUILD_OK" = true ]; then + run_quiet "build llama-server" cmake --build "$LLAMA_BUILD_TMP/llama.cpp/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false + fi + + if [ "$BUILD_OK" = true ]; then + mkdir -p "$SCRIPT_DIR/bin" + if [ -f "$LLAMA_BUILD_TMP/llama.cpp/build/bin/llama-server" ]; then + cp "$LLAMA_BUILD_TMP/llama.cpp/build/bin/llama-server" "$LLAMA_SERVER_BIN" + echo "✅ llama-server built and installed to $LLAMA_SERVER_BIN" + else + echo "⚠️ llama-server binary not found after build — GGUF inference won't be available" + fi + else + echo "⚠️ llama-server build failed — GGUF inference won't be available, but everything else works" + fi + + # Clean up temp build directory + rm -rf "$LLAMA_BUILD_TMP" + fi +fi + +# ── 9. Add shell alias (skip in Colab) ── # Note: venv activation does NOT persist across terminal sessions. # This alias hardcodes the venv python path so users don't need to activate. if [ "$IS_COLAB" = false ]; then diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 494229a087..ff8b75d36a 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -2,8 +2,10 @@ Inference submodule - Inference backend for model loading and generation """ from .inference import InferenceBackend, get_inference_backend +from .llama_cpp import LlamaCppBackend __all__ = [ 'InferenceBackend', 'get_inference_backend', + 'LlamaCppBackend', ] diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py new file mode 100644 index 0000000000..dae8d20489 --- /dev/null +++ b/studio/backend/core/inference/llama_cpp.py @@ -0,0 +1,414 @@ +""" +llama-server inference backend for GGUF models. + +Manages a llama-server subprocess and proxies chat completions +through its /v1/completions endpoint. +""" +import atexit +import json +import logging +import shutil +import signal +import socket +import subprocess +import threading +import time +from pathlib import Path +from typing import Generator, Optional + +import httpx + +logger = logging.getLogger(__name__) + + +class LlamaCppBackend: + """ + Manages a llama-server subprocess for GGUF model inference. + + Lifecycle: + 1. load_model() — starts llama-server with the GGUF file + 2. generate_chat_completion() — formats prompt, proxies to /v1/completions, streams back + 3. unload_model() — terminates llama-server subprocess + """ + + def __init__(self): + self._process: Optional[subprocess.Popen] = None + self._port: Optional[int] = None + self._model_identifier: Optional[str] = None + self._gguf_path: Optional[str] = None + self._healthy = False + self._lock = threading.Lock() + self._chat_template: Optional[str] = None + + atexit.register(self._cleanup) + + # ── Properties ──────────────────────────────────────────────── + + @property + def is_loaded(self) -> bool: + return self._process is not None and self._healthy + + @property + def base_url(self) -> str: + return f"http://127.0.0.1:{self._port}" + + @property + def model_identifier(self) -> Optional[str]: + return self._model_identifier + + # ── Binary discovery ────────────────────────────────────────── + + @staticmethod + def _find_llama_server_binary() -> Optional[str]: + """ + Locate the llama-server binary. + + Search order: + 1. LLAMA_SERVER_PATH environment variable + 2. ./bin/llama-server (built by setup.sh) + 3. llama-server on PATH (system install) + 4. ./llama.cpp/llama-server (unsloth-zoo build output) + """ + import os + + # 1. Env var + env_path = os.environ.get("LLAMA_SERVER_PATH") + if env_path and Path(env_path).is_file(): + return env_path + + # 2. Project bin/ directory (setup.sh output) + project_root = Path(__file__).resolve().parents[3] # core/inference/ → backend/ → studio/ → root + bin_path = project_root / "bin" / "llama-server" + if bin_path.is_file(): + return str(bin_path) + + # 3. System PATH + system_path = shutil.which("llama-server") + if system_path: + return system_path + + # 4. unsloth-zoo build output (from GGUF export) + llama_cpp_path = project_root / "llama.cpp" / "llama-server" + if llama_cpp_path.is_file(): + return str(llama_cpp_path) + + return None + + # ── Port allocation ─────────────────────────────────────────── + + @staticmethod + def _find_free_port() -> int: + """Find an available TCP port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + # ── Lifecycle ───────────────────────────────────────────────── + + def load_model( + self, + gguf_path: str, + model_identifier: str, + n_ctx: int = 4096, + n_gpu_layers: int = -1, + n_threads: Optional[int] = None, + ) -> bool: + """ + Start llama-server with the given GGUF file. + + Args: + gguf_path: Path to the .gguf file + model_identifier: Display identifier for the model + n_ctx: Context window size + n_gpu_layers: Number of layers to offload to GPU (-1 = all) + n_threads: Number of CPU threads (None = auto) + + Returns: + True if server started and health check passed. + """ + with self._lock: + # Kill existing process if any + self._kill_process() + + binary = self._find_llama_server_binary() + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + + self._port = self._find_free_port() + cmd = [ + binary, + "-m", gguf_path, + "--port", str(self._port), + "-c", str(n_ctx), + "-ngl", str(n_gpu_layers), + ] + if n_threads is not None: + cmd.extend(["--threads", str(n_threads)]) + + logger.info(f"Starting llama-server: {' '.join(cmd)}") + + self._process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + self._gguf_path = gguf_path + self._model_identifier = model_identifier + + # Wait for health + if not self._wait_for_health(timeout=120.0): + self._kill_process() + raise RuntimeError( + "llama-server failed to start. " + "Check that the GGUF file is valid and you have enough memory." + ) + + self._healthy = True + + # Try to read chat template from GGUF metadata + self._chat_template = self._read_gguf_chat_template(gguf_path) + + logger.info( + f"llama-server ready on port {self._port} " + f"for model '{model_identifier}'" + ) + return True + + def unload_model(self) -> bool: + """Terminate the llama-server subprocess and clean up state.""" + with self._lock: + self._kill_process() + logger.info(f"Unloaded GGUF model: {self._model_identifier}") + self._model_identifier = None + self._gguf_path = None + self._port = None + self._healthy = False + self._chat_template = None + return True + + def _kill_process(self): + """Terminate the subprocess if running.""" + if self._process is None: + return + try: + self._process.terminate() + self._process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL") + self._process.kill() + self._process.wait(timeout=5) + except Exception as e: + logger.warning(f"Error killing llama-server process: {e}") + finally: + self._process = None + + def _cleanup(self): + """atexit handler to ensure llama-server is terminated.""" + self._kill_process() + + def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool: + """ + Poll llama-server's /health endpoint until it responds 200. + + Also monitors subprocess for early exit/crash. + """ + deadline = time.monotonic() + timeout + url = f"http://127.0.0.1:{self._port}/health" + + while time.monotonic() < deadline: + # Check if process crashed + if self._process.poll() is not None: + # Read remaining output for error info + output = self._process.stdout.read() if self._process.stdout else "" + logger.error( + f"llama-server exited with code {self._process.returncode}. " + f"Output: {output[:2000]}" + ) + return False + + try: + resp = httpx.get(url, timeout=2.0) + if resp.status_code == 200: + return True + except (httpx.ConnectError, httpx.TimeoutException): + pass + + time.sleep(interval) + + logger.error(f"llama-server health check timed out after {timeout}s") + return False + + # ── Chat template ───────────────────────────────────────────── + + @staticmethod + def _read_gguf_chat_template(gguf_path: str) -> Optional[str]: + """ + Try to read the chat_template from GGUF file metadata. + + Uses the gguf Python library if available. + Returns the Jinja2 template string, or None. + """ + try: + from gguf import GGUFReader + + reader = GGUFReader(gguf_path) + for field_name in reader.fields: + if field_name == "tokenizer.chat_template": + field = reader.fields[field_name] + # Field data is an array of bytes + template_bytes = bytes(field.parts[field.data[0]]) + template = template_bytes.decode("utf-8") + logger.info(f"Read chat template from GGUF metadata ({len(template)} chars)") + return template + except ImportError: + logger.debug("gguf library not available, cannot read chat template from GGUF metadata") + except Exception as e: + logger.warning(f"Could not read chat template from GGUF: {e}") + + return None + + def format_prompt(self, messages: list[dict], system_prompt: str = "") -> str: + """ + Format chat messages into a raw prompt string for /v1/completions. + + Attempts to: + 1. Render the GGUF's embedded chat_template with Jinja2 + 2. Fallback to ChatML format + """ + # Build full message list with system prompt + full_messages = [] + if system_prompt: + full_messages.append({"role": "system", "content": system_prompt}) + full_messages.extend(messages) + + # Try Jinja2 rendering if we have a template + if self._chat_template: + try: + return self._render_jinja_template(full_messages) + except Exception as e: + logger.warning(f"Jinja2 template rendering failed, falling back to ChatML: {e}") + + # Fallback: ChatML format + return self._format_chatml(full_messages) + + def _render_jinja_template(self, messages: list[dict]) -> str: + """Render messages using the GGUF's Jinja2 chat template.""" + from jinja2 import BaseLoader, Environment + + env = Environment(loader=BaseLoader(), keep_trailing_newline=True) + # Add common template globals + env.globals["raise_exception"] = lambda msg: (_ for _ in ()).throw(ValueError(msg)) + + template = env.from_string(self._chat_template) + rendered = template.render( + messages=messages, + add_generation_prompt=True, + bos_token="", + eos_token="", + ) + return rendered + + @staticmethod + def _format_chatml(messages: list[dict]) -> str: + """Format messages using ChatML template (universal fallback).""" + parts = [] + for msg in messages: + role = msg.get("role", "user") + content = msg.get("content", "") + parts.append(f"<|im_start|>{role}\n{content}<|im_end|>") + parts.append("<|im_start|>assistant") + return "\n".join(parts) + "\n" + + # ── Generation (proxy to llama-server) ──────────────────────── + + def generate_chat_completion( + self, + prompt: str, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_tokens: int = 512, + repetition_penalty: float = 1.1, + stop: Optional[list[str]] = None, + cancel_event: Optional[threading.Event] = None, + ) -> Generator[str, None, None]: + """ + Send a completion request to llama-server and stream tokens back. + + Uses /v1/completions (NOT /v1/chat/completions) so we control + the prompt format entirely. + + Yields cumulative text (matching InferenceBackend's convention). + """ + if not self.is_loaded: + raise RuntimeError("llama-server is not loaded") + + payload = { + "prompt": prompt, + "stream": True, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k if top_k >= 0 else 0, + "min_p": min_p, + "n_predict": max_tokens, + "repeat_penalty": repetition_penalty, + } + if stop: + payload["stop"] = stop + + url = f"{self.base_url}/v1/completions" + cumulative = "" + + try: + with httpx.Client(timeout=None) as client: + with client.stream("POST", url, json=payload) as response: + if response.status_code != 200: + error_body = response.read().decode() + raise RuntimeError( + f"llama-server returned {response.status_code}: {error_body}" + ) + + buffer = "" + for raw_chunk in response.iter_text(): + if cancel_event is not None and cancel_event.is_set(): + break + + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() + + if not line: + continue + if line == "data: [DONE]": + return + if not line.startswith("data: "): + continue + + try: + data = json.loads(line[6:]) + choices = data.get("choices", []) + if choices: + token = choices[0].get("text", "") + if token: + cumulative += token + yield cumulative + except json.JSONDecodeError: + logger.debug(f"Skipping malformed SSE line: {line[:100]}") + + except httpx.ConnectError: + raise RuntimeError("Lost connection to llama-server") + except Exception as e: + if cancel_event is not None and cancel_event.is_set(): + return + raise diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d2d98d7944..c4a062ae3a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -43,6 +43,7 @@ class LoadResponse(BaseModel): display_name: str = Field(..., description="Display name of the model") is_vision: bool = Field(False, description="Whether model is a vision model") is_lora: bool = Field(False, description="Whether model is a LoRA adapter") + is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp)") inference: dict = Field(..., description="Inference parameters (temperature, top_p, top_k, min_p)") @@ -56,6 +57,7 @@ class InferenceStatusResponse(BaseModel): """Current inference backend status""" active_model: Optional[str] = Field(None, description="Currently active model identifier") is_vision: bool = Field(False, description="Whether the active model is a vision model") + is_gguf: bool = Field(False, description="Whether the active model is a GGUF model (llama.cpp)") loading: List[str] = Field(default_factory=list, description="Models currently being loaded") loaded: List[str] = Field(default_factory=list, description="Models currently loaded") diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 5542db76d5..b3f9b50ed2 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -53,6 +53,7 @@ class ModelDetails(BaseModel): config: Optional[Dict[str, Any]] = Field(None, description="Model configuration dictionary") is_vision: bool = Field(False, description="Whether model is a vision model") is_lora: bool = Field(False, description="Whether model is a LoRA adapter") + is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp format)") base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c30a1638d4..156578b960 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -23,6 +23,7 @@ if str(backend_path) not in sys.path: # Import backend functions try: from core.inference import get_inference_backend + from core.inference.llama_cpp import LlamaCppBackend from utils.models import ModelConfig from utils.inference import load_inference_config except ImportError: @@ -30,6 +31,7 @@ except ImportError: if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from core.inference import get_inference_backend + from core.inference.llama_cpp import LlamaCppBackend from utils.models import ModelConfig from utils.inference import load_inference_config @@ -61,60 +63,111 @@ if not logger.handlers: logger.addHandler(handler) logger.setLevel(logging.INFO) +# GGUF inference backend (llama-server) +_llama_cpp_backend = LlamaCppBackend() + +def get_llama_cpp_backend() -> LlamaCppBackend: + return _llama_cpp_backend + @router.post("/load", response_model=LoadResponse) async def load_model(request: LoadRequest): """ Load a model for inference. - + The model_path should be a clean identifier from GET /models/list. Returns inference configuration parameters (temperature, top_p, top_k, min_p) from the model's YAML config, falling back to default.yaml for missing values. + + GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth. """ try: - backend = get_inference_backend() - # Create config using clean factory method # is_lora is auto-detected from adapter_config.json on disk/HF config = ModelConfig.from_identifier( model_id=request.model_path, hf_token=request.hf_token, ) - + if not config: raise HTTPException( status_code=400, detail=f"Invalid model identifier: {request.model_path}" ) - - # Load the model + + # ── GGUF path: load via llama-server ────────────────────── + if config.is_gguf: + llama_backend = get_llama_cpp_backend() + unsloth_backend = get_inference_backend() + + # Unload any active Unsloth model first to free VRAM + if unsloth_backend.active_model_name: + logger.info(f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF") + unsloth_backend.unload_model(unsloth_backend.active_model_name) + + success = llama_backend.load_model( + gguf_path=config.gguf_file, + model_identifier=config.identifier, + n_ctx=request.max_seq_length, + ) + + if not success: + raise HTTPException( + status_code=500, + detail=f"Failed to load GGUF model: {config.display_name}" + ) + + logger.info(f"Loaded GGUF model via llama-server: {config.identifier}") + + inference_config = load_inference_config(config.identifier) + + return LoadResponse( + status="loaded", + model=config.identifier, + display_name=config.display_name, + is_vision=False, + is_lora=False, + is_gguf=True, + inference=inference_config, + ) + + # ── Standard path: load via Unsloth/transformers ────────── + backend = get_inference_backend() + + # Unload any active GGUF model first + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded: + logger.info("Unloading GGUF model before loading Unsloth model") + llama_backend.unload_model() + success = backend.load_model( config=config, max_seq_length=request.max_seq_length, load_in_4bit=request.load_in_4bit, hf_token=request.hf_token, ) - + if not success: raise HTTPException( status_code=500, detail=f"Failed to load model: {config.display_name}" ) - + logger.info(f"Loaded model: {config.identifier}") - + # Load inference configuration parameters inference_config = load_inference_config(config.identifier) - + return LoadResponse( status="loaded", model=config.identifier, display_name=config.display_name, is_vision=config.is_vision, is_lora=config.is_lora, + is_gguf=False, inference=inference_config, ) - + except HTTPException: raise except Exception as e: @@ -129,13 +182,22 @@ async def load_model(request: LoadRequest): async def unload_model(request: UnloadRequest): """ Unload a model from memory. + Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ try: + # Check if the GGUF backend has this model loaded + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded and llama_backend.model_identifier == request.model_path: + llama_backend.unload_model() + logger.info(f"Unloaded GGUF model: {request.model_path}") + return UnloadResponse(status="unloaded", model=request.model_path) + + # Otherwise, unload from Unsloth backend backend = get_inference_backend() backend.unload_model(request.model_path) logger.info(f"Unloaded model: {request.model_path}") return UnloadResponse(status="unloaded", model=request.model_path) - + except Exception as e: logger.error(f"Error unloading model: {e}", exc_info=True) raise HTTPException( @@ -221,22 +283,37 @@ async def generate_stream(request: GenerateRequest): async def get_status(): """ Get current inference backend status. + Reports whichever backend (Unsloth or llama-server) is currently active. """ try: + llama_backend = get_llama_cpp_backend() + + # If a GGUF model is loaded via llama-server, report that + if llama_backend.is_loaded: + return InferenceStatusResponse( + active_model=llama_backend.model_identifier, + is_vision=False, + is_gguf=True, + loading=[], + loaded=[llama_backend.model_identifier], + ) + + # Otherwise, report Unsloth backend status backend = get_inference_backend() - + is_vision = False if backend.active_model_name: model_info = backend.models.get(backend.active_model_name, {}) is_vision = model_info.get("is_vision", False) - + return InferenceStatusResponse( active_model=backend.active_model_name, is_vision=is_vision, + is_gguf=False, loading=list(getattr(backend, 'loading_models', set())), loaded=list(backend.models.keys()), ) - + except Exception as e: logger.error(f"Error getting status: {e}", exc_info=True) raise HTTPException( @@ -315,29 +392,157 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque Streaming (default): returns SSE chunks matching OpenAI's format. Non-streaming: returns a single ChatCompletion JSON object. - """ - backend = get_inference_backend() - if not backend.active_model_name: - raise HTTPException( - status_code=400, - detail="No model loaded. Call POST /inference/load first.", - ) + Automatically routes to the correct backend: + - GGUF models → llama-server via LlamaCppBackend + - Other models → Unsloth/transformers via InferenceBackend + """ + llama_backend = get_llama_cpp_backend() + using_gguf = llama_backend.is_loaded + + # ── Determine which backend is active ───────────────────── + if using_gguf: + model_name = llama_backend.model_identifier or payload.model + else: + backend = get_inference_backend() + if not backend.active_model_name: + raise HTTPException( + status_code=400, + detail="No model loaded. Call POST /inference/load first.", + ) + model_name = backend.active_model_name or payload.model # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( payload.messages ) - # If no non-system messages were provided, error out if not chat_messages: raise HTTPException( status_code=400, detail="At least one non-system message is required.", ) - # ── Decode image (from content parts OR legacy field) ───── - # Content-part images take priority; fall back to legacy field + # ── GGUF path: format prompt → proxy to llama-server ────── + if using_gguf: + # GGUF models don't support vision + image_b64 = extracted_image_b64 or payload.image_base64 + if image_b64: + raise HTTPException( + status_code=400, + detail="Image provided but GGUF models do not support vision.", + ) + + prompt = llama_backend.format_prompt(chat_messages, system_prompt) + cancel_event = threading.Event() + + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + created = int(time.time()) + + def gguf_generate(): + return llama_backend.generate_chat_completion( + prompt=prompt, + temperature=payload.temperature, + top_p=payload.top_p, + top_k=payload.top_k, + min_p=payload.min_p, + max_tokens=payload.max_tokens or 512, + repetition_penalty=payload.repetition_penalty, + cancel_event=cancel_event, + ) + + if payload.stream: + async def gguf_stream_chunks(): + try: + # First chunk: role + first_chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model_name, + choices=[ChunkChoice( + delta=ChoiceDelta(role="assistant"), + finish_reason=None, + )], + ) + yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + + # Content chunks — llama backend yields cumulative text + prev_text = "" + for cumulative in gguf_generate(): + if await request.is_disconnected(): + cancel_event.set() + return + new_text = cumulative[len(prev_text):] + prev_text = cumulative + if not new_text: + continue + chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model_name, + choices=[ChunkChoice( + delta=ChoiceDelta(content=new_text), + finish_reason=None, + )], + ) + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + + # Final chunk + final_chunk = ChatCompletionChunk( + id=completion_id, + created=created, + model=model_name, + choices=[ChunkChoice( + delta=ChoiceDelta(), + finish_reason="stop", + )], + ) + yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield "data: [DONE]\n\n" + + except asyncio.CancelledError: + cancel_event.set() + raise + except Exception as e: + logger.error(f"Error during GGUF streaming: {e}", exc_info=True) + error_chunk = { + "error": {"message": str(e), "type": "server_error"}, + } + yield f"data: {json.dumps(error_chunk)}\n\n" + + return StreamingResponse( + gguf_stream_chunks(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + else: + try: + full_text = "" + for token in gguf_generate(): + full_text = token + + response = ChatCompletion( + id=completion_id, + created=created, + model=model_name, + choices=[CompletionChoice( + message=CompletionMessage(content=full_text), + finish_reason="stop", + )], + ) + return JSONResponse(content=response.model_dump()) + + except Exception as e: + logger.error(f"Error during GGUF completion: {e}", exc_info=True) + raise HTTPException(status_code=500, detail=str(e)) + + # ── Standard Unsloth path ───────────────────────────────── + + # Decode image (from content parts OR legacy field) image_b64 = extracted_image_b64 or payload.image_base64 image = None @@ -363,7 +568,7 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque except Exception as e: raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}") - # ── Shared generation kwargs ────────────────────────────── + # Shared generation kwargs gen_kwargs = dict( messages=chat_messages, system_prompt=system_prompt, @@ -376,11 +581,10 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque repetition_penalty=payload.repetition_penalty, ) - # ── Choose generation path (adapter-controlled or standard) ── + # Choose generation path (adapter-controlled or standard) cancel_event = threading.Event() if payload.use_adapter is not None: - # Compare mode: toggle adapter state atomically with generation def generate(): return backend.generate_with_adapter_control( use_adapter=payload.use_adapter, @@ -388,11 +592,9 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque **gen_kwargs, ) else: - # Standard path: no adapter toggling def generate(): return backend.generate_chat_response(cancel_event=cancel_event, **gen_kwargs) - model_name = backend.active_model_name or payload.model completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -400,7 +602,6 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque if payload.stream: async def stream_chunks(): try: - # First chunk: send the role first_chunk = ChatCompletionChunk( id=completion_id, created=created, @@ -412,8 +613,6 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque ) yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" - # Content chunks — generate_chat_response yields cumulative - # text, so we diff to get incremental deltas. prev_text = "" for cumulative in generate(): if await request.is_disconnected(): @@ -435,7 +634,6 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque ) yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" - # Final chunk: finish_reason = stop final_chunk = ChatCompletionChunk( id=completion_id, created=created, @@ -475,7 +673,7 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque try: full_text = "" for token in generate(): - full_text = token # generate_stream yields cumulative text + full_text = token response = ChatCompletion( id=completion_id, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 761c04d3e7..881f399c3d 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -88,6 +88,7 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]: or (child / "adapter_config.json").exists() or any(child.glob("*.safetensors")) or any(child.glob("*.bin")) + or any(child.glob("*.gguf")) ) if not has_model_files: continue @@ -104,6 +105,23 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]: updated_at=updated_at, ), ) + # Also scan for standalone .gguf files directly in the models directory + for gguf_file in models_dir.glob("*.gguf"): + if gguf_file.is_file(): + try: + updated_at = gguf_file.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id=str(gguf_file), + display_name=gguf_file.stem, + path=str(gguf_file), + source="models_dir", + updated_at=updated_at, + ), + ) + return found diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index fdf89fce39..27fbe3ee44 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -422,6 +422,32 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: pass +def detect_gguf_model(path: str) -> Optional[str]: + """ + Check if the given path is or contains a GGUF model file. + + Handles three cases: + 1. path is a direct .gguf file path + 2. path is a directory containing .gguf files + 3. path is an HuggingFace repo with GGUF files (not yet — future enhancement) + + Returns the full path to the .gguf file if found, None otherwise. + """ + p = Path(path) + + # Case 1: direct .gguf file + if p.suffix == ".gguf" and p.is_file(): + return str(p.resolve()) + + # Case 2: directory containing .gguf files + if p.is_dir(): + gguf_files = sorted(p.glob("*.gguf"), key=lambda f: f.stat().st_size, reverse=True) + if gguf_files: + return str(gguf_files[0].resolve()) + + return None + + def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]: """ Scan outputs folder for trained LoRA adapters. @@ -595,6 +621,8 @@ class ModelConfig: is_cached: bool # Is this already in HF cache? is_vision: bool # Is this a vision model? is_lora: bool # Is this a lora adapter? + is_gguf: bool = False # Is this a GGUF model? + gguf_file: Optional[str] = None # Full path to the .gguf file base_model: Optional[str] = None # Base model (for LoRAs) @classmethod @@ -675,12 +703,30 @@ class ModelConfig: identifier = model_id.strip() is_local = is_local_path(identifier) path = normalize_path(identifier) if is_local else identifier - + # Add unsloth/ prefix for shorthand HF models if not is_local and "/" not in identifier: identifier = f"unsloth/{identifier}" path = identifier - + + # Auto-detect GGUF models (check before LoRA/vision detection) + if is_local: + gguf_file = detect_gguf_model(path) + if gguf_file: + display_name = Path(gguf_file).stem + logger.info(f"Detected GGUF model: {gguf_file}") + return cls( + identifier=identifier, + display_name=display_name, + path=path, + is_local=True, + is_cached=True, + is_vision=False, + is_lora=False, + is_gguf=True, + gguf_file=gguf_file, + ) + # Auto-detect LoRA for local paths (check adapter_config.json on disk) if not is_lora and is_local: detected_base = get_base_model_from_lora(path) diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 3dca94d7f5..030fe55bff 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -41,11 +41,13 @@ function stripTrailingEpoch(input: string): string { function describeModel(model: { is_lora?: boolean; is_vision?: boolean; + is_gguf?: boolean; }): string | undefined { const tags: string[] = []; + if (model.is_gguf) tags.push("GGUF"); if (model.is_lora) tags.push("LoRA"); if (model.is_vision) tags.push("Vision"); - if (!model.is_lora && !model.is_vision) tags.push("Base"); + if (!model.is_lora && !model.is_vision && !model.is_gguf) tags.push("Base"); return tags.join(" · "); } @@ -54,6 +56,7 @@ function toChatModelSummary(model: { name?: string | null; is_lora?: boolean; is_vision?: boolean; + is_gguf?: boolean; }): ChatModelSummary { return { id: model.id, @@ -61,6 +64,7 @@ function toChatModelSummary(model: { description: describeModel(model), isLora: Boolean(model.is_lora), isVision: Boolean(model.is_vision), + isGguf: Boolean(model.is_gguf), }; } diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index cadcad152d..bc9268e2c2 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -3,6 +3,7 @@ export interface BackendModelDetails { name?: string | null; is_vision?: boolean; is_lora?: boolean; + is_gguf?: boolean; } export interface ListModelsResponse { @@ -35,6 +36,7 @@ export interface LoadModelResponse { display_name: string; is_vision: boolean; is_lora: boolean; + is_gguf?: boolean; inference?: { temperature?: number; top_p?: number; @@ -50,6 +52,7 @@ export interface UnloadModelRequest { export interface InferenceStatusResponse { active_model: string | null; is_vision: boolean; + is_gguf?: boolean; loading: string[]; loaded: string[]; } diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 47478e5aae..9f392fe08d 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -26,6 +26,7 @@ export interface ChatModelSummary { description?: string; isVision: boolean; isLora: boolean; + isGguf?: boolean; } export interface ChatLoraSummary { diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index 6ba70a4d5c..0029f13317 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -11,7 +11,6 @@ export interface HfModelResult { } const EXCLUDED_TAGS = new Set([ - "gguf", "gptq", "awq", "exl2",