Kromoss

Kromoss 是基于 MOSS-Image-Base 在 1024 px 分辨率下进行的完整微调模型。它沿用了 MOSS-Image-SFT 相同的架构、推理管线、文本编码器（Qwen3.5-9B）和 VAE（FLUX.2 small decoder），仓库目录结构也完全一致，因此可直接作为替换使用。

该仓库还附带了两个从 MOSS-Image-Base 蒸馏而来的 8 步 turbo LoRA：8 步，CFG 1，无负向分支。

• 8-steps-dmd2-lora.safetensors（推荐）：在下方 DMD LoRA 的基础上，进一步使用 DMD2 的 GAN 阶段对真实图像进行训练。对比度和细节表现略优，PartiPrompts 得分最高。
• 8-steps-dmd-lora.safetensors：数据空间 DMD 蒸馏。

50 步，CFG 3	8 步，CFG 1 + 8-steps-dmd-lora	8 步，CFG 1 + 8-steps-dmd2-lora

检查点：微调第 11000 步更新（EMA 权重），共 704k 张训练图像。

快速开始

pip install "diffusers>=0.38" "transformers>=5.3" "accelerate>=1.10" "safetensors>=0.6" "Pillow>=10"

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "lodestones/Kromoss",
    custom_pipeline="pipeline_moss_image",
    trust_remote_code=True,
    torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
)
pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)
pipe.to("cuda")

image = pipe(
    "a tabby cat holding a sign that says KROMOSS",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=3.0,
    uncond_skip_layers=[22],
    generator=torch.Generator("cpu").manual_seed(42),
).images[0]
image.save("kromoss.png")

Transformer 和文本编码器使用 bfloat16；VAE 和初始噪声使用 float32。

ComfyUI

ComfyUI 当前 master 分支已原生支持 MOSS-Image。单文件采用与 Comfy-Org/MOSS-Image 相同的布局：融合 proj.qkv，bf16。

kromoss_bf16.safetensors 可直接替换 Comfy-Org 工作流中的 moss_image_sft_bf16.safetensors（50 步，CFG 3）。配合任一 LoRA 使用时，KSampler 设置为 euler / simple，8 步，CFG 1。 ComfyUI 会将 LoRA 中独立的 q/k/v 映射到融合的 qkv 切片上。

convert_to_comfy.py 可将任意 MOSS-Image diffusers transformer/ 文件夹（例如你自己微调的模型）转换为此单文件格式（仅依赖 torch 和 safetensors）：

python convert_to_comfy.py path/to/transformer my_model_bf16.safetensors

关于 MOSS-Image-SFT 的 transformer/ 目录，它能逐字节复现 Comfy-Org 的 moss_image_sft_bf16.safetensors；--compare FILE 会对照参考文件检查键名、形状和数据类型。

kromoss_int8_convrot.safetensors 采用 ConvRot W8A8 格式，与 Comfy-Org 的 moss_image_sft_int8_convrot.safetensors 格式一致。其中 168 个块级线性层（proj.{qkv,o,ffn_up_gate,ffn_down}）以 int8 存储并附带逐行缩放因子，输入维度经过 256 维常规 Hadamard 旋转。ComfyUI 在加载时重建旋转矩阵并在线作用于激活值，因此无需存储旋转矩阵。每个块同时启用 ComfyUI 的 int8 注意力机制。该文件可直接替代 Comfy-Org 的 int8 文件，且与 LoRA 兼容。quantize_convrot.py 可从 bf16 单文件生成该文件；对 moss_image_sft_bf16.safetensors 执行时，可逐字节复现 Comfy-Org 的 int8 文件：

python quantize_convrot.py my_model_bf16.safetensors my_model_int8_convrot.safetensors

8 步 Turbo LoRA

两者均从 MOSS-Image-Base（50 步，CFG 3）蒸馏而来，并迁移至 Kromoss：

• 8-steps-dmd-lora.safetensors：数据空间 DMD，训练 12000 步。
• 8-steps-dmd2-lora.safetensors：在前者基础上继续训练 3000 步，使用 DMD2 的第二阶段——一个在假分数上训练、以真实图像为对抗目标的 GAN 判别器，学习率降至四分之一。

两者均为 fp32 精度，每个 block 的 proj.{q,k,v,o,ffn_up_gate,ffn_down} 层 rank 128 / alpha 128，采用 kohya 键名格式（diffusion_model.<module>.lora_down.weight / lora_up.weight / alpha）。请一次仅加载一个。在 ComfyUI 中，通过 LoRA Loader 以强度 1 加载；在 diffusers 中，以非合并方式应用。请勿将其合并进 bf16 权重：其数值增量已接近 bf16 的舍入精度。4 步同样可用。

from safetensors.torch import load_file

def apply_lora(transformer, path, strength=1.0):
    state = load_file(path, device="cuda")
    modules = dict(transformer.named_modules())
    for key in [k for k in state if k.endswith(".lora_down.weight")]:
        name = key[len("diffusion_model."):-len(".lora_down.weight")]
        down, up = state[key].float(), state[key.replace("lora_down", "lora_up")].float()
        scale = strength * state[f"diffusion_model.{name}.alpha"].item() / down.shape[0]
        def hook(module, args, out, down=down, up=up, scale=scale):
            return out + (scale * (args[0].float() @ down.T) @ up.T).to(out.dtype)
        modules[name].register_forward_hook(hook)

apply_lora(pipe.transformer, "8-steps-dmd2-lora.safetensors")   # hf_hub_download("lodestones/Kromoss", ...)
image = pipe(prompt, height=1024, width=1024, num_inference_steps=8, guidance_scale=1.0,
             generator=torch.Generator("cpu").manual_seed(42)).images[0]

采样设置

PartiPrompts（100 条提示词，CLIP ViT-L/14 得分；Kromoss 行在第 11000 次更新时测量）

下方表格中的各列与上表顺序一致。

训练

• 对 Transformer 的全部参数进行全量微调，分辨率 1024 px（小于 896² 的图像被丢弃）。
• 采用 MOSS-Image 报告中最终阶段的训练配方，缩放至每次更新 64 张图像：
  • 混合 Muon-AdamW 优化器，学习率 9e-6，预热 200 步；
  • 截断对数正态时间步；
  • 最终块 q/k-norm 衰减至一；
  • EMA 0.999（发布的权重即为 EMA 权重）。

许可与致谢

权重及附带的推理代码（pipeline_moss_image.py，源自 MOSS-Image-SFT）基于 Apache License, Version 2.0 发布。本项目基于 OpenMOSS Team 开发的 MOSS-Image 构建。