lollms-client 1.7.10__py3-none-any.whl → 1.8.3__py3-none-any.whl

lollms_client/lollms_ttv_binding.py

@@ -2,7 +2,7 @@
 from abc import abstractmethod
 import importlib
 from pathlib import Path
-from typing import Optional, List, Dict, Any, Union
+from typing import Optional, List, Dict, Any, Union, Callable
 from ascii_colors import trace_exception
 from lollms_client.lollms_base_binding import LollmsBaseBinding
 
@@ -31,6 +31,20 @@ class LollmsTTVBinding(LollmsBaseBinding):
         """
         pass
 
+    def get_zoo(self) -> List[Dict[str, Any]]:
+        """
+        Returns a list of models available for download.
+        each entry is a dict with:
+        name, description, size, type, link
+        """
+        return []
+
+    def download_from_zoo(self, index: int, progress_callback: Callable[[dict], None] = None) -> dict:
+        """
+        Downloads a model from the zoo using its index.
+        """
+        return {"status": False, "message": "Not implemented"}
+
 class LollmsTTVBindingManager:
     """Manages TTV binding discovery and instantiation."""
 
@@ -62,7 +76,7 @@ class LollmsTTVBindingManager:
         binding_class = self.available_bindings.get(binding_name)
         if binding_class:
             try:
-                return binding_class(binding_name=binding_name, **kwargs)
+                return binding_class(**kwargs)
             except Exception as e:
                 trace_exception(e)
                 print(f"Failed to instantiate TTV binding {binding_name}: {str(e)}")

lollms_client/tti_bindings/diffusers/__init__.py

@@ -26,26 +26,20 @@ except ImportError:
 from lollms_client.lollms_tti_binding import LollmsTTIBinding
 from ascii_colors import ASCIIColors
 
-BindingName = "DiffusersBinding"
-
-class DiffusersBinding(LollmsTTIBinding):
-    """
-    Client binding for a dedicated, managed Diffusers server.
-    This architecture prevents multiple models from being loaded into memory
-    in a multi-worker environment, solving OOM errors.
-    """
-    def __init__(self,
-                 **kwargs):
+BindingName = "DiffusersTTIBinding"
+
+class DiffusersTTIBinding(LollmsTTIBinding):
+    def __init__(self, **kwargs):
         # Prioritize 'model_name' but accept 'model' as an alias from config files.
         if 'model' in kwargs and 'model_name' not in kwargs:
             kwargs['model_name'] = kwargs.pop('model')
         super().__init__(binding_name=BindingName, config=kwargs)
 
-
-
+        self.config = kwargs
         self.host = kwargs.get("host", "localhost")
         self.port = kwargs.get("port", 9632)
-        self.auto_start_server = kwargs.get("auto_start_server", True)
+        self.auto_start_server = kwargs.get("auto_start_server", False)
+        self.wait_for_server = kwargs.get("wait_for_server", False)
         self.server_process = None
         self.base_url = f"http://{self.host}:{self.port}"
         self.binding_root = Path(__file__).parent
@@ -53,9 +47,11 @@ class DiffusersBinding(LollmsTTIBinding):
         self.venv_dir = Path("./venv/tti_diffusers_venv")
         self.models_path = Path(kwargs.get("models_path", "./data/models/diffusers_models")).resolve()
         self.extra_models_path = kwargs.get("extra_models_path")
+        self.hf_token = kwargs.get("hf_token", "")  # NEW
         self.models_path.mkdir(exist_ok=True, parents=True)
         if self.auto_start_server:
-            self.ensure_server_is_running()
+            self.ensure_server_is_running(self.wait_for_server)
+
 
     def is_server_running(self) -> bool:
         """Checks if the server is already running and responsive."""
@@ -68,47 +64,21 @@ class DiffusersBinding(LollmsTTIBinding):
         return False
 
 
-    def ensure_server_is_running(self):
+    def ensure_server_is_running(self, wait= False):
         """
         Ensures the Diffusers server is running. If not, it attempts to start it
         in a process-safe manner using a file lock. This method is designed to
         prevent race conditions in multi-worker environments.
         """
         self.server_dir.mkdir(exist_ok=True)
-        # Use a lock file in the binding's server directory for consistency across instances
-        lock_path = self.server_dir / "diffusers_server.lock"
-        lock = FileLock(lock_path)
-
         ASCIIColors.info("Attempting to start or connect to the Diffusers server...")
 
         # First, perform a quick check without the lock to avoid unnecessary waiting.
         if self.is_server_running():
             ASCIIColors.green("Diffusers Server is already running and responsive.")
             return
-
-        try:
-            # Try to acquire the lock with a timeout. If another process is starting
-            # the server, this will wait until it's finished.
-            with lock.acquire(timeout=3):
-                # After acquiring the lock, we MUST re-check if the server is running.
-                # Another process might have started it and released the lock while we were waiting.
-                if not self.is_server_running():
-                    ASCIIColors.yellow("Lock acquired. Starting dedicated Diffusers server...")
-                    self.start_server()
-                    # The process that starts the server is responsible for waiting for it to be ready
-                    # BEFORE releasing the lock. This is the key to preventing race conditions.
-                    self._wait_for_server()
-                else:
-                    ASCIIColors.green("Server was started by another process while we waited. Connected successfully.")
-        except Timeout:
-            # This happens if the process holding the lock takes more than 60 seconds to start the server.
-            # We don't try to start another one. We just wait for the existing one to be ready.
-            ASCIIColors.yellow("Could not acquire lock, another process is taking a long time to start the server. Waiting...")
-            self._wait_for_server(timeout=60) # Give it a longer timeout here just in case.
-
-        # A final verification to ensure we are connected.
-        if not self.is_server_running():
-            raise RuntimeError("Failed to start or connect to the Diffusers server after all attempts.")
+        else:
+            self.start_server(wait)
 
     def install_server_dependencies(self):
         """
