utils/ai_generator.py

# utils/ai_generator.py
import gradio as gr
import os
import time
#from turtle import width  # Added for implementing delays
from torch import cuda
import random
from utils.ai_generator_diffusers_flux import generate_ai_image_local
#from pathlib import Path
from huggingface_hub import InferenceClient
import requests
import io
from PIL import Image
from tempfile import NamedTemporaryFile
import utils.constants as constants

def generate_image_from_text(text, model_name="flax-community/dalle-mini", image_width=768, image_height=512, progress=gr.Progress(track_tqdm=True)):
    # Initialize the InferenceClient
    client = InferenceClient()
    # Generate the image from the text
    response = client(text, model_name)
    # Get the image data
    image_data = response.content
    # Load the image from the data
    image = Image.open(io.BytesIO(image_data))
    # Resize the image
    image = image.resize((image_width, image_height))
    return image

def generate_ai_image(
    map_option,
    prompt_textbox_value,
    neg_prompt_textbox_value,
    model,
    lora_weights=None,
    conditioned_image=None,
    pipeline = "FluxPipeline",
    width=912,
    height=512,
    strength=0.5,
    seed = 0,
    progress=gr.Progress(track_tqdm=True),
    *args,
    **kwargs
):   
    if seed == 0:
        seed = random.randint(0, constants.MAX_SEED)
    if (cuda.is_available() and cuda.device_count() >= 1): # Check if a local GPU is available 
        print("Local GPU available. Generating image locally.")
        if conditioned_image is not None:
            pipeline = "FluxImg2ImgPipeline"
        return generate_ai_image_local(
            map_option,
            prompt_textbox_value,
            neg_prompt_textbox_value,
            model,
            lora_weights=lora_weights,
            seed=seed,
            conditioned_image=conditioned_image,
            pipeline_name=pipeline,
            strength=strength,
            height=height,
            width=width
        )
    else:
        print("No local GPU available. Sending request to Hugging Face API.")
        return generate_ai_image_remote(
            map_option,
            prompt_textbox_value,
            neg_prompt_textbox_value,
            model,
            height=height,
            width=width,
            seed=seed
        )

def generate_ai_image_remote(map_option, prompt_textbox_value, neg_prompt_textbox_value, model, height=512, width=912, num_inference_steps=30, guidance_scale=3.5, seed=777,progress=gr.Progress(track_tqdm=True)):
    max_retries = 3
    retry_delay = 4  # Initial delay in seconds

    try:
        if map_option != "Prompt":
            prompt = constants.PROMPTS[map_option]
            # Convert the negative prompt string to a list
            negative_prompt_str = constants.NEGATIVE_PROMPTS.get(map_option, "")
            negative_prompt = [p.strip() for p in negative_prompt_str.split(',') if p.strip()]
        else:
            prompt = prompt_textbox_value
            # Convert the negative prompt string to a list
            negative_prompt = [p.strip() for p in neg_prompt_textbox_value.split(',') if p.strip()] if neg_prompt_textbox_value else []

        print("Remotely Generating image with the following parameters:")
        print(f"Prompt: {prompt}")
        print(f"Negative Prompt: {negative_prompt}")
        print(f"Height: {height}")
        print(f"Width: {width}")
        print(f"Number of Inference Steps: {num_inference_steps}")
        print(f"Guidance Scale: {guidance_scale}")
        print(f"Seed: {seed}")

        for attempt in range(1, max_retries + 1):
            try:
                if os.getenv("IS_SHARED_SPACE") == "True":
                    client = InferenceClient(
                        model,
                        token=constants.HF_API_TOKEN
                    )
                    image = client.text_to_image(
                        inputs=prompt,
                        parameters={
                            "guidance_scale": guidance_scale,
                            "num_inference_steps": num_inference_steps,
                            "width": width,
                            "height": height,
                            "max_sequence_length":512,                            
                            # Optional: Add 'scheduler' and 'seed' if needed
                            "seed": seed
                        }
                    )
                else:
                    API_URL = f"https://api-inference.huggingface.co/models/{model}"
                    headers = {
                        "Authorization": f"Bearer {constants.HF_API_TOKEN}",
                        "Content-Type": "application/json"
                    }
                    payload = {
                        "inputs": prompt,
                        "parameters": {
                            "guidance_scale": guidance_scale,
                            "num_inference_steps": num_inference_steps,
                            "width": width,
                            "height": height,
                            "max_sequence_length":512,
                            # Optional: Add 'scheduler' and 'seed' if needed
                            "seed": seed
                        }
                    }

                    print(f"Attempt {attempt}: Sending POST request to Hugging Face API...")
                    response = requests.post(API_URL, headers=headers, json=payload, timeout=300)  # Increased timeout to 30 seconds
                    if response.status_code == 200:
                        image_bytes = response.content
                        image = Image.open(io.BytesIO(image_bytes))
                        break  # Exit the retry loop on success
                    elif response.status_code == 400:
                        # Handle 400 Bad Request specifically
                        print(f"Bad Request (400): {response.text}")
                        print("Check your request parameters and payload format.")
                        return None  # Do not retry on 400 errors
                    elif response.status_code in [429, 504]:
                        print(f"Received status code {response.status_code}. Retrying in {retry_delay} seconds...")
                        if attempt < max_retries:
                            time.sleep(retry_delay)
                            retry_delay *= 2  # Exponential backoff
                        else:
                            response.raise_for_status()  # Raise exception after max retries
                    else:
                        print(f"Received unexpected status code {response.status_code}: {response.text}")
                        response.raise_for_status()
            except (requests.exceptions.ReadTimeout, requests.exceptions.ConnectTimeout) as timeout_error:
                print(f"Timeout occurred: {timeout_error}. Retrying in {retry_delay} seconds...")
                if attempt < max_retries:
                    time.sleep(retry_delay)
                    retry_delay *= 2  # Exponential backoff
                else:
                    raise  # Re-raise the exception after max retries
            except requests.exceptions.RequestException as req_error:
                print(f"Request exception: {req_error}. Retrying in {retry_delay} seconds...")
                if attempt < max_retries:
                    time.sleep(retry_delay)
                    retry_delay *= 2  # Exponential backoff
                else:
                    raise  # Re-raise the exception after max retries

        else:
            # If all retries failed
            print("Max retries exceeded. Failed to generate image.")
            return None

        with NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            image.save(tmp.name, format="PNG")
            constants.temp_files.append(tmp.name)
            print(f"Image saved to {tmp.name}")
            return tmp.name

    except Exception as e:
        print(f"Error generating AI image: {e}")
        return None