app.py

import inspect
import math
import os
import random
import re
from typing import Optional

import gradio as gr
import spaces
import torch
from diffusers import Ideogram4Pipeline
import json
import math
import random
import time
from threading import Thread

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModel
import logging
import traceback
import time

logger = logging.getLogger("ideogram")


# ---------------------------------------------------------------------------
# Compatibility shim: diffusers `main` assumes `current_param.shape` is a
# torch.Size (which has `.numel()`), but recent bitsandbytes returns a plain
# tuple from `Params4bit.shape`, so loading the nf4 checkpoint crashes with
# "'tuple' object has no attribute 'numel'". `math.prod(tuple(shape))` is
# version-agnostic and semantically identical. Remove once diffusers pins/fixes
# the bitsandbytes interaction.
# ---------------------------------------------------------------------------
try:
    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 the flattened shape of the current param ({param_name}) "
                f"to be {tuple(loaded_param.shape)} but is {tuple(inferred_shape)}."
            )
        return True

    BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape
except Exception as _exc:  # noqa: BLE001 - shim must never be fatal
    print(f"[startup] bnb shape-check shim not applied: {_exc!r}")

APP_TITLE = "Realism Engine Ideogram 4"

# Official, diffusers-format, nf4-quantised Ideogram 4 (gated -> needs HF_TOKEN + accepted gate).
# nf4 keeps the 9.3B DiT + Qwen3-VL-8B text encoder within a single 24GB A10G (ZeroGPU).
BASE_REPO = "ideogram-ai/ideogram-4-nf4"
DEFAULT_MAIN_GUIDANCE = 7.0
DEFAULT_FINAL_GUIDANCE = 3.0
TEXT_ENCODER_ID = os.environ.get("TEXT_ENCODER_ID", "huihui-ai/Huihui-Qwen3-VL-8B-Instruct-abliterated")
# Realism Engine LoRA lives in the user's own repo (not gated).


# LORA_REPO = "RazzzHF/Realism_Engine_Ideogram_4"
# LORA_WEIGHT = "Realism_Engine_Ideogram_V2.safetensors"
# LORA_ADAPTER = "realism_engine"


LORA_REPO = "RazzzHF/Realism_Engine_Ideogram_4"
LORA_WEIGHT = "Realism_Engine_Ideogram_V4.safetensors"
LORA_ADAPTER = "realism_engine"




HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
MAX_SEED = 2**31 - 1

# Recommended schedules, mirrored from the original ComfyUI workflow notes.
# (Ideogram 4 runs TWO 9.3B transformers per step, so fewer steps ≈ proportionally
# faster — Turbo is the default to keep generation snappy.)
PRESETS = {
    "Ludicrous": {"steps": 8, "mu": 0.5, "std": 1.75},
    "Turbo":     {"steps": 12, "mu": 0.5, "std": 1.75},
    "Default":   {"steps": 20, "mu": 0.0, "std": 1.75},
    "Quality":   {"steps": 48, "mu": 0.0, "std": 1.50},
}


# ---------------------------------------------------------------------------
# Load the pipeline once at startup. On ZeroGPU, `.to("cuda")` here is handled
# by the `spaces` runtime, so the model is resident before the first request.
# ---------------------------------------------------------------------------
# pipe = Ideogram4Pipeline.from_pretrained(
#     BASE_REPO,
#     torch_dtype=torch.bfloat16,
#     token=HF_TOKEN,
# )


if TEXT_ENCODER_ID:
    text_encoder = AutoModel.from_pretrained(
        TEXT_ENCODER_ID,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN,
        low_cpu_mem_usage=True,
    )
    print(f"[model] using alternate text encoder: {TEXT_ENCODER_ID}", flush=True)
else:
    text_encoder = None

    
