lollms-client 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

lollms_client/tti_bindings/diffusers/__init__.py

@@ -4,693 +4,610 @@ import importlib
 from io import BytesIO
 from typing import Optional, List, Dict, Any, Union
 from pathlib import Path
-
-
-try:
-    import pipmaster as pm
-    import platform # For OS detection for torch index
-
-    # Determine initial device preference to guide torch installation
-    preferred_torch_device_for_install = "cpu" # Default assumption
-    
-    # Tentatively set preference based on OS, assuming user might want GPU if available
-    if platform.system() == "Linux" or platform.system() == "Windows":
-        # On Linux/Windows, CUDA is the primary GPU acceleration for PyTorch.
-        # We will try to install a CUDA version of PyTorch.
-        preferred_torch_device_for_install = "cuda"
-    elif platform.system() == "Darwin":
-        # On macOS, MPS is the acceleration. Standard torch install usually handles this.
-        preferred_torch_device_for_install = "mps" # or keep cpu if mps detection is later
-
-    torch_pkgs = ["torch", "torchaudio", "torchvision", "xformers"]
-    diffusers_core_pkgs = ["diffusers", "Pillow", "transformers", "safetensors"]
-
-    torch_index_url = None
-    if preferred_torch_device_for_install == "cuda":
-        # Specify a common CUDA version index. Pip should resolve the correct torch version.
-        # As of late 2023/early 2024, cu118 or cu121 are common. Let's use cu126.
-        # Users with different CUDA setups might need to pre-install torch manually.
-        torch_index_url = "https://download.pytorch.org/whl/cu126"
-        ASCIIColors.info(f"Attempting to ensure PyTorch with CUDA support (target index: {torch_index_url})")
-        # Install torch and torchaudio first from the specific index
-        pm.ensure_packages(torch_pkgs, index_url=torch_index_url)
-        # Then install audiocraft and other dependencies; pip should use the already installed torch
-        pm.ensure_packages(diffusers_core_pkgs)
-    else:
-        # For CPU, MPS, or if no specific CUDA preference was determined for install
-        ASCIIColors.info("Ensuring PyTorch, AudioCraft, and dependencies using default PyPI index.")
-        pm.ensure_packages(torch_pkgs + diffusers_core_pkgs)
-
-    import whisper
-    import torch
-    _whisper_installed = True
-except Exception as e:
-    _whisper_installation_error = str(e)
-    whisper = None
-    torch = None
-
-
+import pipmaster as pm
+# --- Concurrency Imports ---
+import threading
+import queue
+from concurrent.futures import Future
+import time
+import hashlib
+import re
+# -------------------------
+# --- Download Imports ---
+import requests
+from tqdm import tqdm
+# --------------------
+
+pm.ensure_packages(["torch","torchvision"],index_url="https://download.pytorch.org/whl/cu126")
+pm.ensure_packages(["diffusers","pillow","transformers","safetensors", "requests", "tqdm"])
+
+# Attempt to import core dependencies and set availability flag
 try:
     import torch
-    from diffusers import AutoPipelineForText2Image, DiffusionPipeline
-    from diffusers.utils import load_image # Potentially for future img2img etc.
+    from diffusers import AutoPipelineForText2Image, DiffusionPipeline, StableDiffusionPipeline
+    from diffusers.utils import load_image
     from PIL import Image
     DIFFUSERS_AVAILABLE = True
 except ImportError:
     torch = None
     AutoPipelineForText2Image = None
     DiffusionPipeline = None
+    StableDiffusionPipeline = None
     Image = None
     load_image = None
     DIFFUSERS_AVAILABLE = False
-    # Detailed error will be raised in __init__ if user tries to use it
 
 from lollms_client.lollms_tti_binding import LollmsTTIBinding
 from ascii_colors import trace_exception, ASCIIColors
-import json # For potential JSONDecodeError and settings
+import json
+import shutil
 
 # Defines the binding name for the manager
 BindingName = "DiffusersTTIBinding_Impl"
 
+# --- START: Civitai Model Definitions ---
+# Expanded list of popular Civitai models (as single .safetensors files)
+CIVITAI_MODELS = {
+    # Photorealistic
+    "realistic-vision-v6": {
+        "display_name": "Realistic Vision V6.0",
+        "url": "https://civitai.com/api/download/models/130072",
+        "filename": "realisticVisionV60_v60B1.safetensors",
+        "description": "One of the most popular photorealistic models.",
+        "owned_by": "civitai"
+    },
+    "absolute-reality": {
+        "display_name": "Absolute Reality",
+        "url": "https://civitai.com/api/download/models/132760",
+        "filename": "absolutereality_v181.safetensors",
+        "description": "A top-tier model for generating realistic images.",
+        "owned_by": "civitai"
+    },
+    # Artistic / General Purpose
+    "dreamshaper-8": {
+        "display_name": "DreamShaper 8",
+        "url": "https://civitai.com/api/download/models/128713",
+        "filename": "dreamshaper_8.safetensors",
+        "description": "A very popular and versatile general-purpose model.",
+        "owned_by": "civitai"
+    },
+    "juggernaut-xl": {
+        "display_name": "Juggernaut XL",
+        "url": "https://civitai.com/api/download/models/133005",
+        "filename": "juggernautXL_version6Rundiffusion.safetensors",
+        "description": "High-quality artistic model, great for cinematic styles (SDXL-based).",
+        "owned_by": "civitai"
+    },
+    # Anime
+    "anything-v5": {
+        "display_name": "Anything V5",
+        "url": "https://civitai.com/api/download/models/9409",
+        "filename": "anythingV5_PrtRE.safetensors",
+        "description": "A classic and highly popular model for anime-style generation.",
+        "owned_by": "civitai"
+    },
+    "lyriel-v1.6": {
+        "display_name": "Lyriel v1.6",
+        "url": "https://civitai.com/api/download/models/92407",
+        "filename": "lyriel_v16.safetensors",
+        "description": "A popular artistic model for fantasy and stylized images.",
+        "owned_by": "civitai"
+    }
+}
+# --- END: Civitai Model Definitions ---
+
 # Helper for torch.dtype string conversion
 TORCH_DTYPE_MAP_STR_TO_OBJ = {
-    "float16": torch.float16 if torch else "float16", # Keep string if torch not loaded
-    "bfloat16": torch.bfloat16 if torch else "bfloat16",
-    "float32": torch.float32 if torch else "float32",
+    "float16": getattr(torch, 'float16', 'float16'),
+    "bfloat16": getattr(torch, 'bfloat16', 'bfloat16'),
+    "float32": getattr(torch, 'float32', 'float32'),
     "auto": "auto"
 }
 TORCH_DTYPE_MAP_OBJ_TO_STR = {v: k for k, v in TORCH_DTYPE_MAP_STR_TO_OBJ.items()}
-if torch: # Add None mapping if torch is loaded
+if torch:
     TORCH_DTYPE_MAP_OBJ_TO_STR[None] = "None"
 
