import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image

# ============================================================
# Fix: Monkey-patch transformers video_processing_auto bug
# Latest transformers has a bug where VIDEO_PROCESSOR_MAPPING
# is None, causing TypeError in video_processor_class_from_name
# ============================================================
try:
    from transformers.models.auto import video_processing_auto
    _original_func = video_processing_auto.video_processor_class_from_name

    def _patched_video_processor_class_from_name(class_name):
        try:
            return _original_func(class_name)
        except TypeError:
            # VIDEO_PROCESSOR_MAPPING_NAMES is None in some transformers versions
            return None

    video_processing_auto.video_processor_class_from_name = _patched_video_processor_class_from_name
    print("[PATCH] video_processor_class_from_name patched successfully")
except Exception as e:
    print(f"[PATCH] Could not patch video_processing_auto: {e}")

from diffusers import LongCatImageEditPipeline

# --- Load pipeline on CPU at init time ---
print("Loading LongCat-Image-Edit-Turbo pipeline...")
pipe = LongCatImageEditPipeline.from_pretrained(
    "meituan-longcat/LongCat-Image-Edit-Turbo",
    torch_dtype=torch.bfloat16,
)
print("Pipeline loaded on CPU.")


@spaces.GPU(duration=120)
def edit_image(
    input_image,
    prompt,
    negative_prompt="",
    guidance_scale=1.0,
    num_inference_steps=8,
    seed=43,
    randomize_seed=True,
):
    if input_image is None:
        raise gr.Error("이미지를 업로드해주세요 / Please upload an image.")
    if not prompt.strip():
        raise gr.Error("편집 프롬프트를 입력해주세요 / Please enter an editing prompt.")

    if randomize_seed:
        seed = int(np.random.randint(0, 2**31))

    # Enable CPU offload — model (~29GB) exceeds A10G VRAM (24GB)
    # Safe to call multiple times; ensures proper device mapping with ZeroGPU
    pipe.enable_model_cpu_offload()

    img = Image.fromarray(input_image).convert("RGB")

    result = pipe(
        img,
        prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=int(num_inference_steps),
        num_images_per_prompt=1,
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    ).images[0]

    return result, int(seed)


# ========== Gradio UI ==========
css = """
#col-main { max-width: 1200px; margin: 0 auto; }
.title-text { text-align: center; margin-bottom: 0.5em; }
footer { display: none !important; }
"""

with gr.Blocks(css=css, title="LongCat Image Edit Turbo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        <div class="title-text">
        
        # 🐱 LongCat Image Edit Turbo
        **Meituan LongCat** — High-quality image editing with only **8 inference steps**
        
        </div>
        """,
    )

    gr.Markdown(
        """
        > 💡 **Text Rendering Tip**: Wrap target text in quotes → `Change the text to 'Hello World'`  
        > 🌐 Supports **Chinese & English** prompts (中文/英文 프롬프트 지원)
        """
    )

    with gr.Row(elem_id="col-main"):
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="📷 Input Image",
                type="numpy",
                height=512,
            )
            prompt = gr.Textbox(
                label="✏️ Editing Prompt",
                placeholder="e.g. 将猫变成狗 / Add sunglasses / Change background to beach",
                lines=2,
            )
            negative_prompt = gr.Textbox(
                label="🚫 Negative Prompt (optional)",
                placeholder="e.g. blurry, low quality, distorted",
                lines=1,
            )

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        minimum=0.0,
                        maximum=5.0,
                        value=1.0,
                        step=0.1,
                    )
                    num_steps = gr.Slider(
                        label="Inference Steps",
                        minimum=4,
                        maximum=20,
                        value=8,
                        step=1,
                    )
                with gr.Row():
                    seed = gr.Number(label="Seed", value=43, precision=0)
                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)

            run_btn = gr.Button("🚀 Edit Image", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_image = gr.Image(label="🎨 Edited Result", type="pil", height=512)
            used_seed = gr.Number(label="Used Seed", interactive=False)

    # Examples
    gr.Examples(
        examples=[
            ["将猫变成狗"],
            ["Add sunglasses to the person"],
            ["Change the background to a beach"],
            ["Make it look like a watercolor painting"],
            ["把文字改成 'Hello World'"],
            ["Turn it into an anime style illustration"],
        ],
        inputs=[prompt],
        label="💡 Example Prompts",
    )

    run_btn.click(
        fn=edit_image,
        inputs=[input_image, prompt, negative_prompt, guidance_scale, num_steps, seed, randomize_seed],
        outputs=[output_image, used_seed],
    )

    gr.Markdown(
        """
        ---
        **Model**: [meituan-longcat/LongCat-Image-Edit-Turbo](https://huggingface.co/meituan-longcat/LongCat-Image-Edit-Turbo) | 
        **Paper**: [arxiv:2512.07584](https://arxiv.org/abs/2512.07584) | 
        **License**: Apache-2.0
        """
    )

demo.launch()