Qwen-Image-2.1 NF4 + Qwen3.5 Grounded Text Adapter

This repository contains a memory-efficient Qwen-Image-2.1 setup together with a trained Qwen3.5 → Qwen-Image-2.1 text-conditioning adapter.

The Qwen3.5 model is not a drop-in replacement for QwenImage21Pipeline.text_encoder. Qwen-Image-2.1 expects a Qwen3-VL text encoder interface, while this repository uses Qwen3.5 externally and converts its hidden representations into 4096-dimensional Qwen-Image conditioning embeddings.

Quick inference / Hugging Face compatibility

Current repository status: do not call this repository as a stock QwenImage21Pipeline and do not pass Qwen3.5 directly as text_encoder. The Qwen3.5 backbone emits 1024-dimensional hidden states, while Qwen-Image-2.1 consumes 4096-dimensional conditioning. The trained GroundedAdapter5 and token-alignment step must run between them.

The following seemingly standard call is not compatible with the current repository layout:

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "summerMC/qwen-image-2.1-nf4",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")
image = pipe("A cinematic Japanese city at night").images[0]

If the log contains:

Expected types for text_encoder: Qwen3VLForConditionalGeneration, got Qwen3_5ForCausalLM

or generation fails with:

The size of tensor a (1024) must match the size of tensor b (4096)

then Diffusers loaded the stock QwenImage21Pipeline; the grounded adapter was bypassed. trust_remote_code=True alone does not fix this unless the Hub repository actually contains and advertises a custom pipeline implementation.

Intended short API

After this repository is packaged as a Diffusers custom pipeline, the intended user-facing API is:

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "summerMC/qwen-image-2.1-nf4",
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="cuda",
)
image = pipe("A cinematic Japanese city at night").images[0]

For that API to be valid, the repository-side pipeline must internally perform:

Qwen3.5 tokenizer
      ↓
Qwen3.5 lexical + hidden[-9] + hidden[-1]
      ↓ 3072
character-overlap alignment
      ↓
GroundedAdapter5
      ↓ 4096
Qwen-Image-2.1 transformer
      ↓
VAE decode

Until that custom pipeline is present, use the explicit external-conditioning path documented below.

Important

Do not load Qwen3.5 like this:

QwenImage21Pipeline.from_pretrained(
    repo_id,
    text_encoder=qwen35,
)

That produces a type mismatch such as:

Expected Qwen3VLForConditionalGeneration, got Qwen3_5ForCausalLM

Instead, load the image pipeline with:

text_encoder=None

and provide prompt_embeds generated by the grounded adapter.

Conditioning architecture

The trained conditioning path is:

Qwen3.5 token embeddings        1024
Qwen3.5 hidden_states[-9]      1024
Qwen3.5 hidden_states[-1]      1024
                               ----
concatenated features          3072
          │
          ▼
character-overlap alignment
          │
          ▼
GroundedAdapter5
  ├─ whitening statistics
  ├─ 3072 → 4096 direct projection
  ├─ 3072 → 512 context projection
  ├─ 2-layer causal TransformerEncoder
  ├─ 512 → 4096 contextual residual
  └─ learned output RMS
          │
          ▼
Qwen-Image prompt embeddings   4096
          │
          ▼
QwenImage21Pipeline(text_encoder=None)

Adapter dimensions:

The adapter checkpoint contains 37 tensors and should be loaded with strict=True.

Repository files

The relevant files are:

config.json
text_encoder/
qwen35_tokenizer/
qwen35_adapter/
    adapter.safetensors
transformer/
vae/ or the VAE referenced by config.json
processor/

qwen35_adapter/adapter.safetensors is not a simple two-layer MLP checkpoint. Older examples defining a TextAdapter with net.0 / net.2 are incompatible with this checkpoint.

If you see an error similar to:

Missing key(s): net.0.weight, net.0.bias, net.2.weight, net.2.bias
Unexpected key(s): context.layers.*, in_proj.*, linear.weight, ...

you are using an obsolete adapter class.

Correct adapter implementation

import math
import torch
import torch.nn as nn


