Blame for adapter/src/moss_lib.py

#!/usr/bin/env python3
"""Всё про MOSS-Image в одном месте: шаблон промпта, учитель (Qwen3.5-9B), ученик (0.8B), пайплайн.

Задача та же, что в docs/qwen-image-2.1-adapter.md: заменить текстовый энкодер 15.6 ГБ
(Qwen3.5-VL-9B, `language_model`, hidden 4096) на замороженную 0.8B + адаптер. Но здесь
проще и одновременно строже:

  * токенизатор у учителя и ученика ОДИН (248 077 токенов, совпадают и строки, и id),
    поэтому общая сегментация не нужна — id одни и те же;
  * цель — `last_hidden_state`, то есть ПОСТ-норм выход слоя 32 (Qwen3_5TextModel сам
    применяет `norm`). Это не то же самое, что в 2.1 (там был пред-норм выход), и именно
    это видит DiT: text_embedder.proj 4096 -> 3072;
  * учитель и ученик — одно семейство (Qwen3.5 hybrid linear attention), срезы ученика
    берутся после его full-attention слоёв (3,7,11,15,19,23) -> hs[4,8,12,16,20,24];
  * пайплайн принимает `prompt_embeds`/`negative_prompt_embeds` официально, monkeypatch
    не нужен (но метод `from_pretrained(text_encoder=None)` НЕ пропускает загрузку TE —
    пайплайн собираем вручную, см. build_pipeline).
"""
import json
import os
import sys

import torch

MOSS_DIR = os.environ.get("MOSS_DIR", "/home/recoilme/moss")
STUDENT = os.environ.get("MOSS_STUDENT", "Qwen/Qwen3.5-0.8B")
# раскладка папки модели: ученик — в text_encoder/, оригинальный 9B — рядом под явным именем
STUDENT_SUBDIR = "text_encoder"
TEACHER_SUBDIRS = ("text_encoder_qwen3.5-9b", "text_encoder")
# после полных attention-слоёв ученика (0-based 3,7,11,15,19,23) -> hidden_states[4,8,...]
LAYERS = (4, 8, 12, 16, 20, 24)


# ---------------------------------------------------------------- токены и шаблон
def load_tokenizer(moss_dir=None):
    """Токенизатор учителя = токенизатор ученика, берём из папки модели."""
    from transformers import AutoTokenizer
    return AutoTokenizer.from_pretrained(os.path.join(moss_dir or MOSS_DIR, "tokenizer"))


def render(tok, text):
    """Ровно то, что делает `MOSSImagePipeline.encode_prompt` (сырая строка, не id)."""
    return tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False,
                                   add_generation_prompt=True, enable_thinking=True)


def encode(tok, texts, max_len=4096):
    """list[str] -> list[list[int]]: шаблон + обрезка хвоста, как в пайплайне (truncation)."""
    return [tok(render(tok, t), truncation=True, max_length=max_len)["input_ids"] for t in texts]


# ---------------------------------------------------------------- модели
def _has_config(path):
    return os.path.isfile(os.path.join(path, "config.json"))


def teacher_dir(moss_dir=None):
    """Папка оригинального текстового энкодера 9B (не удаляем и не переименовываем в text_encoder)."""
    root = moss_dir or MOSS_DIR
    for name in TEACHER_SUBDIRS:
        path = os.path.join(root, name)
        if _has_config(path):
            return path
    return os.path.join(root, TEACHER_SUBDIRS[0])


def student_dir(moss_dir=None):
    """Ученик: локальная копия в папке модели (`text_encoder/`), иначе HF-репозиторий."""
    local = os.path.join(moss_dir or MOSS_DIR, STUDENT_SUBDIR)
    return local if _has_config(local) else STUDENT


