中文Sign up

p1atdev/MOSS-Image-SFT-KV-Cache-LoRA

Source: 🤗text-to-imageapache-2.0299 MBother✓ 6 checksumsupdated today
or submit a torrent

MOSS-Image-SFT KV Cache LoRA

An experimental correction LoRA for fixed-timestep text KV caching in MOSS-Image-SFT. It reduces the prediction difference caused by computing text-only attention once and reusing its keys and values during image denoising. Long prompts can generate faster; short prompts can be slower.

Recommended: v2, 500 updates. Compared with v1's 10 training prompts, v2 uses 512 prompts covering short descriptions, anime, photography, typography and MOSS JSON layouts. The v1, 100-update checkpoint remains available for older workflows. Version prefixes identify separate training generations; existing weights are not replaced.

Use in ComfyUI

Install comfy-moss-image-kv-cache and drag its v2 example workflow into ComfyUI. An API graph is also provided.

  1. Download the SFT base, text encoder and VAE from Comfy-Org/MOSS-Image.
  2. Place the frozen DMD adapter lora_latest.safetensors and moss_image_sft_kv_cache_v2_step000500.safetensors in ComfyUI/models/loras/.
  3. Use built-in Load LoRA (Model Only) nodes: DMD strength 0.7, then correction strength 1.0.
  4. Add MOSS-Image Text KV Cache after the LoRA loaders, with text_timestep 1000.
  5. Sample with Euler / simple / 8 steps / CFG 1, 1024×1024. The example uses BF16 model and text encoding; launch ComfyUI with --bf16-text-enc for the latter.

The KV node is required; the adapter alone does not enable caching. DMD weights are a separate dependency. No trigger word is needed. To use joint attention, bypass both the correction LoRA and KV node, keeping DMD enabled.

v2 training

  • Format: standard ComfyUI LoRA keys, FP32 weights, rank 16 / alpha 16, approximately 149 MB. All 42 DiT blocks' attention Q/K/V/output and MLP up-gate/down projections are targeted; shared modulation is excluded (37,158,912 trainable parameters).
  • Teacher/student: frozen BF16 joint-attention teacher; NF4 student except shared modulation, with text-only attention at t=1000. Both use the frozen DMD adapter above at strength 0.7. The exported correction is also tested with a BF16 student in native ComfyUI.
  • Data: 512 training and 128 separate validation prompts derived from the three datasets listed above. Only prompt text and structured layout descriptions are used; camera layouts are converted to MOSS JSON. Original dataset images are not used. Of the 512 training prompts, 282 have at most 64 tokens, 102 are anime prompts and 51 are structured MOSS layouts.
  • Objective: prediction MSE on matching teacher trajectory states and timesteps; eight steps, one training noise seed (42), 1024×1024.
  • Optimization: 500 updates with RAdamScheduleFree, learning rate 5e-4, batch 8 × accumulation 2. Weights are exported in optimizer evaluation mode.

On this run's 128 validation prompts, NF4 prediction MSE decreased from 0.20427 to 0.08135 (60.2%); for prompts of at most 64 tokens, 0.19168 → 0.08105 (57.7%). These are teacher-trajectory prediction losses, not image-quality scores or a comparison with v1's different validation set.

SFT training revision: fc5de0d56fc910cbca6f6625a70f5f582b26deb1. DMD revision: f95df83b5e90ae91ea412434b3f9b7f1c8c6314f. See checksum.sha256 for file verification.

ComfyUI comparisons

All images below were generated in native ComfyUI with the BF16 SFT base, BF16 text encoder, DMD 0.7, Euler/simple, 8 steps, CFG 1 and 1024×1024. Columns compare joint attention / KV alone / KV + v2, using the same seed within each comparison. NF4 is not used in these ComfyUI comparisons.

On an RTX 6000 Ada, the 2,941-token layout prompt took 8.03 s with joint attention and 5.26 s with KV + v2 (1.53× faster). The 53-token prompt took 3.98 → 4.39 s. Three-run warm medians include fresh KV construction, sampling, VAE decode and PNG saving; text encoding and initial model loading are excluded.

Full prompts, more comparison images, raw measurements and reproduction scripts are in the custom node README.

Limits

KV caching changes the attention computation. This correction does not reproduce joint-attention pixels, and composition, counting, text and fine details can still differ. Visual checks are small-scale; improvements are not guaranteed for every prompt. The 2,941-token demonstration is longer than the training prompts (up to 928 tokens). Other DMD adapters or strengths, step counts and text timesteps are outside this training recipe. The base model's limitations also apply.

Thanks to OpenMOSS, lodestones, and the prompt dataset authors.