'''
python scripts/gradio_demo.py 
'''

import sys
import os
workspace_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "icedit"))

if workspace_dir not in sys.path:
    sys.path.insert(0, workspace_dir)
    
from diffusers import FluxFillPipeline
import gradio as gr
import numpy as np
import torch
import argparse
import random 
import spaces
from PIL import Image

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

current_lora_scale = 1.0


parser = argparse.ArgumentParser() 
parser.add_argument("--port", type=int, default=7860, help="Port for the Gradio app")
parser.add_argument("--output-dir", type=str, default="gradio_results", help="Directory to save the output image")
parser.add_argument("--flux-path", type=str, default='black-forest-labs/flux.1-fill-dev', help="Path to the model")
parser.add_argument("--lora-path", type=str, default='sanaka87/ICEdit-MoE-LoRA', help="Path to the LoRA weights")
parser.add_argument("--enable-model-cpu-offload", action="store_true", help="Enable CPU offloading for the model")
args = parser.parse_args()

pipe = FluxFillPipeline.from_pretrained(args.flux_path, torch_dtype=torch.bfloat16)
pipe.load_lora_weights(args.lora_path, adapter_name="icedit")
pipe.set_adapters("icedit", 1.0)
    
if args.enable_model_cpu_offload:
    pipe.enable_model_cpu_offload() 
else:
    pipe = pipe.to("cuda")

@spaces.GPU
def infer(edit_images, 
          prompt, 
          seed=666, 
          randomize_seed=False, 
          width=1024, 
          height=1024, 
          guidance_scale=50, 
          num_inference_steps=28, 
          lora_scale=1.0,
          progress=gr.Progress(track_tqdm=True)
):
    

    global current_lora_scale
    
    if lora_scale != current_lora_scale:
        print(f"\033[93m[INFO] LoRA scale changed from {current_lora_scale} to {lora_scale}, reloading LoRA weights\033[0m")
        pipe.set_adapters("icedit", lora_scale)
        current_lora_scale = lora_scale
        
    
    image = edit_images
        
    if image.size[0] != 512:
        print("\033[93m[WARNING] We can only deal with the case where the image's width is 512.\033[0m")
        new_width = 512
        scale = new_width / image.size[0]
        new_height = int(image.size[1] * scale)
        new_height = (new_height // 8) * 8  
        image = image.resize((new_width, new_height))
        print(f"\033[93m[WARNING] Resizing the image to {new_width} x {new_height}\033[0m")
        
    image = image.convert("RGB")
    width, height = image.size
    image = image.resize((512, int(512 * height / width)))
    combined_image = Image.new("RGB", (width * 2, height))
    combined_image.paste(image, (0, 0)) 
    mask_array = np.zeros((height, width * 2), dtype=np.uint8)
    mask_array[:, width:] = 255 
    mask = Image.fromarray(mask_array)
    instruction = f'A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but {prompt}'

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    output_image = pipe(
        prompt=instruction,
        image=combined_image,
        mask_image=mask,
        height=height,
        width=width*2,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=torch.Generator().manual_seed(seed),
    ).images[0]

    w,h = output_image.size
    output_image = output_image.crop((w//2, 0, w, h))

    os.makedirs(args.output_dir, exist_ok=True)
        
    index = len(os.listdir(args.output_dir))
    output_image.save(f"{args.output_dir}/result_{index}.png")
    
    return (image, output_image), seed, lora_scale
    
# 新增的示例，将元组转换为列表
new_examples = [
    ['assets/girl_3.jpg', 'Make it looks like a watercolor painting.', 0, 0.5],
    ['assets/girl.png', 'Make her hair dark green and her clothes checked.', 42, 1.0],
    ['assets/boy.png', 'Change the sunglasses to a Christmas hat.', 27440001, 1.0],
    ['assets/kaori.jpg', 'Make it a sketch.', 329918865, 1.0]
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 1000px;
}
"""

with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# IC-Edit
**Image Editing is worth a single LoRA!** A demo for [IC-Edit](https://river-zhang.github.io/ICEdit-gh-pages/).
More **open-source**, with **lower costs**, **faster speed** (it takes about 9 seconds to process one image), and **powerful performance**.
For more details, check out our [Github Repository](https://github.com/River-Zhang/ICEdit) and [arxiv paper](https://arxiv.org/pdf/2504.20690). If our project resonates with you or proves useful, we'd be truly grateful if you could spare a moment to give it a star.
\n**👑 Feel free to share your results in this [Gallery](https://github.com/River-Zhang/ICEdit/discussions/21)!**
\n🔥 New feature: Try **different LoRA scale**!
""")
        with gr.Row():
            with gr.Column():
                edit_image = gr.Image(
                    label='Upload image for editing',
                    type='pil',
                    sources=["upload", "webcam"],
                    image_mode='RGB',
                    height=600
                )
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                run_button = gr.Button("Run")
            with gr.Column():
                result = gr.ImageSlider(label="Result", show_label=False)
                gr.Markdown("⚠️ If your edit didn't work as desired, **try again with another seed** ! <br> If you use our example, don't forget to uncheck the random seed option. Otherwise, it will still use a random seed.")
        with gr.Accordion("Advanced Settings", open=True):

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():

                width = gr.Slider(
                    label="Width",
                    minimum=512,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                    visible=False
                )

                height = gr.Slider(
                    label="Height",
                    minimum=512,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                    visible=False
                )

            with gr.Row():

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=100,
                    step=0.5,
                    value=50,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

            lora_scale = gr.Slider(
                label="LoRA Scale",
                minimum=0,
                maximum=1.0,
                step=0.01,
                value=1.0,
            )
        gr.Examples(
            examples=new_examples,
            inputs=[edit_image, prompt, seed, lora_scale],
            outputs=[result, seed, lora_scale],
            fn=infer,
            cache_examples=False
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[edit_image, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_scale],
        outputs=[result, seed, lora_scale]
    )

demo.launch(server_port=args.port)