File size: 3,780 Bytes
1e65275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | #!/usr/bin/env python3
"""Инференс MOSS-Image без текстового энкодера 9B: 0.8B + адаптер вместо 15.6 ГБ.
Условие для DiT считается здесь и уходит в пайплайн официальными аргументами
(`prompt_embeds` / `negative_prompt_embeds`), поэтому 15.6 ГБ учителя не нужны вовсе.
CFG включён (guidance 3.0), негатив по умолчанию пустой — его тоже считает адаптер.
python3 src/example.py --prompt "a red fox in a snowy forest at dusk, 85mm" --out fox.png
python3 src/example.py --prompt '{"layout":[{"text":"SALE","box":[40,40,420,120]}]}' --out sale.png
"""
import argparse
import os
import sys
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from adapter_lib import load_adapter # noqa: E402
from moss_lib import (MOSS_DIR, STUDENT, Student, build_pipeline, condition, # noqa: E402
empty_cache, generate, load_tokenizer)
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEVICE = "cuda"
def main():
ap = argparse.ArgumentParser(description="MOSS-Image с адаптером 0.8B вместо энкодера 9B")
ap.add_argument("--prompt", required=True)
ap.add_argument("--negative", default="", help="негативный промпт (по умолчанию пустой, как в пайплайне)")
ap.add_argument("--out", default="out.png")
ap.add_argument("--adapter", default=next(
(p for p in (os.path.join(ROOT, "adapter.safetensors"),
os.path.join(ROOT, "out", "adapter_v1.safetensors")) if os.path.exists(p)),
os.path.join(ROOT, "adapter.safetensors")))
ap.add_argument("--moss-dir", default=MOSS_DIR)
ap.add_argument("--student", default=STUDENT)
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()
tok = load_tokenizer(args.moss_dir)
adapter, meta = load_adapter(args.adapter, device=DEVICE)
print(f"адаптер: {sum(p.numel() for p in adapter.parameters()) / 1e6:.1f}M, шаг "
f"{meta.get('step', '?')}, срезы {meta.get('student_layers')}", flush=True)
student = Student(args.student, [int(x) for x in str(meta["student_layers"]).split(",")])
pos = condition(adapter, student, tok, [args.prompt])[0]
neg = condition(adapter, student, tok, [args.negative])[0]
n_tok = pos[1].sum().item()
print(f"условие: {n_tok} токенов (негатив {neg[1].sum().item()})", flush=True)
if n_tok > int(meta.get("attn_max_len", 10 ** 6)):
raise SystemExit(f"в промпте {n_tok} токенов, а таблица позиций ветки внимания "
f"обучена до {meta['attn_max_len']} — сократите промпт")
del student, adapter # 0.8B и адаптер больше не нужны
empty_cache()
pipe = build_pipeline(args.moss_dir)
pipe.to(DEVICE)
print(f"DiT на карте, VRAM {torch.cuda.memory_allocated() / 2**30:.1f} ГиБ", flush=True)
image = generate(pipe, pos, neg, size=args.size, steps=args.steps,
seed=args.seed, guidance=args.guidance)
image.save(args.out)
print(f"{args.size}px, {args.steps} шагов, сид {args.seed} -> {args.out} {image.size}, "
f"пик VRAM {torch.cuda.max_memory_allocated() / 2**30:.1f} ГиБ", flush=True)
if __name__ == "__main__":
main()
|