unsloth/studio/backend/routes/models.py

# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
# Copyright © 2025 Unsloth AI

"""
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,
        scan_exported_models,
        load_model_defaults,
        get_base_model_from_lora,
        is_vision_model,
        scan_checkpoints,
        list_gguf_variants,
        ModelConfig,
    )
    from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
    from core.inference import get_inference_backend
    from utils.paths import outputs_root, exports_root, resolve_output_dir, resolve_export_dir
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,
        scan_exported_models,
        load_model_defaults,
        get_base_model_from_lora,
        is_vision_model,
        scan_checkpoints,
        list_gguf_variants,
        ModelConfig,
    )
    from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
    from core.inference import get_inference_backend
    from utils.paths import outputs_root, exports_root, resolve_output_dir, resolve_export_dir

from models import (
    CheckpointInfo,
    CheckpointListResponse,
    LocalModelInfo,
    LocalModelListResponse,
    ModelCheckpoints,
    ModelDetails,
    LoRAScanResponse,
    LoRAInfo,
    ModelListResponse,
)
from models.models import GgufVariantDetail, GgufVariantsResponse
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"))
            or any(child.glob("*.gguf"))
        )
        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,
            ),
        )
    # 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


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),
                is_audio=model_data.get("is_audio", False),
                audio_type=model_data.get("audio_type"),
                has_audio_input=model_data.get("has_audio_input", 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,
    hf_token: Optional[str] = Query(None),
    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}")
        from utils.models.model_config import detect_audio_type
        # Load model defaults from backend
        config_dict = load_model_defaults(model_name)

        # Detect model capabilities (pass HF token for gated models)
        is_vision = is_vision_model(model_name)
        audio_type = detect_audio_type(model_name, hf_token=hf_token)

        # Check if it's a LoRA adapter
        is_lora = False
        base_model = None
        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:
            pass

        logger.info(f"Model config result for {model_name}: is_vision={is_vision}, audio_type={audio_type}, is_lora={is_lora}")
        return ModelDetails(
            id=model_name,
            model_name=model_name,
            config=config_dict,
            is_vision=is_vision,
            is_lora=is_lora,
            is_audio=audio_type is not None,
            audio_type=audio_type,
            has_audio_input=is_audio_input_type(audio_type),
            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=str(outputs_root()), description="Directory to scan for LoRA adapters"),
    exports_dir: str = Query(default=str(exports_root()), description="Directory to scan for exported models"),
    current_subject: str = Depends(get_current_subject),
):
    """
    Scan for trained LoRA adapters and exported models.

    Returns both training outputs (from outputs_dir) and exported models
    (from exports_dir) in a single list, distinguished by source field.
    """
    try:
        resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
        resolved_exports_dir = str(resolve_export_dir(exports_dir))
        lora_list = []

        # Scan training outputs
        trained_loras = scan_trained_loras(outputs_dir=resolved_outputs_dir)
        for display_name, adapter_path in trained_loras:
            base_model = get_base_model_from_lora(adapter_path)
            lora_list.append(LoRAInfo(
                display_name=display_name,
                adapter_path=adapter_path,
                base_model=base_model,
                source="training",
            ))

        # Scan exported models (merged, LoRA, base — skips GGUF)
        exported = scan_exported_models(exports_dir=resolved_exports_dir)
        for display_name, model_path, export_type, base_model in exported:
            lora_list.append(LoRAInfo(
                display_name=display_name,
                adapter_path=model_path,
                base_model=base_model,
                source="exported",
                export_type=export_type,
            ))

        return LoRAScanResponse(
            loras=lora_list,
            outputs_dir=resolved_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("/gguf-variants", response_model=GgufVariantsResponse)
async def get_gguf_variants(
    repo_id: str = Query(..., description="HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"),
    hf_token: Optional[str] = Query(None, description="HuggingFace token for private repos"),
    current_subject: str = Depends(get_current_subject),
):
    """
    List available GGUF quantization variants for a HuggingFace repo.

    Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.)
    with file sizes, whether the model supports vision, and the recommended
    default variant.
    """
    try:
        variants, has_vision = list_gguf_variants(repo_id, hf_token=hf_token)

        # Determine default variant
        filenames = [v.filename for v in variants]
        best = _pick_best_gguf(filenames)
        default_variant = _extract_quant_label(best) if best else None

        return GgufVariantsResponse(
            repo_id=repo_id,
            variants=[
                GgufVariantDetail(
                    filename=v.filename,
                    quant=v.quant,
                    size_bytes=v.size_bytes,
                )
                for v in variants
            ],
            has_vision=has_vision,
            default_variant=default_variant,
        )

    except Exception as e:
        logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail=f"Failed to list GGUF variants: {str(e)}",
        )


@router.get("/checkpoints", response_model=CheckpointListResponse)
async def list_checkpoints(
    outputs_dir: str = Query(
        default=str(outputs_root()),
        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:
        resolved_outputs_dir = str(resolve_output_dir(outputs_dir))
        raw_models = scan_checkpoints(outputs_dir=resolved_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=resolved_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)}",
        )