import requests
import gradio as gr
from PIL import Image
import numpy as np

# Replace with your Dream API key
api_key = 'sk-CED85fi0ZhUDMWg4GvFQ5k53o7yoL7WOaPyPQcb8zPi7eDGi'

def generate_image(prompt):
    # Set up the request headers
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    # Set up the request data
    data = {
        "model": "stable-diffusion",
        "prompt": prompt,
        "steps": 100,
        "batch_size": 1,
        "gamma": 0.99,
        "device": "cpu"
    }

    # Send the request to Dream's API
    response = requests.post("https://api.dream.co/stable-diffusion/generate", json=data, headers=headers)
    response.raise_for_status()

    # Extract the image URL from the response
    image_url = response.json()["data"][0]["url"]

    # Download and display the image using PIL and Gradio
    image_bytes = requests.get(image_url).content
    image = Image.open(io.BytesIO(image_bytes))
    image_arr = np.array(image)
    return image_arr

iface = gr.Interface(fn=generate_image, inputs="text", outputs="image", title="bhAI (text to image using Dream Studio's Stable Difusion)", description="Enter a prompt to generate an image.")
iface.launch()