import gradio as gr
import numpy as np
from PIL import Image
from typing import Optional

from models import remix_images

def generate_blank_image(width: int = 512, height: int = 512) -> Image.Image:
    """Generates a blank black PIL Image."""
    return Image.fromarray(np.zeros((height, width, 3), dtype=np.uint8))

with gr.Blocks(css="footer {visibility: hidden}") as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto;">
            <h1 style="font-size: 2.5em; font-weight: bold; margin-bottom: 0.5em;">Image Remixer</h1>
            <p style="font-size: 1.1em; color: #555;">Drag and drop up to three images, provide a text prompt, and let AI remix them!</p>
            <p style="font-size: 0.9em; color: #777;">Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #777; text-decoration: underline;">anycoder</a></p>
        </div>
        """
    )
    with gr.Row():
        with gr.Column():
            canny_input = gr.Image(label="Image for Canny Control (edges)", type="pil", height=256, value=generate_blank_image(), sources=["upload", "clipboard"], interactive=True)
            depth_input = gr.Image(label="Image for Depth Control (structure)", type="pil", height=256, value=generate_blank_image(), sources=["upload", "clipboard"], interactive=True)
            base_input = gr.Image(label="Base Image (optional, for img2img guidance and dimensions)", type="pil", height=256, value=generate_blank_image(), sources=["upload", "clipboard"], interactive=True)
        with gr.Column():
            prompt_input = gr.Textbox(label="Text Prompt", placeholder="A fantastical landscape, highly detailed, oil painting", lines=3, interactive=True)
            guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.1, label="Guidance Scale", interactive=True)
            num_inference_steps = gr.Slider(minimum=10, maximum=100, value=30, step=1, label="Inference Steps", interactive=True)
            
            with gr.Row():
                generate_btn = gr.Button("Remix Images", variant="primary")
                clear_btn = gr.ClearButton([canny_input, depth_input, base_input, prompt_input, output_image])
    
    with gr.Row():
        output_image = gr.Image(label="Remixed Output", interactive=False, height=512)

    generate_btn.click(
        fn=remix_images,
        inputs=[
            prompt_input,
            canny_input,
            depth_input,
            base_input,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=output_image,
        show_progress="minimal",
        api_name="remix_images"
    )

demo.launch(enable_monitoring=True, queue=True)