source/heldout.py

"""Fresh prompts frozen after selecting the dynamic NVFP4 candidate."""
import json,time
from pathlib import Path
import torch
from safetensors.torch import save_file
from nvfp4_runtime import load_pipeline as load_quant
from fp8_runtime import load_pipeline as load_base
from acceleration import accelerate_pipeline
ROOT=Path(__file__).parent
CASES=[
 ('A transparent glass teapot filled with amber tea on a slate table beside sliced dragon fruit, soft window light, crisp reflections, studio photograph, no text.',1024,1024),
 ('Three origami cranes arranged in a row on a pale wooden desk: a red crane on the left, a yellow crane in the center, and a blue crane on the right, precise folded paper, no text.',1024,1024),
 ('An elderly pianist playing a black grand piano in a warmly lit room, both hands visible on the keys, realistic fingers, candid documentary photograph, no text.',1024,1024),
 ('A minimalist travel poster with the exact large headline "SUMMER 2026", a golden sun above a turquoise sea, elegant bold typography.',1024,1024),
 ('一张精美的中国山水海报，清晰准确的四字标题“山海之间”，远山、碧海和细腻的水墨纹理，优雅留白。',1024,1024),
 ('An intricately engraved brass mechanical dragon sculpture on a dark pedestal, delicate interlocking gears, polished metal highlights, museum product photograph, no text.',2048,2048),
 ('This is an RGBA image with transparency. A charming illustrated red panda holding a small green bamboo leaf, clean outlines, fluffy striped tail. The image has alpha channel and the background is transparent.',1024,1024),
 ('A close-up wildlife photograph of a barn owl on a weathered wooden fence, fine speckled feathers, sharp dark eyes, softly blurred spring meadow in the background, no text.',1024,1024),
]
@torch.inference_mode()
def main():
 (ROOT/'heldout-manifest.json').write_text(json.dumps({'frozen_after_candidate_selection':True,'cases':[{'index':i,'prompt':p,'width':w,'height':h,'seed':62000+i,'steps':40} for i,(p,w,h) in enumerate(CASES)]},indent=2,ensure_ascii=False))
 for mode in ['bf16','nvfp4']:
  pipe=load_base('/home/user/models/qwen-image-2.1-b3179ad') if mode=='bf16' else accelerate_pipeline(load_quant('/home/user/models/qwen-image-2.1-b3179ad',ROOT/'release-dynamic/transformer'))
  out=ROOT/('heldout-'+mode);out.mkdir(exist_ok=True)
  for i,(prompt,w,h) in enumerate(CASES):
   latest={}
   def cb(p,s,t,kw):
    if s==39:latest['latents']=kw['latents'].detach()
    return kw
   im=pipe(prompt=prompt,width=w,height=h,num_inference_steps=40,generator=torch.Generator('cuda').manual_seed(62000+i),callback_on_step_end=cb).images[0]
   im.save(out/f'{i:02d}.png');save_file({'latents':latest['latents'].cpu().contiguous()},str(out/f'{i:02d}-latents.safetensors'))
   print('HELDOUT',mode,i,flush=True)
  del pipe;torch.cuda.empty_cache()
main()