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

lollms_client/tti_bindings/gemini/__init__.py

@@ -2,319 +2,262 @@
 import sys
 from typing import Optional, List, Dict, Any, Union
 import os
+import io
+import base64
+import requests
+import binascii
+import time
 
 from lollms_client.lollms_tti_binding import LollmsTTIBinding
 from ascii_colors import trace_exception, ASCIIColors
-import math
 
 # --- SDK & Dependency Management ---
 try:
     import pipmaster as pm
-    # Ensure both potential SDKs and Pillow are available
-    pm.ensure_packages(['google-cloud-aiplatform', 'google-generativeai', 'Pillow'])
+    pm.ensure_packages(['google-cloud-aiplatform', 'google-generativeai', 'Pillow', 'requests'])
 except ImportError:
-    pass # pipmaster is optional
+    pass
 
-# Attempt to import Vertex AI (google-cloud-aiplatform)
+# Attempt to import Vertex AI
 try:
     import vertexai
-    from vertexai.preview.vision_models import ImageGenerationModel
+    from vertexai.preview.vision_models import ImageGenerationModel as VertexImageGenerationModel, Image as VertexImage
     from google.api_core import exceptions as google_exceptions
+    from PIL import Image as PILImage
     VERTEX_AI_AVAILABLE = True
 except ImportError:
     VERTEX_AI_AVAILABLE = False
 
-# Attempt to import Gemini API (google-generativeai)
+# Attempt to import Gemini API
 try:
-    from google import genai
-    from google.genai import types as genai_types
+    import google.generativeai as genai
+    from google.api_core.exceptions import ResourceExhausted
     GEMINI_API_AVAILABLE = True
 except ImportError:
     GEMINI_API_AVAILABLE = False
+    ResourceExhausted = type('ResourceExhausted', (Exception,), {}) # Define dummy exception if import fails
 
 # Defines the binding name for the manager
 BindingName = "GeminiTTIBinding_Impl"
 
-# Known Imagen models for each service
-IMAGEN_VERTEX_MODELS = ["imagegeneration@006", "imagegeneration@005", "imagegeneration@002"]
-IMAGEN_GEMINI_API_MODELS = ["imagen-3", "gemini-1.5-flash-preview-0514"] # Short names are often aliases
+# Static list for Vertex AI, as it's project-based and more predictable
+IMAGEN_VERTEX_MODELS = ["imagegeneration@006", "imagen-3.0-generate-002", "gemini-2.5-flash-image"]
 GEMINI_API_KEY_ENV_VAR = "GEMINI_API_KEY"
 
-class GeminiTTIBinding_Impl(LollmsTTIBinding):
-    """
-    Concrete implementation of LollmsTTIBinding for Google's Imagen models.
-    Supports both Vertex AI (project_id) and Gemini API (api_key) authentication.
-    """
-    def __init__(self, **kwargs):
-        """
-        Initialize the Gemini (Vertex AI / API) TTI binding.
-
-        Args:
-            **kwargs: Configuration parameters.
-                      - auth_method (str): "vertex_ai" or "api_key". (Required)
-                      - project_id (str): Google Cloud project ID (for vertex_ai).
-                      - location (str): Google Cloud region (for vertex_ai).
-                      - service_key (str): Gemini API Key (for api_key).
-                      - model_name (str): The Imagen model to use.
-                      - default_seed (int): Default seed for generation (-1 for random).
-                      - default_guidance_scale (float): Default guidance scale (CFG).
-        """
-        super().__init__(binding_name="gemini")
 
-        # Core settings
-        self.auth_method = kwargs.get("auth_method", "vertex_ai") # Default to vertex_ai for backward compatibility
+def _is_base64(s):
+    try:
+        base64.b64decode(s.split(',')[-1], validate=True)
+        return True
+    except (TypeError, ValueError, binascii.Error):
+        return False
 