-
-# Common Schedulers mapping (User-friendly name to Class name)
+# Common Schedulers mapping
 SCHEDULER_MAPPING = {
-    "default": None,  # Use model's default
-    "ddim": "DDIMScheduler",
-    "ddpm": "DDPMScheduler",
-    "deis_multistep": "DEISMultistepScheduler",
-    "dpm_multistep": "DPMSolverMultistepScheduler", # Alias
-    "dpm_multistep_karras": "DPMSolverMultistepScheduler", # Configured with use_karras_sigmas=True
-    "dpm_single": "DPMSolverSinglestepScheduler",
-    "dpm_adaptive": "DP soluzioniPlusPlusScheduler", # DPM++ 2M Karras in A1111
-    "dpm++_2m": "DPMSolverMultistepScheduler", 
-    "dpm++_2m_karras": "DPMSolverMultistepScheduler", # Configured with use_karras_sigmas=True
-    "dpm++_2s_ancestral": "DPMSolverAncestralDiscreteScheduler",
-    "dpm++_2s_ancestral_karras": "DPMSolverAncestralDiscreteScheduler", # Configured with use_karras_sigmas=True
-    "dpm++_sde": "DPMSolverSDEScheduler",
-    "dpm++_sde_karras": "DPMSolverSDEScheduler", # Configured with use_karras_sigmas=True
-    "euler_ancestral_discrete": "EulerAncestralDiscreteScheduler",
-    "euler_discrete": "EulerDiscreteScheduler",
-    "heun_discrete": "HeunDiscreteScheduler",
-    "heun_karras": "HeunDiscreteScheduler", # Configured with use_karras_sigmas=True
-    "lms_discrete": "LMSDiscreteScheduler",
-    "lms_karras": "LMSDiscreteScheduler", # Configured with use_karras_sigmas=True
-    "pndm": "PNDMScheduler",
-    "unipc_multistep": "UniPCMultistepScheduler",
+    "default": None, "ddim": "DDIMScheduler", "ddpm": "DDPMScheduler", "deis_multistep": "DEISMultistepScheduler",
+    "dpm_multistep": "DPMSolverMultistepScheduler", "dpm_multistep_karras": "DPMSolverMultistepScheduler",
+    "dpm_single": "DPMSolverSinglestepScheduler", "dpm_adaptive": "DPMSolverPlusPlusScheduler",
+    "dpm++_2m": "DPMSolverMultistepScheduler", "dpm++_2m_karras": "DPMSolverMultistepScheduler",
+    "dpm++_2s_ancestral": "DPMSolverAncestralDiscreteScheduler", "dpm++_2s_ancestral_karras": "DPMSolverAncestralDiscreteScheduler",
+    "dpm++_sde": "DPMSolverSDEScheduler", "dpm++_sde_karras": "DPMSolverSDEScheduler",
+    "euler_ancestral_discrete": "EulerAncestralDiscreteScheduler", "euler_discrete": "EulerDiscreteScheduler",
+    "heun_discrete": "HeunDiscreteScheduler", "heun_karras": "HeunDiscreteScheduler",
+    "lms_discrete": "LMSDiscreteScheduler", "lms_karras": "LMSDiscreteScheduler",
+    "pndm": "PNDMScheduler", "unipc_multistep": "UniPCMultistepScheduler",
 }
 SCHEDULER_USES_KARRAS_SIGMAS = [
     "dpm_multistep_karras", "dpm++_2m_karras", "dpm++_2s_ancestral_karras",
     "dpm++_sde_karras", "heun_karras", "lms_karras"
 ]
 
+# --- START: Concurrency and Singleton Management ---
 
-class DiffusersTTIBinding_Impl(LollmsTTIBinding):
+class ModelManager:
     """
-    Concrete implementation of LollmsTTIBinding for Hugging Face Diffusers library.
-    Allows running various text-to-image models locally.
+    Manages a single pipeline instance, its generation queue, and a worker thread.
+    This ensures all interactions with a specific model are thread-safe.
     """
-    DEFAULT_CONFIG = {
-        "model_id_or_path": "stabilityai/stable-diffusion-2-1-base",
-        "device": "auto",  # "auto", "cuda", "mps", "cpu"
-        "torch_dtype_str": "auto",  # "auto", "float16", "bfloat16", "float32"
-        "use_safetensors": True,
-        "scheduler_name": "default",
-        "safety_checker_on": True,  # Note: Diffusers default is ON
-        "num_inference_steps": 25,
-        "guidance_scale": 7.5,
-        "default_width": 768,  # Default for SD 2.1 base
-        "default_height": 768,  # Default for SD 2.1 base
-        "seed": -1,  # -1 for random on each call
-        "enable_cpu_offload": False,
-        "enable_sequential_cpu_offload": False,
-        "enable_xformers": False,  # Explicit opt-in for xformers
-        "hf_variant": None,  # e.g., "fp16"
-        "hf_token": None,
-        "local_files_only": False,
-    }
-
-    def __init__(self,
-                 config: Optional[Dict[str, Any]] = None,
-                 lollms_paths: Optional[Dict[str, Union[str, Path]]] = None,
-                 **kwargs  # Catches other potential parameters like 'service_key' or 'client_id'
-                 ):
-        """
-        Initialize the Diffusers TTI binding.
-
-        Args:
-            config (Optional[Dict[str, Any]]): Configuration dictionary for the binding.
-                                               Overrides DEFAULT_CONFIG.
-            lollms_paths (Optional[Dict[str, Union[str, Path]]]): Dictionary of LOLLMS paths.
-                                                                  Used for model/cache directories.
-            **kwargs: Catches other parameters (e.g. service_key).
-        """
-        super().__init__(binding_name="diffusers")
-
-        if not DIFFUSERS_AVAILABLE:
-            ASCIIColors.error("Diffusers library or its dependencies (torch, Pillow, transformers) are not installed or failed to import.")
-            ASCIIColors.info("Attempting to install/verify packages...")
-            pm.ensure_packages(["torch", "diffusers", "Pillow", "transformers", "safetensors"])
-            try:
-                import torch as _torch
-                from diffusers import AutoPipelineForText2Image as _AutoPipelineForText2Image
-                from diffusers import DiffusionPipeline as _DiffusionPipeline
-                from PIL import Image as _Image
-                globals()['torch'] = _torch
-                globals()['AutoPipelineForText2Image'] = _AutoPipelineForText2Image
-                globals()['DiffusionPipeline'] = _DiffusionPipeline
-                globals()['Image'] = _Image
-
-                # Re-populate torch dtype maps if torch was just loaded
-                global TORCH_DTYPE_MAP_STR_TO_OBJ, TORCH_DTYPE_MAP_OBJ_TO_STR
-                TORCH_DTYPE_MAP_STR_TO_OBJ = {
-                    "float16": _torch.float16, "bfloat16": _torch.bfloat16,
-                    "float32": _torch.float32, "auto": "auto"
-                }
-                TORCH_DTYPE_MAP_OBJ_TO_STR = {v: k for k, v in TORCH_DTYPE_MAP_STR_TO_OBJ.items()}
-                TORCH_DTYPE_MAP_OBJ_TO_STR[None] = "None"
-                ASCIIColors.green("Dependencies seem to be available now.")
-            except ImportError as e:
-                trace_exception(e)
-                raise ImportError(
-                    "Diffusers binding dependencies are still not met after trying to ensure them. "
-                    "Please install torch, diffusers, Pillow, and transformers manually. "
-                    f"Error: {e}"
-                ) from e
-
-        # Merge configs, lollms_paths, and kwargs
-        self.config = {**self.DEFAULT_CONFIG, **(config or {}), **kwargs}
-        self.lollms_paths = {k: Path(v) for k, v in (lollms_paths or {}).items()} if lollms_paths else {}
-
+    def __init__(self, config: Dict[str, Any], models_path: Path):
+        self.config = config
+        self.models_path = models_path
         self.pipeline: Optional[DiffusionPipeline] = None
