import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# outlines_core ships an @torch.compile bitmask kernel dynamo can't trace -> noisy WON'T CONVERT
# spam on every local upsample. We never torch.compile at runtime, so disable dynamo.
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import json
import math
import random
import time

import gradio as gr
import spaces
import torch
from diffusers import Ideogram4Pipeline, Ideogram4Transformer2DModel

try:
    from diffusers import Ideogram4PromptEnhancerHead
    _HAS_HEAD = True
except Exception:  # pragma: no cover - older diffusers without the enhancer head
    _HAS_HEAD = False

# Runtime shim: cu130-era bitsandbytes returns Params4bit.shape as a plain tuple,
# but diffusers' check_quantized_param_shape calls .numel() on it. math.prod handles both.
from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer


def _check_quantized_param_shape(self, param_name, current_param, loaded_param):
    n = math.prod(tuple(current_param.shape))
    inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
    if tuple(loaded_param.shape) != tuple(inferred_shape):
        raise ValueError(
            f"Expected flattened shape of {param_name} to be {inferred_shape}, "
            f"got {tuple(loaded_param.shape)}."
        )
    return True


BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape

MODEL_ID = "fal/ideogram-v4-instant"
COMPONENTS_REPO = "ideogram-ai/ideogram-4-nf4-diffusers"
COMPONENTS_REVISION = "1874bc70267ba2c823a7239e1d70dd308c8d64dc"
LM_HEAD_REPO = "diffusers/qwen3-vl-8b-instruct-lm-head"
HF_TOKEN = os.environ.get("HF_TOKEN")  # components repo is gated -> read it with the Space secret
MAX_SEED = 2**31 - 1

# AoTI: a precompiled ZeroGPUCompiledModel of the conditional transformer,
# produced offline by the compile Space and stored as a .pt2 blob. Loading it
# here and applying with spaces.aoti_apply avoids compiling per-request.
AOTI_REPO = "hugging-apps/ideogram-v4-instant-aoti"
AOTI_PT2_FILENAME = "transformer.pt2"

# Aspect-ratio presets -> (width, height). All multiples of 64.
ASPECT_RATIOS = {
    "1:1 · 1024×1024": (1024, 1024),
    "3:2 · 1216×832": (1216, 832),
    "2:3 · 832×1216": (832, 1216),
    "16:9 · 1344×768": (1344, 768),
    "9:16 · 768×1344": (768, 1344),
}
DEFAULT_RATIO = "1:1 · 1024×1024"

# --- Load model at module scope, .to("cuda") eagerly ---
# The fal/ideogram-v4-instant transformer is BF16 with split QKV weights.
# We load the full NF4 pipeline (which has both transformers), dequantize them,
# then override the conditional transformer with the fal instant checkpoint.
# guidance_schedule=[1.0]*8 makes the unconditional pass a no-op (v = pos_v),
# matching the model card's no-CFG single-branch behavior.

# Optional local prompt enhancer (Qwen3-VL LM head grafted onto the text encoder). Free, on-device.
enhancer_head = None
if _HAS_HEAD:
    try:
        enhancer_head = Ideogram4PromptEnhancerHead.from_pretrained(
            LM_HEAD_REPO, torch_dtype=torch.bfloat16, token=HF_TOKEN
        )
    except Exception as e:
        print(f"[enhancer] LM-head load failed (raw prompt only): {e!r}", flush=True)

