src/tools/image.py

#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.0
#

import aiohttp  # Import aiohttp library for asynchronous HTTP client functionality
from urllib.parse import quote  # Import quote function to URL-encode query parameters safely
from typing import Optional  # Import Optional type hint for parameters that may be None
from src.utils.ip_generator import generate_ip  # Import custom utility to generate random IP addresses for request headers
from src.utils.tools import initialize_tools  # Import custom function to initialize necessary tools and tokens

# Define a class to handle image generation requests asynchronously
class ImageGeneration:
    """
    This class provides a comprehensive interface for generating images asynchronously via HTTP requests.
    It supports multiple image formats with predefined dimensions and allows customization of generation parameters such as model type, seed, and enhancement options.
    The class handles URL encoding, authentication headers, and IP generation to simulate diverse client requests.
    It uses aiohttp for efficient asynchronous HTTP client sessions and manages error handling for invalid inputs and HTTP response failures.
    """
    # Dictionary mapping image format names to their corresponding width and height dimensions
    FORMATS = {
        "default": (1024, 1024),  # Default square format with 1024x1024 pixels
        "square": (1024, 1024),   # Square format identical to default size
        "landscape": (1024, 768),  # Landscape format with width greater than height
        "landscape_large": (1440, 1024),  # Larger landscape format for higher resolution images
        "portrait": (768, 1024),  # Portrait format with height greater than width
        "portrait_large": (1024, 1440),  # Larger portrait format for higher resolution images
    }

    @staticmethod  # Decorator indicating this method does not depend on instance state
    # Asynchronous method to create an image based on given instructions and parameters
    async def create_image(
        generate_image_instruction: str,  # Text instruction describing the image to generate
        image_format: str = "default",  # Desired image format key from FORMATS dictionary
        model: Optional[str] = "flux",  # Optional model identifier controlling generation style or algorithm
        seed: Optional[int] = None,  # Optional random seed to reproduce specific image outputs
        nologo: bool = True,  # Flag to remove logo watermark from generated image
        private: bool = True,  # Flag to mark image generation as private, restricting visibility
        enhance: bool = True,  # Flag to apply enhancement filters to improve image quality
    ) -> str:
        """
        This asynchronous method generates an image URL based on detailed instructions and customization options.
        It validates the requested image format against supported formats, retrieves corresponding dimensions,
        initializes necessary tools and authentication tokens, constructs the request URL with encoded parameters,
        and performs an HTTP GET request asynchronously using aiohttp.
        The method includes headers for authorization and IP address to simulate requests from different clients.
        On successful response (HTTP 200), it returns the final URL of the generated image.
        If the response status is not successful, it raise an exception.
        """
        # Validate if the requested image format is supported, otherwise raise an error
        if image_format not in ImageGeneration.FORMATS:
            raise ValueError("Invalid image format")

        # Retrieve width and height dimensions for the requested image format
        width, height = ImageGeneration.FORMATS[image_format]

        # Initialize external tools and retrieve the image tool URL and authorization token
        _, image_tool, _, poll_token = initialize_tools()

        # Construct the base URL for the image generation request with URL-encoded instruction
        url = f"{image_tool}{quote(generate_image_instruction)}"

        # Prepare query parameters dictionary for the HTTP request with all relevant options
        params = {
            "width": width,  # Image width in pixels
            "height": height,  # Image height in pixels
            "model": model,  # Model identifier string for image generation
            "nologo": str(nologo).lower(),  # Convert boolean to lowercase string for query parameter
            "private": str(private).lower(),  # Convert boolean to lowercase string for query parameter
            "enhance": str(enhance).lower(),  # Convert boolean to lowercase string for query parameter
        }

        # If a seed value is provided, add it to the query parameters to control randomness
        if seed is not None:
            params["seed"] = seed

        # Prepare HTTP headers
        headers = {
            "Authorization": f"Bearer {poll_token}",  # Bearer token for API authentication
            "X-Forwarded-For": generate_ip()  # Random IP address for request header to simulate client origin
        }

        # Create an asynchronous HTTP client session to perform the GET request
        async with aiohttp.ClientSession() as session:
            """
            This context manager ensures that the HTTP session is properly opened and closed asynchronously,
            enabling efficient connection reuse and resource cleanup after the request completes.
            """
            # Perform the HTTP GET request asynchronously with the constructed URL, parameters, and headers
            async with session.get(
                url,  # The fully constructed URL including encoded instruction
                params=params,  # Query parameters controlling image generation options
                headers=headers  # Authorization and client identity headers
            ) as resp:
                """
                This context manager handles the HTTP response asynchronously.
                It waits for the server's response and manages the response lifecycle.
                """
                # Check if the HTTP response status indicates success
                if resp.status == 200 and 'image/' in resp.headers.get('Content-Type', ''):
                    # Return the final URL of the generated image as a string
                    return str(resp.url)
                # If the response is not successful or content type is unexpected, raise an exception
                raise Exception()