app.py

import gradio as gr
import time
from dotenv import load_dotenv
import os
from openai import OpenAI
import secrets

load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

if client.api_key is None:
    raise ValueError("Die Umgebungsvariable OPENAI_API_KEY ist nicht gesetzt.")

APP_PASSWORD = os.getenv("APP_PASSWORD")
if APP_PASSWORD is None:
    raise ValueError("Die Umgebungsvariable APP_PASSWORD ist nicht gesetzt.")

GENERATION_TIMEOUT = float(os.getenv("GENERATION_TIMEOUT", 60))

MAGIC_PROMPTS = [
    """You are an expert DALL-E 2 image generation prompt optimizer. Your task is to take the user's initial prompt idea and enhance it to produce the best possible results from DALL-E 2. Follow these steps:
Analyze the user's base prompt.
Expand the prompt by adding specific details about:
Subject/main elements
Setting/background
Lighting and atmosphere
Color palette
Artistic style (e.g. photorealistic, oil painting, digital art)
Composition and perspective
Incorporate descriptive adjectives and evocative language.
Add relevant artistic/technical terms (e.g. macro shot, fisheye lens, chiaroscuro).
Include quality boosters like "highly detailed", "award-winning", "stunning".
Specify image type if relevant (e.g. digital illustration, 35mm photograph).
Mention any desired emotions or moods.
Avoid negative language - focus on what should be included rather than excluded.
Keep text simple and minimal if text is required in the image.
Aim for a prompt length of 40-60 words.
Format the final prompt clearly, using proper punctuation and capitalization.
Only return the improved prompt without any additional comments or messages.""",
    """You are an expert DALL-E 3 image generation prompt optimizer. Enhance and refine anything the user sends you as an image generation prompt for DALL-E 3:

Provide an improved version of the user message by following these guidelines:
Adds more specific details about the scene, subjects, and atmosphere
Incorporates precise descriptors for colors, textures, and lighting
Specifies the artistic style or medium (e.g., oil painting, digital art, photography)
Includes relevant compositional elements (foreground, background, perspective)
Adds any missing context or setting information
Removes any redundant or vague language
Ensures the prompt is coherent and follows a logical structure
Incorporates relevant technical terms or jargon related to art or photography
Suggests any additional elements that could enhance the overall image
Optimizes the prompt length for DALL-E 3's capabilities (aim for 40-60 words)
Only return the improved prompt without any additional comments or messages.
""",
"Verbessere den folgenden Bildgenerierungsprompt, um ihn detaillierter und kreativer zu machen."
]

def enhance_prompt(original_prompt, img_model="dall-e-2"):
    system_message = MAGIC_PROMPTS[0] if img_model == "dall-e-2" else MAGIC_PROMPTS[1]

    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": original_prompt}
            ],
            max_tokens=300
        )
        enhanced_prompt = response.choices[0].message.content.strip()
        return enhanced_prompt
    except Exception as e:
        print("Ein Fehler ist beim Verbessern des Prompts aufgetreten:", e)
        return original_prompt

def generate_image(prompt, use_magic_prompt, style, password, user_id, last_generation_time, model):
    print(f"\nGenerierungsversuch von User {user_id} mit Modell {model} mit Prompt:\n{prompt}")

    if prompt=="": 
        return "Bitte gib einen Bildprompt ein.", None, last_generation_time
    
    if password != APP_PASSWORD:
        print("Generierung abgebrochen: Falsches Passwort.")
        return "Falsches Passwort. Bitte versuche es erneut.", None,  last_generation_time

    current_time = time.time()
    if model == "dall-e-3" and current_time - last_generation_time < GENERATION_TIMEOUT:
        remaining_time = max(0, int(GENERATION_TIMEOUT - (current_time - last_generation_time)))
        print("Generierung abgebrochen: Cooldown noch nicht abgelaufen")
        return f"Bitte warte noch {remaining_time} Sekunden bis zur nächsten Bildgenerierungen.\n\nNutze die Wartezeit, um deinen Prompt zu verfeinern. &#128521;", None,  last_generation_time

    if use_magic_prompt:
        prompt = enhance_prompt(prompt, model)
        print("\nFinaler Prompt:", prompt)

    try:
        response = client.images.generate(
            model=model,
            prompt=prompt,
            size="1024x1024",
            quality="standard",
            style=style,
            n=1
        )
        
        image_url = response.data[0].url
        
        # Update the last generation time for this user
        last_generation_time = current_time if model == "dall-e-3" else last_generation_time
        
        return prompt, image_url, last_generation_time
    except Exception as e:
        error_message = str(e)
        if "content_policy_violation" in error_message:
            return "Leider können wir deinen Prompt nicht verarbeiten, da er evtl. gegen die Content Policy von OpenAI verstösst. Wenn du denkst, dass das ein Fehler ist, versuche den Prompt so umzuformulieren, dass er jugendfreundlich ist.", None
        print("Ein Fehler ist aufgetreten:", e)
        return f"Ein Fehler ist beim Generieren des Bildes aufgetreten: {error_message}", None, last_generation_time

