| |
| """Инференс 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 |
| from moss_lib import (MOSS_DIR, STUDENT, Student, build_pipeline, condition, |
| 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 |
| 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() |
|
|
