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