import os
import gradio as gr
import torch

PIPELINE = None


def _get_pipeline():
    global PIPELINE
    if PIPELINE is None:
        if not torch.cuda.is_available():
            raise RuntimeError("This demo requires a CUDA-capable GPU runtime.")

        from diffusers import QwenImageEditPlusPipeline

        if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
            torch_dtype = torch.bfloat16
        else:
            torch_dtype = torch.float16

        PIPELINE = QwenImageEditPlusPipeline.from_pretrained(
            "Qwen/Qwen-Image-Edit-2511",
            torch_dtype=torch_dtype,
            use_safetensors=True,
        )
        if torch.cuda.is_available():
            PIPELINE.to("cuda")
        PIPELINE.set_progress_bar_config(disable=False)

    return PIPELINE


def generate_image(image1, image2, prompt, negative_prompt, guidance_scale, num_inference_steps, seed):
    if image1 is None:
        raise gr.Error("Please upload at least one input image before generating.")

    if not torch.cuda.is_available():
        raise gr.Error("This demo requires a CUDA-capable GPU runtime. Please run it on a GPU-backed Space or local machine.")

    images = [img for img in [image1, image2] if img is not None]
    if len(images) == 1:
        images.append(images[0])

    prompt_text = (prompt or "Turn this scene into a cinematic fantasy poster with glowing lanterns and soft mist.").strip()
    negative_text = negative_prompt or " "

    pipe = _get_pipeline()
    if torch.cuda.is_available():
        generator = torch.Generator(device="cuda").manual_seed(int(seed))
    else:
        generator = torch.Generator(device="cpu").manual_seed(int(seed))

    output = pipe(
        image=images,
        prompt=prompt_text,
        negative_prompt=negative_text,
        true_cfg_scale=float(guidance_scale),
        guidance_scale=1.0,
        num_inference_steps=int(num_inference_steps),
        num_images_per_prompt=1,
        generator=generator,
    )
    return output.images[0]


with gr.Blocks(title="Qwen Image Edit 2511 Demo") as demo:
    gr.Markdown("# Qwen Image Edit 2511 Demo")
    gr.Markdown(
        "This Hugging Face Space demonstrates the Qwen Image Edit 2511 model for guided image editing "
        "from one or two reference images. The first run may take a few minutes while the model loads."
    )

    with gr.Row():
        image_a = gr.Image(label="Reference image A", type="pil")
        image_b = gr.Image(label="Reference image B (optional)", type="pil")

    with gr.Row():
        prompt_box = gr.Textbox(
            label="Edit prompt",
            value="Turn this scene into a cinematic fantasy poster with glowing lanterns and soft mist.",
            lines=2,
        )
        negative_box = gr.Textbox(
            label="Negative prompt",
            value="blurry, low quality, text, watermark, distorted anatomy",
            lines=2,
        )

    with gr.Row():
        guidance_scale = gr.Slider(minimum=1.0, maximum=6.0, step=0.5, value=4.0, label="True CFG scale")
        num_inference_steps = gr.Slider(minimum=10, maximum=60, step=1, value=40, label="Inference steps")
        seed = gr.Number(value=0, precision=0, label="Seed")

    submit_btn = gr.Button("Generate edit")
    output_image = gr.Image(label="Edited image", type="pil")

    submit_btn.click(
        fn=generate_image,
        inputs=[image_a, image_b, prompt_box, negative_box, guidance_scale, num_inference_steps, seed],
        outputs=output_image,
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
    )