import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before any CUDA-touching import
import torch
import gradio as gr
import random
import numpy as np
from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
from huggingface_hub import hf_hub_download

dtype = torch.bfloat16
device = "cuda"

# Tiny VAE for fast preview, good VAE for final output (same pattern as official FLUX.1-dev space)
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype
).to(device)

# Load base FLUX.1-dev with tiny VAE
BASE_MODEL = "black-forest-labs/FLUX.1-dev"
CMO_LORA = "Bruece/FLUX.1-dev-CMO"
LORA_ALPHA = 128
LORA_R = 64
LORA_SCALE = LORA_ALPHA / LORA_R  # = 2.0

pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL, torch_dtype=dtype, vae=taef1
).to(device)

# Manually load and merge the CMO LoRA adapter weights.
# Use safetensors.safe_open (numpy backend, no torch) to avoid ZeroGPU's
# torch patching which fails at module scope (no CUDA available yet).
from safetensors import safe_open

_lora_path = hf_hub_download(CMO_LORA, "adapter_model.safetensors", repo_type="model")

# Load all LoRA weights as numpy arrays first, then merge into the transformer
_lora_a_pairs = {}  # module_path -> A weight (numpy)
_lora_b_pairs = {}  # module_path -> B weight (numpy)

with safe_open(_lora_path, framework="pt", device="cpu") as _f:
    for _key in _f.keys():
        if not _key.startswith("base_model.model."):
            continue
        _rest = _key[len("base_model.model."):]
        if _rest.endswith(".lora_A.weight"):
            _module_path = _rest[: -len(".lora_A.weight")]
            _lora_a_pairs[_module_path] = _f.get_tensor(_key)
        elif _rest.endswith(".lora_B.weight"):
            _module_path = _rest[: -len(".lora_B.weight")]
            _lora_b_pairs[_module_path] = _f.get_tensor(_key)

# Merge LoRA weights into the transformer: w_new = w_orig + scale * (B @ A)
_merge_count = 0
for _module_path, _a_tensor in _lora_a_pairs.items():
    if _module_path not in _lora_b_pairs:
        continue
    _b_tensor = _lora_b_pairs[_module_path]
    # Navigate to the module in the transformer
    _module = pipe.transformer
    for _part in _module_path.split("."):
        _module = getattr(_module, _part)
    # Merge: w_orig + scale * (B @ A)
    _delta = (_b_tensor.float() @ _a_tensor.float()) * LORA_SCALE
    _module.weight.data.add_(_delta.to(_module.weight.data.dtype))
    _merge_count += 1

print(f"CMO LoRA: merged {_merge_count} adapter pairs into FLUX.1-dev transformer")
del _lora_a_pairs, _lora_b_pairs
torch.cuda.empty_cache()

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


def _estimate(prompt, num_inference_steps, *args, **kwargs):
    return min(180, 30 + int(num_inference_steps * 4))


@spaces.GPU(duration=_estimate)
def generate(
    prompt: str,
    seed: int = 42,
    randomize_seed: bool = True,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 4.5,
    num_inference_steps: int = 40,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt using FLUX.1-dev fine-tuned with CMO.

    CMO (Correlation-Weighted Multi-Reward Optimization) improves compositional
    text-to-image generation by adaptively weighting conflicting concept rewards
    (object existence, attributes, numeracy, size, spatial relations).

    Args:
        prompt: Text description of the image to generate.
        seed: RNG seed for reproducibility.
        randomize_seed: If True, pick a random seed each run.
        width: Output image width in pixels.
        height: Output image height in pixels.
        guidance_scale: Classifier-free guidance scale.
        num_inference_steps: Number of denoising steps.
    """
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]
    return image, seed


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["A red apple is on the left of the yellow banana"],
    ["Two Tyrannosaurus rexes engaged in a boxing match"],
    ["a photo of a black kite and a green bear"],
    ["A brown cow wearing yellow sunglasses in a pastel chalk drawing"],
    ["a cat holding a sign that says hello world"],
    ["The green plant was on top of the white nightstand"],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """# FLUX.1-dev CMO — Compositional Text-to-Image
FLUX.1-dev fine-tuned with **Correlation-Weighted Multi-Reward Optimization (CMO)** for improved compositional generation.
[[Paper](https://huggingface.co/papers/2603.18528)] [[Code](https://github.com/TheDarkKnight-21th/CMO)] [[Model](https://huggingface.co/Bruece/FLUX.1-dev-CMO)]
"""
        )
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt (e.g. 'A red apple is on the left of the yellow banana')",
                container=False,
                scale=4,
            )
            run_button = gr.Button("Run", variant="primary", scale=1)

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=42,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=4.5,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=40,
                )

        gr.Examples(
            examples=EXAMPLES,
            fn=generate,
            inputs=[prompt],
            outputs=[result, seed],
            cache_examples=True,
            cache_mode="lazy",
        )

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

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)