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 fancy():
    # try some https://huggingface.co/spaces/stabilityai/stable-diffusion/blob/main/app.py trix
    # for some bizarre reason this results in huggingface.co 'Temporary failure in name resolution'
    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 stretched 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="256px")
            
            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()
    block.queue().launch()

def plain():
    # trix from https://tmabraham.github.io/blog/gradio_hf_spaces_tutorial
    title = "Tiki Diffusion Model"
    description = '''
        Diffusion model trained on random tiki images. 
        FIXME:
            - runs very slow 120s - 240s
            - image is weirdly stretched
            - did this change make it to the front-end?
    '''
    article = gr.HTML('''
        <p><img src="blob/main/tiki-600e.png"></p>
        <p>Trained with <a href="https://github.com/huggingface/diffusers">diffusers</a>.</p>
    ''')
    
    # ValueError: The parameter `examples` must either be a string directory or a list(if there is only 1 input component) or (more generally), a nested list, where each sublist represents a set of inputs.
    examples = ['tiki-600e.png']
    
    interpretation = 'default' # no idea...
    enable_queue = True
    
    # https://github.com/gradio-app/gradio/issues/287
    css = '''
        .output_image   {height: 40rem !important; width: 100% !important;  } 
        .object-contain {height: 256px !important; width: 256px !important; }
    '''
    # css = ".output-image, .input-image, .image-preview {height: 600px !important}"
    
    inputs = None
    outputs = "pil"
    
    gr.Interface(
        fn=imagine, 
        inputs=inputs,     
        outputs=outputs,
        title=title,
        description=description,
        article=article,
        css=css,
        #examples=examples,
        interpretation=interpretation,
        enable_queue=enable_queue
    ).launch()

def main():
    if True:
        return fancy()
    plain()

        
main()