MOSS4GB — MOSS-Image-SFT on 4 GB cards (5B DiT, 4-bit, no text encoder)

Run MOSS-Image-SFT on a 4 GB card with the 9B text encoder never resident.

Piece	Where it runs	Footprint
DiT MOSSImageTransformer	GPU, NF4 4-bit	~2.6 GB
Text encoder Qwen3.5-9B	nowhere at inference — replaced by cached embeddings	0
VAE decode	CPU, fp32	250 MB RAM
Latents / sampling state	GPU, fp32	<10 MB

The DiT is 5.08 B parameters (10.16 GB bf16) and the encoder is 8.39 B (16.78 GB bf16). Only the encoder is genuinely impossible on this card, so the pipeline is built around deleting it from the runtime path.

Measured on a 1050 Ti (4 GB, sm_61, Windows)

Metric	Value
DiT load, streaming from disk	349-513 s (I/O bound, drive dependent)
VRAM resident	2.77 GB (262 Linear NF4, 7 fp16, 5.08 B params)
256x256, cfg 1	2.91 s/step
256x256, cfg 3 (two forwards/step)	13.8 s/step
Peak VRAM at 256x256	2.90 GB
VAE decode 256x256 / 512x512, CPU	7.8-28.5 s / 51.6 s
Decode vs upstream pipeline	max abs diff 0.0

Consequence for real use: a 20-step 256x256 render at cfg 3 is about 4.5 minutes of sampling on top of the load. Keep a warm process if you are iterating; the load-time NF4 cache (<model-dir>/.moss4gb-cache/) is written on the first generate and read back in about a minute after that.

Install from this repo

hf download AEUPH/moss4gb --local-dir moss4gb-hf
cd moss4gb-hf
pip install torch diffusers>=0.38 'bitsandbytes>=0.49' 'transformers>=5.3' safetensors huggingface_hub
# weights come from the upstream repo (transformer + vae only; ~9.7 GB)
hf download OpenMOSS-Team/MOSS-Image-SFT --include 'transformer/*' --include 'vae/*' --include 'pipeline_moss_image.py' --include 'scheduler/*' --local-dir models/MOSS-Image-SFT
python -m moss4gb doctor

Why this works at all

pipeline_moss_image.py already declares _optional_components = ['text_encoder', 'tokenizer', 'vae'] and accepts prompt_embeds: [B, S, 4096] instead of a string. The transformer only requires S >= 1. So the encoder's outputs are the real interface, and they are 300 KB each.

Quick start

# 1. weights (~10.5 GB; the 17 GB text encoder is only needed by `bank build`)
python -m moss4gb doctor

# 2. a placeholder bank so the path is runnable immediately (no conditioning)
python -m moss4gb bank null

# 3. render (256 to stay small; 512 needs the card almost entirely free)
python -m moss4gb generate --height 256 --width 256 --steps 20 --seed 42 --out out/first.png

Commands here are written as single lines on purpose: PowerShell's line continuation is a backtick, not a trailing \, so a bash-style \ becomes a literal argument (error: unrecognized arguments: \).

The null bank produces an unconditioned sample: useful for verifying VRAM, speed and that the VAE decode is correct, not for judging quality. With a one-token zero embedding the model produces structured noise, and because the seed and the conditioning are both constant the repeat is byte-identical.

Building a real prompt bank

This is the only step that does not fit on a 4 GB card. Full runbook: docs/BUILD_BANK_GPU.md.

python -m moss4gb bank build --device cuda --dtype bfloat16 --prompts-file prompts/characters.txt
python -m moss4gb bank build --device cuda --dtype bfloat16 --load-in-8bit --prompt "a woman with pink hair"

The DiT consumes Qwen3.5-9B hidden states, so a real bank needs that encoder once. It needs transformers>=5.3 and roughly 20 GB of RAM (or a GPU with ~12 GB free), and 17 GB of disk. That does not fit on the 4 GB / 16 GB target box, so build it elsewhere and copy the result.

Check the prompt formatting first — it costs seconds and no model download:

python -m moss4gb bank build --dry-run --prompt "a wide view of a red bicycle"
# -> 17 tokens, <|im_start|>user\n...<|im_end|>\n<|im_start|>assistant\n thinking\n

The repo's own tokenizer/ has no chat template (apply_chat_template raises), and the template decides exactly which tokens the DiT was trained on, so qwen35_chat_template.jinja is attached automatically. Don't substitute your own.