pipe = Ideogram4Pipeline.from_pretrained(
    BASE_REPO,
    text_encoder=text_encoder,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
pipe.transformer.dequantize()
pipe.unconditional_transformer.dequantize()

pipe.to("cuda")

# Attach the Realism Engine LoRA. It's an ai-toolkit LoRA trained on Ideogram 4
# (keys like "diffusion_model.layers.N.attention.qkv.lora_A.weight"), so the inner
# module names are diffusers-native — only the "diffusion_model." prefix needs
# stripping. This diffusers build's Ideogram4Pipeline has no pipeline-level
# load_lora_weights, so we inject on the transformer via PEFT. Any failure falls
# back to the (working) base model.
LORA_OK = False
LORA_HOW = ""
_LORA_TARGETS = []  # transformer modules the adapter was injected into


def _convert_ai_toolkit_lora(raw: dict) -> dict:
    """Map ai-toolkit Ideogram 4 LoRA keys onto the diffusers transformer layout:
      - strip the "diffusion_model." prefix
      - attention.o            -> attention.to_out.0
      - attention.qkv (fused)  -> attention.to_q / to_k / to_v
            lora_A is shared; lora_B's [3H, r] rows split into thirds (q, k, v).
      - adaln_modulation, feed_forward.w1/w2/w3 already match -> pass through.
    """
    qkv_re = re.compile(r"^(.*\.attention)\.qkv\.(lora_[AB])\.weight$")
    o_re = re.compile(r"^(.*\.attention)\.o\.(lora_[AB])\.weight$")
    out = {}
    for k, v in raw.items():
        key = k[len("diffusion_model."):] if k.startswith("diffusion_model.") else k

        m = qkv_re.match(key)
        if m:
            base, ab = m.group(1), m.group(2)
            if ab == "lora_A":
                for proj in ("to_q", "to_k", "to_v"):
                    out[f"{base}.{proj}.lora_A.weight"] = v.contiguous()
            else:  # lora_B: [3H, r] -> three [H, r] row blocks
                third = v.shape[0] // 3
                for i, proj in enumerate(("to_q", "to_k", "to_v")):
                    out[f"{base}.{proj}.lora_B.weight"] = v[i * third:(i + 1) * third].contiguous()
            continue

        m = o_re.match(key)
        if m:
            out[f"{m.group(1)}.to_out.0.{m.group(2)}.weight"] = v.contiguous()
            continue

        out[key] = v
    return out


def _load_realism_lora():
    global LORA_HOW, _LORA_TARGETS
    if hasattr(pipe, "load_lora_weights"):
        pipe.load_lora_weights(
            LORA_REPO, weight_name=LORA_WEIGHT, adapter_name=LORA_ADAPTER, token=HF_TOKEN
        )
        LORA_HOW = "pipeline"
        return

    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    path = hf_hub_download(LORA_REPO, LORA_WEIGHT, token=HF_TOKEN)
    state_dict = _convert_ai_toolkit_lora(load_file(path))

    # Ideogram 4 uses dual-branch CFG: separate `transformer` (conditional) and
    # `unconditional_transformer`. The original ComfyUI workflow applied the LoRA
    # to BOTH — and since the model's post-training "safety" refusal lives in the
    # weights, applying the Realism Engine LoRA fully (attention included) to both
    # branches is what actually suppresses the "Image blocked by safety filter".
    targets = [pipe.transformer]
    extra = getattr(pipe, "unconditional_transformer", None)
    if extra is not None:
        targets.append(extra)
    for t in targets:
        t.load_lora_adapter(dict(state_dict), adapter_name=LORA_ADAPTER, prefix=None)
    _LORA_TARGETS = targets
    LORA_HOW = "transformer"


try:
    _load_realism_lora()
    LORA_OK = True
    print(f"[startup] Loaded Realism Engine LoRA via {LORA_HOW} onto {max(len(_LORA_TARGETS), 1)} module(s)")
except Exception as exc:  # noqa: BLE001 - we want any failure to be non-fatal
    print(f"[startup] WARNING: could not load LoRA, using base model only: {exc!r}")


def _set_lora_scale(strength: float) -> None:
    if not LORA_OK:
        return
    try:
        if LORA_HOW == "pipeline":
            pipe.set_adapters([LORA_ADAPTER], adapter_weights=[float(strength)])
        else:
            for t in _LORA_TARGETS:
                t.set_adapters([LORA_ADAPTER], weights=[float(strength)])
    except Exception as exc:  # noqa: BLE001
        print(f"[generate] WARNING: could not set LoRA strength: {exc!r}")


# Optional "fast mode". Ideogram 4 runs BOTH a conditional and an unconditional
# 9.3B transformer every step, blended as v = gw*v_pos + (1-gw)*v_neg. With
# guidance_scale forced to 1.0, the unconditional term is weighted by 0, so we can
# skip that entire 9.3B pass for ~2x speed — at the cost of CFG/prompt adherence.
# The wrapper is installed once and only short-circuits when _FAST is set per call.
_FAST = False
_uncond = getattr(pipe, "unconditional_transformer", None)
if _uncond is not None:
    _orig_uncond_forward = _uncond.forward

    def _uncond_forward(*args, **kwargs):
        if _FAST:
            # Multiplied by (1 - 1.0) = 0 downstream, so a zero scalar is exact.
            return (torch.zeros(1, device=_uncond.device, dtype=_uncond.dtype),)
        return _orig_uncond_forward(*args, **kwargs)

    _uncond.forward = _uncond_forward


# Figure out which keyword args this diffusers build's __call__ actually accepts,
# so we never crash by passing an unsupported one (e.g. negative_prompt / mu / std).
def _call_param_names() -> set:
    try:
        params = inspect.signature(pipe.__call__).parameters
    except (TypeError, ValueError):
        return set()
    if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()):
        return set()  # accepts **kwargs -> don't filter
    return set(params)


