import os
import gc
import logging
import importlib
from enum import Enum
import traceback
from typing import Optional, Any, Dict, List, Tuple, Union, Type, Set
import psutil
from collections import OrderedDict
import time
import sys
import threading

import torch
import diffusers
import numpy as np
import asyncio

from diffusers import (
    DiffusionPipeline,
    FluxInpaintPipeline,
    FluxPipeline,
)

from diffusers.loaders import FromSingleFileMixin
from huggingface_hub.constants import HF_HUB_CACHE
from huggingface_hub.file_download import repo_folder_name
from huggingface_hub.utils import EntryNotFoundError
from accelerate import utils as accelerate_utils

from .utils.config import get_config, get_environment
from .utils.globals import (
    # get_hf_model_manager,
    get_architectures,
    get_available_torch_device,
    set_available_torch_device,
    get_model_downloader,
)
from .utils.model_downloader import ModelSource
from .utils.load_models import load_state_dict_from_file
from .utils.base_types.config import PipelineConfig
from .utils import diffusers_fix

from ..model_interface import ModelManagementInterface, DownloaderType

from .utils.flashpack_loader import FlashPackLoader

# Configure logging
logger = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=sys.stdout,
)

# Constants
VRAM_SAFETY_MARGIN_GB = 3.5
DEFAULT_MAX_VRAM_BUFFER_GB = 2.0
RAM_SAFETY_MARGIN_GB = 10.0


# Keys correspond to diffusers pipeline classes
MODEL_COMPONENTS = {
    "FluxPipeline": [
        "vae",
        "transformer",
        "text_encoder",
        "text_encoder_2",
        "scheduler",
        "tokenizer",
        "tokenizer_2",
    ],
    "StableDiffusionXLPipeline": [
        "vae",
        "unet",
        "text_encoder",
        "text_encoder_2",
        "scheduler",
        "tokenizer",
        "tokenizer_2",
    ],
    "StableDiffusionPipeline": [
        "vae",
        "unet",
        "text_encoder",
        "scheduler",
        "tokenizer",
    ],
}


def get_pipeline_class(
    class_name: Union[str, Tuple[str, str]],
) -> Tuple[Type[DiffusionPipeline], Optional[str]]:
    """Get the appropriate pipeline class based on class_name configuration.
    Learn about custom-pipelines here: https://huggingface.co/docs/diffusers/v0.6.0/en/using-diffusers/custom_pipelines
    Note that from_single_file does not support custom pipelines, only from_pretrained does.

    Args:
        class_name: Either a string naming a diffusers class, or a tuple of (package, class)

    Returns:
        Tuple of (Pipeline class to use, custom_pipeline to include as a kwarg to the pipeline)
    """
    # class_name is in the form of [package, class]
    if isinstance(class_name, tuple):
        # Load from custom package
        package, cls = class_name
        module = importlib.import_module(package)
        return (getattr(module, cls), None)

    # Try loading class_name as a diffusers class
    try:
        pipeline_class = getattr(importlib.import_module("diffusers"), class_name)
        if not issubclass(pipeline_class, DiffusionPipeline):
            print("custompipeline2", class_name)
            raise TypeError(f"{class_name} does not inherit from DiffusionPipeline")
        return (pipeline_class, None)

    except (ImportError, AttributeError):
        # Assume the class name is the name of a custom pipeline
        print("custompipeline1", class_name)
        return (DiffusionPipeline, class_name)


class LRUCache:
    """
    Least Recently Used (LRU) Cache for tracking model usage.

    Maintains separate tracking for GPU and CPU cached models using timestamps
    to determine usage patterns and inform memory management decisions.
    """

    def __init__(self):
        """Initialize empty GPU and CPU caches."""
        self.gpu_cache: OrderedDict = OrderedDict()  # model_id -> last_used_timestamp
        self.cpu_cache: OrderedDict = OrderedDict()  # model_id -> last_used_timestamp

    def access(self, model_id: str, cache_type: str = "gpu") -> None:
        """
        Record access of a model, updating its position in the LRU cache.

        Args:
            model_id: Unique identifier for the model
            cache_type: Type of cache to update ("gpu" or "cpu")
        """
        cache = self.gpu_cache if cache_type == "gpu" else self.cpu_cache
        cache.pop(model_id, None)  # Remove if exists
        cache[model_id] = time.time()  # Add to end (most recently used)

    def remove(self, model_id: str, cache_type: str = "gpu") -> None:
        """
        Remove a model from cache tracking.

        Args:
            model_id: Unique identifier for the model to remove
            cache_type: Type of cache to remove from ("gpu" or "cpu")
        """
        cache = self.gpu_cache if cache_type == "gpu" else self.cpu_cache
        cache.pop(model_id, None)

    def get_lru_models(
        self, cache_type: str = "gpu", count: Optional[int] = None
    ) -> List[str]:
        """
        Get least recently used models.

        Args:
            cache_type: Type of cache to query ("gpu" or "cpu")
            count: Optional number of models to return. If None, returns all models

        Returns:
            List of model IDs ordered by least recently used first
        """
        cache = self.gpu_cache if cache_type == "gpu" else self.cpu_cache
        model_list = list(cache.keys())
        return model_list if count is None else model_list[:count]


