import sys,time,json,argparse
from pathlib import Path
import torch
from PIL import Image
from safetensors.torch import save_file
from nvfp4_runtime import load_pipeline
from acceleration import accelerate_pipeline
ROOT=Path(__file__).resolve().parent
OLD=Path('/home/user/.local/share/rtx-pro-apps/qwen-image-2.1-fp8')
sys.path.insert(0,str(OLD));from prompts import EVALUATION
ap=argparse.ArgumentParser();ap.add_argument('--indices',default='all');ap.add_argument('--out',default='evaluation');ap.add_argument('--quant',default='release/transformer');ap.add_argument('--eager',action='store_true');args=ap.parse_args()
out=ROOT/args.out;out.mkdir(exist_ok=True)
@torch.inference_mode()
def main():
 pipe=load_pipeline('/home/user/models/qwen-image-2.1-b3179ad',ROOT/args.quant)
 if not args.eager:accelerate_pipeline(pipe)
 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')
  assert torch.isfinite(latest['latents']).all()
  save_file({'latents':latest['latents'].cpu().contiguous()},str(out/f'{i:02d}-latents.safetensors'))
  row={'index':i,'width':w,'height':h,'seconds_including_any_compilation':sec,'seed':20000+i};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(OLD/'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('EVALUATION_COMPLETE',flush=True)
main()