Memory profile of text_to_image prevents consumer-GPU inference; also image_resolution 512/256 appear broken
Thanks for releasing this — the masked-discrete-diffusion approach is genuinely interesting and the model runs. I've been getting it working on an RTX 3070 (8GB) in ComfyUI via a custom GGUF quantization of the dLM, and hit a few things worth reporting.
- image_resolution=512 and 256 seem broken as shipped.
text_to_image defaults to image_resolution=512, n_tokens=1024. At that setting, _t2i_wte raises a shape mismatch: the placeholder mask selects 1024 positions per batch item, but call_gen_embedding returns 256 embeddings after downsample_gen (2×2, 32×32 → 16×16). The 4:1 ratio is invariant to n_tokens.
Only image_resolution=1024, n_tokens=4096 is self-consistent (64×64 → 32×32 → matches the reserved placeholder count). Is 1024 the only supported resolution, or is n_tokens_txt supposed to be n_tokens // 4?
- The sampling tail is the memory wall, not the transformer.
The 34-layer body streams fine in 4-bit. What doesn't fit is the tail:
gen_predictor outputs [batch, 4096, 131072] — 1GB in bf16 per CFG branch2GB)
probs = logits.softmax(dim=-1) (line 792) allocates another full-vocab tensor
dists.Categorical(logits / temperature).sample() → torch.multinomial, which upcasts probabilities to fp32 (
the CFG combine (1+gs)logits - gslogits_un materializes several more full-size temporaries
Peak is 5+ vocab-sized tensors simultaneously. On 8GB this OOMs even with all weights offloaded to system RAM.
Workarounds that made it fit (all mathematically equivalent, happy to share):
Gumbel-max sampling — argmax(logits + gumbel) instead of multinomial, no fp32 probs
x0_p via exp(logits[x0] − logsumexp(logits, -1)), chunked — avoids the full softmax while keeping confidence_policy='mmada' correct
in-place CFG (mul_/sub_)
sequential rather than batched cond/uncond forwards
Would you consider upstreaming something like this? A low_memory=True path would make the model accessible to a much larger audience, and none of it changes the output distribution.
- Step time grows monotonically.
Per-step time climbs from ~20s to ~130s over 20 steps with use_cache/cache_prompt at their defaults. Since this is bidirectional masked diffusion rather than autoregressive decoding, is the KV cache across diffusion steps intentional? Setting both to False didn't fully flatten it, so I may be missing where else state accumulates.
- Quantization sensitivity (observation, not a bug report).
Quantizing gen_predictor / diffusion_head degrades output very differently from quantizing the transformer body. Since those are 131k-way classifiers over the codebook, weight error flips the argmax to a different token, and the VQ decoder renders an entirely different patch — so errors appear as structural corruption (malformed keyboard keys, distorted anatomy) rather than the softening you'd see in a latent diffusion model. Keeping the classifier at Q8 while running the body at Q4 helps.
Is there a recommended minimum precision for those tensors? And are there reference settings beyond the defaults (n_steps=20, guidance_scale=5.0) for text rendering specifically? Text glyphs are consistently well-formed but incorrect in my runs, which may be expected for the base model — I'd like to know whether I'm leaving quality on the table or seeing it as intended.
Happy to open PRs for any of the memory items if useful.
Example outputs on Q4 gguf:
The sign should read "NL-Diffusion-Image" and the clouds should read "nvidia"
type writer paper should read "Hello World"