README.md

metadata

license: unlicense
tags:
  - flux2
  - flux
  - image-to-image
  - image-upscaling
  - latent-upscaling
  - super-resolution
  - diffusion
  - flow-matching
  - rectified-flow
  - generative-models
  - image-generation
  - pytorch
library_name: pytorch
base_model:
  - TensorForger/FlowUpscaler
pipeline_tag: image-to-image

Flow Upscaler

Flow Upscaler is a fast latent upscaler model that works in the Flux.2 latent space.

Under the hood, it is a lightweight Rectified Flow model with 59M parameters that generates upscaled latents in a single denoising step.

Inference pipeline (flow_upscaler_pipeline.py)

flow_upscaler_pipeline.py is a standalone, diffusers-based image-in / image-out wrapper around the model — the plain-Python equivalent of the ComfyUI node. FlowUpscalerPipeline combines the UNet defined in upscaler_unet.py (weights: flow_upscaler.safetensors) with a shared Flux.2 VAE and runs the full loop:

1. Encode (encode_image): the input image is resized to a multiple of 16, VAE-encoded, and the 32-channel latents are normalized with the VAE's BatchNorm running stats — the same convention the model was trained with.
2. Upscale (upscale_latents): pure noise is denoised into the 2x latent in a single FlowMatchEuler step, conditioned on the small latent. Passes chain for 4x/8x because a pass's output lives in the same normalized latent space as its conditioning input.
3. Decode + color fix (upscale_image): latents are denormalized and VAE-decoded. With the default color_fix=True, each pass ends with a latent low-frequency transplant and the decoded image gets a low-band L/a/b transplant from the input — the two-stage drift correction detailed in the next section.

Memory is bounded for large outputs: the VAE switches to tiling above 2048px, the UNet runs through a chunked inference executor (upscaler_unet.py), and outputs are capped at MAX_OUTPUT_SIDE = 16384 px per side (excess passes are dropped with a warning). This keeps even 16K outputs within a 24GB-class GPU in bf16.

import torch
from diffusers import AutoencoderKLFlux2  # any Flux.2 checkpoint's VAE
from PIL import Image

from flow_upscaler_pipeline import FlowUpscalerPipeline

vae = AutoencoderKLFlux2.from_pretrained(
    "black-forest-labs/FLUX.2-dev", subfolder="vae", torch_dtype=torch.bfloat16
).to("cuda")
pipeline = FlowUpscalerPipeline("flow_upscaler.safetensors", vae, dtype=torch.bfloat16)

image = Image.open("input.png")
result = pipeline.upscale_image(image, num_passes=2)  # 2 passes = 4x
result.save("output.png")

For the VAE, consider MageFlow-VAE-diffusers — an efficient drop-in replacement for the Flux.2 VAE that operates in the same latent space (~10x faster encode, ~6x faster decode, roughly half the peak VRAM, with a chunked decoder that keeps memory near-constant at large resolutions). Since upscaling cost here is dominated by the encode/decode ends, it speeds up the whole loop.

upscale_image knobs: num_passes (2x per pass), seed, color_fix (master switch for both correction stages), and the per-stage transplant_strength / color_fix_strength sliders described below. transplant_strength also acts as a fidelity/realism control — lower values track the input more faithfully, higher values give the model's crisper rendition.