@@ -175,43 +145,70 @@ class DiffusersBinding(LollmsTTIBinding):
 
         ASCIIColors.green("Server dependencies are satisfied.")
 
-    def start_server(self):
+    def start_server(self, wait=True, timeout_s=20):
         """
-        Installs dependencies and launches the FastAPI server as a background subprocess.
+        Launches the FastAPI server in a background thread and returns immediately.
         This method should only be called from within a file lock.
         """
-        server_script = self.server_dir / "main.py"
-        if not server_script.exists():
-            # Fallback for old structure
-            server_script = self.binding_root / "server.py"
-            if not server_script.exists():
-                raise FileNotFoundError(f"Server script not found at {server_script}. Make sure it's in a 'server' subdirectory.")
-        if not self.venv_dir.exists():
-            self.install_server_dependencies()
-
-        if sys.platform == "win32":
-            python_executable = self.venv_dir / "Scripts" / "python.exe"
-        else:
-            python_executable = self.venv_dir / "bin" / "python"
-
-        command = [
-            str(python_executable),
-            str(server_script),
-            "--host", self.host,
-            "--port", str(self.port),
-            "--models-path", str(self.models_path.resolve()) # Pass models_path to server
-        ]
-
-        if self.extra_models_path:
-            resolved_extra_path = Path(self.extra_models_path).resolve()
-            command.extend(["--extra-models-path", str(resolved_extra_path)])
+        import threading
+        
 
