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

model_id = "stabilityai/stable-diffusion-2"

# Use the Euler scheduler here instead
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
pipe.enable_attention_slicing()
prompt = "a young doctor talk with a patient in the cartoon style"


num_imgs=3
scale=7.5
steps=100
i=0
def gen(num_imgs, scale, steps):
    image = pipe(num_imgs*[prompt], height=768, width=768,guidance_scale=scale, num_inference_steps=steps).images
    return image
gr.Interface(fn=gen, inputs=['text', gr.Slider(1, 100, 20), gr.Slider(1, maximum=20, value=10, step=1), gr.Slider(minimum=1, step=1, maximum=999999999999999999, randomize=True)], outputs='image', title="Stable Diffusion 2.0 ZQL", description="SD 2.0. <b>WARNING:</b> Extremely Slow. 130s/Iteration. Expect 25-50mins an image for 10-20 iterations respectively.", article = "Code Monkey: <a href=\"https://huggingface.co/qianli\">千里马</a>").launch()

for i in range(num_imgs):
    file_name="doctor"+str(i)+".png"
    image[i].save(file_name)