#!/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()