-        # Use DETACHED_PROCESS on Windows to allow the server to run independently of the parent process.
-        # On Linux/macOS, the process will be daemonized enough to not be killed with the worker.
-        creationflags = subprocess.DETACHED_PROCESS if sys.platform == "win32" else 0
+        def _start_server_background():
+            """Helper method to start the server in a background thread."""
+            # Use a lock file in the binding's server directory for consistency across instances
+            lock_path = self.server_dir / "diffusers_server.lock"
+            lock = FileLock(lock_path)
+            with lock.acquire(timeout=0):
+                try:
+                    server_script = self.server_dir / "main.py"
+                    if not server_script.exists():
+                        # Fallback for old structure
+                        server_script = self.binding_root / "server.py"
+                        if not server_script.exists():
+                            raise FileNotFoundError(f"Server script not found at {server_script}. Make sure it's in a 'server' subdirectory.")
+                    if not self.venv_dir.exists():
+                        self.install_server_dependencies()
+
+                    if sys.platform == "win32":
+                        python_executable = self.venv_dir / "Scripts" / "python.exe"
+                    else:
+                        python_executable = self.venv_dir / "bin" / "python"
+
+                    command = [
+                        str(python_executable),
+                        str(server_script),
+                        "--host", self.host,
+                        "--port", str(self.port),
+                        "--models-path", str(self.models_path.resolve())
+                    ]
+
+                    if self.extra_models_path:
+                        resolved_extra_path = Path(self.extra_models_path).resolve()
+                        command.extend(["--extra-models-path", str(resolved_extra_path)])
+
+                    if self.hf_token:
+                        command.extend(["--hf-token", self.hf_token])
+
+                    if self.extra_models_path:
+                        resolved_extra_path = Path(self.extra_models_path).resolve()
+                        command.extend(["--extra-models-path", str(resolved_extra_path)])
+
+                    creationflags = subprocess.DETACHED_PROCESS if sys.platform == "win32" else 0
+                    self.server_process = subprocess.Popen(command, creationflags=creationflags)
+                    ASCIIColors.info("Diffusers server process launched in the background.")
+                    while(not self.is_server_running()):
+                        time.sleep(1)
+                    
+                except Exception as e:
+                    ASCIIColors.error(f"Failed to start Diffusers server: {e}")
+                    raise
+
+        # Start the server in a background thread
+        thread = threading.Thread(target=_start_server_background, daemon=True)
+        thread.start()
+        if wait:
+            thread.join()
 
-        self.server_process = subprocess.Popen(command, creationflags=creationflags)
-        ASCIIColors.info("Diffusers server process launched in the background.")
 
     def _wait_for_server(self, timeout=30):
         """Waits for the server to become responsive."""
@@ -288,6 +285,7 @@ class DiffusersBinding(LollmsTTIBinding):
         pass
 
     def generate_image(self, prompt: str, negative_prompt: str = "", **kwargs) -> bytes:
+        self.ensure_server_is_running(True)
         params = kwargs.copy()
         if "model_name" not in params and self.config.get("model_name"):
             params["model_name"] = self.config["model_name"]
@@ -300,6 +298,7 @@ class DiffusersBinding(LollmsTTIBinding):
         return response.content
 
     def edit_image(self, images: Union[str, List[str], "Image.Image", List["Image.Image"]], prompt: str, **kwargs) -> bytes:
+        self.ensure_server_is_running(True)
         images_b64 = []
         if not isinstance(images, list):
             images = [images]
@@ -342,24 +341,78 @@ class DiffusersBinding(LollmsTTIBinding):
         response = self._post_json_request("/edit_image", data=json_payload)
         return response.content
 
-    def list_models(self) -> List[Dict[str, Any]]:
-        return self._get_request("/list_models").json()
+    def list_models(self) -> list:
+        """
+        Lists only models that are available locally on disk.
+
+        The Diffusers server scans `models_path` and `extra_models_path` for:
+          - Diffusers pipeline folders (with model_index.json, etc.)
+          - .safetensors checkpoints and associated configs.
+
+        Returns list of dicts: {"model_name": str, "display_name": str, "description": str}
+        """
+        self.ensure_server_is_running(True)
+        try:
+            response = self._get_request("/list_models")
+            data = response.json()
+            if not isinstance(data, list):
+                return []
+            return data
+        except Exception as e:
+            ASCIIColors.warning(f"Failed to list local Diffusers models: {e}")
+            return []
+
 
     def list_local_models(self) -> List[str]:
