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

api = replicate.Client(api_token="r8_9BTRdfQjCrVHkVMyQ6xAYJS52S6mLzx4YP6VA")

# Setting up Replicate AI API for stable diffusion
def generate_image(input_text, width=768, height=768, guidance_scale=7.5, num_inference_steps=50):
    guidance_scale = float(guidance_scale)
    output = api.run(
        "stability-ai/stable-diffusion:ac732df83cea7fff18b8472768c88ad041fa750ff7682a21affe81863cbe77e4",
        input={
            "width": width,
            "height": height,
            "prompt": input_text,
            "scheduler": "K_EULER",
            "num_outputs": 1,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps
        }
    )

    image_url = output[0]
    response = requests.get(image_url)
    image = Image.open(BytesIO(response.content))
    return image

# Setting up a gradio interface

iface = gr.Interface(
    fn=generate_image,
    inputs=[gr.Textbox(label="Prompt"),
            gr.Slider(label="Width", minimum = 64,  maximum = 768, step = 64),
            gr.Slider(label="Height", minimum = 64,  maximum = 768, step = 64),
            gr.Slider(label="Guidance Scale", minimum = 0,  maximum = 20, step = 0.5),
            gr.Radio([10,20,30,40,50], label="Inference Steps", info="Choose the number of Inference Steps")],
    outputs="image")

iface.launch(debug=True)