import json
import math
import os
import random
import time
from threading import Lock, Thread

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

import gradio as gr
import spaces
import torch
from diffusers import Ideogram4Pipeline
from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer
from huggingface_hub import hf_hub_download


MODEL_ID = "ideogram-ai/ideogram-4-nf4"
LM_HEAD_REPO = "multimodalart/qwen3-vl-8b-instruct-lm-head"
AOTI_REPO = "multimodalart/i4-block-aoti"
MAX_SEED = 2**31 - 1


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

pipe = None
pipe_lock = Lock()
_AOTI_APPLIED = False
_AOTI_OK = False


def get_pipe():
    global pipe, _AOTI_OK
    if pipe is not None:
        return pipe

    with pipe_lock:
        if pipe is None:
            print(f"Loading {MODEL_ID} pipeline...", flush=True)
            t = time.perf_counter()
            token = os.environ.get("HF_TOKEN") or None
            loaded_pipe = Ideogram4Pipeline.from_pretrained(
                MODEL_ID,
                torch_dtype=torch.bfloat16,
                token=token,
            )
            loaded_pipe.transformer.dequantize()
            loaded_pipe.unconditional_transformer.dequantize()
            loaded_pipe.to("cuda")
            pipe = loaded_pipe
            print(f"Pipeline loaded in {time.perf_counter() - t:.1f}s.", flush=True)

            try:
                hf_hub_download(AOTI_REPO, "package.pt2", subfolder="Ideogram4TransformerBlock")
                from torch._inductor.cpu_vec_isa import valid_vec_isa_list

                valid_vec_isa_list()
                _AOTI_OK = True
            except Exception as e:
                _AOTI_OK = False
                print(f"[aoti] prefetch/prewarm failed, running eager: {e!r}", flush=True)

    return pipe


def apply_aoti():
    global _AOTI_APPLIED
    if _AOTI_APPLIED or not _AOTI_OK or pipe is None:
        return

    try:
        t = time.perf_counter()
        spaces.aoti_blocks_load(pipe.transformer, AOTI_REPO)
        spaces.aoti_blocks_load(pipe.unconditional_transformer, AOTI_REPO)
        _AOTI_APPLIED = True
        print(f"AOTI blocks loaded in {time.perf_counter() - t:.2f}s.", flush=True)
    except Exception as e:
        print(f"[aoti] apply failed, running eager: {e!r}", flush=True)