+        self.ensure_server_is_running(True)
         return self._get_request("/list_local_models").json()
 
     def list_available_models(self) -> List[str]:
+        self.ensure_server_is_running(True)
         return self._get_request("/list_available_models").json()
 
     def list_services(self, **kwargs) -> List[Dict[str, str]]:
+        self.ensure_server_is_running(True)
         return self._get_request("/list_models").json()
 
     def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
+        self.ensure_server_is_running(True)
         # The server holds the state, so we fetch it.
         return self._get_request("/get_settings").json()
+    
+    def get_zoo(self):
+        return [
+            {"name": "Stable Diffusion 1.5", "description": "The classic and versatile SD1.5 base model.", "size": "4GB", "type": "checkpoint", "link": "runwayml/stable-diffusion-v1-5"},
+            {"name": "Stable Diffusion 2.1", "description": "The 768x768 base model from the SD2.x series.", "size": "5GB", "type": "checkpoint", "link": "stabilityai/stable-diffusion-2-1"},
+            {"name": "Stable Diffusion XL 1.0", "description": "Official 1024x1024 text-to-image model from Stability AI.", "size": "7GB", "type": "checkpoint", "link": "stabilityai/stable-diffusion-xl-base-1.0"},
+            {"name": "SDXL Turbo", "description": "A fast, real-time text-to-image model based on SDXL.", "size": "7GB", "type": "checkpoint", "link": "stabilityai/sdxl-turbo"},
+            {"name": "Kandinsky 3", "description": "A powerful multilingual model with strong prompt understanding.", "size": "Unknown", "type": "checkpoint", "link": "kandinsky-community/kandinsky-3"},
+            {"name": "Playground v2.5", "description": "A high-quality model focused on aesthetic outputs.", "size": "Unknown", "type": "checkpoint", "link": "playgroundai/playground-v2.5-1024px-aesthetic"},
+            {"name": "epiCRealism", "description": "A popular community model for generating photorealistic images.", "size": "2GB", "type": "checkpoint", "link": "emilianJR/epiCRealism"},
+            {"name": "Realistic Vision 5.1", "description": "One of the most popular realistic models, great for portraits.", "size": "2GB", "type": "checkpoint", "link": "SG161222/Realistic_Vision_V5.1_noVAE"},
+            {"name": "Photon", "description": "A model known for high-quality, realistic images with good lighting.", "size": "2GB", "type": "checkpoint", "link": "Photon-v1"},
+            {"name": "Waifu Diffusion 1.4", "description": "A widely-used model for generating high-quality anime-style images.", "size": "2GB", "type": "checkpoint", "link": "hakurei/waifu-diffusion"},
+            {"name": "Counterfeit V3.0", "description": "A strong model for illustrative and 2.5D anime styles.", "size": "2GB", "type": "checkpoint", "link": "gsdf/Counterfeit-V3.0"},
+            {"name": "Animagine XL 3.0", "description": "A state-of-the-art anime model on the SDXL architecture.", "size": "7GB", "type": "checkpoint", "link": "cagliostrolab/animagine-xl-3.0"},
+            {"name": "DreamShaper 8", "description": "Versatile SD1.5 style model (CivitAI).", "size": "2GB", "type": "checkpoint", "link": "https://civitai.com/api/download/models/128713"},
+            {"name": "Juggernaut XL", "description": "Artistic SDXL (CivitAI).", "size": "7GB", "type": "checkpoint", "link": "https://civitai.com/api/download/models/133005"},
+            {"name": "Stable Diffusion 3 Medium", "description": "SOTA model with advanced prompt understanding (Gated).", "size": "Unknown", "type": "checkpoint", "link": "stabilityai/stable-diffusion-3-medium-diffusers"},
+            {"name": "FLUX.1 Schnell", "description": "Powerful and fast next-gen model (Gated).", "size": "Unknown", "type": "checkpoint", "link": "black-forest-labs/FLUX.1-schnell"},
+            {"name": "FLUX.1 Dev", "description": "Larger developer version of FLUX.1 (Gated).", "size": "Unknown", "type": "checkpoint", "link": "black-forest-labs/FLUX.1-dev"},
+        ]
+
+    def download_from_zoo(self, index: int, progress_callback: Callable[[dict], None] = None) -> dict:
+        zoo = self.get_zoo()
+        if index < 0 or index >= len(zoo):
+            msg = "Index out of bounds"
+            ASCIIColors.error(msg)
+            return {"status": False, "message": msg}
+        item = zoo[index]
+        return self.pull_model(item["link"], progress_callback=progress_callback)
 
     def set_settings(self, settings: Union[Dict[str, Any], List[Dict[str, Any]]], **kwargs) -> bool:
