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

api_key = 'sk-CED85fi0ZhUDMWg4GvFQ5k53o7yoL7WOaPyPQcb8zPi7eDGi' # your Stability AI API key

def stable_diffusion(user_input):
    prompt = user_input
    print("Prompt:", prompt)
    
    # Create the Stable Difusion image using the Stability AI API
    response = requests.post(
        'https://api.stability.ai/v1/image/generate',
        headers={'Authorization': f'Bearer {api_key}'},
        json={
            'text': prompt,
            'model': 'sd-c2',
            'steps': 200,
            'size': 512,
            'start_size': 64,
            'start_scale': 1,
            'end_scale': 0.25,
            'deterministic': False,
            'clip': False,
            'top_k': 0,
            'top_p': 1,
            'temperature': 1,
            'seed': None,
            'noise_seed': None,
        }
    )
    response.raise_for_status()
    
    # Convert the image to a format that Gradio can display
    img_bytes = io.BytesIO(response.content)
    img = Image.open(img_bytes)
    img_arr = np.array(img)
    
    return img_arr

iface = gr.Interface(fn=stable_diffusion, inputs="text", outputs="image", title="bhAI (text to image)")
iface.launch()