import os

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

# spaces MUST be imported before torch / diffusers so ZeroGPU can
# monkey-patch CUDA. This is how the working Ideogram 4 Spaces boot.
import spaces
import math
import sys

import torch
import gradio as gr
from diffusers import Ideogram4Pipeline
from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer
from PIL import Image

# cu130-era bitsandbytes returns Params4bit.shape as a plain tuple.
# diffusers' check_quantized_param_shape calls .numel() on it and crashes
# with: AttributeError: 'tuple' object has no attribute 'numel'
# Same shim as linoyts/ideogram-4-turbotime.
def _check_quantized_param_shape(self, param_name, current_param, loaded_param):
    current_shape = getattr(current_param, "shape", current_param)
    loaded_shape = getattr(loaded_param, "shape", loaded_param)
    n = math.prod(tuple(current_shape))
    inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
    if tuple(loaded_shape) != tuple(inferred_shape):
        raise ValueError(
            f"Expected flattened shape of {param_name} to be {inferred_shape}, "
            f"got {tuple(loaded_shape)}."
        )
    return True


BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape

MODEL_ID = "ideogram-ai/ideogram-4-nf4"
LORA_REPO = "ostris/ideogram_4_turbotime_lora"
LORA_FILE = "ideogram_4_turbotime_v1.safetensors"
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def log(msg: str) -> None:
    print(msg, file=sys.stderr, flush=True)


if not HF_TOKEN:
    raise RuntimeError(
        "Set an HF_TOKEN secret. Accept the gate on "
        "https://huggingface.co/ideogram-ai/ideogram-4-nf4"
    )

log("startup: loading diffusers Ideogram4Pipeline nf4")
pipe = Ideogram4Pipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
# Working ZeroGPU Spaces dequantize nf4 -> bf16 then move to emulated cuda
# at import. ZeroGPU packs those weights and streams them on first GPU call.
if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "dequantize"):
    pipe.transformer.dequantize()
if hasattr(pipe, "unconditional_transformer") and hasattr(pipe.unconditional_transformer, "dequantize"):
    pipe.unconditional_transformer.dequantize()
pipe.to("cuda")

try:
    pipe.transformer.load_lora_adapter(
        LORA_REPO,
        weight_name=LORA_FILE,
        prefix="diffusion_model",
        token=HF_TOKEN,
    )
    log("startup: TurboTime LoRA loaded (2-8 steps, guidance_scale=1)")
except Exception as exc:
    log(f"startup: LoRA load failed, will run base CFG path: {exc!r}")

log("startup: pipeline on cuda (ZeroGPU-emulated until Generate)")


def get_duration(prompt, width, height, steps, progress=None):
    # Must accept the same args as generate(), including Gradio's progress.
    tokens = (int(width) // 16) * (int(height) // 16)
    seconds = 25 + int(steps) * max(0.35, 0.00035 * tokens)
    return int(min(max(seconds, 60), 180))


@spaces.GPU(duration=get_duration, size="xlarge")
def generate(prompt, width, height, steps, progress=gr.Progress(track_tqdm=True)):
    log(
        f"generate() cuda={torch.cuda.is_available()} "
        f"{int(width)}x{int(height)} steps={int(steps)}"
    )
    if not (prompt or "").strip():
        raise gr.Error("Prompt is empty.")
    if not torch.cuda.is_available():
        raise gr.Error("ZeroGPU did not attach CUDA to this call.")

    # Magic prompt OFF: feed the textbox verbatim (same as --no-magic-prompt).
    # This diffusers build defaults guidance_schedule to the 48-step CFG
    # curve. Passing guidance_scale at the same time raises:
    # "Only one of guidance_scale and guidance_schedule may be set."
    # TurboTime is tuned for a constant scale of 1.0 (no CFG).
    kwargs = dict(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(steps),
        guidance_scale=1.0,
        guidance_schedule=None,
        generator=torch.Generator(device="cuda").manual_seed(0),
    )
    # Some diffusers builds use prompt_upsampling=; keep it off if present.
    try:
        images = pipe(**kwargs, prompt_upsampling=False).images
    except TypeError:
        images = pipe(**kwargs).images
    image = images[0]
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    log(f"generate() ok {image.size}")
    return image


with gr.Blocks(title="Ideogram 4 (ZeroGPU)") as demo:
    gr.Markdown(
        "# Ideogram 4 on ZeroGPU\n"
        "This is the stack the working Spaces use: **diffusers + nf4 + "
        "TurboTime LoRA + `size='xlarge'`**. Magic prompt is **off** — "
        "your prompt is used as-is.\n\n"
        "Needs an `HF_TOKEN` whose account accepted "
        "[ideogram-ai/ideogram-4-nf4](https://huggingface.co/ideogram-ai/ideogram-4-nf4). "
        "First Space boot can take a long time while weights download."
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                placeholder="A photo of a cat holding a sign that says hello world",
            )
            with gr.Row():
                width = gr.Slider(256, 2048, value=1024, step=16, label="Width")
                height = gr.Slider(256, 2048, value=1024, step=16, label="Height")
            steps = gr.Slider(2, 12, value=8, step=1, label="Steps (TurboTime LoRA)")
            run_btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            output = gr.Image(label="Result")
    run_btn.click(
        fn=generate,
        inputs=[prompt, width, height, steps],
        outputs=output,
        api_name="predict",
    )

demo.queue().launch(ssr_mode=False, show_error=True)