import gradio as gr 
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch

model_id = "stabilityai/stable-diffusion-2"
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
device = troch.device("cuda" if torch.cuda.is_avaliable() else "cpu")

pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16")
pipe = pipe.to(device)

def txt2img(prompt):
  image = pipe(prompt, height=768, width=768, guidance_scale = 10).images[0]
  # image.save("sd_image.png")
  return image
    
gr.Interface(txt2img, gr.Text(), gr.Image(), title = 'Stable Diffusion 2.0 Colab with Gradio UI').launch(share = False, debug = True)