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

lollms_client/__init__.py

@@ -8,7 +8,7 @@ from lollms_client.lollms_utilities import PromptReshaper # Keep general utiliti
 from lollms_client.lollms_mcp_binding import LollmsMCPBinding, LollmsMCPBindingManager
 from lollms_client.lollms_llm_binding import LollmsLLMBindingManager
 
-__version__ = "1.6.4" # Updated version
+__version__ = "1.6.6" # Updated version
 
 # Optionally, you could define __all__ if you want to be explicit about exports
 __all__ = [

lollms_client/lollms_core.py

@@ -1,4 +1,6 @@
 # lollms_client/lollms_core.py
+# author: ParisNeo
+# description: LollmsClient definition file
 import requests
 from ascii_colors import ASCIIColors, trace_exception
 from lollms_client.lollms_types import MSG_TYPE, ELF_COMPLETION_FORMAT
@@ -519,7 +521,7 @@ class LollmsClient():
             Union[str, dict]: Generated text or error dictionary if failed.
         """
         if self.llm:
-            
+            images = [str(image) for image in images] if images else None
             ctx_size = ctx_size if ctx_size is not None else self.llm.default_ctx_size if self.llm.default_ctx_size else None
             if ctx_size is None:
                 ctx_size = self.llm.get_ctx_size()

lollms_client/tti_bindings/diffusers/__init__.py

@@ -44,15 +44,17 @@ class DiffusersBinding(LollmsTTIBinding):
             kwargs['model_name'] = kwargs.pop('model')
 
         self.config = kwargs
-        self.host = kwargs.get("host", "localhost") 
+        self.host = kwargs.get("host", "localhost")
         self.port = kwargs.get("port", 9630)
         self.auto_start_server = kwargs.get("auto_start_server", True)
         self.server_process = None
         self.base_url = f"http://{self.host}:{self.port}"
         self.binding_root = Path(__file__).parent
         self.server_dir = self.binding_root / "server"
-        self.venv_dir = Path("./venv/tti_diuffusers_venv")
-        self.models_path = Path(kwargs.get("models_path", "./diffusers_models")).resolve()
+        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.models_path.mkdir(exist_ok=True, parents=True)
         if self.auto_start_server:
             self.ensure_server_is_running()
 
@@ -66,27 +68,48 @@ class DiffusersBinding(LollmsTTIBinding):
             return False
         return False
 
+
     def ensure_server_is_running(self):
         """
         Ensures the Diffusers server is running. If not, it attempts to start it
-        in a process-safe manner using a file lock.
+        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, timeout=60) # Increased timeout for long installs
+        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:
-            with lock:
+            # 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=60):
+                # 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. Connected successfully.")
+                    ASCIIColors.green("Server was started by another process while we waited. Connected successfully.")
         except Timeout:
-            ASCIIColors.yellow("Could not acquire lock. Another process is likely starting the server. Waiting for it to become available...")
+            # 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=300) # Give it a longer timeout here just in case.
 
