Instructions to use AiArtLab/Ra with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Diffusers
How to use AiArtLab/Ra with Diffusers:
pip install -U diffusers transformers accelerate
import torch from diffusers import DiffusionPipeline # switch to "mps" for apple devices pipe = DiffusionPipeline.from_pretrained("AiArtLab/Ra", dtype=torch.bfloat16, device_map="cuda") prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" image = pipe(prompt).images[0] - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- Draw Things
- DiffusionBee
Download docs/ARCH.md from AiArtLab/Ra: direct link, hf CLI and curl.
- Browser
- Download file 6.03 kB
-
https://huggingface.co/AiArtLab/Ra/resolve/main/docs/ARCH.md
- Command line
-
hf download hf://AiArtLab/Ra/docs/ARCH.md
-
curl -L -o ARCH.md https://huggingface.co/AiArtLab/Ra/resolve/main/docs/ARCH.md
Архитектура
Три части, каждая со своей задачей. Всё, что ниже, проверено запуском.
1. Тело: пирамидальный U-трансформер
Принцип: ширина растёт там, где токенов мало. Стоимость слоя ~ токены × C², поэтому на
высоком разрешении держим узкие свёрточные блоки, а ширину и полное внимание — внизу.
| уровень | токенов при 512² | ширина | блоки | механика |
|---|---|---|---|---|
| L0 ÷1 | 4096 (512²) / 1024 (256²) | 192 | 6 enc / 4 dec + 1 self-attn с каждой стороны | depthwise 3×3 + expand 4× (ConvNeXt-стиль), кросс-внимание к тексту в каждом блоке |
| L1 ÷2 | 1024 / 256 | 512 | 4 enc / 3 dec | трансформер: self-attn + cross-attn к тексту + FFN |
| L2 ÷4 | 256 / 64 | 1024 | 3 enc / 2 dec | то же, полное внимание дёшево |
| L3 ÷8 | 64 / 16 | 1536 | 2 (боттлнек) | то же; широкий низ = глобальные ограничения (симметрия, композиция) |
Переходы и связи:
- вниз: mean-pool 2×2 +
Linear, вверх:Linear+ nearest 2×; - скипы аддитивные (SD-стиль): выходы
e0/e1/e2прибавляются к декодеру того же уровня; ширина на входе и выходе уровня совпадает, поэтому сложение без проекций; - RoPE на каждом уровне, фаза в латентных пикселях (L1 сдвинут вдвое, L2 вчетверо), поэтому одна фаза означает одно расстояние независимо от уровня;
- модуляция: общий вектор
t + pooled-текст→(shift, scale, gate)на уровень; в conv-блоках L0 пять чанков (три на conv-путь, два на вход кросс-внимания); - zero-init: выход conv-блоков и FFN обнулён, поэтому новые ветки не ломают старт.
Размеры: тело 143.3M. Шаг обучения 97 мс при bs8 на 256² (8.1 ГБ). Для сравнения, прежний
плоский стек ширины 1024 на всех уровнях: 196 мс и hf 0.18 против 0.54 у пирамиды (архитектура
удалена в tmp/archive).
2. Текст
- Ettin-150M (
text_encoder/): ModernBERT, 768, 22 слоя, byte-level BPE, MIT, заморожен. - SliceMixer (
src/text.py, 10.04M обучаемых): 7 срезов (выходы после глобальных слоёв 2, 5, 8, 11, 14, 17, 20) → пер-срезный LayerNorm (нормы растут 52 → 548, без нормировки ранние слои не влияют) → 2 блока attention по оси срезов →Linear(7→1)→ zero-init добавка к среднему по срезам. Идея изmoss/adapter, где 0.8B + 158M голова заменили 9B-энкодер. - Вход в тело: кросс-внимание на всех уровнях (общие ключи/значения на уровень) плюс pooled-текст в модуляции.
3. Хвост
Декодер AE без изменений архитектуры: патчи 8×8, 6 блоков @512, adaLN с низкоранговыми дельтами,
long skip из блока n−2, финальный PixelShuffle + conv3×3 ×3. 23.87M.
Обучается как обычный flow по пикселям: случайные t, случайный шум, цель — картинка, условие —
z0 + (ошибка тела).detach(). Условие без градиента: рендерер видит реальную ошибку тела, но не
может «научиться исправлять мусор» (именно это раньше увело его с 29.6 на 21.3 дБ).
4. Сэмплинг
z ~ N(0,1)
4–8 шагов Эйлера по телу: x0 = тело(z, t, текст); z ← z − (z − x0)/(t·N)
CFG 1.5 (выше 2 при недообученной безусловной ветви ухудшает латент)
1 + 3 шага декодера AE: изображение
При 256² сэмпл стоит ~35 мс на 4800-й карте; при 512² — примерно вчетверо дороже.
5. Бюджет
| часть | параметры |
|---|---|
| тело | 143.3M |
| выбор слоёв текста | 10.04M |
| хвост | 23.87M |
| итого | 177.3M |
6. Что зафиксировано и почему
- x0, а не velocity. Смена параметризации тела не затрагивает декодер (это отдельный flow со своей целью), но требует переучивания тела с нуля — поэтому оставлено как есть.
- Текст не экономим. Кросс-внимание на L0 добавило 2.5M параметров и 9 мс,
switchвырос 0.43 → 0.78 при том же уровне деталей. - Никаких фиксированных шумов и клипов в обучении — только стандартный flow-matching.
- Удалено как тупиковое: плоский 1024-стек, резидуальная U-ветка, Mix-FFN поверх плоского стека, velocity-переключение посреди обучения, LPIPS в лоссе, кропы одной картинки вместо набора.