scripts/audit_vae_decode.py

"""Decode the exact saved latent with one VAE setting at a time (no denoising).

Use benchmark_edits_v2 --save-latents. Saved latents are trusted local tensors.
This is diagnostic evidence, not a replacement benchmark image.
"""
import argparse
import inspect
import json
import time
from pathlib import Path

import torch
from diffusers import AutoencoderKLQwenImage21
from diffusers.image_processor import VaeImageProcessor
from scripts.integrity import sha256
from scripts.editing_protocol import image_diagnostics, save_diagnostics


def main():
    ap = argparse.ArgumentParser(__doc__)
    ap.add_argument('--model', required=True)
    ap.add_argument('--latents', required=True)
    ap.add_argument('--output', required=True)
    ap.add_argument('--dtype', choices=['bf16', 'fp32'], default='bf16')
    ap.add_argument('--tiling', action='store_true')
    args = ap.parse_args()
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=False)
    dtype = torch.bfloat16 if args.dtype == 'bf16' else torch.float32
    vae = AutoencoderKLQwenImage21.from_pretrained(args.model, subfolder='vae',
                                                 torch_dtype=dtype, local_files_only=True).to('cuda')
    if args.tiling:
        vae.enable_tiling()
    latents = torch.load(args.latents, weights_only=True, map_location='cpu').to('cuda', dtype=dtype)
    torch.cuda.synchronize()
    started = time.perf_counter()
    with torch.inference_mode():
        decoded = vae.decode(latents, return_dict=False)[0][:, :, 0]
        image = VaeImageProcessor(vae_scale_factor=16).postprocess(decoded, output_type='pil')[0]
    torch.cuda.synchronize()
    seconds = time.perf_counter()-started
    image.save(out/'decoded.png')
    save_diagnostics(image, out, 'decoded')
    record = dict(dtype=args.dtype, tiling=args.tiling, seconds=seconds,
                  latents_sha256=sha256(args.latents), image_sha256=sha256(out/'decoded.png'),
                  vae_source_sha256=sha256(inspect.getfile(type(vae))),
                  vae_weights=[dict(path=p.name, sha256=sha256(p)) for p in
                               sorted((Path(args.model)/'vae').glob('*.safetensors'))],
                  **image_diagnostics(image))
    (out/'record.json').write_text(json.dumps(record, indent=2), encoding='utf-8')
    print(json.dumps(record, indent=2))


if __name__ == '__main__':
    main()