huggingface-cli download OpenMOSS-Team/MOSS-Image-SFT \
    --include 'text_encoder/*' --local-dir models/MOSS-Image-SFT
python -m moss4gb bank build --prompt "a wide view of a red bicycle"
python -m moss4gb bank build --prompts-file prompts.txt

# on the 4 GB box
python -m moss4gb bank import --source /path/to/bank
python -m moss4gb generate --bank-id p_157d4f7a84a3 --height 256 --width 256 --steps 20 --cfg 3 --seed 42 --out out/sample.png

A bank entry is ~4 MB fp16 per 500 tokens, so a few hundred prompts is a few hundred MB.

Notes and limits

• Precision. Pascal has no bf16, so weights are cast to fp16 and quantized to NF4 with compute_dtype=float16. If a run produces NaN/black output, the residual stream is overflowing fp16; retry with --compute-dtype float32 (about 2x slower), or swap in KBlueLeaf's FP16-fixed SFT transformer, whose residual rescaling exists for exactly this class of card.
• Resolution. 256x256 is comfortable. 512x512 leaves only a few hundred MB free once the DiT is resident, so it will page on a machine with a browser open; 1024x1024 is out of reach. --height/--width must be multiples of 16.
• Steps. 15-25 steps is the intended range. With --cfg 3 that is two transformer forwards per step; --cfg 1 halves the time at some quality cost. --uncond-skip 22 keeps the default negative-branch layer skip.
• System RAM. The loader is streaming: it never materializes the fp16 model, so peak RSS stays around 6.8 GB (one shard of evictable file pages) instead of 10.2 GB. This matters on a 16 GB box with ~6 GB free.
• Prompt bank. assets/bank currently holds the null bank, so renders are unconditioned. Building a real one needs the 9B encoder on a machine with ~20 GB RAM; see below.
• The conditioning is frozen. With no text encoder there is no instruction-following beyond what is in the bank. To get image-input conditioning you need a trained bridge into the text slot, not a flag.

Tests

python -m pytest tests -q      # 26 tests, ~21 s, no GPU needed

Covers the bank format, chat-template formatting and token counts, the NF4 skip list, and byte-level parity between our CPU decode and the upstream pipeline's own decode.

Files

moss4gb/quant.py    streaming NF4 loader (meta skeleton -> per-tensor -> GPU)
moss4gb/decode.py   CPU VAE decode: BatchNorm fold + 2x2 channel unshuffle
moss4gb/bank.py     embedding bank format, null bank, encoder-side builder
moss4gb/cli.py      doctor / bank / generate
moss4gb/metadata_cache.py   caches diffusers' import-time package scan
tests/              CPU-only suite

Trainable lane — one image in, one LoRA out

This pack renders for free, and the adapter lane is free too: a single-image loop (one image -> one 3 MB LoRA) with a measured recipe, plus the one loader bug (alpha / r, not r) that decides whether an adapter helps or turns the render into confetti:

TRAINABLE-REFERENCE.md

Offload to ZeroGPU — no local GPU required

Everything heavy runs on the free ZeroGPU slot. The only things computed off-device are the 2.77 GB NF4 cache download and the prompt bank storage.

Endpoint	Runs on the GPU	Notes
/generate	NF4 → dense bf16 expansion, sampling, VAE decode	256²/12 steps ≈ 1 s here; reserve scales with steps × area × CFG
/bank_train	Qwen3.5-9B encoder forward (one pass per new prompt)	bank entry ≈ 300 KB, reused for every later call
/bank_export	nothing	returns bank.safetensors for the local 4 GB path
/status	nothing	readiness, precisions, bank size

Python client (same auth tricks apply to any HF token):

from gradio_client import Client
c = Client("AEUPH/moss4gb", hf_token="hf_...")
img, meta = c.predict(
    prompt="a wide view of a red bicycle leaning on a white wall",
    negative_prompt="", height=256, width=256,
    steps=12, guidance_scale=1.0, seed=42, uncond_skip="22",
    api_name="/generate",
)
c.predict("a red bicycle on cobblestones", api_name="/bank_train")   # encoder pass
print(c.predict(api_name="/status"))

Agent/MCP: launch(mcp_server=True) exposes the same four calls as tools at https://aeuph-moss4gb.hf.space/gradio_api/mcp/.

Economy note: reservations are computed per request (reserve_seconds) and clamped to 120 s, so an anonymous visitor never burns the daily quota on a call it did not ask for.