import gradio as gr
from image_generator import ImageGenerator

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

def greet(prompt, mix_prompt, mix_ratio, negative_prompt, steps, init_image ):

    print(f"{prompt=} {mix_prompt=} {mix_ratio=} {negative_prompt=} {steps=} {init_image=} ")
    generated_image, latents = ig.generate(
        prompt=prompt,
        secondary_prompt=mix_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

iface = gr.Interface(
    fn=greet, 
    inputs=[
        gr.Textbox(value="a cute dog", label="Prompt", info="primary prompt used to generate an image"),
        gr.Textbox(value=None, label="Secondary Prompt",  info="secondary prompt to mix with the primary embeddings"),
        gr.Slider(0, 1, value=0.5, label="Mix Ratio", info="mix ratio between primary and secondary prompt. 0 = primary only. 1 = secondary only"),
        gr.Textbox(value=None, label="Negative Prompt", info="remove certain aspect from the picture"),
        gr.Slider(10, 50, value=30, step=1, label="Generation Steps", info="How many steps are used to generate the picture"),
        gr.Image(type="pil", value=None, label="Starting Image",), # info="starting image from this image as opposed to random noise"
        ],
    outputs=[
        gr.Image(type="pil", label="Generated Image",),
        gr.Image(type="pil", label="Starting Image with Added Noise",)])
iface.launch()