hf-diffusers / app.py
mrm8488's picture
Set UI stuff
b0c5c4b
raw history blame
No virus
1.15 kB
from diffusers import DDPMPipeline
import gradio as gr
from ui import title, description, examples
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='Choose a model')
],
outputs=["image"],
title=title,
description=description
).launch()