File size: 2,585 Bytes
5156e7a
 
 
e7aa175
3f98781
 
 
 
 
 
e7aa175
3f98781
 
e7aa175
5156e7a
 
e7aa175
5156e7a
3f98781
e7aa175
5156e7a
b725492
3f98781
0abf9df
5156e7a
3f98781
 
0d83943
3f98781
0d83943
3f98781
b725492
 
0d83943
 
 
3f98781
 
 
 
 
 
b725492
 
 
3f98781
b725492
 
 
 
3f98781
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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()