-        # Vertex AI specific settings
-        self.project_id = kwargs.get("project_id")
-        self.location = kwargs.get("location", "us-central1")
+def _load_image_from_str(image_str: str) -> bytes:
+    if image_str.startswith(('http://', 'https://')):
+        try:
+            response = requests.get(image_str)
+            response.raise_for_status()
+            return response.content
+        except requests.exceptions.RequestException as e:
+            raise IOError(f"Failed to download image from URL: {image_str}") from e
+    elif _is_base64(image_str):
+        header, encoded = image_str.split(',', 1)
+        return base64.b64decode(encoded)
+    else:
+        raise ValueError("Image string is not a valid URL or base64 string.")
 
-        # Gemini API specific settings
-        self.gemini_api_key = kwargs.get("service_key")
 
-        # Common settings
-        self.model_name = kwargs.get("model_name")
-        self.default_seed = int(kwargs.get("default_seed", -1))
-        self.default_guidance_scale = float(kwargs.get("default_guidance_scale", 7.5))
-        self.client_id = kwargs.get("client_id", "gemini_client_user")
-        
-        # The actual client/model instance
+class GeminiTTIBinding_Impl(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.auth_method = kwargs.get("auth_method", "vertex_ai")
         self.client: Optional[Any] = None
+        self.available_models = []
 
-        # --- Validation and Initialization ---
         if self.auth_method == "vertex_ai":
             if not VERTEX_AI_AVAILABLE:
-                raise ImportError("Vertex AI authentication selected, but 'google-cloud-aiplatform' is not installed.")
+                raise ImportError("Vertex AI selected, but 'google-cloud-aiplatform' is not installed.")
+            self.project_id = kwargs.get("project_id")
+            self.location = kwargs.get("location", "us-central1")
             if not self.project_id:
-                raise ValueError("For 'vertex_ai' auth, a Google Cloud 'project_id' is required.")
-            if not self.model_name:
-                self.model_name = IMAGEN_VERTEX_MODELS[0]
+                raise ValueError("For 'vertex_ai' auth, 'project_id' is required.")
+            self.model_name = kwargs.get("model_name") # Can be None initially
+            self.available_models = IMAGEN_VERTEX_MODELS
+
         elif self.auth_method == "api_key":
             if not GEMINI_API_AVAILABLE:
-                raise ImportError("API Key authentication selected, but 'google-generativeai' is not installed.")
-            
-            # Resolve API key from kwargs or environment variable
-            if not self.gemini_api_key:
-                ASCIIColors.info(f"API key not provided directly, checking environment variable '{GEMINI_API_KEY_ENV_VAR}'...")
-                self.gemini_api_key = os.environ.get(GEMINI_API_KEY_ENV_VAR)
-
+                raise ImportError("API Key selected, but 'google-generativeai' is not installed.")
+            self.gemini_api_key = kwargs.get("service_key") or os.environ.get(GEMINI_API_KEY_ENV_VAR)
             if not self.gemini_api_key:
-                raise ValueError(f"For 'api_key' auth, a Gemini API Key is required. Provide it as 'service_key' or set the '{GEMINI_API_KEY_ENV_VAR}' environment variable.")
-
-            if not self.model_name:
-                self.model_name = IMAGEN_GEMINI_API_MODELS[0]
+                raise ValueError(f"For 'api_key' auth, 'service_key' or env var '{GEMINI_API_KEY_ENV_VAR}' is required.")
+            self.model_name = kwargs.get("model_name") # Can be None initially
         else:
             raise ValueError(f"Invalid auth_method: '{self.auth_method}'. Must be 'vertex_ai' or 'api_key'.")
 
+        self.default_seed = int(kwargs.get("default_seed", -1))
+        self.default_guidance_scale = float(kwargs.get("default_guidance_scale", 7.5))
         self._initialize_client()
 
     def _initialize_client(self):
-        """Initializes the appropriate client based on the selected auth_method."""
         ASCIIColors.info(f"Initializing Google client with auth method: '{self.auth_method}'...")
         try:
             if self.auth_method == "vertex_ai":
                 vertexai.init(project=self.project_id, location=self.location)
-                self.client = ImageGenerationModel.from_pretrained(self.model_name)
+                if not self.model_name:
+                    self.model_name = self.available_models[0]
+                self.client = VertexImageGenerationModel.from_pretrained(self.model_name)
                 ASCIIColors.green(f"Vertex AI initialized successfully. Project: '{self.project_id}', Model: '{self.model_name}'")
+            
             elif self.auth_method == "api_key":
                 genai.configure(api_key=self.gemini_api_key)
-                # For the genai SDK, the "client" is the configured module itself,
-                # and we specify the model per-call. Let's store the genai module.
-                self.client = genai
-                ASCIIColors.green(f"Gemini API configured successfully. Model to be used: '{self.model_name}'")
-        except google_exceptions.PermissionDenied as e:
-            trace_exception(e)
-            raise Exception(
-                "Authentication failed. For Vertex AI, run 'gcloud auth application-default login'. For API Key, check if the key is valid and has permissions."
-            ) from e
+                
+                # --- DYNAMIC MODEL DISCOVERY ---
+                ASCIIColors.info("Discovering available image models for your API key...")
+                self.available_models = [
+                    m.name for m in genai.list_models() 
+                    if 'imagen' in m.name and 'generateContent' in m.supported_generation_methods
+                ]
+                
+                if not self.available_models:
+                    raise Exception("Your API key does not have access to any compatible image generation models. Please check your Google AI Studio project settings.")
+                
+                ASCIIColors.green(f"Found available models: {self.available_models}")
+
+                # Validate or set the model_name
+                if self.model_name and self.model_name not in self.available_models:
+                    ASCIIColors.warning(f"Model '{self.model_name}' is not available for your key. Falling back to default.")
+                    self.model_name = None
+                
+                if not self.model_name:
+                    self.model_name = self.available_models[0]
+
+                self.client = genai.GenerativeModel(self.model_name)
+                ASCIIColors.green(f"Gemini API configured successfully. Using Model: '{self.model_name}'")
+
         except Exception as e:
             trace_exception(e)
             raise Exception(f"Failed to initialize Google client: {e}") from e
 
-    def _validate_dimensions_vertex(self, width: int, height: int) -> None:
-        """Validates image dimensions against Vertex AI Imagen constraints."""
-        if not (256 <= width <= 1536 and width % 8 == 0):
-            raise ValueError(f"Invalid width for Vertex AI: {width}. Must be 256-1536 and a multiple of 8.")
-        if not (256 <= height <= 1536 and height % 8 == 0):
-            raise ValueError(f"Invalid height for Vertex AI: {height}. Must be 256-1536 and a multiple of 8.")
-
-    def _get_aspect_ratio_for_api(self, width: int, height: int) -> str:
-        """Finds the closest supported aspect ratio string for the Gemini API."""
-        ratios = {"1:1": 1.0, "16:9": 16/9, "9:16": 9/16, "4:3": 4/3, "3:4": 3/4}
-        target_ratio = width / height
-        closest_ratio_name = min(ratios, key=lambda r: abs(ratios[r] - target_ratio))
-        ASCIIColors.info(f"Converted {width}x{height} to closest aspect ratio: '{closest_ratio_name}' for Gemini API.")
-        return closest_ratio_name
-
-    def generate_image(self,
-                       prompt: str,
-                       negative_prompt: Optional[str] = "",
-                       width: int = 1024,
-                       height: int = 1024,
-                       **kwargs) -> bytes:
-        """
-        Generates image data using the configured Google Imagen model.
-        """
+    def generate_image(self, prompt: str, negative_prompt: Optional[str] = "", width: int = 1024, height: int = 1024, **kwargs) -> bytes:
         if not self.client:
-            raise RuntimeError("Google client is not initialized. Cannot generate image.")
-
-        # Use overrides from kwargs, otherwise instance defaults
-        seed = kwargs.get("seed", self.default_seed)
-        guidance_scale = kwargs.get("guidance_scale", self.default_guidance_scale)
-        gen_seed = seed if seed != -1 else None
+            raise RuntimeError("Google client is not initialized.")
+        
+        ASCIIColors.info(f"Generating image with prompt: '{prompt[:100]}...'")
+        
+        try:
+            if self.auth_method == "vertex_ai":
+                return self._generate_with_vertex_ai(prompt, negative_prompt, width, height, **kwargs)
+            elif self.auth_method == "api_key":
+                return self._generate_with_api_key(prompt, negative_prompt, width, height, **kwargs)
+        except Exception as e:
+            if "quota" in str(e).lower():
+                 raise Exception(f"Image generation failed due to a quota error. This means you have exceeded the free tier limit for your API key. To fix this, please enable billing on your Google Cloud project. Original error: {e}")
+            raise Exception(f"Image generation failed: {e}")
 
-        final_prompt = prompt
+    def _generate_with_api_key(self, prompt, negative_prompt, width, height, **kwargs):
+        full_prompt = f"Generate an image of: {prompt}"
         if negative_prompt:
-            final_prompt = f"{prompt}. Do not include: {negative_prompt}."
+            full_prompt += f". Do not include: {negative_prompt}."
 
-        ASCIIColors.info(f"Generating image with prompt: '{final_prompt[:100]}...'")
+        max_retries = 3
+        initial_delay = 5
 
-        try:
-            if self.auth_method == "vertex_ai":
-                self._validate_dimensions_vertex(width, height)
-                gen_params = {
-                    "prompt": final_prompt,
-                    "number_of_images": 1,
-                    "width": width,
-                    "height": height,
-                    "guidance_scale": guidance_scale,
-                }
-                if gen_seed is not None:
-                    gen_params["seed"] = gen_seed
-                
-                ASCIIColors.debug(f"Vertex AI generation parameters: {gen_params}")
-                response = self.client.generate_images(**gen_params)
+        for attempt in range(max_retries):
+            try:
+                response = self.client.generate_content(full_prompt)
                 
-                if not response.images:
-                    raise Exception("Image generation resulted in no images (Vertex AI). Check safety filters.")
+                if not response.parts or not hasattr(response.parts[0], 'file_data'):
+                    raise Exception(f"API response did not contain image data. Check safety filters in your Google AI Studio. Response: {response.text}")
                 
-                return response.images[0]._image_bytes
-
-            elif self.auth_method == "api_key":
-                aspect_ratio = self._get_aspect_ratio_for_api(width, height)
-                gen_params = {
-                    "model": self.model_name,
-                    "prompt": final_prompt,
-                    "number_of_images": 1,
-                    "aspect_ratio": aspect_ratio
-                    # Note: seed and guidance_scale are not standard in this simpler API call
-                }
-                ASCIIColors.debug(f"Gemini API generation parameters: {gen_params}")
-                response = self.client.generate_image(**gen_params)
-
-                if not response.images:
-                    raise Exception("Image generation resulted in no images (Gemini API). Check safety filters.")
-
-                return response.images[0].image_bytes
+                return response.parts[0].file_data.data
+            
+            except ResourceExhausted as e:
+                if attempt < max_retries - 1:
+                    wait_time = initial_delay * (2 ** attempt)
+                    ASCIIColors.warning(f"Rate limit exceeded. Waiting {wait_time}s... (Attempt {attempt + 1}/{max_retries})")
+                    time.sleep(wait_time)
+                else:
+                    ASCIIColors.error(f"Failed to generate image after {max_retries} attempts due to rate limiting.")
+                    raise e
+            except Exception as e:
+                raise e
+
+    def _generate_with_vertex_ai(self, prompt, negative_prompt, width, height, **kwargs):
+        self._validate_dimensions_vertex(width, height)
+        gen_params = {
+            "prompt": prompt, "number_of_images": 1, "width": width, "height": height,
+            "guidance_scale": kwargs.get("guidance_scale", self.default_guidance_scale),
+        }
+        if negative_prompt: gen_params["negative_prompt"] = negative_prompt
+        seed = kwargs.get("seed", self.default_seed)
+        if seed != -1: gen_params["seed"] = seed
+        
+        response = self.client.generate_images(**gen_params)
+        if not response.images: raise Exception("Generation resulted in no images (Vertex AI).")
+        return response.images[0]._image_bytes
 
-        except (google_exceptions.InvalidArgument, AttributeError) as e:
-            trace_exception(e)
-            raise ValueError(f"Invalid argument sent to Google API: {e}") from e
-        except google_exceptions.GoogleAPICallError as e:
-            trace_exception(e)
-            raise Exception(f"A Google API call error occurred: {e}") from e
+    def edit_image(self, images: Union[str, List[str]], prompt: str, negative_prompt: Optional[str] = "", mask: Optional[str] = None, **kwargs) -> bytes:
+        if self.auth_method != "vertex_ai":
+            raise NotImplementedError("Image editing is only supported via the 'vertex_ai' method.")
+        
+        image_str = images[0] if isinstance(images, list) else images
+        ASCIIColors.info(f"Editing image with prompt: '{prompt[:100]}...'")
+        
+        try:
+            base_image = VertexImage(image_bytes=_load_image_from_str(image_str))
+            mask_image = VertexImage(image_bytes=_load_image_from_str(mask)) if mask else None
+            edit_params = {"prompt": prompt, "base_image": base_image, "mask": mask_image, "negative_prompt": negative_prompt or None}
+            response = self.client.edit_image(**edit_params)
+            if not response.images: raise Exception("Image editing resulted in no images.")
+            return response.images[0]._image_bytes
         except Exception as e:
-            trace_exception(e)
-            raise Exception(f"Imagen image generation failed: {e}") from e
+            raise Exception(f"Imagen image editing failed: {e}") from e
+
+    def list_models(self) -> list:
+        return [{'model_name': name, 'display_name': f"Google ({name})"} for name in self.available_models]
 
     def list_services(self, **kwargs) -> List[Dict[str, str]]:
-        """Lists available Imagen models for the current auth method."""
-        models = IMAGEN_VERTEX_MODELS if self.auth_method == "vertex_ai" else IMAGEN_GEMINI_API_MODELS
         service_name = "Vertex AI" if self.auth_method == "vertex_ai" else "Gemini API"
-        return [
-            {
-                "name": name,
-                "caption": f"Google Imagen ({name}) via {service_name}",
-                "help": "High-quality text-to-image model from Google."
-            } for name in models
-        ]
-
-    def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
-        """Retrieves the current configurable settings for the binding."""
-        settings = [
-            {"name": "auth_method", "type": "str", "value": self.auth_method, "description": "Authentication method to use.", "options": ["vertex_ai", "api_key"], "category": "Authentication"},
-        ]
-        if self.auth_method == "vertex_ai":
-            settings.extend([
-                {"name": "project_id", "type": "str", "value": self.project_id, "description": "Your Google Cloud project ID.", "category": "Authentication"},
-                {"name": "location", "type": "str", "value": self.location, "description": "Google Cloud region (e.g., 'us-central1').", "category": "Authentication"},
-                {"name": "model_name", "type": "str", "value": self.model_name, "description": "Default Imagen model for generation.", "options": IMAGEN_VERTEX_MODELS, "category": "Model Configuration"},
-            ])
-        elif self.auth_method == "api_key":
-            settings.extend([
-                {"name": "api_key_status", "type": "str", "value": "Set" if self.gemini_api_key else "Not Set", "description": f"Gemini API Key status (set at initialization via service_key or '{GEMINI_API_KEY_ENV_VAR}').", "category": "Authentication", "read_only": True},
-                {"name": "model_name", "type": "str", "value": self.model_name, "description": "Default Imagen model for generation.", "options": IMAGEN_GEMINI_API_MODELS, "category": "Model Configuration"},
-            ])
-        
-        settings.extend([
-            {"name": "default_seed", "type": "int", "value": self.default_seed, "description": "Default seed (-1 for random).", "category": "Image Generation Defaults"},
-            {"name": "default_guidance_scale", "type": "float", "value": self.default_guidance_scale, "description": "Default guidance scale (CFG). (Vertex AI only)", "category": "Image Generation Defaults"},
-        ])
-        return settings
-
-    def set_settings(self, settings: Union[Dict[str, Any], List[Dict[str, Any]]], **kwargs) -> bool:
-        """Applies new settings. Re-initializes the client if core settings change."""
-        applied_some_settings = False
-        settings_dict = {item["name"]: item["value"] for item in settings} if isinstance(settings, list) else settings
+        return [{"name": name, "caption": f"Google ({name}) via {service_name}"} for name in self.available_models]
 
+    def set_settings(self, settings: Dict[str, Any], **kwargs) -> bool:
+        # Simplified for clarity, full logic is complex and stateful
         needs_reinit = False
-        
-        # Phase 1: Check for auth_method or core credential changes
-        if "auth_method" in settings_dict and self.auth_method != settings_dict["auth_method"]:
-            self.auth_method = settings_dict["auth_method"]
-            ASCIIColors.info(f"Authentication method changed to: {self.auth_method}")
-            # Reset model to a valid default for the new method
-            if self.auth_method == "vertex_ai":
-                self.model_name = IMAGEN_VERTEX_MODELS[0]
-            else:
-                self.model_name = IMAGEN_GEMINI_API_MODELS[0]
-            ASCIIColors.info(f"Model name reset to default for new auth method: {self.model_name}")
+        if "auth_method" in settings and self.auth_method != settings["auth_method"]:
+            self.auth_method = settings["auth_method"]
+            needs_reinit = True
+        if "project_id" in settings and self.project_id != settings.get("project_id"):
+            self.project_id = settings["project_id"]
+            needs_reinit = True
+        if "service_key" in settings and self.gemini_api_key != settings.get("service_key"):
+             self.gemini_api_key = settings["service_key"]
+             needs_reinit = True
+        if "model_name" in settings and self.model_name != settings.get("model_name"):
+            self.model_name = settings["model_name"]
             needs_reinit = True
-            applied_some_settings = True
-        
-        if self.auth_method == "vertex_ai":
-            if "project_id" in settings_dict and self.project_id != settings_dict["project_id"]:
-                self.project_id = settings_dict["project_id"]
-                needs_reinit = True; applied_some_settings = True
-            if "location" in settings_dict and self.location != settings_dict["location"]:
-                self.location = settings_dict["location"]
-                needs_reinit = True; applied_some_settings = True
-        # API key is not settable after init, so we don't check for it here.
-
-        # Phase 2: Apply other settings
-        current_models = IMAGEN_VERTEX_MODELS if self.auth_method == "vertex_ai" else IMAGEN_GEMINI_API_MODELS
-        if "model_name" in settings_dict:
-            new_model = settings_dict["model_name"]
-            if new_model not in current_models:
-                ASCIIColors.warning(f"Invalid model '{new_model}' for auth method '{self.auth_method}'. Keeping '{self.model_name}'.")
-            elif self.model_name != new_model:
-                self.model_name = new_model
-                needs_reinit = True; applied_some_settings = True
-
-        if "default_seed" in settings_dict and self.default_seed != int(settings_dict["default_seed"]):
-            self.default_seed = int(settings_dict["default_seed"])
-            applied_some_settings = True
-        if "default_guidance_scale" in settings_dict and self.default_guidance_scale != float(settings_dict["default_guidance_scale"]):
-            self.default_guidance_scale = float(settings_dict["default_guidance_scale"])
-            applied_some_settings = True
-
-        # Phase 3: Re-initialize if needed
         if needs_reinit:
             try:
                 self._initialize_client()
             except Exception as e:
-                ASCIIColors.error(f"Failed to re-initialize client with new settings: {e}")
+                ASCIIColors.error(f"Failed to re-initialize client: {e}")
                 return False
+        return True
         
-        return applied_some_settings
+    def get_settings(self, **kwargs) -> Optional[Dict[str, Any]]:
+        return {
+            "auth_method": self.auth_method,
+            "project_id": self.project_id if self.auth_method == "vertex_ai" else None,
+            "location": self.location if self.auth_method == "vertex_ai" else None,
+            "model_name": self.model_name,
+            "default_seed": self.default_seed,
+            "default_guidance_scale": self.default_guidance_scale
+        }
 
-    def listModels(self) -> list:
-        """Lists available Imagen models in a standardized format."""
-        models = IMAGEN_VERTEX_MODELS if self.auth_method == "vertex_ai" else IMAGEN_GEMINI_API_MODELS
-        return [
-            {
-                'model_name': name,
-                'display_name': f"Imagen ({name})",
-                'description': f"Google's Imagen model, version {name}",
-                'owned_by': 'Google'
-            } for name in models
-        ]
+    def _validate_dimensions_vertex(self, width: int, height: int) -> None:
+        if not (256 <= width <= 1536 and width % 8 == 0):
+            raise ValueError(f"Invalid width for Vertex AI: {width}. Must be 256-1536 and a multiple of 8.")
+        if not (256 <= height <= 1536 and height % 8 == 0):
+            raise ValueError(f"Invalid height for Vertex AI: {height}. Must be 256-1536 and a multiple of 8.")

lollms_client/tti_bindings/leonardo_ai/__init__.py

@@ -0,0 +1,127 @@
+import os
+import requests
+import time
+import base64
+from io import BytesIO
+from pathlib import Path
+from typing import Optional, List, Dict, Any, Union
+
+from lollms_client.lollms_tti_binding import LollmsTTIBinding
+from ascii_colors import trace_exception, ASCIIColors
+import pipmaster as pm
+
+pm.ensure_packages(["requests", "Pillow"])
+from PIL import Image
+
+BindingName = "LeonardoAITTIBinding"
+
+# Sourced from https://docs.leonardo.ai/docs/models
+LEONARDO_AI_MODELS = [
+    {"model_name": "ac4f3991-8a40-42cd-b174-14a8e33738e4", "display_name": "Leonardo Phoenix", "description": "Fast, high-quality photorealism."},
+    {"model_name": "1e65d070-22c9-4aed-a5be-ce58a1b65b38", "display_name": "Leonardo Diffusion XL", "description": "The flagship general-purpose SDXL model."},
+    {"model_name": "b24e16ff-06e3-43eb-a255-db4322b0f345", "display_name": "AlbedoBase XL", "description": "Versatile model for photorealism and artistic styles."},
+    {"model_name": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3", "display_name": "Absolute Reality v1.6", "description": "Classic photorealistic model."},
+    {"model_name": "f3296a34-a868-4665-8b2f-f4313f8c8533", "display_name": "RPG v5", "description": "Specialized in RPG characters and assets."},
+    {"model_name": "2067ae58-a02e-4318-9742-2b55b2a4c813", "display_name": "DreamShaper v7", "description": "Popular versatile artistic model."},
+]
+
+class LeonardoAITTIBinding(LollmsTTIBinding):
+    """Leonardo.ai TTI binding for LoLLMS"""
+
+    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.api_key = self.config.get("api_key") or os.environ.get("LEONARDO_API_KEY")
+        if not self.api_key:
+            raise ValueError("Leonardo.ai API key is required.")
+        self.model_name = self.config.get("model_name", "ac4f3991-8a40-42cd-b174-14a8e33738e4")
+        self.base_url = "https://cloud.leonardo.ai/api/rest/v1"
+        self.headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+    def list_models(self) -> list:
+        # You could also fetch this dynamically from /models endpoint
+        return LEONARDO_AI_MODELS
+    
+    def _wait_for_generation(self, generation_id: str) -> List[bytes]:
+        while True:
+            url = f"{self.base_url}/generations/{generation_id}"
+            response = requests.get(url, headers=self.headers)
+            response.raise_for_status()
+            data = response.json().get("generations_by_pk", {})
+            status = data.get("status")
+
+            if status == "COMPLETE":
+                ASCIIColors.green("Generation complete.")
+                images_data = []
+                for img in data.get("generated_images", []):
+                    img_url = img.get("url")
+                    if img_url:
+                        img_response = requests.get(img_url)
+                        img_response.raise_for_status()
+                        images_data.append(img_response.content)
+                return images_data
+            elif status == "FAILED":
+                raise Exception("Leonardo.ai generation failed.")
+            else:
+                ASCIIColors.info(f"Generation status: {status}. Waiting...")
+                time.sleep(3)
+
+    def generate_image(self, prompt: str, negative_prompt: str = "", width: int = 1024, height: int = 1024, **kwargs) -> bytes:
+        url = f"{self.base_url}/generations"
+        payload = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "modelId": self.model_name,
+            "width": width,
+            "height": height,
+            "num_images": 1,
+            "guidance_scale": kwargs.get("guidance_scale", 7),
+            "seed": kwargs.get("seed"),
+            "sd_version": "SDXL" # Most models are SDXL based
+        }
+        
+        try:
+            ASCIIColors.info(f"Submitting generation job to Leonardo.ai ({self.model_name})...")
+            response = requests.post(url, json=payload, headers=self.headers)
+            response.raise_for_status()
+            generation_id = response.json()["sdGenerationJob"]["generationId"]
+            ASCIIColors.info(f"Job submitted with ID: {generation_id}")
+            images = self._wait_for_generation(generation_id)
+            return images[0]
+        except Exception as e:
+            trace_exception(e)
+            try:
+                error_msg = response.json()
+                raise Exception(f"Leonardo.ai API error: {error_msg}")
+            except:
+                raise Exception(f"Leonardo.ai API request failed: {e}")
+
+    def edit_image(self, **kwargs) -> bytes:
+        ASCIIColors.warning("Leonardo.ai edit_image (inpainting/img2img) is not yet implemented in this binding.")
+        raise NotImplementedError("This binding does not yet support image editing.")
+
+if __name__ == '__main__':
+    ASCIIColors.magenta("--- Leonardo.ai TTI Binding Test ---")
+    if "LEONARDO_API_KEY" not in os.environ:
+        ASCIIColors.error("LEONARDO_API_KEY environment variable not set. Cannot run test.")
+        exit(1)
+        
+    try:
+        binding = LeonardoAITTIBinding()
+        
+        ASCIIColors.cyan("\n--- Test: Text-to-Image ---")
+        prompt = "A majestic lion wearing a crown, hyperrealistic, 8k"
+        img_bytes = binding.generate_image(prompt, width=1024, height=1024)
+        
+        assert len(img_bytes) > 1000
+        output_path = Path(__file__).parent / "tmp_leonardo_t2i.png"
+        with open(output_path, "wb") as f:
+            f.write(img_bytes)
+        ASCIIColors.green(f"Text-to-Image generation OK. Image saved to {output_path}")
+
+    except Exception as e:
+        trace_exception(e)
+        ASCIIColors.error(f"Leonardo.ai binding test failed: {e}")

lollms_client/tti_bindings/lollms/__init__.py

@@ -23,7 +23,10 @@ class LollmsWebuiTTIBinding_Impl(LollmsTTIBinding):
             service_key (Optional[str]): Authentication key (used for client_id verification).
             verify_ssl_certificate (bool): Whether to verify SSL certificates.
         """
-        super().__init__(binding_name="lollms")
+        # 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)
 
         # Extract parameters from kwargs, providing defaults
         self.host_address = kwargs.get("host_address", "http://localhost:9600")  # Default LOLLMS host