Fix model card image links with absolute resolve URLs

7475a87 verified 26 days ago

6.31 kB

---
license: apache-2.0
pipeline_tag: text-to-image
library_name: diffusion-single-file
base_model: OpenMOSS-Team/MOSS-Image-SFT
base_model_relation: adapter
datasets:
  - Photoroom/midjourney-v6-recap
  - jaddai/openbrush
  - terminusresearch/pseudo-camera-10k-structured-json
tags:
  - comfyui
  - lora
  - moss-image
  - kv-cache
---
# 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](moss_image_sft_kv_cache_v2_step000500.safetensors).** 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](moss_image_sft_kv_cache_v1_step000100.safetensors) 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](https://github.com/p1atdev/comfy-moss-image-kv-cache) and drag its [v2 example workflow](https://github.com/p1atdev/comfy-moss-image-kv-cache/blob/main/examples/moss_kv_v2.json) into ComfyUI. An [API graph](https://github.com/p1atdev/comfy-moss-image-kv-cache/blob/main/examples/moss_kv_v2_api.json) is also provided.

1. Download the **SFT** base, text encoder and VAE from [Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image).
2. Place the frozen DMD adapter [lora_latest.safetensors](https://huggingface.co/lodestones/MOSS-Image-SFT-8-steps-dmd/resolve/f95df83b5e90ae91ea412434b3f9b7f1c8c6314f/lora_latest.safetensors) and [moss_image_sft_kv_cache_v2_step000500.safetensors](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](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.

![Anime prompt comparisons in ComfyUI](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/anime_comparison_v2.jpg)

![Photo, object count and typography comparisons in ComfyUI](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/prompt_comparison_v2.jpg)

![Long JSON prompt comparison in ComfyUI](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/long_prompt_comparison.jpg)

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.

![BF16 KV-cache benchmark with v2](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/kv_cache_benchmark.png)

Full prompts, more comparison images, raw measurements and reproduction scripts are in the [custom node README](https://github.com/p1atdev/comfy-moss-image-kv-cache#speed-with-long-prompts).

## 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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT), [lodestones](https://huggingface.co/lodestones/MOSS-Image-SFT-8-steps-dmd), and the prompt dataset authors.

	---
	license: apache-2.0
	pipeline_tag: text-to-image
	library_name: diffusion-single-file
	base_model: OpenMOSS-Team/MOSS-Image-SFT
	base_model_relation: adapter
	datasets:
	- Photoroom/midjourney-v6-recap
	- jaddai/openbrush
	- terminusresearch/pseudo-camera-10k-structured-json
	tags:
	- comfyui
	- lora
	- moss-image
	- kv-cache
	---

	# 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](moss_image_sft_kv_cache_v2_step000500.safetensors). 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](moss_image_sft_kv_cache_v1_step000100.safetensors) 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](https://github.com/p1atdev/comfy-moss-image-kv-cache) and drag its [v2 example workflow](https://github.com/p1atdev/comfy-moss-image-kv-cache/blob/main/examples/moss_kv_v2.json) into ComfyUI. An [API graph](https://github.com/p1atdev/comfy-moss-image-kv-cache/blob/main/examples/moss_kv_v2_api.json) is also provided.

	1. Download the SFT base, text encoder and VAE from [Comfy-Org/MOSS-Image](https://huggingface.co/Comfy-Org/MOSS-Image).
	2. Place the frozen DMD adapter [lora_latest.safetensors](https://huggingface.co/lodestones/MOSS-Image-SFT-8-steps-dmd/resolve/f95df83b5e90ae91ea412434b3f9b7f1c8c6314f/lora_latest.safetensors) and [moss_image_sft_kv_cache_v2_step000500.safetensors](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](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.

	![Anime prompt comparisons in ComfyUI](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/anime_comparison_v2.jpg)

	![Photo, object count and typography comparisons in ComfyUI](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/prompt_comparison_v2.jpg)

	![Long JSON prompt comparison in ComfyUI](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/long_prompt_comparison.jpg)

	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.

	![BF16 KV-cache benchmark with v2](https://huggingface.co/p1atdev/MOSS-Image-SFT-KV-Cache-LoRA/resolve/main/assets/kv_cache_benchmark.png)

	Full prompts, more comparison images, raw measurements and reproduction scripts are in the [custom node README](https://github.com/p1atdev/comfy-moss-image-kv-cache#speed-with-long-prompts).

	## 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](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT), [lodestones](https://huggingface.co/lodestones/MOSS-Image-SFT-8-steps-dmd), and the prompt dataset authors.