src/tools/image.py

#
# SPDX-FileCopyrightText: Hadad <[email protected]>
# SPDX-License-Identifier: Apache-2.0
#

import httpx  # Import httpx library for performing asynchronous HTTP requests efficiently
from urllib.parse import quote  # Import quote function to safely encode strings for use in URLs
from typing import Optional  # Import Optional type hint for parameters that can 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 utility function to initialize and retrieve tool endpoints

# Define a class named ImageGeneration to encapsulate functionalities related to generating image content
class ImageGeneration:
    # This class provides methods to create image files based on text instructions

    """
    A class to handle image generation requests to an external image generation service.

    Attributes:
        FORMATS (dict): A dictionary mapping image format names to their (width, height) dimensions.

    Methods:
        create_image: Asynchronously generates an image based on a textual instruction and parameters,
                      returning the URL of the generated image.
    """

    # Image formats
    FORMATS = {
        "default": (1024, 1024),  # Default square image size (width x height)
        "square": (1024, 1024),  # Square image format with equal width and height
        "landscape": (1024, 768),  # Landscape format with wider width than height
        "landscape_large": (1440, 1024),  # Larger landscape format with increased resolution
        "portrait": (768, 1024),  # Portrait format with taller height than width
        "portrait_large": (1024, 1440),  # Larger portrait format with increased resolution
    }

    @staticmethod  # Decorator indicating that the following method does not depend on instance state and can be called on the class itself
    # Define an asynchronous method to create image from a text instruction
    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-realism",  # Optional model name for image generation; defaults to 'flux-realism'
        seed: Optional[int] = None,  # Optional seed value for randomization control in image generation
        nologo: bool = True,  # Whether to generate image without logo watermark; defaults to True
        private: bool = True,  # Whether the generated image should be private; defaults to True
        enhance: bool = True,  # Whether to apply enhancement filters to the generated image; defaults to True
    ) -> str:
        """
        Asynchronously generate an image URL by sending a request to the image generation service.
        This method will keep retrying until a successful response with status code 200 is received.

        Args:
            generate_image_instruction (str): The textual instruction or description for the desired image.
            image_format (str, optional): The format key specifying image dimensions. Defaults to "default".
            model (Optional[str], optional): The image generation model to use. Defaults to "flux-realism".
            seed (Optional[int], optional): Seed for randomization to reproduce images. Defaults to None.
            nologo (bool, optional): Flag to exclude logo watermark. Defaults to True.
            private (bool, optional): Flag to mark image as private. Defaults to True.
            enhance (bool, optional): Flag to apply image enhancement. Defaults to True.

        Returns:
            str: The URL of the generated image if the request is successful.

        Raises:
            ValueError: If the specified image_format is not supported.
            Exception: If the image generation continuously fails (currently infinite retry).
        """
        # Validate that the requested image format exists in the FORMATS dictionary
        if image_format not in ImageGeneration.FORMATS:
            raise ValueError("Invalid image format.")

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

        # Initialize tools and retrieve the image generation service endpoint
        _, image_tool, _ = initialize_tools()

        # Encode the image instruction to safely include it in the URL path
        generate_image_instruct = quote(generate_image_instruction)

        # Construct the full URL for the image generation request by appending the encoded instruction
        url = f"{image_tool}{generate_image_instruct}"  # Full endpoint URL for image generation

        # Prepare query parameters including image dimensions, model, and flags converted to string "true"/"false"
        params = {
            "width": width,  # Image width parameter
            "height": height,  # Image height parameter
            "model": model,  # Model name for image generation
            "nologo": "true" if nologo else "false",  # Flag to exclude logo watermark as string
            "private": "true" if private else "false",  # Flag to mark image as private as string
            "enhance": "true" if enhance else "false"  # Flag to apply enhancement as string
        }

        # Include seed parameter if provided to control randomness in image generation
        if seed is not None:
            params["seed"] = seed  # Add seed to parameters to reproduce images

        # Prepare HTTP headers with a generated random IP to simulate different client origins
        headers = {
            "X-Forwarded-For": generate_ip()  # Random IP address for request header to simulate client origin
        }

        # Create an asynchronous HTTP client with no timeout limit to perform the request
        async with httpx.AsyncClient(timeout=None) as client:
            # Keep retrying the request until a successful response with status 200 is received
            while True:
                # Send a GET request to the image generation service with URL, parameters, and headers
                resp = await client.get(url, params=params, headers=headers)

                # Check if the response status code indicates success
                if resp.status_code == 200:
                    # Return the URL of the generated image as a string
                    return str(resp.url)
                else:
                    # Wait briefly before retrying to avoid overwhelming the server
                    await asyncio.sleep(15)  # Pause 15 second before retrying