import gradio as gr
from diffusers import DDPMPipeline, DDIMPipeline, PNDMPipeline

MODEL_ID = "verkaDerkaDerk/tiki-based-128"
MODEL_ID = "verkaDerkaDerk/tiki-64"
PIPELINE = DDPMPipeline.from_pretrained(MODEL_ID)
# doesn't work... PIPELINE = PIPELINE.to("cuda")

#def imagine(name):
def imagine():
    image = PIPELINE()["sample"][0]
    return image

def main():
    if True:
        # try some https://huggingface.co/spaces/stabilityai/stable-diffusion/blob/main/app.py trix
        block = gr.Blocks()
        with block:
            gr.HTML('''
                <pre>
                This is an unconditioned diffusion model trained on around 500 
                miscellaneous tiki images from around the web.
                
                It was trained for around 1300 epochs with a loss of 0.0249
                or something... still running fix this later... uhm...
                
                It may take a really long time like 120s for some reason.
                
                The image size is hugely stretch atm... still trying to figure out 
                how to control it...
                </pre>
            ''')
            with gr.Group():
                with gr.Box():
                    with gr.Row():
                        btn = gr.Button("Generate image")
                       
                gallery = gr.Gallery(
                    label="Generated images", show_label=False, elem_id="gallery"
                ).style(grid=[1], height="128px")
                btn.click(imagine, inputs=None, outputs=gallery)
                gr.HTML('''
                    Trained with <a href="https://github.com/huggingface/diffusers">diffusers</a>.
                ''')
        block.queue(max_size=40).launch()
    else:
        inputs = "text"
        inputs = None
        outputs = "pil"
        gr.Interface(fn=imagine, inputs=inputs, outputs=outputs).launch()

        
main()