_CALL_PARAMS = _call_param_names()


def _filter_kwargs(kwargs: dict) -> dict:
    if not _CALL_PARAMS:
        return kwargs
    return {k: v for k, v in kwargs.items() if k in _CALL_PARAMS}


def _round16(value: int, lo: int = 512, hi: int = 1536) -> int:
    value = int(max(lo, min(int(value), hi)))
    return max(lo, (value // 16) * 16)


@spaces.GPU(duration=70)
def generate_image(
    prompt: str,
    negative_prompt: str = "",
    width: int = 1024,
    height: int = 1024,
    steps: int = 12,
    guidance: float = 7.0,
    mu: float = 0.5,
    std: float = 1.75,
    lora_strength: float = 0.9,
    seed: Optional[int] = -1,
    fast_mode: bool = False,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Prompt is required.")

    start_time = time.time()

    logger.info("=" * 80)
    logger.info("START IMAGE GENERATION")
    logger.info("Prompt: %s", prompt)
    logger.info("Negative Prompt: %s", negative_prompt)
    print(prompt)

    global _FAST
    _FAST = bool(fast_mode) and _uncond is not None
    if _FAST:
        guidance = 1.0

    width = _round16(width)
    height = _round16(height)
    steps = int(max(4, min(int(steps), 60)))

    if seed is None or int(seed) < 0:
        seed = random.randint(0, MAX_SEED)

    seed = int(seed)

    logger.info(
        "Params | width=%s height=%s steps=%s guidance=%s mu=%s std=%s lora=%s seed=%s fast=%s",
        width,
        height,
        steps,
        guidance,
        mu,
        std,
        lora_strength,
        seed,
        fast_mode,
    )

    if torch.cuda.is_available():
        logger.info(
            "GPU BEFORE | allocated=%.2f GB reserved=%.2f GB",
            torch.cuda.memory_allocated() / 1024**3,
            torch.cuda.memory_reserved() / 1024**3,
        )

    _set_lora_scale(lora_strength)

    generator = torch.Generator("cuda").manual_seed(seed)

    call_kwargs = _filter_kwargs(
        {
            "negative_prompt": negative_prompt.strip() or None,
            "height": height,
            "width": width,
            "num_inference_steps": steps,
            "guidance_scale": float(guidance),
            "guidance_schedule": None,
            "mu": float(mu),
            "std": float(std),
            "generator": generator,
        }
    )

    logger.info("PIPELINE KWARGS:")
    for k, v in call_kwargs.items():
        if k == "generator":
            logger.info("  %s=<torch.Generator>", k)
        else:
            logger.info("  %s=%s", k, v)

    try:
        logger.info("Calling pipe()...")

        result = pipe(prompt, **call_kwargs)

        logger.info("pipe() completed successfully")
        logger.info("Result type: %s", type(result))

        if hasattr(result, "images"):
            logger.info("Images returned: %s", len(result.images))
        else:
            logger.warning("Result has no images attribute")

        image = result.images[0]

        elapsed = round(time.time() - start_time, 2)

        if torch.cuda.is_available():
            logger.info(
                "GPU AFTER | allocated=%.2f GB reserved=%.2f GB",
                torch.cuda.memory_allocated() / 1024**3,
                torch.cuda.memory_reserved() / 1024**3,
            )

        logger.info("SUCCESS in %.2f seconds", elapsed)
        logger.info("=" * 80)

        return image, seed

    except Exception as e:
        elapsed = round(time.time() - start_time, 2)

        logger.error("=" * 80)
        logger.error("IMAGE GENERATION FAILED")
        logger.error("Elapsed: %.2f seconds", elapsed)
        logger.error("Exception Type: %s", type(e).__name__)
        logger.error("Exception Message: %s", str(e))

        err = str(e).lower()

        if "blocked" in err:
            logger.error("CONTENT FILTER DETECTED")

        if "safety" in err:
            logger.error("SAFETY FILTER DETECTED")

        if "policy" in err:
            logger.error("POLICY FILTER DETECTED")

        if "moderation" in err:
            logger.error("MODERATION FILTER DETECTED")

        if "cuda out of memory" in err:
            logger.error("CUDA OOM DETECTED")

        logger.error("PROMPT:")
        logger.error(prompt)

        logger.error("NEGATIVE PROMPT:")
        logger.error(negative_prompt)

        logger.error("FULL TRACEBACK:")
        logger.error(traceback.format_exc())

        if torch.cuda.is_available():
            logger.error(
                "GPU FAILURE | allocated=%.2f GB reserved=%.2f GB",
                torch.cuda.memory_allocated() / 1024**3,
                torch.cuda.memory_reserved() / 1024**3,
            )

        logger.error("=" * 80)

        raise

def _apply_preset(name: str):
    p = PRESETS.get(name, PRESETS["Default"])
    return p["steps"], p["mu"], p["std"]

DEFAULT_CAPTION =    {"high_level_description":"A candid photograph of Super Mario and Princess Zelda laughing together over coffee at a sunlit outdoor cafe, with Link watching them from another table in the background with a frustrated expression.","compositional_deconstruction":{"background":"A bright, airy outdoor Parisian-style cafe terrace with a white wrought-iron fence and a blurred street scene in the distance. Natural diffused daylight creates soft shadows on a light grey stone pavement.","elements":[{"type":"obj","bbox":[350,150,850,450],"desc":"Super Mario (Nintendo character), sitting at a small round cafe table. He wears his signature red cap, blue overalls with yellow buttons, and a red long-sleeved shirt. He is leaning back in a white metal chair, mouth open in a hearty laugh, holding a white ceramic espresso cup."},{"type":"obj","bbox":[350,450,850,750],"desc":"Princess Zelda (Nintendo character), sitting opposite Mario. She wears her royal pink and white gown with gold embroidery and a small gold crown atop her long blonde hair. She is laughing with her hand partially covering her mouth, looking at Mario."},{"type":"obj","bbox":[550,300,650,500],"desc":"A small round white marble cafe table between Mario and Zelda, holding two white ceramic coffee cups on matching saucers and a small silver sugar bowl."},{"type":"obj","bbox":[400,750,700,900],"desc":"Link (Nintendo character), sitting at a separate small table in the mid-ground. He wears his green tunic and pointed green cap. He is leaning forward with his chin resting on one hand, eyes narrowed and brow furrowed in a visible angry scowl, staring toward Mario and Zelda."},{"type":"obj","bbox":[600,800,650,850],"desc":"A single white coffee cup sitting untouched on Link's table."},{"type":"text","bbox":[200,600,300,800],"text":"CAFÉ\nCÉLESTE","desc":"A black wrought-iron hanging sign above the cafe entrance, featuring elegant white serif typography."}]}}


def dumps_caption(caption):
    return json.dumps(caption, ensure_ascii=False, separators=(",", ":"), indent=2)


def normalize_caption(raw_caption):
    try:
        caption = json.loads(raw_caption, strict=False)
    except Exception as e:
        raise gr.Error(f"JSON parse error: {e}") from e

    if not isinstance(caption, dict):
        raise gr.Error("Top-level JSON must be an object.")

    if "compositional_deconstruction" not in caption:
        gr.Warning("compositional_deconstruction is missing. The model accepts any string, but this is outside the usual Ideogram 4 caption format.")

    return json.dumps(caption, ensure_ascii=False, separators=(",", ":")), caption


def build_preset(mode, main_guidance=DEFAULT_MAIN_GUIDANCE, final_guidance=DEFAULT_FINAL_GUIDANCE):
    preset = dict(MODES.get(mode, MODES["Default · 20 steps"]))
    steps = int(preset.pop("num_inference_steps"))
    final_steps = min(int(preset.pop("final_guidance_steps")), steps)
    main_steps = steps - final_steps
    guidance_schedule = (float(main_guidance),) * main_steps + (float(final_guidance),) * final_steps
    preset.update(num_inference_steps=steps, guidance_schedule=guidance_schedule)
    return preset

with gr.Blocks(title=APP_TITLE) as demo:
    lora_note = (
        "Realism Engine LoRA: **active**."
        if LORA_OK
        else "Realism Engine LoRA: **not loaded** (running base Ideogram 4)."
    )
    
    with gr.Row():
        with gr.Column():
            caption = gr.Textbox(label="JSON caption", value=dumps_caption(DEFAULT_CAPTION), lines=28)
            prompt = gr.Textbox(
                label="Prompt",
                lines=6,
                
                placeholder="Describe the image, or paste an Ideogram structured JSON caption...",
            )
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                lines=2,
                value="low quality, blurry, distorted, bad anatomy",
            )

            preset = gr.Dropdown(
                choices=list(PRESETS.keys()),
                value="Turbo",
                label="Quality preset (sets steps / mu / std)",
            )

            fast_mode = gr.Checkbox(
                value=False,
                label="⚡ Fast mode — skip the negative pass (~2× faster, lower prompt adherence)",
            )

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

            with gr.Row():
                steps = gr.Slider(4, 60, value=12, step=1, label="Steps")
                guidance = gr.Slider(1.0, 12.0, value=7.0, step=0.1, label="Guidance")

            with gr.Row():
                mu = gr.Slider(-1.0, 1.0, value=0.5, step=0.05, label="mu (schedule shift)")
                std = gr.Slider(0.5, 3.0, value=1.75, step=0.05, label="std (schedule spread)")

            lora_strength = gr.Slider(
                0.0, 1.0, value=0.9, step=0.05,
                label="LoRA strength (sweet spot 0.5-0.9)",
                interactive=LORA_OK,
            )
            seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
            btn = gr.Button("Generate", variant="primary")

        with gr.Column():
            output = gr.Image(label="Generated image", type="pil")
            used_seed = gr.Number(label="Seed used", interactive=False)
    gr.Examples(
        examples=[
            [dumps_caption(DEFAULT_CAPTION)],
            [
                dumps_caption(
                    {
                        "high_level_description": "A square package label for a fictional tea brand called BLUE HARBOR.",
                        "style_description": {
                            "aesthetics": "premium, calm, balanced, Japanese-inspired packaging design",
                            "lighting": "even studio light",
                            "medium": "graphic_design",
                            "art_style": "flat vector label design with refined serif typography",
                            "color_palette": ["#F8FAFC", "#0F172A", "#2563EB", "#94A3B8", "#EAB308"],
                        },
                        "compositional_deconstruction": {
                            "background": "A clean ivory square label with a thin navy border.",
                            "elements": [
                                {
                                    "type": "text",
                                    "bbox": [170, 180, 300, 820],
                                    "text": "BLUE HARBOR",
                                    "desc": "Elegant navy serif uppercase brand name centered at the top.",
                                    "color_palette": ["#0F172A"],
                                },
                                {
                                    "type": "obj",
                                    "bbox": [360, 320, 650, 680],
                                    "desc": "A simple blue line illustration of ocean waves inside a gold circular seal.",
                                    "color_palette": ["#2563EB", "#EAB308"],
                                },
                                {
                                    "type": "text",
                                    "bbox": [720, 250, 810, 750],
                                    "text": "EARL GREY",
                                    "desc": "Small spaced navy sans-serif product text centered near the bottom.",
                                    "color_palette": ["#0F172A"],
                                },
                            ],
                        },
                    }
                )
            ],
        ],
        inputs=[caption],
    )
    preset.change(_apply_preset, inputs=preset, outputs=[steps, mu, std])

    btn.click(
        fn=generate_image,
        inputs=[
            prompt,
            negative_prompt,
            width,
            height,
            steps,
            guidance,
            mu,
            std,
            lora_strength,
            seed,
            fast_mode,
        ],
        outputs=[output, used_seed],
        api_name="generate",
    )

demo.queue(max_size=20)

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