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')