diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index da5b11c60d..231566cc22 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -2,6 +2,7 @@ """ Export backend - handles model exporting in various formats """ +import json import logging import os from pathlib import Path @@ -200,6 +201,18 @@ class ExportBackend: logger.error(traceback.format_exc()) return False, f"Failed to load checkpoint: {str(e)}" + def _write_export_metadata(self, save_directory: str): + """Write export_metadata.json with base model info for Chat page discovery.""" + try: + base_model = get_base_model_from_lora(self.current_checkpoint) if self.current_checkpoint else None + metadata = {"base_model": base_model} + metadata_path = os.path.join(save_directory, "export_metadata.json") + with open(metadata_path, "w") as f: + json.dump(metadata, f, indent=2) + logger.info(f"Wrote export metadata to {metadata_path}") + except Exception as e: + logger.warning(f"Could not write export metadata: {e}") + def export_merged_model(self, save_directory: str, format_type: str = "16-bit (FP16)", @@ -244,6 +257,9 @@ class ExportBackend: self.current_tokenizer, save_method=save_method ) + + # Write export metadata so the Chat page can identify the base model + self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") # Push to hub if requested @@ -297,6 +313,9 @@ class ExportBackend: self.current_model.save_pretrained(save_directory) self.current_tokenizer.save_pretrained(save_directory) + + # Write export metadata so the Chat page can identify the base model + self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") # Push to hub if requested diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 5542db76d5..10d87d7825 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -57,10 +57,12 @@ class ModelDetails(BaseModel): class LoRAInfo(BaseModel): - """LoRA adapter information""" + """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") + 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' or 'merged' (for exports)") class LoRAScanResponse(BaseModel): diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 761c04d3e7..8c96ee5f28 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -18,6 +18,7 @@ from auth.authentication import get_current_subject try: from utils.models import ( scan_trained_loras, + scan_exported_models, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -32,6 +33,7 @@ except ImportError: 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, @@ -289,35 +291,45 @@ async def get_model_config( @router.get("/loras") async def scan_loras( outputs_dir: str = Query(default="./outputs", description="Directory to scan for LoRA adapters"), + exports_dir: str = Query(default="./exports", description="Directory to scan for exported models"), 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. + 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: - # Call backend scan function - trained_loras = scan_trained_loras(outputs_dir=outputs_dir) - - # Convert to LoRAInfo objects lora_list = [] + + # Scan training outputs + trained_loras = scan_trained_loras(outputs_dir=outputs_dir) 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( + lora_list.append(LoRAInfo( display_name=display_name, adapter_path=adapter_path, - base_model=base_model - ) - lora_list.append(lora_info) - + base_model=base_model, + source="training", + )) + + # Scan exported models (merged, LoRA, base — skips GGUF) + exported = scan_exported_models(exports_dir=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=outputs_dir ) - + except Exception as e: logger.error(f"Error scanning LoRAs: {e}", exc_info=True) raise HTTPException( diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 505fd35edd..006deb99c0 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -5,6 +5,7 @@ from .model_config import ( ModelConfig, is_vision_model, scan_trained_loras, + scan_exported_models, load_model_defaults, get_base_model_from_lora, load_model_config, @@ -17,6 +18,7 @@ __all__ = [ 'ModelConfig', 'is_vision_model', 'scan_trained_loras', + 'scan_exported_models', 'load_model_defaults', 'get_base_model_from_lora', 'load_model_config', diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index fdf89fce39..f7b95fad11 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -465,6 +465,90 @@ def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]: logger.error(f"Error scanning outputs folder: {e}") return [] +def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str, str, Optional[str]]]: + """ + Scan exports folder for exported models (merged, LoRA, base). + Skips GGUF-only exports (not loadable by Unsloth inference backend). + + The exports directory is two levels deep: {run}/{checkpoint}/ + + Returns: + List of tuples: [(display_name, model_path, export_type, base_model), ...] + export_type: "lora" | "merged" + """ + results = [] + exports_path = Path(exports_dir) + + if not exports_path.exists(): + return results + + try: + for run_dir in exports_path.iterdir(): + if not run_dir.is_dir(): + continue + for checkpoint_dir in run_dir.iterdir(): + if not checkpoint_dir.is_dir(): + continue + + adapter_config = checkpoint_dir / "adapter_config.json" + config_file = checkpoint_dir / "config.json" + has_weights = ( + any(checkpoint_dir.glob("*.safetensors")) + or any(checkpoint_dir.glob("*.bin")) + ) + has_gguf = any(checkpoint_dir.glob("*.gguf")) + + base_model = None + export_type = None + + if adapter_config.exists(): + export_type = "lora" + try: + cfg = json.loads(adapter_config.read_text()) + base_model = cfg.get("base_model_name_or_path") + except Exception: + pass + elif config_file.exists() and has_weights: + export_type = "merged" + # Read base model from export_metadata.json (written at export time) + export_meta = checkpoint_dir / "export_metadata.json" + try: + if export_meta.exists(): + meta = json.loads(export_meta.read_text()) + base_model = meta.get("base_model") + except Exception: + pass + elif has_gguf: + # GGUF-only — not loadable by current inference backend + continue + else: + continue + + # Fallback: read base model from the original training run's + # adapter_config.json in ./outputs/{run_name}/ + if not base_model: + outputs_adapter_cfg = Path("./outputs") / run_dir.name / "adapter_config.json" + try: + if outputs_adapter_cfg.exists(): + cfg = json.loads(outputs_adapter_cfg.read_text()) + base_model = cfg.get("base_model_name_or_path") + except Exception: + pass + + display_name = f"{run_dir.name} / {checkpoint_dir.name}" + model_path = str(checkpoint_dir) + results.append((display_name, model_path, export_type, base_model)) + logger.debug(f"Found exported model: {display_name} ({export_type})") + + results.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True) + logger.info(f"Found {len(results)} exported models in {exports_dir}") + return results + + except Exception as e: + logger.error(f"Error scanning exports folder: {e}") + return [] + + def get_base_model_from_lora(lora_path: str) -> Optional[str]: """ Read the base model name from a LoRA adapter's config. diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 80e9dd0b6b..441fdf0f02 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -365,15 +365,26 @@ export function LoraModelPicker({
{index > 0 ?
: null} {baseModel} - {adapters.map((adapter) => ( - onSelect(adapter.id, { source: "lora", isLora: true })} - /> - ))} + {adapters.map((adapter) => { + const isExported = adapter.source === "exported"; + const isMerged = adapter.exportType === "merged"; + const tag = isExported + ? isMerged ? "Merged" : "LoRA" + : "LoRA"; + const meta = isExported ? `${tag} · Exported` : tag; + return ( + onSelect(adapter.id, { + source: isExported ? "exported" : "lora", + isLora: !isMerged, + })} + /> + ); + })}
)) )} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index dcf110bfb7..b8df0f6c6c 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -10,10 +10,12 @@ export interface ModelOption { export interface LoraModelOption extends ModelOption { baseModel?: string; updatedAt?: number; + source?: "training" | "exported"; + exportType?: "lora" | "merged"; } export interface ModelSelectorChangeMeta { - source: "hub" | "lora"; + source: "hub" | "lora" | "exported"; isLora: boolean; } diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index c363704d61..6e37468e71 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -410,6 +410,8 @@ export function ChatPage(): ReactElement { name: lora.name, baseModel: lora.baseModel, updatedAt: lora.updatedAt, + source: lora.source, + exportType: lora.exportType, })), [lorasFromStore], ); 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..1c8922e929 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 @@ -68,6 +68,8 @@ function toLoraSummary(lora: { display_name: string; adapter_path: string; base_model?: string | null; + source?: "training" | "exported" | null; + export_type?: "lora" | "merged" | null; }): ChatLoraSummary { const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? ""; const updatedAt = @@ -78,6 +80,8 @@ function toLoraSummary(lora: { name: stripTrailingEpoch(lora.display_name), baseModel: lora.base_model || "Unknown base model", updatedAt, + source: lora.source ?? undefined, + exportType: lora.export_type ?? undefined, }; } diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index cadcad152d..003c3b3629 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -14,6 +14,8 @@ export interface BackendLoraInfo { display_name: string; adapter_path: string; base_model?: string | null; + source?: "training" | "exported" | null; + export_type?: "lora" | "merged" | null; } export interface ListLorasResponse { diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 47478e5aae..87eb4c565c 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -33,4 +33,6 @@ export interface ChatLoraSummary { name: string; baseModel: string; updatedAt?: number; + source?: "training" | "exported"; + exportType?: "lora" | "merged"; }