784a9ed71c
Round 40 review findings (5 P1 + 1 P2 + 3 P3): P1: 1. routes/export.py: wrap /export/merged, /export/base, /export/gguf, /export/lora in a public-load window so backend.export_*() running in a worker thread cannot be torn down by a concurrent workload that sees is_export_active() == False during the pre-active gap. 2. utils/datasets/llm_assist.py: add public_load_pending_for(workload) helper. routes/inference.py: _release_export_for now refuses 503 when export is mid-handoff. 3. models/models.py: AddScanFolderRequest.path now rejects control characters and embedded hf_ tokens before being logged or reflected. 4. models/training.py: local_datasets and local_eval_datasets list entries get the same control-char / embedded-token validators that model_name / hf_dataset already have. 5. models/training.py: format_type joins the validator list (copied into training_kwargs and into trainer log lines). 6. models/export.py: _validate_save_directory now rejects embedded hf_ tokens (already covered other identifier fields). P2: 7. images-page.tsx:162: defer the mount fetchAndUpdateStatus call through setTimeout(..., 0) so it does not trip react-hooks/set-state-in-effect on scoped lint. P3 cleanup: 8. core/inference/diffusion.py: drop unused gguf_basename assignment. 9. core/inference/diffusion.py + routes/inference.py: drop unused owned_names computation from the chat-release helpers; the final sweep intentionally no longer filters by that snapshot.
295 lines
10 KiB
Python
295 lines
10 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Pydantic schemas for Model Management API
|
|
"""
|
|
|
|
from pydantic import BaseModel, Field, field_validator
|
|
from typing import Optional, List, Dict, Any, Literal
|
|
|
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
|
|
|
ModelType = Literal["text", "vision", "audio", "embeddings"]
|
|
|
|
|
|
class CheckpointInfo(BaseModel):
|
|
"""Information about a discovered checkpoint directory."""
|
|
|
|
display_name: str = Field(
|
|
..., description = "User-friendly checkpoint name (folder name)"
|
|
)
|
|
path: str = Field(..., description = "Full path to the checkpoint directory")
|
|
loss: Optional[float] = Field(None, description = "Training loss at this checkpoint")
|
|
|
|
|
|
class ModelCheckpoints(BaseModel):
|
|
"""A training run and its associated checkpoints."""
|
|
|
|
name: str = Field(..., description = "Training run folder name")
|
|
checkpoints: List[CheckpointInfo] = Field(
|
|
default_factory = list,
|
|
description = "List of checkpoints for this training run (final + intermediate)",
|
|
)
|
|
base_model: Optional[str] = Field(
|
|
None,
|
|
description = "Base model name from adapter_config.json or config.json",
|
|
)
|
|
peft_type: Optional[str] = Field(
|
|
None,
|
|
description = "PEFT type (e.g. LORA) if adapter training, None for full fine-tune",
|
|
)
|
|
lora_rank: Optional[int] = Field(
|
|
None,
|
|
description = "LoRA rank (r) if applicable",
|
|
)
|
|
is_quantized: bool = Field(
|
|
False,
|
|
description = "Whether the model uses BNB quantization (e.g. bnb-4bit)",
|
|
)
|
|
|
|
|
|
class CheckpointListResponse(BaseModel):
|
|
"""Response for listing available checkpoints in an outputs directory."""
|
|
|
|
outputs_dir: str = Field(..., description = "Directory that was scanned")
|
|
models: List[ModelCheckpoints] = Field(
|
|
default_factory = list,
|
|
description = "List of training runs with their checkpoints",
|
|
)
|
|
|
|
|
|
class ModelDetails(BaseModel):
|
|
"""Detailed model configuration and metadata - can be used for both list and detail views"""
|
|
|
|
id: str = Field(..., description = "Model identifier")
|
|
model_name: Optional[str] = Field(
|
|
None, description = "Model identifier (alias for id, for backward compatibility)"
|
|
)
|
|
name: Optional[str] = Field(None, description = "Display name for the model")
|
|
config: Optional[Dict[str, Any]] = Field(
|
|
None, description = "Model configuration dictionary"
|
|
)
|
|
is_vision: bool = Field(False, description = "Whether model is a vision model")
|
|
is_embedding: bool = Field(
|
|
False, description = "Whether model is an embedding/sentence-transformer 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)"
|
|
)
|
|
is_audio: bool = Field(False, description = "Whether model is a TTS audio model")
|
|
audio_type: Optional[str] = Field(
|
|
None, description = "Audio codec type: snac, csm, bicodec, dac"
|
|
)
|
|
has_audio_input: bool = Field(
|
|
False, description = "Whether model accepts audio input (ASR)"
|
|
)
|
|
model_type: Optional[ModelType] = Field(
|
|
None, description = "Collapsed model modality: text, vision, audio, or embeddings"
|
|
)
|
|
base_model: Optional[str] = Field(
|
|
None, description = "Base model if this is a LoRA adapter"
|
|
)
|
|
max_position_embeddings: Optional[int] = Field(
|
|
None, description = "Maximum context length supported by the model"
|
|
)
|
|
model_size_bytes: Optional[int] = Field(
|
|
None, description = "Total size of model weight files in bytes"
|
|
)
|
|
|
|
|
|
class LoRAInfo(BaseModel):
|
|
"""LoRA adapter or exported model information"""
|
|
|
|
display_name: str = Field(..., description = "Display name for the LoRA")
|
|
adapter_path: str = Field(
|
|
..., description = "Path to the LoRA adapter or exported model"
|
|
)
|
|
base_model: Optional[str] = Field(None, description = "Base model identifier")
|
|
source: Optional[str] = Field(None, description = "'training' or 'exported'")
|
|
export_type: Optional[str] = Field(
|
|
None, description = "'lora', 'merged', or 'gguf' (for exports)"
|
|
)
|
|
|
|
|
|
class LoRAScanResponse(BaseModel):
|
|
"""Response schema for scanning trained LoRA adapters"""
|
|
|
|
loras: List[LoRAInfo] = Field(
|
|
default_factory = list, description = "List of found LoRA adapters"
|
|
)
|
|
outputs_dir: str = Field(..., description = "Directory that was scanned")
|
|
|
|
|
|
class ModelListResponse(BaseModel):
|
|
"""Response schema for listing models"""
|
|
|
|
models: List[ModelDetails] = Field(
|
|
default_factory = list, description = "List of models"
|
|
)
|
|
default_models: List[str] = Field(
|
|
default_factory = list, description = "List of default model IDs"
|
|
)
|
|
|
|
|
|
class GgufVariantDetail(BaseModel):
|
|
"""A single GGUF quantization variant in a HuggingFace repo."""
|
|
|
|
filename: str = Field(
|
|
..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')"
|
|
)
|
|
quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')")
|
|
size_bytes: int = Field(0, description = "File size in bytes")
|
|
downloaded: bool = Field(
|
|
False, description = "Whether this variant is already in the local HF cache"
|
|
)
|
|
|
|
|
|
class GgufVariantsResponse(BaseModel):
|
|
"""Response for listing GGUF quantization variants in a HuggingFace repo."""
|
|
|
|
repo_id: str = Field(..., description = "HuggingFace repo ID")
|
|
variants: List[GgufVariantDetail] = Field(
|
|
default_factory = list, description = "Available GGUF variants"
|
|
)
|
|
has_vision: bool = Field(
|
|
False, description = "Whether the model has vision support (mmproj files)"
|
|
)
|
|
default_variant: Optional[str] = Field(
|
|
None, description = "Recommended default quantization variant"
|
|
)
|
|
|
|
|
|
class LocalModelInfo(BaseModel):
|
|
"""Discovered local model candidate."""
|
|
|
|
id: str = Field(..., description = "Identifier to use for loading/training")
|
|
display_name: str = Field(..., description = "Display label")
|
|
path: str = Field(..., description = "Local path where model data was discovered")
|
|
source: Literal["models_dir", "hf_cache", "lmstudio", "custom"] = Field(
|
|
...,
|
|
description = "Discovery source",
|
|
)
|
|
model_id: Optional[str] = Field(
|
|
None,
|
|
description = "HF repo id for cached models, e.g. org/model",
|
|
)
|
|
updated_at: Optional[float] = Field(
|
|
None,
|
|
description = "Unix timestamp of latest observed update",
|
|
)
|
|
|
|
|
|
class LocalModelListResponse(BaseModel):
|
|
"""Response schema for listing local/cached models."""
|
|
|
|
models_dir: str = Field(
|
|
..., description = "Directory scanned for custom local models"
|
|
)
|
|
hf_cache_dir: Optional[str] = Field(
|
|
None,
|
|
description = "HF cache root that was scanned",
|
|
)
|
|
lmstudio_dirs: List[str] = Field(
|
|
default_factory = list,
|
|
description = "LM Studio model directories that were scanned",
|
|
)
|
|
models: List[LocalModelInfo] = Field(
|
|
default_factory = list,
|
|
description = "Discovered local/cached models",
|
|
)
|
|
|
|
|
|
class AddScanFolderRequest(BaseModel):
|
|
"""Request body for adding a custom scan folder."""
|
|
|
|
path: str = Field(
|
|
..., description = "Absolute or relative directory path to scan for models"
|
|
)
|
|
|
|
# path is reflected back in /scan-folders error details and logged
|
|
# via add_scan_folder_endpoint when the directory is missing, so
|
|
# apply the same identifier hardening used on other logged paths.
|
|
@field_validator("path")
|
|
@classmethod
|
|
def _no_path_control_chars(cls, v, info):
|
|
return _no_control_chars(v, info.field_name)
|
|
|
|
@field_validator("path")
|
|
@classmethod
|
|
def _no_path_embedded_hf_tokens(cls, v, info):
|
|
return _reject_embedded_hf_token(v, info.field_name)
|
|
|
|
|
|
class ScanFolderInfo(BaseModel):
|
|
"""A registered custom model scan folder."""
|
|
|
|
id: int = Field(..., description = "Database row ID")
|
|
path: str = Field(..., description = "Normalized absolute path")
|
|
created_at: str = Field(..., description = "ISO 8601 creation timestamp")
|
|
|
|
|
|
class BrowseEntry(BaseModel):
|
|
"""A directory entry surfaced by the folder browser."""
|
|
|
|
name: str = Field(..., description = "Entry name (basename, not full path)")
|
|
has_models: bool = Field(
|
|
False,
|
|
description = (
|
|
"Hint that the directory likely contains models "
|
|
"(*.gguf, *.safetensors, config.json, or HF-style "
|
|
"`models--*` subfolders). Used by the UI to highlight "
|
|
"promising candidates; the scanner itself is authoritative."
|
|
),
|
|
)
|
|
hidden: bool = Field(
|
|
False,
|
|
description = "Name starts with a dot (e.g. `.cache`)",
|
|
)
|
|
|
|
|
|
class BrowseFoldersResponse(BaseModel):
|
|
"""Response schema for the folder browser endpoint."""
|
|
|
|
current: str = Field(..., description = "Absolute path of the directory just listed")
|
|
parent: Optional[str] = Field(
|
|
None,
|
|
description = (
|
|
"Parent directory of `current`, or null if `current` is the "
|
|
"filesystem root. The frontend uses this to render an `Up` row."
|
|
),
|
|
)
|
|
entries: List[BrowseEntry] = Field(
|
|
default_factory = list,
|
|
description = (
|
|
"Subdirectories of `current`. Sorted with model-bearing "
|
|
"directories first, then alphabetically case-insensitive; "
|
|
"hidden entries come last within each group."
|
|
),
|
|
)
|
|
suggestions: List[str] = Field(
|
|
default_factory = list,
|
|
description = (
|
|
"Handy starting points (home, HF cache, already-registered "
|
|
"scan folders). Rendered as quick-pick chips above the list."
|
|
),
|
|
)
|
|
truncated: bool = Field(
|
|
False,
|
|
description = (
|
|
"True when the listing was capped because the directory had "
|
|
"more subfolders than the server is willing to enumerate in "
|
|
"one request. The UI should show a hint telling the user to "
|
|
"narrow their path."
|
|
),
|
|
)
|
|
model_files_here: int = Field(
|
|
0,
|
|
description = (
|
|
"Count of GGUF/safetensors files immediately inside "
|
|
"``current``. Used by the UI to surface a hint on leaf "
|
|
"model directories (which otherwise look `empty` because "
|
|
"they contain only files, no subdirectories)."
|
|
),
|
|
)
|