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()