def load_text_model(path, dtype=torch.bfloat16, load_8bit=False, device="cuda"):
    """Qwen3.5 в любой раскладке — класс берём из `config.json`.

    Две раскладки одного и того же семейства: `Qwen3_5Model` (ключи `language_model.*` — так лежит
    текстовый энкодер MOSS) и `Qwen3_5ForConditionalGeneration` (ключи `model.*` — родной чекпойнт
    Qwen3.5-0.8B). Угадывать нельзя: не та раскладка грузится без ошибки, но со случайными весами
    и только предупреждением в логе — поэтому смотрим architectures и падаем, если не узнали.
    """
    import transformers
    from transformers import AutoConfig
    arch = (AutoConfig.from_pretrained(path, local_files_only=os.path.isdir(path)).architectures or [""])[0]
    classes = {"Qwen3_5Model": transformers.Qwen3_5Model,
               "Qwen3_5ForConditionalGeneration": transformers.Qwen3_5ForConditionalGeneration}
    if arch not in classes:
        raise SystemExit(f"{path}: не знаю, как грузить architectures={arch!r}")
    local = {"local_files_only": os.path.isdir(path)}
    if load_8bit:
        from transformers import BitsAndBytesConfig
        return classes[arch].from_pretrained(
            path, quantization_config=BitsAndBytesConfig(load_in_8bit=True),
            device_map={"": 0}, **local).eval()
    return classes[arch].from_pretrained(path, dtype=dtype, **local).to(device).eval()


def text_stack(model):
    """Текстовая стопка без vision-tower: у MOSS TE это `language_model`, у 0.8B — `model.language_model`."""
    for name in ("language_model", "model.language_model"):
        obj = model
        for part in name.split("."):
            obj = getattr(obj, part, None)
            if obj is None:
                break
        if obj is not None and hasattr(obj, "norm"):
            return obj
    raise SystemExit("не нашёл текстовую модель с .norm")


class Teacher:
    """Qwen3.5-VL-9B (текстовая часть): цель = post-norm `last_hidden_state`, (B,L,4096)."""

    def __init__(self, path=None, device="cuda", dtype=torch.bfloat16, load_8bit=False):
        path = path or teacher_dir()
        self.base = load_text_model(path, dtype=dtype, load_8bit=load_8bit, device=device)
        self.text = text_stack(self.base)
        self.quant = "8bit" if load_8bit else str(dtype).split(".")[-1]
        cfg = self.text.config
        print(f"учитель  : {sum(p.numel() for p in self.base.parameters()) / 1e9:.2f}B, "
              f"слоёв {cfg.num_hidden_layers}, hidden {cfg.hidden_size}, {self.quant}, "
              f"VRAM {torch.cuda.memory_allocated() / 2**30:.1f} ГиБ", flush=True)

    @torch.no_grad()
    def last_hidden(self, ids, mask=None):
        """(B,L) -> (B,L,4096) float32, пост-норм выход последнего слоя."""
        return self.text(input_ids=ids, attention_mask=mask, use_cache=False).last_hidden_state.float()


def load_teacher(device="cuda", prefer_8bit=False, path=None):
    """Учитель целиком, а если не влезает в карту — 8 бит. Возвращает (teacher, режим).

    После неудачной попытки надо отпустить traceback: он держит кадр с недогруженной моделью,
    и без gc/empty_cache вторая попытка падает на уже занятой памяти (проверено).
    """
    order = [True, False] if prefer_8bit else [False, True]
    errors = []
    for load_8bit in order:
        name = "8-бит" if load_8bit else "bf16"
        try:
            return Teacher(path=path, device=device, load_8bit=load_8bit), \
                ("8bit" if load_8bit else "bf16")
        except (torch.OutOfMemoryError, RuntimeError) as exc:
            if "out of memory" not in str(exc).lower():
                raise
            free_gib = torch.cuda.mem_get_info()[0] / 2**30
            print(f"  {name} не влез ({free_gib:.1f} ГиБ свободно) — пробую дальше", flush=True)
            errors.append(f"{name}: {str(exc).splitlines()[0][:100]}")
            del exc
            empty_cache()
    raise SystemExit("учитель не грузится: " + " | ".join(errors))


