import gradio as gr
import requests
import base64
from io import BytesIO
from PIL import Image

def generate_image(api_key, prompt, size):
    sizes = {
        "512x512": (512, 512),
        "1024x1024": (1024, 1024),
        "512x624": (512, 624),
        "624x512": (624, 512)
    }
    
    height, width = sizes[size]
    
    url = 'https://api.stability.ai/v1/generation/stable-diffusion-v1-6/text-to-image'
    
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    
    data = {
        'text_prompts': [{'text': prompt}],
        'cfg_scale': 7,
        'height': height,
        'width': width,
        'samples': 1,
        'steps': 30
    }
    
    response = requests.post(url, headers=headers, json=data)
    
    if response.status_code == 200:
        response_data = response.json()
        image_base64 = response_data['artifacts'][0]['base64']
        image = Image.open(BytesIO(base64.b64decode(image_base64)))
        return image
    else:
        return f"Error: {response.json()['error']}"

api_key_input = gr.Textbox(label="API Key", type="password")
prompt_input = gr.Textbox(label="Prompt")
size_input = gr.Dropdown(choices=["512x512", "1024x1024", "512x624", "624x512"], label="Image Size")

iface = gr.Interface(
    fn=generate_image,
    inputs=[api_key_input, prompt_input, size_input],
    outputs=gr.Image(type="pil"),
    title="Stable Diffusion Image Generator",
    description="Enter your Stability AI API key, a prompt, and select the image size to generate an image using Stable Diffusion 1.6."
)

iface.launch()