import gradio as gr
from image_generator import ImageGenerator
import os


header = """Hi! This HuggingFace Space is a demo for the homework from the [10th lesson](https://course.fast.ai/Lessons/lesson10.html) of the fast.ai course. You can pick some of the examples below and click the "Generate Image" Button.

The code demonstrates:
* how to use an existing image as a starting point for the output image generation, in addition to the prompt
* how to use negative prompt
* how to capture latents through the generation
* how to mix prompt embeddings

![a cute dog mixed with a bird](/file=ai-math.png "a cute dog mixed with a bird")

"""


ig = ImageGenerator()
print(ig)
ig.load_models()
ig.load_scheduler()

def call(prompt, secondary_prompt, mix_ratio, negative_prompt, steps, init_image ):

    print(f"{prompt=} {secondary_prompt=} {mix_ratio=} {negative_prompt=} {steps=} {init_image=} ")
    generated_image, latents = ig.generate(
        prompt=prompt,
        secondary_prompt=secondary_prompt,
        prompt_mix_ratio=mix_ratio,    
        negative_prompt=negative_prompt,
        steps=steps,
        init_image=init_image, 
        latent_callback_mod=None )

    if init_image is not None:
        noisy_latent = latents[1]
    else:
        noisy_latent = None

    return generated_image, noisy_latent

def update_noisy_image_visibility(init_image):
    if init_image is None:
        print("update_noisy_image_visibility: hide noisy image")
        return gr.Image(type="pil", label="Starting Image with Added Noise", visible=False)
    else:
        print("update_noisy_image_visibility: show noisy image")
        return gr.Image(type="pil", label="Starting Image with Added Noise", visible=True)

def run_inference(prompt="", secondary_prompt="", mix_ratio=0.5, negative_prompt="", guidance=7.5, steps=10, init_image=None, progress=gr.Progress()): #, mix_ratio, negative_prompt, steps, starting_image, load_set_btn, 
    print(f"{prompt=} {secondary_prompt=} {mix_ratio=} {negative_prompt=} {steps=} {init_image=} ")

    generated_image, latents = ig.generate( 
        prompt=prompt,
        secondary_prompt=secondary_prompt,
        prompt_mix_ratio=mix_ratio,    
        negative_prompt=negative_prompt,
        guidance=guidance,
        steps=steps,
        init_image=init_image, 
        latent_callback_mod=1,
        progress_tqdm=progress.tqdm )

    if init_image is not None:
        noisy_latent = latents[1]
    else:
        noisy_latent = None

    return generated_image, noisy_latent, ig.image_grid(latents)

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown(value=header)
    with gr.Row():
        with gr.Column(scale=1):        
            prompt = gr.Textbox(value="a cute dog", label="Prompt", info="primary prompt used to generate an image")
            secondary_prompt = gr.Textbox(value=None, label="Secondary Prompt",  info="secondary prompt to mix with the primary embeddings")
            mix_ratio = gr.Slider(0, 1, value=0.5, label="Mix Ratio", info="mix ratio between primary and secondary prompt. 0 = primary only. 1 = secondary only")
            negative_prompt = gr.Textbox(value=None, label="Negative Prompt", info="remove certain aspect from the picture")
            guidance = gr.Slider(0, 14, value=7.5, label="Guidance", info="how closely the model should follow the prompt (higher the closer)")
            steps = gr.Slider(10, 50, value=10, step=1, label="Generation Steps", info="How many steps are used to generate the picture")
            init_image = gr.Image(type="pil", value=None, label="Starting Image",) # info="starting image from this image as opposed to random noise"            
            generate_image_btn = gr.Button("Generate Image")

        with gr.Column(scale=1):
            output_image = gr.Image(type="pil", label="Generated Image",)
            noisy_image = gr.Image(type="pil", label="Starting Image with Added Noise", visible=False)
            noisy_image.change(fn=update_noisy_image_visibility, inputs=init_image, outputs=noisy_image)
            latent_images = gr.Image(type="pil", label="Latents through the denoising process", visible=True)
    
    with gr.Row():
        # broken images should be fixed soon https://github.com/gradio-app/gradio/issues/5067
        gr.Examples(
            examples=[
                # simple prompt
                ["a cute dog", "", "", "", 7.5, 10, None],

                # negative prompt
                ["a beautiful tree", "", "", "green", 7.5, 10, None],

                # with base image
                ["a painting of Paris at night in the style of Pierre Auguste Renoir", "", "", "", 7.5, 50, os.path.join( os.path.dirname(__file__), "examples/ex4.jpg")],

                # with prompt
                ["a sloth", "a jaguar", 0.5, "", 7.5, 30, None],
            ],
            inputs=[prompt, secondary_prompt, mix_ratio, negative_prompt, guidance, steps, init_image],
            outputs=[output_image, noisy_image, latent_images],
            fn=run_inference,
            cache_examples=False)

    generate_image_btn.click(
        fn=run_inference, 
        inputs=[prompt, secondary_prompt, mix_ratio, negative_prompt, guidance, steps, init_image],
        outputs=[output_image, noisy_image, latent_images])

demo.launch(allowed_paths=["./"]) #allowed_paths=["/tmp/"]