class Student:
    """Qwen3.5-0.8B (заморожен): срезы hidden_states -> (B, L, K*hidden)."""

    def __init__(self, path=None, layers=LAYERS, device="cuda", dtype=torch.bfloat16):
        path = path or student_dir()
        self.path = path
        self.base = load_text_model(path, dtype=dtype, device=device)
        self.base.requires_grad_(False)
        self.text = text_stack(self.base)
        if self.text.config.hidden_size >= 4096:
            raise SystemExit(f"{path}: hidden {self.text.config.hidden_size} — это учитель, а не "
                             f"ученик (у Qwen3.5-0.8B hidden 1024)")
        self.layers = list(layers)
        self.device = device
        cfg = self.text.config
        print(f"ученик   : {sum(p.numel() for p in self.base.parameters()) / 1e9:.2f}B, "
              f"слоёв {cfg.num_hidden_layers}, hidden {cfg.hidden_size}, срезы {self.layers}",
              flush=True)

    @property
    def in_dim(self):
        return len(self.layers) * self.text.config.hidden_size

    @torch.no_grad()
    def hidden(self, ids, mask=None):
        out = self.text(input_ids=ids, attention_mask=mask, output_hidden_states=True, use_cache=False)
        hs = out.hidden_states
        assert hs is not None and len(hs) > max(self.layers), f"скрытых состояний {len(hs)}"
        x = torch.stack([hs[i] for i in self.layers], dim=1)                 # (B, K, L, d)
        return x.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[2], -1).float()


# ---------------------------------------------------------------- инференс
def build_pipeline(moss_dir=None, dit_dtype=torch.bfloat16, vae_dtype=torch.float32):
    """MOSSImagePipeline без текстового энкодера.

    `DiffusionPipeline.from_pretrained(..., text_encoder=None)` здесь НЕ работает: text_encoder
    объявлен в __init__ с дефолтом None, поэтому diffusers не считает его «переданным» и лезет
    грузить 15.6 ГБ. Собираем компоненты руками — их три, и все локальные.
    """
    moss_dir = moss_dir or MOSS_DIR
    if moss_dir not in sys.path:
        sys.path.insert(0, moss_dir)
    from diffusers import AutoencoderKLFlux2
    from pipeline_moss_image import MOSSImagePipeline, MOSSImageScheduler, MOSSImageTransformer

    transformer = MOSSImageTransformer.from_pretrained(
        os.path.join(moss_dir, "transformer"), torch_dtype=dit_dtype, local_files_only=True)
    vae = AutoencoderKLFlux2.from_pretrained(
        os.path.join(moss_dir, "vae"), torch_dtype=vae_dtype, local_files_only=True)
    cfg = json.load(open(os.path.join(moss_dir, "scheduler", "scheduler_config.json")))
    keys = ("num_train_timesteps", "shift", "training_shift")
    scheduler = MOSSImageScheduler(**{k: cfg[k] for k in keys if k in cfg})
    return MOSSImagePipeline(transformer.eval(), scheduler, text_encoder=None, tokenizer=None,
                             vae=vae.eval())


@torch.no_grad()
def condition(adapter, student, tok, texts, max_len=4096):
    """list[str] -> list[(embeds [L,4096] bf16, mask [L] bool)]: по одному промпту без паддинга.

    Пайплайн кодирует каждый промпт отдельно (`unpadded individual forwards`), поэтому и здесь
    так же: длину маскирует сам DiT, а не токенизатор.
    """
    device = next(adapter.parameters()).device
    out = []
    for ids in encode(tok, texts, max_len):
        si = torch.tensor([ids], dtype=torch.long, device=device)
        mask = torch.ones_like(si, dtype=torch.bool)
        y = adapter(student.hidden(si, mask), mask)
        out.append((y[0].to(torch.bfloat16).cpu(), mask[0].cpu()))
    return out


@torch.no_grad()
def generate(pipe, pos, neg, size=1024, steps=30, seed=1234, guidance=3.0, device="cuda"):
    """pos/neg: (embeds [L,4096], mask [L]) -> PIL. CFG включён всегда (uncond_skip_layers)."""
    emb, mask = pos
    nemb, nmask = neg
    gen = torch.Generator(device).manual_seed(seed)
    return pipe(prompt_embeds=emb[None].to(device), prompt_attention_mask=mask[None].to(device),
                negative_prompt_embeds=nemb[None].to(device),
                negative_prompt_attention_mask=nmask[None].to(device),
                height=size, width=size, num_inference_steps=steps, guidance_scale=guidance,
                generator=gen, output_type="pil").images[0]


def empty_cache():
    """Вернуть освобождённую память карты. Сами объекты надо удалить у себя (`del`): `del` внутри
    функции снимает только локальную ссылку, объект остаётся жив, пока на него смотрит вызывающий."""
    import gc
    gc.collect()
    torch.cuda.empty_cache()