1import torch
2from diffusers.models import AutoencoderKLFlux2
3
4from capacitor_decoder import CapacitorDecoder, CapacitorDecoderInferenceConfig
5
6
7def flux2_patchify_and_whiten(
8 latents: torch.Tensor,
9 vae: AutoencoderKLFlux2,
10) -> torch.Tensor:
11 b, c, h, w = latents.shape
12 if h % 2 != 0 or w % 2 != 0:
13 raise ValueError(f"Expected even FLUX.2 latent grid, got H={h}, W={w}")
14 z = latents.reshape(b, c, h // 2, 2, w // 2, 2)
15 z = z.permute(0, 1, 3, 5, 2, 4).reshape(b, c * 4, h // 2, w // 2)
16 mean = vae.bn.running_mean.view(1, -1, 1, 1).to(device=z.device, dtype=torch.float32)
17 var = vae.bn.running_var.view(1, -1, 1, 1).to(device=z.device, dtype=torch.float32)
18 std = torch.sqrt(var + float(vae.config.batch_norm_eps))
19 return (z.to(torch.float32) - mean) / std
20
21
22device = "cuda"
23flux2 = AutoencoderKLFlux2.from_pretrained(
24 "BiliSakura/VAEs",
25 subfolder="FLUX2-VAE",
26 torch_dtype=torch.bfloat16,
27).to(device)
28decoder = CapacitorDecoder.from_pretrained(
29 "data-archetype/capacitor_decoder",
30 device=device,
31 dtype=torch.bfloat16,
32)
33
34image = ... # [1, 3, H, W] in [-1, 1], with H and W divisible by 16
35
36with torch.inference_mode():
37 posterior = flux2.encode(image.to(device=device, dtype=torch.bfloat16))
38 latent_mean = posterior.latent_dist.mean
39
40 # Default path: whiten in float32, then cast back to model dtype before decode.
41 latents = flux2_patchify_and_whiten(latent_mean, flux2).to(dtype=torch.bfloat16)
42 recon = decoder.decode(
43 latents,
44 height=int(image.shape[-2]),
45 width=int(image.shape[-1]),
46 inference_config=CapacitorDecoderInferenceConfig(num_steps=1),
47 )
