Ideogram4 Optimizations KJ(Ideogram4OptimizationsKJ)

This node applies two experimental VRAM-reduction optimizations to the Ideogram4 model forward pass. It is part of the "KJNodes/experimental" category and may alter model outputs — use only when peak VRAM is a limiting factor and after verifying output quality.

Description

The Ideogram4 Optimizations KJ node patches the Ideogram4 model to reduce peak GPU memory during inference by targeting the two largest transient tensors in each transformer block: the SwiGLU activations and the RoPE computation. Both optimizations are controlled by discrete boolean and integer inputs, allowing fine-grained control over memory savings vs. computational overhead.

Inputs

• model (MODEL, required) – The Ideogram4 base model to which the optimizations will be applied.

• chunk_ffn (BOOLEAN, default True) – When enabled, the feedforward (SwiGLU) activations are split over the token sequence dimension, capping the intermediate tensor size to (B, chunk_size, hidden) instead of (B, L, hidden). This reduces peak memory at the cost of a small increase in computation.

• ffn_chunks (INT, default 2, range 1–64, step 1) – The number of sequence chunks to split the feedforward into. Higher values yield lower peak memory but increase overhead (more kernel launches). A value of 1 effectively disables chunking. Start with 2 and increase if VRAM is still insufficient.

• ffn_seq_threshold (INT, default 1024, range 256–65536, step 256) – Chunking is only applied when the input token sequence length exceeds this value. For short sequences (e.g., below 1024 tokens), the baseline memory is already small enough that chunking would add unnecessary overhead. Set this threshold based on your typical generation length.

• bf16_rope (BOOLEAN, default True) – When enabled, Rotary Position Embedding (RoPE) is applied in the model’s working dtype (typically bfloat16 or float16) instead of being upcasted to float32. This roughly halves the RoPE activation memory and matches the precision used in Hugging Face reference implementations. Output quality may differ slightly from the default fp32 RoPE path.

Outputs

• MODEL – The same model with the selected optimizations patched in. It is meant to be used as a drop-in replacement for the original unpatched model in a ComfyUI workflow.

Usage Notes

• Both optimizations are experimental. Always test output quality and compare with the unpatched model before deploying in a production workflow.
• The chunk_ffn and bf16_rope optimizations are independent; enable only what you need. For maximum memory savings, enable both.
• When using chunk_ffn, start with the default ffn_chunks = 2 and only increase if VRAM overflow persists, as more chunks introduce greater overhead.
• The ffn_seq_threshold is most useful when your workflow processes both very short and very long sequences (e.g., image conditioning vs. full image generation). Tune it to avoid overhead on short prompts.
• These optimizations are designed specifically for Ideogram4; they will have no effect on other model architectures and may cause errors if applied. Ensure the input model is an Ideogram4 variant.