Port Ideogram 4.0 to JAX / MaxDiffusion

[mbohlool]

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]

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

[mbohlool]

/gemini review

[github-actions]

🤖 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]

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