File size: 2,472 Bytes
f5b159e | 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 | 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)
# Keep snapshots on CPU, outside timed qualification; trace request boundaries.
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')
|