File size: 1,046 Bytes
90da81e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from diffusers import DDPMPipeline
import gradio as gr



models = [
{'type': 'pokemon', 'res': 64, 'id': 'mrm8488/ddpm-ema-pokemon-64'},
{'type': 'flowers', 'res': 64, 'id': 'mrm8488/ddpm-ema-flower-64'},
{'type': 'anime_faces', 'res': 128, 'id': 'mrm8488/mrm8488/ddpm-ema-anime-256'},
{'type': 'butterflies', 'res': 128, 'id': 'mrm8488/ddpm-ema-butterflies-128'},
{'type': 'human_faces', 'res': 256, 'id': 'fusing/ddpm-celeba-hq'}
]

for model in models:
    pipeline = DDPMPipeline.from_pretrained(model['id'])
    pipeline.save_pretrained(model['id'])

def predict(type):
    model_id = None
    for model in models:
        if model['type'] == type:
            model_id = model['id']
            break
    # load model and scheduler
    pipeline = DDPMPipeline.from_pretrained(model_id)

    # run pipeline in inference 
    image = pipeline()["sample"]

    return image[0]

gr.Interface(
    predict,
    inputs=[gr.components.Dropdown(choices = [model['type'] for model in models], label='Models')
    ],
    outputs=["image"]
).launch()