# Import necessary libraries
from together import Together
import gradio as gr
import os

# Retrieve API key from Hugging Face Secrets
API_KEY = os.getenv("MY_API_KEY")

client = Together(api_key=API_KEY)

def generate_image(prompt, model="black-forest-labs/FLUX.1-schnell", steps=4):
    try:
        if not prompt.strip():
            return None  # Avoid empty prompts
        print(f"Generating image for: {prompt}")

        
        response = client.images.generate(
            prompt=prompt,
            model=model,
            steps=steps
        )

        # Extract the image URL
        image_url = response.data[0].url
        print("Generated Image URL:", image_url)
        return image_url
    except Exception as e:
        print("Error generating image:", e)
        return None

# Gradio Interface Setup
with gr.Blocks() as app:
    gr.Markdown("## Real-Time AI Image Generator 🚀")
    gr.Markdown("Start typing a prompt below to generate images in real-time using FLUX Schnell model.")

    with gr.Row():
        prompt_box = gr.Textbox(label="Enter your prompt", placeholder="A horse on the moon")
        image_output = gr.Image(label="Generated Image")

    # Automatically trigger image generation when input changes
    prompt_box.change(fn=generate_image, inputs=prompt_box, outputs=image_output)

# Launch the Gradio app
app.launch()