recoilme/moss_te_adapter

MOSS-Image text encoder → Qwen3.5-0.8B + adapter

MOSS-Image conditions its DiT on Qwen3.5-VL-9B text states — 15.63 GiB in bf16, which is why the full pipeline needs ~25.4 GiB of VRAM. This adapter replaces that encoder with a frozen Qwen/Qwen3.5-0.8B plus a 158M regression head that reproduces the same (L, 4096) condition. DiT, VAE and scheduler are untouched.

	native	with adapter
text encoder	15.63 GiB	1.7 GB + 0.63 GB
peak VRAM, 1024 px / 30 steps	~25.4 GiB (does not fit 16 GB)	12.7 GiB

Metrics

Cosine and relative error against the native encoder, text positions, same prompts:

	cos	rel
5 fixed probes (prose / sign / tags / JSON / empty)	0.905	33.0%
128 held-out captions	0.924	33.4%

A/B with one seed, 1024 px, 30 steps — cosine of the adapter condition vs the native one: prose 0.931, sign 0.902, tags 0.940, JSON layout 0.756.

Usage

adapter.safetensors holds only the head. It is useless without the frozen student Qwen/Qwen3.5-0.8B; metadata in the file (student_layers, norm, hidden, attention, attn_max_len, …) describes the head completely.

pip install torch transformers diffusers safetensors pillow
python src/example.py --prompt "a red fox in a snowy forest at dusk, 85mm" \
    --moss-dir /path/to/MOSS-Image-Base --out fox.png

The pipeline loads without its text encoder and takes the condition from the adapter:

from adapter_lib import load_adapter
from moss_lib import Student, build_pipeline, condition, generate, load_tokenizer

tok = load_tokenizer(moss_dir)                                     # tokenizer of the teacher
adapter, meta = load_adapter("adapter.safetensors", device="cuda")
layers = [int(x) for x in str(meta["student_layers"]).split(",")]
student = Student("Qwen/Qwen3.5-0.8B", layers)

pos = condition(adapter, student, tok, [prompt])[0]                # (L, 4096) bf16 + mask
neg = condition(adapter, student, tok, [""])[0]                    # empty prompt = CFG negative
del student, adapter

pipe = build_pipeline(moss_dir).to("cuda")                         # no text_encoder component
image = generate(pipe, pos, neg, size=1024, steps=30, guidance=3.0, seed=1234)
image.save("fox.png")

Files

adapter.safetensors	the head, 158.4M params, bf16 weights
src/example.py	inference: 0.8B + adapter instead of the 9B encoder
src/moss_lib.py, src/adapter_lib.py	prompt template, models, head architecture, pipeline assembly
src/cache_targets.py, src/train_adapter.py	teacher targets → disk, then training on the cache
src/prompts.py, src/native_embeds.py, src/sample_ab.py	prompt pool, native condition, A/B grids
samples/	native encoder (left) vs adapter (right), one seed

Training

• Target: last_hidden_state of the teacher — post-norm output of layer 32, (L, 4096).
• Student input: hidden states [4, 8, 12, 16, 20, 24] of Qwen3.5-0.8B (after its full-attention layers 3, 7, 11, 15, 19, 23), per-slice LayerNorm, concatenated to in_dim 6144.
• Head: 6144 → 6144 → 6144 → 4096 MLP, plus 2 self-attention blocks over tokens and 2 blocks of attention across the slice axis (both residual, zero-init, added to the MLP output).
• Loss: per-token MSE weighted by 1/max(‖y‖², 128²); padding weight 0.1.
• Data: 14,507 captions, 2 epochs (1,798 steps), batch 16, lr 3e-4 cosine → 1e-5. The empty prompt and typical negatives are mixed into every batch (empty prompt is the CFG negative).
• Teacher: ran in 8-bit, bf16 does not fit on a 16 GB card.

Notes and limitations

• Teacher and student share the tokenizer (248,077 tokens, identical ids), so no token remapping: the prompt is apply_chat_template([{"role": "user", "content": text}], add_generation_prompt=True, enable_thinking=True), all tokens go to the DiT, nothing is dropped.
• The targets come from an 8-bit teacher (bitsandbytes computes matmuls in fp16), so they carry ~1% error against bf16 — small next to the 33% regression error, but not a bit-exact reference.
• JSON-layout prompts are the weakest case (0.756): such prompts are absent from the training pool.
• Trained on MOSS-Image-Base; not tested on MOSS-Image-SFT.
• Ceiling is the native encoder itself: the adapter reproduces its states, it does not make the DiT understand prompts it never understood.
• Apache-2.0. Upstream licenses of MOSS-Image and Qwen3.5-0.8B apply.

Samples

Native encoder (left) vs adapter (right), one seed, 1024 px, 30 steps:

t2i with the adapter only: