import os
import numpy as np
import gradio as gr
from utils.i2i import i2i_gen

MAX_SEED = np.iinfo(np.int32).max

with gr.Blocks(
    title="🪄 LayerDiffuse - Flux version (Image to Image)",
    theme="CultriX/gradio-theme"
) as demo:
    gr.Markdown(
        """
        # 🪄 LayerDiffuse - Flux version (Image to Image)

        A Flux version implementation of LayerDiffuse for image-to-image generation.
        
        Upload an image with transparency (PNG with alpha channel) and transform it with a text prompt.
        The output image will maintain the same aspect ratio as your input image.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Input Image",
                type="pil",
                sources=["upload", "webcam", "clipboard"],
                image_mode="RGBA",
                height=512
            )
            
            prompt = gr.Text(
                label="Prompt",
                info="Describe what you want to generate",
                placeholder="E.g: a beautiful landscape with mountains, high quality"
            )
            
            with gr.Row():
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                    randomize=True
                )
                
                strength = gr.Slider(
                    label="Strength",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.01,
                    value=0.5,
                    info="How much to transform the image (0 = no change, 1 = complete change)"
                )
            
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=1,
                    maximum=20,
                    step=0.1,
                    value=7.0,
                )
                
                num_inference_steps = gr.Slider(
                    label="Steps",
                    minimum=10,
                    maximum=100,
                    step=1,
                    value=50,
                )
            
            generate_btn = gr.Button("Generate", variant="primary")
        
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Result",
                type="pil",
                height=512
            )
    
    # Set up the generation flow
    gr.on(
        triggers=[generate_btn.click],
        fn=lambda: gr.update(interactive=False, value="Generating..."),
        outputs=generate_btn,
        api_name=False
    ).then(
        fn=i2i_gen,
        inputs=[
            input_image,
            prompt,
            seed,
            guidance_scale,
            num_inference_steps,
            strength
        ],
        outputs=output_image
    ).then(
        fn=lambda: gr.update(interactive=True, value="Generate"),
        outputs=generate_btn,
        api_name=False
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_error=True, share=True)