-        # Normalize settings from list of dicts to a single dict if needed
+        self.ensure_server_is_running(True)
+            # Normalize settings from list of dicts to a single dict if needed
         parsed_settings = settings if isinstance(settings, dict) else {s["name"]: s["value"] for s in settings if "name" in s and "value" in s}
         response = self._post_json_request("/set_settings", data=parsed_settings)
         return response.json().get("success", False)
@@ -448,4 +501,4 @@ class DiffusersBinding(LollmsTTIBinding):
     def __del__(self):
         # The client destructor does not stop the server,
         # as it is a shared resource for all worker processes.
-        pass
+        pass

lollms_client/tti_bindings/diffusers/server/main.py

@@ -835,7 +835,7 @@ async def edit_image(request: EditRequestJSON):
 def pull_model_endpoint(payload: PullModelRequest):
     if not payload.hf_id and not payload.safetensors_url:
         raise HTTPException(status_code=400, detail="Provide either 'hf_id' or 'safetensors_url'.")
-
+    
     # 1) Pull Hugging Face model into a folder
     if payload.hf_id:
         model_id = payload.hf_id.strip()
@@ -902,83 +902,85 @@ def pull_model_endpoint(payload: PullModelRequest):
 @router.get("/list_local_models")
 def list_local_models_endpoint():
     local_models = set()
+    models_root = Path(args.models_path)
+    extra_root = Path(args.extra_models_path) if args.extra_models_path else None
 
-    # 1) Single-file models: top-level *.safetensors only
-    for f in state.models_path.glob("*.safetensors"):
-        local_models.add(f.name)
-
-    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)
-
-    # 2) Folder-based HF/diffusers models: treat folder name as the model
-    def add_folder_models(base: Path):
-        if not base or not base.exists():
+    def scan_root(root: Path):
+        if not root or not root.exists():
             return
-        for entry in base.iterdir():
-            if not entry.is_dir():
+        
+        # 1. Diffusers folders (Recursive)
+        for model_index in root.rglob("model_index.json"):
+             # For listing just the name, we probably want the folder name or relative path
+             # Keeping it simple: folder name.
+             local_models.add(model_index.parent.name)
+
+        # 2. Safetensors files (Recursive)
+        for safepath in root.rglob("*.safetensors"):
+            if (safepath.parent / "model_index.json").exists():
                 continue
-            has_index = (entry / "model_index.json").exists()
-            has_safetensors = any(entry.glob("*.safetensors"))
-            if has_index or has_safetensors:
-                local_models.add(entry.name)
+            local_models.add(safepath.name)
+            
+    scan_root(models_root)
+    scan_root(extra_root)
+    
+    return sorted(list(local_models))
 
-    add_folder_models(state.models_path)
-    add_folder_models(state.extra_models_path)
+@app.get("/list_models")
+def list_models() -> list[dict]:
+    models_root = Path(args.models_path)
+    extra_root = Path(args.extra_models_path) if args.extra_models_path else None
+    result = []
+    seen_paths = set()
 
