import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from torch import autocast
def generate_image(prompt, model_id, guidance_scale):
  """Generate image"""
  img_path = "image.png"
  with autocast(device):
    img = pipes[model_id](prompt, guidance_scale=guidance_scale)["images"][0]
  img.save(img_path)
  return img_path

# the model id
model_ids = ["CompVis/stable-diffusion-v1-4"]
# check for device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"The code is runing on {device}")
# load the p model pipeline
pipes = {
    model_id:StableDiffusionPipeline.from_pretrained(model_id, variant="fp16", torch_dtype=torch.float32).to(device)
    for model_id in model_ids
}

example_prompt = "a photo of an astronaut riding a horse on mars"
# gradio interface
interface = gr.Interface(
    fn=generate_image,
    inputs= [gr.Textbox(label="Prompt", value=example_prompt),
             gr.Dropdown(model_ids,label="Select the model ID"),
             gr.Slider(minimum=0.1,maximum=10, value=7.5, label="Guidance Scale")],
    outputs="image",
    live=False,
    title="Text to Image",
    description="Enter the prompt in text input, click generate to generete the image")
# lauch the gradio
interface.launch(share=True, debug=True) # set share to False if you dont want to make public