hf-diffusers / app.py
mrm8488's picture
Initial commit
90da81e
raw
history blame
1.05 kB
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()