File size: 4,499 Bytes
47ea602 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """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()
|