Added the training and models routes

backend/backend/__init__.py

@@ -0,0 +1,44 @@
+"""
+Unified backend module for Unsloth
+"""
+
+# Inference
+from .inference import InferenceBackend
+
+# Training
+from .trainer import UnslothTrainer, get_trainer
+from .training import TrainingBackend, get_training_backend, create_training_handlers
+
+# Configuration
+from .model_config import is_vision_model, ModelConfig, scan_trained_loras
+# Utilities
+from .path_utils import normalize_path, is_local_path, is_model_cached
+from .utils import without_hf_auth, format_error_message, get_gpu_memory_info, search_hf_models
+from .dataset_utils import format_and_template_dataset
+
+__all__ = [
+    # Inference
+    'InferenceBackend',
+
+    # Training
+    'UnslothTrainer',
+    'get_trainer',
+    'get_training_backend',
+    'TrainingBackend',
+    "create_training_handlers",
+
+    # Config
+    'ModelConfig',
+    'is_vision_model',
+    'scan_trained_loras',
+
+    # Utils
+    'search_hf_models',
+    'format_and_template_dataset',
+    'normalize_path',
+    'is_local_path',
+    'is_model_cached',
+    'without_hf_auth',
+    'format_error_message',
+    'get_gpu_memory_info',
+]

backend/backend/export.py

