| import sys,json,hashlib |
| from pathlib import Path |
| import argparse |
| PACKAGE=Path(__file__).resolve().parent.parent |
| sys.path.insert(0,str(PACKAGE)) |
| p=argparse.ArgumentParser() |
| p.add_argument('--model-dir',type=Path,required=True) |
| p.add_argument('--output',type=Path,required=True) |
| a=p.parse_args() |
| from qwen_image21.runtime import ImageEngine |
| import torch |
| from PIL import Image |
| out=a.output;out.mkdir(parents=True,exist_ok=False) |
| e=ImageEngine(a.model_dir) |
| |
| captured={} |
| def capture_tensor(name,t): |
| if isinstance(t,tuple):t=t[0] |
| if not isinstance(t,torch.Tensor):return |
| captured[name]=t.detach().cpu().contiguous() |
| orig=e.pipeline.encode_prompt |
| def encode(*args,**kwargs): |
| r=orig(*args,**kwargs);capture_tensor('prompt',r[0]);return r |
| e.pipeline.encode_prompt=encode |
| e.pipeline.vae.decoder.register_forward_pre_hook(lambda m,a: capture_tensor('vae_input',a[0])) |
| e.pipeline.vae.decoder.register_forward_hook(lambda m,a,o: capture_tensor('vae_output',o)) |
| prompt='A blue ceramic teapot on a wooden table, soft daylight, studio photograph' |
| records=[];control=None;image=None |
| for name,mode in [('a0','text-to-image'),('a1','text-to-image'),('b_rgba','rgba'),('a2','text-to-image'),('b_edit','edit'),('a3','text-to-image'),('a4','text-to-image')]: |
| captured={} |
| if mode=='text-to-image':p=prompt;seed=51;im=None |
| elif mode=='rgba':p='A cute cartoon dragon sticker.';seed=52;im=None |
| else:p='Change the blue teapot to bright red. Keep the table and composition.';seed=53;im=image |
| result,png,metrics=e.generate(p,width=1024,height=1024,seed=seed,mode=mode,image=im) |
| (out/(name+'.png')).write_bytes(png) |
| if image is None:image=result |
| row={'name':name,'sha':hashlib.sha256(png).hexdigest(),'metrics':metrics} |
| if mode=='text-to-image': |
| if control is None:control=captured |
| row['stages']={} |
| for key,t in captured.items(): |
| diff=(t.float()-control[key].float()).abs() |
| row['stages'][key]={'exact':torch.equal(t,control[key]),'max':diff.max().item(),'mae':diff.mean().item()} |
| records.append(row);print(json.dumps(row),flush=True) |
| (out/'result.json').write_text(json.dumps(records,indent=2)+'\n') |
|
|
| assert all(row['stages'][stage]['exact'] for row in records if 'stages' in row |
| for stage in ('prompt','vae_input')), 'Request conditioning or denoising state leaked across requests' |
| print('PASS: exact prompt embeddings and final denoising latents across immediate and A-B-A repeats') |
|
|
