Port Ideogram 4.0 to JAX / MaxDiffusion by mbohlool · Pull Request #442 · AI-Hypercomputer/maxdiffusion

mbohlool · 2026-07-08T21:37:10Z

Title: Port Ideogram 4.0 to JAX / MaxDiffusion

Summary:
This PR ports the Ideogram 4.0 pipeline to JAX. It introduces the native Flax NNX Ideogram4Transformer, integrates the TorchaxQwen3VLTextEncoder for complex prompt processing, and implements the required asymmetric CFG denoising loop. It also includes caption key re-ordering in the pipeline to enforce Ideogram's strict JSON-based prompt schema (high_level_description, style_description, compositional_deconstruction).

Generated Examples:

Prompt:

{
  "high_level_description": "A sleek red sports car driving through a futuristic cyberpunk city in the rain, with neon reflections on the wet asphalt.",
  "style_description": {
    "aesthetics": "cyberpunk, futuristic, moody, dynamic",
    "lighting": "neon lights, cinematic lighting, reflections",
    "photo": "photorealistic, motion blur on background",
    "medium": "digital photography",
    "color_palette": ["#E74C3C", "#3498DB", "#9B59B6", "#1ABC9C", "#2C3E50"]
  },
  "compositional_deconstruction": {
    "background": "A bustling cyberpunk cityscape at night, towering skyscrapers, heavy rain.",
    "elements": [
      {
        "type": "obj",
        "bbox": [400, 100, 900, 900],
        "desc": "A highly detailed, sleek red futuristic sports car speeding forward."
      },
      {
        "type": "obj",
        "bbox": [800, 0, 1000, 1000],
        "desc": "Wet asphalt road reflecting vibrant pink and blue neon lights."
      }
    ]
  }
}

Prompt:

{
  "high_level_description": "A ginger cat wearing a tiny wizard hat reading a glowing spellbook.",
  "style_description": {
    "aesthetics": "magical, fantasy, whimsical",
    "lighting": "warm glowing light, magical illumination from the book",
    "photo": "highly detailed, macro photography, shallow depth of field",
    "medium": "digital photography, fantasy illustration",
    "color_palette": ["#E67E22", "#8E44AD", "#F1C40F", "#2C3E50", "#D35400"]
  },
  "compositional_deconstruction": {
    "background": "A dimly lit, dusty wizard's study filled with old scrolls and potions.",
    "elements": [
      {
        "type": "obj",
        "bbox": [100, 200, 800, 800],
        "desc": "A fluffy ginger cat wearing a detailed purple wizard hat with gold stars."
      },
      {
        "type": "obj",
        "bbox": [700, 200, 1000, 800],
        "desc": "An ancient, open spellbook glowing with magical yellow light."
      }
    ]
  }
}

Prompt:

{
  "high_level_description": "A stunning cinematic shot of a futuristic city at night with flying cars and bright neon lights.",
  "style_description": {
    "aesthetics": "cinematic, futuristic, sci-fi",
    "lighting": "neon lighting, night time, glowing cityscapes",
    "photo": "photorealistic, hyper-detailed, wide angle",
    "medium": "digital photography",
    "color_palette": ["#00FFFF", "#FF00FF", "#000080", "#FFD700", "#191970"]
  },
  "compositional_deconstruction": {
    "background": "A dark night sky overlooking a sprawling futuristic megalopolis.",
    "elements": [
      {
        "type": "obj",
        "bbox": [200, 0, 1000, 1000],
        "desc": "Dense cityscape with towering skyscrapers illuminated by intricate neon signs."
      },
      {
        "type": "obj",
        "bbox": [100, 200, 400, 800],
        "desc": "Several futuristic flying cars with glowing thrusters streaking across the sky."
      }
    ]
  }
}

github-actions · 2026-07-08T21:37:38Z

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

mbohlool · 2026-07-08T23:35:20Z

/gemini review

github-actions · 2026-07-09T06:20:55Z

🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-07-09T06:26:38Z

🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details.

mbohlool · 2026-07-09T18:33:11Z

