Lens on a single RTX 4090: quantizing the DiT to run 1440 without offload

Lens is a 3.8B-parameter foundational text-to-image model from Microsoft, built for training efficiency and high-resolution generation. Architecturally it is unusual: instead of a CLIP/T5 text encoder, it conditions a 48-block double-stream MMDiT on multi-layer hidden states pulled from a GPT-OSS-20B language model (layers 5/11/17/23, concatenated), and it decodes latents through the FLUX.2 VAE. Three checkpoints ship: the RL-tuned microsoft/Lens (20 steps, cfg 5.0), a distilled 4-step Lens-Turbo, and the supervised Lens-Base.

The interesting consequence for a local run is the memory shape. That GPT-OSS-20B encoder is enormous for a text encoder — but it is shipped in MXFP4 (4-bit), so it lands at roughly 11 GB instead of ~40 GB. The denoiser and VAE are comparatively small. So the question on a 24 GB card is never "does it load" — it is "how much room is left for the diffusion activations once that 4-bit encoder is sitting on the GPU."

The reference path just works. Load the MXFP4 encoder, assemble the pipeline, move it to CUDA, and generate. On a 4090 a 1024² image at 20 steps takes ~9 s once the model is warm. There are two non-obvious setup details worth flagging up front, because both cost time before the first image appears:

• A pre-release dependency. diffusers 0.38 pins a pre-release safetensors (0.8.0rc), so the install needs uv pip install --prerelease=allow — a plain install fails to resolve.
• MXFP4 on Ada. MXFP4 matmul kernels officially target Hopper (sm_90); the 4090 is Ada (sm_89). In practice the transformers + kernels path still runs the 4-bit encoder fine here (you get a harmless kernels deprecation warning). The naive "dequantize to bf16" escape hatch is a trap — a dequantized 20B encoder is ~40 GB and will not fit in 24 GB at all.

# install (note the pre-release flag)
uv pip install torch==2.11.0+cu126 torchvision==0.26.0+cu126 \
    --index-url https://download.pytorch.org/whl/cu126
uv pip install --prerelease=allow -r requirements.txt

# generate (with CPU offload — the out-of-the-box path, any resolution)
python inference.py --prompt "a cinematic mountain lake at sunrise" \
    --base_resolution 1440 --aspect_ratio 16:9 --steps 20 --cfg 5.0 --offload

CPU offload works, but it is not free — it shuttles whole modules between CPU and GPU, and I wanted the encoder and denoiser to simply stay resident. Dropping --offload and calling pipe.to("cuda") exposes the real budget. The 4-bit encoder (~11 GB) plus the bf16 DiT (~7.6 GB) is ~19 GB before any image, leaving ~5 GB for activations. That is enough for 1024 (it peaks right at the 24 GB ceiling) but not for 1440: every 1440 image hit a hard CUDA out of memory, the bf16 path needing ~27 GB.

Before the fix that worked, two reasonable-sounding ideas failed, and they are worth recording because both look correct on paper:

Both dead-ends point at the same conclusion: you cannot claw the 11 GB encoder back at runtime. If 1440 is going to fit, the savings have to come from the only component still at full precision — the DiT.

The text encoder is already 4-bit; there is nothing left to take there. The DiT, however, was running in bf16 (~7.6 GB). Quantizing it weight-only to FP8 via torchao (diffusers exposes this through TorchAoConfig) roughly halves it to ~3.8 GB. The 4090 has FP8 tensor cores, and weight-only FP8 uses native PyTorch ops — so it runs even though two of torchao's Hopper-only prebuilt kernels fail to load (they aren't needed here).

That ~3.8 GB saving is far more than the ~1–2 GB by which 1440 was overshooting. The 1440 peak drops from ~27 GB (OOM) to ~23 GB, and Lens generates 1440 fully on the GPU, no offload. I wired it into a --quant_transformer flag; the model's own attention already uses memory-efficient SDPA, so the only change needed was the weight precision.

uv pip install torchao

# 1440, no offload, on a 24 GB 4090 — FP8-quantized DiT + MXFP4 encoder
python inference.py --prompt "..." \
    --base_resolution 1440 --aspect_ratio 16:9 --steps 30 --cfg 4.5 \
    --quant_transformer float8_weight_only        # NOTE: no --offload

Every image below was generated no offload at 1024 (the peak-VRAM edge case), one per category, to see where Lens is strong and where it isn't. Each caption is the exact configuration — base resolution, aspect ratio, steps, cfg and seed — followed by the prompt. Same prompt + seed + config is reproducible.

With the FP8 DiT, the three categories that benefit most from resolution run at 1440 with the encoder and denoiser both resident. Captions show the measured peak VRAM — all comfortably under the 24 GB ceiling that bf16 blew past. The extra resolution shows: the portrait in particular gains real micro-detail in the skin and beard versus its 1024 version above.

The alpine landscape in the gallery came out painterly, and the culprit was the prompt. "National Geographic style" reads to the model as an aesthetic, not a camera. Swapping the art language for concrete photographic cues — a body, a lens, an aperture, "RAW photo", "high dynamic range", "no haze", "extreme detail" — and nudging cfg down to 4.5 pushes Lens firmly into photoreal territory.

Lens is a genuinely strong text-to-image model that punches above its 3.8B size, and the GPT-OSS encoder clearly pays off where it matters most: text rendering, labelled diagrams, and prompt following are all better than I expected from a model this size. Portraits, illustration and comic styling are reliably good. Where it stumbles is the usual T2I weak spot — exact object counts ("two moons" became three) — and, like any model, it does what the prompt literally says, so vague aesthetic words can quietly steer you away from photorealism.

On hardware, the headline is the memory story. The 4-bit encoder is a clever way to ship a 20B conditioner, but it also fixes ~11 GB of your budget in place. On a 24 GB 4090 that makes 1024 the comfortable no-offload ceiling out of the box — and a one-line FP8 quantization of the DiT is what turns "1440 needs offload" into "1440 runs fully on the GPU," with no quality I could see. That is the most reusable takeaway here: when the encoder is already quantized, quantize the denoiser.