-        self.current_model_id_or_path = None  # To track if model needs reload
-
-        # Resolve auto settings for device and dtype
-        if self.config["device"].lower() == "auto":
-            if torch.cuda.is_available():
-                self.config["device"] = "cuda"
-            elif torch.backends.mps.is_available():
-                self.config["device"] = "mps"
-            else:
-                self.config["device"] = "cpu"
-
-        if self.config["torch_dtype_str"].lower() == "auto":
-            if self.config["device"] == "cpu":
-                self.config["torch_dtype_str"] = "float32"  # CPU usually float32
-            else:
-                self.config["torch_dtype_str"] = "float16"  # Common default for GPU
+        self.ref_count = 0
+        self.lock = threading.Lock()
+        self.queue = queue.Queue()
+        self.worker_thread = threading.Thread(target=self._generation_worker, daemon=True)
+        self._stop_event = threading.Event()
+        self.is_loaded = False
+        
+        self.worker_thread.start()
 
-        self.torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower(), torch.float32)
-        if self.torch_dtype == "auto":  # Should have been resolved above
-            self.torch_dtype = torch.float16 if self.config["device"] != "cpu" else torch.float32
-            self.config["torch_dtype_str"] = TORCH_DTYPE_MAP_OBJ_TO_STR.get(self.torch_dtype, "float32")
+    def acquire(self):
+        with self.lock:
+            self.ref_count += 1
+            return self
 
-        # For potential lollms client specific features
-        self.client_id = kwargs.get("service_key", kwargs.get("client_id", "diffusers_client_user"))
+    def release(self):
+        with self.lock:
+            self.ref_count -= 1
+            return self.ref_count
 
-        self.load_model()
+    def stop(self):
+        self._stop_event.set()
+        self.queue.put(None)
+        self.worker_thread.join(timeout=5)
 
+    def _load_pipeline(self):
+        if self.pipeline:
+            return
 
-    def _resolve_model_path(self, model_id_or_path: str) -> str:
-        """Resolves a model name/path against lollms_paths if not absolute."""
-        if os.path.isabs(model_id_or_path):
-            return model_id_or_path
-        
-        # Check personal_models_path/diffusers_models/<name>
-        if self.lollms_paths.get('personal_models_path'):
-            personal_diffusers_path = self.lollms_paths['personal_models_path'] / "diffusers_models" / model_id_or_path
-            if personal_diffusers_path.exists() and personal_diffusers_path.is_dir():
-                ASCIIColors.info(f"Found local model in personal_models_path: {personal_diffusers_path}")
-                return str(personal_diffusers_path)
-
-        # Check models_zoo_path/diffusers_models/<name> (if different from personal)
-        if self.lollms_paths.get('models_zoo_path') and \
-           self.lollms_paths.get('models_zoo_path') != self.lollms_paths.get('personal_models_path'):
-            zoo_diffusers_path = self.lollms_paths['models_zoo_path'] / "diffusers_models" / model_id_or_path
-            if zoo_diffusers_path.exists() and zoo_diffusers_path.is_dir():
-                ASCIIColors.info(f"Found local model in models_zoo_path: {zoo_diffusers_path}")
-                return str(zoo_diffusers_path)
+        model_name = self.config.get("model_name", "")
+        if not model_name:
+            raise ValueError("Model name cannot be empty for loading.")
+            
+        ASCIIColors.info(f"Loading Diffusers model: {model_name}")
+        model_path = self._resolve_model_path(model_name)
+        torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower())
         
-        ASCIIColors.info(f"Assuming '{model_id_or_path}' is a Hugging Face Hub ID or already fully qualified.")
-        return model_id_or_path
-
-    def load_model(self):
-        """Loads the Diffusers pipeline based on current configuration."""
-        ASCIIColors.info("Loading Diffusers model...")
-        if self.pipeline is not None:
-            self.unload_model() # Ensure old model is cleared
-
         try:
-            model_path = self._resolve_model_path(self.config["model_id_or_path"])
-            self.current_model_id_or_path = model_path # Store what's actually loaded
-
-            load_args = {
-                "torch_dtype": self.torch_dtype,
-                "use_safetensors": self.config["use_safetensors"],
-                "token": self.config["hf_token"],
-                "local_files_only": self.config["local_files_only"],
-            }
-            if self.config["hf_variant"]:
-                load_args["variant"] = self.config["hf_variant"]
-            
-            if not self.config["safety_checker_on"]:
-                load_args["safety_checker"] = None
-            
-            if self.lollms_paths.get("shared_cache_path"):
-                load_args["cache_dir"] = str(self.lollms_paths["shared_cache_path"] / "huggingface_diffusers")
+            if str(model_path).endswith(".safetensors"):
+                ASCIIColors.info(f"Loading from single safetensors file: {model_path}")
+                try:
+                    # Modern, preferred method for newer diffusers versions
+                    self.pipeline = AutoPipelineForText2Image.from_single_file(
+                        model_path,
+                        torch_dtype=torch_dtype,
+                        cache_dir=self.config.get("hf_cache_path")
+                    )
+                except AttributeError:
+                    # Fallback for older diffusers versions
+                    ASCIIColors.warning("AutoPipelineForText2Image.from_single_file not found. Falling back to StableDiffusionPipeline.")
+                    ASCIIColors.warning("Consider updating diffusers for better compatibility: pip install --upgrade diffusers")
+                    self.pipeline = StableDiffusionPipeline.from_single_file(
+                        model_path,
+                        torch_dtype=torch_dtype,
+                        cache_dir=self.config.get("hf_cache_path")
+                    )
+            else:
+                ASCIIColors.info(f"Loading from pretrained folder/repo: {model_path}")
+                load_args = {
+                    "torch_dtype": torch_dtype, "use_safetensors": self.config["use_safetensors"],
+                    "token": self.config["hf_token"], "local_files_only": self.config["local_files_only"],
+                }
+                if self.config["hf_variant"]: load_args["variant"] = self.config["hf_variant"]
+                if not self.config["safety_checker_on"]: load_args["safety_checker"] = None
+                if self.config.get("hf_cache_path"): load_args["cache_dir"] = str(self.config["hf_cache_path"])
+                self.pipeline = AutoPipelineForText2Image.from_pretrained(model_path, **load_args)
 
+        except Exception as e:
+            error_str = str(e).lower()
+            if "401" in error_str or "gated" in error_str or "authorization" in error_str:
+                auth_error_msg = (
+                    f"AUTHENTICATION FAILED for model '{model_name}'. This is likely a 'gated' model on Hugging Face.\n"
+                    "Please ensure you have accepted its license and provided a valid HF Access Token in the settings."
+                )
+                raise RuntimeError(auth_error_msg) from e
+            else:
+                raise e
 
-            # Use AutoPipelineForText2Image for flexibility
-            pipeline_class_to_load = AutoPipelineForText2Image
-            custom_pipeline_class_name = self.config.get("pipeline_class_name")
+        self._set_scheduler()
+        self.pipeline.to(self.config["device"])
 
-            if custom_pipeline_class_name:
-                try:
-                    diffusers_module = importlib.import_module("diffusers")
-                    pipeline_class_to_load = getattr(diffusers_module, custom_pipeline_class_name)
-                    ASCIIColors.info(f"Using specified pipeline class: {custom_pipeline_class_name}")
-                except (ImportError, AttributeError) as e:
-                    ASCIIColors.warning(f"Could not load custom pipeline class {custom_pipeline_class_name}: {e}. Falling back to AutoPipelineForText2Image.")
-                    pipeline_class_to_load = AutoPipelineForText2Image
-            
-            self.pipeline = pipeline_class_to_load.from_pretrained(model_path, **load_args)
-            
-            # Scheduler
-            self._set_scheduler()
+        if self.config["enable_xformers"]:
+            try:
+                self.pipeline.enable_xformers_memory_efficient_attention()
+            except Exception as e:
+                ASCIIColors.warning(f"Could not enable xFormers: {e}.")
+        
+        if self.config["enable_cpu_offload"] and self.config["device"] != "cpu":
+            self.pipeline.enable_model_cpu_offload()
+        elif self.config["enable_sequential_cpu_offload"] and self.config["device"] != "cpu":
+            self.pipeline.enable_sequential_cpu_offload()
+        
+        self.is_loaded = True
+        ASCIIColors.green(f"Model '{model_name}' loaded successfully on '{self.config['device']}'.")
 