-    return sorted(list(local_models))
+    def scan_root(root: Path):
+        if not root or not root.exists():
+            return
+        
+        # 1. Diffusers folders (Recursive)
+        # We look for model_index.json
+        for model_index in root.rglob("model_index.json"):
+             folder = model_index.parent
+             resolved_path = str(folder.resolve())
+             if resolved_path in seen_paths:
+                 continue
+             seen_paths.add(resolved_path)
+             
+             result.append({
+                "model_name": resolved_path,
+                "display_name": folder.name,
+                "description": "Local Diffusers pipeline"
+             })
+
+        # 2. Safetensors files (Recursive)
+        for safepath in root.rglob("*.safetensors"):
+            # Skip if part of a diffusers folder
+            if (safepath.parent / "model_index.json").exists():
+                continue
+            
+            resolved_path = str(safepath.resolve())
+            if resolved_path in seen_paths:
+                continue
+            seen_paths.add(resolved_path)
 
-@router.get("/list_models")
-def list_models_endpoint():
-    models = []
-
-    # 1) Local models - ensure dict format
-    local_files = list_local_models_endpoint()
-    for model_name in local_files:
-        models.append({
-            "model_name": model_name,
-            "display_name": model_name,
-            "description": "(Local) Folder model" if not model_name.endswith(".safetensors") else "(Local) Local safetensors file",
-            "owned_by": "local_user"
-        })
-
-    # 2) HF Public models - already dicts from HF_PUBLIC_MODELS
-    for category, hf_models in HF_PUBLIC_MODELS.items():
-        for model_info in hf_models:
-            models.append({
-                "model_name": model_info["model_name"],
-                "display_name": model_info["display_name"],
-                "description": f"({category}) {model_info['desc']}",
-                "owned_by": "huggingface"
+            result.append({
+                "model_name": resolved_path,
+                "display_name": safepath.stem,
+                "description": "Local .safetensors checkpoint"
             })
 
-    # 3) Gated models - same
-    if state.config.get("hf_token"):
-        for category, gated_models in HF_GATED_MODELS.items():
-            for model_info in gated_models:
-                models.append({
-                    "model_name": model_info["model_name"],
-                    "display_name": model_info["display_name"],
-                    "description": f"({category}) {model_info['desc']}",
-                    "owned_by": "huggingface"
-                })
+    scan_root(models_root)
+    scan_root(extra_root)
+    return result
 
-    # 4) Civitai models - ensure dict format
-    for key, info in CIVITAI_MODELS.items():
-        models.append({
-            "model_name": key,
-            "display_name": info["display_name"],
-            "description": f"(Civitai) {info['description']}",
-            "owned_by": info["owned_by"]
-        })
-
-    return models  # Plain list of dicts - JSON serializable
 
 
 
 @router.get("/list_available_models")
 def list_available_models_endpoint():
-    discoverable = [m['model_name'] for m in list_models_endpoint()]
+    # Use list_models() to get all available models (dicts) then extract names
+    models_dicts = list_models()
+    discoverable = [m['model_name'] for m in models_dicts]
     return sorted(list(set(discoverable)))
 
 @router.get("/get_settings")
@@ -1032,12 +1034,21 @@ if __name__ == "__main__":
     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.")
+    parser.add_argument(
+        "--hf-token",
+        type=str,
+        default=None,
+        help="Optional Hugging Face access token used to download private or gated repos."
+    )
+
     args = parser.parse_args()
 
     MODELS_PATH = Path(args.models_path)
     EXTRA_MODELS_PATH = Path(args.extra_models_path) if args.extra_models_path else None
     state = ServerState(MODELS_PATH, EXTRA_MODELS_PATH)
-    
+    if args.hf_token:
+        state.config["hf_token"] = args.hf_token
+        ASCIIColors.info("Hugging Face token received via CLI and stored in server config.")
     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()}")