import gradio as gr
from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

#funtion to call model and show images
#function to make API call writing
def show_image(prompt):
    num_images = 2
    scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True)
    model_id = "runwayml/stable-diffusion-v1-5"

    device = "cuda"
    remove_safety = False

    pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16, revision="fp16", use_auth_token=True)
    if remove_safety:
        pipe.safety_checker = lambda images, clip_input: (images, False)

    pipe = pipe.to(device)

    prompts = [ prompt ] * num_images

    with autocast("cuda"):
        images = pipe(prompts, guidance_scale=7.5, num_inference_steps=50).images

    images[0].save("output.jpg")
    print(type(images[0]))
    return images[0]


demo = gr.Interface(fn=show_image, inputs="textbox", outputs=gr.Image(label = "Output Image Component"))

if __name__ == "__main__":
    demo.launch()