-            self.pipeline.to(self.config["device"])
+    def _unload_pipeline(self):
+        if self.pipeline:
+            del self.pipeline
+            self.pipeline = None
+            if torch and torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            self.is_loaded = False
+            ASCIIColors.info(f"Model '{self.config.get('model_name')}' unloaded.")
 
-            if self.config["enable_xformers"]:
+    def _generation_worker(self):
+        while not self._stop_event.is_set():
+            try:
+                job = self.queue.get(timeout=1)
+                if job is None:
+                    break
+                future, pipeline_args = job
                 try:
-                    self.pipeline.enable_xformers_memory_efficient_attention()
-                    ASCIIColors.info("xFormers memory efficient attention enabled.")
+                    with self.lock:
+                        if not self.pipeline:
+                            self._load_pipeline()
+                    with torch.no_grad():
+                        pipeline_output = self.pipeline(**pipeline_args)
+                    pil_image: Image.Image = pipeline_output.images[0]
+                    img_byte_arr = BytesIO()
+                    pil_image.save(img_byte_arr, format="PNG")
+                    future.set_result(img_byte_arr.getvalue())
                 except Exception as e:
-                    ASCIIColors.warning(f"Could not enable xFormers: {e}. Proceeding without it.")
+                    trace_exception(e)
+                    future.set_exception(e)
+                finally:
+                    self.queue.task_done()
+            except queue.Empty:
+                continue
+
+    def _download_civitai_model(self, model_key: str):
+        model_info = CIVITAI_MODELS[model_key]
+        url = model_info["url"]
+        filename = model_info["filename"]
+        dest_path = self.models_path / filename
+        temp_path = dest_path.with_suffix(".temp")
+        
+        ASCIIColors.cyan(f"Downloading '{filename}' from Civitai...")
+        try:
+            with requests.get(url, stream=True) as r:
+                r.raise_for_status()
+                total_size = int(r.headers.get('content-length', 0))
+                with open(temp_path, 'wb') as f, tqdm(
+                    total=total_size, unit='iB', unit_scale=True, desc=f"Downloading {filename}"
+                ) as bar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        bar.update(len(chunk))
             
-            if self.config["enable_cpu_offload"] and self.config["device"] != "cpu":
-                self.pipeline.enable_model_cpu_offload()
-                ASCIIColors.info("Model CPU offload enabled.")
-            elif self.config["enable_sequential_cpu_offload"] and self.config["device"] != "cpu":
-                self.pipeline.enable_sequential_cpu_offload() # More aggressive
-                ASCIIColors.info("Sequential CPU offload enabled.")
-
-
-            ASCIIColors.green(f"Diffusers model '{model_path}' loaded successfully on device '{self.config['device']}' with dtype '{self.config['torch_dtype_str']}'.")
-
+            shutil.move(temp_path, dest_path)
+            ASCIIColors.green(f"Model '{filename}' downloaded successfully.")
         except Exception as e:
-            trace_exception(e)
-            self.pipeline = None
-            raise RuntimeError(f"Failed to load Diffusers model '{self.config['model_id_or_path']}': {e}") from e
+            if temp_path.exists():
+                temp_path.unlink()
+            raise Exception(f"Failed to download model {filename}: {e}") from e
+
+    def _resolve_model_path(self, model_name: str) -> Union[str, Path]:
+        path_obj = Path(model_name)
+        if path_obj.is_absolute() and path_obj.exists():
+            return model_name
+        
+        if model_name in CIVITAI_MODELS:
+            filename = CIVITAI_MODELS[model_name]["filename"]
+            local_path = self.models_path / filename
+            if not local_path.exists():
+                self._download_civitai_model(model_name)
+            return local_path
+        
+        local_path = self.models_path / model_name
+        if local_path.exists():
+            return local_path
+            
+        return model_name
 
     def _set_scheduler(self):
         if not self.pipeline: return
-        
         scheduler_name_key = self.config["scheduler_name"].lower()
-        if scheduler_name_key == "default":
-            ASCIIColors.info(f"Using model's default scheduler: {self.pipeline.scheduler.__class__.__name__}")
-            return
+        if scheduler_name_key == "default": return
 
         scheduler_class_name = SCHEDULER_MAPPING.get(scheduler_name_key)
         if scheduler_class_name:
             try:
-                scheduler_module = importlib.import_module("diffusers.schedulers")
-                SchedulerClass = getattr(scheduler_module, scheduler_class_name)
-                
+                SchedulerClass = getattr(importlib.import_module("diffusers.schedulers"), scheduler_class_name)
                 scheduler_config = self.pipeline.scheduler.config
-                if scheduler_name_key in SCHEDULER_USES_KARRAS_SIGMAS:
-                    scheduler_config["use_karras_sigmas"] = True
-                else: # Ensure it's False if not a karras variant for this scheduler
-                    if "use_karras_sigmas" in scheduler_config:
-                         scheduler_config["use_karras_sigmas"] = False
-
-
+                scheduler_config["use_karras_sigmas"] = scheduler_name_key in SCHEDULER_USES_KARRAS_SIGMAS
                 self.pipeline.scheduler = SchedulerClass.from_config(scheduler_config)
-                ASCIIColors.info(f"Switched scheduler to {scheduler_name_key} ({scheduler_class_name}).")
             except Exception as e:
-                trace_exception(e)
                 ASCIIColors.warning(f"Could not switch scheduler to {scheduler_name_key}: {e}. Using current default.")
-        else:
-            ASCIIColors.warning(f"Unknown scheduler name: {self.config['scheduler_name']}. Using model default.")
 
