import gradio as gr
import torch

from model import generate_image
from utils import load_model, load_image, im_convert, device


model = load_model(d=device)

max_image_size = 400


def generate(content: torch.Tensor, style: torch.Tensor, alpha_slider: float):

    content_img = load_image(image=content, max_size=max_image_size).to(device)
    style_img = load_image(image=style, max_size=max_image_size, shape=content_img.shape[-2:]).to(device)

    target_img = content_img.clone().requires_grad_(True).to(
        device)  # Initialize the target image as a clone of the original content image

    steps = 300
    for target, status in generate_image(model=model, content=content_img, style=style_img, target=target_img, steps = steps, content_wt=alpha_slider):
        yield im_convert(target), status

    # return target


def check_inputs(img1, img2):
    """Enable the submit button only if both images are uploaded."""
    if img1 is not None and img2 is not None:
        return gr.update(interactive=True)  # Enable button
    return gr.update(interactive=False)  # Keep button disabled

with gr.Blocks() as demo:
    gr.Markdown("Transfer Image Style with the VGG19 model.")

    with gr.Row():
        content_image = gr.Image(type="pil", label="Original Image")
        style_image = gr.Image(type="pil", label="Style Reference Image")

    # Examples Section
    gr.Examples(
            examples=[
                ["./images/input-image-1.jpg", "./images/style-image-1.jpg"],
                ["./images/input-image-2.jpeg", "./images/style-image-2.jpg"]
            ],
            inputs=[content_image, style_image]
    )

    alpha_slider = gr.Slider(0, 1, value=1, step=0.1, label="Blending Ratio")
    submit_button = gr.Button("Blend Images", "Generate", variant="primary", interactive=False)

    output = gr.Image(type="pil", label="Blended Image")
    status = gr.Textbox(label="Progress")

    submit_button.click(generate, inputs=[content_image, style_image, alpha_slider], outputs=[output, status])

    # When images change, check if both are uploaded to enable the button
    content_image.change(fn=check_inputs, inputs=[content_image, style_image], outputs=submit_button)
    style_image.change(fn=check_inputs, inputs=[content_image, style_image], outputs=submit_button)


# Launch the demo!
demo.launch()

# if __name__ == "__main__":
#     demo.launch()