# 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