import torch
from diffusers import DDIMPipeline, DDPMPipeline, PNDMPipeline
from diffusers import DDIMScheduler, DDPMScheduler, PNDMScheduler
from diffusers import UNet2DModel
import gradio as gr
import PIL.Image
import numpy as np
import random

model_id = "google/ddpm-celebahq-256"
model = UNet2DModel.from_pretrained(model_id)

# load model and scheduler
ddpm_scheduler = DDPMScheduler.from_config(model_id)
ddpm_pipeline = DDPMPipeline(unet=model, scheduler=ddpm_scheduler)

ddim_scheduler = DDIMScheduler.from_config(model_id)
ddim_pipeline = DDIMPipeline(unet=model, scheduler=ddim_scheduler)

pndm_scheduler = PNDMScheduler.from_config(model_id)
pndm_pipeline = PNDMPipeline(unet=model, scheduler=pndm_scheduler)

# run pipeline in inference (sample random noise and denoise)
def predict(steps=100, seed=42,scheduler="ddim"):
    torch.cuda.empty_cache()
    generator = torch.manual_seed(seed)
    if(scheduler == "ddim"):
        images = ddim_pipeline(generator=generator, num_inference_steps=steps)["sample"]
    elif(scheduler == "ddpm"):
        images = ddpm_pipeline(generator=generator)["sample"]
    elif(scheduler == "pndm"):
        if(steps > 100):
            steps = 100
        images = pndm_pipeline(generator=generator, num_inference_steps=steps)["sample"]

    return(images[0])


random_seed = random.randint(0, 2147483647)
gr.Interface(
    predict,
    inputs=[
        gr.inputs.Slider(1, 1000, label='Inference Steps (ignored for the ddpm scheduler, that diffuses for 1000 steps - limited to 100 steps max for pndm)', default=20, step=1),
        gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed, step=1),
        gr.inputs.Radio(["ddpm", "ddim", "pndm"], default="ddpm",label="Diffusion scheduler")
    ],
    outputs=gr.Image(shape=[256,256], type="pil", elem_id="output_image"),
    css="#output_image{width: 256px}",
    title="ddpm-celebahq-256 diffusion - 🧨 diffusers library",
    description="This Spaces contains an unconditional diffusion process for the <a href=\"https://huggingface.co/google/ddpm-celebahq-256\">ddpm-celebahq-256</a> face generator model by <a href=\"https://huggingface.co/google\">Google</a> using the <a href=\"https://github.com/huggingface/diffusers\">diffusers library</a>. You can try the diffusion process not only with the default <code>ddpm</code> scheduler but also with <code>ddim</code> and <code>pndm</code>, showcasing the modularity of the library. <a href=\"https://github.com/huggingface/diffusers/tree/main/src/diffusers/schedulers\">Learn more about schedulers here.</a>",
).launch()