MOSS-Image-heal

Полный файнтюн MOSS-Image-Base от KBlueLeaf: трансформер обучен так, чтобы убрать огромные выбросы в residual stream. Остальные компоненты — текстовый энкодер, токенизатор, VAE, scheduler и код пайплайна — взяты из официальной базовой модели. Размер трансформера — около 5,08 млрд параметров.

🔎 Ключевые особенности

• ⬇ Максимальное значение в stream снижено с 58 411 до 3 083; до fp16-потолка остаётся 21,2x запаса.
• ⚙ RMS stream в среднем по блокам упал с 27,75 до 3,70 — примерно в 7,5 раза.
• 🛡 Веса аккуратнее: ни один элемент не превышает 64x RMS своего тензора (у базы таких было 349).
• 🎯 Выходы близки к базе: flow loss 0,4289 против 0,4280, velocity cosine 0,9983.
• 💻 Есть файлы для diffusers и ComfyUI: bf16, fp16 и int8 ConvRot.

🧑‍💻 Пример

Загрузка пайплайна с официальной базой; трансформер затем подставляется из KBlueLeaf/MOSS-Image-heal.

import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
    "OpenMOSS-Team/MOSS-Image-Base",
    custom_pipeline="pipeline_moss_image",
    trust_remote_code=True, transformer=None,
    torch_dtype=torch.bfloat16,
)

Подробнее

Обучение: 50 000 шагов × 64 сэмпла, около 3,2 млн сэмплов, AdamW, lr 1e-5; примерно 200 GPU-часов на B300.
Данные: около 1 млн постов danbooru и GBC-подписи; в карточке есть отчёты по residual stream и точности файлов.
VRAM не указан; статистика stream измерена на 6 text-image сэмплах.