+class PipelineRegistry:
+    _instance = None
+    _lock = threading.Lock()
+    
+    def __new__(cls, *args, **kwargs):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super().__new__(cls)
+                cls._instance._managers = {}
+                cls._instance._registry_lock = threading.Lock()
+        return cls._instance
+
+    def _get_config_key(self, config: Dict[str, Any]) -> str:
+        critical_keys = [
+            "model_name", "device", "torch_dtype_str", "use_safetensors", 
+            "safety_checker_on", "hf_variant", "enable_cpu_offload", 
+            "enable_sequential_cpu_offload", "enable_xformers",
+            "local_files_only", "hf_cache_path"
+        ]
+        key_data = tuple(sorted((k, config.get(k)) for k in critical_keys))
+        return hashlib.sha256(str(key_data).encode('utf-8')).hexdigest()
+
+    def get_manager(self, config: Dict[str, Any], models_path: Path) -> ModelManager:
+        key = self._get_config_key(config)
+        with self._registry_lock:
+            if key not in self._managers:
+                self._managers[key] = ModelManager(config.copy(), models_path)
+            return self._managers[key].acquire()
+
+    def release_manager(self, config: Dict[str, Any]):
+        key = self._get_config_key(config)
+        with self._registry_lock:
+            if key in self._managers:
+                manager = self._managers[key]
+                ref_count = manager.release()
+                if ref_count == 0:
+                    ASCIIColors.info(f"Reference count for model '{config.get('model_name')}' is zero. Cleaning up.")
+                    manager.stop()
+                    manager._unload_pipeline()
+                    del self._managers[key]
+
+class DiffusersTTIBinding_Impl(LollmsTTIBinding):
+    DEFAULT_CONFIG = {
+        "model_name": "", "device": "auto", "torch_dtype_str": "auto", "use_safetensors": True,
+        "scheduler_name": "default", "safety_checker_on": True, "num_inference_steps": 25,
+        "guidance_scale": 7.0, "default_width": 512, "default_height": 512, "seed": -1,
+        "enable_cpu_offload": False, "enable_sequential_cpu_offload": False, "enable_xformers": False,
+        "hf_variant": None, "hf_token": None, "hf_cache_path": None, "local_files_only": False,
+    }
+
+    def __init__(self, **kwargs):
+        super().__init__(binding_name=BindingName)
+
+        if not DIFFUSERS_AVAILABLE:
+            raise ImportError(
+                "Diffusers or its dependencies not installed. "
+                "Please run: pip install torch torchvision diffusers Pillow transformers safetensors requests tqdm"
+            )
+
+        self.config = {**self.DEFAULT_CONFIG, **kwargs}
+        self.model_name = self.config.get("model_name", "")
+        self.models_path = Path(kwargs.get("models_path", Path(__file__).parent / "models"))
+        self.models_path.mkdir(parents=True, exist_ok=True)
+        
+        self.registry = PipelineRegistry()
+        self.manager: Optional[ModelManager] = None
+        
+        self._resolve_device_and_dtype()
+        if self.model_name:
+            self._acquire_manager()
+
+    def _acquire_manager(self):
+        if self.manager:
+            self.registry.release_manager(self.manager.config)
+        self.manager = self.registry.get_manager(self.config, self.models_path)
+        ASCIIColors.info(f"Binding instance acquired manager for '{self.config['model_name']}'.")
+
+    def _resolve_device_and_dtype(self):
+        if self.config["device"].lower() == "auto":
+            self.config["device"] = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+        
+        if self.config["torch_dtype_str"].lower() == "auto":
+            self.config["torch_dtype_str"] = "float16" if self.config["device"] != "cpu" else "float32"
+
+    def list_safetensor_models(self) -> List[str]:
+        if not self.models_path.exists(): return []
+        return sorted([f.name for f in self.models_path.iterdir() if f.is_file() and f.suffix == ".safetensors"])
+
+    def listModels(self) -> list:
+        # Start with hardcoded Civitai and Hugging Face models
+        civitai_list = [
+            {'model_name': key, 'display_name': info['display_name'], 'description': info['description'], 'owned_by': info['owned_by']}
+            for key, info in CIVITAI_MODELS.items()
+        ]
+        hf_default_list = [
+            # SDXL Models (1024x1024 native)
+            {'model_name': "stabilityai/stable-diffusion-xl-base-1.0", 'display_name': "Stable Diffusion XL 1.0", 'description': "Official SDXL base model from Stability AI. Native resolution is 1024x1024.", 'owned_by': 'HuggingFace'},
+            {'model_name': "playgroundai/playground-v2.5-1024px-aesthetic", 'display_name': "Playground v2.5", 'description': "Known for high aesthetic quality. Native resolution is 1024x1024.", 'owned_by': 'HuggingFace'},
+            # SD 1.5 Models (512x512 native)
+            {'model_name': "runwayml/stable-diffusion-v1-5", 'display_name': "Stable Diffusion 1.5", 'description': "A popular and versatile open-access text-to-image model.", 'owned_by': 'HuggingFace'},
+            {'model_name': "dataautogpt3/OpenDalleV1.1", 'display_name': "OpenDalle v1.1", 'description': "An open-source reproduction of DALL-E 3, good for prompt adherence.", 'owned_by': 'HuggingFace'},
+            {'model_name': "stabilityai/stable-diffusion-2-1-base", 'display_name': "Stable Diffusion 2.1 (512px)", 'description': "A 512x512 resolution model from Stability AI.", 'owned_by': 'HuggingFace'},
+            {'model_name': "CompVis/stable-diffusion-v1-4", 'display_name': "Stable Diffusion 1.4 (Gated)", 'description': "Original SD v1.4. Requires accepting license on Hugging Face and an HF token.", 'owned_by': 'HuggingFace'}
+        ]
+        
+        # Discover local .safetensors files
+        custom_local_models = []
+        civitai_filenames = {info['filename'] for info in CIVITAI_MODELS.values()}
+        local_safetensors = self.list_safetensor_models()
+
+        for filename in local_safetensors:
+            if filename not in civitai_filenames:
+                custom_local_models.append({
+                    'model_name': filename,
+                    'display_name': filename,
+                    'description': 'Local safetensors file from your models folder.',
+                    'owned_by': 'local_user'
+                })
+
+        return civitai_list + hf_default_list + custom_local_models
+
+    def load_model(self):
+        ASCIIColors.info("load_model() called. Loading is now automatic.")
+        if self.model_name and not self.manager:
+             self._acquire_manager()
 
     def unload_model(self):
-        if self.pipeline is not None:
-            del self.pipeline
-            self.pipeline = None
-            ASCIIColors.info("Diffusers pipeline unloaded.")
-        if torch and torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    def generate_image(self,
-                       prompt: str,
-                       negative_prompt: Optional[str] = "",
-                       width: Optional[int] = None, # Uses default from config if None
-                       height: Optional[int] = None, # Uses default from config if None
-                       **kwargs) -> bytes:
-        """
-        Generates image data using the Diffusers pipeline.
-
-        Args:
-            prompt (str): The positive text prompt.
-            negative_prompt (Optional[str]): The negative prompt.
-            width (int): Image width. Overrides default.
-            height (int): Image height. Overrides default.
-            **kwargs: Additional parameters for the pipeline:
-                      - num_inference_steps (int)
-                      - guidance_scale (float)
-                      - seed (int)
-                      - eta (float, for DDIM)
-                      - num_images_per_prompt (int, though this binding returns one)
-                      - clip_skip (int, if supported by pipeline - advanced)
-        Returns:
-            bytes: The generated image data (PNG format).
-        Raises:
-            Exception: If the request fails or image generation fails.
-        """
-        if not self.pipeline:
-            raise RuntimeError("Diffusers pipeline is not loaded. Cannot generate image.")
-
-        # Use call-specific or configured defaults
-        _width = width if width is not None else self.config["default_width"]
-        _height = height if height is not None else self.config["default_height"]
+        if self.manager:
+            ASCIIColors.info(f"Binding instance releasing manager for '{self.manager.config['model_name']}'.")
+            self.registry.release_manager(self.manager.config)
+            self.manager = None
+
+    def generate_image(self, prompt: str, negative_prompt: str = "", width: int = None, height: int = None, **kwargs) -> bytes:
+        if not self.model_name:
+            raise RuntimeError("No model_name configured. Please select a model in settings.")
+        
+        if not self.manager:
+            self._acquire_manager()
+        
+        _width = width or self.config["default_width"]
+        _height = height or self.config["default_height"]
         _num_inference_steps = kwargs.get("num_inference_steps", self.config["num_inference_steps"])
         _guidance_scale = kwargs.get("guidance_scale", self.config["guidance_scale"])
         _seed = kwargs.get("seed", self.config["seed"])
 
-        generator = None
-        if _seed != -1: # -1 means random seed
-            generator = torch.Generator(device=self.config["device"]).manual_seed(_seed)
+        generator = torch.Generator(device=self.config["device"]).manual_seed(_seed) if _seed != -1 else None
         
         pipeline_args = {
-            "prompt": prompt,
-            "negative_prompt": negative_prompt if negative_prompt else None,
-            "width": _width,
-            "height": _height,
-            "num_inference_steps": _num_inference_steps,
-            "guidance_scale": _guidance_scale,
-            "generator": generator,
-            "num_images_per_prompt": kwargs.get("num_images_per_prompt", 1)
+            "prompt": prompt, "negative_prompt": negative_prompt or None, "width": _width,
+            "height": _height, "num_inference_steps": _num_inference_steps,
+            "guidance_scale": _guidance_scale, "generator": generator,
         }
