429 lines
13 KiB
Python
429 lines
13 KiB
Python
"""
|
|
Model Management API routes
|
|
"""
|
|
import sys
|
|
from pathlib import Path
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from typing import List, Optional
|
|
import logging
|
|
|
|
# Add backend directory to path
|
|
backend_path = Path(__file__).parent.parent.parent
|
|
if str(backend_path) not in sys.path:
|
|
sys.path.insert(0, str(backend_path))
|
|
|
|
from auth.authentication import get_current_subject
|
|
|
|
# Import backend functions
|
|
try:
|
|
from utils.models import (
|
|
scan_trained_loras,
|
|
load_model_defaults,
|
|
get_base_model_from_lora,
|
|
is_vision_model,
|
|
scan_checkpoints,
|
|
ModelConfig,
|
|
)
|
|
from core.inference import get_inference_backend
|
|
except ImportError:
|
|
# Fallback: try to import from parent directory
|
|
parent_backend = backend_path.parent / "backend"
|
|
if str(parent_backend) not in sys.path:
|
|
sys.path.insert(0, str(parent_backend))
|
|
from utils.models import (
|
|
scan_trained_loras,
|
|
load_model_defaults,
|
|
get_base_model_from_lora,
|
|
is_vision_model,
|
|
scan_checkpoints,
|
|
ModelConfig,
|
|
)
|
|
from core.inference import get_inference_backend
|
|
|
|
from models import (
|
|
CheckpointInfo,
|
|
CheckpointListResponse,
|
|
LocalModelInfo,
|
|
LocalModelListResponse,
|
|
ModelCheckpoints,
|
|
ModelDetails,
|
|
LoRAScanResponse,
|
|
LoRAInfo,
|
|
ModelListResponse,
|
|
)
|
|
from models.responses import LoRABaseModelResponse, VisionCheckResponse
|
|
|
|
router = APIRouter()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Configure logger
|
|
if not logger.handlers:
|
|
handler = logging.StreamHandler()
|
|
handler.setLevel(logging.INFO)
|
|
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
handler.setFormatter(formatter)
|
|
logger.addHandler(handler)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
def _resolve_hf_cache_dir() -> Path:
|
|
"""Resolve local HF cache root used by hub downloads."""
|
|
try:
|
|
from huggingface_hub.constants import HF_HUB_CACHE
|
|
return Path(HF_HUB_CACHE)
|
|
except Exception:
|
|
return Path.home() / ".cache" / "huggingface" / "hub"
|
|
|
|
|
|
def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
|
|
if not models_dir.exists() or not models_dir.is_dir():
|
|
return []
|
|
|
|
found: List[LocalModelInfo] = []
|
|
for child in models_dir.iterdir():
|
|
if not child.is_dir():
|
|
continue
|
|
has_model_files = (
|
|
(child / "config.json").exists()
|
|
or (child / "adapter_config.json").exists()
|
|
or any(child.glob("*.safetensors"))
|
|
or any(child.glob("*.bin"))
|
|
)
|
|
if not has_model_files:
|
|
continue
|
|
try:
|
|
updated_at = child.stat().st_mtime
|
|
except OSError:
|
|
updated_at = None
|
|
found.append(
|
|
LocalModelInfo(
|
|
id=str(child),
|
|
display_name=child.name,
|
|
path=str(child),
|
|
source="models_dir",
|
|
updated_at=updated_at,
|
|
),
|
|
)
|
|
return found
|
|
|
|
|
|
def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]:
|
|
if not cache_dir.exists() or not cache_dir.is_dir():
|
|
return []
|
|
|
|
found: List[LocalModelInfo] = []
|
|
for repo_dir in cache_dir.glob("models--*"):
|
|
if not repo_dir.is_dir():
|
|
continue
|
|
|
|
repo_name = repo_dir.name[len("models--"):]
|
|
if not repo_name:
|
|
continue
|
|
model_id = repo_name.replace("--", "/")
|
|
|
|
try:
|
|
updated_at = repo_dir.stat().st_mtime
|
|
except OSError:
|
|
updated_at = None
|
|
|
|
found.append(
|
|
LocalModelInfo(
|
|
id=model_id,
|
|
model_id=model_id,
|
|
display_name=model_id.split("/")[-1],
|
|
path=str(repo_dir),
|
|
source="hf_cache",
|
|
updated_at=updated_at,
|
|
),
|
|
)
|
|
return found
|
|
|
|
|
|
@router.get("/local", response_model=LocalModelListResponse)
|
|
async def list_local_models(
|
|
models_dir: str = Query(default="./models", description="Directory to scan for local model folders"),
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
List local model candidates from custom models dir and HF cache.
|
|
"""
|
|
try:
|
|
models_root = Path(models_dir).expanduser().resolve()
|
|
hf_cache_dir = _resolve_hf_cache_dir()
|
|
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
|
|
|
|
deduped: dict[str, LocalModelInfo] = {}
|
|
for model in local_models:
|
|
if model.id not in deduped:
|
|
deduped[model.id] = model
|
|
|
|
models = sorted(
|
|
deduped.values(),
|
|
key=lambda item: (item.updated_at or 0),
|
|
reverse=True,
|
|
)
|
|
|
|
return LocalModelListResponse(
|
|
models_dir=str(models_root),
|
|
hf_cache_dir=str(hf_cache_dir),
|
|
models=models,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error listing local models: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to list local models: {str(e)}",
|
|
)
|
|
|
|
|
|
|
|
|
|
@router.get("/list")
|
|
async def list_models(
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
List available models (default models and loaded models).
|
|
|
|
This endpoint returns the default models and any currently loaded models.
|
|
"""
|
|
try:
|
|
inference_backend = get_inference_backend()
|
|
|
|
# Get default models
|
|
default_models = inference_backend.default_models
|
|
|
|
# Get loaded models
|
|
loaded_models = []
|
|
for model_name, model_data in inference_backend.models.items():
|
|
model_info = ModelDetails(
|
|
id=model_name,
|
|
name=model_name.split("/")[-1] if "/" in model_name else model_name,
|
|
is_vision=model_data.get("is_vision", False),
|
|
is_lora=model_data.get("is_lora", False)
|
|
)
|
|
loaded_models.append(model_info)
|
|
|
|
# Combine default and loaded models
|
|
all_models = []
|
|
seen_ids = set()
|
|
|
|
# Add default models
|
|
for model_id in default_models:
|
|
if model_id not in seen_ids:
|
|
model_info = ModelDetails(
|
|
id=model_id,
|
|
name=model_id.split("/")[-1] if "/" in model_id else model_id
|
|
)
|
|
all_models.append(model_info)
|
|
seen_ids.add(model_id)
|
|
|
|
# Add loaded models
|
|
for model_info in loaded_models:
|
|
if model_info.id not in seen_ids:
|
|
all_models.append(model_info)
|
|
seen_ids.add(model_info.id)
|
|
|
|
return ModelListResponse(
|
|
models=all_models,
|
|
default_models=default_models
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error listing models: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to list models: {str(e)}"
|
|
)
|
|
|
|
|
|
@router.get("/config/{model_name:path}")
|
|
async def get_model_config(
|
|
model_name: str,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Get configuration for a specific model.
|
|
|
|
This endpoint wraps the backend load_model_defaults function.
|
|
"""
|
|
try:
|
|
logger.info(f"Getting model config for: {model_name}")
|
|
# Load model defaults from backend
|
|
config_dict = load_model_defaults(model_name)
|
|
|
|
# Check if it's a vision model
|
|
is_vision = is_vision_model(model_name)
|
|
|
|
# Check if it's a LoRA adapter
|
|
is_lora = False
|
|
base_model = None
|
|
|
|
# Try to create ModelConfig to get more info
|
|
try:
|
|
model_config = ModelConfig.from_identifier(model_name)
|
|
is_lora = model_config.is_lora
|
|
base_model = model_config.base_model if is_lora else None
|
|
except Exception:
|
|
# If ModelConfig creation fails, use defaults
|
|
pass
|
|
|
|
logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_lora={is_lora}, base_model={base_model}")
|
|
return ModelDetails(
|
|
id=model_name,
|
|
model_name=model_name,
|
|
config=config_dict,
|
|
is_vision=is_vision,
|
|
is_lora=is_lora,
|
|
base_model=base_model,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error getting model config: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to get model config: {str(e)}"
|
|
)
|
|
|
|
|
|
@router.get("/loras")
|
|
async def scan_loras(
|
|
outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters"),
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Scan for trained LoRA adapters in the outputs directory.
|
|
|
|
This endpoint wraps the backend scan_trained_loras function.
|
|
"""
|
|
try:
|
|
# Call backend scan function
|
|
trained_loras = scan_trained_loras(outputs_dir=outputs_dir)
|
|
|
|
# Convert to LoRAInfo objects
|
|
lora_list = []
|
|
for display_name, adapter_path in trained_loras:
|
|
# Get base model if available
|
|
base_model = get_base_model_from_lora(adapter_path)
|
|
|
|
lora_info = LoRAInfo(
|
|
display_name=display_name,
|
|
adapter_path=adapter_path,
|
|
base_model=base_model
|
|
)
|
|
lora_list.append(lora_info)
|
|
|
|
return LoRAScanResponse(
|
|
loras=lora_list,
|
|
outputs_dir=outputs_dir
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error scanning LoRAs: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to scan LoRA adapters: {str(e)}"
|
|
)
|
|
|
|
|
|
@router.get("/loras/{lora_path:path}/base-model", response_model=LoRABaseModelResponse)
|
|
async def get_lora_base_model(
|
|
lora_path: str,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Get the base model for a LoRA adapter.
|
|
|
|
This endpoint wraps the backend get_base_model_from_lora function.
|
|
"""
|
|
try:
|
|
base_model = get_base_model_from_lora(lora_path)
|
|
|
|
if base_model is None:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"Could not determine base model for LoRA: {lora_path}"
|
|
)
|
|
|
|
return LoRABaseModelResponse(
|
|
lora_path=lora_path,
|
|
base_model=base_model,
|
|
)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error getting LoRA base model: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to get base model: {str(e)}"
|
|
)
|
|
|
|
|
|
@router.get("/check-vision/{model_name:path}", response_model=VisionCheckResponse)
|
|
async def check_vision_model(
|
|
model_name: str,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Check if a model is a vision model.
|
|
|
|
This endpoint wraps the backend is_vision_model function.
|
|
"""
|
|
try:
|
|
logger.info(f"Checking if vision model: {model_name}")
|
|
is_vision = is_vision_model(model_name)
|
|
|
|
logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
|
|
return VisionCheckResponse(
|
|
model_name=model_name,
|
|
is_vision=is_vision,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error checking vision model: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to check vision model: {str(e)}"
|
|
)
|
|
|
|
@router.get("/checkpoints", response_model=CheckpointListResponse)
|
|
async def list_checkpoints(
|
|
outputs_dir: str = Query(
|
|
default="./outputs",
|
|
description="Directory to scan for checkpoints",
|
|
),
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
List available checkpoints in the outputs directory.
|
|
|
|
Scans the outputs folder for training runs and their checkpoints.
|
|
"""
|
|
try:
|
|
raw_models = scan_checkpoints(outputs_dir=outputs_dir)
|
|
|
|
models = [
|
|
ModelCheckpoints(
|
|
name=model_name,
|
|
checkpoints=[
|
|
CheckpointInfo(display_name=display_name, path=path, loss=loss)
|
|
for display_name, path, loss in checkpoints
|
|
],
|
|
base_model=metadata.get("base_model"),
|
|
peft_type=metadata.get("peft_type"),
|
|
lora_rank=metadata.get("lora_rank"),
|
|
)
|
|
for model_name, checkpoints, metadata in raw_models
|
|
]
|
|
|
|
return CheckpointListResponse(
|
|
outputs_dir=outputs_dir,
|
|
models=models,
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error listing checkpoints: {e}", exc_info=True)
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to list checkpoints: {str(e)}",
|
|
)
|