-        self._wait_for_server()
+        # 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.")
 
     def install_server_dependencies(self):
         """
@@ -97,6 +120,24 @@ class DiffusersBinding(LollmsTTIBinding):
         pm_v = pm.PackageManager(venv_path=str(self.venv_dir))
 
         # --- PyTorch Installation ---
+        ASCIIColors.info(f"Installing server dependencies")
+        pm_v.ensure_packages([
+            "requests", "uvicorn", "fastapi", "python-multipart", "filelock"
+        ])
+        ASCIIColors.info(f"Installing parisneo libraries")
+        pm_v.ensure_packages([
+            "ascii_colors","pipmaster"
+        ])
+        ASCIIColors.info(f"Installing misc libraries (numpy, tqdm...)")
+        pm_v.ensure_packages([
+            "tqdm", "numpy"
+        ])
+        ASCIIColors.info(f"Installing Pillow")
+        pm_v.ensure_packages([
+            "pillow"
+        ])
+        
+        ASCIIColors.info(f"Installing pytorch")
         torch_index_url = None
         if sys.platform == "win32":
             try:
@@ -104,21 +145,27 @@ class DiffusersBinding(LollmsTTIBinding):
                 result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, check=True)
                 ASCIIColors.green("NVIDIA GPU detected. Installing CUDA-enabled PyTorch.")
                 # Using a common and stable CUDA version. Adjust if needed.
-                torch_index_url = "https://download.pytorch.org/whl/cu121" 
+                torch_index_url = "https://download.pytorch.org/whl/cu128"
             except (FileNotFoundError, subprocess.CalledProcessError):
                 ASCIIColors.yellow("`nvidia-smi` not found or failed. Installing standard PyTorch. If you have an NVIDIA GPU, please ensure drivers are installed and in PATH.")
-        
+
         # Base packages including torch. pm.ensure_packages handles verbose output.
         pm_v.ensure_packages(["torch", "torchvision"], index_url=torch_index_url)
 
-
         # Standard dependencies
+        ASCIIColors.info(f"Installing transformers dependencies")
         pm_v.ensure_packages([
-            "pillow", "transformers", "safetensors", "requests", "tqdm", "numpy",
-            "accelerate", "uvicorn", "fastapi", "python-multipart", "filelock", "ascii_colors"
+            "transformers", "safetensors", "accelerate"
         ])
-        
+        ASCIIColors.info(f"[Optional] Installing xformers")
+        try:
+            pm_v.ensure_packages([
+                "xformers"
+            ])
+        except:
+            pass
         # Git-based diffusers to get the latest version
+        ASCIIColors.info(f"Installing diffusers library from github")
         pm_v.ensure_packages([
             {
                 "name": "diffusers",
@@ -127,14 +174,6 @@ class DiffusersBinding(LollmsTTIBinding):
             }
         ])
 
-        # XFormers (optional but recommended for NVIDIA)
-        if torch_index_url: # Only try to install xformers if CUDA is likely present
-            try:
-                ASCIIColors.info("Attempting to install xformers for performance optimization...")
-                pm_v.ensure_packages(["xformers"], upgrade=True)
-            except Exception as e:
-                ASCIIColors.warning(f"Could not install xformers. It's optional but recommended for performance on NVIDIA GPUs. Error: {e}")
-
         ASCIIColors.green("Server dependencies are satisfied.")
 
     def start_server(self):
@@ -164,10 +203,14 @@ class DiffusersBinding(LollmsTTIBinding):
             "--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)])
+
         # 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
-        
+
         self.server_process = subprocess.Popen(command, creationflags=creationflags)
         ASCIIColors.info("Diffusers server process launched in the background.")
 
@@ -191,11 +234,11 @@ class DiffusersBinding(LollmsTTIBinding):
             time.sleep(2)
         raise RuntimeError("Failed to connect to the Diffusers server within the specified timeout.")
 
-    def _post_request(self, endpoint: str, data: Optional[dict] = None, files: Optional[dict] = None) -> requests.Response:
-        """Helper to make POST requests to the server."""
+    def _post_json_request(self, endpoint: str, data: Optional[dict] = None) -> requests.Response:
+        """Helper to make POST requests with a JSON body."""
         try:
             url = f"{self.base_url}{endpoint}"
-            response = requests.post(url, json=data, files=files, timeout=3600) # Long timeout for generation
+            response = requests.post(url, json=data, timeout=3600) # Long timeout for generation
             response.raise_for_status()
             return response
         except requests.exceptions.RequestException as e:
@@ -208,6 +251,24 @@ class DiffusersBinding(LollmsTTIBinding):
                     ASCIIColors.error(f"Server raw response: {e.response.text}")
             raise RuntimeError("Communication with the Diffusers server failed.") from e
 
+    def _post_multipart_request(self, endpoint: str, data: Optional[dict] = None, files: Optional[list] = None) -> requests.Response:
+        """Helper to make multipart/form-data POST requests for file uploads."""
+        try:
+            url = f"{self.base_url}{endpoint}"
+            response = requests.post(url, data=data, files=files, timeout=3600)
+            response.raise_for_status()
+            return response
+        except requests.exceptions.RequestException as e:
+            # (Error handling is the same as above)
+            ASCIIColors.error(f"Failed to communicate with Diffusers server at {url}.")
+            ASCIIColors.error(f"Error details: {e}")
+            if hasattr(e, 'response') and e.response:
+                try:
+                    ASCIIColors.error(f"Server response: {e.response.json().get('detail', e.response.text)}")
+                except json.JSONDecodeError:
+                    ASCIIColors.error(f"Server raw response: {e.response.text}")
+            raise RuntimeError("Communication with the Diffusers server failed.") from e
+        
     def _get_request(self, endpoint: str, params: Optional[dict] = None) -> requests.Response:
         """Helper to make GET requests to the server."""
         try:
@@ -222,69 +283,78 @@ class DiffusersBinding(LollmsTTIBinding):
     def unload_model(self):
         ASCIIColors.info("Requesting server to unload the current model...")
         try:
-            self._post_request("/unload_model")
+            self._post_json_request("/unload_model")
         except Exception as e:
             ASCIIColors.warning(f"Could not send unload request to server: {e}")
         pass
 
     def generate_image(self, prompt: str, negative_prompt: str = "", **kwargs) -> bytes:
-        response = self._post_request("/generate_image", data={
+        params = kwargs.copy()
+        if "model_name" not in params and self.config.get("model_name"):
+            params["model_name"] = self.config["model_name"]
+            
+        response = self._post_json_request("/generate_image", data={
             "prompt": prompt,
             "negative_prompt": negative_prompt,
-            "params": kwargs
+            "params": params
         })
         return response.content
 
     def edit_image(self, images: Union[str, List[str], "Image.Image", List["Image.Image"]], prompt: str, **kwargs) -> bytes:
-        files = {}
-        image_paths = []
-        
+        images_b64 = []
         if not isinstance(images, list):
             images = [images]
 
-        for i, img in enumerate(images):
-            if hasattr(img, 'save'): # PIL Image
+
+        for img in images:
+            # Case 1: Input is a PIL Image object
+            if hasattr(img, 'save'):
                 buffer = BytesIO()
                 img.save(buffer, format="PNG")
-                buffer.seek(0)
-                files[f"image_{i}"] = (f"image_{i}.png", buffer, "image/png")
-            elif isinstance(img, str) and Path(img).is_file():
-                # The server will load this path directly
-                image_paths.append(img)
-            elif isinstance(img, str): # Handle base64 strings
+                b64_string = base64.b64encode(buffer.getvalue()).decode('utf-8')
+                images_b64.append(b64_string)
+            
+            # Case 2: Input is a string (could be path or already base64)
+            elif isinstance(img, str):
                 try:
-                    # Simple base64 check
-                    if img.startswith("data:image/") and ";base64," in img:
-                        b64_data = img.split(";base64,")[1]
-                        img_bytes = base64.b64decode(b64_data)
-                        files[f"image_{i}"] = (f"image_{i}.png", img_bytes, "image/png")
+                    b64_string = img.split(";base64,")[1] if ";base64," in img else img
+                    base64.b64decode(b64_string) # Validate
+                    images_b64.append(b64_string)
                 except Exception:
-                     raise ValueError(f"Unsupported string image format in edit_image: {img[:100]}")
+                    ASCIIColors.warning(f"Warning: A string input was not a valid file path or base64. Skipping.")
             else:
-                 raise ValueError(f"Unsupported image type in edit_image: {type(img)}")
+                raise ValueError(f"Unsupported image type in edit_image: {type(img)}")
+        if not images_b64:
+            raise ValueError("No valid images were provided to the edit_image function.")
+        
+        params = kwargs.copy()
+        if "model_name" not in params and self.config.get("model_name"):
+            params["model_name"] = self.config["model_name"]
 
-        data_payload = {
+        # Translate "mask" to "mask_image" for server compatibility
+        if "mask" in params and params["mask"]:
+            params["mask_image"] = params.pop("mask")
+
+        json_payload = {
             "prompt": prompt,
-            "image_paths": image_paths,
-            "params": kwargs
+            "images_b64": images_b64,
+            "params": params
         }
-
-        # FastAPI needs separate form fields for json and files
-        response = self._post_request("/edit_image", data={"json_payload": json.dumps(data_payload)}, files=files)
+        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_local_models(self) -> List[str]:
         return self._get_request("/list_local_models").json()
-        
+
     def list_available_models(self) -> List[str]:
         return self._get_request("/list_available_models").json()
-    
+
     def list_services(self, **kwargs) -> List[Dict[str, str]]:
         return self._get_request("/list_models").json()
-    
+
     def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
         # The server holds the state, so we fetch it.
         return self._get_request("/get_settings").json()
@@ -292,7 +362,7 @@ class DiffusersBinding(LollmsTTIBinding):
     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
         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_request("/set_settings", data=parsed_settings)
+        response = self._post_json_request("/set_settings", data=parsed_settings)
         return response.json().get("success", False)
 
     def ps(self) -> List[dict]: