| |
| """A/B картинками: родной энкодер MOSS (Qwen3.5-9B) против 0.8B + адаптер. |
| |
| Один промпт, один сид, один размер — картинки встают рядом, слева родной, справа адаптер. |
| Условие родного энкодера должно быть посчитано заранее (`native_embeds.py`): 15.6 ГБ учителя и |
| DiT в 16 ГБ не влезают одновременно, а условие от промпта не зависит. |
| |
| python3 src/native_embeds.py --out out/native_embeds.pt |
| python3 src/sample_ab.py --adapter out/adapter_v1.safetensors --out-dir out/samples |
| """ |
| import argparse |
| import os |
| import sys |
|
|
| import torch |
| import torch.nn.functional as F |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from adapter_lib import load_adapter |
| from moss_lib import (MOSS_DIR, STUDENT, Student, build_pipeline, condition, |
| empty_cache, generate, load_tokenizer) |
| from prompts import AB |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf" |
|
|
|
|
| def grid(rows, out_path, gap=8, header=40, caption=28, bg=(24, 24, 24), fg=(235, 235, 235)): |
| """rows: [(подпись, prompt, [PIL, PIL])]; колонки = родной / адаптер.""" |
| try: |
| font_h, font_c = ImageFont.truetype(FONT, 20), ImageFont.truetype(FONT, 14) |
| except OSError: |
| font_h = font_c = ImageFont.load_default() |
| width = max(sum(im.width for im in r[2]) + gap * (len(r[2]) - 1) for r in rows) |
| height = header + sum(caption + gap + max(im.height for im in r[2]) for r in rows) |
| canvas = Image.new("RGB", (width, height), bg) |
| draw = ImageDraw.Draw(canvas) |
| draw.text((8, 10), "родной Qwen3.5-9B (слева) / 0.8B + адаптер (справа)", font=font_h, fill=fg) |
| y = header |
| for name, prompt, images in rows: |
| draw.text((8, y), f"{name}: {prompt[:110]}", font=font_c, fill=fg) |
| y += caption |
| x = 0 |
| for im in images: |
| canvas.paste(im, (x, y)) |
| x += im.width + gap |
| y += max(im.height for im in images) + gap |
| canvas.save(out_path) |
| return out_path |
|
|
|
|
| def cosine(a, b): |
| n = min(a.shape[0], b.shape[0]) |
| return F.cosine_similarity(a[:n].float(), b[:n].float(), dim=-1).mean().item() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="A/B: родной энкодер MOSS против адаптера") |
| ap.add_argument("--adapter", default=os.path.join(ROOT, "adapter.safetensors")) |
| ap.add_argument("--native", default=os.path.join(ROOT, "out", "native_embeds.pt")) |
| ap.add_argument("--moss-dir", default=MOSS_DIR) |
| ap.add_argument("--student", default=STUDENT) |
| ap.add_argument("--out-dir", default=os.path.join(ROOT, "out", "samples")) |
| ap.add_argument("--size", type=int, default=1024) |
| ap.add_argument("--steps", type=int, default=30) |
| ap.add_argument("--guidance", type=float, default=3.0) |
| ap.add_argument("--seed", type=int, default=1234) |
| args = ap.parse_args() |
| if not os.path.exists(args.native): |
| raise SystemExit(f"нет {args.native} — сначала запустите src/native_embeds.py") |
|
|
| native = torch.load(args.native, map_location="cpu", weights_only=True) |
| os.makedirs(args.out_dir, exist_ok=True) |
| tok = load_tokenizer(args.moss_dir) |
| adapter, meta = load_adapter(args.adapter, device="cuda") |
| student = Student(args.student, [int(x) for x in str(meta["student_layers"]).split(",")]) |
| pos = dict(zip([n for n, _ in AB], condition(adapter, student, tok, [t for _, t in AB]))) |
| neg = condition(adapter, student, tok, [""])[0] |
| for name, emb in pos.items(): |
| print(f" {name:10s} cos(условие адаптера, родное) = {cosine(emb[0], native['items'][name]):.3f}", |
| flush=True) |
| del student, adapter |
| empty_cache() |
|
|
| pipe = build_pipeline(args.moss_dir) |
| pipe.to("cuda") |
| rows = [] |
| for name, prompt in AB: |
| adapter_img = generate(pipe, pos[name], neg, size=args.size, steps=args.steps, |
| seed=args.seed, guidance=args.guidance) |
| nat_emb = (native["items"][name], torch.ones(native["items"][name].shape[0], dtype=torch.bool)) |
| nat_img = generate(pipe, nat_emb, (native["items"]["negative"], |
| torch.ones(native["items"]["negative"].shape[0], dtype=torch.bool)), |
| size=args.size, steps=args.steps, seed=args.seed, guidance=args.guidance) |
| adapter_img.save(os.path.join(args.out_dir, f"{name}_adapter.png")) |
| nat_img.save(os.path.join(args.out_dir, f"{name}_native.png")) |
| rows.append((name, prompt, [nat_img, adapter_img])) |
| print(f" {name}: готово", flush=True) |
|
|
| out = grid(rows, os.path.join(args.out_dir, "ab_native_vs_adapter.png")) |
| print(f"готово: {out}, пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