class DefaultModelManager:
    """
    Manages loading, unloading and memory allocation of machine learning models.

    Handles dynamic movement of models between GPU and CPU memory based on
    available resources and usage patterns. Implements optimization strategies
    for efficient memory usage and model performance.
    """

    def __init__(self):
        """Initialize the model memory manager with empty states and default configurations."""
        # Model tracking
        self.current_model: Optional[str] = None
        self.base_pipelines: Dict[str, DiffusionPipeline] = {}
        self.loaded_models: Dict[str, DiffusionPipeline] = {}
        self.cpu_models: Dict[str, DiffusionPipeline] = {}
        self.model_sizes: Dict[str, float] = {}
        self.model_types: Dict[str, torch.dtype] = {}
        self.loaded_model: Optional[DiffusionPipeline] = None

        self.pipeline_locks: Dict[str, threading.Lock] = {}

        # Memory tracking
        self.vram_usage: float = 0
        self.ram_usage: float = 0
        self.max_vram: float = self._get_total_vram() - DEFAULT_MAX_VRAM_BUFFER_GB
        self.system_ram: float = psutil.virtual_memory().total / (1024**3)

        # State flags
        self.is_in_device: bool = False
        self.is_startup_load: bool = False
        self.environment: str = os.getenv("ENVIRONMENT", "dev")

        # Managers and caches
        self.model_downloader = get_model_downloader()
        self.cache_dir = HF_HUB_CACHE
        self.lru_cache = LRUCache()
        self.flashpack_loader = FlashPackLoader()

        self.allowed_model_ids: Optional[Set[str]] = None
        self.supports_all_known_models: bool = False # Default to explicit list
        self._download_lock = asyncio.Lock() # Lock for ensuring only one prioritized download happens at a time

        device = "cuda" if torch.cuda.is_available() else "cpu"
        set_available_torch_device(device)


    async def process_supported_models_config(
        self, supported_model_ids: List[str], downloader_instance: Optional[DownloaderType] = None # Can ignore if using self.model_downloader
        ) -> None:
        logger.info(f"DMM: Processing scheduler-provided supported_model_ids: {supported_model_ids}")
        if len(supported_model_ids) == 1 and supported_model_ids[0] == "*": # "*" means all
            self.supports_all_known_models = True
            self.allowed_model_ids = None 
            logger.info("DMM: Configured to support ALL known models (on-demand) due to '*' from scheduler.")
        elif not supported_model_ids: # Empty list means no specific dynamic models for this worker config
            self.supports_all_known_models = False
            self.allowed_model_ids = set() 
            logger.info("DMM: Configured with an empty supported model list from scheduler (supports no dynamic models).")
        else: # Specific list
            self.supports_all_known_models = False
            self.allowed_model_ids = set(supported_model_ids)
            logger.info(f"DMM: Configured to support specific models: {self.allowed_model_ids}")

    async def load_model_into_vram(self, model_id: str) -> bool:
        logger.info(f"DefaultModelManager: Request to load '{model_id}' into VRAM.")
        pipeline = await self.load(model_id) 
        if pipeline:
            self.base_pipelines[model_id] = pipeline
            self.loaded_model = pipeline
            self.current_model = model_id
            self.loaded_models[model_id] = pipeline # this is for backward compatibility (will be removed in the future)
            return True
        return False

    async def get_active_pipeline(self, model_id: str) -> Optional[Any]:
        logger.info(f"DefaultModelManager: Request to get active pipeline for '{model_id}'.")
        base_pipeline = await self.load(model_id)
        if not base_pipeline:
            logger.error(f"Failed to load base pipeline for {model_id}")
            return None
        
        # Create a task-specific pipeline instance with fresh scheduler state
        task_pipeline = self.get_pipeline_for_task(model_id, "worker_task")
        if not task_pipeline:
            logger.error(f"Failed to create thread-safe pipeline for {model_id}")
            return None
        
        return task_pipeline

    def get_vram_loaded_models(self) -> List[str]:
        # This should return a list of model_ids that are currently in VRAM
        return list[str](self.loaded_models.keys())
    
    # def get_pipeline_for_task(self, model_id: str, run_id: str = None) -> Optional[DiffusionPipeline]:
    #     """
    #     Create a task-specific pipeline instance with fresh scheduler state.
    #     This prevents thread safety issues in concurrent executions.
        
    #     This is the KEY FIX for the IndexError: index 29 is out of bounds issue.
    #     """
    #     try:
    #         # Get the base pipeline (shared components)
    #         base_pipeline = self.get_base_pipeline(model_id)
    #         if not base_pipeline:
    #             logger.error(f"No base pipeline found for model {model_id}")
    #             return None

    #         # Create a fresh scheduler for this task to prevent state corruption
    #         fresh_scheduler = base_pipeline.scheduler.from_config(
    #             base_pipeline.scheduler.config
    #         )
            
    #         # Create task-specific pipeline with shared components but fresh scheduler
    #         task_pipeline = base_pipeline.__class__.from_pipe(
    #             base_pipeline,
    #             scheduler=fresh_scheduler
    #         )
            
    #         logger.info(f"[ModelManager] Created thread-safe pipeline for task {run_id} model {model_id}")
    #         return task_pipeline
            
    #     except Exception as e:
    #         logger.error(f"Error creating task-specific pipeline for {model_id}: {e}")
    #         return None

    def get_pipeline_for_task(self, model_id: str, run_id: str = None) -> Optional[DiffusionPipeline]:
        """
        Return the base pipeline directly - no recreation needed.
        Scheduler state issues are rare with proper usage.
        """
        base_pipeline = self.get_base_pipeline(model_id)
        if not base_pipeline:
            logger.error(f"No base pipeline found for model {model_id}")
            return None
        
        # Return the base pipeline directly - it's thread-safe enough for inference
        logger.info(f"[ModelManager] Returning base pipeline for model {model_id}")
        return base_pipeline
        
    def get_base_pipeline(self, model_id: str) -> Optional[DiffusionPipeline]:
        """
        Get the base pipeline for component sharing.
        This contains the loaded model components but should NOT be used directly for inference.
        """
        # Check base pipelines first
        if model_id in self.base_pipelines:
            return self.base_pipelines[model_id]
        
        # Check loaded models (for backward compatibility)
        if model_id in self.loaded_models:
            return self.loaded_models[model_id]
            
        # Check CPU models
        if model_id in self.cpu_models:
            return self.cpu_models[model_id]
            
        return None


    # If the scheduler is not set in `pipeline_defs`, then we'll rely on diffusers to pick a default
    # scheduler.
    def _setup_scheduler(self, pipeline: DiffusionPipeline, model_id: str) -> None:
        """Setup scheduler from component config"""
        config = get_config()
        model_config = config.pipeline_defs.get(model_id)
        if not model_config:
            logger.error(f"Model {model_id} not found in configuration")
            return
        components = model_config.get("components", {})
        scheduler_config = components.get("scheduler", {}) if components else {}
        if not scheduler_config:
            return

        scheduler_class = scheduler_config.get("class_name")
        scheduler_kwargs = scheduler_config.get("kwargs", {})

        try:
            INVALID_SCHEDULER_CONFIG_KEYS = {
            'do_classifier_free_guidance',
            'guidance_scale',
            }

            # Only pass kwargs that are valid for scheduler initialization
            filtered_kwargs = {
                k: v for k, v in scheduler_kwargs.items() 
                if k.lower() not in {key.lower() for key in INVALID_SCHEDULER_CONFIG_KEYS}
            }
            
            if len(scheduler_kwargs) != len(filtered_kwargs):
                removed = set(scheduler_kwargs.keys()) - set(filtered_kwargs.keys())
                logger.warning(f"Removed invalid scheduler config keys for {model_id}: {removed}")  

            new_scheduler = getattr(diffusers, scheduler_class).from_config(
                pipeline.scheduler.config, **filtered_kwargs
                )
                
            pipeline.scheduler = new_scheduler

            print(f"Successfully set scheduler to {scheduler_class}")
        except Exception as e:
            logger.error(f"Error setting scheduler for {model_id}: {e}")

    def _get_memory_info(self) -> Tuple[float, float]:
        """
        Get current memory availability.

        Returns:
            Tuple containing:
            - Available RAM in GB
            - Available VRAM in GB
        """
        ram = self._get_available_ram()
        vram = self._get_available_vram()
        return ram, vram

    def _get_available_ram(self) -> float:
        """
        Get available system RAM in GB.

        Returns:
            Available RAM in GB
        """
        return psutil.virtual_memory().available / (1024**3)


    def _get_total_vram(self) -> float:
        """
        Get total VRAM available on the system.

        Returns:
            Total VRAM in GB, or 0 if no CUDA device is available
        """
        if torch.cuda.is_available():
            return torch.cuda.get_device_properties(0).total_memory / (1024**3)
        return 0

    def _get_available_vram(self) -> float:
        """
        Get currently available VRAM.

        Returns:
            Available VRAM in GB
        """
        if torch.cuda.is_available():
            available_vram_gb = (
                torch.cuda.get_device_properties(0).total_memory
                - torch.cuda.memory_allocated()
            ) / (1024**3)
            logger.debug(f"Available VRAM: {available_vram_gb:.2f} GB")
            return available_vram_gb
        return 100
    
    def _need_optimization(self, model_size: float) -> bool:
        """
        Check if a model needs optimization to fit in GPU memory.

        Args:
            model_size: Size of the model in GB

        Returns:
            True if the model needs optimization, False otherwise
        """
        print(f"model_size: {model_size} GB")
        print(f"system_ram: {self.system_ram} GB")
        return model_size < self.system_ram
    
    def _can_fit_gpu(self, model_size: float) -> bool:
        """
        Check if a model can fit in GPU memory.

        Args:
            model_size: Size of the model in GB

        Returns:
            True if the model can fit in GPU memory, False otherwise
        """
        print(f"model_size: {model_size} GB")
        print(f"available_inference_memory: {self._get_available_vram() - VRAM_SAFETY_MARGIN_GB} GB")
        return (model_size <= (self._get_available_vram() - VRAM_SAFETY_MARGIN_GB))
    

    async def _ensure_model_files_on_disk(self, model_id: str, model_config_dict: Dict[str, Any]) -> bool:
        """
        Internal helper: Ensures a specific model and its components are downloaded.
        This is called by self.load() before attempting to load from disk.
        This method is now the primary point for on-demand downloads.
        Returns True if all necessary files are confirmed/downloaded, False otherwise.
        """
        logger.info(f"DMM._ensure_model_files_on_disk: Checking/downloading '{model_id}'...")
        all_parts_ok = True
        async with self.model_downloader as downloader: # Use context manager for session
            # 1. Main model source
            main_source_str = model_config_dict.get("source")
            if main_source_str:
                is_main_downloaded, _ = await downloader.is_downloaded(model_id, model_config=model_config_dict)
                if not is_main_downloaded:
                    logger.info(f"DMM: Main files for '{model_id}' (source: {main_source_str}) not on disk. Downloading...")
                    try:
                        await downloader.download_model(model_id, ModelSource(main_source_str))
                        is_main_downloaded, _ = await downloader.is_downloaded(model_id, model_config=model_config_dict) # Verify
                        if not is_main_downloaded:
                            logger.error(f"DMM: Download verification FAILED for main files of '{model_id}'.")
                            all_parts_ok = False
                        else:
                            logger.info(f"DMM: Main files for '{model_id}' downloaded and verified.")
                    except Exception as e:
                        logger.exception(f"DMM: Exception during download of main files for '{model_id}': {e}")
                        all_parts_ok = False
                else:
                    logger.info(f"DMM: Main files for '{model_id}' already on disk.")
            else:
                logger.error(f"DMM: No source found for model '{model_id}' in its config. Cannot download.")
                return False # Cannot proceed without a source

            # 2. Components (if main download was okay or not needed, and components exist)
            if all_parts_ok and model_config_dict.get("components") and isinstance(model_config_dict["components"], dict):
                for comp_name, comp_details in model_config_dict["components"].items():
                    comp_source_str = None
                    if isinstance(comp_details, dict): comp_source_str = comp_details.get("source")
                    elif isinstance(comp_details, str): comp_source_str = comp_details
                    
                    if comp_source_str:
                        comp_dl_id = f"{model_id}::{comp_name}"
                        # TODO: Need a robust way to check if component is downloaded via downloader.is_downloaded
                        # For now, we rely on downloader.download_model being idempotent.
                        logger.info(f"DMM: Ensuring component '{comp_name}' for '{model_id}' (source: {comp_source_str}) is downloaded...")
                        try:
                            await downloader.download_model(comp_dl_id, ModelSource(comp_source_str))
                            # Add verification for component if possible
                            logger.info(f"DMM: Component '{comp_name}' for '{model_id}' download attempt finished.")
                        except Exception as e:
                            logger.exception(f"DMM: Exception downloading component '{comp_name}' for '{model_id}': {e}")
                            all_parts_ok = False; break # Stop if a component fails
        return all_parts_ok
    

    async def load(
        self, model_id: str, gpu: Optional[int] = None, pipe_type: Optional[str] = None
    ) -> Optional[DiffusionPipeline]:
        """
        Load a model into memory, handling placement and optimization.

        This method implements the main model loading logic, including:
        - Checking if model is already loaded
        - Managing memory allocation between GPU and CPU
        - Applying optimizations as needed
        - Handling model movement between devices

        Args:
            model_id: Identifier for the model to load
            gpu: Optional GPU device number
            pipe_type: Optional pipeline type specification

        Returns:
            Loaded pipeline or None if loading failed
        """
        logger.info(f"Processing model {model_id}")

        if not self.supports_all_known_models:
            if self.allowed_model_ids is None: # Should have been set if not supports_all
                 logger.error(f"DMM.load: Inconsistent state - worker not configured for all models, but allowed_model_ids is None. Denying load for '{model_id}'.")
                 return None
            if model_id not in self.allowed_model_ids:
                logger.error(f"DMM.load: Model '{model_id}' is NOT in the allowed list for this worker instance. Allowed: {self.allowed_model_ids}. Denying load.")
                return None
            
        # If self.supports_all_known_models is True, any model_id is principally allowed, contingent on it being defined in the global pipeline_defs.

        logger.debug(f"DMM.load: Model '{model_id}' is permitted for this worker instance.")

        # Check existing loaded models
        if model_id in self.loaded_models:
            logger.info(f"Model {model_id} already loaded in GPU")
            self.lru_cache.access(model_id, "gpu")
            return self.loaded_models[model_id]

        # CPU-offloaded models
        if model_id in self.cpu_models:
            logger.info(f"Model {model_id} is CPU-offloaded and ready")
            self.lru_cache.access(model_id, "cpu")
            return self.cpu_models[model_id]
        
        
        # On-Demand model download
        cfg = get_config()
        if not cfg or not cfg.pipeline_defs:
            logger.error("DefaultModelManager: Global config or pipeline_defs not loaded. Cannot download models.")
            return None
        
        model_config_dict = cfg.pipeline_defs.get(model_id)
        if not model_config_dict:
            logger.error(f"Model {model_id} not found in configuration")
            return None
        
        async with self._download_lock: # Ensure only one download process at a time for this MMM instance
            if not await self._ensure_model_files_on_disk(model_id, model_config_dict):
                logger.error(f"Failed to ensure model files for {model_id} are on disk")
                return None

        # Load new model
        return await self._load_new_model(model_id, gpu, pipe_type)


    async def _load_new_model(
        self, model_id: str, gpu: Optional[int] = None, pipe_type: Optional[str] = None
    ) -> Optional[DiffusionPipeline]:
        """
        Load a new model that isn't currently in memory.

        Handles the complete loading process including:
        - Configuration validation
        - Memory allocation
        - Model loading
        - Optimization application

        Args:
            model_id: Identifier for the model
            gpu: Optional GPU device number
            pipe_type: Optional pipeline type specification (deprecated, use class_name in config)

        Returns:
            Loaded pipeline or None if loading failed
        """
        try:
            config = get_config()
            model_config = config.pipeline_defs.get(model_id)

            if not model_config:
                logger.error(f"Model {model_id} not found in configuration")
                return None

            # Prepare memory
            estimated_size = await self._get_model_size(model_config, model_id) + 1.0 # add 1GB for safety margin (model inference overhead in memory)

            print(f"estimated_size for model {model_id}: {estimated_size} GB")

            # try direct gpu load
            if self._can_fit_gpu(estimated_size):
                pipeline = await self._load_model_by_source(model_id, model_config)
                if pipeline is None:
                    return None
                
                self._setup_scheduler(pipeline, model_id)
                if self._move_model_to_gpu(pipeline, model_id):
                    self.loaded_models[model_id] = pipeline
                    self.model_sizes[model_id] = estimated_size
                    self.vram_usage += estimated_size
                    self.lru_cache.access(model_id, "gpu")
                    # Background prefetch other deployment models
                    if hasattr(self, 'flashpack_loader') and self.flashpack_loader.local_cache:
                        if self.allowed_model_ids and len(self.allowed_model_ids) > 1:
                            sources = {}
                            for mid in self.allowed_model_ids:
                                if mid != model_id:
                                    model_cfg = config.pipeline_defs.get(mid)
                                    if model_cfg:
                                        sources[mid] = model_cfg.get("source") if isinstance(model_cfg, dict) else model_cfg.source
                            if sources:
                                logger.info(f"🔄 Starting background prefetch for {len(sources)} other models")
                                asyncio.create_task(
                                    self.flashpack_loader.prefetch_deployment_models(
                                        list(self.allowed_model_ids), sources, exclude_model_id=model_id
                                    )
                                )
                    return pipeline
                else:
                    logger.error(f"Failed to move {model_id} to GPU")
                    return None
                
            # TODO: Change this ugly condition
            if estimated_size <= self.max_vram - (VRAM_SAFETY_MARGIN_GB - DEFAULT_MAX_VRAM_BUFFER_GB) and (len(self.loaded_models) > 0 or len(self.cpu_models) > 0):
                # if not enough space, try to make space
                self._free_space_for_model(estimated_size)
                if self._can_fit_gpu(estimated_size):
                    pipeline = await self._load_model_by_source(model_id, model_config)
                    if pipeline is None:
                        return None
                    
                    self._setup_scheduler(pipeline, model_id)
                    if self._move_model_to_gpu(pipeline, model_id):
                        self.loaded_models[model_id] = pipeline
                        self.model_sizes[model_id] = estimated_size
                        self.vram_usage += estimated_size
                        self.lru_cache.access(model_id, "gpu")
                        # Background prefetch other deployment models
                        if hasattr(self, 'flashpack_loader') and self.flashpack_loader.local_cache:
                            if self.allowed_model_ids and len(self.allowed_model_ids) > 1:
                                sources = {}
                                for mid in self.allowed_model_ids:
                                    if mid != model_id:
                                        model_cfg = config.pipeline_defs.get(mid)
                                        if model_cfg:
                                            sources[mid] = model_cfg.get("source") if isinstance(model_cfg, dict) else model_cfg.source
                                if sources:
                                    logger.info(f"🔄 Starting background prefetch for {len(sources)} other models")
                                    asyncio.create_task(
                                        self.flashpack_loader.prefetch_deployment_models(
                                            list(self.allowed_model_ids), sources, exclude_model_id=model_id
                                        )
                                    )
                        return pipeline
                    else:
                        logger.error(f"Failed to move {model_id} to GPU")
                        return None
            
            # if still not enough VRAM, apply optimization (CPU offload)
            if self.environment != "prod":
                if self._need_optimization(estimated_size):
                    logger.info("Unloading all models for large model loading")

                    for model_id_to_unload in list(self.loaded_models.keys()):
                        self._unload_model_for_space(model_id_to_unload, self.model_sizes[model_id_to_unload], "gpu")
                    for model_id_to_unload in list(self.cpu_models.keys()):
                        self._unload_model_for_space(model_id_to_unload, self.model_sizes[model_id_to_unload], "cpu")

                    pipeline = await self._load_model_by_source(model_id, model_config)
                    if pipeline is None:
                        return None
                    
                    self._setup_scheduler(pipeline, model_id)
                    logger.info(f"Applying optimizations for {model_id}")
                    self.apply_optimizations(pipeline, model_id, True)

                    self.cpu_models[model_id] = pipeline
                    self.model_sizes[model_id] = estimated_size
                    self.lru_cache.access(model_id, "cpu")
                    # Background prefetch other deployment models
                    if hasattr(self, 'flashpack_loader') and self.flashpack_loader.local_cache:
                        if self.allowed_model_ids and len(self.allowed_model_ids) > 1:
                            sources = {}
                            for mid in self.allowed_model_ids:
                                if mid != model_id:
                                    model_cfg = config.pipeline_defs.get(mid)
                                    if model_cfg:
                                        sources[mid] = model_cfg.get("source") if isinstance(model_cfg, dict) else model_cfg.source
                            if sources:
                                logger.info(f"🔄 Starting background prefetch for {len(sources)} other models")
                                asyncio.create_task(
                                    self.flashpack_loader.prefetch_deployment_models(
                                        list(self.allowed_model_ids), sources, exclude_model_id=model_id
                                    )
                                )
                    return pipeline
            
            logger.error(f"Insufficient memory to load model {model_id}")
            return None
        
        except Exception as e:
            logger.error("Error loading new model {}: {}".format(model_id, str(e)))
            traceback.print_exc()
            return None

    async def _load_model_by_source(
        self, model_id: str, model_config: PipelineConfig) -> Optional[DiffusionPipeline]:
        """
        Load a model based on its source configuration.

        Args:
            model_id: Model identifier
            model_config: Model configuration from PipelineConfig
            gpu: Optional GPU device number

        Returns:
            Loaded pipeline or None if loading failed
        """

        if isinstance(model_config, dict):
            source = model_config.get("source")
            class_name = model_config.get("class_name")
        else:
            source = model_config.source
            class_name = model_config.class_name

        prefix, path = source.split(":", 1)

        try:
            # === LOCAL CACHE + FLASHPACK ===
            # Try FlashPack first
            flashpack_path = self.flashpack_loader.get_flashpack_path(model_id, source)
            if flashpack_path:
                (pipeline_class, _) = get_pipeline_class(class_name)
                pipeline = await self.flashpack_loader.load_from_flashpack(
                    model_id, flashpack_path, pipeline_class
                )
                if pipeline:
                    return pipeline
                logger.warning(f"FlashPack loading failed for {model_id}, falling back")
            
            # No FlashPack - copy safetensors to local cache anyway
            if self.flashpack_loader.local_cache:
                local_path = await self.flashpack_loader.local_cache.ensure_local(
                    model_id, source, priority=True
                )
                if local_path:
                    logger.info(f"📂 Loading from local cache: {local_path}")
                    if local_path.suffix == ".safetensors" or local_path.is_file():
                        path = str(local_path)
                        prefix = "file"
                    else:
                        # It's a directory (HF model)
                        path = str(local_path)
                        prefix = "hf"
            
            # === STANDARD LOADING: Fallback to standard loading if FlashPack fails ===
            if prefix == "hf":
                is_downloaded, variant = await self.model_downloader.is_downloaded(
                    model_id,
                    model_config
                )

                if not is_downloaded:
                    logger.info(f"Model {model_id} not downloaded")
                    return None

                # Get model index and use as fallback in case class_name is unspecified
                if class_name is None or class_name == "":
                    model_index = await self.model_downloader.get_diffusers_multifolder_components(
                        path
                    )
                    if model_index and "_class_name" in model_index:
                        class_name = model_index["_class_name"]
                    else:
                        logger.error(f"Unknown diffusers class_name for {model_id}")
                        return None

                return await self.load_huggingface_model(
                    model_id, path, class_name, variant, model_config
                )
            elif prefix in ["file", "ct"]:
                return await self.load_single_file_model(
                    model_id, path, prefix, class_name
                )
            elif source.startswith(("http://", "https://")):
                # Handle Civitai/direct download models
                source_obj = ModelSource(source)
                is_downloaded = await self.model_downloader.is_downloaded(model_id)
                if not is_downloaded:
                    logger.info(f"Model {model_id} not downloaded")
                    return None

                cached_path = await self.model_downloader._get_cache_path(
                    model_id, source_obj
                )
                if not os.path.exists(cached_path):
                    logger.error(f"Cached model file not found at {cached_path}")
                    return None

                return await self.load_single_file_model(
                    model_id, cached_path, "file", class_name
                )
            else:
                logger.error(f"Unsupported model source prefix: {prefix}")
                return None
        except Exception as e:
            logger.error("Error loading model from source: {}".format(str(e)))
            return None
        
    def _free_space_for_model(self, model_size: float) -> None:
        available_vram = self._get_available_vram() - VRAM_SAFETY_MARGIN_GB
        if available_vram >= model_size:
            return
        
        space_needed = model_size - available_vram
        logger.info(f"Need to free {space_needed:.2f} GB of VRAM")

        gpu_models = [
            (mid, self.model_sizes[mid])
            for mid in self.lru_cache.get_lru_models("gpu")
        ]

        freed_space = 0
        for model_id_to_unload, size in gpu_models:
            if freed_space >= space_needed:
                break

            freed_space += self._unload_model_for_space(model_id_to_unload, size, "gpu")

            if self._get_available_vram() - VRAM_SAFETY_MARGIN_GB >= model_size:
                break


    def _unload_model_for_space(
        self, model_id: str, model_size: float, device: str
    ) -> float:
        """
        Completely unload a model and free up GPU/CPU memory.

        Args:
            model_id: Model identifier
            model_size: Size of the model in GB
            device: Device to unload from ("gpu" or "cpu")
        Returns:
            Amount of space freed in GB
        """
        logger.info(f"Unloading {model_id} from memory")

        try:
            if device == "gpu" and model_id in self.loaded_models:
                pipeline = self.loaded_models.pop(model_id)

                if hasattr(pipeline, "remove_all_hooks"):
                    logger.info(f"Removing all hooks for model {model_id}")
                    pipeline.remove_all_hooks()

                # # Move model to CPU first to clear CUDA memory
                # if hasattr(pipeline, "to"):
                #     pipeline.to("cpu", silence_dtype_warnings=True)

                # Explicitly delete model components
                for attr in [
                    "vae",
                    "unet",
                    "text_encoder",
                    "text_encoder_2",
                    "tokenizer",
                    "scheduler",
                    "transformer",
                    "tokenizer_2",
                    "text_encoder_3",
                    "tokenizer_3",
                ]:
                    if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None:
                        component = getattr(pipeline, attr)
                        delattr(pipeline, attr)
                        del component

                # Delete pipeline reference
                # del self.loaded_models[model_id]
                self.vram_usage -= model_size
                self.lru_cache.remove(model_id, "gpu")
                del pipeline

            if model_id in self.cpu_models:
                logger.info(f"Unloading {model_id} from CPU memory")
                cpu_pipeline = self.cpu_models.pop(model_id)

                if hasattr(cpu_pipeline, "remove_all_hooks"):
                    logger.info(f"Removing all hooks for CPU-offloaded model {model_id}")
                    cpu_pipeline.remove_all_hooks()

                self.lru_cache.remove(model_id, "cpu")
                del cpu_pipeline

            # Force garbage collection and memory clearing
            self.flush_memory()

            return model_size

        except Exception as e:
            logger.error(
                "Error during model unloading for {}: {}".format(model_id, str(e))
            )
            # Still try to clean up references even if error occurs
            self.loaded_models.pop(model_id, None)
            self.cpu_models.pop(model_id, None)
            self.vram_usage -= model_size
            self.lru_cache.remove(model_id, "gpu")
            self.lru_cache.remove(model_id, "cpu")
            self.flush_memory()
            return model_size

    def _move_model_to_gpu(self, pipeline: DiffusionPipeline, model_id: str) -> bool:
        """
        Safely move a model to GPU memory with proper dtype handling.

        Args:
            pipeline: The pipeline to move
            model_id: Identifier of the model

        Returns:
            Boolean indicating success of the operation
        """
        # Skip if it's OmniGen pipeline
        if pipeline.__class__.__name__ == "OmniGenPipeline":
            return True

        try:
            device = get_available_torch_device()
            print("Moving to device", device)

            # check if pippeline supports vae tiling and slicing
            if hasattr(pipeline, "enable_vae_tiling"):
                logger.info("Enabling vae tiling")
                pipeline.enable_vae_tiling()
            if hasattr(pipeline, "enable_vae_slicing"):
                logger.info("Enabling vae slicing")
                pipeline.enable_vae_slicing()

            pipeline = pipeline.to(device=device)

            print("Done with moving to GPU")

            self.flush_memory()
            return True

        except RuntimeError as e:
            if "out of memory" in str(e):
                logger.error(f"GPU out of memory while moving model {model_id} to GPU")
            else:
                logger.error(
                    "Runtime error moving model {} to GPU: {}".format(model_id, str(e))
                )
            self.flush_memory()
            return False

    async def load_huggingface_model(
        self,
        model_id: str,
        repo_id: str,
        class_name: Union[str, Tuple[str, str]],
        variant: Optional[str] = None,
        model_config: Optional[PipelineConfig] = None,
    ) -> Optional[DiffusionPipeline]:
        """
        Load a model from HuggingFace.

        Args:
            model_id: Model identifier
            repo_id: HuggingFace repository ID
            gpu: Optional GPU device number
            class_name: Pipeline class name or (package, class) tuple
            variant: Optional model variant
            model_config: Optional model configuration

        Returns:
            Loaded pipeline or None if loading failed
        """
        try:
            pipeline_kwargs = await self._prepare_pipeline_kwargs(model_config, variant)
            variant = None if variant == "" else variant
            # TO DO: make this more robust
            torch_dtype = (
                torch.bfloat16 if "flux" in model_id.lower() else torch.float16
            )

            # Get appropriate pipeline class
            (pipeline_class, custom_pipeline) = get_pipeline_class(class_name)
            print("custompipeline", custom_pipeline)

            if custom_pipeline is not None:
                pipeline_kwargs["custom_pipeline"] = custom_pipeline

            print(
                f"repo_id={repo_id},torch_dtype={torch_dtype},local_files_only=True,variant={variant},pipeline_kwargs={pipeline_kwargs},"
            )
            try:
                pipeline = await asyncio.to_thread(pipeline_class.from_pretrained,
                    repo_id,
                    torch_dtype=torch_dtype,
                    local_files_only=True,
                    variant=variant,
                    **pipeline_kwargs,
                )
            except EntryNotFoundError as e:
                print(f"Custom pipeline '{custom_pipeline}' not found: {e}")
                print("Falling back to the default pipeline...")
                del pipeline_kwargs["custom_pipeline"]
                pipeline = await asyncio.to_thread(pipeline_class.from_pretrained,
                    repo_id,
                    variant=variant,
                    torch_dtype=torch_dtype,
                    **pipeline_kwargs,
                )

            self.flush_memory()
            self.loaded_model = pipeline
            self.current_model = model_id
            logger.info(f"Model {model_id} loaded successfully")
            return pipeline

        except Exception as e:
            traceback.print_exc()
            logger.error("Failed to load model {}: {}".format(model_id, str(e)))
            return None

    async def _prepare_pipeline_kwargs(
        self, model_config: Optional[PipelineConfig], variant: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Prepare kwargs for pipeline initialization.

        Args:
            model_config: Model configuration from PipelineConfig

        Returns:
            Dictionary of pipeline initialization arguments
        """
        if not model_config:
            return {}

        pipeline_kwargs = {}
        try:
            if isinstance(model_config, dict):
                class_name = model_config.get("class_name")
                components = model_config.get("components")
                custom_pipeline = model_config.get("custom_pipeline")
            else:
                class_name = model_config.class_name
                components = model_config.components

            if custom_pipeline:
                pipeline_kwargs["custom_pipeline"] = custom_pipeline

            if components:
                for key, component in components.items():
                    main_model_source = model_config.get("source")
                    # Check if component and also id component has .source attribute
                    if component:
                        component_source = component.get("source", None)
                        if component_source:
                            pipeline_kwargs[key] = await self._prepare_component(
                                main_model_source, component, class_name, key, variant
                            )

            return pipeline_kwargs

        except Exception as e:
            logger.error("Error preparing pipeline kwargs: {}".format(str(e)))
            return {}

    async def _prepare_component(
        self,
        main_model_source: str,
        component: PipelineConfig,
        model_class_name: Optional[Union[str, Tuple[str, str]]],
        key: str,
        variant: Optional[str] = None,
    ) -> Any:
        """
        Prepare a model component based on its configuration.

        Args:
            component: Component configuration
            model_class_name: Class name of the parent model
            key: Component key

        Returns:
            Loaded component or None if loading failed
        """
        try:
            if isinstance(component, dict):
                source = component.get("source")
            else:
                source = component.source

            if not source.endswith((".safetensors", ".bin", ".ckpt", ".pt")):
                # check if the url has more than 2 forward slashes. If it does, the last one is the subfolder, the source is the first part
                # e.g. hf:cozy-creator/Flux.1-schnell-8bit/transformer this will be, source = hf:cozy-creator/Flux.1-schnell-8bit, subfolder = transformer
                if source.count("/") > 1:
                    repo_id = "/".join(source.split("/")[:-1])
                    subfolder = source.split("/")[-1]
                    return await self._load_diffusers_component(
                        main_model_source.replace("hf:", ""),
                        repo_id.replace("hf:", ""),
                        subfolder,
                        variant,
                    )
                else:
                    return await self._load_diffusers_component(
                        main_model_source.replace("hf:", ""),
                        source.replace("hf:", ""),
                        variant,
                    )
            else:
                return self._load_custom_component(source, model_class_name, key)
        except Exception as e:
            logger.error("Error preparing component {}: {}".format(key, str(e)))
            return None

    async def load_single_file_model(
        self,
        model_id: str,
        path: str,
        prefix: str,
        class_name: Optional[Union[str, Tuple[str, str]]] = None,
    ) -> Optional[DiffusionPipeline]:
        """
        Load a model from a single file.

        Args:
            model_id: Model identifier
            path: Path to model file
            prefix: Source prefix (file/ct)
            gpu: Optional GPU device number
            class_name: Pipeline class name or (package, class) tuple

        Returns:
            Loaded pipeline or None if loading failed
        """
        logger.info(f"Loading single file model {model_id}")

        # TO DO: we could try inferring the class using our old detect_model code here!
        if class_name is None or class_name == "":
            logger.error("Model class_name must be specified for single file models")
            return None

        (pipeline_class, custom_pipeline) = get_pipeline_class(class_name)

        try:
            print(f"Model path: {path}")
            if prefix != "file":
                model_path = self._get_model_path(path, prefix)
                if not os.path.exists(model_path):
                    logger.error(f"Model file not found: {model_path}")
                    return None
            else:
                model_path = path

            pipeline = await self._load_pipeline_from_file(
                pipeline_class, model_path, model_id, class_name
            )
            if pipeline:
                self.loaded_model = pipeline
                self.current_model = model_id
            return pipeline

        except Exception as e:
            logger.error("Error loading single file model: {}".format(str(e)))
            return None

    def _get_model_path(self, path: str, prefix: str) -> str:
        if prefix == "ct" and ("http" in path or "https" in path):
            path = path.split("/")[-1]
        return os.path.join(get_config().models_path, path)

    async def _load_pipeline_from_file(
        self,
        pipeline_class: Type[DiffusionPipeline],
        model_path: str,
        model_id: str,
        class_name: Optional[Union[str, Tuple[str, str]]],
    ) -> Optional[DiffusionPipeline]:
        """
        Load a pipeline from a file using appropriate loading method.

        Args:
            pipeline_class: Class to instantiate pipeline
            model_path: Path to model file
            model_id: Model identifier
            class_name: Pipeline class name or (package, class) tuple

        Returns:
            Loaded pipeline or None if loading failed
        """
        try:
            if issubclass(pipeline_class, FromSingleFileMixin):
                return self._load_from_single_file(pipeline_class, model_path, model_id)
            else:
                # Get the actual class name string if it's a tuple
                class_str = (
                    class_name[1] if isinstance(class_name, tuple) else class_name
                )
                return self._load_custom_architecture(
                    pipeline_class, model_path, class_str
                )
        except Exception as e:
            logger.error("Error loading pipeline from file: {}".format(str(e)))
            return None

    def _load_from_single_file(
        self, pipeline_class: Any, path: str, model_id: str
    ) -> DiffusionPipeline:
        """
        Load a model pipeline from a single file using the FromSingleFileMixin.

        Uses different torch datatypes based on model type:
        - bfloat16 for Flux models
        - float16 for other models

        Args:
            pipeline_class: The pipeline class to instantiate
            path: Path to the model file
            model_id: Model identifier (used to determine model type)

        Returns:
            Loaded pipeline instance

        Raises:
            Exception: If loading fails
        """
        try:
            # Determine appropriate dtype based on model type
            torch_dtype = torch.bfloat16 if "flux" in model_id.lower() else torch.float16

            # Load the model with appropriate dtype
            pipeline = pipeline_class.from_single_file(path, torch_dtype=torch_dtype)

            logger.info(f"Successfully loaded single file model {model_id}")
            return pipeline

        except Exception as e:
            logger.error("Error loading model from single file: {}".format(str(e)))
            raise

    def _load_custom_architecture(
        self, pipeline_class: Any, path: str, type: str
    ) -> DiffusionPipeline:
        """
        Load a model with custom architecture configuration.

        Handles loading of individual components and assembling them into
        a complete pipeline.

        Args:
            pipeline_class: The pipeline class to instantiate
            path: Path to the model file
            type: Model type (determines which components to load)

        Returns:
            Assembled pipeline instance

        Raises:
            Exception: If loading fails or architecture is not found
        """
        try:
            # Load the complete state dict
            state_dict = load_state_dict_from_file(path)

            # Create empty pipeline instance
            pipeline = pipeline_class()

            # Load each component specified for this model type
            for component_name in MODEL_COMPONENTS[type]:
                # Skip certain components that don't need loading
                if component_name in ["scheduler", "tokenizer", "tokenizer_2"]:
                    continue

                # Construct architecture key and get class
                arch_key = f"core_extension_1.{type}_{component_name}"
                architecture_class = get_architectures().get(arch_key)

                if not architecture_class:
                    logger.error(f"Architecture not found for {arch_key}")
                    continue

                try:
                    # Initialize and load the component
                    architecture = architecture_class()
                    architecture.load(state_dict)

                    # Set the component in the pipeline
                    setattr(pipeline, component_name, architecture.model)
                    logger.debug(f"Loaded component {component_name} for {type}")

                except Exception as component_error:
                    logger.error(
                        "Error loading component {}: {}".format(
                            component_name, str(component_error)
                        )
                    )
                    continue

            logger.info("Successfully loaded custom architecture model")
            return pipeline

        except Exception as e:
            logger.error("Error loading custom architecture: {}".format(str(e)))
            raise

    def apply_optimizations(
        self,
        pipeline: DiffusionPipeline,
        model_id: str,
        force_full_optimization: bool = False,
    ) -> None:
        """
        Apply memory and performance optimizations to a pipeline.
        Uses existing optimization functions when model exceeds total VRAM.

        Args:
            pipeline: The pipeline to optimize
            model_id: Model identifier
            force_full_optimization: Whether to force optimization
        """
        # Skip if it's OmniGen pipeline
        if pipeline.__class__.__name__ == "OmniGenPipeline":
            return

        if self.loaded_model is not None and self.is_in_device:
            logger.info(f"Model {model_id} already optimized")
            return

        device = get_available_torch_device()

        # Only optimize if model is bigger than total VRAM
        model_size = self.model_sizes.get(model_id, 0)
        if model_size > self.max_vram or force_full_optimization:
            logger.info(f"Applying optimizations for {model_id}")

            # Get list of optimizations
            optimizations = self._get_full_optimizations()

            # Apply the optimizations
            self._apply_optimization_list(pipeline, optimizations, device)
        else:
            logger.info(f"No optimization needed for {model_id}")

    def _get_full_optimizations(self) -> List[Tuple[str, str, Dict[str, Any]]]:
        """
        Get list of all available optimizations.

        Returns:
            List of tuples containing (optimization_function, name, parameters)
        """
        optimizations = [
            ("enable_vae_slicing", "VAE Sliced", {}),
            ("enable_vae_tiling", "VAE Tiled", {}),
            (
                "enable_model_cpu_offload",
                "CPU Offloading",
                {"device": get_available_torch_device()},
            ),
        ]

        if not isinstance(self.loaded_model, (FluxPipeline, FluxInpaintPipeline)):
            optimizations.append(
                (
                    "enable_xformers_memory_efficient_attention",
                    "Memory Efficient Attention",
                    {},
                )
            )

        return optimizations

    def _apply_optimization_list(
        self,
        pipeline: DiffusionPipeline,
        optimizations: List[Tuple[str, str, Dict[str, Any]]],
        device: torch.device,
    ) -> None:
        """
        Apply a list of optimizations to a pipeline.

        Args:
            pipeline: The pipeline to optimize
            optimizations: List of optimization specifications
            device: Target device
        """
        device_type = device if isinstance(device, str) else device.type

        if device_type == "mps":
            setattr(torch, "mps", torch.backends.mps)

        for opt_func, opt_name, kwargs in optimizations:
            try:
                getattr(pipeline, opt_func)(**kwargs)
                logger.info("{} enabled".format(opt_name))
            except Exception as e:
                logger.error("Error enabling {}: {}".format(opt_name, str(e)))

        if device_type == "mps":
            delattr(torch, "mps")

    def flush_memory(self) -> None:
        """Clear unused memory from GPU and perform garbage collection."""
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
            torch.cuda.synchronize()
            logger.info("GPU memory flushed successfully")
        elif torch.backends.mps.is_available():
            pass  # MPS doesn't need explicit cache clearing

    async def initialize_startup_models(self, model_ids: List[str]) -> None:
        """Inititalize models at startup with randomization"""
        if not model_ids:
            logger.info("No models configured for enabled models")
            return

        self.is_startup_load = True
        try:
            model_configs = {}
            model_sizes = {}

            for model_id in model_ids:
                model_config = get_config().pipeline_defs[model_id]
                if not model_config:
                    logger.warning(f"Model {model_id} not found in pipeline_defs")
                    continue

                try:
                    size = await self._get_model_size(model_config, model_id) + 1.0 # add 1GB for safety margin (model inference overhead in memory)
                    model_sizes[model_id] = size 
                    model_configs[model_id] = model_config
                except Exception as e:
                    logger.error(f"Error getting model size for {model_id}: {e}")
                    continue

            # Create randomized order of models
            available_models = list(model_ids)
            # available_vram = self._get_available_vram() - VRAM_SAFETY_MARGIN_GB

            rng = np.random.default_rng()
            random_models = rng.permutation(available_models).tolist()

            logger.info(f"Loading models in random order: {random_models}")
            total_loaded = 0
            total_size = 0

            for model_id in random_models:
                estimated_size = model_sizes[model_id]

                # Check if exceed available VRAM
                if estimated_size > (
                    self._get_available_vram() - VRAM_SAFETY_MARGIN_GB
                ):
                    logger.info(
                        f"Stopping model loading: Next model {model_id} "
                        f"({estimated_size:.2f} GB) would exceed available inference Memory "
                        f"({self._get_available_vram() - VRAM_SAFETY_MARGIN_GB:.2f} GB)"
                    )
                    break

                try:
                    pipeline = await self.load(model_id)
                    if pipeline is not None:
                        total_loaded += 1
                        total_size += estimated_size
                        logger.info(
                            f"Successfully loaded model {model_id} "
                            f"({total_loaded}/{len(random_models)}). "
                            f"Total VRAM used: {total_size:.2f} GB"
                        )

                        # warm up the model
                        logger.info(f"Warming up model {model_id}")
                        await self.warmup_pipeline(model_id)
                    else:
                        logger.warning(f"Failed to load model {model_id}")

                except Exception as e:
                    logger.error(f"Error loading model {model_id}: {e}")

            logger.info(
                f"Startup loading complete. Loaded and warmed up {total_loaded}/{len(random_models)} "
                f"models using {total_size:.2f}GB/{self.max_vram - VRAM_SAFETY_MARGIN_GB:.2f}GB available VRAM"
            )
        except Exception as e:
            logger.error(f"Error initializing startup models: {e}")
        finally:
            self.is_startup_load = False

    async def warmup_pipeline(self, model_id: str) -> None:
        """
        Warm up a pipeline by running a test inference.

        Args:
            model_id: Model identifier
        """
        if model_id not in self.loaded_models:
            if not self.is_startup_load:
                logger.info(f"Loading model {model_id} for warm-up")
                pipeline = await self.load(model_id)
                if pipeline is None:
                    logger.warning(f"Failed to load model {model_id} for warm-up")
                    return
            else:
                logger.warning(f"Model {model_id} is not loaded")
                return
        else:
            pipeline = self.loaded_models[model_id]

        logger.info(f"Warming up pipeline for model {model_id}")

        try:
            with torch.no_grad():
                if isinstance(pipeline, DiffusionPipeline) and callable(pipeline):
                    _ = pipeline(
                        prompt="This is a warm-up prompt",
                        num_inference_steps=4,
                        output_type="pil",
                    )
                else:
                    logger.warning(
                        f"Unsupported pipeline type for warm-up: {type(pipeline)}"
                    )
        except Exception as e:
            logger.error(f"Error during warm-up for model {model_id}: {str(e)}")

        self.flush_memory()
        logger.info(f"Warm-up completed for model {model_id}")

    def get_all_model_ids(self) -> List[str]:
        """
        Get list of all available model IDs from config.

        Returns:
            List of model identifiers
        """
        config = get_config()
        return list(config.pipeline_defs.keys())

    def get_enabled_models(self) -> List[str]:
        """
        Get list of models that should be warmed up.

        Returns:
            List of model identifiers to be used for generation
        """
        config = get_config()
        return config.enabled_models

    def unload(self, model_id: str) -> None:
        """
        Unload a model from memory and clean up associated resources.

        Handles unloading from both GPU and CPU memory, updates memory tracking,
        and cleans up LRU cache entries.

        Args:
            model_id: Model identifier to unload
        """
        try:
            # Unload from GPU if present
            if model_id in self.loaded_models:
                model_size = self.model_sizes.get(model_id, 0)
                pipeline = self.loaded_models.pop(model_id)

                del pipeline
                self.vram_usage -= model_size
                self.lru_cache.remove(model_id, "gpu")
                logger.info(f"Model {model_id} unloaded from GPU")

            # Unload from CPU if present
            if model_id in self.cpu_models:
                model_size = self.model_sizes.get(model_id, 0)
                pipeline = self.cpu_models.pop(model_id)

                del pipeline
                self.ram_usage -= model_size
                self.lru_cache.remove(model_id, "cpu")
                logger.info(f"Model {model_id} unloaded from CPU")

            # Clean up current model reference if it matches
            if model_id == self.current_model:
                self.loaded_model = None
                self.current_model = None

            # Remove from model sizes tracking
            if model_id in self.model_sizes:
                del self.model_sizes[model_id]

            # Remove from model types tracking
            if model_id in self.model_types:
                del self.model_types[model_id]

            # Force memory cleanup
            self.flush_memory()

        except Exception as e:
            logger.error(f"Error unloading model {model_id}: {str(e)}")

    def is_loaded(self, model_id: str) -> bool:
        """
        Check if a model is currently loaded.

        Args:
            model_id: Model identifier

        Returns:
            Boolean indicating if model is loaded
        """
        return model_id == self.current_model and self.loaded_model is not None

    def get_model(self, model_id: str) -> Optional[DiffusionPipeline]:
        """
        Get a loaded model pipeline.

        Args:
            model_id: Model identifier

        Returns:
            Pipeline if model is loaded, None otherwise
        """
        if not self.is_loaded(model_id):
            return None
        return self.loaded_model

    def get_model_device(self, model_id: str) -> Optional[torch.device]:
        """
        Get the device where a model is loaded.

        Args:
            model_id: Model identifier

        Returns:
            Device if model is loaded, None otherwise
        """
        model = self.get_model(model_id)
        return model.device if model else None
    

    async def _get_model_size(
        self, model_config: PipelineConfig, model_id: str
    ) -> float:
        """
        Calculate the total size of a model including all its components.

        Args:
            model_config: Model configuration dictionary from pipeline_defs
            model_id: The model identifier (key from pipeline_defs)

        Returns:
            Size in GB
        """

        if isinstance(model_config, dict):
            source = model_config.get("source")
        else:
            source = model_config.source

        if source.startswith(("http://", "https://")):
            # For downloaded HTTP(S) models, get size from cache
            try:
                source_obj = ModelSource(source)  # Create ModelSource object
                cache_path = await self.model_downloader._get_cache_path(
                    model_id, source_obj
                )
                print(f"Cache Path: {cache_path}")
                if os.path.exists(cache_path):
                    return os.path.getsize(cache_path) / (1024**3)
                else:
                    logger.warning(
                        f"Cache file not found for {model_id}, assuming default size"
                    )
                    return 7.0  # Default size assumption (never going to be used)
            except Exception as e:
                logger.error(f"Error getting cached model size: {e}")
                return (
                    7.0  # Default fallback size (never going to be used for anything)
                )
        elif source.startswith("file:"):
            path = source.replace("file:", "")
            return os.path.getsize(path) / (1024**3) if os.path.exists(path) else 0
        elif source.startswith("hf:"):
            # Handle HuggingFace models as before
            repo_id = source.replace("hf:", "")
            return self._calculate_repo_size(repo_id, model_config)
        else:
            logger.error(f"Unsupported source type for size calculation: {source}")
            return 0

    def _calculate_repo_size(self, repo_id: str, model_config: PipelineConfig) -> float:
        """
        Calculate the total size of a HuggingFace repository.

        Args:
            repo_id: Repository identifier
            model_config: Model configuration dictionary

        Returns:
            Total size in GB
        """
        total_size = self._get_size_for_repo(repo_id)

        if isinstance(model_config, dict):
            components = model_config.get("components")
        else:
            components = model_config.components

        if components:
            for key, component in components.items():
                if isinstance(component, dict) and "source" in component:
                    component_size = self._calculate_component_size(
                        component, repo_id, key
                    )
                    total_size += component_size

        total_size_gb = total_size / (1024**3)
        logger.debug(f"Total size: {total_size_gb:.2f} GB")
        print(f"Total size: {total_size_gb:.2f} GB")
        return total_size_gb

    def _calculate_component_size(
        self, component: Dict[str, Any], repo_id: str, key: str
    ) -> float:
        """
        Calculate the size of a model component.

        Args:
            component: Component configuration
            repo_id: Repository identifier
            key: Component key

        Returns:
            Component size in bytes
        """
        component_source = component["source"]
        if len(component_source.split("/")) > 2:
            component_repo = "/".join(component_source.split("/")[0:2]).replace(
                "hf:", ""
            )
        else:
            component_repo = component_source.replace("hf:", "")

        component_name = (
            key
            if not component_source.endswith((".safetensors", ".bin", ".ckpt", ".pt"))
            else component_source.split("/")[-1]
        )

        total_size = self._get_size_for_repo(
            component_repo, component_name
        ) - self._get_size_for_repo(repo_id, key)

        # total_size = self._get_size_for_repo(component_repo, component_name)

        return total_size

    def _get_size_for_repo(
        self, repo_id: str, component_name: Optional[str] = None
    ) -> int:
        """
        Get the size of a specific repository or component.

        Args:
            repo_id: Repository identifier
            component_name: Optional component name

        Returns:
            Size in bytes
        """
        if component_name == "scheduler":
            return 0

        print(f"Getting size for {repo_id} {component_name}")
        storage_folder = os.path.join(
            self.cache_dir, repo_folder_name(repo_id=repo_id, repo_type="model")
        )

        if not os.path.exists(storage_folder):
            logger.warning(f"Storage folder for {repo_id} not found")
            return 0

        commit_hash = self._get_commit_hash(storage_folder)
        if not commit_hash:
            return 0

        snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash)
        if component_name:
            snapshot_folder = os.path.join(snapshot_folder, component_name)

        return self._calculate_folder_size(snapshot_folder)

    def _get_commit_hash(self, storage_folder: str) -> Optional[str]:
        """
        Get the commit hash for a repository.

        Args:
            storage_folder: Path to the repository storage folder

        Returns:
            Commit hash string or None if not found
        """
        refs_path = os.path.join(storage_folder, "refs", "main")
        if not os.path.exists(refs_path):
            logger.warning("No commit hash found for {}".format(storage_folder))
            return None
        try:
            with open(refs_path, "r") as f:
                return f.read().strip()
        except Exception as e:
            logger.error("Error reading commit hash: {}".format(str(e)))
            return None

    def _calculate_folder_size(self, folder: str) -> int:
        """
        Calculate the total size of model files in a folder.

        Args:
            folder: Path to the folder

        Returns:
            Total size in bytes
        """
        if not os.path.isdir(folder):
            if os.path.exists(folder):
                return os.path.getsize(folder)
            else:
                return 0

        variants = ["bf16", "fp8", "fp16", ""]
        selected_variant = next(
            (v for v in variants if self._check_variant_files(folder, v)), None
        )

        if selected_variant is None:
            return 0

        total_size = 0
        for root, _, files in os.walk(folder):
            # skip if in root
            if root == folder:
                continue
            for file in files:
                if self._is_valid_file(file, selected_variant):
                    total_size += os.path.getsize(os.path.join(root, file))
                    # break

        return total_size

    def _check_variant_files(self, folder: str, variant: str) -> bool:
        """
        Check if a folder contains files of a specific variant.

        Args:
            folder: Path to the folder
            variant: Variant to check for

        Returns:
            Boolean indicating if variant files exist
        """
        for root, _, files in os.walk(folder):
            if any(self._is_valid_file(f, variant) for f in files):
                return True
        return False

    def _is_valid_file(self, file: str, variant: str) -> bool:
        """
        Check if a file is a valid model file of a specific variant.

        Args:
            file: Filename to check
            variant: Variant to check for

        Returns:
            Boolean indicating if file is valid
        """
        if variant:
            # check if the variant is in the file name
            if variant in file:
                return True
            else:
                return False
        return file.endswith((".safetensors", ".bin", ".ckpt"))


    async def _load_diffusers_component(
        self,
        main_model_repo: str,
        component_repo: str,
        component_name: Optional[str] = None,
        variant: Optional[str] = None,
    ) -> Any:
        """
        Load a diffusers component.

        Args:
            component_repo: Repository identifier
            component_name: Name of the component

        Returns:
            Loaded component
        """
        try:
            model_index = (
                await self.model_downloader.get_diffusers_multifolder_components(
                    main_model_repo
                )
            )
            if model_index is None:
                raise ValueError(f"model_index does not exist for {main_model_repo}")

            component_info = model_index.get(component_name)
            if not component_info:
                raise ValueError(f"Invalid component info for {component_name}")

            module_path, class_name = component_info
            module = importlib.import_module(module_path)
            model_class = getattr(module, class_name)

            if component_name:
                if variant:
                    component = await asyncio.to_thread(model_class.from_pretrained,
                        component_repo,
                        subfolder=component_name,
                        variant=variant,
                        torch_dtype=torch.bfloat16
                        if "flux" in component_repo.lower()
                        else torch.float16,
                    )
                else:
                    component = await asyncio.to_thread(model_class.from_pretrained,
                        component_repo,
                        subfolder=component_name,
                        torch_dtype=torch.bfloat16
                        if "flux" in component_repo.lower()
                        else torch.float16,
                    )
            else:
                if variant:
                    component = await asyncio.to_thread(model_class.from_pretrained,
                        component_repo,
                        variant=variant,
                        torch_dtype=torch.bfloat16
                        if "flux" in component_repo.lower()
                        else torch.float16,
                    )
                else:
                    component = await asyncio.to_thread(model_class.from_pretrained,
                        component_repo,
                        torch_dtype=torch.bfloat16
                        if "flux" in component_repo.lower()
                        else torch.float16,
                    )


            return component

        except Exception as e:
            logger.error(
                "Error loading component {} from {}: {}".format(
                    component_name, component_repo, str(e)
                )
            )
            raise

    def _load_custom_component(
        self, repo_id: str, category: str, component_name: str
    ) -> Any:
        """
        Load a custom component.

        Args:
            repo_id: Repository identifier
            category: Component category
            component_name: Name of the component

        Returns:
            Loaded component
        """
        try:
            file_path = self._get_component_file_path(repo_id)
            state_dict = load_state_dict_from_file(file_path)

            architectures = get_architectures()
            arch_key = f"core_extension_1.{category.lower()}_{component_name.lower()}"
            architecture_class = architectures.get(arch_key)

            if not architecture_class:
                raise ValueError(f"Architecture not found for {arch_key}")

            architecture = architecture_class()
            architecture.load(state_dict)
            model = architecture.model

            return model

        except Exception as e:
            logger.error("Error loading custom component: {}".format(str(e)))
            raise

    def _get_component_file_path(self, repo_id: str) -> str:
        """
        Get the file path for a component.

        Args:
            repo_id: Repository identifier

        Returns:
            Path to component file
        """
        repo_folder = os.path.dirname(repo_id.replace("hf:", ""))
        weights_name = repo_id.split("/")[-1]

        model_folder = os.path.join(
            self.cache_dir, repo_folder_name(repo_id=repo_folder, repo_type="model")
        )

        if not os.path.exists(model_folder):
            model_folder = os.path.join(self.cache_dir, repo_folder)
            if not os.path.exists(model_folder):
                raise FileNotFoundError(f"Cache folder for {repo_id} not found")
            return os.path.join(model_folder, weights_name)

        commit_hash = self._get_commit_hash(model_folder) or ""
        return os.path.join(model_folder, "snapshots", commit_hash, weights_name)
