---
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](https://bfl.ai/models/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.

```python
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](https://huggingface.co/MinhNH232331M/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.

### Visual demo

One-image walkthrough of the two knobs (`demo_images/`): a 320x480 portrait
upscaled 2 passes (4x) at seed 42.

**Color fix** — both runs at `transplant_strength = 0` so the output
structure matches the original and only color differs. Uncorrected, the skin
drifts visibly pale/green (R −3.1, G −0.5, B −2.8 /255 mean cast; mean
chroma +3.5%); with the fix, means pin to within 0.2/255 and chroma to +0.1%:

![color fix comparison](demo_images/compare_color_fix.png)

**Latent transplant** — both runs at `color_fix_strength = 1`. Strength 1
nearly doubles measured sharpness (mean abs Laplacian 1.40 -> 2.71) at the
cost of input fidelity (PSNR to input 43.1 -> 40.8 dB) — the
fidelity/realism tradeoff, visible in the lashes and iris detail:

![latent transplant comparison](demo_images/compare_latent_transplant.png)

Inputs and full outputs for all three settings are in `demo_images/`. Color
metrics for this demo were measured through the gradio demo's lossy webp
output, which inflates the residual percentiles slightly; the tables above
use the lossless notebook protocol.

### 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](https://github.com/TensorForger/comfyui-flow-upscaler)**

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](https://github.com/tensorforger/CTGMWorkshop).

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

![comparison](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/upscaler_comparison.png)

## 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.

![example](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/flow_upscaler_architecture.PNG)

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.

![example](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/flow_upscaler_training_approach.PNG)


## 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](https://github.com/tensorforger/CTGMWorkshop)