examples = [
    "Astronaut im Dschungel, kalte Farbpalette, gedämpfte Farben, detailliert, 8k",
    "Ein Astronaut reitet auf einem grünen Pferd",
    "Ein köstliches Ceviche-Käsekuchenstück",
    "Ein wunderschönes Glumanda in einer blühenden Wiese",
    "Ein futuristisches Stadtbild bei Nacht mit Neonlichtern",
    "Ein Roboter, der in einem antiken Tempel tanzt",
    "Ein magisches Einhorn, das über einen Regenbogen springt",
    "Ein mittelalterlicher Ritter, der gegen einen Drachen kämpft"
]

def update_button_state(last_generation_time):
    current_time = time.time()
    remaining_time = max(0, int(GENERATION_TIMEOUT - (current_time - last_generation_time)))
    if remaining_time > 0:
        return gr.update(value=f"Bitte warte {remaining_time} Sekunden", interactive=False)
    else:
        return gr.update(value="Bild generieren", interactive=True)
    
def update_welcome_message(user_id):
    return f"Willkommen, {user_id}!"

def generate_user_id():
    return secrets.token_urlsafe(16)

def initialize_session():
    new_user_id = generate_user_id()
    return new_user_id, update_welcome_message(new_user_id), 0

with gr.Blocks() as demo:
    user_id = gr.State(generate_user_id)
    last_generation_time = gr.State(0)

    gr.Markdown("# codora DALL-E Bildgenerator")
    welcome_message = gr.Markdown("Willkommen!")
    gr.Markdown(f"Gib einen Bildprompt ein und verwende optional die magische Prompt-Funktion, um ihn zu verbessern. Mit DALL·E 3 kannst du ein Bild alle {GENERATION_TIMEOUT} Sekunden generieren. Mit DALL·E 2 kannst du so viele Bilder generieren, wie du möchtest.")

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Bildprompt")
            use_magic_prompt = gr.Checkbox(label="Magischen Prompt verwenden")
            style = gr.Radio(
                label="Stil",
                choices=["vivid", "natural"],
                value="vivid"
            )
            model = gr.Radio(
                label="Modell",
                choices=["dall-e-2", "dall-e-3"],
                value="dall-e-2"
            )
            password = gr.Textbox(label="App Passwort", type="password")
            generate_button = gr.Button("Bild generieren")

        with gr.Column(scale=1):
            final_prompt = gr.Textbox(label="Finaler Prompt")
            generated_image = gr.Image(label="Generiertes Bild")  # Removed scale argument

    # Place examples below the generation button
    gr.Examples(
        examples=examples,
        inputs=[prompt]
    )

    generate_button.click(
        fn=generate_image,
        inputs=[prompt, use_magic_prompt, style, password, user_id, last_generation_time, model],
        outputs=[final_prompt, generated_image, last_generation_time]
    )

    demo.load(fn=initialize_session, outputs=[user_id, welcome_message, last_generation_time])

    demo.load(
        fn=update_button_state,
        inputs=[last_generation_time],
        outputs=[generate_button],
        every=1
    )

demo.launch()