@@ -0,0 +1,506 @@
+# backend/export.py
+"""
+Export backend - handles model exporting in various formats
+"""
+import logging
+import os
+from pathlib import Path
+from typing import Optional, Tuple, List
+from peft import PeftModel, PeftModelForCausalLM
+from unsloth import FastLanguageModel, FastVisionModel
+from huggingface_hub import HfApi, ModelCard
+from transformers.modeling_utils import PushToHubMixin
+import torch
+
+from .model_config import is_vision_model, get_base_model_from_lora
+from .inference import get_inference_backend
+
+logger = logging.getLogger(__name__)
+
+# Model card template
+MODEL_CARD = \
+"""---
+base_model: {base_model}
+tags:
+- text-generation-inference
+- transformers
+- unsloth
+- {model_type}
+- {extra}
+license: apache-2.0
+language:
+- en
+---
+
+# Uploaded finetuned {method} model
+
+- **Developed by:** {username}
+- **License:** apache-2.0
+- **Finetuned from model :** {base_model}
+
+This {model_type} model was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth) and Huggingface's TRL library.
+
+[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)
+"""
+
+class ExportBackend:
+    """Handles model export operations"""
+
+    def __init__(self):
+        self.inference_backend = get_inference_backend()
+        self.current_checkpoint = None
+        self.current_model = None
+        self.current_tokenizer = None
+        self.is_vision = False
+        self.is_peft = False
+
+    def cleanup_memory(self):
+        """Offload and delete all models from memory"""
+        try:
+            logger.info("Starting memory cleanup...")
+
+            # Unload all models from inference backend
+            model_names = list(self.inference_backend.models.keys())
+            for model_name in model_names:
+                self.inference_backend.unload_model(model_name)
+
+            # Clear current export state
+            self.current_model = None
+            self.current_tokenizer = None
+            self.current_checkpoint = None
+
+            # Force garbage collection
+            import gc
+            gc.collect()
+
+            # Clear CUDA cache
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+
+            logger.info("Memory cleanup completed successfully")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error during memory cleanup: {e}")
+            return False
+
+    def scan_checkpoints(self, outputs_dir: str = "./outputs") -> List[Tuple[str, str]]:
+        """
+        Scan outputs folder for model checkpoints.
+
+        Returns:
+            List of tuples: [(display_name, checkpoint_path), ...]
+        """
+        checkpoints = []
+        outputs_path = Path(outputs_dir)
+
+        if not outputs_path.exists():
+            logger.warning(f"Outputs directory not found: {outputs_dir}")
+            return checkpoints
+
+        try:
+            for item in outputs_path.iterdir():
+                if item.is_dir():
+                    # Check if this directory contains a model
+                    config_file = item / "config.json"
+                    adapter_config = item / "adapter_config.json"
+
+                    if config_file.exists() or adapter_config.exists():
+                        # This is a valid checkpoint
+                        display_name = item.name
+                        checkpoint_path = str(item)
+                        checkpoints.append((display_name, checkpoint_path))
+                        logger.debug(f"Found checkpoint: {display_name}")
+
+            # Sort by modification time (newest first)
+            checkpoints.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
+
+            logger.info(f"Found {len(checkpoints)} checkpoints in {outputs_dir}")
+            return checkpoints
+
+        except Exception as e:
+            logger.error(f"Error scanning checkpoints: {e}")
+            return []
+
+    def load_checkpoint(self,
+                       checkpoint_path: str,
+                       max_seq_length: int = 2048,
+                       load_in_4bit: bool = True) -> Tuple[bool, str]:
+        """
+        Load a checkpoint for export.
+
+        Returns:
+            Tuple of (success: bool, message: str)
+        """
+        try:
+            logger.info(f"Loading checkpoint: {checkpoint_path}")
+
+            # First, cleanup existing models
+            self.cleanup_memory()
+
+            # Detect if vision model
+            checkpoint_path_obj = Path(checkpoint_path)
+
+            # Check if it's a LoRA adapter
+            adapter_config = checkpoint_path_obj / "adapter_config.json"
+            if adapter_config.exists():
+                # It's a LoRA - get base model to check vision
+                base_model = get_base_model_from_lora(checkpoint_path)
+                if base_model:
+                    self.is_vision = is_vision_model(base_model)
+                else:
+                    return False, "Could not determine base model for adapter"
+            else:
+                # Check the model itself
+                self.is_vision = is_vision_model(checkpoint_path)
+
+            # Load model based on type
+            if self.is_vision:
+                logger.info("Loading as vision model...")
+                model, processor = FastVisionModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    max_seq_length=max_seq_length,
+                    dtype=None,
+                    load_in_4bit=load_in_4bit,
+                )
+                tokenizer = processor  # For vision models, processor acts as tokenizer
+            else:
+                logger.info("Loading as text model...")
+                model, tokenizer = FastLanguageModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    max_seq_length=max_seq_length,
+                    dtype=None,
+                    load_in_4bit=load_in_4bit,
+                )
+
+            # Check if PEFT model
+            self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM))
+
+            # Store loaded model
+            self.current_model = model
+            self.current_tokenizer = tokenizer
+            self.current_checkpoint = checkpoint_path
+
+            model_type = "Vision" if self.is_vision else "Text"
+            peft_info = " (PEFT Adapter)" if self.is_peft else " (Merged Model)"
+
+            logger.info(f"Successfully loaded {model_type} model{peft_info}")
+            return True, f"Loaded {model_type} model{peft_info} successfully"
+
+        except Exception as e:
+            logger.error(f"Error loading checkpoint: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return False, f"Failed to load checkpoint: {str(e)}"
+
+    def export_merged_model(self,
+                           save_directory: str,
+                           format_type: str = "16-bit (FP16)",
+                           push_to_hub: bool = False,
+                           repo_id: Optional[str] = None,
+                           hf_token: Optional[str] = None,
+                           private: bool = False) -> Tuple[bool, str]:
+        """
+        Export merged model (for PEFT models).
+
+        Args:
+            save_directory: Local directory to save model
+            format_type: "16-bit (FP16)" or "4-bit (FP4)"
+            push_to_hub: Whether to push to Hugging Face Hub
+            repo_id: Hub repository ID (username/model-name)
+            hf_token: Hugging Face token
+            private: Whether to make the repo private
+
+        Returns:
+            Tuple of (success: bool, message: str)
+        """
+        if not self.current_model or not self.current_tokenizer:
+            return False, "No model loaded. Please select a checkpoint first."
+
+        if not self.is_peft:
+            return False, "This is not a PEFT model. Use 'Export Base Model' instead."
+
+        try:
+            # Determine save method
+            if format_type == "4-bit (FP4)":
+                save_method = "merged_4bit_forced"
+            else:  # 16-bit (FP16)
+                save_method = "merged_16bit"
+
+            # Save locally if requested
+            if save_directory:
+                logger.info(f"Saving merged model locally to: {save_directory}")
+                os.makedirs(save_directory, exist_ok=True)
+
+                self.current_model.save_pretrained_merged(
+                    save_directory,
+                    self.current_tokenizer,
+                    save_method=save_method
+                )
+                logger.info(f"Model saved successfully to {save_directory}")
+
+            # Push to hub if requested
+            if push_to_hub:
+                if not repo_id or not hf_token:
+                    return False, "Repository ID and Hugging Face token required for Hub upload"
+
+                logger.info(f"Pushing merged model to Hub: {repo_id}")
+
+                self.current_model.push_to_hub_merged(
+                    repo_id,
+                    self.current_tokenizer,
+                    save_method=save_method,
+                    token=hf_token,
+                    private=private
+                )
+                logger.info(f"Model pushed successfully to {repo_id}")
+
+            return True, "Model exported successfully"
+
+        except Exception as e:
+            logger.error(f"Error exporting merged model: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return False, f"Export failed: {str(e)}"
+
+    def export_base_model(self,
+                         save_directory: str,
+                         push_to_hub: bool = False,
+                         repo_id: Optional[str] = None,
+                         hf_token: Optional[str] = None,
+                         private: bool = False) -> Tuple[bool, str]:
+        """
+        Export base model (for non-PEFT models).
+
+        Returns:
+            Tuple of (success: bool, message: str)
+        """
+        if not self.current_model or not self.current_tokenizer:
+            return False, "No model loaded. Please select a checkpoint first."
+
+        if self.is_peft:
+            return False, "This is a PEFT model. Use 'Merged Model' export type instead."
+
+        try:
+            # Save locally if requested
+            if save_directory:
+                logger.info(f"Saving base model locally to: {save_directory}")
+                os.makedirs(save_directory, exist_ok=True)
+
+                self.current_model.save_pretrained(save_directory)
+                self.current_tokenizer.save_pretrained(save_directory)
+                logger.info(f"Model saved successfully to {save_directory}")
+
+            # Push to hub if requested
+            if push_to_hub:
+                if not repo_id or not hf_token:
+                    return False, "Repository ID and Hugging Face token required for Hub upload"
+
+                logger.info(f"Pushing base model to Hub: {repo_id}")
+
+                # Get base model name
+                base_model = self.current_model.config._name_or_path
+
+                # Create repo
+                hf_api = HfApi(token=hf_token)
+                repo_id = PushToHubMixin._create_repo(
+                    PushToHubMixin,
+                    repo_id=repo_id,
+                    private=private,
+                    token=hf_token,
+                )
+                username = repo_id.split("/")[0]
+
+                # Create and push model card
+                content = MODEL_CARD.format(
+                    username=username,
+                    base_model=base_model,
+                    model_type=self.current_model.config.model_type,
+                    method="",
+                    extra="unsloth",
+                )
+                card = ModelCard(content)
+                card.push_to_hub(repo_id, token=hf_token, commit_message="Unsloth Model Card")
+
+                # Upload model files
+                if save_directory:
+                    hf_api.upload_folder(
+                        folder_path=save_directory,
+                        repo_id=repo_id,
+                        repo_type="model"
+                    )
+                    logger.info(f"Model pushed successfully to {repo_id}")
+                else:
+                    return False, "Local save directory required for Hub upload"
+
+            return True, "Model exported successfully"
+
+        except Exception as e:
+            logger.error(f"Error exporting base model: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return False, f"Export failed: {str(e)}"
+
+
+    def export_gguf(self,
+                    save_directory: str,
+                    quantization_method: str = "Q4_K_M",
+                    push_to_hub: bool = False,
+                    repo_id: Optional[str] = None,
+                    hf_token: Optional[str] = None) -> Tuple[bool, str]:
+        """
+        Export model in GGUF format.
+
+        Args:
+            save_directory: Local directory to save model
+            quantization_method: GGUF quantization method (e.g., "Q4_K_M")
+            push_to_hub: Whether to push to Hugging Face Hub
+            repo_id: Hub repository ID
+            hf_token: Hugging Face token
+
+        Returns:
+            Tuple of (success: bool, message: str)
+        """
+        if not self.current_model or not self.current_tokenizer:
+            return False, "No model loaded. Please select a checkpoint first."
+
+        try:
+            # Convert quantization method to lowercase for unsloth
+            quant_method = quantization_method.lower()
+
+            # Save locally if requested
+            if save_directory:
+                logger.info(f"Saving GGUF model locally to: {save_directory}")
+
+                # Create the directory if it doesn't exist
+                os.makedirs(save_directory, exist_ok=True)
+
+                # Get the base filename for the GGUF file
+                import shutil
+                original_dir = os.getcwd()
+
+                try:
+                    # Change to target directory
+                    os.chdir(save_directory)
+                    logger.info(f"Changed directory to: {save_directory}")
+
+                    # Now save (will save in current directory)
+                    self.current_model.save_pretrained_gguf(
+                        "model",  # Base filename
+                        self.current_tokenizer,
+                        quantization_method=quant_method
+                    )
+
+                    logger.info(f"GGUF model saved successfully in {save_directory}")
+
+                    # Check if llama.cpp directory was created here
+                    llama_cpp_in_target = os.path.join(save_directory, "llama.cpp")
+                    llama_cpp_in_original = os.path.join(original_dir, "llama.cpp")
+
+                    if os.path.exists(llama_cpp_in_target):
+                        logger.info(f"Found llama.cpp directory in {save_directory}")
+
+                        # Remove llama.cpp from original directory if it exists
+                        if os.path.exists(llama_cpp_in_original):
+                            logger.info(f"Removing existing llama.cpp in {original_dir}")
+                            shutil.rmtree(llama_cpp_in_original)
+
+                        # Move llama.cpp back to original directory
+                        logger.info(f"Moving llama.cpp to {original_dir}")
+                        shutil.move(llama_cpp_in_target, llama_cpp_in_original)
+                        logger.info(f"Successfully moved llama.cpp back to original directory")
+
+                finally:
+                    # Always change back to original directory
+                    os.chdir(original_dir)
+                    logger.info(f"Changed back to original directory: {original_dir}")
+
+            # Push to hub if requested
+            if push_to_hub:
+                if not repo_id or not hf_token:
+                    return False, "Repository ID and Hugging Face token required for Hub upload"
+
+                logger.info(f"Pushing GGUF model to Hub: {repo_id}")
+
+                self.current_model.push_to_hub_gguf(
+                    repo_id,
+                    self.current_tokenizer,
+                    quantization_method=quant_method,
+                    token=hf_token
+                )
+                logger.info(f"GGUF model pushed successfully to {repo_id}")
+
+            return True, f"GGUF model exported successfully ({quantization_method})"
+
+        except Exception as e:
+            logger.error(f"Error exporting GGUF model: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return False, f"GGUF export failed: {str(e)}"
+
+    def export_lora_adapter(self,
+                           save_directory: str,
+                           push_to_hub: bool = False,
+                           repo_id: Optional[str] = None,
+                           hf_token: Optional[str] = None,
+                           private: bool = False) -> Tuple[bool, str]:
+        """
+        Export LoRA adapter only (not merged).
+
+        Returns:
+            Tuple of (success: bool, message: str)
+        """
+        if not self.current_model or not self.current_tokenizer:
+            return False, "No model loaded. Please select a checkpoint first."
+
+        if not self.is_peft:
+            return False, "This is not a PEFT model. No adapter to export."
+
+        try:
+            # Save locally if requested
+            if save_directory:
+                logger.info(f"Saving LoRA adapter locally to: {save_directory}")
+                os.makedirs(save_directory, exist_ok=True)
+
+                self.current_model.save_pretrained(save_directory)
+                self.current_tokenizer.save_pretrained(save_directory)
+                logger.info(f"Adapter saved successfully to {save_directory}")
+
+            # Push to hub if requested
+            if push_to_hub:
+                if not repo_id or not hf_token:
+                    return False, "Repository ID and Hugging Face token required for Hub upload"
+
+                logger.info(f"Pushing LoRA adapter to Hub: {repo_id}")
+
+                self.current_model.push_to_hub(
+                    repo_id,
+                    token=hf_token,
+                    private=private
+                )
+                self.current_tokenizer.push_to_hub(
+                    repo_id,
+                    token=hf_token,
+                    private=private
+                )
+                logger.info(f"Adapter pushed successfully to {repo_id}")
+
+            return True, "LoRA adapter exported successfully"
+
+        except Exception as e:
+            logger.error(f"Error exporting LoRA adapter: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
+            return False, f"Adapter export failed: {str(e)}"
+
+
+# Global export backend instance
+_export_backend = None
+
+def get_export_backend() -> ExportBackend:
+    """Get or create the global export backend instance"""
+    global _export_backend
+    if _export_backend is None:
+        _export_backend = ExportBackend()
+    return _export_backend

backend/backend/model_config.py

@@ -0,0 +1,704 @@
+"""
+Model and LoRA configuration handling
+"""
+from transformers import AutoConfig
+from dataclasses import dataclass
+from typing import Optional, Dict, Any
+from .path_utils import normalize_path, is_local_path, is_model_cached
+from .utils import without_hf_auth
+import logging
+from pathlib import Path
+from typing import List, Tuple
+import json
+import yaml
+
+
+logger = logging.getLogger(__name__)
+
+# Model name mapping: maps all equivalent model names to their canonical YAML config file
+# Format: "canonical_model_name.yaml": [list of all equivalent model names]
+# Based on the model mapper provided - canonical filename is based on the first model name in the mapper
+MODEL_NAME_MAPPING = {
+    "unsloth_answerdotai_ModernBERT-large.yaml": [
+        "answerdotai/ModernBERT-large",
+    ],
+    "unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml": [
+        "unsloth/Qwen2.5-Coder-7B-Instruct-bnb-4bit",
+        "unsloth/Qwen2.5-Coder-7B-Instruct",
+        "Qwen/Qwen2.5-Coder-7B-Instruct",
+    ],
+    "unsloth_codegemma-7b-bnb-4bit.yaml": [
+        "unsloth/codegemma-7b-bnb-4bit",
+        "unsloth/codegemma-7b",
+        "google/codegemma-7b",
+    ],
+    "unsloth_ERNIE-4.5-21B-A3B-PT.yaml": [
+        "unsloth/ERNIE-4.5-21B-A3B-PT",
+    ],
+    "unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml": [
+        "unsloth/ERNIE-4.5-VL-28B-A3B-PT",
+    ],
+   
+    "tiiuae_Falcon-H1-0.5B-Instruct.yaml": [
+        "tiiuae/Falcon-H1-0.5B-Instruct",
+        "unsloth/Falcon-H1-0.5B-Instruct",
+    ],
+    "unsloth_functiongemma-270m-it.yaml": [
+        "unsloth/functiongemma-270m-it-unsloth-bnb-4bit",
+        "google/functiongemma-270m-it",
+        "unsloth/functiongemma-270m-it-unsloth-bnb-4bit",
+    ],
+    "unsloth_gemma-2-2b.yaml": [
+        "unsloth/gemma-2-2b-bnb-4bit",
+        "google/gemma-2-2b",
+    ],
+    "unsloth_gemma-2-27b-bnb-4bit.yaml": [
+        "unsloth/gemma-2-9b-bnb-4bit",
+        "unsloth/gemma-2-9b",
+        "google/gemma-2-9b",
+        "unsloth/gemma-2-27b",
+        "google/gemma-2-27b",
+    ],
+    "unsloth_gemma-3-4b-pt.yaml": [
+        "unsloth/gemma-3-4b-pt-unsloth-bnb-4bit",
+        "google/gemma-3-4b-pt",
+        "unsloth/gemma-3-4b-pt-bnb-4bit",
+    ],
+    "unsloth_gemma-3-4b-it.yaml": [
+        "unsloth/gemma-3-4b-it-unsloth-bnb-4bit",
+        "google/gemma-3-4b-it",
+        "unsloth/gemma-3-4b-it-bnb-4bit",
+    ],
+    "unsloth_gemma-3-27b-it.yaml": [
+        "unsloth/gemma-3-27b-it-unsloth-bnb-4bit",
+        "google/gemma-3-27b-it",
+        "unsloth/gemma-3-27b-it-bnb-4bit",
+    ],
+    "unsloth_gemma-3-270m-it.yaml": [
+        "unsloth/gemma-3-270m-it-unsloth-bnb-4bit",
+        "google/gemma-3-270m-it",
+        "unsloth/gemma-3-270m-it-bnb-4bit",
+    ],
+    "unsloth_gemma-3n-E4B-it.yaml": [
+        "unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
+        "google/gemma-3n-E4B-it",
+        "unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
+    ],
+    "unsloth_gemma-3n-E4B.yaml": [
+        "unsloth/gemma-3n-E4B-unsloth-bnb-4bit",
+        "google/gemma-3n-E4B",
+    ],
+    "unsloth_gpt-oss-20b.yaml": [
+        "openai/gpt-oss-20b",
+        "unsloth/gpt-oss-20b-unsloth-bnb-4bit",
+        "unsloth/gpt-oss-20b-BF16",
+    ],
+    
+    "unsloth_gpt-oss-120b.yaml": [
+        "openai/gpt-oss-120b",
+        "unsloth/gpt-oss-120b-unsloth-bnb-4bit",
+    ],
+    "unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml": [
+        "unsloth/granite-4.0-350m",
+        "ibm-granite/granite-4.0-350m",
+        "unsloth/granite-4.0-350m-bnb-4bit",
+    ],
+    "unsloth_granite-4.0-h-micro.yaml": [
+        "ibm-granite/granite-4.0-h-micro",
+        "unsloth/granite-4.0-h-micro-bnb-4bit",
+        "unsloth/granite-4.0-h-micro-unsloth-bnb-4bit",
+    ],
+    "unsloth_LFM2-1.2B.yaml": [
+        "unsloth/LFM2-1.2B",
+    ],
+    "unsloth_llama-3-8b-bnb-4bit.yaml": [
+        "unsloth/llama-3-8b",
+        "meta-llama/Meta-Llama-3-8B",
+    ],
+    "unsloth_llama-3-8b-Instruct-bnb-4bit.yaml": [
+        "unsloth/llama-3-8b-Instruct",
+        "meta-llama/Meta-Llama-3-8B-Instruct",
+    ],
+    "unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml": [
+        "unsloth/Meta-Llama-3.1-8B-bnb-4bit",
+        "unsloth/Meta-Llama-3.1-8B-unsloth-bnb-4bit",
+        "meta-llama/Meta-Llama-3.1-8B",
+        "unsloth/Meta-Llama-3.1-70B-bnb-4bit",
+        "unsloth/Meta-Llama-3.1-8B",
+        "unsloth/Meta-Llama-3.1-70B",
+        "meta-llama/Meta-Llama-3.1-70B",
+        "unsloth/Meta-Llama-3.1-405B-bnb-4bit",
+        "meta-llama/Meta-Llama-3.1-405B",
+    ],
+    
+    "unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml": [
+        "unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
+        "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
+        "meta-llama/Meta-Llama-3.1-8B-Instruct",
+        "unsloth/Meta-Llama-3.1-8B-Instruct",
+        "RedHatAI/Llama-3.1-8B-Instruct-FP8",
+        "unsloth/Llama-3.1-8B-Instruct-FP8-Block",
+        "unsloth/Llama-3.1-8B-Instruct-FP8-Dynamic",
+    ],
+    "unsloth_Llama-3.2-3B-Instruct.yaml": [
+        "unsloth/Llama-3.2-3B-Instruct-unsloth-bnb-4bit",
+        "meta-llama/Llama-3.2-3B-Instruct",
+        "unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
+        "RedHatAI/Llama-3.2-3B-Instruct-FP8",
+        "unsloth/Llama-3.2-3B-Instruct-FP8-Block",
+        "unsloth/Llama-3.2-3B-Instruct-FP8-Dynamic",
+    ],
+    "unsloth_Llama-3.2-1B-Instruct.yaml": [
+        "unsloth/Llama-3.2-1B-Instruct-unsloth-bnb-4bit",
+        "meta-llama/Llama-3.2-1B-Instruct",
+        "unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
+        "RedHatAI/Llama-3.2-1B-Instruct-FP8",
+        "unsloth/Llama-3.2-1B-Instruct-FP8-Block",
+        "unsloth/Llama-3.2-1B-Instruct-FP8-Dynamic",
+    ],
+    "unsloth_Llama-3.2-11B-Vision-Instruct.yaml": [
+        "unsloth/Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit",
+        "meta-llama/Llama-3.2-11B-Vision-Instruct",
+        "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
+    ],
+    "unsloth_Llama-3.3-70B-Instruct.yaml": [
+        "unsloth/Llama-3.3-70B-Instruct-unsloth-bnb-4bit",
+        "meta-llama/Llama-3.3-70B-Instruct",
+        "unsloth/Llama-3.3-70B-Instruct-bnb-4bit",
+        "RedHatAI/Llama-3.3-70B-Instruct-FP8",
+        "unsloth/Llama-3.3-70B-Instruct-FP8-Block",
+        "unsloth/Llama-3.3-70B-Instruct-FP8-Dynamic",
+    ],
+    "unsloth_Llasa-3B.yaml": [
+        "HKUSTAudio/Llasa-1B",
+        "unsloth/Llasa-3B",
+    ],
+    "unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml": [
+        "unsloth/Magistral-Small-2509",
+        "mistralai/Magistral-Small-2509",
+        "unsloth/Magistral-Small-2509-bnb-4bit",
+    ],
+    "unsloth_Ministral-3-3B-Instruct-2512.yaml": [
+        "unsloth/Ministral-3-3B-Instruct-2512",
+    ],
+    "unsloth_mistral-7b-v0.3-bnb-4bit.yaml": [
+        "unsloth/mistral-7b-v0.3-bnb-4bit"
+        "unsloth/mistral-7b-v0.3",
+        "mistralai/Mistral-7B-v0.3",
+    ],
+    "unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml": [
+        "unsloth/Mistral-Nemo-Base-2407-bnb-4bit",
+        "unsloth/Mistral-Nemo-Base-2407",
+        "mistralai/Mistral-Nemo-Base-2407",
+        "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", 
+        "unsloth/Mistral-Nemo-Instruct-2407", 
+        "mistralai/Mistral-Nemo-Instruct-2407",
+
+    ],
+    "unsloth_Mistral-Small-Instruct-2409.yaml": [
+        "unsloth/Mistral-Small-Instruct-2409-bnb-4bit",
+        "mistralai/Mistral-Small-Instruct-2409",
+    ],
+    "unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml": [
+        "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
+        "unsloth/mistral-7b-instruct-v0.3",
+        "mistralai/Mistral-7B-Instruct-v0.3",
+    ],
+    "unsloth_Qwen2.5-1.5B-Instruct.yaml": [
+        "unsloth/Qwen2.5-1.5B-Instruct-unsloth-bnb-4bit",
+        "Qwen/Qwen2.5-1.5B-Instruct",
+        "unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit",
+    ],
+    "unsloth_Nemotron-3-Nano-30B-A3B.yaml": [
+        "unsloth/Nemotron-3-Nano-30B-A3B",
+    ],
+    "unsloth_orpheus-3b-0.1-ft.yaml": [
+        "unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit",
+        "canopylabs/orpheus-3b-0.1-ft",
+        "unsloth/orpheus-3b-0.1-ft-bnb-4bit",
+    ],
+    "OuteAI_Llama-OuteTTS-1.0-1B.yaml": [
+        "OuteAI/Llama-OuteTTS-1.0-1B",
+    ],
+    "unsloth_PaddleOCR-VL.yaml": [
+        "unsloth/PaddleOCR-VL",
+    ],
+    "unsloth_Phi-3-medium-4k-instruct.yaml": [
+        "unsloth/Phi-3-medium-4k-instruct-bnb-4bit",
+        "microsoft/Phi-3-medium-4k-instruct",
+    ],
+    "unsloth_Phi-3.5-mini-instruct.yaml": [
+        "unsloth/Phi-3.5-mini-instruct-bnb-4bit",
+        "microsoft/Phi-3.5-mini-instruct",
+    ],
+    "unsloth_Phi-4.yaml": [
+        "unsloth/phi-4-unsloth-bnb-4bit",
+        "microsoft/phi-4",
+        "unsloth/phi-4-bnb-4bit",
+    ],
+    "unsloth_Pixtral-12B-2409.yaml": [
+        "unsloth/Pixtral-12B-2409-unsloth-bnb-4bit",
+        "mistralai/Pixtral-12B-2409",
+        "unsloth/Pixtral-12B-2409-bnb-4bit",
+    ],
+    "unsloth_Qwen2-7B.yaml": [
+        "unsloth/Qwen2-7B-bnb-4bit",
+        "Qwen/Qwen2-7B",
+    ],
+    "unsloth_Qwen2-VL-7B-Instruct.yaml": [
+        "unsloth/Qwen2-VL-7B-Instruct-unsloth-bnb-4bit",
+        "Qwen/Qwen2-VL-7B-Instruct",
+        "unsloth/Qwen2-VL-7B-Instruct-bnb-4bit",
+    ],
+    "unsloth_Qwen2.5-7B.yaml": [
+        "unsloth/Qwen2.5-7B-unsloth-bnb-4bit",
+        "Qwen/Qwen2.5-7B",
+        "unsloth/Qwen2.5-7B-bnb-4bit",
+    ],
+    "unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml": [
+        "unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit",
+        "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+    ],
+    "unsloth_Qwen2.5-Coder-14B-Instruct.yaml": [
+        "unsloth/Qwen2.5-Coder-14B-Instruct-bnb-4bit",
+        "Qwen/Qwen2.5-Coder-14B-Instruct",
+    ],
+    "unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml": [
+        "unsloth/Qwen2.5-VL-7B-Instruct",
+        "Qwen/Qwen2.5-VL-7B-Instruct",
+        "unsloth/Qwen2.5-VL-7B-Instruct-unsloth-bnb-4bit",
+    ],
+    "unsloth_Qwen3-0.6B.yaml": [
+        "unsloth/Qwen3-0.6B-unsloth-bnb-4bit",
+        "Qwen/Qwen3-0.6B",
+        "unsloth/Qwen3-0.6B-bnb-4bit",
+        "Qwen/Qwen3-0.6B-FP8",
+        "unsloth/Qwen3-0.6B-FP8",
+    ],
+    "unsloth_Qwen3-4B-Instruct-2507.yaml": [
+        "unsloth/Qwen3-4B-Instruct-2507-unsloth-bnb-4bit",
+        "Qwen/Qwen3-4B-Instruct-2507",
+        "unsloth/Qwen3-4B-Instruct-2507-bnb-4bit",
+        "Qwen/Qwen3-4B-Instruct-2507-FP8",
+        "unsloth/Qwen3-4B-Instruct-2507-FP8",
+    ],
+    "unsloth_Qwen3-4B-Thinking-2507.yaml": [
+        "unsloth/Qwen3-4B-Thinking-2507-unsloth-bnb-4bit",
+        "Qwen/Qwen3-4B-Thinking-2507",
+        "unsloth/Qwen3-4B-Thinking-2507-bnb-4bit",
+        "Qwen/Qwen3-4B-Thinking-2507-FP8",
+        "unsloth/Qwen3-4B-Thinking-2507-FP8",
+    ],
+    "unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml": [
+        "unsloth/Qwen3-14B-Base",
+        "Qwen/Qwen3-14B-Base",
+        "unsloth/Qwen3-14B-Base-bnb-4bit",
+    ],
+    "unsloth_Qwen3-14B.yaml": [
+        "unsloth/Qwen3-14B-unsloth-bnb-4bit",
+        "Qwen/Qwen3-14B",
+        "unsloth/Qwen3-14B-bnb-4bit",
+        "Qwen/Qwen3-14B-FP8",
+        "unsloth/Qwen3-14B-FP8",
+    ],
+    "unsloth_Qwen3-32B.yaml": [
+        "unsloth/Qwen3-32B-unsloth-bnb-4bit",
+        "Qwen/Qwen3-32B",
+        "unsloth/Qwen3-32B-bnb-4bit",
+        "Qwen/Qwen3-32B-FP8",
+        "unsloth/Qwen3-32B-FP8",
+    ],
+    "unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml": [
+        "Qwen/Qwen3-VL-8B-Instruct-FP8",
+        "unsloth/Qwen3-VL-8B-Instruct-FP8",
+        "unsloth/Qwen3-VL-8B-Instruct",
+        "Qwen/Qwen3-VL-8B-Instruct",
+        "unsloth/Qwen3-VL-8B-Instruct-bnb-4bit",
+    ],
+    "sesame_csm-1b.yaml": [
+        "sesame/csm-1b",
+    ],
+    "Spark-TTS-0.5B_LLM.yaml": [
+        "Spark-TTS-0.5B/LLM",
+    ],
+    "unsloth_tinyllama-bnb-4bit.yaml": [
+        "unsloth/tinyllama",
+        "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
+    ],
+    "unsloth_whisper-large-v3.yaml": [
+        "unsloth/whisper-large-v3",
+        "openai/whisper-large-v3",
+    ],
+}
+
+# Reverse mapping for quick lookup: model_name -> canonical_filename
+_REVERSE_MODEL_MAPPING = {}
+for canonical_file, model_names in MODEL_NAME_MAPPING.items():
+    for model_name in model_names:
+        _REVERSE_MODEL_MAPPING[model_name] = canonical_file
+
+def load_model_config(model_name: str, use_auth: bool = False, token: Optional[str] = None):
+    """
+    Load model config with optional authentication control.
+    """
+
+    if token:
+        # Explicit token provided - use it
+        return AutoConfig.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            token=token
+        )
+
+    if not use_auth:
+        # Load without any authentication (for public model checks)
+        with without_hf_auth():
+            return AutoConfig.from_pretrained(
+                model_name,
+                trust_remote_code=True,
+                token=None
+            )
+
+    # Use default authentication (cached tokens)
+    return AutoConfig.from_pretrained(
+        model_name,
+        trust_remote_code=True
+    )
+pass
+
+
+def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
+    """
+    Detect vision models by checking architecture in config.
+    Works for fine-tuned models since they inherit the base architecture.
+
+    Args:
+        model_name: Model identifier (HF repo or local path)
+        hf_token: Optional HF token for accessing gated/private models
+    """
+    try:
+        config = load_model_config(model_name, token=hf_token)
+
+        # Check vision arch
+        if hasattr(config, 'architectures'):
+            is_vlm = any(
+                x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
+                for x in config.architectures
+            )
+            if is_vlm:
+                logger.info(f"Model {model_name} detected as vision model: architecture {config.architectures}")
+                return True
+
+        # Quick check for vision config as backup
+        if hasattr(config, 'vision_config'):
+            logger.info(f"Model {model_name} detected as vision model: has vision_config")
+            return True
+
+        return False
+
+    except Exception as e:
+        logger.warning(f"Could not determine if {model_name} is vision model: {e}")
+        return False
+pass
+
+
+def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]:
+    """
+    Scan outputs folder for trained LoRA adapters.
+
+    Returns:
+        List of tuples: [(display_name, adapter_path), ...]
+
+    Example:
+        [
+            ("unsloth_Meta-Llama-3.1_...", "./outputs/unsloth_Meta-Llama-3.1_.../"),
+            ("my_finetuned_model", "./outputs/my_finetuned_model/"),
+        ]
+    """
+    trained_loras = []
+    outputs_path = Path(outputs_dir)
+
+    if not outputs_path.exists():
+        logger.warning(f"Outputs directory not found: {outputs_dir}")
+        return trained_loras
+
+    try:
+        for item in outputs_path.iterdir():
+            if item.is_dir():
+                # Check if this directory contains a LoRA adapter
+                adapter_config = item / "adapter_config.json"
+                adapter_model = item / "adapter_model.safetensors"
+
+                if adapter_config.exists() or adapter_model.exists():
+                    display_name = item.name
+                    adapter_path = str(item)
+                    trained_loras.append((display_name, adapter_path))
+                    logger.debug(f"Found trained LoRA: {display_name}")
+
+        # Sort by modification time (newest first)
+        trained_loras.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
+
+        logger.info(f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}")
+        return trained_loras
+
+    except Exception as e:
+        logger.error(f"Error scanning outputs 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.
+
+    Args:
+        lora_path: Path to the LoRA adapter directory
+
+    Returns:
+        Base model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit")
+        or None if not found
+
+    Example:
+        >>> get_base_model_from_lora("./outputs/unsloth_Meta-Llama-3.1_.../")
+        "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
+    """
+    try:
+        lora_path_obj = Path(lora_path)
+
+        # Try adapter_config.json first
+        adapter_config_path = lora_path_obj / "adapter_config.json"
+        if adapter_config_path.exists():
+            with open(adapter_config_path, 'r') as f:
+                config = json.load(f)
+                base_model = config.get("base_model_name_or_path")
+                if base_model:
+                    logger.info(f"Detected base model from adapter_config.json: {base_model}")
+                    return base_model
+
+        # Fallback: try training_args.bin (requires torch)
+        training_args_path = lora_path_obj / "training_args.bin"
+        if training_args_path.exists():
+            try:
+                import torch
+                training_args = torch.load(training_args_path)
+                if hasattr(training_args, 'model_name_or_path'):
+                    base_model = training_args.model_name_or_path
+                    logger.info(f"Detected base model from training_args.bin: {base_model}")
+                    return base_model
+            except Exception as e:
+                logger.warning(f"Could not load training_args.bin: {e}")
+
+        # Last resort: parse from directory name
+        # Format: unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit_timestamp
+        dir_name = lora_path_obj.name
+        if dir_name.startswith("unsloth_"):
+            # Remove timestamp suffix (usually _1234567890)
+            parts = dir_name.split("_")
+            # Reconstruct model name
+            if len(parts) >= 2:
+                model_parts = parts[1:-1]  # Skip "unsloth" and timestamp
+                base_model = "unsloth/" + "_".join(model_parts)
+                logger.info(f"Detected base model from directory name: {base_model}")
+                return base_model
+
+        logger.warning(f"Could not detect base model for LoRA: {lora_path}")
+        return None
+
+    except Exception as e:
+        logger.error(f"Error reading base model from LoRA config: {e}")
+        return None
+pass
+
+# Status indicators that appear in UI dropdowns
+UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "]
+
+def load_model_defaults(model_name: str) -> Dict[str, Any]:
+    """
+    Load default training parameters for a model from YAML file.
+    
+    Args:
+        model_name: Model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-bnb-4bit")
+    
+    Returns:
+        Dictionary with default parameters from YAML file, or empty dict if not found
+    
+    The function looks for a YAML file in configs/model_defaults/ (including subfolders) 
+    based on the model name or its aliases from MODEL_NAME_MAPPING.
+    If no specific file exists, it falls back to default.yaml.
+    """
+    try:
+        # Get the script directory to locate configs
+        script_dir = Path(__file__).parent.parent
+        defaults_dir = script_dir / "configs" / "model_defaults"
+        
+        # First, check if model is in the mapping
+        if model_name in _REVERSE_MODEL_MAPPING:
+            canonical_file = _REVERSE_MODEL_MAPPING[model_name]
+            # Search in subfolders and root
+            for config_path in defaults_dir.rglob(canonical_file):
+                if config_path.is_file():
+                    with open(config_path, 'r', encoding='utf-8') as f:
+                        config = yaml.safe_load(f) or {}
+                        logger.info(f"Loaded model defaults from {config_path} (via mapping)")
+                        return config
+        
+        # Try exact model name match (for backward compatibility)
+        model_filename = model_name.replace("/", "_") + ".yaml"
+        # Search in subfolders and root
+        for config_path in defaults_dir.rglob(model_filename):
+            if config_path.is_file():
+                with open(config_path, 'r', encoding='utf-8') as f:
+                    config = yaml.safe_load(f) or {}
+                    logger.info(f"Loaded model defaults from {config_path}")
+                    return config
+        
+        # Fall back to default.yaml
+        default_config_path = defaults_dir / "default.yaml"
+        if default_config_path.exists():
+            with open(default_config_path, 'r', encoding='utf-8') as f:
+                config = yaml.safe_load(f) or {}
+                logger.info(f"Loaded default model defaults from {default_config_path}")
+                return config
+        
+        logger.warning(f"No default config found for model {model_name}")
+        return {}
+        
+    except Exception as e:
+        logger.error(f"Error loading model defaults for {model_name}: {e}")
+        return {}
+
+@dataclass
+class ModelConfig:
+    """Configuration for a model to load"""
+    identifier: str      # Clean model identifier (org/name or path)
+    display_name: str    # Original UI display name
+    path: str           # Normalized filesystem path
+    is_local: bool      # Is this a local file vs HF model?
+    is_cached: bool     # Is this already in HF cache?
+    is_vision: bool     # Is this a vision model?
+    is_lora: bool       # Is this a lora adapter?
+    base_model: Optional[str] = None  # Base model (for LoRAs)
+
+    @classmethod
+    def from_lora_path(cls, lora_path: str, hf_token: Optional[str] = None) -> Optional['ModelConfig']:
+        """
+        Create ModelConfig from a local LoRA adapter path.
+
+        Automatically detects the base model from adapter config.
+
+        Args:
+            lora_path: Path to LoRA adapter (e.g., "./outputs/unsloth_Meta-Llama-3.1_.../")
+            hf_token: HF token for vision detection
+
+        Returns:
+            ModelConfig for the LoRA adapter
+        """
+        try:
+            lora_path_obj = Path(lora_path)
+
+            if not lora_path_obj.exists():
+                logger.error(f"LoRA path does not exist: {lora_path}")
+                return None
+
+            # Get base model
+            base_model = get_base_model_from_lora(lora_path)
+            if not base_model:
+                logger.error(f"Could not determine base model for LoRA: {lora_path}")
+                return None
+
+            # Check if base model is vision
+            is_vision = is_vision_model(base_model, hf_token=hf_token)
+
+            display_name = lora_path_obj.name
+            identifier = lora_path  # Use path as identifier for local LoRAs
+
+            return cls(
+                identifier=identifier,
+                display_name=display_name,
+                path=lora_path,
+                is_local=True,
+                is_cached=True,  # Local LoRAs are always "cached"
+                is_vision=is_vision,
+                is_lora=True,
+                base_model=base_model,
+            )
+
+        except Exception as e:
+            logger.error(f"Error creating ModelConfig from LoRA path: {e}")
+            return None
+
+
+    @classmethod
+    def from_ui_selection(cls,
+                          dropdown_value: Optional[str],
+                          search_value: Optional[str],
+                          local_models: list = None,
+                          hf_token: Optional[str] = None,
+                          is_lora: bool = False) -> Optional['ModelConfig']:
+        """
+        Create a universal ModelConfig from UI dropdown/search selections.
+        Handles base models and LoRA adapters.
+        """
+        selected = None
+        if search_value and search_value.strip():
+            selected = search_value.strip()
+        elif dropdown_value:
+            selected = dropdown_value
+
+        if not selected:
+            return None
+
+        display_name = selected
+
+        #  Use the correct 'local_models' parameter to resolve display names
+        if " (Active)" in selected or " (Ready)" in selected:
+            clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
+            if local_models:
+                for local_display, local_path in local_models:
+                    if local_display == clean_display_name:
+                        selected = local_path
+                        break
+
+        # Clean all UI status indicators to get the final identifier
+        identifier = selected
+        for status in UI_STATUS_INDICATORS:
+            identifier = identifier.replace(status, "")
+        identifier = identifier.strip()
+
+        is_local = is_local_path(identifier)
+        path = normalize_path(identifier) if is_local else identifier
+
+        # Add unsloth/ prefix for shorthand HF models
+        if not is_local and "/" not in identifier:
+            identifier = f"unsloth/{identifier}"
+            path = identifier
+
+        # --- Logic for Base Model and Vision Detection ---
+        base_model = None
+        is_vision = False
+
+        if is_lora:
+            # For a LoRA, we MUST find its base model.
+            base_model = get_base_model_from_lora(path)
+            if not base_model:
+                logger.warning(f"Could not determine base model for LoRA '{path}'. Cannot create config.")
+                return None # Cannot proceed without a base model
+
+            # A LoRA's vision capability is determined by its base model.
+            is_vision = is_vision_model(base_model, hf_token=hf_token)
+        else:
+            # For a base model, just check its own vision status.
+            is_vision = is_vision_model(identifier, hf_token=hf_token)
+
+        from .path_utils import is_model_cached
+        is_cached = is_model_cached(identifier) if not is_local else True
+
+        return cls(
+            identifier=identifier,
+            display_name=display_name,
+            path=path,
+            is_local=is_local,
+            is_cached=is_cached,
+            is_vision=is_vision,
+            is_lora=is_lora,
+            base_model=base_model, # This will be None for base models, and populated for LoRAs
+        )
+    pass

backend/backend/path_utils.py

@@ -0,0 +1,78 @@
+"""
+Path utilities for model and dataset handling
+"""
+import os
+from pathlib import Path
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def normalize_path(path: str) -> str:
+    """
+    Convert Windows paths to WSL format if needed.
+
+    Examples:
+        C:\\Users\\... -> /mnt/c/Users/...
+        /home/user/... -> /home/user/... (unchanged)
+    """
+    if not path:
+        return path
+
+    # Handle Windows drive letters (C:\\ or c:\\)
+    if len(path) >= 3 and path[1] == ':' and path[2] in ('\\', '/'):
+        drive = path[0].lower()
+        rest = path[3:].replace('\\', '/')
+        return f'/mnt/{drive}/{rest}'
+
+    # Already Unix-style or relative
+    return path.replace('\\', '/')
+pass
+
+def is_local_path(path: str) -> bool:
+    """
+    Check if path is a local filesystem path vs HuggingFace model identifier.
+
+    Examples:
+        True: /home/user/model, C:\\models, ./model, ~/model
+        False: unsloth/llama-3.1-8b, microsoft/phi-2
+    """
+    if not path:
+        return False
+
+    # Obvious HF patterns
+    if path.count('/') == 1 and not path.startswith(('/', '.', '~')):
+        return False  # Looks like org/model format
+
+    # Filesystem indicators
+    return (
+        path.startswith(('/', '.', '~')) or  # Unix absolute/relative
+        ':' in path or                        # Windows drive or URL
+        '\\' in path or                       # Windows separator
+        os.path.isabs(path)                  # System-absolute
+    )
+pass
+
+def get_cache_path(model_name: str) -> Optional[Path]:
+    """Get HuggingFace cache path for a model if it exists."""
+    cache_dir = Path.home() / '.cache' / 'huggingface' / 'hub'
+    model_cache_name = model_name.replace("/", "--")
+    model_cache_path = cache_dir / f'models--{model_cache_name}'
+
+    return model_cache_path if model_cache_path.exists() else None
+pass
+
+def is_model_cached(model_name: str) -> bool:
+    """Check if model is downloaded in HuggingFace cache."""
+    cache_path = get_cache_path(model_name)
+    if not cache_path:
+        return False
+
+    # Check for actual model files
+    for suffix in ['.safetensors', '.bin', '.json']:
+        if list(cache_path.rglob(f'*{suffix}')):
+            return True
+
+    return False
+pass

backend/backend/trainer.py

@@ -0,0 +1,864 @@
+"""
+Unsloth Training Backend
+Integrates Unsloth training capabilities with the Gradio UI
+"""
+import torch
+torch._dynamo.config.recompile_limit = 64
+from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
+from unsloth.chat_templates import get_chat_template
+
+import os
+import json
+import threading
+import math
+import logging
+from typing import Optional, Callable
+from dataclasses import dataclass
+import pandas as pd
+from datasets import Dataset, load_dataset
+
+# Add the parent directory to sys.path to import unsloth modules
+#sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+from .model_config import is_vision_model
+from .dataset_utils import format_and_template_dataset
+from .dataset_utils import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
+from trl import SFTTrainer, SFTConfig
+
+# Import Unsloth trainers
+#from unsloth_compiled_cache.UnslothSFTTrainer import _UnslothSFTTrainer as SFTTrainer
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TrainingProgress:
+    """Training progress tracking"""
+    epoch: int = 0
+    step: int = 0
+    total_steps: int = 0
+    loss: float = 0.0
+    learning_rate: float = 0.0
+    is_training: bool = False
+    is_completed: bool = False
+    error: Optional[str] = None
+    status_message: str = "Ready to train"  # Current stage message
+
+class UnslothTrainer:
+    """
+    Unsloth Training Backend for Gradio UI Integration
+    """
+
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.trainer = None
+        self.training_thread = None
+        self.training_progress = TrainingProgress()
+        self.progress_callbacks = []
+        self.is_training = False
+        self.should_stop = False
+
+        # Model state tracking
+        self.is_vlm = False
+        self.model_name = None
+
+        # Thread safety
+        self._lock = threading.Lock()
+
+        # Store training context for later transfer
+        self.training_context = {
+            'base_model_name': None,
+            'output_dir': None,
+            'is_lora': True,  # Default to LoRA
+        }
+
+    def add_progress_callback(self, callback: Callable[[TrainingProgress], None]):
+        """Add callback for training progress updates"""
+        self.progress_callbacks.append(callback)
+
+    def _update_progress(self, **kwargs):
+        """Update training progress and notify callbacks"""
+        with self._lock:
+            for key, value in kwargs.items():
+                if hasattr(self.training_progress, key):
+                    setattr(self.training_progress, key, value)
+
+            # Notify all callbacks
+            for callback in self.progress_callbacks:
+                try:
+                    callback(self.training_progress)
+                except Exception as e:
+                    logger.error(f"Error in progress callback: {e}")
+
+    def load_model(self,
+                   model_name: str,
+                   max_seq_length: int = 2048,
+                   load_in_4bit: bool = True,
+                   hf_token: Optional[str] = None) -> bool:
+        """Load model for training (supports both text and vision models)"""
+        try:
+            print("\nClearing GPU memory before training...")
+            torch.cuda.empty_cache()
+            import gc
+            gc.collect()
+
+            # Detect if this is a vision model first
+            self.is_vlm = is_vision_model(model_name)
+            self.model_name = model_name
+
+            logger.info(f"Model type detected: {'Vision' if self.is_vlm else 'Text'}")
+
+            # Reset training state for new run
+            self._update_progress(
+                is_training=True,
+                is_completed=False,
+                error=None,
+                step=0,
+                loss=0.0,
+                epoch=0
+            )
+
+            # Update UI immediately with loading message
+            model_display = model_name.split('/')[-1] if '/' in model_name else model_name
+            self._update_progress(
+                status_message=f"Loading {'vision' if self.is_vlm else 'text'} model... {model_display}"
+            )
+
+            print(f"\nLoading {'vision' if self.is_vlm else 'text'} model: {model_name}")
+
+            # Set HF token if provided
+            if hf_token:
+                os.environ["HF_TOKEN"] = hf_token
+
+
+            # Branch based on model type
+            if self.is_vlm:
+                # Load vision model - returns (model, tokenizer)
+                self.model, self.tokenizer = FastVisionModel.from_pretrained(
+                    model_name=model_name,
+                    max_seq_length=max_seq_length,
+                    dtype=None,  # Auto-detect
+                    load_in_4bit=load_in_4bit,
+                    token=hf_token,
+                )
+                logger.info("Loaded vision model")
+            else:
+                # Load text model - returns (model, tokenizer)
+                self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+                    model_name=model_name,
+                    max_seq_length=max_seq_length,
+                    dtype=None,  # Auto-detect
+                    load_in_4bit=load_in_4bit,
+                    token=hf_token,
+                )
+                logger.info("Loaded text model")
+
+            if self.should_stop:
+                return False
+
+            self._update_progress(status_message="Model loaded successfully")
+            print("Model loaded successfully")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error loading model: {e}")
+            self._update_progress(error=str(e), is_training=False)
+            return False
+
+    def prepare_model_for_training(self,
+                                   use_lora: bool = True,
+                                   # Vision-specific LoRA parameters (only used if is_vlm=True)
+                                   finetune_vision_layers: bool = True,
+                                   finetune_language_layers: bool = True,
+                                   finetune_attention_modules: bool = True,
+                                   finetune_mlp_modules: bool = True,
+                                   # Standard LoRA parameters
+                                   target_modules: list = None,
+                                   lora_r: int = 16,
+                                   lora_alpha: int = 16,
+                                   lora_dropout: float = 0.0,
+                                   use_gradient_checkpointing: str = "unsloth",
+                                   use_rslora: bool = False,
+                                   use_loftq: bool = False) -> bool:
+        """
+        Prepare model for training (with optional LoRA).
+        """
+        try:
+            if self.model is None:
+                raise ValueError("Model not loaded. Call load_model() first.")
+
+
+            # Full finetuning mode - skip PEFT entirely
+            if not use_lora:
+                self._update_progress(status_message="Full finetuning mode - no LoRA adapters")
+                print("Full finetuning mode - training all parameters\n")
+                return True
+
+            # LoRA/QLoRA mode - apply PEFT
+            if target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
+                target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                                "gate_proj", "up_proj", "down_proj"]
+
+            # Validate and normalize gradient_checkpointing
+            # Must be one of: True, False, or "unsloth"
+            if isinstance(use_gradient_checkpointing, str):
+                use_gradient_checkpointing = use_gradient_checkpointing.strip().lower()
+                if use_gradient_checkpointing == "" or use_gradient_checkpointing == "unsloth":
+                    use_gradient_checkpointing = "unsloth"
+                elif use_gradient_checkpointing in ("true", "1", "yes"):
+                    use_gradient_checkpointing = True
+                elif use_gradient_checkpointing in ("false", "0", "no"):
+                    use_gradient_checkpointing = False
+                else:
+                    # Invalid value, default to "unsloth"
+                    logger.warning(f"Invalid gradient_checkpointing value: {use_gradient_checkpointing}, defaulting to 'unsloth'")
+                    use_gradient_checkpointing = "unsloth"
+            elif use_gradient_checkpointing not in (True, False, "unsloth"):
+                # Invalid type or value, default to "unsloth"
+                logger.warning(f"Invalid gradient_checkpointing type/value: {use_gradient_checkpointing}, defaulting to 'unsloth'")
+                use_gradient_checkpointing = "unsloth"
+
+            # Verify model is loaded
+            if self.model is None:
+                error_msg = "Model is None - model was not loaded properly"
+                logger.error(error_msg)
+                self._update_progress(error=error_msg)
+                return False
+
+            # Check if model has the expected attributes
+            if not hasattr(self.model, 'config'):
+                error_msg = "Model does not have config attribute - model may not be loaded correctly"
+                logger.error(error_msg)
+                self._update_progress(error=error_msg)
+                return False
+
+            print(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n")
+            print(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n")
+
+            # Branch based on vision vs text
+            if self.is_vlm:
+                # Vision model LoRA
+                print(f"Vision model LoRA configuration:")
+                print(f"  - Finetune vision layers: {finetune_vision_layers}")
+                print(f"  - Finetune language layers: {finetune_language_layers}")
+                print(f"  - Finetune attention modules: {finetune_attention_modules}")
+                print(f"  - Finetune MLP modules: {finetune_mlp_modules}\n")
+
+                self.model = FastVisionModel.get_peft_model(
+                    self.model,
+                    finetune_vision_layers=finetune_vision_layers,
+                    finetune_language_layers=finetune_language_layers,
+                    finetune_attention_modules=finetune_attention_modules,
+                    finetune_mlp_modules=finetune_mlp_modules,
+                    r=lora_r,
+                    target_modules=target_modules,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    bias="none",
+                    use_gradient_checkpointing=use_gradient_checkpointing,
+                    random_state=3407,
+                    use_rslora=use_rslora,
+                    loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
+                )
+            else:
+                # Text model LoRA
+                print(f"Text model LoRA configuration:")
+                print(f"  - Target modules: {target_modules}\n")
+
+                self.model = FastLanguageModel.get_peft_model(
+                    self.model,
+                    r=lora_r,
+                    target_modules=target_modules,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    bias="none",
+                    use_gradient_checkpointing=use_gradient_checkpointing,
+                    random_state=3407,
+                    use_rslora=use_rslora,
+                    loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None,
+                )
+
+            # Check if stopped during LoRA preparation
+            if self.should_stop:
+                print("Stopped during LoRA configuration\n")
+                return False
+
+            self._update_progress(status_message="LoRA adapters configured")
+            print("LoRA adapters configured successfully\n")
+            return True
+
+        except Exception as e:
+            import traceback
+            import sys
+            error_details = f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)"
+            full_traceback = traceback.format_exc()
+            logger.error(f"Error preparing model: {error_details}")
+            logger.error(f"Full traceback:\n{full_traceback}")
+            print(f"\n[ERROR] Error preparing model: {error_details}", file=sys.stderr, flush=True)
+            print(f"[ERROR] Full traceback:\n{full_traceback}", file=sys.stderr, flush=True)
+            self._update_progress(error=error_details)
+            return False
+
+    def load_and_format_dataset(self,
+                     dataset_source: str,
+                     format_type: str = "auto",
+                     local_datasets: list = None) -> Optional[Dataset]:
+        """
+        Load and prepare dataset for training
+        """
+        try:
+            dataset = None
+
+            if local_datasets:
+                # Load local datasets
+                all_data = []
+                for dataset_file in local_datasets:
+                    file_path = os.path.join("datasets", dataset_file)
+                    if file_path.endswith('.json'):
+                        with open(file_path, 'r', encoding='utf-8') as f:
+                            data = json.load(f)
+                            if isinstance(data, list):
+                                all_data.extend(data)
+                            else:
+                                all_data.append(data)
+                    elif file_path.endswith('.csv'):
+                        df = pd.read_csv(file_path)
+                        all_data.extend(df.to_dict('records'))
+
+                if all_data:
+                    dataset = Dataset.from_list(all_data)
+
+                    # Check if stopped during dataset loading
+                    if self.should_stop:
+                        print("Stopped during dataset loading\n")
+                        return None
+
+                    self._update_progress(status_message=f"Loaded {len(all_data)} samples from local files")
+                    print(f"Loaded {len(all_data)} samples from local files\n")
+
+            elif dataset_source:
+                # Load from Hugging Face
+                dataset = load_dataset(dataset_source, split="train")
+
+                # Check if stopped during dataset loading
+                if self.should_stop:
+                    print("Stopped during dataset loading\n")
+                    return None
+
+                self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}")
+                print(f"Loaded dataset from Hugging Face: {dataset_source}\n")
+
+            if dataset is None:
+                raise ValueError("No dataset provided")
+
+            # Check if stopped before applying template
+            if self.should_stop:
+                print("Stopped before applying chat template\n")
+                return None
+
+            # NEW: Use unified format_and_template_dataset
+            print(f"Formatting dataset with format_type='{format_type}'...\n")
+
+            #breakpoint()
+            dataset_info = format_and_template_dataset(
+                dataset,
+                model_name=self.model_name,
+                tokenizer=self.tokenizer,  # Works for both text and vision models
+                is_vlm=self.is_vlm,
+                format_type=format_type,  # "auto", "alpaca", "chatml", "sharegpt"
+                dataset_name=dataset_source,
+            )
+
+            # Check if stopped during formatting
+            if self.should_stop:
+                print("Stopped during dataset formatting\n")
+                return None
+
+            self._update_progress(status_message=f"Dataset formatted and ready for training")
+            print(f"Dataset formatted successfully\n")
+            return dataset_info
+
+        except Exception as e:
+            logger.error(f"Error loading dataset: {e}")
+            self._update_progress(error=str(e))
+            return None
+
+    def start_training(self,
+                       dataset: Dataset,
+                       output_dir: str = "./outputs",
+                       num_epochs: int = 3,
+                       learning_rate: float = 5e-5,
+                       batch_size: int = 2,
+                       gradient_accumulation_steps: int = 4,
+                       warmup_steps: int = None,
+                       warmup_ratio: float = None,
+                       max_steps: int = 0,
+                       save_steps: int = 0,
+                       weight_decay: float = 0.01,
+                       random_seed: int = 3407,
+                       packing: bool = False,
+                       train_on_completions: bool = False,
+                       enable_wandb: bool = False,
+                       wandb_project: str = "unsloth-training",
+                       wandb_token: str = None,
+                       enable_tensorboard: bool = False,
+                       tensorboard_dir: str = "runs",
+                       **kwargs) -> bool:
+        """Start training in a separate thread"""
+
+        if self.is_training:
+            logger.warning("Training already in progress")
+            return False
+
+
+        if self.model is None or self.tokenizer is None:
+            self._update_progress(error="Model not loaded")
+            return False
+
+        # Start training in separate thread
+        self.training_thread = threading.Thread(
+            target=self._train_worker,
+            args=(dataset,),
+            kwargs={
+                'output_dir': output_dir,
+                'num_epochs': num_epochs,
+                'learning_rate': learning_rate,
+                'batch_size': batch_size,
+                'gradient_accumulation_steps': gradient_accumulation_steps,
+                'warmup_steps': warmup_steps,
+                'warmup_ratio': warmup_ratio,
+                'max_steps': max_steps,
+                'save_steps': save_steps,
+                'weight_decay': weight_decay,
+                'random_seed': random_seed,
+                'packing': packing,
+                'train_on_completions': train_on_completions,
+                'enable_wandb': enable_wandb,
+                'wandb_project': wandb_project,
+                'wandb_token': wandb_token,
+                'enable_tensorboard': enable_tensorboard,
+                'tensorboard_dir': tensorboard_dir,
+                **kwargs
+            }
+        )
+
+        self.should_stop = False
+        self.training_thread.start()
+        return True
+
+    def _train_worker(self, dataset: Dataset, **training_args):
+        """Worker function for training (runs in separate thread)"""
+        try:
+            self._update_progress(is_training=True, error=None)
+
+            # Setup logging
+            if training_args.get('enable_wandb', False) and training_args.get('wandb_token'):
+                os.environ["WANDB_API_KEY"] = training_args['wandb_token']
+                import wandb
+                wandb.init(project=training_args.get('wandb_project', 'unsloth-training'))
+
+            # Create output directory
+            output_dir = training_args.get('output_dir', './outputs')
+            os.makedirs(output_dir, exist_ok=True)
+
+            # ========== DATA COLLATOR SELECTION ==========
+            # Detect special model types
+            model_name_lower = self.model_name.lower()
+            is_deepseek_ocr = "deepseek" in model_name_lower and "ocr" in model_name_lower
+
+            print("Configuring data collator...\n")
+
+            data_collator = None  # Default to built-in data collator
+            if is_deepseek_ocr:
+                # Special DeepSeek OCR collator - auto-install if needed
+                print("Detected DeepSeek OCR model\n")
+                # Ensure DeepSeek OCR module is installed
+                if not _ensure_deepseek_ocr_installed():
+                    error_msg = (
+                        "Failed to install DeepSeek OCR module. "
+                        "Please install manually: "
+                        "from huggingface_hub import snapshot_download; "
+                        "snapshot_download('unsloth/DeepSeek-OCR', local_dir='deepseek_ocr')"
+                    )
+                    logger.error(error_msg)
+                    self._update_progress(error=error_msg, is_training=False)
+                    return
+
+                try:
+                    from backend.data_utils import DeepSeekOCRDataCollator
+
+                    print("Configuring DeepSeek OCR data collator...\n")
+                    FastVisionModel.for_training(self.model)
+                    data_collator = DeepSeekOCRDataCollator(
+                        tokenizer=self.tokenizer,
+                        model=self.model,
+                        image_size=640,
+                        base_size=1024,
+                        crop_mode=True,
+                        train_on_responses_only=training_args.get('train_on_completions', False),
+                    )
+                    print("DeepSeek OCR data collator configured successfully\n")
+
+                except Exception as e:
+                    logger.error(f"Failed to configure DeepSeek OCR collator: {e}")
+                    error_msg = f"Error configuring DeepSeek OCR: {str(e)}"
+                    self._update_progress(error=error_msg, is_training=False)
+                    return
+
+            elif self.is_vlm:
+                # Standard VLM collator
+                print("Using UnslothVisionDataCollator for vision model\n")
+                from unsloth.trainer import UnslothVisionDataCollator
+
+                FastVisionModel.for_training(self.model)
+                data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
+                print("Vision data collator configured\n")
+
+            # ========== TRAINING CONFIGURATION ==========
+            # Handle epochs vs max_steps properly
+            max_steps_val = training_args.get('max_steps', 0)
+            num_epochs_val = training_args.get('num_epochs', 3)
+
+            # Handle warmup_steps vs warmup_ratio
+            warmup_steps_val = training_args.get('warmup_steps', None)
+            warmup_ratio_val = training_args.get('warmup_ratio', None)
+            
+            config_args = {
+                "per_device_train_batch_size": training_args.get('batch_size', 2),
+                "gradient_accumulation_steps": training_args.get('gradient_accumulation_steps', 4),
+                "num_train_epochs": training_args.get('num_epochs', 3),  # Default to epochs
+                "learning_rate": training_args.get('learning_rate', 2e-4),
+                "fp16": not is_bfloat16_supported(),
+                "bf16": is_bfloat16_supported(),
+                "logging_steps": 1,
+                "weight_decay": training_args.get('weight_decay', 0.01),
+                "seed": training_args.get('random_seed', 3407),
+                "output_dir": output_dir,
+                "report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
+            }
+            
+            # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps
+            if warmup_ratio_val is not None:
+                config_args["warmup_ratio"] = warmup_ratio_val
+                print(f"Using warmup_ratio: {warmup_ratio_val}\n")
+            elif warmup_steps_val is not None:
+                config_args["warmup_steps"] = warmup_steps_val
+                print(f"Using warmup_steps: {warmup_steps_val}\n")
+            else:
+                # Default to warmup_steps if neither provided
+                config_args["warmup_steps"] = 5
+                print(f"Using default warmup_steps: 5\n")
+
+            #  If max_steps is specified, use it instead of epochs
+            max_steps_val = training_args.get('max_steps', 0)
+            if max_steps_val and max_steps_val > 0:
+                del config_args["num_train_epochs"]  # Remove epochs
+                config_args["max_steps"] = max_steps_val  # Use steps instead
+                print(f"Training for {max_steps_val} steps\n")
+            else:
+                print(f"Training for {config_args['num_train_epochs']} epochs\n")
+
+            # Add model-specific parameters
+            # Use optim and lr_scheduler_type from training_args if provided, otherwise use defaults
+            optim_value = training_args.get('optim', "adamw_8bit")
+            lr_scheduler_type_value = training_args.get('lr_scheduler_type', "linear")
+            
+            if self.is_vlm:
+                # Vision-specific config
+                print("Configuring vision model training parameters\n")
+                # Use provided values or defaults for vision models
+                optim_value = training_args.get('optim', "adamw_torch_fused")
+                lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine")
+                config_args.update({
+                    "optim": optim_value,
+                    "lr_scheduler_type": lr_scheduler_type_value,
+                    "gradient_checkpointing": True,
+                    "gradient_checkpointing_kwargs": {"use_reentrant": False},
+                    "max_grad_norm": 0.3,  # Recommended for vision models
+                    "remove_unused_columns": False,
+                    "dataset_text_field": "",
+                    "dataset_kwargs": {"skip_prepare_dataset": True},
+                    "max_length": training_args.get('max_seq_length', 2048),
+                })
+            else:
+                print("Configuring text model training parameters\n")
+                config_args.update({
+                    "optim": optim_value,
+                    "lr_scheduler_type": lr_scheduler_type_value,
+                    "dataset_text_field": "text",
+                })
+
+                # Only add packing for text models (not DeepSeek OCR which is VLM)
+                if not is_deepseek_ocr:
+                    packing_enabled = training_args.get('packing', False)
+                    config_args["packing"] = packing_enabled
+                    print(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n")
+
+            print(f"The configuration is: {config_args}")
+
+            print("Training configuration prepared\n")
+            # ========== TRAINER INITIALIZATION ==========
+            if self.is_vlm:
+                self.trainer = SFTTrainer(
+                    model=self.model,
+                    train_dataset=dataset['dataset'],
+                    processing_class = self.tokenizer.tokenizer,
+                    data_collator=data_collator,
+                    args=SFTConfig(**config_args),
+                )
+            else:
+                self.trainer = SFTTrainer(
+                    model=self.model,
+                    tokenizer=self.tokenizer,
+                    train_dataset=dataset['dataset'],
+                    data_collator=data_collator,
+                    args=SFTConfig(**config_args),
+                )
+            print("Trainer initialized\n")
+
+            # ========== TRAIN ON RESPONSES ONLY ==========
+            # Determine if we should train on responses only
+            instruction_part = None
+            response_part = None
+            train_on_responses_enabled = training_args.get('train_on_completions', False)
+
+            # DeepSeek OCR handles this internally in its collator, so skip
+            if train_on_responses_enabled and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
+                try:
+                    print("Configuring train on responses only...\n")
+
+                    # Get the template mapping for this model
+                    model_name_lower = self.model_name.lower()
+
+                    if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
+                        template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
+                        print(f"Detected template: {template_name}\n")
+
+                        if template_name in TEMPLATE_TO_RESPONSES_MAPPER:
+                            instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["instruction"]
+                            response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"]
+
+                            print(f"Instruction marker: {instruction_part[:50]}...\n")
+                            print(f"Response marker: {response_part[:50]}...\n")
+                        else:
+                            print(f"No response mapping found for template: {template_name}\n")
+                            train_on_responses_enabled = False
+                    else:
+                        print(f"No template mapping found for model: {self.model_name}\n")
+                        train_on_responses_enabled = False
+
+                except Exception as e:
+                    logger.warning(f"Could not configure train on responses: {e}")
+                    train_on_responses_enabled = False
+
+            # Apply train on responses only if we have valid parts
+            if train_on_responses_enabled and instruction_part and response_part and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'):
+                try:
+                    from unsloth.chat_templates import train_on_responses_only
+
+                    self.trainer = train_on_responses_only(
+                        self.trainer,
+                        instruction_part=instruction_part,
+                        response_part=response_part,
+                    )
+                    print("Train on responses only configured successfully\n")
+                except Exception as e:
+                    logger.warning(f"Failed to apply train on responses only: {e}")
+                    train_on_responses_enabled = False
+            else:
+                if train_on_responses_enabled and is_deepseek_ocr:
+                    print("Train on responses handled by DeepSeek OCR collator\n")
+                else:
+                    print("Training on full sequences (including prompts)\n")
+
+            # Add custom callback for progress tracking
+            from transformers import TrainerCallback
+
+            class ProgressCallback(TrainerCallback):
+                def __init__(self, trainer_instance):
+                    self.trainer_instance = trainer_instance
+
+                def on_train_begin(self, args, state, control, **kwargs):
+                    """Called at the beginning of training"""
+                    pass
+
+                def on_log(self, args, state, control, logs=None, **kwargs):
+                    """Called when logging occurs"""
+                    if logs:
+                        # Get loss from either 'loss' or 'train_loss' key
+                        loss_value = logs.get('loss', logs.get('train_loss', 0.0))
+                        self.trainer_instance._update_progress(
+                            step=state.global_step,
+                            epoch=round(state.epoch, 2) if state.epoch else 0,  # Round epoch to 2 decimals
+                            loss=loss_value,
+                            learning_rate=logs.get('learning_rate', 0.0),
+                            status_message=""  # Clear status message so metrics show
+                        )
+
+                def on_epoch_end(self, args, state, control, **kwargs):
+                    """Called at the end of each epoch"""
+                    self.trainer_instance._update_progress(
+                        epoch=state.epoch,
+                        step=state.global_step
+                    )
+
+                def on_step_end(self, args, state, control, **kwargs):
+                    """Called at the end of each step"""
+                    # Check if we should stop training
+                    if self.trainer_instance.should_stop:
+                        print(f"Stop detected at step {state.global_step}\n")
+                        control.should_training_stop = True
+                        return control
+
+            # ========== PROGRESS TRACKING ==========
+            progress_callback = ProgressCallback(self)
+            self.trainer.add_callback(progress_callback)
+
+            num_samples = len(dataset["dataset"])
+            batch_size = training_args.get('batch_size', 2)
+            grad_accum = training_args.get('gradient_accumulation_steps', 4)
+            num_epochs = training_args.get('num_epochs', 3)
+            max_steps_val = training_args.get('max_steps', 0)
+
+            # Step 1: Calculate dataloader length (number of batches)
+            len_dataloader = math.ceil(num_samples / batch_size)
+
+            # Step 2: Calculate steps per epoch (following transformers logic)
+            num_update_steps_per_epoch = max(
+                len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0),
+                1
+            )
+
+            # Step 3: Determine total steps based on max_steps or epochs
+            if max_steps_val and max_steps_val > 0:
+                # Use max_steps if specified
+                total_steps = max_steps_val
+                print(f"Progress tracking: {total_steps} steps (max_steps)\n")
+            else:
+                # Calculate from epochs
+                total_steps = num_update_steps_per_epoch * num_epochs
+                print(f"Progress tracking: {total_steps} steps ({num_epochs} epochs × {num_update_steps_per_epoch} steps/epoch)\n")
+
+            self._update_progress(total_steps=total_steps)
+
+            # ========== START TRAINING ==========
+            self._update_progress(status_message="Starting training...")
+            print("Starting training...\n")
+            self.trainer.train()
+
+            # ========== SAVE MODEL ==========
+            self.trainer.save_model()
+            self.tokenizer.save_pretrained(output_dir)
+            print(f"\nTraining completed! Model saved to {output_dir}\n")
+
+            self._update_progress(
+                is_training=False,
+                is_completed=True,
+                #status_message=status_msg
+                status_message=f"Training completed! Model saved to {output_dir}",
+            )
+
+        except Exception as e:
+            logger.error(f"Training error: {e}")
+            self._update_progress(is_training=False, error=str(e))
+
+        finally:
+            self.is_training = False
+
+    def stop_training(self):
+        """Stop ongoing training"""
+        print("\nStopping training...")
+        self.should_stop = True
+        self.is_training = False
+        # Clear the status message so timer doesn't show stale status
+        self._update_progress(is_training=False, status_message="")
+
+        # If trainer exists, try to stop it gracefully
+        if self.trainer:
+            try:
+                # The callback will catch should_stop flag and stop the training loop
+                print("Training will stop at next step...\n")
+            except Exception as e:
+                logger.error(f"Error stopping trainer: {e}")
+
+    def get_training_progress(self) -> TrainingProgress:
+        """Get current training progress"""
+        with self._lock:
+            return self.training_progress
+
+    def cleanup(self):
+        """Cleanup resources"""
+        if self.trainer:
+            self.trainer = None
+        if self.model:
+            self.model = None
+        if self.tokenizer:
+            self.tokenizer = None
+
+        # Clear GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+def _ensure_deepseek_ocr_installed():
+    """
+    Auto-install DeepSeek OCR module if not available.
+    Downloads from HuggingFace hub as a local module.
+
+    Returns:
+        bool: True if available (either already installed or just installed)
+    """
+    try:
+        # Try importing to see if already available
+        from deepseek_ocr.modeling_deepseekocr import format_messages
+        logger.info("DeepSeek OCR module already available")
+        return True
+    except ImportError:
+        pass
+
+    try:
+        logger.info("DeepSeek OCR module not found. Auto-installing from HuggingFace...")
+        print("\n Downloading DeepSeek OCR module from HuggingFace...\n")
+
+        from huggingface_hub import snapshot_download
+        import sys
+        import os
+
+        # Get the script directory to install locally
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        parent_dir = os.path.dirname(script_dir)  # Go up to project root
+
+        # Download to project root as 'deepseek_ocr' folder
+        local_dir = os.path.join(parent_dir, "deepseek_ocr")
+
+        snapshot_download(
+            "unsloth/DeepSeek-OCR",
+            local_dir=local_dir,
+            local_dir_use_symlinks=False
+        )
+
+        # Add to sys.path if not already there
+        if parent_dir not in sys.path:
+            sys.path.insert(0, parent_dir)
+
+        # Try importing again
+        from deepseek_ocr.modeling_deepseekocr import format_messages
+
+        logger.info("DeepSeek OCR module installed successfully")
+        print("DeepSeek OCR module installed successfully!\n")
+        return True
+
+    except Exception as e:
+        logger.error(f"Failed to install DeepSeek OCR module: {e}")
+        print(f"\n❌ Failed to install DeepSeek OCR module: {e}\n")
+        return False
+
+# Global trainer instance
+_trainer_instance = None
+
+def get_trainer() -> UnslothTrainer:
+    """Get global trainer instance"""
+    global _trainer_instance
+    if _trainer_instance is None:
+        _trainer_instance = UnslothTrainer()
+    return _trainer_instance

backend/backend/training.py

@@ -0,0 +1,680 @@
+"""
+Training backend and UI integration
+"""
+import gradio as gr
+import matplotlib.pyplot as plt
+from typing import Dict, Any, Generator, Tuple
+import logging
+
+from .trainer import get_trainer, TrainingProgress
+
+logger = logging.getLogger(__name__)
+
+# Plot styling constants
+PLOT_WIDTH = 8  # Inches
+PLOT_HEIGHT = 3.5  # Inches
+
+
+class TrainingBackend:
+    """
+    Training orchestration and UI integration.
+    Handles both text and vision models, LoRA and full finetuning.
+    """
+
+    def __init__(self):
+        self.trainer = get_trainer()
+
+        # Training Metrics
+        self.loss_history = []
+        self.lr_history = []
+        self.step_history = []
+        self.current_theme = "light"
+
+        self.trainer.add_progress_callback(self._on_progress_update)
+
+        logger.info("TrainingBackend initialized")
+
+    def _on_progress_update(self, progress: TrainingProgress):
+        """Callback for progress updates"""
+        if progress.step > 0 and progress.loss > 0:
+            self.loss_history.append(progress.loss)
+            self.lr_history.append(progress.learning_rate)
+            self.step_history.append(progress.step)
+
+    def start_training(self,
+                       # Model parameters
+                       model_name: str,
+                       training_type: str,  # NEW: "LoRA/QLoRA" or "Full Finetuning"
+                       hf_token: str,
+                       load_in_4bit: bool,
+                       max_seq_length: int,
+
+                       # Dataset parameters
+                       hf_dataset: str,
+                       local_datasets: list,
+                       format_type: str,  # CHANGED: was data_template
+
+                       # Training parameters
+                       num_epochs: int,
+                       learning_rate: str,
+                       batch_size: int,
+                       gradient_accumulation_steps: int,
+                       warmup_steps: int,          # May be None even without default
+                       warmup_ratio: float,        # May be None even without default
+                       max_steps: int,
+                       save_steps: int,
+                       weight_decay: float,
+                       random_seed: int,
+                       packing: bool,
+
+                       # LoRA parameters
+                       use_lora: bool,  # Should be derived from training_type
+                       lora_r: int,
+                       lora_alpha: int,
+                       lora_dropout: float,
+                       target_modules: list,
+                       gradient_checkpointing: str,
+                       use_rslora: bool,
+                       use_loftq: bool,
+                       train_on_completions: bool,
+
+                       # NEW: Vision-specific LoRA parameters
+                       finetune_vision_layers: bool,
+                       finetune_language_layers: bool,
+                       finetune_attention_modules: bool,
+                       finetune_mlp_modules: bool,
+
+                       # Logging parameters
+                       enable_wandb: bool,
+                       wandb_token: str,
+                       wandb_project: str,
+                       enable_tensorboard: bool,
+                       tensorboard_dir: str,
+                       optim: str = "adamw_8bit",
+                       lr_scheduler_type: str = "linear") -> Generator[Tuple, None, None]:
+        """
+        Start training - yields UI updates as generator.
+
+        Yields:
+            Tuple of (start_btn_update, stop_btn_update, progress_visible, config_visible)
+        """
+        try:
+            # Reset stop flag and clear history
+            self.trainer.should_stop = False
+            self.loss_history = []
+            self.lr_history = []
+            self.step_history = []
+            import time
+            output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}"
+
+            # NEW: Derive use_lora from training_type
+            use_lora_actual = (training_type == "LoRA/QLoRA")
+            if use_lora_actual: print("using Lora")
+            else: print("using full finetuning")
+            logger.info(f"Starting training - Type: {training_type}, Model: {model_name}")
+
+            # Yield initial status - buttons toggle immediately
+            yield (
+                gr.update(interactive=False),  # Start button disabled
+                gr.update(interactive=True),   # Stop button enabled
+                gr.update(visible=True),       # Training progress visible
+                #gr.update(visible=False)       # Config selection hidden
+            )
+
+            # ========== LOAD MODEL ==========
+            logger.info("Loading model...")
+            success = self.trainer.load_model(
+                model_name=model_name,
+                max_seq_length=max_seq_length,
+                load_in_4bit=load_in_4bit if use_lora_actual else False,  # Only 4bit for LoRA
+                hf_token=hf_token if hf_token.strip() else None
+            )
+
+            if not success or self.trainer.should_stop:
+                logger.error("Failed to load model or stopped by user")
+                return
+
+            # Capture if this is a vision model
+            #self.current_training_session['is_vlm'] = self.trainer.is_vlm
+
+            yield (
+                gr.update(interactive=False),
+                gr.update(interactive=True),
+                gr.update(visible=True),
+                #gr.update(visible=False)
+            )
+
+            # ========== PREPARE MODEL FOR TRAINING ==========
+            if use_lora_actual:
+                logger.info("Preparing model with LoRA...")
+                success = self.trainer.prepare_model_for_training(
+                    use_lora=True,
+                    # Vision-specific parameters
+                    finetune_vision_layers=finetune_vision_layers,
+                    finetune_language_layers=finetune_language_layers,
+                    finetune_attention_modules=finetune_attention_modules,
+                    finetune_mlp_modules=finetune_mlp_modules,
+                    # Standard LoRA parameters
+                    target_modules=target_modules,
+                    lora_r=lora_r,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    use_gradient_checkpointing=gradient_checkpointing,
+                    use_rslora=use_rslora,
+                    use_loftq=use_loftq
+                )
+            else:
+                logger.info("Preparing model for full finetuning...")
+                success = self.trainer.prepare_model_for_training(
+                    use_lora=False  # Full finetuning
+                )
+
+            if not success or self.trainer.should_stop:
+                logger.error("Failed to prepare model or stopped by user")
+                return
+
+            yield (
+                gr.update(interactive=False),
+                gr.update(interactive=True),
+                gr.update(visible=True),
+                #gr.update(visible=False)
+            )
+
+            # ========== LOAD DATASET ==========
+            logger.info("Loading dataset...")
+            #breakpoint()
+            dataset = self.trainer.load_and_format_dataset(
+                dataset_source=hf_dataset if hf_dataset.strip() else None,
+                format_type=format_type,
+                local_datasets=local_datasets if local_datasets else None
+            )
+
+            if dataset is None or self.trainer.should_stop:
+                logger.error("Failed to load dataset or stopped by user")
+                return
+
+            yield (
+                gr.update(interactive=False),
+                gr.update(interactive=True),
+                gr.update(visible=True),
+                #gr.update(visible=False)
+            )
+
+            # ========== START TRAINING ==========
+            # Convert learning rate string to float
+            try:
+                lr_value = float(learning_rate)
+            except ValueError:
+                logger.error(f"Invalid learning rate: {learning_rate}")
+                self.trainer._update_progress(
+                    error=f"Invalid learning rate: {learning_rate}",
+                    is_training=False
+                )
+                return
+
+            logger.info("Starting training worker thread...")
+            success = self.trainer.start_training(
+                dataset=dataset,
+                #output_dir=f"./outputs/{model_name.replace('/', '_')}_{int(__import__('time').time())}",
+                output_dir=output_dir,
+                num_epochs=num_epochs,
+                learning_rate=lr_value,
+                batch_size=batch_size,
+                gradient_accumulation_steps=gradient_accumulation_steps,
+                warmup_steps=warmup_steps,
+                warmup_ratio=warmup_ratio,
+                max_steps=max_steps if max_steps > 0 else 0,
+                save_steps=save_steps if save_steps > 0 else 0,
+                weight_decay=weight_decay,
+                random_seed=random_seed,
+                packing=packing,
+                train_on_completions=train_on_completions,
+                enable_wandb=enable_wandb,
+                wandb_project=wandb_project,
+                wandb_token=wandb_token if wandb_token.strip() else None,
+                enable_tensorboard=enable_tensorboard,
+                tensorboard_dir=tensorboard_dir,
+                max_seq_length=max_seq_length,  # Pass through for config
+                optim=optim,
+                lr_scheduler_type=lr_scheduler_type,
+            )
+
+            if not success:
+                logger.error("Failed to start training")
+                yield (
+                    gr.update(interactive=True),
+                    gr.update(interactive=False),
+                    gr.update(visible=False),
+                    #gr.update(visible=True)
+                )
+
+        except Exception as e:
+            logger.error(f"Error in start_training: {e}", exc_info=True)
+            self.trainer._update_progress(
+                error=str(e),
+                is_training=False
+            )
+            yield (
+                gr.update(interactive=True),
+                gr.update(interactive=False),
+                gr.update(visible=False),
+                #gr.update(visible=True)
+            )
+
+    def stop_training(self) -> Tuple:
+        """
+        Stop ongoing training.
+
+        Returns:
+            Tuple of (start_btn_update, stop_btn_update, progress_visible, config_visible)
+        """
+        try:
+            logger.info("Stopping training...")
+            self.trainer.stop_training()
+
+            return (
+                gr.update(interactive=True),   # Start button enabled
+                gr.update(interactive=False),  # Stop button disabled
+                gr.update(visible=False),      # Training progress hidden
+                #gr.update(visible=True)        # Config selection visible
+            )
+        except Exception as e:
+            logger.error(f"Error stopping training: {e}")
+            return (
+                gr.update(interactive=True),
+                gr.update(interactive=False),
+                gr.update(visible=False),
+                #gr.update(visible=True)
+            )
+
+    def get_training_status(self, theme: str = "light") -> Tuple[plt.Figure, gr.update, gr.update, gr.update]:
+        """
+        Get current training status and loss plot.
+
+        Args:
+            theme: "light" or "dark" for plot styling
+
+        Returns:
+            Tuple of (plot, start_btn, stop_btn, progress_visible)
+        """
+
+        try:
+            progress = self.trainer.get_training_progress()
+
+            # If not training and not completed, return no updates
+            if not (progress.is_training or progress.is_completed or progress.error):
+                return (None, gr.update(), gr.update(), gr.update())
+
+            # Generate plot
+            plot = self._create_loss_plot(progress, theme)
+
+            # If completed or error, enable start button
+            if progress.is_completed or progress.error:
+                return (
+                    plot,
+                    gr.update(interactive=True),   # Start button enabled
+                    gr.update(interactive=False),  # Stop button disabled
+                    gr.update(visible=True),       # Training progress visible
+                )
+
+            # Still training - no button updates
+            return (plot, gr.update(), gr.update(), gr.update())
+
+        except Exception as e:
+            logger.error(f"Error getting training status: {e}")
+            return (None, gr.update(), gr.update(), gr.update())
+
+    def refresh_plot_for_theme(self, theme: str) -> plt.Figure:
+        """
+        Refresh plot with new theme.
+
+        Args:
+            theme: "light" or "dark"
+
+        Returns:
+            Updated matplotlib figure
+        """
+        if theme and isinstance(theme, str) and theme in ['light', 'dark']:
+            self.current_theme = theme
+
+        # Always generate plot if we have loss history
+        if self.loss_history:
+            progress = self.trainer.get_training_progress()
+            return self._create_loss_plot(progress, self.current_theme)
+
+        return None
+
+    def is_training_active(self) -> bool:
+        """
+        Check if training is currently active (from load_model start to completion/error).
+        
+        Returns:
+            True if training is in progress, False otherwise
+        """
+        try:
+            progress = self.trainer.get_training_progress()
+            # Training is active if is_training is True
+            # Also check if we're in loading/preparation phase (status_message indicates activity)
+            is_active = progress.is_training
+            # Also consider it active if we have a status message indicating loading/preparation
+            # but haven't completed or errored yet
+            if not is_active and not progress.is_completed and not progress.error:
+                status = progress.status_message or ""
+                if any(keyword in status.lower() for keyword in ["loading", "preparing", "training"]):
+                    is_active = True
+            return is_active
+        except Exception as e:
+            logger.error(f"Error checking training state: {e}")
+            return False
+
+    def _create_loss_plot(self, progress: TrainingProgress, theme: str = "light") -> plt.Figure:
+            """
+            Create training loss plot with theme-aware styling.
+
+            Args:
+                progress: Current training progress
+                theme: "light" or "dark"
+
+            Returns:
+                Matplotlib figure
+            """
+            plt.close('all')
+
+            # Theme-specific styling
+            LIGHT_STYLE = {
+                "facecolor": "#ffffff",
+                "grid_color": "#d1d5db",
+                "line": "#16b88a",
+                "text": "#1f2937",
+                "empty_text": "#6b7280"
+            }
+            DARK_STYLE = {
+                "facecolor": "#292929",
+                "grid_color": "#404040",
+                "line": "#4ade80",
+                "text": "#e5e7eb",
+                "empty_text": "#9ca3af"
+            }
+
+            style = LIGHT_STYLE if theme == "light" else DARK_STYLE
+
+            fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT))
+            fig.patch.set_facecolor(style["facecolor"])
+            ax.set_facecolor(style["facecolor"])
+
+            if self.loss_history:
+                steps = self.step_history
+                losses = self.loss_history
+                scatter_color = "#60a5fa"
+                # Scatter plot for raw loss points
+                ax.scatter(
+                    steps,
+                    losses,
+                    s=16,
+                    alpha=0.6,
+                    color=scatter_color,
+                    linewidths=0,
+                    label="Training Loss (raw)",
+                )
+
+                # Moving average line overlay (trailing window)
+                MA_WINDOW = 20  # adjust smoothing aggressiveness
+                window = min(MA_WINDOW, len(losses))
+
+                if window >= 2:
+                    cumsum = [0.0]
+                    for v in losses:
+                        cumsum.append(cumsum[-1] + float(v))
+
+                    ma = []
+                    for i in range(len(losses)):
+                        start = max(0, i - window + 1)
+                        denom = i - start + 1
+                        ma.append((cumsum[i + 1] - cumsum[start]) / denom)
+
+                    ax.plot(
+                        steps,
+                        ma,
+                        color=style["line"],
+                        linewidth=2.5,
+                        alpha=0.95,
+                        label=f"Moving Avg ({ma[-1]:.4f})",
+                    )
+
+                    leg = ax.legend(frameon=False, fontsize=9)
+                    for t in leg.get_texts():
+                        t.set_color(style["text"])
+
+                ax.set_xlabel('Steps', fontsize=10, color=style["text"])
+                ax.set_ylabel('Loss', fontsize=10, color=style["text"])
+
+                # Build status message for title
+                if progress.error:
+                    title = f"Error: {progress.error}"
+                elif progress.is_completed:
+                    title = f"Training completed! Final loss: {progress.loss:.4f}"
+                elif progress.status_message:
+                    title = progress.status_message
+                elif progress.step > 0:
+                    title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {progress.loss:.4f}"
+                else:
+                    title = "Training Loss"
+
+                ax.set_title(title, fontsize=11, fontweight='bold',
+                            pad=10, color=style["text"])
+
+                # Style grid and spines
+                ax.grid(True, alpha=0.4, linestyle='--', color=style["grid_color"])
+                ax.tick_params(colors=style["text"], which='both')
+                ax.spines['top'].set_visible(False)
+                ax.spines['right'].set_visible(False)
+                ax.spines['bottom'].set_color(style["text"])
+                ax.spines['left'].set_color(style["text"])
+            else:
+                display_msg = progress.status_message if progress.status_message else 'Waiting for training data...'
+                ax.text(0.5, 0.5, display_msg,
+                       ha='center', va='center', fontsize=16,
+                       color=style["empty_text"],
+                       transform=ax.transAxes)
+                ax.set_xticks([])
+                ax.set_yticks([])
+                for spine in ax.spines.values():
+                    spine.set_visible(False)
+
+            fig.tight_layout()
+            return fig
+
+    def _transfer_to_inference_backend(self) -> bool:
+        """
+        Transfer the trained model to InferenceBackend.
+        Called automatically when training completes.
+        """
+        print("=" * 60)
+        print("DEBUG: _transfer_to_inference_backend() CALLED")
+        print("=" * 60)
+
+        try:
+            from .inference import get_inference_backend
+
+            session = self.current_training_session
+
+            # Check if already transferred
+            if session.get('transferred', False):
+                print("DEBUG: Already transferred, returning True")
+                logger.info("Model already transferred, skipping")
+                return True
+
+            # Validate session data
+            if not session.get('base_model_name') or not session.get('output_dir'):
+                logger.warning("Training session incomplete, cannot transfer")
+                logger.warning(f"Session data: {session}")
+                return False
+
+            inference_backend = get_inference_backend()
+
+            base_model_name = session['base_model_name']
+            output_dir = session['output_dir']
+            is_lora = session['is_lora']
+            is_vlm = session['is_vlm']
+
+            logger.info(f"=" * 60)
+            logger.info(f"TRANSFERRING MODEL TO INFERENCE BACKEND")
+            logger.info(f"=" * 60)
+            logger.info(f"  Base model: {base_model_name}")
+            logger.info(f"  Output dir: {output_dir}")
+            logger.info(f"  Is LoRA: {is_lora}")
+            logger.info(f"  Is VLM: {is_vlm}")
+
+            # Transfer the model object directly from trainer memory.
+            # If is_lora is True, self.trainer.model is a PeftModel (Base + Adapter).
+            # If is_lora is False, it is the finetuned Base Model.
+            inference_backend.models[base_model_name] = {
+                "model": self.trainer.model,
+                "tokenizer": self.trainer.tokenizer,
+                "is_vision": is_vlm,
+                "is_lora": is_lora,
+                "model_path": base_model_name,
+                "base_model": None,
+                "loaded_adapters": {},
+                # Unsloth/PEFT training keeps the active adapter named 'default' in memory
+                "active_adapter": "default" if is_lora else None,
+            }
+
+            # For vision models, also transfer processor
+            if is_vlm:
+                if hasattr(self.trainer, 'tokenizer'):
+                    inference_backend.models[base_model_name]["processor"] = self.trainer.tokenizer
+                    logger.info("  Transferred processor for vision model")
+
+            # Load chat template info
+            inference_backend._load_chat_template_info(base_model_name)
+
+            # If it was LoRA, register the output path.
+            # This ensures the Eval UI dropdown (which lists files) knows that
+            # the model currently in memory corresponds to this specific output directory.
+            if is_lora:
+                inference_backend.models[base_model_name]["last_trained_adapter"] = output_dir
+                logger.info(f"Marked trained LoRA adapter: {output_dir}")
+
+            # Set as active model
+            inference_backend.active_model_name = base_model_name
+            logger.info(f"Set active model: {base_model_name}")
+
+            return True
+
+        except Exception as e:
+            logger.error(f"Error transferring model to inference backend: {e}")
+            import traceback
+            traceback.print_exc()
+            return False
+
+
+# ========== GLOBAL INSTANCE ==========
+_training_backend = None
+
+def get_training_backend() -> TrainingBackend:
+    """Get global training backend instance"""
+    global _training_backend
+    if _training_backend is None:
+        _training_backend = TrainingBackend()
+    return _training_backend
+
+
+# ========== UI HANDLER CREATION ==========
+def create_training_handlers(train_components: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Create training event handlers for Gradio UI components.
+
+    Args:
+        train_components: Dictionary of Gradio components from train page
+
+    Returns:
+        Dictionary of handler functions
+    """
+    backend = get_training_backend()
+
+    def start_training_handler(*args):
+        """Handler for start training button - yields status updates"""
+        try:
+            # Extract parameters in the order they're passed from the UI
+            (model_name, training_type, hf_token, load_4bit, max_seq_length,
+             hf_dataset, local_datasets, format_type,
+             num_epochs, learning_rate, batch_size, gradient_accumulation_steps,
+             warmup_steps, warmup_ratio, max_steps, save_steps, weight_decay, random_seed, packing,
+             use_lora, lora_r, lora_alpha, lora_dropout, target_modules,
+             gradient_checkpointing, use_rslora, use_loftq, train_on_completions,
+             finetune_vision_layers, finetune_language_layers,
+             finetune_attention_modules, finetune_mlp_modules,
+             enable_wandb, wandb_token, wandb_project,
+             enable_tensorboard, tensorboard_dir, optim, lr_scheduler_type) = args
+
+            # Start training with correctly named parameters - this is a generator
+            for update_tuple in backend.start_training(
+                model_name=model_name,
+                training_type=training_type,
+                hf_token=hf_token,
+                load_in_4bit=load_4bit,
+                max_seq_length=max_seq_length,
+                hf_dataset=hf_dataset,
+                local_datasets=local_datasets,
+                format_type=format_type,
+                num_epochs=num_epochs,
+                learning_rate=learning_rate,
+                batch_size=batch_size,
+                gradient_accumulation_steps=gradient_accumulation_steps,
+                warmup_steps=warmup_steps,
+                warmup_ratio=warmup_ratio,
+                max_steps=max_steps,
+                save_steps=save_steps,
+                weight_decay=weight_decay,
+                random_seed=random_seed,
+                packing=packing,
+                use_lora=use_lora,
+                lora_r=lora_r,
+                lora_alpha=lora_alpha,
+                lora_dropout=lora_dropout,
+                target_modules=target_modules,
+                gradient_checkpointing=gradient_checkpointing,
+                use_rslora=use_rslora,
+                use_loftq=use_loftq,
+                train_on_completions=train_on_completions,
+                finetune_vision_layers=finetune_vision_layers,
+                finetune_language_layers=finetune_language_layers,
+                finetune_attention_modules=finetune_attention_modules,
+                finetune_mlp_modules=finetune_mlp_modules,
+                enable_wandb=enable_wandb,
+                wandb_token=wandb_token,
+                wandb_project=wandb_project,
+                enable_tensorboard=enable_tensorboard,
+                tensorboard_dir=tensorboard_dir
+            ):
+                # Yield each status update to Gradio
+                yield update_tuple
+
+        except Exception as e:
+            logger.error(f"Error in start_training_handler: {e}", exc_info=True)
+            yield (
+                gr.update(interactive=True),   # Start button
+                gr.update(interactive=False),  # Stop button
+                gr.update(visible=False),      # Training progress
+                #gr.update(visible=True)        # Config selection
+            )
+
+    def stop_training_handler():
+        """Handler for stop training button"""
+        return backend.stop_training()
+
+    def update_training_status():
+        """Periodic update of training status and plot"""
+        return backend.get_training_status(backend.current_theme)
+
+    def refresh_plot_for_theme(theme):
+        """Refresh plot with new theme"""
+        return backend.refresh_plot_for_theme(theme)
+
+    return {
+        'start_training': start_training_handler,
+        'stop_training': stop_training_handler,
+        'update_status': update_training_status,
+        'refresh_plot': refresh_plot_for_theme
+    }

backend/backend/utils.py

@@ -0,0 +1,208 @@
+"""
+Shared backend utilities
+"""
+import gradio as gr
+import os
+import logging
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Optional, Dict, Any
+import shutil
+import tempfile
+
+
+logger = logging.getLogger(__name__)
+
+@contextmanager
+def without_hf_auth():
+    """
+    Context manager to temporarily disable HuggingFace authentication.
+
+    Usage:
+        with without_hf_auth():
+            # Code that should run without cached tokens
+            model_info(model_name, token=None)
+    """
+    # Save environment variables
+    saved_env = {}
+    env_vars = ['HF_TOKEN', 'HUGGINGFACE_HUB_TOKEN', 'HF_HOME']
+    for var in env_vars:
+        if var in os.environ:
+            saved_env[var] = os.environ[var]
+            del os.environ[var]
+
+    # Save disable flag
+    saved_disable = os.environ.get('HF_HUB_DISABLE_IMPLICIT_TOKEN')
+    os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = '1'
+
+    # Move token files temporarily
+    token_files = []
+    token_locations = [
+        Path.home() / '.cache' / 'huggingface' / 'token',
+        Path.home() / '.huggingface' / 'token'
+    ]
+
+    for token_loc in token_locations:
+        if token_loc.exists():
+            temp = tempfile.NamedTemporaryFile(delete=False)
+            temp.close()
+            shutil.move(str(token_loc), temp.name)
+            token_files.append((token_loc, temp.name))
+
+    try:
+        yield
+    finally:
+        # Restore tokens
+        for original, temp in token_files:
+            try:
+                original.parent.mkdir(parents=True, exist_ok=True)
+                shutil.move(temp, str(original))
+            except Exception as e:
+                logger.error(f"Failed to restore token {original}: {e}")
+
+        # Restore environment
+        for var, value in saved_env.items():
+            os.environ[var] = value
+
+        if saved_disable is not None:
+            os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = saved_disable
+        else:
+            os.environ.pop('HF_HUB_DISABLE_IMPLICIT_TOKEN', None)
+pass
+
+def format_error_message(error: Exception, model_name: str) -> str:
+    """
+    Format user-friendly error messages for common issues.
+
+    Args:
+        error: The exception that occurred
+        model_name: Name of the model being loaded
+
+    Returns:
+        User-friendly error string
+    """
+    error_str = str(error).lower()
+    model_short = model_name.split('/')[-1] if '/' in model_name else model_name
+
+    if "repository not found" in error_str or "404" in error_str:
+        return f"Model '{model_short}' not found. Check the model name."
+
+    if "401" in error_str or "unauthorized" in error_str:
+        return f"Authentication failed for '{model_short}'. Please provide a valid HF token."
+
+    if "gated" in error_str or "access to model" in error_str:
+        return f"Model '{model_short}' requires authentication. Please provide a valid HF token."
+
+    if "invalid user token" in error_str:
+        return "Invalid HF token. Please check your token and try again."
+
+    if "memory" in error_str or "cuda" in error_str or "out of memory" in error_str:
+        return f"Not enough GPU memory to load '{model_short}'. Try a smaller model or free GPU memory."
+
+    # Generic fallback
+    return str(error)
+pass
+
+def get_gpu_memory_info() -> Dict[str, Any]:
+    """Get GPU memory information."""
+    import torch
+
+    if not torch.cuda.is_available():
+        return {"available": False}
+
+    try:
+        device = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(device)
+
+        total = props.total_memory
+        allocated = torch.cuda.memory_allocated(device)
+        reserved = torch.cuda.memory_reserved(device)
+
+        return {
+            "available": True,
+            "device": device,
+            "total_gb": total / (1024**3),
+            "allocated_gb": allocated / (1024**3),
+            "reserved_gb": reserved / (1024**3),
+            "free_gb": (total - allocated) / (1024**3),
+            "utilization_pct": (allocated / total) * 100
+        }
+    except Exception as e:
+        logger.error(f"Error getting GPU info: {e}")
+        return {"available": False, "error": str(e)}
+pass
+
+def log_gpu_memory(context: str):
+    """Log GPU memory usage with context."""
+    memory_info = get_gpu_memory_info()
+    if memory_info.get("available"):
+        logger.info(
+            f"GPU Memory [{context}]: "
+            f"{memory_info['allocated_gb']:.2f}GB/{memory_info['total_gb']:.2f}GB "
+            f"({memory_info['utilization_pct']:.1f}% used, "
+            f"{memory_info['free_gb']:.2f}GB free)"
+        )
+    else:
+        logger.info(f"GPU Memory [{context}]: No CUDA GPU available")
+pass
+
+"""
+Model utility functions - search, discovery, etc.
+"""
+
+
+def search_hf_models(search_query: str, hf_token: Optional[str] = None):
+    """
+    Search HuggingFace model hub.
+    """
+    import requests
+
+    if not search_query or not search_query.strip():
+        return gr.update(choices=[])
+
+    # Simple debouncing: only search if query is at least 2 characters
+    if len(search_query.strip()) < 2:
+        return gr.update(choices=[])
+
+    try:
+        headers = {}
+        if hf_token and hf_token.strip():
+            headers["Authorization"] = f"Bearer {hf_token.strip()}"
+
+        url = "https://huggingface.co/api/models"
+        params = {
+            "search": search_query,
+            "pipeline_tag": "text-generation",
+            "library": "transformers",
+            "limit": 15,
+            "sort": "downloads",
+            "direction": -1
+        }
+
+        response = requests.get(url, headers=headers, params=params, timeout=10)
+
+        if response.status_code == 200:
+            models = response.json()
+            unsloth_results = []
+            other_results = []
+
+            for model in models:
+                model_id = model.get("modelId", "")
+                if model_id and "gguf" not in model_id.lower():
+                    result = (f"{model_id}", model_id)
+
+                    if model_id.startswith("unsloth/"):
+                        unsloth_results.append(result)
+                    else:
+                        other_results.append(result)
+
+            # Combine with unsloth models first
+            search_results = unsloth_results + other_results
+            return gr.update(choices=search_results)
+        else:
+            logger.warning(f"HF API returned status {response.status_code}")
+            return gr.update(choices=[])
+
+    except Exception as e:
+        logger.warning(f"Model search failed: {e}")
+        return gr.update(choices=[])

backend/main.py

@@ -0,0 +1,113 @@
+"""
+Main FastAPI application for Unsloth UI Backend
+"""
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+from pathlib import Path
+from datetime import datetime
+
+# Import routers
+from routes import training_router, models_router
+
+# Create FastAPI app
+app = FastAPI(
+    title="Unsloth UI Backend",
+    version="1.0.0",
+    description="Backend API for Unsloth UI - Training and Model Management"
+)
+
+# CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, specify allowed origins
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ============ Register API Routes ============
+
+# Register routers
+app.include_router(training_router, prefix="/api/train", tags=["training"])
+app.include_router(models_router, prefix="/api/models", tags=["models"])
+
+
+# ============ Health and System Endpoints ============
+
+@app.get("/api/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "timestamp": datetime.now().isoformat(),
+        "service": "Unsloth UI Backend"
+    }
+
+
+@app.get("/api/system")
+async def get_system_info():
+    """Get system information"""
+    import torch
+    import platform
+    import psutil
+
+    # GPU Info
+    gpu_info = {"available": False, "devices": []}
+    if torch.cuda.is_available():
+        gpu_info["available"] = True
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            gpu_info["devices"].append(
+                {
+                    "index": i,
+                    "name": props.name,
+                    "memory_total_gb": round(props.total_memory / 1e9, 2),
+                }
+            )
+
+    # CPU & Memory
+    memory = psutil.virtual_memory()
+
+    return {
+        "platform": platform.platform(),
+        "python_version": platform.python_version(),
+        "cpu_count": psutil.cpu_count(),
+        "memory": {
+            "total_gb": round(memory.total / 1e9, 2),
+            "available_gb": round(memory.available / 1e9, 2),
+            "percent_used": memory.percent,
+        },
+        "gpu": gpu_info,
+    }
+
+
+# ============ Serve Frontend (Optional) ============
+
+def setup_frontend(app: FastAPI, build_path: Path):
+    """Mount frontend static files (optional)"""
+    if build_path.exists():
+        # Mount assets
+        assets_dir = build_path / "assets"
+        if assets_dir.exists():
+            app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
+
+        @app.get("/")
+        async def serve_root():
+            return FileResponse(build_path / "index.html")
+
+        @app.get("/{full_path:path}")
+        async def serve_frontend(full_path: str):
+            if full_path.startswith("api"):
+                return {"error": "API endpoint not found"}
+
+            file_path = build_path / full_path
+            if file_path.is_file():
+                return FileResponse(file_path)
+
+            return FileResponse(build_path / "index.html")
+
+        return True
+    return False
+

backend/models/__init__.py

@@ -0,0 +1,37 @@
+"""
+Pydantic models for API request/response schemas
+"""
+from .training import (
+    TrainingStartRequest,
+    TrainingStartResponse,
+    TrainingStatusResponse,
+    TrainingMetricsResponse,
+    TrainingProgressResponse,
+)
+from .models import (
+    ModelSearchRequest,
+    ModelSearchResponse,
+    ModelListResponse,
+    ModelConfigResponse,
+    LoRAScanResponse,
+    LoRAInfo,
+    ModelInfo,
+)
+
+__all__ = [
+    # Training schemas
+    "TrainingStartRequest",
+    "TrainingStartResponse",
+    "TrainingStatusResponse",
+    "TrainingMetricsResponse",
+    "TrainingProgressResponse",
+    # Model management schemas
+    "ModelSearchRequest",
+    "ModelSearchResponse",
+    "ModelListResponse",
+    "ModelConfigResponse",
+    "LoRAScanResponse",
+    "LoRAInfo",
+    "ModelInfo",
+]
+

backend/models/models.py

@@ -0,0 +1,56 @@
+"""
+Pydantic schemas for Model Management API
+"""
+from pydantic import BaseModel, Field
+from typing import Optional, List, Dict, Any
+
+
+class ModelSearchRequest(BaseModel):
+    """Request schema for searching HuggingFace models"""
+    query: str = Field(..., description="Search query")
+    hf_token: Optional[str] = Field(None, description="HuggingFace token for authenticated searches")
+
+
+class ModelInfo(BaseModel):
+    """Model information"""
+    id: str = Field(..., description="Model identifier")
+    name: Optional[str] = Field(None, description="Display name")
+    description: Optional[str] = Field(None, description="Model description")
+    size: Optional[str] = Field(None, description="Model size")
+    is_vision: bool = Field(False, description="Whether model is a vision model")
+    is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
+
+
+class ModelSearchResponse(BaseModel):
+    """Response schema for model search"""
+    models: List[ModelInfo] = Field(default_factory=list, description="List of matching models")
+    total: int = Field(0, description="Total number of results")
+
+
+class ModelListResponse(BaseModel):
+    """Response schema for listing available models"""
+    models: List[ModelInfo] = Field(default_factory=list, description="List of available models")
+    default_models: List[str] = Field(default_factory=list, description="List of default model IDs")
+
+
+class ModelConfigResponse(BaseModel):
+    """Response schema for model configuration"""
+    model_name: str = Field(..., description="Model identifier")
+    config: Dict[str, Any] = Field(..., description="Model configuration dictionary")
+    is_vision: bool = Field(False, description="Whether model is a vision model")
+    is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
+    base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter")
+
+
+class LoRAInfo(BaseModel):
+    """LoRA adapter information"""
+    display_name: str = Field(..., description="Display name for the LoRA")
+    adapter_path: str = Field(..., description="Path to the LoRA adapter")
+    base_model: Optional[str] = Field(None, description="Base model identifier")
+
+
+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")
+

backend/models/training.py

@@ -0,0 +1,96 @@
+"""
+Pydantic schemas for Training API
+"""
+from pydantic import BaseModel, Field
+from typing import Optional, List
+
+
+class TrainingStartRequest(BaseModel):
+    """Request schema for starting training"""
+    # Model parameters
+    model_name: str = Field(..., description="Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')")
+    training_type: str = Field(..., description="Training type: 'LoRA/QLoRA' or 'Full Finetuning'")
+    hf_token: Optional[str] = Field(None, description="HuggingFace token")
+    load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
+    max_seq_length: int = Field(2048, description="Maximum sequence length")
+
+    # Dataset parameters
+    hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier")
+    local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths")
+    format_type: str = Field(..., description="Dataset format type")
+
+    # Training parameters
+    num_epochs: int = Field(1, description="Number of training epochs")
+    learning_rate: str = Field("2e-4", description="Learning rate")
+    batch_size: int = Field(1, description="Batch size")
+    gradient_accumulation_steps: int = Field(1, description="Gradient accumulation steps")
+    warmup_steps: Optional[int] = Field(None, description="Warmup steps")
+    warmup_ratio: Optional[float] = Field(None, description="Warmup ratio")
+    max_steps: Optional[int] = Field(None, description="Maximum training steps")
+    save_steps: int = Field(100, description="Steps between checkpoints")
+    weight_decay: float = Field(0.01, description="Weight decay")
+    random_seed: int = Field(42, description="Random seed")
+    packing: bool = Field(False, description="Enable sequence packing")
+
+    # LoRA parameters
+    use_lora: bool = Field(True, description="Use LoRA (derived from training_type)")
+    lora_r: int = Field(16, description="LoRA rank")
+    lora_alpha: int = Field(16, description="LoRA alpha")
+    lora_dropout: float = Field(0.0, description="LoRA dropout")
+    target_modules: List[str] = Field(default_factory=list, description="Target modules for LoRA")
+    gradient_checkpointing: str = Field("", description="Gradient checkpointing setting")
+    use_rslora: bool = Field(False, description="Use RSLoRA")
+    use_loftq: bool = Field(False, description="Use LoftQ")
+    train_on_completions: bool = Field(False, description="Train on completions only")
+
+    # Vision-specific LoRA parameters
+    finetune_vision_layers: bool = Field(False, description="Finetune vision layers")
+    finetune_language_layers: bool = Field(False, description="Finetune language layers")
+    finetune_attention_modules: bool = Field(False, description="Finetune attention modules")
+    finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules")
+
+    # Logging parameters
+    enable_wandb: bool = Field(False, description="Enable Weights & Biases logging")
+    wandb_token: Optional[str] = Field(None, description="W&B token")
+    wandb_project: Optional[str] = Field(None, description="W&B project name")
+    enable_tensorboard: bool = Field(False, description="Enable TensorBoard logging")
+    tensorboard_dir: Optional[str] = Field(None, description="TensorBoard directory")
+    optim: str = Field("adamw_8bit", description="Optimizer")
+    lr_scheduler_type: str = Field("linear", description="Learning rate scheduler type")
+
+
+class TrainingStartResponse(BaseModel):
+    """Response schema for training start"""
+    status: str = Field(..., description="Status: 'started' or 'error'")
+    job_id: Optional[str] = Field(None, description="Training job ID")
+    message: str = Field(..., description="Status message")
+    error: Optional[str] = Field(None, description="Error message if status is 'error'")
+
+
+class TrainingStatusResponse(BaseModel):
+    """Response schema for training status"""
+    status: str = Field(..., description="Status: 'idle', 'preparing', 'training', 'stopping', 'error'")
+    is_active: bool = Field(..., description="Whether training is currently active (actual training running)")
+    message: str = Field(..., description="Status message")
+    current_step: Optional[int] = Field(None, description="Current training step")
+    total_steps: Optional[int] = Field(None, description="Total training steps")
+
+
+class TrainingMetricsResponse(BaseModel):
+    """Response schema for training metrics"""
+    loss_history: List[float] = Field(default_factory=list, description="Loss values")
+    lr_history: List[float] = Field(default_factory=list, description="Learning rate values")
+    step_history: List[int] = Field(default_factory=list, description="Step numbers")
+    current_loss: Optional[float] = Field(None, description="Current loss value")
+    current_lr: Optional[float] = Field(None, description="Current learning rate")
+    current_step: Optional[int] = Field(None, description="Current step")
+
+
+class TrainingProgressResponse(BaseModel):
+    """Response schema for training progress updates"""
+    step: int = Field(..., description="Current step")
+    loss: float = Field(..., description="Current loss")
+    learning_rate: float = Field(..., description="Current learning rate")
+    status_message: str = Field(..., description="Status message")
+    progress_percent: Optional[float] = Field(None, description="Progress percentage")
+

backend/requirements.txt

@@ -0,0 +1,7 @@
+fastapi>=0.100.0
+uvicorn>=0.27.0
+pydantic>=2.0
+torch
+psutil
+nest-asyncio>=1.5.8
+

backend/routes/__init__.py

@@ -0,0 +1,8 @@
+"""
+API Routes
+"""
+
+from routes.training import router as training_router
+from routes.models import router as models_router
+
+__all__ = ["training_router", "models_router"]

backend/routes/models.py

@@ -0,0 +1,310 @@
+"""
+Model Management API routes
+"""
+import sys
+from pathlib import Path
+from fastapi import APIRouter, HTTPException, Query
+from typing import 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))
+
+# Import backend functions
+try:
+    from backend.utils import search_hf_models
+    from backend.model_config import (
+        scan_trained_loras,
+        load_model_defaults,
+        get_base_model_from_lora,
+        is_vision_model,
+        ModelConfig,
+    )
+    from backend.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 backend.utils import search_hf_models
+    from backend.model_config import (
+        scan_trained_loras,
+        load_model_defaults,
+        get_base_model_from_lora,
+        is_vision_model,
+        ModelConfig,
+    )
+    from backend.inference import get_inference_backend
+
+from models.models import (
+    ModelSearchRequest,
+    ModelSearchResponse,
+    ModelInfo,
+    ModelListResponse,
+    ModelConfigResponse,
+    LoRAScanResponse,
+    LoRAInfo,
+)
+
+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)
+
+
+@router.post("/search")
+async def search_models(request: ModelSearchRequest):
+    """
+    Search for models on HuggingFace Hub.
+    
+    This endpoint wraps the backend search_hf_models function.
+    """
+    try:
+        # Call backend search function
+        gradio_update = search_hf_models(
+            search_query=request.query,
+            hf_token=request.hf_token
+        )
+        
+        # Convert Gradio update to list of model IDs
+        model_list = []
+        if gradio_update and hasattr(gradio_update, 'choices'):
+            choices = gradio_update.choices
+        elif isinstance(gradio_update, dict) and 'choices' in gradio_update:
+            choices = gradio_update['choices']
+        elif isinstance(gradio_update, list):
+            choices = gradio_update
+        else:
+            choices = []
+        
+        # Process choices - they may be tuples (display_name, model_id) or just strings
+        for choice in choices:
+            if isinstance(choice, tuple) and len(choice) >= 2:
+                # Format: (display_name, model_id)
+                model_id = choice[1] if len(choice) > 1 else choice[0]
+                display_name = choice[0]
+                model_info = ModelInfo(
+                    id=model_id,
+                    name=display_name
+                )
+            elif isinstance(choice, str):
+                # Just a model ID string
+                model_info = ModelInfo(id=choice)
+            else:
+                continue
+            model_list.append(model_info)
+        
+        return ModelSearchResponse(
+            models=model_list,
+            total=len(model_list)
+        )
+        
+    except Exception as e:
+        logger.error(f"Error searching models: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to search models: {str(e)}"
+        )
+
+
+@router.get("/list")
+async def list_models():
+    """
+    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 = ModelInfo(
+                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 = ModelInfo(
+                    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):
+    """
+    Get configuration for a specific model.
+    
+    This endpoint wraps the backend load_model_defaults function.
+    """
+    try:
+        # 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
+        
+        return ModelConfigResponse(
+            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")
+):
+    """
+    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")
+async def get_lora_base_model(lora_path: str):
+    """
+    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 {
+            "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}")
+async def check_vision_model(model_name: str):
+    """
+    Check if a model is a vision model.
+    
+    This endpoint wraps the backend is_vision_model function.
+    """
+    try:
+        is_vision = is_vision_model(model_name)
+        
+        return {
+            "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)}"
+        )
+

backend/routes/training.py

@@ -0,0 +1,437 @@
+"""
+Training API routes
+"""
+import sys
+from pathlib import Path
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
+from typing import Dict
+import logging
+import asyncio
+from datetime import datetime
+import threading
+
+# Add backend directory to path
+# The backend code should be in the same directory structure
+backend_path = Path(__file__).parent.parent.parent
+if str(backend_path) not in sys.path:
+    sys.path.insert(0, str(backend_path))
+
+# Import backend functions
+try:
+    from backend.training import get_training_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 backend.training import get_training_backend
+
+from models.training import (
+    TrainingStartRequest,
+    TrainingStartResponse,
+    TrainingStatusResponse,
+    TrainingMetricsResponse,
+    TrainingProgressResponse,
+)
+
+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)
+
+
+@router.post("/start")
+async def start_training(request: TrainingStartRequest):
+    """
+    Start a training job.
+    
+    This endpoint initiates training in the background and returns immediately.
+    Use the /status endpoint to check training progress.
+    """
+    try:
+        logger.info(f"Starting training job with model: {request.model_name}")
+        backend = get_training_backend()
+        
+        # Check if training is already active
+        if backend.is_training_active():
+            return TrainingStartResponse(
+                status="error",
+                message="Training is already in progress. Stop current training before starting a new one.",
+                error="Training already active"
+            )
+        
+        # Validate dataset paths if provided
+        if request.local_datasets:
+            validated_datasets = []
+            # Get the backend directory (where this file is located)
+            backend_dir = Path(__file__).parent.parent
+            utils_datasets_dir = backend_dir / "utils" / "datasets"
+            
+            for dataset_path in request.local_datasets:
+                dataset_file = Path(dataset_path)
+                
+                # If not absolute, try multiple locations
+                if not dataset_file.is_absolute():
+                    # First try: relative to current working directory
+                    candidate = Path.cwd() / dataset_path
+                    if not candidate.exists():
+                        # Second try: relative to utils/datasets folder
+                        candidate = utils_datasets_dir / dataset_path
+                    if not candidate.exists():
+                        # Third try: just the filename in utils/datasets
+                        candidate = utils_datasets_dir / dataset_file.name
+                    dataset_file = candidate
+                
+                if not dataset_file.exists():
+                    logger.warning(f"Dataset file not found: {dataset_path} (resolved: {dataset_file})")
+                else:
+                    logger.info(f"Found dataset file: {dataset_file}")
+                validated_datasets.append(str(dataset_file))
+            request.local_datasets = validated_datasets
+        
+        # Convert request to kwargs for backend
+        training_kwargs = {
+            "model_name": request.model_name,
+            "training_type": request.training_type,
+            "hf_token": request.hf_token or "",
+            "load_in_4bit": request.load_in_4bit,
+            "max_seq_length": request.max_seq_length,
+            "hf_dataset": request.hf_dataset or "",
+            "local_datasets": request.local_datasets,
+            "format_type": request.format_type,
+            "num_epochs": request.num_epochs,
+            "learning_rate": request.learning_rate,
+            "batch_size": request.batch_size,
+            "gradient_accumulation_steps": request.gradient_accumulation_steps,
+            "warmup_steps": request.warmup_steps,
+            "warmup_ratio": request.warmup_ratio,
+            "max_steps": request.max_steps,
+            "save_steps": request.save_steps,
+            "weight_decay": request.weight_decay,
+            "random_seed": request.random_seed,
+            "packing": request.packing,
+            "use_lora": request.use_lora,
+            "lora_r": request.lora_r,
+            "lora_alpha": request.lora_alpha,
+            "lora_dropout": request.lora_dropout,
+            "target_modules": request.target_modules if request.target_modules else None,
+            "gradient_checkpointing": request.gradient_checkpointing.strip() if request.gradient_checkpointing and request.gradient_checkpointing.strip() else "unsloth",
+            "use_rslora": request.use_rslora,
+            "use_loftq": request.use_loftq,
+            "train_on_completions": request.train_on_completions,
+            "finetune_vision_layers": request.finetune_vision_layers,
+            "finetune_language_layers": request.finetune_language_layers,
+            "finetune_attention_modules": request.finetune_attention_modules,
+            "finetune_mlp_modules": request.finetune_mlp_modules,
+            "enable_wandb": request.enable_wandb,
+            "wandb_token": request.wandb_token or "",
+            "wandb_project": request.wandb_project or "",
+            "enable_tensorboard": request.enable_tensorboard,
+            "tensorboard_dir": request.tensorboard_dir or "",
+            "optim": request.optim,
+            "lr_scheduler_type": request.lr_scheduler_type,
+        }
+        
+        # Generate job ID
+        job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        
+        # Set initial "preparing" state
+        try:
+            backend.trainer._update_progress(
+                status_message="Initializing training...",
+                is_training=False
+            )
+        except:
+            pass
+        
+        def run_training():
+            try:
+                logger.info(f"Starting training job {job_id} with model {request.model_name}")
+                
+                # Update status to show we're loading model
+                try:
+                    backend.trainer._update_progress(status_message="Loading model...")
+                except Exception as e:
+                    logger.error(f"Error updating progress: {e}")
+                
+                # Consume the generator - this actually runs the training
+                update_count = 0
+                for update_tuple in backend.start_training(**training_kwargs):
+                    update_count += 1
+                    if update_count % 10 == 0:
+                        logger.info(f"Training progress update #{update_count}")
+                
+                logger.info(f"Training job {job_id} completed successfully")
+                
+            except Exception as e:
+                logger.error(f"Training error in job {job_id}: {e}", exc_info=True)
+                try:
+                    backend.trainer._update_progress(
+                        error=str(e),
+                        is_training=False
+                    )
+                except Exception as update_error:
+                    logger.error(f"Failed to update progress: {update_error}")
+        
+        # Start training in a daemon thread
+        training_thread = threading.Thread(target=run_training, daemon=True, name=f"Training-{job_id}")
+        training_thread.start()
+        
+        # Store thread reference for status checking
+        backend._training_thread = training_thread
+        
+        # Give it a moment to start
+        import time
+        time.sleep(0.5)
+        
+        # Verify training thread is alive
+        if not training_thread.is_alive():
+            logger.warning(f"Training thread died immediately for job {job_id}")
+            return TrainingStartResponse(
+                status="error",
+                message="Training thread failed to start. Check server logs for details.",
+                error="Thread not alive"
+            )
+        
+        return TrainingStartResponse(
+            status="started",
+            job_id=job_id,
+            message="Training job started successfully"
+        )
+        
+    except Exception as e:
+        logger.error(f"Error starting training: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to start training: {str(e)}"
+        )
+
+
+@router.post("/stop")
+async def stop_training():
+    """
+    Stop the currently running training job.
+    """
+    try:
+        backend = get_training_backend()
+        
+        if not backend.is_training_active():
+            return {
+                "status": "idle",
+                "message": "No training job is currently running"
+            }
+        
+        # Call backend stop method
+        backend.stop_training()
+        
+        return {
+            "status": "stopped",
+            "message": "Training job stopped successfully"
+        }
+        
+    except Exception as e:
+        logger.error(f"Error stopping training: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to stop training: {str(e)}"
+        )
+
+
+@router.get("/status")
+async def get_training_status():
+    """
+    Get the current training status.
+    """
+    try:
+        backend = get_training_backend()
+        
+        # Check if training is active
+        is_active = backend.is_training_active()
+        
+        # Check if there's a training thread running (preparation phase)
+        has_thread = hasattr(backend, '_training_thread') and backend._training_thread and backend._training_thread.is_alive()
+        
+        # Get progress info
+        try:
+            progress = backend.trainer.get_training_progress()
+            status_message = progress.status_message or "Ready to train"
+        except:
+            progress = None
+            status_message = "Unknown"
+        
+        if is_active:
+            # Actual training is running
+            trainer = backend.trainer
+            current_step = getattr(trainer.training_progress, 'step', None) or (progress.step if progress else None)
+            total_steps = getattr(trainer.training_progress, 'total_steps', None) or (progress.total_steps if progress else None)
+            
+            return TrainingStatusResponse(
+                status="training",
+                is_active=True,
+                message=status_message or "Training is in progress",
+                current_step=current_step,
+                total_steps=total_steps
+            )
+        elif has_thread or (progress and status_message and any(keyword in status_message.lower() for keyword in ["loading", "preparing", "initializing"])):
+            # Training thread is running but not yet in active training phase
+            return TrainingStatusResponse(
+                status="preparing",
+                is_active=False,
+                message=status_message or "Preparing training...",
+                current_step=None,
+                total_steps=None
+            )
+        else:
+            return TrainingStatusResponse(
+                status="idle",
+                is_active=False,
+                message="No training job is currently running"
+            )
+            
+    except Exception as e:
+        logger.error(f"Error getting training status: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to get training status: {str(e)}"
+        )
+
+
+@router.get("/metrics")
+async def get_training_metrics():
+    """
+    Get training metrics (loss, learning rate, steps).
+    """
+    try:
+        backend = get_training_backend()
+        
+        # Get metrics from backend
+        loss_history = backend.loss_history
+        lr_history = backend.lr_history
+        step_history = backend.step_history
+        
+        # Get current values
+        current_loss = loss_history[-1] if loss_history else None
+        current_lr = lr_history[-1] if lr_history else None
+        current_step = step_history[-1] if step_history else None
+        
+        return TrainingMetricsResponse(
+            loss_history=loss_history,
+            lr_history=lr_history,
+            step_history=step_history,
+            current_loss=current_loss,
+            current_lr=current_lr,
+            current_step=current_step
+        )
+        
+    except Exception as e:
+        logger.error(f"Error getting training metrics: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to get training metrics: {str(e)}"
+        )
+
+
+@router.get("/progress")
+async def stream_training_progress():
+    """
+    Stream training progress updates using Server-Sent Events (SSE).
+    
+    This endpoint provides real-time updates on training progress.
+    """
+    async def event_generator():
+        backend = get_training_backend()
+        
+        # Send initial status
+        is_active = backend.is_training_active()
+        initial_message = 'Connecting...' if is_active else 'No training in progress'
+        yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message=initial_message).model_dump_json()}\n\n"
+        
+        # If not active, check if there's any history
+        if not is_active:
+            if backend.step_history:
+                # Training completed - send final metrics
+                final_step = backend.step_history[-1]
+                final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
+                final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
+                yield f"data: {TrainingProgressResponse(step=final_step, loss=final_loss, learning_rate=final_lr, status_message='Training completed').model_dump_json()}\n\n"
+            else:
+                yield f"data: {TrainingProgressResponse(step=-1, loss=0.0, learning_rate=0.0, status_message='No training in progress').model_dump_json()}\n\n"
+            return
+        
+        # Poll for updates while training is active
+        last_step = -1
+        no_update_count = 0
+        max_no_updates = 300  # Timeout after 5 minutes
+        
+        while backend.is_training_active():
+            try:
+                # Get current metrics
+                if backend.step_history:
+                    current_step = backend.step_history[-1]
+                    current_loss = backend.loss_history[-1] if backend.loss_history else 0.0
+                    current_lr = backend.lr_history[-1] if backend.lr_history else 0.0
+                    
+                    # Only send if step changed
+                    if current_step != last_step:
+                        progress = TrainingProgressResponse(
+                            step=current_step,
+                            loss=current_loss,
+                            learning_rate=current_lr,
+                            status_message=f"Training step {current_step}"
+                        )
+                        yield f"data: {progress.model_dump_json()}\n\n"
+                        last_step = current_step
+                        no_update_count = 0
+                    else:
+                        no_update_count += 1
+                        # Send heartbeat every 10 seconds
+                        if no_update_count % 10 == 0:
+                            progress = TrainingProgressResponse(
+                                step=current_step,
+                                loss=current_loss,
+                                learning_rate=current_lr,
+                                status_message=f"Training step {current_step} (waiting for next update...)"
+                            )
+                            yield f"data: {progress.model_dump_json()}\n\n"
+                else:
+                    # No steps yet, but training is active
+                    no_update_count += 1
+                    if no_update_count % 5 == 0:
+                        yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message='Preparing training...').model_dump_json()}\n\n"
+                
+                # Timeout check
+                if no_update_count > max_no_updates:
+                    logger.warning("Progress stream timeout - no updates received")
+                    yield f"data: {TrainingProgressResponse(step=last_step, loss=0.0, learning_rate=0.0, status_message='Progress timeout - training may have stopped').model_dump_json()}\n\n"
+                    break
+                
+                await asyncio.sleep(1)  # Poll every second
+                
+            except Exception as e:
+                logger.error(f"Error in progress stream: {e}", exc_info=True)
+                yield f"data: {TrainingProgressResponse(step=0, loss=0.0, learning_rate=0.0, status_message=f'Error: {str(e)}').model_dump_json()}\n\n"
+                break
+        
+        # Send final status
+        final_step = backend.step_history[-1] if backend.step_history else last_step
+        final_loss = backend.loss_history[-1] if backend.loss_history else 0.0
+        final_lr = backend.lr_history[-1] if backend.lr_history else 0.0
+        yield f"data: {TrainingProgressResponse(step=final_step, loss=final_loss, learning_rate=final_lr, status_message='Training completed').model_dump_json()}\n\n"
+    
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+        }
+    )
+

backend/run.py

@@ -0,0 +1,94 @@
+"""
+Run script for Unsloth UI Backend.
+Works independently and can be moved to any directory.
+"""
+import sys
+from pathlib import Path
+
+# Add the backend directory to Python path
+backend_dir = Path(__file__).parent
+if str(backend_dir) not in sys.path:
+    sys.path.insert(0, str(backend_dir))
+
+
+def run_server(
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    frontend_path: Path = None,
+    silent: bool = False,
+):
+    """
+    Start the FastAPI server.
+
+    Args:
+        host: Host to bind to
+        port: Port to bind to
+        frontend_path: Path to frontend build directory (optional)
+        silent: Suppress startup messages
+    """
+    import nest_asyncio
+
+    nest_asyncio.apply()
+
+    import asyncio
+    from threading import Thread
+    import time
+    import uvicorn
+
+    from main import app, setup_frontend
+
+    # Setup frontend if path provided
+    if frontend_path:
+        if setup_frontend(app, frontend_path):
+            if not silent:
+                print(f"✅ Frontend loaded from {frontend_path}")
+        else:
+            if not silent:
+                print(f"⚠️ Frontend not found at {frontend_path}")
+
+    # Run server
+    def _run():
+        config = uvicorn.Config(app, host=host, port=port, log_level="info")
+        server = uvicorn.Server(config)
+        asyncio.run(server.serve())
+
+    thread = Thread(target=_run, daemon=True)
+    thread.start()
+    time.sleep(3)
+
+    if not silent:
+        print("")
+        print("=" * 50)
+        print(f"🦥 Unsloth UI Backend is running on port {port}")
+        print(f"   API: http://{host}:{port}/api")
+        print(f"   Health: http://{host}:{port}/api/health")
+        print("=" * 50)
+
+    return app
+
+
+# For direct execution
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run Unsloth UI Backend server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
+    parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
+    parser.add_argument(
+        "--frontend", type=str, default=None, help="Path to frontend build"
+    )
+    parser.add_argument("--silent", action="store_true", help="Suppress output")
+
+    args = parser.parse_args()
+
+    frontend_path = Path(args.frontend) if args.frontend else None
+    run_server(
+        host=args.host, port=args.port, frontend_path=frontend_path, silent=args.silent
+    )
+
+    # Keep running
+    import time
+
+    while True:
+        time.sleep(1)
+