import gradio as gr
import requests
import os

# Set your Stability AI API key
API_KEY = "your_stability_ai_api_key"

# Function to generate images using Stability AI API
def generate_image(prompt, steps, guidance):
    url = "https://api.stability.ai/v2beta/stable-image/generate/core"
    headers = {"Authorization": f"Bearer {API_KEY}"}
    payload = {
        "prompt": prompt,
        "steps": steps,
        "cfg_scale": guidance,
        "width": 512,
        "height": 512,
        "samples": 1
    }
    response = requests.post(url, headers=headers, json=payload)
    
    if response.status_code == 200:
        image_url = response.json()["artifacts"][0]["base64"]
        return f"data:image/png;base64,{image_url}"
    else:
        return f"Error: {response.json()}"

# Create Gradio UI
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter your prompt"),
        gr.Slider(minimum=10, maximum=50, value=20, step=1, label="Number of Steps"),
        gr.Slider(minimum=1, maximum=20, value=7.5, step=0.5, label="Guidance Scale"),
    ],
    outputs=gr.Image(label="Generated Image"),
    title="AI Image Generator",
    description="Generate images using Stability AI API. Enter a prompt and tweak settings.",
)

# Launch the app
iface.launch()