"""Generate an image with the fused INT8 Ideogram 4.0 DiT on a single RTX 3090.

  python download_deps.py                       # one time (gated base-repo access required)
  python usage.py "a poster that says HELLO"

Needs: an Ampere GPU with INT8 tensor cores (RTX 3090), `triton`, the `ideogram4` package,
and access to the gated base repo `ideogram-ai/ideogram-4-fp8` plus the INT8 W8A8 weights
at `transformerlab/ideogram-4-int8-w8a8` (this repo carries the fused kernel, not the
weights). On a 24 GB card the INT8 build fits at 1024px.
"""
import sys
import torch
from huggingface_hub import hf_hub_download
from ideogram4 import Ideogram4Pipeline, Ideogram4PipelineConfig
from fused_int8 import load_fused_int8

prompt = sys.argv[1] if len(sys.argv) > 1 else 'a storefront sign that says "FRESH COFFEE"'

# 1) base pipeline (text encoder + VAE + the two FP8 DiT branches, custom inference code)
pipe = Ideogram4Pipeline.from_pretrained(
    config=Ideogram4PipelineConfig(weights_repo="ideogram-ai/ideogram-4-fp8"),
    device="cuda", dtype=torch.bfloat16)

# 2) the INT8 W8A8 weights live in our int8-w8a8 repo; this repo provides the fused kernel
weights = hf_hub_download("transformerlab/ideogram-4-int8-w8a8",
                          "ideogram4-int8-w8a8.safetensors")
n_fused, n_prot = load_fused_int8(pipe, weights)
print(f"installed fused INT8 kernel on {n_fused} linears (+{n_prot} bf16-protected)")

# 3) generate (single 24 GB RTX 3090; first denoise step autotunes the kernel)
img = pipe(prompt, num_steps=48, height=1024, width=1024, seed=1000)[0]
img.save("out.png")
print("saved out.png")