t0 = time.perf_counter()
pipe = Ideogram4Pipeline.from_pretrained(
    COMPONENTS_REPO,
    revision=COMPONENTS_REVISION,
    prompt_enhancer_head=enhancer_head,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
pipe.transformer.dequantize()
pipe.unconditional_transformer.dequantize()

# Override the conditional transformer with the fal instant checkpoint
instant_transformer = Ideogram4Transformer2DModel.from_pretrained(
    MODEL_ID,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
pipe.transformer = instant_transformer

pipe.to("cuda")
print(f"[timing] pipeline load: {time.perf_counter() - t0:.1f}s", flush=True)

# --- Load the AoTI-compiled conditional transformer and apply it before serving ---
# Download the precompiled ZeroGPUCompiledModel (.pt2 blob) and swap it in with
# spaces.aoti_apply, so requests never pay a compile cost. Falls back to the eager
# transformer if the artifact is unavailable or incompatible.
try:
    import pickle

    from huggingface_hub import hf_hub_download

    t_aoti = time.perf_counter()
    _pt2_path = hf_hub_download(
        repo_id=AOTI_REPO,
        filename=AOTI_PT2_FILENAME,
        repo_type="model",
    )
    with open(_pt2_path, "rb") as _f:
        _compiled_transformer = pickle.load(_f)
    spaces.aoti_apply(_compiled_transformer, pipe.transformer)
    print(
        f"[timing] AoTI apply: {time.perf_counter() - t_aoti:.1f}s "
        f"({type(_compiled_transformer).__name__})",
        flush=True,
    )
except Exception as e:
    print(f"AoTI load failed ({e!r}); running eager transformer", flush=True)

# No-CFG guidance schedule: all ones means v = pos_v (conditional only)
NO_CFG_SCHEDULE = (1.0,) * 8


def _looks_like_json(text):
    s = (text or "").strip()
    return s.startswith("{") and s.endswith("}")


def _build_prompt(text_prompt: str) -> str:
    """Wrap a natural-language prompt into Ideogram 4's JSON caption format."""
    return json.dumps(
        {"high_level_description": text_prompt},
        ensure_ascii=False,
        separators=(",", ":"),
    )


# --- Warm the local prompt enhancer on the startup worker (forks inherit the graft) --------------------
@spaces.GPU(duration=120)
def _warmup():
    if enhancer_head is not None:
        pipe.upsample_prompt("a red apple on a wooden table", height=1024, width=1024)


if enhancer_head is not None:
    try:
        _warmup()
        print("[enhancer] prompt enhancer grafted", flush=True)
    except Exception as e:
        print(f"[enhancer] warmup failed (will graft lazily on first request): {e!r}", flush=True)


def _gpu_duration(prompt, aspect_ratio=DEFAULT_RATIO, enhance=True, seed=0, randomize_seed=True, *args, **kwargs):
    """Estimate GPU duration based on image resolution (8 fixed steps, 2 transformer passes)."""
    _TOK_1024 = (1024 // 16) ** 2  # 4096 image tokens
    _TOK_2048 = (2048 // 16) ** 2  # 16384
    _PS_1024 = 1.0 / 1.10  # ~0.91 s/it per transformer
    _PS_2048 = 6.0  # 6 s/it per transformer
    _PS_B = (_PS_2048 - _PS_1024) / (_TOK_2048 - _TOK_1024)
    _PS_A = _PS_1024 - _PS_B * _TOK_1024
    width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_RATIO])
    tok = (int(width) // 16) * (int(height) // 16)
    per_step = max(0.2, _PS_A + _PS_B * tok)
    # 2 transformer calls per step (conditional + unconditional, even though uncond is a no-op)
    budget = 8 * per_step * 2 + 12  # 12s for text encoding + VAE + cold-start overhead
    if enhance:
        budget += 20  # local prompt upsampling (grafted Qwen head)
    return max(30, min(240, int(math.ceil(budget * 1.3))))


@spaces.GPU(duration=_gpu_duration)
def generate(
    prompt: str,
    aspect_ratio: str = DEFAULT_RATIO,
    enhance: bool = True,
    seed: int = 0,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt using Ideogram 4 Instant.

    An 8-step distilled text-to-image model by fal with no runtime CFG, producing
    high-quality images—including text rendering—in seconds. Ideogram 4 is trained
    on structured JSON captions, so a plain prompt is optionally expanded into one
    on-device (Qwen3-VL) before generation — or paste your own JSON caption.

    Args:
        prompt: A plain-text prompt (expanded into Ideogram's JSON caption when `enhance` is on) or a
            complete structured JSON caption fed to the model verbatim.
        aspect_ratio: One of the preset aspect-ratio / resolution labels.
        enhance: Expand a plain-text prompt into Ideogram's structured JSON caption before generation.
        seed: RNG seed (ignored when `randomize_seed` is on).
        randomize_seed: Draw a fresh random seed each run.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if randomize_seed or seed is None or int(seed) < 0:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_RATIO])
    generator = torch.Generator(device="cuda").manual_seed(seed)

    # Ideogram 4 is trained on structured JSON captions. If the user typed JSON, honour it verbatim;
    # otherwise (when `enhance` is on and the enhancer is available) upsample the plain prompt into a
    # native caption with the on-device Qwen3-VL head. Toggle off to feed a minimal JSON wrapper.
    text = prompt.strip()
    if _looks_like_json(text):
        final_prompt = text  # already a JSON caption
    elif enhance and enhancer_head is not None:
        progress(0.0, desc="✍️ Writing the JSON caption…")
        try:
            final_prompt = pipe.upsample_prompt(
                text, height=height, width=width, generator=generator
            )[0]
        except Exception as e:
            print(f"[enhancer] failed, using raw prompt: {e!r}", flush=True)
            gr.Warning("Prompt enhancer unavailable — generating from the raw prompt.")
            final_prompt = _build_prompt(text)
    else:
        final_prompt = _build_prompt(text)

    progress(0.0, desc="🎨 Generating…")
    t = time.perf_counter()
    image = pipe(
        prompt=final_prompt,
        height=height,
        width=width,
        num_inference_steps=8,
        guidance_schedule=NO_CFG_SCHEDULE,
        mu=0.0,
        std=1.75,
        generator=generator,
    ).images[0]
    dt = time.perf_counter() - t
    print(f"[timing] diffusion (8 steps, {width}x{height}): {dt:.2f}s", flush=True)

    try:
        caption = json.loads(final_prompt)
    except (TypeError, ValueError):
        caption = {"prompt": final_prompt}
    return image, seed, caption, f"8 steps · {dt:.1f}s"


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

DEFAULT_PROMPT = "a ginger cat wearing a tiny wizard hat reading a spellbook"

with gr.Blocks(title="Ideogram 4 Instant · by fal") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Ideogram 4 Instant ⚡ — by fal\n"
            "[**fal/ideogram-v4-instant**](https://huggingface.co/fal/ideogram-v4-instant) is a speed-distilled "
            "Ideogram 4 checkpoint: **8 steps, no runtime CFG**. Ideogram 4 is trained on "
            "**structured JSON captions**, so a plain prompt is expanded into one on-device (Qwen3-VL) before "
            "generation — or paste your own JSON caption.\n\n"
            "[Model](https://huggingface.co/fal/ideogram-v4-instant) · "
            "[Base Ideogram 4](https://huggingface.co/ideogram-ai/ideogram-4-nf4-diffusers) · "
            "[fal blog](https://blog.fal.ai/serving-sub-second-ideogram-v4-without-quality-loss/)"
        )

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    value=DEFAULT_PROMPT,
                    lines=3,
                    info="Plain text (auto-expanded to a JSON caption) or a full structured JSON caption.",
                )
                run = gr.Button("Generate", variant="primary")
                aspect_ratio = gr.Radio(
                    choices=list(ASPECT_RATIOS.keys()), value=DEFAULT_RATIO, label="Aspect ratio"
                )
                with gr.Accordion("Advanced settings", open=False):
                    enhance = gr.Checkbox(
                        label="Enhance prompt → JSON caption",
                        value=True,
                        info="Ideogram 4 is trained on structured captions. On = best quality (recommended). "
                        "Ignored when the prompt is already JSON.",
                    )
                    with gr.Row():
                        seed = gr.Number(label="Seed", value=0, precision=0)
                        randomize = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Column():
                out_image = gr.Image(label="Output", type="pil")
                with gr.Row():
                    out_seed = gr.Number(label="Seed used", precision=0, interactive=False)
                    out_time = gr.Textbox(label="Generation", interactive=False)
                out_caption = gr.JSON(label="Caption fed to the model")

        gr.Examples(
            examples=[
                ["a ginger cat wearing a tiny wizard hat reading a spellbook"],
                ["A bold typographic poster with the words HELLO WORLD in vibrant gradient colors"],
                ["an isometric illustration of a tiny city floating in the clouds"],
                ["a cinematic photo of a golden retriever puppy in a field of sunflowers at golden hour"],
            ],
            inputs=[prompt],
            outputs=[out_image, out_seed, out_caption, out_time],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        generate,
        inputs=[prompt, aspect_ratio, enhance, seed, randomize],
        outputs=[out_image, out_seed, out_caption, out_time],
        api_name="generate",
    )

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