diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 7a767659ac..66df893f6c 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -87,51 +87,8 @@ class ExportBackend: Returns: List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...] """ - models = [] - outputs_path = Path(outputs_dir) - - if not outputs_path.exists(): - logger.warning(f"Outputs directory not found: {outputs_dir}") - return models - - try: - for item in outputs_path.iterdir(): - if not item.is_dir(): - continue - - config_file = item / "config.json" - adapter_config = item / "adapter_config.json" - - if not (config_file.exists() or adapter_config.exists()): - continue - - # This is a valid training run - checkpoints = [] - - # Add the final model checkpoint - checkpoints.append((item.name, str(item))) - - # Scan for intermediate checkpoints (checkpoint-N subdirs) - for sub in sorted(item.iterdir()): - if not sub.is_dir() or not sub.name.startswith("checkpoint-"): - continue - sub_config = sub / "config.json" - sub_adapter = sub / "adapter_config.json" - if sub_config.exists() or sub_adapter.exists(): - checkpoints.append((sub.name, str(sub))) - - models.append((item.name, checkpoints)) - logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)") - - # Sort by modification time (newest first) - models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True) - - logger.info(f"Found {len(models)} training runs in {outputs_dir}") - return models - - except Exception as e: - logger.error(f"Error scanning checkpoints: {e}") - return [] + from utils.models.checkpoints import scan_checkpoints + return scan_checkpoints(outputs_dir=outputs_dir) def load_checkpoint(self, checkpoint_path: str, diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index aad07fd8b0..b66efe7093 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -8,6 +8,9 @@ from .training import ( TrainingProgress, ) from .models import ( + CheckpointInfo, + ModelCheckpoints, + CheckpointListResponse, ModelDetails, LoRAInfo, LoRAScanResponse, @@ -20,9 +23,6 @@ from .auth import ( AuthStatusResponse, ) from .export import ( - CheckpointInfo, - ModelCheckpoints, - CheckpointListResponse, LoadCheckpointRequest, ExportStatusResponse, ExportOperationResponse, diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 779ddb1416..4a3dd664ee 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -5,33 +5,6 @@ from pydantic import BaseModel, Field from typing import List, Optional, Literal, Dict, Any -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") - - -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)", - ) - - -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 LoadCheckpointRequest(BaseModel): """Request for loading a checkpoint into the export backend.""" diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 8460d94200..01f161516b 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -5,6 +5,33 @@ from pydantic import BaseModel, Field from typing import Optional, List, Dict, Any +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") + + +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)", + ) + + +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") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 34df78e276..162a71b2c4 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -26,9 +26,6 @@ except ImportError: # Import Pydantic models from models import ( - CheckpointInfo, - CheckpointListResponse, - ModelCheckpoints, LoadCheckpointRequest, ExportStatusResponse, ExportOperationResponse, @@ -51,44 +48,6 @@ if not logger.handlers: logger.setLevel(logging.INFO) -@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. - - Wraps ExportBackend.scan_checkpoints. - """ - try: - backend = get_export_backend() - raw_models = backend.scan_checkpoints(outputs_dir=outputs_dir) - - models = [ - ModelCheckpoints( - name=model_name, - checkpoints=[ - CheckpointInfo(display_name=display_name, path=path) - for display_name, path in checkpoints - ], - ) - for model_name, checkpoints 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)}", - ) @router.post("/load-checkpoint", response_model=ExportOperationResponse) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1bfd08f38c..bba6f7ae34 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -21,6 +21,7 @@ try: load_model_defaults, get_base_model_from_lora, is_vision_model, + scan_checkpoints, ModelConfig, ) from core.inference import get_inference_backend @@ -34,11 +35,15 @@ except ImportError: 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, + ModelCheckpoints, ModelDetails, LoRAScanResponse, LoRAInfo, @@ -266,3 +271,40 @@ async def check_vision_model( 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) + for display_name, path in checkpoints + ], + ) + for model_name, checkpoints 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)}", + ) diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 8344d21a89..505fd35edd 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -11,6 +11,7 @@ from .model_config import ( MODEL_NAME_MAPPING, UI_STATUS_INDICATORS, ) +from .checkpoints import scan_checkpoints __all__ = [ 'ModelConfig', @@ -21,4 +22,5 @@ __all__ = [ 'load_model_config', 'MODEL_NAME_MAPPING', 'UI_STATUS_INDICATORS', + 'scan_checkpoints', ] diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py new file mode 100644 index 0000000000..3e7bd9cd6c --- /dev/null +++ b/studio/backend/utils/models/checkpoints.py @@ -0,0 +1,62 @@ +""" +Checkpoint scanning utilities for discovering training runs and their checkpoints. +""" +import logging +from pathlib import Path +from typing import List, Tuple + +logger = logging.getLogger(__name__) + + +def scan_checkpoints(outputs_dir: str = "./outputs") -> List[Tuple[str, List[Tuple[str, str]]]]: + """ + Scan outputs folder for training runs and their checkpoints. + + Returns: + List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...] + """ + models = [] + outputs_path = Path(outputs_dir) + + if not outputs_path.exists(): + logger.warning(f"Outputs directory not found: {outputs_dir}") + return models + + try: + for item in outputs_path.iterdir(): + if not item.is_dir(): + continue + + config_file = item / "config.json" + adapter_config = item / "adapter_config.json" + + if not (config_file.exists() or adapter_config.exists()): + continue + + # This is a valid training run + checkpoints = [] + + # Add the final model checkpoint + checkpoints.append((item.name, str(item))) + + # Scan for intermediate checkpoints (checkpoint-N subdirs) + for sub in sorted(item.iterdir()): + if not sub.is_dir() or not sub.name.startswith("checkpoint-"): + continue + sub_config = sub / "config.json" + sub_adapter = sub / "adapter_config.json" + if sub_config.exists() or sub_adapter.exists(): + checkpoints.append((sub.name, str(sub))) + + models.append((item.name, checkpoints)) + logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)") + + # Sort by modification time (newest first) + models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True) + + logger.info(f"Found {len(models)} training runs in {outputs_dir}") + return models + + except Exception as e: + logger.error(f"Error scanning checkpoints: {e}") + return []