import json,time,sys
from pathlib import Path
import torch
ROOT=Path(__file__).parent;APP=ROOT.parent
sys.path.insert(0,str(APP/'release'))
from fp8_runtime import load_pipeline
from acceleration import accelerate_pipeline
from fast_decode import enable
PROMPT='A small corgi puppy sitting upright, full body, soft detailed fur, isolated on a transparent background, no shadow, no text.'

@torch.inference_mode()
def main():
    p=accelerate_pipeline(load_pipeline('/home/user/models/qwen-image-2.1-b3179ad',APP/'release/transformer'))
    records=[]
    def run():
        torch.cuda.synchronize();start=time.perf_counter()
        out=p(prompt=PROMPT,width=1024,height=1024,num_inference_steps=40,generator=torch.Generator('cuda').manual_seed(71001),output_type='latent').images
        torch.cuda.synchronize();return time.perf_counter()-start,out.clone()
    for mode in ['baseline','metadata','graph']:
        if mode!='baseline':enable(p,graph=mode=='graph')
        warm,_=run(); times=[]
        for i in range(3):
            sec,y=run();times.append(sec)
            if mode=='baseline':ref=y
            print(mode,i,sec,flush=True)
        row={'mode':mode,'warmup_seconds':warm,'times':times,'latent_equal':torch.equal(y,ref),'latent_nrmse':((y.float()-ref.float()).square().mean()/ref.float().square().mean()).sqrt().item()}
        records.append(row);(ROOT/'decode-trial.json').write_text(json.dumps(records,indent=2));print(row,flush=True)
main()