_TOK_1024 = (1024 // 16) ** 2
_TOK_2048 = (2048 // 16) ** 2
_PS_1024 = 1.0 / 1.10
_PS_2048 = 6.0
_PS_B = (_PS_2048 - _PS_1024) / (_TOK_2048 - _TOK_1024)
_PS_A = _PS_1024 - _PS_B * _TOK_1024
LOCAL_UPSAMPLE_S = 15
DIFFUSION_OVERHEAD_S = 10
DURATION_MARGIN = 1.35


def per_step(width, height):
    tokens = (int(width) // 16) * (int(height) // 16)
    return max(0.2, _PS_A + _PS_B * tokens)


def gpu_duration(
    final_prompt,
    width,
    height,
    num_inference_steps,
    guidance_scale,
    image_count,
    seed,
    use_magic_prompt,
    progress=None,
):
    budget = int(num_inference_steps) * per_step(width, height) * int(image_count)
    budget += DIFFUSION_OVERHEAD_S
    if use_magic_prompt:
        budget += LOCAL_UPSAMPLE_S
    return max(60, int(math.ceil(budget * DURATION_MARGIN)))


def normalize_seed(seed):
    base_seed = int(seed) % MAX_SEED
    if base_seed < 0:
        base_seed += MAX_SEED
    return base_seed


def build_prompt(prompt, negative_prompt):
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    negative_prompt = (negative_prompt or "").strip()
    if negative_prompt:
        return f"{prompt}\nAvoid: {negative_prompt}"

    return prompt


def build_sampler_kwargs(num_inference_steps, guidance_scale):
    steps = max(1, int(num_inference_steps))
    guidance = float(guidance_scale)
    polish_steps = max(1, round(steps * 0.1))
    main_steps = max(0, steps - polish_steps)
    guidance_schedule = (guidance,) * main_steps + (3.0,) * polish_steps
    mu = 0.5 if steps <= 12 else 0.0
    std = 1.5 if steps >= 48 else 1.75
    return {
        "num_inference_steps": steps,
        "guidance_schedule": guidance_schedule,
        "mu": mu,
        "std": std,
    }


@spaces.GPU(duration=gpu_duration, size="xlarge")
def gpu_generate(
    final_prompt,
    width,
    height,
    num_inference_steps,
    guidance_scale,
    image_count,
    seed,
    use_magic_prompt,
    progress=gr.Progress(track_tqdm=True),
):
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available inside the ZeroGPU worker.")

    ideogram_pipe = get_pipe()

    aoti_thread = Thread(target=apply_aoti, daemon=True)
    aoti_thread.start()

    caption = final_prompt
    if use_magic_prompt:
        progress(0.0, desc="Upsampling prompt with local Qwen...")
        try:
            t = time.perf_counter()
            caption = ideogram_pipe.upsample_prompt(
                final_prompt,
                height=int(height),
                width=int(width),
                lm_head_repo_id=LM_HEAD_REPO,
            )[0]
            print(f"Prompt upsampled in {time.perf_counter() - t:.2f}s.", flush=True)
        except Exception as e:
            print(f"[upsample] local failed: {e!r}", flush=True)
            gr.Warning("Local prompt upsampler failed. Generating from the raw prompt.")

    aoti_thread.join()

    images = []
    seeds = [(normalize_seed(seed) + i) % MAX_SEED for i in range(int(image_count))]
    sampler_kwargs = build_sampler_kwargs(num_inference_steps, guidance_scale)

    for current_seed in seeds:
        progress(0.0, desc=f"Generating image with seed {current_seed}...")
        generator = torch.Generator(device="cuda").manual_seed(int(current_seed))
        image = ideogram_pipe(
            prompt=caption,
            width=int(width),
            height=int(height),
            generator=generator,
            **sampler_kwargs,
        ).images[0]
        images.append(image)

    try:
        parsed_caption = json.loads(caption)
    except Exception:
        parsed_caption = {"prompt": caption}

    return images, ", ".join(str(s) for s in seeds), parsed_caption


def generate_image(
    prompt,
    negative_prompt,
    height,
    width,
    num_inference_steps,
    guidance_scale,
    image_count,
    seed,
    randomize_seed,
    use_magic_prompt,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    final_prompt = build_prompt(prompt, negative_prompt)
    return gpu_generate(
        final_prompt,
        int(width),
        int(height),
        int(num_inference_steps),
        float(guidance_scale),
        int(image_count),
        normalize_seed(seed),
        bool(use_magic_prompt),
    )


examples = [
    [
        'A clean product poster for a premium coffee brand, exact text "MORNING RITUAL", warm studio lighting, elegant typography, realistic packaging'
    ],
    [
        'A movie poster for a retro sci-fi film, exact title text "ORBIT CITY", bold typography, neon lights, cinematic composition'
    ],
    [
        "A detailed editorial photo of a ceramic artist studio, handmade bowls, soft window light, natural colors, realistic texture"
    ],
    [
        'A square app icon for a weather app, exact text "SKYCAST", minimal vector style, blue and yellow palette, crisp edges'
    ],
]


with gr.Blocks(title="Ideogram 4 NF4 Demo") as demo:
    gr.Markdown(
        """
        # Ideogram 4 NF4 Demo

        Generate images using [ideogram-ai/ideogram-4-nf4](https://huggingface.co/ideogram-ai/ideogram-4-nf4) on ZeroGPU.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Enter your image description...",
                lines=4,
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Things you want the model to avoid...",
                lines=3,
            )

            with gr.Row():
                height = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Height",
                )
                width = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Width",
                )

            with gr.Row():
                num_inference_steps = gr.Slider(
                    minimum=8,
                    maximum=48,
                    value=20,
                    step=1,
                    label="Inference Steps",
                    info="Official presets are 12, 20, and 48 steps.",
                )
                image_count = gr.Slider(
                    minimum=1,
                    maximum=2,
                    value=1,
                    step=1,
                    label="Images",
                )

            guidance_scale = gr.Slider(
                minimum=1.0,
                maximum=10.0,
                value=7.0,
                step=0.1,
                label="CFG Guidance Scale",
                info="Ideogram 4 default main guidance is 7.0.",
            )

            use_magic_prompt = gr.Checkbox(
                label="Use Local Magic Prompt",
                value=True,
                info="Rewrites plain text into Ideogram's structured JSON prompt.",
            )

            with gr.Row():
                seed = gr.Number(
                    label="Seed",
                    value=0,
                    precision=0,
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize Seed",
                    value=True,
                )

            generate_btn = gr.Button("Generate", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_images = gr.Gallery(
                label="Generated Images",
                columns=2,
                rows=2,
                preview=True,
            )
            used_seeds = gr.Textbox(
                label="Seeds Used",
                interactive=False,
            )
            caption_json = gr.JSON(
                label="Caption Sent To Model",
            )

    gr.Markdown("### Example Prompts")
    gr.Examples(
        examples=examples,
        inputs=[prompt],
        cache_examples=False,
    )

    gr.Markdown(
        "Model by Ideogram. Ideogram 4 NF4 is gated and uses the Ideogram 4 Non-Commercial license."
    )

    inputs = [
        prompt,
        negative_prompt,
        height,
        width,
        num_inference_steps,
        guidance_scale,
        image_count,
        seed,
        randomize_seed,
        use_magic_prompt,
    ]
    outputs = [output_images, used_seeds, caption_json]

    generate_btn.click(
        fn=generate_image,
        inputs=inputs,
        outputs=outputs,
    )

    prompt.submit(
        fn=generate_image,
        inputs=inputs,
        outputs=outputs,
    )


if __name__ == "__main__":
    demo.launch(show_error=True)