mlx-community/FlowUpscaler-bf16

TensorForger/FlowUpscaler converted to bf16 (118.6 MB). It is a ×2 latent upscaler for FLUX.2's latent space: a 59M-parameter attention-free U-Net (SDXL-style ResNet blocks, GroupNorm, FiLM conditioning from the low-resolution latents) that predicts a velocity, used in one rectified-flow step from noise. It was distilled from FLUX.2 klein-4B generations at 512 → 1024. Passes chain: two passes give ×4.

The weights and the reference code are released by their author under the Unlicense (public domain). This conversion is too; the original LICENSE is kept in this repo.

Conversion

• Every tensor of flow_upscaler.safetensors (sha256 91eb93b40179441e569e01fe04f20fd9b951e434f88c84dd0ba523f830a81839, fp32, 237 MB) cast to bf16. The key names and PyTorch layouts are unchanged (OIHW conv weights); consumers transpose at load.
• These weights are modified (cast to bf16) from the original.
• Measured against the PyTorch fp32 reference on a FLUX.2 klein-4B 1024² draw upscaled to 2048² and decoded with the exact FLUX.2 VAE: these bf16 weights in fp32 compute are 67.7 dB PSNR from the reference's image (latents relL2 2.8e-3); bf16 compute is 60.3 dB. The Swift port reads them in fp32 compute by default.

Use

The model works on bn-normalised, unpatchified FLUX.2 latents (32 channels at 1/8): patchify to 128 channels, subtract the VAE's bn.running_mean, divide by sqrt(bn.running_var + 1e-4), unpatchify. The packed latents the FLUX.2 and FLUX 3 klein DiTs produce are already normalised: unpatchify them and go.

noise = randn(B, 32, 2h, 2w)
x = noise - unet(sample=noise, timestep=[1000], latents_small=latents)   # one step, t = 1000 -> 0

Swift/MLX: xocialize/flux2-vae-mlx-swift (0.8.0+), FlowUpscaler:

let up = try FlowUpscaler.load(url: dir.appendingPathComponent("flow_upscaler.safetensors"))
let packed2x = up.upscalePacked(packed, passes: 1, seed: 0)       // [B,128,h,w] -> [B,128,2h,2w]
let image = vae.decodePackedLatents(packed2x, tile: .standard)    // the exact row-band decode

Parity of the Swift port against TensorForger's upscaler_unet.py (PyTorch CPU fp32, same noise and latents, MLX GPU fp32 with the exact conv route): cos 1.0000000, relL2 1.1e-6 on square, non-square and odd latent sizes; 126.9 dB on the decoded 2048² image.

Notes

• It has only seen 512 → 1024. A second pass (to 4×) mostly sharpens; it does not invent structure.
• Decode with an exact FLUX.2 VAE decode. diffusers' tiled FLUX.2 decode normalises each tile on its own and leaves seams and a colour cast at these sizes.
• Not a general photo upscaler: on photos it sits between RealPLKSR and VOSR2 and rewrites small text.