Color drift and corrections (this repo's pipeline)

Chained single-step passes drift measurably. The local wrapper (flow_upscaler_pipeline.py, used by app_local.py) measures and corrects three drift modes, all enabled by the default color_fix=True:

Drift mode	Magnitude (uncorrected)	Cause
Per-channel latent mean drift	~0.1 sigma/pass -> 2-5/255 RGB cast after 2 passes	single-step Euler DC error (VAE round-trip is clean)
Chroma amplification	~8%/pass mean chroma; p99 +14/255 after 3 passes, concentrated on already-colorful regions (lips, skin)	model bias, rides on high-band latent content through the nonlinear decoder
Tone-curve stretch	dL_std +1.1, shadows -2.75, highlights +2 after 3 passes	side effect of the latent transplant below (raw model output is slightly flat)

Two-stage correction:

1. Latent low-frequency transplant (_lowfreq_transplant), after every pass: the output latent's low band (antialiased downsample by LOWFREQ_FACTOR = 8) is replaced with the conditioning's. Fixes the mean drift and keeps the next pass's conditioning on-manifold, which measurably sharpens results. Cost: ~1.5 ms / +33 MiB at an 8K-UHD output. transplant_strength interpolates this stage from off (softer, truest to input structure) to full (default, crisper); color is unaffected at any setting. Beyond drift correction, this knob doubles as a fidelity/realism control: lower values stay measurably closer to the input (PSNR to input 43.9 dB at 0.0 vs 38.4 dB at 1.0 in the 2-pass test below), higher values favor the model's sharper, more confident rendition.
2. Pixel-space low-band L/a/b transplant (_lab_stat_match), after the final decode: Lab diff maps computed at the input's resolution (area-downsampled output vs input, fp32), bilinearly upsampled and added to the decoded image in fp16 row strips on GPU. Below the input's Nyquist the input is ground truth, so regional color and tone are pinned exactly while the model keeps all detail above it. Cost incl. uint8 download: 13 / 54 / 235 ms and +51 / +128 / +404 MiB transient VRAM at 2K / 4K / 8K. color_fix_strength interpolates this stage linearly; residual drift scales with 1 - strength.

Measured results

Protocol: output area-downsampled (INTER_AREA) back to input size and compared against the input in Lab; worst cases from a 512px-base photo set, 2-3 pass chains, seed 42, NVIDIA L4, bf16. Reproduce with test_flow_upscaler.ipynb (one section per design choice, with plots).

Metric (at input scale)	Uncorrected	Both fixes on
RGB mean cast, 2 passes	up to 5/255 per channel	<= 0.05/255
Mean chroma (Lab C_mean)	+7...+15%	within +/-0.5%
Chroma p99 (worst case, 3 passes)	+13.8/255	+0.6/255
Tone contrast (L_std), 3 passes	+1.08	+0.03
Shadows p5 / highlights p95	-2.75 / +1.96	0.00 / 0.00

Knob response (2-pass restore-photo test, other knob at its 1.0 default — the two stages are orthogonal by construction and by measurement):

transplant_strength	1.0	0.5	0.0
chroma error dC	-0.02	-0.02	-0.02
sharpness (mean abs Laplacian)	7.94	4.86	3.28
PSNR to input (dB)	38.4	41.7	43.9

color_fix_strength	1.0	0.5	0.0
chroma error dC	-0.02	+0.53	+1.16
sharpness	7.94	7.98	8.07

Compounding is per pass: before the pixel-space fix, mean chroma ran +8.6% after 1 pass and +14.8% after 2; the latent mean drift is ~0.1 sigma per pass on average with 0.35-0.48 sigma outlier channels. The VAE round-trip itself is clean (chroma -0.02, RGB means <= 0.9/255) — every drift mode above is created by the flow UNet or by the correction stack itself.

Nyquist / frequency-band analysis

Band bookkeeping, with frequencies written as fractions of each grid's Nyquist: one latent pixel decodes to 8 image pixels, so the latent grid's Nyquist sits at 1/16 cycles/px in pixel space; a conditioning at half the output resolution owns everything below 0.5x the output-latent Nyquist. Each correction stage operates in the widest band where its reference is ground truth and its domain assumption holds:

• Stage 1 corrects below output-latent Nyquist / LOWFREQ_FACTOR (= /8). The limiting factor is not the conditioning (ground truth up to 0.5) but how faithfully latent-space bands map to pixel-space bands — see below.
• Stage 2 corrects below the input image's Nyquist, in pixel space, where the low band is the pixel-space low band by definition — no mapping assumption needed. Everything above the input Nyquist is model-invented detail with no reference, and is deliberately left untouched.

The latent->pixel mapping was measured directly on the Mage decoder: inject a band-limited 0.1-sigma perturbation into one latent spectral band, decode, and integrate the luma response spectrum per pixel band (bands as fractions of the latent Nyquist; pixel content above 1.0 is finer than the latent grid can represent spatially):

Latent band	-> same pixel band	-> finer than latent Nyquist
0 - 1/16	87.7%	7.8%
1/16 - 1/8	90.6%	4.2%
1/8 - 1/4	77.1%	15.0%
1/4 - 1/2	64.7%	26.2%
1/2 - 1	38.0%	58.1%

Three design rules fall out of this table:

1. Latent band surgery is sound only in the top two rows (~88-91% correspondence) — exactly where LOWFREQ_FACTOR 8-16 operates. Below 77% the spliced content scatters across unrelated pixel frequencies and reads as ghosting (matches the failed wide-band experiments below).
2. The ~8% upward leakage from the lowest band is the latent transplant's sharpening / tone-stretch side effect, quantified — the energy it injects into finer-than-latent-Nyquist pixel detail.
3. Energy fractions understate statistical bias: high-band latent chroma texture rectifies through the nonlinear decoder into a mean-saturation shift carrying almost no low-band energy (the mechanism behind drift mode 2). Exact guarantees therefore only come from the pixel-space stage.

ComfyUI Node

Features:

• Upscaling from 512x512 to 1024x1024 takes 8ms*
• The model is trained for 2X upscaling, but multiple passes can be chained to reach up to 8K resolution
• A full pipeline with Flux generation, upscaling to 8K, and decoding runs in just 25 seconds (on RTX 5090)
• The training process uses Flow Distillation with Flux.2 as a teacher, forcing the model to learn strong image semantics

*On RTX 5090, in latent space, without decoding, see benchmark here.

Here is one 4X upscaled image (two passes):

How it works

Architecturally, Flow Upscaler is a U-Net with SDXL-style ResNet blocks. It takes a noisy sample as input and predicts velocity as output. The generation process happens directly in high-resolution latent space.

The low-resolution latents are passed through a separate conditioning encoder that produces control signals, which are injected into the main U-Net encoder using FiLM conditioning.

No attention layers are used, so compute scales linearly with image area. This makes generation at 8K resolution possible.

The model is trained using Flow Distillation with Flux.2-klein-4B as a teacher. We generated 20K diverse images with Flux, storing the initial noise, generated latents, and downscaled latents used for conditioning.

The downscaled latents are created by decoding high-resolution latents, downscaling them in pixel space, and encoding them back into latents. Direct latent downscaling introduces artifacts and breaks latent patterns, resulting in blurry decoded images.

Training code

If you want to explore the training code or use the model outside ComfyUI, see:

notebooks/flow_upscaler in https://github.com/tensorforger/CTGMWorkshop