import gradio as gr
from together import Together
import base64
from PIL import Image
import io
import logging
import requests
from datetime import datetime
from dotenv import load_dotenv
import os
import json
import psycopg2
from psycopg2 import Error
from urllib.parse import urlparse
from functools import lru_cache
import time
from typing import Tuple, Any, Optional
import genarationData
# Load environment variables
load_dotenv()
# Logging configuration
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# Initialize Together client
api_key = os.getenv("TOGETHER_API_KEY")
client = Together(api_key=api_key)
# Configuration constants
IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")
POSTGRES_URL = os.getenv("POSTGRES_URL")
DEBUG = os.getenv("DEBUG", "False").lower() == "true"
CACHE_TIMEOUT = 3600
MAX_RETRIES = 3
RETRY_DELAY = 1
@lru_cache(maxsize=1)
def get_db_config() -> dict:
"""Parse PostgreSQL URL and return connection configuration."""
parsed = urlparse(POSTGRES_URL)
return {
'host': parsed.hostname,
'user': parsed.username,
'password': parsed.password,
'database': parsed.path.strip('/'),
'port': parsed.port,
}
def get_db_connection() -> psycopg2.extensions.connection:
"""Establish a connection to the PostgreSQL database with retry logic."""
for attempt in range(MAX_RETRIES):
try:
config = get_db_config()
connection = psycopg2.connect(**config)
if connection:
logger.info('Successfully connected to PostgreSQL database')
return connection
except Error as e:
if attempt == MAX_RETRIES - 1:
logger.error(f"Final attempt failed to connect to PostgreSQL: {e}")
raise
time.sleep(RETRY_DELAY)
raise Exception("Failed to connect to database after retries")
def init_db():
"""Initialize the database schema if it doesn't exist."""
try:
with get_db_connection() as connection:
with connection.cursor() as cursor:
cursor.execute('''
CREATE TABLE IF NOT EXISTS generated_images (
id SERIAL PRIMARY KEY,
generation_prompt TEXT NOT NULL,
generation_timestamp TIMESTAMP NOT NULL,
generation_width INT NOT NULL,
generation_height INT NOT NULL,
generation_steps INT NOT NULL,
imgbb_id VARCHAR(255) NOT NULL,
imgbb_title VARCHAR(255),
imgbb_url_viewer TEXT,
imgbb_url TEXT,
imgbb_display_url TEXT,
imgbb_width VARCHAR(50),
imgbb_height VARCHAR(50),
imgbb_size VARCHAR(50),
imgbb_time VARCHAR(50),
imgbb_expiration VARCHAR(50),
delete_url TEXT,
raw_response TEXT,
user_id VARCHAR(255)
);
''')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_timestamp ON generated_images (generation_timestamp);')
cursor.execute('CREATE INDEX IF NOT EXISTS idx_imgbb_id ON generated_images (imgbb_id);')
connection.commit()
logger.info("Database initialized successfully")
except Error as e:
logger.error(f"Error initializing database: {e}")
raise
def retry_with_backoff(func):
"""Decorator for functions to retry with exponential backoff."""
def wrapper(*args, **kwargs):
for i in range(MAX_RETRIES):
try:
return func(*args, **kwargs)
except Exception as e:
if i == MAX_RETRIES - 1:
raise
wait = (2 ** i) * RETRY_DELAY
logger.warning(f"Attempt {i+1} failed, retrying in {wait} seconds...")
time.sleep(wait)
return wrapper
@retry_with_backoff
#def upload_to_imgbb(image_bytes: bytes) -> dict:
# """Upload image to ImgBB and return the response."""
# img_base64 = base64.b64encode(image_bytes).decode('utf-8')
# url = "https://api.imgbb.com/1/upload"
# payload = {
# "key": IMGBB_API_KEY,
# "image": img_base64
# }
# response = requests.post(url, payload, timeout=30)
# response.raise_for_status()
# logger.info("Successfully uploaded to ImgBB")
# return response.json()
@retry_with_backoff
def generate_image(prompt: str, width: int, height: int, steps: int) -> Tuple[Image.Image, bytes]:
"""Generate an image using the Together API."""
if not prompt.strip():
raise ValueError("Please enter a prompt")
logger.info(f"Generating image with parameters: width={width}, height={height}, steps={steps}")
response = client.images.generate(
prompt=prompt,
model="black-forest-labs/FLUX.1-schnell-Free",
width=width,
height=height,
steps=steps,
n=1,
response_format="b64_json"
)
image_bytes = base64.b64decode(response.data[0].b64_json)
image = Image.open(io.BytesIO(image_bytes))
logger.info("Image generated successfully")
return image, image_bytes
def handle_generation(prompt: str, width: int, height: int, steps: int) -> Tuple[Optional[Image.Image], str]:
"""Handle the entire process of image generation, upload, and database save."""
try:
image, image_bytes = generate_image(prompt, width, height, steps)
imgbb_response = upload_to_imgbb(image_bytes)
if save_to_database(prompt, width, height, steps, imgbb_response):
return image, "Image generated successfully!"
else:
return image, "Image generated and uploaded, but database save failed!"
except Exception as e:
logger.error(f"Error in handle_generation: {str(e)}")
return None, f"Error: {str(e)}"
def save_to_database(prompt: str, width: int, height: int, steps: int, imgbb_response: dict) -> bool:
"""Save image generation details to the database."""
try:
with get_db_connection() as connection:
with connection.cursor() as cursor:
data = imgbb_response['data']
insert_query = '''
INSERT INTO generated_images (
generation_prompt, generation_timestamp, generation_width, generation_height,
generation_steps, imgbb_id, imgbb_title, imgbb_url_viewer, imgbb_url,
imgbb_display_url, imgbb_width, imgbb_height, imgbb_size, imgbb_time,
imgbb_expiration, delete_url, raw_response, user_id
) VALUES (
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
)
'''
values = (
prompt, datetime.now(), width, height, steps,
data.get('id'), data.get('title'), data.get('url_viewer'),
data.get('url'), data.get('display_url'), data.get('width'),
data.get('height'), data.get('size'), data.get('time'),
data.get('expiration'), data.get('delete_url'),
json.dumps(imgbb_response), None # Assuming user_id is optional
)
cursor.execute(insert_query, values)
connection.commit()
logger.info("Successfully saved to database")
return True
except Error as e:
logger.error(f"Database error: {e}")
return False
def create_demo():
"""Create and return the Gradio demo interface."""
with gr.Blocks(css="style.css", theme="NoCrypt/miku", title="Elixir Craft Image Generator") as demo:
gr.HTML("""
Elixir Craft Image Generator
Elixir Craft Image Generator 🖼️
Welcome to the AI Image Generator powered by the FLUX.1 [schnell] model! 🎨
Enter a description of any scene, character, or object you'd like to see come to life, adjust image dimensions,
and select the number of steps to control image detail. Click "Generate Image" to create your
custom artwork in seconds!
Features:
Generate high-quality images from text descriptions
Discover all images generated with the FLUX.1 AI Image Generator. Each creation is stored in the gallery for
you to view, share, or download. Every image includes the prompt details and settings.