# # SPDX-FileCopyrightText: Hadad # 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