-        if "eta" in kwargs: pipeline_args["eta"] = kwargs["eta"]
-        if "clip_skip" in kwargs and hasattr(self.pipeline, "clip_skip"): # Handle clip_skip if supported
-            pipeline_args["clip_skip"] = kwargs["clip_skip"]
-
-
-        ASCIIColors.info(f"Generating image with prompt: '{prompt[:100]}...'")
-        ASCIIColors.debug(f"Pipeline args: {pipeline_args}")
-
+        
+        future = Future()
+        self.manager.queue.put((future, pipeline_args))
+        ASCIIColors.info(f"Job for prompt '{prompt[:50]}...' queued. Waiting...")
+        
         try:
-            with torch.no_grad(): # Important for inference
-                 pipeline_output = self.pipeline(**pipeline_args)
-            
-            pil_image: Image.Image = pipeline_output.images[0]
-
-            # Convert PIL Image to bytes (PNG)
-            img_byte_arr = BytesIO()
-            pil_image.save(img_byte_arr, format="PNG")
-            img_bytes = img_byte_arr.getvalue()
-            
+            image_bytes = future.result()
             ASCIIColors.green("Image generated successfully.")
-            return img_bytes
-
+            return image_bytes
         except Exception as e:
-            trace_exception(e)
-            raise Exception(f"Diffusers image generation failed: {e}") from e
+            raise Exception(f"Image generation failed: {e}") from e
+
+    def list_local_models(self) -> List[str]:
+        if not self.models_path.exists(): return []
+        
+        folders = [
+            d.name for d in self.models_path.iterdir()
+            if d.is_dir() and ((d / "model_index.json").exists() or (d / "unet" / "config.json").exists())
+        ]
+        safetensors = self.list_safetensor_models()
+        return sorted(folders + safetensors)
+    
+    def list_available_models(self) -> List[str]:
+        discoverable_models = [m['model_name'] for m in self.listModels()]
+        local_models = self.list_local_models()
+        
+        combined_list = sorted(list(set(local_models + discoverable_models)))
+        return combined_list
 
     def list_services(self, **kwargs) -> List[Dict[str, str]]:
-        """
-        Lists the currently loaded model as the available service.
-        Future: Could scan local model directories or list known HF models.
-        """
-        if self.pipeline and self.current_model_id_or_path:
-            return [{
-                "name": os.path.basename(self.current_model_id_or_path),
-                "caption": f"Diffusers: {os.path.basename(self.current_model_id_or_path)}",
-                "help": (f"Currently loaded model. Path/ID: {self.current_model_id_or_path}. "
-                         f"Device: {self.config['device']}. DType: {self.config['torch_dtype_str']}. "
-                         f"Scheduler: {self.pipeline.scheduler.__class__.__name__}.")
-            }]
-        return [{"name": "diffusers_unloaded", "caption": "No Diffusers model loaded", "help": "Configure a model in settings."}]
+        models = self.list_available_models()
+        local_models = self.list_local_models()
+
+        if not models:
+            return [{"name": "diffusers_no_models", "caption": "No models found", "help": f"Place models in '{self.models_path.resolve()}'."}]
+
+        services = []
+        for m in models:
+            help_text = "Hugging Face model ID"
+            if m in local_models:
+                 help_text = f"Local model from: {self.models_path.resolve()}"
+            elif m in CIVITAI_MODELS:
+                filename = CIVITAI_MODELS[m]['filename']
+                help_text = f"Civitai model (downloads as {filename})"
+            
+            services.append({"name": m, "caption": f"Diffusers: {m}", "help": help_text})
+        return services
 
     def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
