Create flux_lora_tool.py

flux_lora_tool.py

@@ -0,0 +1,614 @@
+"""
+Zhou Protocol FLUX-LoRA Integration Tool
+
+This module provides a Smolagents Tool implementation for interacting with FLUX-LoRA-DLC API.
+It enables agents to generate high-quality images with customizable LoRA models.
+
+Usage:
+    flux_tool = FluxLoRATool()
+    agent = CodeAgent(tools=[flux_tool], ...)
+"""
+
+import os
+import uuid
+import tempfile
+import logging
+from typing import Dict, Any, Optional, List, Union, Tuple
+from dataclasses import dataclass
+import contextlib
+from pathlib import Path
+
+# Third-party
+import requests
+from PIL import Image
+from gradio_client import Client
+
+# Smolagents
+from smolagents import Tool
+
+# -----------------------------------------------------------------------------
+# CONSTANTS AND TYPE DEFINITIONS
+# -----------------------------------------------------------------------------
+
+@dataclass
+class LoRAModelInfo:
+    """Value object representing LoRA model information."""
+    name: str
+    description: Optional[str] = None
+    example_image_url: Optional[str] = None
+
+
+@dataclass
+class ImageGenerationResult:
+    """Value object representing a generated image result."""
+    image_path: str
+    seed: int
+    metadata: Optional[Dict[str, Any]] = None
+
+
+# -----------------------------------------------------------------------------
+# CORE TOOL IMPLEMENTATION
+# -----------------------------------------------------------------------------
+
+class FluxLoRATool(Tool):
+    """
+    Tool for generating images using FLUX-LoRA-DLC API.
+    
+    This tool implements the Zhou Protocol integration patterns to provide
+    a clean, efficient interface for image generation using LoRA models.
+    """
+    
+    name = "flux_lora_generator"
+    description = """
+    Generates high-quality images using FLUX-LoRA models.
+    Can use custom LoRA models, adjust image parameters, and handle image inputs.
+    """
+    inputs = {
+        "prompt": {
+            "type": "string", 
+            "description": "Detailed description of the desired image."
+        },
+        "image_input": {
+            "type": "string", 
+            "description": "Optional URL or file path to input image for img2img generation.",
+            "optional": True
+        },
+        "image_strength": {
+            "type": "float",
+            "description": "Strength of input image influence (0.0-1.0), where 1.0 maintains more of original image.",
+            "optional": True,
+            "default": 0.75
+        },
+        "cfg_scale": {
+            "type": "float",
+            "description": "Guidance scale for prompt adherence (1.0-30.0).",
+            "optional": True,
+            "default": 3.5
+        },
+        "steps": {
+            "type": "integer",
+            "description": "Number of sampling steps (10-100).",
+            "optional": True,
+            "default": 28
+        },
+        "seed": {
+            "type": "integer",
+            "description": "Random seed for reproducibility. Use -1 for random seed.",
+            "optional": True,
+            "default": -1
+        },
+        "width": {
+            "type": "integer",
+            "description": "Image width in pixels.",
+            "optional": True,
+            "default": 1024
+        },
+        "height": {
+            "type": "integer",
+            "description": "Image height in pixels.",
+            "optional": True,
+            "default": 1024
+        },
+        "lora_scale": {
+            "type": "float",
+            "description": "LoRA influence scale (0.0-1.0).",
+            "optional": True,
+            "default": 0.95
+        },
+        "custom_lora": {
+            "type": "string",
+            "description": "Custom LoRA model to use. Leave empty for default.",
+            "optional": True
+        }
+    }
+    output_type = "string"
+    
+    def __init__(
+        self, 
+        api_url: str = "xkerser/FLUX-LoRA-DLC",
+        image_save_dir: Optional[str] = None,
+        connection_timeout: int = 60,
+        verbose: bool = False
+    ):
+        """
+        Initialize the FLUX-LoRA Tool with Zhou Protocol connection patterns.
+        
+        Args:
+            api_url: URL or endpoint ID for the FLUX-LoRA-DLC API
+            image_save_dir: Directory to save generated images (created if doesn't exist)
+            connection_timeout: API connection timeout in seconds
+            verbose: Enable detailed logging
+        """
+        super().__init__()
+        
+        # Initialize logging
+        self.logger = logging.getLogger("flux_lora_tool")
+        self.logger.setLevel(logging.DEBUG if verbose else logging.INFO)
+        
+        # Set up client and storage directories
+        self.api_url = api_url
+        self.connection_timeout = connection_timeout
+        self._client = None  # Lazy initialization
+        
+        # Set up image storage directory
+        self.image_save_dir = image_save_dir or os.path.join(tempfile.gettempdir(), "flux_lora_images")
+        os.makedirs(self.image_save_dir, exist_ok=True)
+        self.logger.info(f"FluxLoRATool initialized. Images will be saved to: {self.image_save_dir}")
+    
+    @property
+    def client(self) -> Client:
+        """
+        Get or initialize the Gradio client with proper connection handling.
+        
+        Returns:
+            Initialized Gradio client
+        
+        Raises:
+            ConnectionError: If client initialization fails
+        """
+        if self._client is None:
+            try:
+                self._client = Client(
+                    self.api_url,
+                    timeout=self.connection_timeout
+                )
+                self.logger.debug(f"Gradio client initialized for: {self.api_url}")
+            except Exception as e:
+                error_msg = f"Failed to initialize FLUX-LoRA client: {str(e)}"
+                self.logger.error(error_msg)
+                raise ConnectionError(error_msg) from e
+        
+        return self._client
+    
+    def _validate_inputs(self, **kwargs) -> Dict[str, Any]:
+        """
+        Validate and normalize input parameters with Zhou Protocol validation patterns.
+        
+        Args:
+            **kwargs: Input parameters
+            
+        Returns:
+            Validated and normalized parameters
+            
+        Raises:
+            ValueError: If input validation fails
+        """
+        validated = {}
+        
+        # Required parameter: prompt
+        if not kwargs.get("prompt"):
+            raise ValueError("Prompt is required for image generation")
+        validated["prompt"] = kwargs["prompt"]
+        
+        # Image input handling
+        if "image_input" in kwargs and kwargs["image_input"]:
+            input_image = kwargs["image_input"]
+            # Handle URL vs. local file
+            if input_image.startswith(("http://", "https://")):
+                # We'll need to download and process this
+                validated["image_input"] = self._download_image(input_image)
+            else:
+                # Check if file exists
+                if not os.path.exists(input_image):
+                    raise ValueError(f"Image file not found: {input_image}")
+                validated["image_input"] = input_image
+        
+        # Numeric parameter validation with constraints
+        numeric_params = {
+            "image_strength": {"min": 0.0, "max": 1.0, "default": 0.75},
+            "cfg_scale": {"min": 1.0, "max": 30.0, "default": 3.5},
+            "steps": {"min": 10, "max": 100, "default": 28},
+            "width": {"min": 128, "max": 2048, "default": 1024},
+            "height": {"min": 128, "max": 2048, "default": 1024},
+            "lora_scale": {"min": 0.0, "max": 1.0, "default": 0.95}
+        }
+        
+        for param, constraints in numeric_params.items():
+            if param in kwargs and kwargs[param] is not None:
+                value = kwargs[param]
+                
+                # Type conversion if needed
+                if param in ["steps", "width", "height"]:
+                    try:
+                        value = int(value)
+                    except (ValueError, TypeError):
+                        raise ValueError(f"Parameter '{param}' must be an integer")
+                else:
+                    try:
+                        value = float(value)
+                    except (ValueError, TypeError):
+                        raise ValueError(f"Parameter '{param}' must be a number")
+                
+                # Range validation
+                if value < constraints["min"] or value > constraints["max"]:
+                    raise ValueError(
+                        f"Parameter '{param}' must be between {constraints['min']} and {constraints['max']}"
+                    )
+                
+                validated[param] = value
+            else:
+                validated[param] = constraints["default"]
+        
+        # Special handling for seed
+        if "seed" in kwargs and kwargs["seed"] is not None:
+            try:
+                seed = int(kwargs["seed"])
+                # -1 indicates random seed
+                if seed == -1:
+                    try:
+                        seed = self._get_random_seed()
+                    except Exception as e:
+                        self.logger.warning(f"Failed to get random seed from API: {e}")
+                        # Fallback to Python's random
+                        import random
+                        seed = random.randint(0, 2**32 - 1)
+                validated["seed"] = seed
+            except (ValueError, TypeError):
+                raise ValueError("Seed must be an integer")
+        else:
+            # Default to random seed
+            validated["seed"] = self._get_random_seed()
+        
+        # Custom LoRA handling
+        if "custom_lora" in kwargs and kwargs["custom_lora"]:
+            validated["custom_lora"] = kwargs["custom_lora"]
+        
+        return validated
+    
+    def _download_image(self, url: str) -> str:
+        """
+        Download image from URL and save to local file.
+        
+        Args:
+            url: Image URL
+            
+        Returns:
+            Local file path
+            
+        Raises:
+            ConnectionError: If download fails
+        """
+        try:
+            response = requests.get(url, stream=True, timeout=30)
+            response.raise_for_status()
+            
+            # Generate temporary file path
+            file_ext = self._guess_extension(response.headers.get("Content-Type", ""))
+            temp_path = os.path.join(
+                self.image_save_dir, 
+                f"input_{uuid.uuid4().hex}{file_ext}"
+            )
+            
+            # Save image
+            with open(temp_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            
+            self.logger.debug(f"Downloaded image from {url} to {temp_path}")
+            return temp_path
+            
+        except Exception as e:
+            error_msg = f"Failed to download image from {url}: {str(e)}"
+            self.logger.error(error_msg)
+            raise ConnectionError(error_msg) from e
+    
+    def _guess_extension(self, content_type: str) -> str:
+        """
+        Guess file extension from content type.
+        
+        Args:
+            content_type: HTTP Content-Type header
+            
+        Returns:
+            File extension (with dot)
+        """
+        content_type = content_type.lower()
+        if "jpeg" in content_type or "jpg" in content_type:
+            return ".jpg"
+        elif "png" in content_type:
+            return ".png"
+        elif "webp" in content_type:
+            return ".webp"
+        elif "gif" in content_type:
+            return ".gif"
+        else:
+            return ".png"  # Default to PNG
+    
+    def _get_random_seed(self) -> int:
+        """
+        Get a random seed from the API.
+        
+        Returns:
+            Random seed value
+            
+        Raises:
+            RuntimeError: If random seed retrieval fails
+        """
+        try:
+            result = self.client.predict(api_name="/get_random_value")
+            if isinstance(result, (int, float)):
+                return int(result)
+            else:
+                raise ValueError(f"Unexpected result type: {type(result)}")
+        except Exception as e:
+            # Just log and re-raise as we have fallback in the validation method
+            self.logger.warning(f"Failed to get random seed: {e}")
+            raise
+    
+    def _handle_custom_lora(self, custom_lora: Optional[str]) -> None:
+        """
+        Add or remove custom LoRA model.
+        
+        Args:
+            custom_lora: Custom LoRA model string
+            
+        Raises:
+            RuntimeError: If LoRA handling fails
+        """
+        if not custom_lora:
+            # Remove any existing custom LoRA
+            try:
+                self.client.predict(api_name="/remove_custom_lora")
+                self.logger.debug("Removed custom LoRA")
+            except Exception as e:
+                error_msg = f"Failed to remove custom LoRA: {str(e)}"
+                self.logger.error(error_msg)
+                raise RuntimeError(error_msg) from e
+        else:
+            # Add custom LoRA
+            try:
+                self.client.predict(
+                    custom_lora=custom_lora,
+                    api_name="/add_custom_lora"
+                )
+                self.logger.debug(f"Added custom LoRA: {custom_lora}")
+            except Exception as e:
+                error_msg = f"Failed to add custom LoRA '{custom_lora}': {str(e)}"
+                self.logger.error(error_msg)
+                raise RuntimeError(error_msg) from e
+    
+    def forward(
+        self,
+        prompt: str,
+        image_input: Optional[str] = None,
+        image_strength: Optional[float] = None,
+        cfg_scale: Optional[float] = None,
+        steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        lora_scale: Optional[float] = None,
+        custom_lora: Optional[str] = None
+    ) -> str:
+        """
+        Generate an image with FLUX-LoRA.
+        
+        Args:
+            prompt: Text description of the desired image
+            image_input: Optional path or URL to input image for img2img
+            image_strength: Strength of input image influence (0.0-1.0)
+            cfg_scale: Guidance scale (1.0-30.0)
+            steps: Number of sampling steps (10-100)
+            seed: Random seed (-1 for random)
+            width: Image width in pixels (128-2048)
+            height: Image height in pixels (128-2048)
+            lora_scale: LoRA influence scale (0.0-1.0)
+            custom_lora: Custom LoRA model to use
+            
+        Returns:
+            Formatted string with image generation results
+            
+        Raises:
+            ValueError: If input validation fails
+            ConnectionError: If API communication fails
+            RuntimeError: If image generation fails
+        """
+        # Step 1: Validate and normalize inputs
+        try:
+            params = self._validate_inputs(
+                prompt=prompt,
+                image_input=image_input,
+                image_strength=image_strength,
+                cfg_scale=cfg_scale,
+                steps=steps,
+                seed=seed,
+                width=width,
+                height=height,
+                lora_scale=lora_scale,
+                custom_lora=custom_lora
+            )
+            self.logger.debug(f"Validated parameters: {params}")
+        except ValueError as e:
+            return f"Parameter validation failed: {str(e)}"
+        
+        # Step 2: Handle custom LoRA if specified
+        if "custom_lora" in params:
+            try:
+                custom_lora_value = params.pop("custom_lora")
+                self._handle_custom_lora(custom_lora_value)
+            except RuntimeError as e:
+                return f"Custom LoRA setup failed: {str(e)}"
+        
+        # Step 3: Generate image
+        try:
+            # Prepare image input if provided
+            img_param = None
+            if "image_input" in params and params["image_input"]:
+                from gradio_client import handle_file
+                img_param = handle_file(params.pop("image_input"))
+            
+            # Call the API
+            generation_args = {
+                "prompt": params["prompt"],
+                "image_strength": params["image_strength"],
+                "cfg_scale": params["cfg_scale"],
+                "steps": params["steps"],
+                "randomize_seed": False,  # We handle seed explicitly
+                "seed": params["seed"],
+                "width": params["width"],
+                "height": params["height"],
+                "lora_scale": params["lora_scale"],
+            }
+            
+            # Add image input if available
+            if img_param:
+                generation_args["image_input"] = img_param
+            
+            self.logger.info(f"Generating image with params: {generation_args}")
+            result = self.client.predict(
+                api_name="/run_lora",
+                **generation_args
+            )
+            
+            # Process result
+            if isinstance(result, tuple) and len(result) >= 2:
+                image_path, actual_seed = result[0], result[1]
+                
+                # Save image to our directory
+                try:
+                    output_path = self._save_image(image_path)
+                    image_result = ImageGenerationResult(
+                        image_path=output_path,
+                        seed=int(actual_seed)
+                    )
+                    return self._format_result(image_result, params["prompt"])
+                except Exception as e:
+                    self.logger.error(f"Failed to save generated image: {e}")
+                    return f"Image generated but failed to save: {str(e)}"
+            else:
+                raise ValueError(f"Unexpected API response format: {result}")
+                
+        except Exception as e:
+            error_msg = f"Image generation failed: {str(e)}"
+            self.logger.error(error_msg)
+            return error_msg
+    
+    def _save_image(self, image_path: str) -> str:
+        """
+        Save generated image to specified directory.
+        
+        Args:
+            image_path: Path to generated image from API
+            
+        Returns:
+            Path to saved image
+            
+        Raises:
+            IOError: If image saving fails
+        """
+        try:
+            # Load the image
+            img = Image.open(image_path)
+            
+            # Generate timestamp-based filename
+            timestamp = uuid.uuid4().hex[:8]
+            output_filename = f"flux_lora_{timestamp}.png"
+            output_path = os.path.join(self.image_save_dir, output_filename)
+            
+            # Save to our directory
+            img.save(output_path)
+            self.logger.debug(f"Saved image to {output_path}")
+            
+            return output_path
+            
+        except Exception as e:
+            error_msg = f"Failed to save image: {str(e)}"
+            self.logger.error(error_msg)
+            raise IOError(error_msg) from e
+    
+    def _format_result(self, result: ImageGenerationResult, prompt: str) -> str:
+        """
+        Format the image generation result as a string.
+        
+        Args:
+            result: Image generation result
+            prompt: Original prompt
+            
+        Returns:
+            Formatted string with generation details
+        """
+        lines = [
+            f"📷 Image generated successfully!",
+            f"🖼️ Image saved to: {result.image_path}",
+            f"🌱 Seed used: {result.seed}",
+            f"📝 Original prompt: {prompt}",
+        ]
+        
+        # Add metadata if available
+        if result.metadata:
+            lines.append("📊 Additional metadata:")
+            for key, value in result.metadata.items():
+                lines.append(f"   - {key}: {value}")
+        
+        return "\n".join(lines)
+
+
+# -----------------------------------------------------------------------------
+# UTILITY FUNCTIONS
+# -----------------------------------------------------------------------------
+
+def download_image(url: str, output_dir: Optional[str] = None) -> str:
+    """
+    Standalone utility to download an image from a URL.
+    
+    Args:
+        url: Image URL
+        output_dir: Directory to save image (created if doesn't exist)
+        
+    Returns:
+        Path to downloaded image
+        
+    Raises:
+        ValueError: If URL is invalid
+        ConnectionError: If download fails
+        IOError: If saving fails
+    """
+    if not url.startswith(("http://", "https://")):
+        raise ValueError(f"Invalid URL: {url}")
+    
+    # Setup output directory
+    if output_dir is None:
+        output_dir = os.path.join(tempfile.gettempdir(), "flux_lora_images")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    try:
+        # Download image
+        response = requests.get(url, stream=True, timeout=30)
+        response.raise_for_status()
+        
+        # Determine file extension
+        content_type = response.headers.get("Content-Type", "")
+        ext = ".jpg" if "jpeg" in content_type.lower() else ".png"
+        
+        # Save image
+        output_path = os.path.join(output_dir, f"download_{uuid.uuid4().hex}{ext}")
+        with open(output_path, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        
+        return output_path
+        
+    except requests.RequestException as e:
+        raise ConnectionError(f"Failed to download image: {str(e)}")
+    except IOError as e:
+        raise IOError(f"Failed to save image: {str(e)}")