/gemini review

github-actions · 2026-07-09T18:33:43Z

🤖 Hi @mbohlool, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-07-09T18:36:30Z

🤖 I'm sorry @mbohlool, but I was unable to process your request. Please see the logs for more details.

mbohlool · 2026-07-20T17:00:50Z

/gemini review

eltsai · 2026-07-21T17:50:35Z

+  # 1. Warmup Compilation
+  config.get_keys()["enable_profiler"] = False
+  config.get_keys()["enable_ml_diagnostics"] = False
+  config.get_keys()["num_inference_steps"] = 2  # lower for warmup


since we have z, _, _ = jax.lax.fori_loop(0, num_steps, denoise_step, init_val) in src/maxdiffusion/pipelines/ideogram/ideogram_pipeline.py, setting the warmup step as 2 means that the later 50 step inference time includes compile

I think this is the reason why we see a 50 sec generation time. We could either:

go with what WAN does: only compile one step and do a plain for loop

warmup with 50 steps

I am not sure which would be faster, it's worth testing

Addressed: Updated the warmup compilation pass in generate_ideogram.py to use num_inference_steps = original_num_steps so XLA compiles the full 50-step loop during warmup, removing the compilation overhead from the timed inference pass. But the inference time is is ~50s. We can improve it in a follow up PR.

eltsai · 2026-07-21T17:58:15Z

+    q, k = _apply_rotary_pos_emb(q, k, cos, sin)
+
+    # Block-diagonal mask from segment ids
+    attn_mask = jnp.expand_dims(segment_ids, axis=2) == jnp.expand_dims(segment_ids, axis=1)


I am not quite sure why we need masks to do this...we can leave this for future optimization

Acknowledged: The segment mask handles variable-length sequence padding vs. active token boundaries. We'll leave specialized splash kernel optimizations for a follow-up PR.

eltsai · 2026-07-21T18:00:33Z

+
+class Ideogram4Attention(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs, hidden_size: int, num_heads: int, eps: float = 1e-5, dtype=jnp.float32):


dtype=jnp.float32 default is never overridden, so weights_dtype/activations_dtype: bfloat16 in the config are inert.

Addressed: Made dtype configurable in Ideogram4Config and Ideogram4Transformer, and ensured IdeogramPipeline passes config.weights_dtype (bfloat16) to the transformer instances.

- Implements native Ideogram4Transformer using flax.nnx with configurable dtype support and FSDP sharding rules - Implements IdeogramPipeline orchestrating conditional/unconditional transformers, Qwen3VLTextEncoder, and AutoEncoder with an asymmetric CFG denoising loop - Adds Qwen3VLTextEncoder for multi-layer intermediate hidden state extraction and FP8 quantized weight loading - Pre-dequantizes FP8 linear weights at load time to avoid per-step overhead - Implements Ideogram 4.0 resolution-dependent noise schedule and vectorizes sigma step calculations - Adds canonical JSON prompt schema parsing and key re-ordering for structured prompts - Adds robust offline cache loading for multi-component pipeline tokenizers - Adds generate_ideogram.py generation entrypoint with full-inference warmup compilation - Includes unit tests (generate_ideogram_smoke_test.py and generate_ideogram_end_to_end.py)

github-actions · 2026-08-25T23:26:22Z

🤖 Hi @eltsai, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-08-25T23:29:27Z

🤖 I'm sorry @eltsai, but I was unable to process your request. Please see the logs for more details.

github-actions · 2026-08-26T19:10:52Z

🤖 I'm sorry @eltsai, but I was unable to process your request. Please see the logs for more details.

mbohlool requested a review from entrpn as a code owner July 8, 2026 21:37

mbohlool force-pushed the ideogram_jax branch 9 times, most recently from 23fa1c2 to ab6318b Compare July 8, 2026 23:35

Perseus14 added the gemini-review label Jul 9, 2026

mbohlool added gemini-review and removed gemini-review labels Jul 9, 2026

eltsai self-requested a review July 14, 2026 21:50