File size: 2,385 Bytes
e9190ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | """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()
|