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

…r and Qwen text encoder - Implements `Ideogram4Transformer` using `flax.nnx` - Ports `TorchaxQwen3VLTextEncoder` to handle Ideogram's LLM prompt conditions - Adds strict JSON schema caption key re-ordering in `IdeogramPipeline` - Adds full asymmetric CFG denoising loop

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

eltsai · 2026-07-21T17:52:00Z

@@ -0,0 +1,138 @@
+import jax


this file and some other files are missing the license header, this would incur errors when doing copy-bara

eltsai · 2026-07-21T17:53:53Z

@@ -0,0 +1,138 @@
+import jax
+


also this file is named for torchax but there's no torchax here (I see device=torch.device("cpu"))

eltsai · 2026-07-21T17:54:59Z

+      self.bias = None
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    w = self.weight.to(x.dtype) * self.weight_scale.to(x.dtype).unsqueeze(1)


this line dequantizes the full weight tensor on every forward. We should dequantize once at load and cache

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

eltsai · 2026-07-21T17:59:42Z

+
+    # JAX scaled dot product attention
+    scale = 1.0 / math.sqrt(self.head_dim)
+    attn_weights = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale


in base_ideogram.yml we have attention: "flash" but here it uses a naive attention, can be optimized later

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.

eltsai · 2026-07-21T18:01:18Z

+    schedule_fn = get_schedule_for_resolution((height, width), known_mean=0.5)
+    step_intervals = make_step_intervals(num_steps)
+    # Compute schedule evaluated at all intervals. schedule_fn returns a continuous scalar mapping.
+    sigmas = jnp.array(


this can be vectorized

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

mbohlool force-pushed the ideogram_jax branch 8 times, most recently from 83fcf30 to 23fa1c2 Compare July 8, 2026 23:27

mbohlool force-pushed the ideogram_jax branch 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

eltsai reviewed Jul 21, 2026

View reviewed changes


		class Ideogram4Attention(nnx.Module):

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

Uh oh!

Conversation

mbohlool commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Jul 8, 2026

Uh oh!

mbohlool commented Jul 8, 2026

Uh oh!

github-actions Bot commented Jul 9, 2026

Uh oh!

github-actions Bot commented Jul 9, 2026

Uh oh!

mbohlool commented Jul 9, 2026

Uh oh!

github-actions Bot commented Jul 9, 2026

Uh oh!

github-actions Bot commented Jul 9, 2026

Uh oh!

mbohlool commented Jul 20, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mbohlool commented Jul 8, 2026 •

edited

Loading