-        """
-        Retrieves the current configurable settings for the Diffusers binding.
-        """
-        # Actual device and dtype after auto-resolution
-        resolved_device = self.config['device']
-        resolved_dtype_str = self.config['torch_dtype_str']
-        
-        # For display, show the original 'auto' if it was set that way, plus the resolved value
-        display_device = self.config['device'] if self.config['device'].lower() != 'auto' else f"auto ({resolved_device})"
-        display_dtype = self.config['torch_dtype_str'] if self.config['torch_dtype_str'].lower() != 'auto' else f"auto ({resolved_dtype_str})"
-
-        settings = [
-            {"name": "model_id_or_path", "type": "str", "value": self.config["model_id_or_path"], "description": "Hugging Face model ID or local path to Diffusers model directory."},
-            {"name": "device", "type": "str", "value": self.config["device"], "description": f"Device for inference. Current resolved: {resolved_device}", "options": ["auto", "cuda", "mps", "cpu"]},
-            {"name": "torch_dtype_str", "type": "str", "value": self.config["torch_dtype_str"], "description": f"Torch dtype for model. Current resolved: {resolved_dtype_str}", "options": ["auto", "float16", "bfloat16", "float32"]},
-            {"name": "hf_variant", "type": "str", "value": self.config["hf_variant"], "description": "Model variant (e.g., 'fp16', 'bf16'). Optional."},
-            {"name": "use_safetensors", "type": "bool", "value": self.config["use_safetensors"], "description": "Prefer loading models from .safetensors files."},
-            {"name": "scheduler_name", "type": "str", "value": self.config["scheduler_name"], "description": "Scheduler to use for diffusion.", "options": list(SCHEDULER_MAPPING.keys())},
-            {"name": "safety_checker_on", "type": "bool", "value": self.config["safety_checker_on"], "description": "Enable the safety checker (if model has one)."},
+        available_models = self.list_available_models()
+        return [
+            {"name": "model_name", "type": "str", "value": self.model_name, "description": "Local, Civitai, or Hugging Face model.", "options": available_models},
+            {"name": "device", "type": "str", "value": self.config["device"], "description": f"Inference device. Resolved: {self.config['device']}", "options": ["auto", "cuda", "mps", "cpu"]},
+            {"name": "torch_dtype_str", "type": "str", "value": self.config["torch_dtype_str"], "description": f"Torch dtype. Resolved: {self.config['torch_dtype_str']}", "options": ["auto", "float16", "bfloat16", "float32"]},
+            {"name": "hf_variant", "type": "str", "value": self.config["hf_variant"], "description": "HF model variant (e.g., 'fp16')."},
+            {"name": "use_safetensors", "type": "bool", "value": self.config["use_safetensors"], "description": "Prefer .safetensors when loading from Hugging Face."},
+            {"name": "scheduler_name", "type": "str", "value": self.config["scheduler_name"], "description": "Scheduler for diffusion.", "options": list(SCHEDULER_MAPPING.keys())},
+            {"name": "safety_checker_on", "type": "bool", "value": self.config["safety_checker_on"], "description": "Enable the safety checker."},
             {"name": "enable_cpu_offload", "type": "bool", "value": self.config["enable_cpu_offload"], "description": "Enable model CPU offload (saves VRAM, slower)."},
             {"name": "enable_sequential_cpu_offload", "type": "bool", "value": self.config["enable_sequential_cpu_offload"], "description": "Enable sequential CPU offload (more VRAM savings, much slower)."},
-            {"name": "enable_xformers", "type": "bool", "value": self.config["enable_xformers"], "description": "Enable xFormers memory efficient attention (if available)."},
-            {"name": "default_width", "type": "int", "value": self.config["default_width"], "description": "Default width for generated images."},
-            {"name": "default_height", "type": "int", "value": self.config["default_height"], "description": "Default height for generated images."},
-            {"name": "num_inference_steps", "type": "int", "value": self.config["num_inference_steps"], "description": "Default number of inference steps."},
+            {"name": "enable_xformers", "type": "bool", "value": self.config["enable_xformers"], "description": "Enable xFormers memory efficient attention."},
+            {"name": "default_width", "type": "int", "value": self.config["default_width"], "description": "Default image width. Note: SDXL models prefer 1024."},
+            {"name": "default_height", "type": "int", "value": self.config["default_height"], "description": "Default image height. Note: SDXL models prefer 1024."},
+            {"name": "num_inference_steps", "type": "int", "value": self.config["num_inference_steps"], "description": "Default inference steps."},
             {"name": "guidance_scale", "type": "float", "value": self.config["guidance_scale"], "description": "Default guidance scale (CFG)."},
-            {"name": "seed", "type": "int", "value": self.config["seed"], "description": "Default seed for generation (-1 for random)."},
-            {"name": "hf_token", "type": "str", "value": self.config["hf_token"], "description": "Hugging Face API token (for private/gated models). Set to 'None' or empty if not needed. Store securely.", "is_secret": True},
-            {"name": "local_files_only", "type": "bool", "value": self.config["local_files_only"], "description": "Only use local files, do not try to download."},
+            {"name": "seed", "type": "int", "value": self.config["seed"], "description": "Default seed (-1 for random)."},
+            {"name": "hf_token", "type": "str", "value": self.config["hf_token"], "description": "HF API token (for private/gated models).", "is_secret": True},
+            {"name": "hf_cache_path", "type": "str", "value": self.config["hf_cache_path"], "description": "Path to HF cache."},
+            {"name": "local_files_only", "type": "bool", "value": self.config["local_files_only"], "description": "Do not download from Hugging Face."},
         ]
-        return settings
 
     def set_settings(self, settings: Union[Dict[str, Any], List[Dict[str, Any]]], **kwargs) -> bool:
-        """
-        Applies new settings to the Diffusers binding. Some settings may trigger a model reload.
-        """
-        if isinstance(settings, list): # Convert from ConfigTemplate list format
-            parsed_settings = {item["name"]: item["value"] for item in settings if "name" in item and "value" in item}
-        elif isinstance(settings, dict):
-            parsed_settings = settings
-        else:
-            ASCIIColors.error("Invalid settings format. Expected a dictionary or list of dictionaries.")
-            return False
-
-        old_config = self.config.copy()
-        needs_reload = False
+        parsed_settings = settings if isinstance(settings, dict) else \
+                          {item["name"]: item["value"] for item in settings if "name" in item and "value" in item}
 
-        for key, value in parsed_settings.items():
-            if key in self.config:
-                if self.config[key] != value:
-                    self.config[key] = value
-                    ASCIIColors.info(f"Setting '{key}' changed to: {value}")
-                    if key in ["model_id_or_path", "device", "torch_dtype_str", 
-                               "use_safetensors", "safety_checker_on", "hf_variant",
-                               "enable_cpu_offload", "enable_sequential_cpu_offload", "enable_xformers",
-                               "hf_token", "local_files_only"]:
-                        needs_reload = True
-                    elif key == "scheduler_name" and self.pipeline: # Scheduler can be changed on loaded pipeline
-                        self._set_scheduler() # Attempt to apply immediately
-            else:
-                ASCIIColors.warning(f"Unknown setting '{key}' ignored.")
+        critical_keys = self.registry._get_config_key({}).__self__.critical_keys
+        needs_manager_swap = False
 
-        if needs_reload:
-            ASCIIColors.info("Reloading model due to settings changes...")
-            try:
-                # Resolve auto device/dtype again if they were part of the change
-                if "device" in parsed_settings and self.config["device"].lower() == "auto":
-                    if torch.cuda.is_available(): self.config["device"] = "cuda"
-                    elif torch.backends.mps.is_available(): self.config["device"] = "mps"
-                    else: self.config["device"] = "cpu"
-                
-                if "torch_dtype_str" in parsed_settings and self.config["torch_dtype_str"].lower() == "auto":
-                    self.config["torch_dtype_str"] = "float16" if self.config["device"] != "cpu" else "float32"
+        for key, value in parsed_settings.items():
+            if self.config.get(key) != value:
+                ASCIIColors.info(f"Setting '{key}' changed to: {value}")
+                self.config[key] = value
+                if key == "model_name": self.model_name = value
+                if key in critical_keys: needs_manager_swap = True
                 
-                # Update torch_dtype object from string
-                self.torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower(), torch.float32)
-                if self.torch_dtype == "auto": # Should be resolved by now
-                    self.torch_dtype = torch.float16 if self.config["device"] != "cpu" else torch.float32
-                    self.config["torch_dtype_str"] = TORCH_DTYPE_MAP_OBJ_TO_STR.get(self.torch_dtype, "float32")
-
+        if needs_manager_swap and self.model_name:
+            ASCIIColors.info("Critical settings changed. Swapping model manager...")
+            self._resolve_device_and_dtype()
+            self._acquire_manager()
+        
+        if not needs_manager_swap and self.manager:
+            self.manager.config.update(parsed_settings)
+            if 'scheduler_name' in parsed_settings and self.manager.pipeline:
+                 with self.manager.lock:
+                    self.manager._set_scheduler()
 
-                self.load_model()
-                ASCIIColors.green("Model reloaded successfully with new settings.")
-            except Exception as e:
-                trace_exception(e)
-                ASCIIColors.error(f"Failed to reload model with new settings: {e}. Reverting critical settings.")
-                # Revert critical settings and try to reload with old config
-                self.config = old_config 
-                self.torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower(), torch.float32)
-                try:
-                    self.load_model()
-                    ASCIIColors.info("Reverted to previous model configuration.")
-                except Exception as e_revert:
-                    trace_exception(e_revert)
-                    ASCIIColors.error(f"Failed to revert to previous model configuration: {e_revert}. Binding may be unstable.")
-                return False
         return True
 
     def __del__(self):
         self.unload_model()
 
-# Example Usage (for testing within this file)
+# Example Usage
 if __name__ == '__main__':
     ASCIIColors.magenta("--- Diffusers TTI Binding Test ---")
     
     if not DIFFUSERS_AVAILABLE:
-        ASCIIColors.error("Diffusers or its dependencies are not available. Cannot run test.")
-        # Attempt to guide user for installation
-        print("Please ensure PyTorch, Diffusers, Pillow, and Transformers are installed.")
-        print("For PyTorch with CUDA: visit https://pytorch.org/get-started/locally/")
-        print("Then: pip install diffusers Pillow transformers safetensors")
+        ASCIIColors.error("Diffusers not available. Cannot run test.")
         exit(1)
 
-    # --- Configuration ---
-    # Small, fast model for testing. Replace with a full model for real use.
-    # "CompVis/stable-diffusion-v1-4" is ~5GB
-    # "google/ddpm-cat-256" is smaller, but a DDPM, not Stable Diffusion.
-    # Using a tiny SD model if one exists, or a small variant.
-    # For a quick test, let's try a small LCM LoRA with SD1.5 if possible or just a base model.
-    # Note: "runwayml/stable-diffusion-v1-5" is a good standard test model.
-    # For a *very* quick CI-like test, one might use a dummy model or a very small one.
-    # Let's use a smaller SD variant if available, otherwise default to 2.1-base.
-    test_model_id = "runwayml/stable-diffusion-v1-5" # ~4GB download. Use a smaller one if you have it locally.
-    # test_model_id = "hf-internal-testing/tiny-stable-diffusion-pipe" # Very small, for testing structure
-    
-    # Create dummy lollms_paths
     temp_paths_dir = Path(__file__).parent / "temp_lollms_paths_diffusers"
