File size: 2,494 Bytes
1081be0 | 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 | import sys,time,json,argparse,collections
from pathlib import Path
import torch
from PIL import Image
from safetensors.torch import save_file
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from fp8_runtime import load_pipeline,CalibratedFP8Linear
from prompts import EVALUATION
ap=argparse.ArgumentParser();ap.add_argument('--mode',choices=['exact','sage'],default='exact');ap.add_argument('--dynamic',action='store_true');ap.add_argument('--cache',type=float,default=0);ap.add_argument('--indices',default='all');args=ap.parse_args()
out=Path(__file__).parent/(f'validation-{args.mode}-{args.cache}'+('-dynamic' if args.dynamic else ''));out.mkdir(exist_ok=True)
@torch.inference_mode()
def main():
pipe=load_pipeline('/home/user/models/qwen-image-2.1-b3179ad',ROOT/'release/transformer')
if args.dynamic:
from acceleration import accelerate_pipeline
accelerate_pipeline(pipe)
else:raise ValueError('This release validates --mode exact --dynamic only')
assert args.mode=='exact' and not args.cache
cache=None
rows=[]
indices=list(range(16)) if args.indices=='all' else [int(x) for x in args.indices.split(',')]
for i in indices:
prompt=EVALUATION[i];w=h=2048 if i%4==0 else 1024
if i==12:w,h=1536,864
if i==13:w,h=864,1536
latest={}
def cb(p,step,t,kw):
if step==39:latest['latents']=kw['latents'].detach()
return kw
torch.cuda.synchronize();t=time.perf_counter()
im=pipe(prompt=prompt,width=w,height=h,num_inference_steps=40,generator=torch.Generator('cuda').manual_seed(20000+i),callback_on_step_end=cb).images[0]
torch.cuda.synchronize();sec=time.perf_counter()-t;im.save(out/f'{i:02d}.png')
save_file({'latents':latest['latents'].cpu().contiguous()},str(out/f'{i:02d}-latents.safetensors'))
row={'index':i,'seconds':sec,'reused':cache.skips if cache else 0};rows.append(row);(out/'runs.json').write_text(json.dumps(rows,indent=2));print(json.dumps(row),flush=True)
if args.indices=='all':
for j,prompt in enumerate(['Replace the background with a blooming spring garden and preserve the animal.','Turn this room into a warm evening scene with lamps switched on, preserving its furniture.']):
im=Image.open(ROOT/'evaluation/bf16'/f'{[0,3][j]:02d}.png').resize((1024,1024))
result=pipe(prompt=prompt,image=im,width=1024,height=1024,num_inference_steps=40,generator=torch.Generator('cuda').manual_seed(21000+j)).images[0]
result.save(out/f'edit-{j}.png');print('EDIT',j,flush=True)
print('COMPLETE',flush=True)
main()
|