app.py

import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv
import os
import requests
import base64
from PIL import Image
from io import BytesIO

load_dotenv()

# Initialize OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# System prompt (to be updated)
SYSTEM_PROMPT = """
You are tasked with enhancing user prompts to generate clear, detailed, and creative descriptions for a sticker creation AI. The final prompt should be imaginative, visually rich, and aligned with the goal of producing a cute, stylized, and highly personalized sticker based on the user's input.

Instructions:

Visual Clarity: The enhanced prompt must provide clear visual details that can be directly interpreted by the image generation model. Break down and elaborate on specific elements of the scene, object, or character based on the user input.

Example: If the user says "A girl with pink hair," elaborate by adding features like "short wavy pink hair with soft, pastel hues."
Style & Theme:

Emphasize that the final output should reflect a cute, playful, and approachable style.
Add terms like "adorable," "cartoonish," "sticker-friendly," or "chibi-like" to guide the output to a lighter, cuter aesthetic.
Include styling prompts like “minimalistic lines,” “soft shading,” and “vibrant yet soothing colors.”
Personalization:

If a reference or context is given, enhance it to make the sticker feel personalized.
Add context-appropriate descriptors like “wearing a cozy blue hoodie,” “soft pink blush on cheeks,” or “a playful expression.”
Expression & Pose:

Where applicable, refine prompts with suggestions about facial expressions or body language. For example, “Smiling softly with big sparkling eyes” or “A cute wink and a slight tilt of the head.”
Background & Accessories:

Optionally suggest simple, complementary backgrounds or accessories, depending on user input. For instance, "A light pastel background with small, floating hearts" or "Holding a tiny, sparkling star."
Colors:

Emphasize the color scheme based on the user's description, making sure it's consistent with a cute, playful style.
Use descriptors like “soft pastels,” “bright and cheerful,” or “warm and friendly hues” to set the mood.
Avoid Overcomplication:

Keep the descriptions short enough to be concise and not overly complex, as the output should retain a sticker-friendly quality.
Avoid unnecessary details that could clutter the design.
Tone and Language:

The tone should be light, imaginative, and fun, matching the playful nature of stickers.

Example:
User Input:
"A girl with pink hair wearing a hoodie."

Enhanced Prompt:
"An adorable girl with short, wavy pink hair in soft pastel hues, wearing a cozy light blue hoodie. She has a sweet smile with big, sparkling eyes, and a playful expression. The sticker style is cartoonish with minimalistic lines and soft shading. The background is a simple light pastel color with small floating hearts, creating a cute and inviting look."
"""

# Function to enhance the user's prompt
def enhance_prompt(user_prompt) -> str:
    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt}
        ]
    )
    ep = completion.choices[0].message.content
    print('Enhanced Prompt:', ep)
    return ep

# Function to generate images using the selected models
def generate_images(user_prompt, selected_models):
    enhanced_prompt = enhance_prompt(user_prompt)
    images = []
    headers = {
        "Authorization": f"Bearer {os.getenv('AIMLAPI_API_KEY')}",
    }
    for model in selected_models:
        try:
            payload = {
                "prompt": enhanced_prompt,
                "model": model,
            }
            response = requests.post(
                "https://api.aimlapi.com/images/generations", headers=headers, json=payload
            )
            if response.status_code == 201:
                response_json = response.json()
                print(f"Response for model {model}: {response_json}")
                # Handle OpenAI models differently (Aspect 2)
                if model in ["dall-e-3", "dall-e-2"]:
                    if 'data' in response_json and 'url' in response_json['data'][0]:
                        image_url = response_json['data'][0]['url']
                        image_response = requests.get(image_url)
                        image = Image.open(BytesIO(image_response.content))
                        images.append(image)
                    else:
                        print(f"No URL found for model {model}")
                else:
                    # Handle other models (Aspect 1)
                    if 'images' in response_json and 'url' in response_json['images'][0]:
                        image_url = response_json['images'][0]['url']
                        image_response = requests.get(image_url)
                        image = Image.open(BytesIO(image_response.content))
                        images.append(image)
                    else:
                        print(f"No URL found for model {model}")
            else:
                print(f"Error with model {model}: {response.text}")
        except Exception as e:
            print(f"Exception occurred with model {model}: {e}")
            continue
    return images

# List of available image generation models
model_list = [
    "stable-diffusion-v35-large",
    "flux-pro/v1.1",
    "dall-e-3",
    "stable-diffusion-v3-medium",
    "runwayml/stable-diffusion-v1-5",
    "stabilityai/stable-diffusion-xl-base-1.0",
    "stabilityai/stable-diffusion-2-1",
    "SG161222/Realistic_Vision_V3.0_VAE",
    "prompthero/openjourney",
    "wavymulder/Analog-Diffusion",
    "flux-pro",
    "flux-realism",
    "dall-e-2",
]

# Gradio Interface
with gr.Blocks() as demo:
    # Title and links
    with gr.Row():
        gr.Markdown("""
            # Let's Generate Cutesy AI Sticker!
            <p align="center">
            <a title="Page" href="https://ai-sticker-maker.vercel.app/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                <img src="https://img.shields.io/badge/Project-Website-pink?logo=googlechrome&logoColor=white">
            </a>
            <a title="arXiv" href="https://rebrand.ly/aistickermakerpaper" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                <img src="https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=white">
            </a>
            <a title="Github" href="https://github.com/abdibrokhim/ai-sticker-maker" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                <img src="https://img.shields.io/github/stars/EnVision-Research/Lotus?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
            </a>
            <a title="Social" href="https://x.com/abdibrokhim" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
            </a>
            <a title="Social" href="https://x.com/haodongli00/status/1839524569058582884" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
            </a>
            <br>
            <strong>Please consider starring <span style="color: orange">&#9733;</span> the <a href="https://github.com/abdibrokhim/ai-sticker-maker" target="_blank" rel="noopener noreferrer">GitHub Repo</a> if you find this useful!</strong>
        """)
    with gr.Row():
        with gr.Column(scale=1):
            # Model selection
            selected_models = gr.CheckboxGroup(
                choices=model_list,
                label="Select Image Generation Models",
                value=["stable-diffusion-v35-large"]
            )
        with gr.Column(scale=2):
            # User prompt input
            # Example propt: a very cutesy panda sitting and easting a pink very creamy ice cream
            user_prompt = gr.Textbox(
                placeholder="A girl with short pink hair wearing an oversize hoodie...",
                label="Enter your prompt here"
            )
            # Generate button
            generate_button = gr.Button("Generate Images")

    # Outputs
    image_outputs = gr.Gallery(label="Generated Images", columns=[3], rows=[1], elem_id="gallery")

    # Function to run on button click
    def on_click(user_prompt, selected_models):
        images = generate_images(user_prompt, selected_models)
        # Filter out None values in case of errors
        return [img for img in images if img is not None]

    # Event binding
    generate_button.click(
        fn=on_click,
        inputs=[user_prompt, selected_models],
        outputs=image_outputs
    )

# Launch the Gradio app
demo.launch()