code/rapid_moss/sample.py

"""Side-by-side evaluation: release teacher vs data-space generators (lodestones-format dirs).

    python -m rapid_moss.sample --model DIR --cache CACHE --prompts eval.jsonl --out OUT \
        --method teacher:50 --method dmd:/path/lodestones:8 --method dmd:/path/lodestones:4
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import torch
from safetensors import safe_open

from .model import load_base, sample_teacher
from .nets import Generator
from .textenc import Cache, read_prompts
from .train import decode


def load_generator(model, path, steps, device):
    with safe_open(Path(path) / 'lora.safetensors', 'pt') as f:
        key = next(k for k in f.keys() if k.endswith('.lora_down.weight'))
        rank = f.get_slice(key).get_shape()[0]
    gen = Generator(model, rank, rank, steps).to(device).requires_grad_(False)
    gen.load_lodestones(path)
    gen.lora.freeze_cache()
    return gen


def label_grid(columns, labels, ids, path, cell=512):
    from PIL import Image, ImageDraw
    pad = 28
    sheet = Image.new('RGB', (len(columns) * cell, len(ids) * (cell + pad)), 'white')
    draw = ImageDraw.Draw(sheet)
    for c, (images, label) in enumerate(zip(columns, labels)):
        for r, im in enumerate(images):
            im = Image.fromarray(im)
            im.thumbnail((cell, cell))
            y = r * (cell + pad)
            draw.text((c * cell + 6, y + 6), f'{label} | {ids[r]}', fill='black')
            sheet.paste(im, (c * cell + (cell - im.width) // 2, y + pad + (cell - im.height) // 2))
    sheet.save(path, quality=88)


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--cache', required=True)
    ap.add_argument('--prompts', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--method', action='append', required=True)
    ap.add_argument('--seed', type=int, default=1234)
    ap.add_argument('--limit', type=int, default=0)
    args = ap.parse_args()
    dev = torch.device('cuda')
    torch.backends.cuda.matmul.allow_tf32 = True
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    model = load_base(args.model, dev)
    from diffusers import AutoencoderKLFlux2
    vae = AutoencoderKLFlux2.from_pretrained(Path(args.model) / 'vae', torch_dtype=torch.float32,
                                             local_files_only=True).to(dev).eval()
    cache = Cache(args.cache)
    rows = read_prompts([args.prompts])
    rows = rows[:args.limit] if args.limit else rows
    neg = model.context(cache.get('').to(dev))
    columns, labels, timings = [], [], {}
    for spec in args.method:
        kind, *rest = spec.split(':')
        if kind == 'teacher':
            steps = int(rest[0]) if rest else 50
            cfg = float(rest[1]) if len(rest) > 1 else 3.0
            tag, gen = f'teacher{steps}-cfg{cfg:g}', None
        elif kind == 'dmd':
            path, steps = rest[0], int(rest[1])
            tag = f'{Path(path).name}-{steps}step'
            gen = load_generator(model, path, steps, dev)
        else:
            raise SystemExit(f'unknown method {spec}')
        (out / tag).mkdir(exist_ok=True)
        images, secs = [], []
        for row in rows:
            ctx = model.context(cache.get(row['id']).to(dev))
            g = torch.Generator('cpu').manual_seed(args.seed)
            z = torch.randn(1, model.tr.config.latent_channels, row['h'] // 16, row['w'] // 16, generator=g).to(dev)
            torch.cuda.synchronize()
            t0 = time.time()
            x0 = sample_teacher(model, z, ctx, neg, steps, cfg) if gen is None else gen.rollout(model, z, ctx, steps)
            torch.cuda.synchronize()
            secs.append(time.time() - t0)
            images.append(decode(vae, x0))
            from PIL import Image
            Image.fromarray(images[-1]).save(out / tag / f"{row['id']}.png")
        del gen
        columns.append(images)
        labels.append(tag)
        timings[tag] = round(sum(secs) / len(secs), 3)
        print(json.dumps({'method': tag, 'dit_s_per_image': timings[tag]}), flush=True)
    (out / 'timings.json').write_text(json.dumps(timings, indent=1))
    ids = [r['id'] for r in rows]
    for start in range(0, len(ids), 6):
        label_grid([c[start:start + 6] for c in columns], labels, ids[start:start + 6],
                   out / f'grid-{start // 6:02d}.jpg')


if __name__ == '__main__':
    main()