-    temp_paths_dir.mkdir(parents=True, exist_ok=True)
-    mock_lollms_paths = {
-        "personal_models_path": temp_paths_dir / "personal_models",
-        "models_zoo_path": temp_paths_dir / "models_zoo",
-        "shared_cache_path": temp_paths_dir / "shared_cache", # For Hugging Face cache
-    }
-    for p in mock_lollms_paths.values(): Path(p).mkdir(parents=True, exist_ok=True)
-    (Path(mock_lollms_paths["personal_models_path"]) / "diffusers_models").mkdir(exist_ok=True)
-
-
-    binding_config = {
-        "model_id_or_path": test_model_id,
-        "device": "auto", # Let it auto-detect
-        "torch_dtype_str": "auto",
-        "num_inference_steps": 10, # Faster for testing
-        "default_width": 256, # Smaller for faster testing
-        "default_height": 256,
-        "safety_checker_on": False, # Often disabled for local use flexibility
-        "hf_variant": "fp16" if test_model_id == "runwayml/stable-diffusion-v1-5" else None, # SD 1.5 has fp16 variant
-    }
-
-    try:
-        ASCIIColors.cyan("\n1. Initializing DiffusersTTIBinding_Impl...")
-        binding = DiffusersTTIBinding_Impl(config=binding_config, lollms_paths=mock_lollms_paths)
-        ASCIIColors.green("Initialization successful.")
-        ASCIIColors.info(f"Loaded model: {binding.current_model_id_or_path}")
-        ASCIIColors.info(f"Device: {binding.config['device']}, DType: {binding.config['torch_dtype_str']}")
-        ASCIIColors.info(f"Scheduler: {binding.pipeline.scheduler.__class__.__name__ if binding.pipeline else 'N/A'}")
-
-
-        ASCIIColors.cyan("\n2. Listing services...")
-        services = binding.list_services()
-        ASCIIColors.info(json.dumps(services, indent=2))
-        assert services and services[0]["name"] == os.path.basename(binding.current_model_id_or_path)
-
-        ASCIIColors.cyan("\n3. Getting settings...")
-        settings_list = binding.get_settings()
-        ASCIIColors.info(json.dumps(settings_list, indent=2, default=str)) # default=str for Path objects if any
-        # Find model_id_or_path in settings
-        found_model_setting = any(s['name'] == 'model_id_or_path' and s['value'] == test_model_id for s in settings_list)
-        assert found_model_setting, "Model ID not found or incorrect in get_settings"
-
-
-        ASCIIColors.cyan("\n4. Generating an image...")
-        test_prompt = "A vibrant cat astronaut exploring a neon galaxy"
-        test_negative_prompt = "blurry, low quality, text, watermark"
+    temp_models_path = temp_paths_dir / "models"
+    
+    if temp_paths_dir.exists(): shutil.rmtree(temp_paths_dir)
+    temp_models_path.mkdir(parents=True, exist_ok=True)
         
-        # Use smaller dimensions for test if default are large
-        gen_width = min(binding.config["default_width"], 256)
-        gen_height = min(binding.config["default_height"], 256)
-
-        image_bytes = binding.generate_image(
-            prompt=test_prompt,
-            negative_prompt=test_negative_prompt,
-            width=gen_width, height=gen_height,
-            num_inference_steps=8 # Even fewer for speed
-        )
-        assert image_bytes and isinstance(image_bytes, bytes)
-        ASCIIColors.green(f"Image generated successfully (size: {len(image_bytes)} bytes).")
-        # Save the image for verification
-        test_image_path = Path(__file__).parent / "test_diffusers_image.png"
-        with open(test_image_path, "wb") as f:
-            f.write(image_bytes)
-        ASCIIColors.info(f"Test image saved to: {test_image_path.resolve()}")
-
-
-        ASCIIColors.cyan("\n5. Setting new settings (changing scheduler and guidance_scale)...")
-        new_settings_dict = {
-            "scheduler_name": "ddim", # Change scheduler
-            "guidance_scale": 5.0,   # Change guidance scale
-            "num_inference_steps": 12 # Change inference steps
-        }
-        binding.set_settings(new_settings_dict)
-        assert binding.config["scheduler_name"] == "ddim"
-        assert binding.config["guidance_scale"] == 5.0
-        assert binding.config["num_inference_steps"] == 12
-        ASCIIColors.info(f"New scheduler (intended): ddim, Actual: {binding.pipeline.scheduler.__class__.__name__}")
-        ASCIIColors.info(f"New guidance_scale: {binding.config['guidance_scale']}")
+    try:
+        ASCIIColors.cyan("\n--- Test: Loading a Hugging Face model ---")
+        # Using a very small model for fast testing
+        binding_config = {"models_path": str(temp_models_path), "model_name": "hf-internal-testing/tiny-stable-diffusion-torch"}
+        binding = DiffusersTTIBinding_Impl(**binding_config)
         
-        ASCIIColors.cyan("\n6. Generating another image with new settings...")
-        image_bytes_2 = binding.generate_image(
-            prompt="A serene landscape with a crystal river",
-            width=gen_width, height=gen_height
-        )
-        assert image_bytes_2 and isinstance(image_bytes_2, bytes)
-        ASCIIColors.green(f"Second image generated successfully (size: {len(image_bytes_2)} bytes).")
-        test_image_path_2 = Path(__file__).parent / "test_diffusers_image_2.png"
-        with open(test_image_path_2, "wb") as f:
-            f.write(image_bytes_2)
-        ASCIIColors.info(f"Second test image saved to: {test_image_path_2.resolve()}")
-
-        # Test model reload by changing a critical parameter (e.g. safety_checker_on)
-        # This requires a different model or a config that can be easily toggled.
-        # For now, assume reload on critical param change works if no error is thrown.
-        ASCIIColors.cyan("\n7. Testing settings change requiring model reload (safety_checker_on)...")
-        current_safety_on = binding.config["safety_checker_on"]
-        binding.set_settings({"safety_checker_on": not current_safety_on})
-        assert binding.config["safety_checker_on"] == (not current_safety_on)
-        ASCIIColors.green("Model reload due to safety_checker_on change seems successful.")
+        img_bytes = binding.generate_image("a tiny robot", width=64, height=64, num_inference_steps=2)
+        assert len(img_bytes) > 1000, "Image generation from HF model should succeed."
+        ASCIIColors.green("HF model loading and generation successful.")
 
+        del binding
+        time.sleep(0.1)
 
     except Exception as e:
         trace_exception(e)
         ASCIIColors.error(f"Diffusers binding test failed: {e}")
     finally:
-        ASCIIColors.cyan("\nCleaning up...")
-        if 'binding' in locals() and binding:
-            binding.unload_model()
-        
-        # Clean up temp_lollms_paths
-        import shutil
+        ASCIIColors.cyan("\nCleaning up temporary directories...")
         if temp_paths_dir.exists():
-            try:
-                shutil.rmtree(temp_paths_dir)
-                ASCIIColors.info(f"Cleaned up temporary directory: {temp_paths_dir}")
-            except Exception as e_clean:
-                ASCIIColors.warning(f"Could not fully clean up {temp_paths_dir}: {e_clean}")
-        if 'test_image_path' in locals() and test_image_path.exists():
-             # os.remove(test_image_path) # Keep for manual check
-             pass
-        if 'test_image_path_2' in locals() and test_image_path_2.exists():
-             # os.remove(test_image_path_2) # Keep for manual check
-             pass
+            shutil.rmtree(temp_paths_dir)
         ASCIIColors.magenta("--- Diffusers TTI Binding Test Finished ---")