class GroundedAdapter5(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("x_mean", torch.zeros(3072))
        self.register_buffer("x_std", torch.ones(3072))
        self.register_buffer("y_mean", torch.zeros(4096))
        self.register_buffer("y_std", torch.ones(4096))

        self.linear = nn.Linear(3072, 4096, bias=False)
        self.lognorm = nn.Linear(3072, 1)

        self.in_proj = nn.Sequential(
            nn.Linear(3072, 512),
            nn.LayerNorm(512),
            nn.GELU(),
        )

        block = nn.TransformerEncoderLayer(
            d_model=512,
            nhead=8,
            dim_feedforward=1536,
            dropout=0.05,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )

        self.context = nn.TransformerEncoder(
            block,
            num_layers=2,
            enable_nested_tensor=False,
        )

        self.residual = nn.Linear(512, 4096)

    def forward(self, features, valid=None):
        if features.ndim == 2:
            features = features[None]

        x = (features.float() - self.x_mean) / self.x_std

        if valid is None:
            valid = torch.ones(
                x.shape[:2],
                device=x.device,
                dtype=torch.bool,
            )

        z = self.in_proj(x)
        length = z.shape[1]

        pos = torch.arange(length, device=z.device).float()[:, None]
        freq = torch.exp(
            torch.arange(0, 512, 2, device=z.device).float()
            * (-math.log(10000) / 512)
        )

        pe = torch.zeros(length, 512, device=z.device, dtype=z.dtype)
        pe[:, 0::2] = torch.sin(pos * freq)
        pe[:, 1::2] = torch.cos(pos * freq)

        causal = torch.triu(
            torch.ones(length, length, device=z.device, dtype=torch.bool),
            diagonal=1,
        )

        z = self.context(
            z + pe[None],
            mask=causal,
            src_key_padding_mask=~valid,
            is_causal=True,
        )

        raw = (self.linear(x) + self.residual(z)).float()
        raw = raw * self.y_std + self.y_mean

        rms = torch.exp(self.lognorm(x).float().clamp(-2, 5))
        raw_rms = (
            raw.square()
            .mean(-1, keepdim=True)
            .sqrt()
            .clamp_min(0.05)
        )

        return raw / raw_rms * rms

Load it with:

from safetensors.torch import load_file

adapter = GroundedAdapter5().cuda()
adapter.load_state_dict(
    load_file("qwen35_adapter/adapter.safetensors", device="cpu"),
    strict=True,
)
adapter.eval()

A correct checkpoint load should report no missing or unexpected keys.

Token alignment

The adapter was trained using full rendered character-overlap alignment. The prompt is first inserted into the configured chat template, then student and teacher token offsets are aligned by character overlap.

The current configuration uses:

alignment = full_rendered_character_overlap_v4
drop      = 14

Do not silently substitute a different tokenizer. In particular, the Qwen3.5 student tokenizer must produce token IDs compatible with the Qwen3.5 text encoder embedding table.

A useful sanity check is:

ids = student_tokenizer(
    template.format("a red cat"),
    add_special_tokens=False,
)["input_ids"]

assert len(ids) > 0
assert max(ids) < student_model.get_input_embeddings().num_embeddings

Feature extraction

The adapter input is constructed from three 1024-dimensional Qwen3.5 representations:

out = student.model(
    **inputs,
    output_hidden_states=True,
    use_cache=False,
)

lexical = student.get_input_embeddings()(inputs["input_ids"])

features = torch.cat(
    [
        lexical,
        out.hidden_states[-9],
        out.hidden_states[-1],
    ],
    dim=-1,
)

This gives a 3072-dimensional feature vector per student token. Character-overlap alignment then maps these features onto the teacher token sequence before the adapter is applied.

Pipeline loading

Qwen3.5 must remain external to Diffusers:

vae = AutoencoderKLQwenImage21.from_pretrained(
    config["teacher"],
    subfolder="vae",
    revision=config["teacher_revision"],
    dtype=torch.bfloat16,
)

vae.to("cuda")
vae.enable_tiling()

pipe = QwenImage21Pipeline.from_pretrained(
    config["student"],
    revision=config["student_revision"],
    text_encoder=None,
    vae=vae,
    dtype=torch.bfloat16,
    device_map="cuda",
)

The explicit CUDA placement of the VAE is important. If the VAE remains on CPU while latents are on CUDA, decoding fails with an error such as:

Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type) should be the same

Generation

After creating the external prompt embeddings:

prompt_embeds = adapter(aligned_features).to(torch.bfloat16)

prompt_embeds_mask = torch.ones(
    prompt_embeds.shape[:2],
    device="cuda",
    dtype=torch.long,
)

image = pipe(
    prompt_embeds=prompt_embeds,
    prompt_embeds_mask=prompt_embeds_mask,
    width=768,
    height=768,
    output_resolution=768,
    num_inference_steps=30,
    true_cfg_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(42),
    use_kv_cache=True,
    output_type="pil",
).images[0]

Validated invariants

The following structural checks have been validated against the current adapter checkpoint:

checkpoint tensors       37
missing keys             0
unexpected keys          0
input feature shape      [T, 3072]
adapter output shape     [1, T, 4096]
strict checkpoint load   PASS
finite forward output    PASS

These checks establish checkpoint/module compatibility. Image quality additionally depends on using the exact tokenizer, template, token alignment, student model revision, feature layers, normalization statistics, and generation configuration used by the trained adapter.

Troubleshooting

Expected types for text_encoder ... got Qwen3_5ForCausalLM

Qwen3.5 was passed directly as the Diffusers text encoder. Load the pipeline with text_encoder=None and pass externally generated prompt_embeds.

Missing key(s): net.0.weight ...

An obsolete MLP TextAdapter implementation is being used. Use GroundedAdapter5 above.

Student tokenizer returns zero tokens

The tokenizer files are incomplete or the wrong tokenizer loader/path is being used. Verify the Qwen3.5 tokenizer independently and ensure its IDs fit inside the student embedding vocabulary.

Purple/green noise or structurally meaningless images

If diffusion completes but the image is meaningless, inspect conditioning before changing the image model. Check:

student tokenizer identity
student token count
raw Qwen3.5 feature mean/std
aligned feature mean/std
adapter output mean/std
adapter token RMS
prompt embedding shape = [1, T, 4096]

A successful diffusion loop does not prove that the conditioning embeddings are semantically calibrated.

torchao shared-library warnings

Warnings about optional torchao CUTLASS/MXFP8 extensions are not, by themselves, evidence that this NF4/bitsandbytes inference path has failed. Diagnose them separately unless checkpoint loading explicitly requires a torchao tensor type.

Hardware

The tested path requires a CUDA GPU with BF16 support. NVIDIA L4 is suitable for the configuration described above.

Notes

• Prompt length is validated against a 224-token student envelope.
• No silent prompt truncation is performed.
• NF4 Linear4bit modules use BF16 compute where applicable.
• The VAE can use tiling to reduce peak memory use.
• strict=True is intentionally used for the adapter checkpoint.

Base model

Qwen-Image-2.1 is provided by Qwen. This repository contains a modified/quantized inference configuration and an experimental text-conditioning adapter. Refer to the upstream model repository for the base model's license, usage conditions, and model documentation.