lollms-client 1.6.4__py3-none-any.whl → 1.6.6__py3-none-any.whl

lollms_client/tti_bindings/diffusers/server/main.py

@@ -18,9 +18,12 @@ import gc
 import argparse
 import uvicorn
 from fastapi import FastAPI, APIRouter, HTTPException, UploadFile, Form
+from fastapi import Request, Response
 from fastapi.responses import Response
 from pydantic import BaseModel, Field
 import sys
+import platform
+import inspect
 
 # Add binding root to sys.path to ensure local modules can be imported if structured that way.
 binding_root = Path(__file__).resolve().parent.parent
@@ -53,13 +56,13 @@ except ImportError as e:
 # --- Server Setup ---
 app = FastAPI(title="Diffusers TTI Server")
 router = APIRouter()
-MODELS_PATH = Path("./models") # Default, will be overridden by command-line arg
+MODELS_PATH = Path("./models")
 
 # --- START: Core Logic (Complete and Unabridged) ---
 CIVITAI_MODELS = {
     "realistic-vision-v6": {
         "display_name": "Realistic Vision V6.0", "url": "https://civitai.com/api/download/models/501240?type=Model&format=SafeTensor&size=pruned&fp=fp16",
-        "filename": "realisticVisionV60_v60B1.safetensors", "description": "Photorealistic SD1.5 checkpoint.", "owned_by": "civitai"
+        "filename": "realisticVisionV60_v60B1.safensors", "description": "Photorealistic SD1.5 checkpoint.", "owned_by": "civitai"
     },
     "absolute-reality": {
         "display_name": "Absolute Reality", "url": "https://civitai.com/api/download/models/132760?type=Model&format=SafeTensor&size=pruned&fp=fp16",
@@ -119,6 +122,47 @@ CIVITAI_MODELS = {
     },
 }
 
+HF_PUBLIC_MODELS = {
+    "General Purpose & SDXL": [
+        {"model_name": "stabilityai/stable-diffusion-xl-base-1.0", "display_name": "Stable Diffusion XL 1.0", "desc": "Official 1024x1024 text-to-image model from Stability AI."},
+        {"model_name": "stabilityai/sdxl-turbo", "display_name": "SDXL Turbo", "desc": "A fast, real-time text-to-image model based on SDXL."},
+        {"model_name": "kandinsky-community/kandinsky-3", "display_name": "Kandinsky 3", "desc": "A powerful multilingual model with strong prompt understanding and aesthetic quality."},
+        {"model_name": "playgroundai/playground-v2.5-1024px-aesthetic", "display_name": "Playground v2.5", "desc": "A high-quality model focused on aesthetic outputs."},
+    ],
+    "Photorealistic": [
+        {"model_name": "emilianJR/epiCRealism", "display_name": "epiCRealism", "desc": "A popular community model for generating photorealistic images."},
+        {"model_name": "SG161222/Realistic_Vision_V5.1_noVAE", "display_name": "Realistic Vision 5.1", "desc": "One of the most popular realistic models, great for portraits and scenes."},
+        {"model_name": "Photon-v1", "display_name": "Photon", "desc": "A model known for high-quality, realistic images with good lighting and detail."},
+    ],
+    "Anime & Illustration": [
+        {"model_name": "hakurei/waifu-diffusion", "display_name": "Waifu Diffusion 1.4", "desc": "A widely-used model for generating high-quality anime-style images."},
+        {"model_name": "gsdf/Counterfeit-V3.0", "display_name": "Counterfeit V3.0", "desc": "A strong model for illustrative and 2.5D anime styles."},
+        {"model_name": "cagliostrolab/animagine-xl-3.0", "display_name": "Animagine XL 3.0", "desc": "A state-of-the-art anime model on the SDXL architecture."},
+    ],
+    "Artistic & Stylized": [
+        {"model_name": "wavymulder/Analog-Diffusion", "display_name": "Analog Diffusion", "desc": "Creates images with a vintage, analog film aesthetic."},
+        {"model_name": "dreamlike-art/dreamlike-photoreal-2.0", "display_name": "Dreamlike Photoreal 2.0", "desc": "Produces stunning, artistic, and photorealistic images."},
+    ],
+    "Image Editing Tools": [
+        {"model_name": "stabilityai/stable-diffusion-xl-refiner-1.0", "display_name": "SDXL Refiner 1.0", "desc": "A dedicated refiner model to improve details in SDXL generations."},
+        {"model_name": "Qwen/Qwen-Image-Edit", "display_name": "Qwen Image Edit", "desc": "An instruction-based model for various image editing tasks."},
+        {"model_name": "Qwen/Qwen-Image-Edit-2509", "display_name": "Qwen Image Edit Plus", "desc": "Advanced multi-image editing, fusion, and pose transfer."},
+    ],
+    "Legacy & Base Models": [
+        {"model_name": "runwayml/stable-diffusion-v1-5", "display_name": "Stable Diffusion 1.5", "desc": "The classic and versatile SD1.5 base model."},
+        {"model_name": "stabilityai/stable-diffusion-2-1", "display_name": "Stable Diffusion 2.1", "desc": "The 768x768 base model from the SD2.x series."},
+    ]
+}
+
+HF_GATED_MODELS = {
+    "Next-Generation (Gated Access Required)": [
+        {"model_name": "stabilityai/stable-diffusion-3-medium-diffusers", "display_name": "Stable Diffusion 3 Medium", "desc": "State-of-the-art model with advanced prompt understanding. Requires free registration."},
+        {"model_name": "black-forest-labs/FLUX.1-schnell", "display_name": "FLUX.1 Schnell", "desc": "A powerful and extremely fast next-generation model. Requires access request."},
+        {"model_name": "black-forest-labs/FLUX.1-dev", "display_name": "FLUX.1 Dev", "desc": "The larger developer version of the FLUX.1 model. Requires access request."},
+    ]
+}
+
+
 TORCH_DTYPE_MAP_STR_TO_OBJ = {
     "float16": getattr(torch, 'float16', 'float16'), "bfloat16": getattr(torch, 'bfloat16', 'bfloat16'),
     "float32": getattr(torch, 'float32', 'float32'), "auto": "auto"
@@ -141,6 +185,7 @@ SCHEDULER_USES_KARRAS_SIGMAS = [
     "dpm++_2m_sde_karras","dpm2_karras","dpm2_a_karras"
 ]
 
+
 class ModelManager:
     def __init__(self, config: Dict[str, Any], models_path: Path, registry: 'PipelineRegistry'):
         self.config = config
@@ -159,6 +204,7 @@ class ModelManager:
         self._stop_monitor_event = threading.Event()
         self._unload_monitor_thread = None
         self._start_unload_monitor()
+        self.supported_args: Optional[set] = None
 
     def acquire(self):
         with self.lock:
@@ -208,9 +254,25 @@ class ModelManager:
             if not local_path.exists():
                 self._download_civitai_model(model_name)
             return local_path
+
+        # Search in extra models path
+        if state.extra_models_path and state.extra_models_path.exists():
+            found_paths = list(state.extra_models_path.rglob(model_name))
+            if found_paths:
+                ASCIIColors.info(f"Found model in extra path: {found_paths[0]}")
+                return found_paths[0]
+
+        # Search in primary models path
+        found_paths = list(self.models_path.rglob(model_name))
+        if found_paths:
+            ASCIIColors.info(f"Found model in primary path: {found_paths[0]}")
+            return found_paths[0]
+
+        # Fallback for HF hub models that are folders, not single files.
         local_path = self.models_path / model_name
         if local_path.exists():
             return local_path
+
         return model_name
 
     def _download_civitai_model(self, model_key: str):
@@ -219,7 +281,7 @@ class ModelManager:
         filename = model_info["filename"]
         dest_path = self.models_path / filename
         temp_path = dest_path.with_suffix(".temp")
-        ASCIIColors.cyan(f"Downloading '{filename}' from Civitai...")
+        ASCIIColors.cyan(f"Downloading '{filename}' from Civitai... to {dest_path}")
         try:
             with requests.get(url, stream=True) as r:
                 r.raise_for_status()
@@ -233,13 +295,13 @@ class ModelManager:
         except Exception as e:
             if temp_path.exists():
                 temp_path.unlink()
-            raise Exception(f"Failed to download model {filename}: {e}") 
+            raise Exception(f"Failed to download model {filename}: {e}")
 
     def _set_scheduler(self):
         if not self.pipeline:
             return
-        if "Qwen" in self.config.get("model_name", ""):
-            ASCIIColors.info("Qwen model detected, skipping custom scheduler setup.")
+        if "Qwen" in self.config.get("model_name", "") or "FLUX" in self.config.get("model_name", ""):
+            ASCIIColors.info("Special model detected, skipping custom scheduler setup.")
             return
         scheduler_name_key = self.config["scheduler_name"].lower()
         if scheduler_name_key == "default":
@@ -256,71 +318,116 @@ class ModelManager:
                 ASCIIColors.warning(f"Could not switch scheduler to {scheduler_name_key}: {e}. Using current default.")
 
     def _execute_load_pipeline(self, task: str, model_path: Union[str, Path], torch_dtype: Any):
-        model_name = self.config.get("model_name", "")
+        if platform.system() == "Windows":
+            os.environ["HF_HUB_ENABLE_SYMLINKS"] = "0"
+        
+        model_name_from_config = self.config.get("model_name", "")
+        use_device_map = False
+
         try:
-            load_args = {}
+            load_params = {}
             if self.config.get("hf_cache_path"):
-                load_args["cache_dir"] = str(self.config["hf_cache_path"])
-            if str(model_path).endswith(".safetensors"):
-                if task == "text2image":
-                    try:
-                        self.pipeline = AutoPipelineForText2Image.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
-                    except AttributeError:
-                        self.pipeline = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
-                elif task == "image2image":
-                    self.pipeline = AutoPipelineForImage2Image.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
-                elif task == "inpainting":
-                    self.pipeline = AutoPipelineForInpainting.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
-            else:
-                common_args = {
-                    "torch_dtype": torch_dtype,
+                load_params["cache_dir"] = str(self.config["hf_cache_path"])
+            load_params["torch_dtype"] = torch_dtype
+
+            is_qwen_model = "Qwen" in model_name_from_config
+            is_flux_model = "FLUX" in model_name_from_config
+
+            if is_qwen_model or is_flux_model:
+                ASCIIColors.info(f"Special model '{model_name_from_config}' detected. Using dedicated pipeline loader.")
+                load_params.update({
                     "use_safetensors": self.config["use_safetensors"],
                     "token": self.config["hf_token"],
                     "local_files_only": self.config["local_files_only"]
-                }
+                })
                 if self.config["hf_variant"]:
-                    common_args["variant"] = self.config["hf_variant"]
+                    load_params["variant"] = self.config["hf_variant"]
                 if not self.config["safety_checker_on"]:
-                    common_args["safety_checker"] = None
-                if self.config.get("hf_cache_path"):
-                    common_args["cache_dir"] = str(self.config["hf_cache_path"])
-
-                if "Qwen-Image-Edit-2509" in str(model_path):
-                    self.pipeline = QwenImageEditPlusPipeline.from_pretrained(model_path, **common_args)
-                elif "Qwen-Image-Edit" in str(model_path):
-                    self.pipeline = QwenImageEditPipeline.from_pretrained(model_path, **common_args)
-                elif "Qwen/Qwen-Image" in str(model_path):
-                    self.pipeline = DiffusionPipeline.from_pretrained(model_path, **common_args)
-                elif task == "text2image":
-                    self.pipeline = AutoPipelineForText2Image.from_pretrained(model_path, **common_args)
-                elif task == "image2image":
-                    self.pipeline = AutoPipelineForImage2Image.from_pretrained(model_path, **common_args)
-                elif task == "inpainting":
-                    self.pipeline = AutoPipelineForInpainting.from_pretrained(model_path, **common_args)
+                    load_params["safety_checker"] = None
+                
+                should_offload = self.config["enable_cpu_offload"] or self.config["enable_sequential_cpu_offload"]
+                if should_offload:
+                    ASCIIColors.info(f"Offload enabled. Forcing device_map='auto' for {model_name_from_config}.")
+                    use_device_map = True
+                    load_params["device_map"] = "auto"
+                
+                if is_flux_model:
+                    self.pipeline = AutoPipelineForText2Image.from_pretrained(model_name_from_config, **load_params)
+                elif "Qwen-Image-Edit-2509" in model_name_from_config:
+                    self.pipeline = QwenImageEditPlusPipeline.from_pretrained(model_name_from_config, **load_params)
+                elif "Qwen-Image-Edit" in model_name_from_config:
+                    self.pipeline = QwenImageEditPipeline.from_pretrained(model_name_from_config, **load_params)
+                elif "Qwen/Qwen-Image" in model_name_from_config:
+                    self.pipeline = DiffusionPipeline.from_pretrained(model_name_from_config, **load_params)
+            
+            else:
+                is_safetensors_file = str(model_path).endswith(".safetensors")
+                if is_safetensors_file:
+                    ASCIIColors.info(f"Loading standard model from local .safetensors file: {model_path}")
+                    try:
+                        self.pipeline = AutoPipelineForText2Image.from_single_file(model_path, **load_params)
+                    except Exception as e:
+                        ASCIIColors.warning(f"Failed to load with AutoPipeline, falling back to StableDiffusionPipeline: {e}")
+                        self.pipeline = StableDiffusionPipeline.from_single_file(model_path, **load_params)
+                else:
+                    ASCIIColors.info(f"Loading standard model from Hub: {model_path}")
+                    load_params.update({
+                        "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_params["variant"] = self.config["hf_variant"]
+                    if not self.config["safety_checker_on"]:
+                        load_params["safety_checker"] = None
+                    
+                    is_large_model = "stable-diffusion-3" in str(model_path)
+                    should_offload = self.config["enable_cpu_offload"] or self.config["enable_sequential_cpu_offload"]
+                    if is_large_model and should_offload:
+                        ASCIIColors.info(f"Large model '{model_path}' detected with offload enabled. Using device_map='auto'.")
+                        use_device_map = True
+                        load_params["device_map"] = "auto"
+
+                    if task == "text2image":
+                        self.pipeline = AutoPipelineForText2Image.from_pretrained(model_path, **load_params)
+                    elif task == "image2image":
+                        self.pipeline = AutoPipelineForImage2Image.from_pretrained(model_path, **load_params)
+                    elif task == "inpainting":
+                        self.pipeline = AutoPipelineForInpainting.from_pretrained(model_path, **load_params)
+        
         except Exception as e:
             error_str = str(e).lower()
             if "401" in error_str or "gated" in error_str or "authorization" in error_str:
-                msg = (
-                    f"AUTHENTICATION FAILED for model '{model_name}'. "
-                    "Please ensure you accepted the model license and provided a valid HF token."
-                )
+                msg = (f"AUTHENTICATION FAILED for model '{model_name_from_config}'. Please ensure you accepted the model license and provided a valid HF token.")
                 raise RuntimeError(msg) 
             raise e
+
         self._set_scheduler()
-        self.pipeline.to(self.config["device"])
-        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()
+
+        if not use_device_map:
+            self.pipeline.to(self.config["device"])
+            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()
+        else:
+             ASCIIColors.info("Device map handled device placement. Skipping manual pipeline.to() and offload calls.")
+
+        if self.pipeline:
+            sig = inspect.signature(self.pipeline.__call__)
+            self.supported_args = {p.name for p in sig.parameters.values()}
+            ASCIIColors.info(f"Pipeline supported arguments detected: {self.supported_args}")
+
         self.is_loaded = True
         self.current_task = task
         self.last_used_time = time.time()
-        ASCIIColors.green(f"Model '{model_name}' loaded successfully on '{self.config['device']}' for task '{task}'.")
+        ASCIIColors.green(f"Model '{model_name_from_config}' loaded successfully using '{'device_map' if use_device_map else 'standard'}' mode for task '{task}'.")
 
     def _load_pipeline_for_task(self, task: str):
         if self.pipeline and self.current_task == task:
@@ -346,10 +453,7 @@ class ModelManager:
         
         ASCIIColors.warning(f"Failed to load '{model_name}' due to OOM. Attempting to unload other models to free VRAM.")
         
-        candidates_to_unload = [
-            m for m in self.registry.get_all_managers()
-            if m is not self and m.is_loaded
-        ]
+        candidates_to_unload = [m for m in self.registry.get_all_managers() if m is not self and m.is_loaded]
         candidates_to_unload.sort(key=lambda m: m.last_used_time)
 
         if not candidates_to_unload:
@@ -378,6 +482,7 @@ class ModelManager:
             model_name = self.config.get('model_name', 'Unknown')
             del self.pipeline
             self.pipeline = None
+            self.supported_args = None
             gc.collect()
             if torch and torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -398,8 +503,16 @@ class ModelManager:
                         self.last_used_time = time.time()
                         if not self.is_loaded or self.current_task != task:
                             self._load_pipeline_for_task(task)
+                    
+                    if self.supported_args:
+                        filtered_args = {k: v for k, v in pipeline_args.items() if k in self.supported_args}
+                    else:
+                        ASCIIColors.warning("Supported argument set not found. Using unfiltered arguments.")
+                        filtered_args = pipeline_args
+
                     with torch.no_grad():
-                        output = self.pipeline(**pipeline_args)
+                        output = self.pipeline(**filtered_args)
+                    
                     pil = output.images[0]
                     buf = BytesIO()
                     pil.save(buf, format="PNG")
@@ -409,7 +522,6 @@ class ModelManager:
                     future.set_exception(e)
                 finally:
                     self.queue.task_done()
-                    # Aggressive cleanup
                     if output is not None:
                         del output
                     gc.collect()
@@ -465,24 +577,25 @@ class PipelineRegistry:
             return list(self._managers.values())
 
 class ServerState:
-    def __init__(self, models_path: Path):
+    def __init__(self, models_path: Path, extra_models_path: Optional[Path] = None):
         self.models_path = models_path
+        self.extra_models_path = extra_models_path
         self.models_path.mkdir(parents=True, exist_ok=True)
+        if self.extra_models_path:
+            self.extra_models_path.mkdir(parents=True, exist_ok=True)
         self.config_path = self.models_path.parent / "diffusers_server_config.json"
         self.registry = PipelineRegistry()
         self.manager: Optional[ModelManager] = None
         self.config = {}
-        self.load_config() # This will set self.config
+        self.load_config()
         self._resolve_device_and_dtype()
-
-        # Eagerly acquire manager at startup if a model is configured
         if self.config.get("model_name"):
             try:
                 ASCIIColors.info(f"Acquiring initial model manager for '{self.config['model_name']}' on startup.")
                 self.manager = self.registry.get_manager(self.config, self.models_path)
             except Exception as e:
                 ASCIIColors.error(f"Failed to acquire model manager on startup: {e}")
-                self.manager = None # Ensure manager is None on failure
+                self.manager = None
 
     def get_default_config(self) -> Dict[str, Any]:
         return {
@@ -495,7 +608,6 @@ class ServerState:
         }
 
     def save_config(self):
-        """Saves the current configuration to a JSON file."""
         try:
             with open(self.config_path, 'w') as f:
                 json.dump(self.config, f, indent=4)
@@ -504,13 +616,11 @@ class ServerState:
             ASCIIColors.error(f"Failed to save server config: {e}")
 
     def load_config(self):
-        """Loads configuration from JSON file, falling back to defaults."""
         default_config = self.get_default_config()
         if self.config_path.exists():
             try:
                 with open(self.config_path, 'r') as f:
                     loaded_config = json.load(f)
-                # Merge loaded config into defaults to ensure all keys are present
                 default_config.update(loaded_config)
                 self.config = default_config
                 ASCIIColors.info(f"Loaded server configuration from {self.config_path}")
@@ -519,53 +629,45 @@ class ServerState:
                 self.config = default_config
         else:
             self.config = default_config
-        # Save back to ensure file exists and is up-to-date with all keys
         self.save_config()
 
     def _resolve_device_and_dtype(self):
         if self.config.get("device", "auto").lower() == "auto":
             self.config["device"] = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
         
-        # Prioritize bfloat16 for Qwen models on supported hardware, as it's more stable
-        if "Qwen" in self.config.get("model_name", "") and self.config["device"] == "cuda":
+        if ("Qwen" in self.config.get("model_name", "") or "FLUX" in self.config.get("model_name", "")) and self.config["device"] == "cuda":
             if hasattr(torch.cuda, 'is_bf16_supported') and torch.cuda.is_bf16_supported():
                 self.config["torch_dtype_str"] = "bfloat16"
-                ASCIIColors.info("Qwen model detected on compatible hardware. Forcing dtype to bfloat16 for stability.")
+                ASCIIColors.info("Special model detected on compatible hardware. Forcing dtype to bfloat16 for stability.")
                 return
 
         if self.config["torch_dtype_str"].lower() == "auto":
             self.config["torch_dtype_str"] = "float16" if self.config["device"] != "cpu" else "float32"
 
     def update_settings(self, new_settings: Dict[str, Any]):
-        """Updates settings, swaps the manager if critical settings change, and saves the config."""
         if 'model' in new_settings and 'model_name' not in new_settings:
             new_settings['model_name'] = new_settings.pop('model')
 
-        # Safeguard: If a model is already configured and the new settings don't specify one,
-        # keep the old one. This prevents a misconfigured client from wiping a valid server state.
         if self.config.get("model_name") and not new_settings.get("model_name"):
             ASCIIColors.info("Incoming settings have no model_name. Preserving existing model.")
             new_settings["model_name"] = self.config["model_name"]
 
-        # Release old manager if it exists
         if self.manager:
             self.registry.release_manager(self.manager.config)
             self.manager = None
 
-        # Update the config in memory
         self.config.update(new_settings)
         ASCIIColors.info(f"Server config updated. Current model_name: {self.config.get('model_name')}")
         
         self._resolve_device_and_dtype()
 
-        # Acquire new manager with the updated config
         if self.config.get("model_name"):
             ASCIIColors.info("Acquiring model manager with updated configuration...")
             self.manager = self.registry.get_manager(self.config, self.models_path)
         else:
             ASCIIColors.warning("No model_name in config after update, manager not acquired.")
         
-        self.save_config() # Persist the new state
+        self.save_config()
         return True
 
     def get_active_manager(self) -> ModelManager:
@@ -586,77 +688,208 @@ class EditRequestPayload(BaseModel):
     image_paths: List[str] = Field(default_factory=list)
     params: Dict[str, Any] = Field(default_factory=dict)
 
+class EditRequestJSON(BaseModel):
+    prompt: str
+    images_b64: List[str] = Field(description="A list of Base64 encoded image strings.")
+    params: Dict[str, Any] = Field(default_factory=dict)
+def get_sanitized_request_for_logging(request_data: Any) -> Dict[str, Any]:
+    """
+    Takes a request object (Pydantic model or dict) and returns a 'safe' dictionary
+    for logging, with long base64 strings replaced by placeholders.
+    """
+    import copy
+
+    try:
+        if hasattr(request_data, 'model_dump'):
+            data = request_data.model_dump()
+        elif isinstance(request_data, dict):
+            data = copy.deepcopy(request_data)
+        else:
+            return {"error": "Unsupported data type for sanitization"}
+
+        # Sanitize the main list of images
+        if 'images_b64' in data and isinstance(data['images_b64'], list):
+            count = len(data['images_b64'])
+            data['images_b64'] = f"[<{count} base64 image(s) truncated>]"
+
+        # Sanitize a potential mask in the 'params' dictionary
+        if 'params' in data and isinstance(data.get('params'), dict):
+            if 'mask_image' in data['params'] and isinstance(data['params']['mask_image'], str):
+                original_len = len(data['params']['mask_image'])
+                data['params']['mask_image'] = f"[<base64 mask truncated, len={original_len}>]"
+        
+        return data
+    except Exception:
+        return {"error": "Failed to sanitize request data."}
+
 # --- API Endpoints ---
 @router.post("/generate_image")
 async def generate_image(request: T2IRequest):
+    manager = None
+    temp_config = None
     try:
-        manager = state.get_active_manager()
         params = request.params
-        seed = int(params.get("seed", state.config.get("seed", -1)))
+        
+        # Determine which model manager to use for this specific request
+        if "model_name" in params and params["model_name"]:
+            temp_config = state.config.copy()
+            temp_config["model_name"] = params.pop("model_name") # Remove from params to avoid being passed to pipeline
+            manager = state.registry.get_manager(temp_config, state.models_path)
+            ASCIIColors.info(f"Using per-request model: {temp_config['model_name']}")
+        else:
+            manager = state.get_active_manager()
+            ASCIIColors.info(f"Using session-configured model: {manager.config.get('model_name')}")
+
+        seed = int(params.get("seed", manager.config.get("seed", -1)))
         generator = None
         if seed != -1:
-            generator = torch.Generator(device=state.config["device"]).manual_seed(seed)
-
+            generator = torch.Generator(device=manager.config["device"]).manual_seed(seed)
+        
+        width = int(params.get("width", manager.config.get("width", 512)))
+        height = int(params.get("height", manager.config.get("height", 512)))
+        
         pipeline_args = {
-            "prompt": request.prompt, "negative_prompt": request.negative_prompt,
-            "width": int(params.get("width", state.config.get("width", 512))),
-            "height": int(params.get("height", state.config.get("height", 512))),
-            "num_inference_steps": int(params.get("num_inference_steps", state.config.get("num_inference_steps", 25))),
-            "guidance_scale": float(params.get("guidance_scale", state.config.get("guidance_scale", 7.0))),
+            "prompt": request.prompt,
+            "negative_prompt": request.negative_prompt,
+            "width": width,
+            "height": height,
+            "num_inference_steps": int(params.get("num_inference_steps", manager.config.get("num_inference_steps", 25))),
+            "guidance_scale": float(params.get("guidance_scale", manager.config.get("guidance_scale", 7.0))),
             "generator": generator
         }
+        pipeline_args.update(params)
+        
+        model_name = manager.config.get("model_name", "")
+        task = "text2image"
+        
+        if "Qwen-Image-Edit" in model_name:
+            rng_seed = seed if seed != -1 else None
+            rng = np.random.default_rng(seed=rng_seed)
+            random_pixels = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
+            placeholder_image = Image.fromarray(random_pixels, 'RGB')
+            pipeline_args["image"] = placeholder_image
+            pipeline_args["strength"] = float(params.get("strength", 1.0))
+            task = "image2image" 
         
         future = Future()
-        manager.queue.put((future,"text2image", pipeline_args))
+        manager.queue.put((future, task, pipeline_args))
         result_bytes = future.result()
         return Response(content=result_bytes, media_type="image/png")
     except Exception as e:
         trace_exception(e)
         raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        if temp_config and manager:
+            state.registry.release_manager(temp_config)
+            ASCIIColors.info(f"Released per-request model: {temp_config['model_name']}")
+
 
 @router.post("/edit_image")
-async def edit_image(json_payload: str = Form(...), files: List[UploadFile] = []):
+async def edit_image(request: EditRequestJSON):
+    manager = None
+    temp_config = None
     try:
-        data = EditRequestPayload.parse_raw(json_payload)
-        manager = state.get_active_manager()
+        params = request.params
+
+        if "model_name" in params and params["model_name"]:
+            temp_config = state.config.copy()
+            temp_config["model_name"] = params.pop("model_name")
+            manager = state.registry.get_manager(temp_config, state.models_path)
+            ASCIIColors.info(f"Using per-request model: {temp_config['model_name']}")
+        else:
+            manager = state.get_active_manager()
+            ASCIIColors.info(f"Using session-configured model: {manager.config.get('model_name')}")
+
+        model_name = manager.config.get("model_name", "")
 
         pil_images = []
-        for file in files:
-            contents = await file.read()
-            pil_images.append(Image.open(BytesIO(contents)).convert("RGB"))
-        
-        for path in data.image_paths:
-             pil_images.append(load_image(path).convert("RGB"))
+        for b64_string in request.images_b64:
+            b64_data = b64_string.split(";base64,")[1] if ";base64," in b64_string else b64_string
+            image_bytes = base64.b64decode(b64_data)
+            pil_images.append(Image.open(BytesIO(image_bytes)).convert("RGB"))
 
-        if not pil_images:
-            raise HTTPException(status_code=400, detail="No images provided for editing.")
+        if not pil_images: raise HTTPException(status_code=400, detail="No valid images provided.")
+
+        pipeline_args = {"prompt": request.prompt}
+        seed = int(params.get("seed", -1))
+        if seed != -1: pipeline_args["generator"] = torch.Generator(device=manager.config["device"]).manual_seed(seed)
+        
+        if "mask_image" in params and params["mask_image"]:
+            b64_mask = params["mask_image"]
+            b64_data = b64_mask.split(";base64,")[1] if ";base64," in b64_mask else b64_mask
+            mask_bytes = base64.b64decode(b64_data)
+            pipeline_args["mask_image"] = Image.open(BytesIO(mask_bytes)).convert("L")
         
-        task = "inpainting" if data.params.get("mask") else "image2image"
+        task = "inpainting" if "mask_image" in pipeline_args else "image2image"
+        
+        if "Qwen-Image-Edit-2509" in model_name:
+            task = "image2image"
+            pipeline_args.update({"true_cfg_scale": 4.0, "guidance_scale": 1.0, "num_inference_steps": 40, "negative_prompt": " "})
+            edit_mode = params.get("edit_mode", "fusion")
+            if edit_mode == "fusion": pipeline_args["image"] = pil_images
+        else:
+            pipeline_args.update({"image": pil_images[0], "strength": 0.8, "guidance_scale": 7.5, "num_inference_steps": 25})
 
-        pipeline_args = {
-            "prompt": data.prompt,
-            "image": pil_images[0], # Simple i2i for now
-            "strength": float(data.params.get("strength", 0.8)),
-            # Add other params like mask etc.
-        }
+        pipeline_args.update(params)
         
-        future = Future()
-        manager.queue.put((future, task, pipeline_args))
-        result_bytes = future.result()
-        return Response(content=result_bytes, media_type="image/png")
+        future = Future(); manager.queue.put((future, task, pipeline_args))
+        return Response(content=future.result(), media_type="image/png")
     except Exception as e:
+        sanitized_payload = get_sanitized_request_for_logging(request)
+        ASCIIColors.error(f"Exception in /edit_image. Sanitized Payload: {json.dumps(sanitized_payload, indent=2)}")
         trace_exception(e)
         raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        if temp_config and manager:
+            state.registry.release_manager(temp_config)
+            ASCIIColors.info(f"Released per-request model: {temp_config['model_name']}")
+
 
 @router.get("/list_models")
 def list_models_endpoint():
-    civitai = [{'model_name': key, 'display_name': info['display_name'], 'description': info['description'], 'owned_by': info['owned_by']} for key, info in CIVITAI_MODELS.items()]
-    local = [{'model_name': f.name, 'display_name': f.stem, 'description': 'Local safetensors file.', 'owned_by': 'local_user'} for f in state.models_path.glob("*.safetensors")]
-    return civitai + local
+    huggingface_models = []
+    # Add public models, organized by category
+    for category, models in HF_PUBLIC_MODELS.items():
+        for model_info in models:
+            huggingface_models.append({
+                'model_name': model_info['model_name'], 
+                'display_name': model_info['display_name'], 
+                'description': f"({category}) {model_info['desc']}", 
+                'owned_by': 'huggingface'
+            })
+
+    # Conditionally add gated models if an HF token is provided in the server config
+    if state.config.get("hf_token"):
+        ASCIIColors.info("HF token detected, including gated models in the list.")
+        for category, models in HF_GATED_MODELS.items():
+            for model_info in models:
+                huggingface_models.append({
+                    'model_name': model_info['model_name'], 
+                    'display_name': model_info['display_name'], 
+                    'description': f"({category}) {model_info['desc']}", 
+                    'owned_by': 'huggingface'
+                })
+    else:
+        ASCIIColors.info("No HF token found, showing public models only.")
+
+    civitai_models = [{'model_name': key, 'display_name': info['display_name'], 'description': f"(Civitai) {info['description']}", 'owned_by': info['owned_by']} for key, info in CIVITAI_MODELS.items()]
+    
+    local_files = list_local_models_endpoint()
+    local_models = [{'model_name': filename, 'display_name': Path(filename).stem, 'description': '(Local) Local safetensors file.', 'owned_by': 'local_user'} for filename in local_files]
+
+    return huggingface_models + civitai_models + local_models
 
 @router.get("/list_local_models")
 def list_local_models_endpoint():
-    return sorted([f.name for f in state.models_path.glob("*.safetensors")])
+    local_models = set()
+    # Main models path
+    for f in state.models_path.glob("**/*.safetensors"):
+        local_models.add(f.name)
+    # Extra models path
+    if state.extra_models_path and state.extra_models_path.exists():
+        for f in state.extra_models_path.glob("**/*.safetensors"):
+            local_models.add(f.name)
+    return sorted(list(local_models))
 
 @router.get("/list_available_models")
 def list_available_models_endpoint():
@@ -666,7 +899,6 @@ def list_available_models_endpoint():
 @router.get("/get_settings")
 def get_settings_endpoint():
     settings_list = []
-    # Add options for dropdowns
     available_models = list_available_models_endpoint()
     schedulers = list(SCHEDULER_MAPPING.keys())
     config_to_display = state.config or state.get_default_config()
@@ -714,14 +946,18 @@ if __name__ == "__main__":
     parser.add_argument("--host", type=str, default="localhost", help="Host to bind to.")
     parser.add_argument("--port", type=int, default=9630, help="Port to bind to.")
     parser.add_argument("--models-path", type=str, required=True, help="Path to the models directory.")
+    parser.add_argument("--extra-models-path", type=str, default=None, help="Path to an extra models directory.")
     args = parser.parse_args()
 
     MODELS_PATH = Path(args.models_path)
-    state = ServerState(MODELS_PATH)
+    EXTRA_MODELS_PATH = Path(args.extra_models_path) if args.extra_models_path else None
+    state = ServerState(MODELS_PATH, EXTRA_MODELS_PATH)
     
     ASCIIColors.cyan(f"--- Diffusers TTI Server ---")
     ASCIIColors.green(f"Starting server on http://{args.host}:{args.port}")
     ASCIIColors.green(f"Serving models from: {MODELS_PATH.resolve()}")
+    if EXTRA_MODELS_PATH:
+        ASCIIColors.green(f"Serving extra models from: {EXTRA_MODELS_PATH.resolve()}")
     if not DIFFUSERS_AVAILABLE:
         ASCIIColors.error("Diffusers or its dependencies are not installed correctly in the server's environment!")
     else: