MinhNH232331M/FlowUpscaler-diffusers

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)

This repo is in diffusers format: FlowUpscalerPipeline is a DiffusionPipeline whose components are the UNet (unet/upscaler_unet.py, a diffusers ModelMixin, weights in unet/diffusion_pytorch_model.safetensors), a shared Flux.2 VAE, and a FlowMatchEulerDiscreteScheduler — the plain-Python equivalent of the ComfyUI node, loadable with DiffusionPipeline.from_pretrained. The VAE is intentionally not bundled (it is shared with any Flux.2 checkpoint, and FLUX.2-dev is gated), so it is passed in at load time; trust_remote_code=True is required because the pipeline and UNet classes are loaded from this repo. The pipeline 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 (__call__): 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 (unet/upscaler_unet.py, detailed in the next section), 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, DiffusionPipeline
from PIL import Image

vae = AutoencoderKLFlux2.from_pretrained(  # any Flux.2 checkpoint's VAE
    "black-forest-labs/FLUX.2-dev", subfolder="vae", torch_dtype=torch.bfloat16
)
pipeline = DiffusionPipeline.from_pretrained(
    "MinhNH232331M/FlowUpscaler-diffusers",
    vae=vae, torch_dtype=torch.bfloat16, trust_remote_code=True,
).to("cuda")

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

Two compatibility paths are kept. flow_upscaler.safetensors at the repo root is the same flat checkpoint the ComfyUI node consumes, and FlowUpscalerPipeline.from_single_file("flow_upscaler.safetensors", vae, dtype=...) replicates the pre-diffusers constructor for local checkouts. The seed-based pipeline.upscale_image(image, num_passes, seed=42, ...) method also remains, and a given seed reproduces the pre-conversion outputs bit-for-bit (__call__ takes a generator instead, diffusers-style).

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.

Pipeline call knobs: num_passes (2x per pass), generator (RNG for the per-pass noise; upscale_image takes a seed instead), color_fix (master switch for both correction stages), the per-stage transplant_strength / color_fix_strength sliders described below, and output_type ("pil" default, "np", "pt"). transplant_strength also acts as a fidelity/realism control — lower values track the input more faithfully, higher values give the model's crisper rendition.

UNet inference optimizations (unet/upscaler_unet.py)

The network is attention-free, so compute scales linearly with area — but the stock forward's memory did not scale gracefully: at large latents each residual block allocated 4-8 full-resolution temporaries (GroupNorm outputs, SiLU copies, FiLM scale/shift maps, conv workspaces), and the skip-concat up-block briefly held several 768-channel full-res tensors. unet/upscaler_unet.py now layers four inference-only optimizations on the same weights. All of them are gated on torch.is_grad_enabled(), so training behavior is bit-for-bit untouched.

Measured on an NVIDIA L4 (23 GB), 3840x2160 -> 7680x4320 (one 2x pass at a 960x540x32 latent; one full-res 384-channel activation = 0.37 GiB in bf16):

Configuration	UNet-pass peak	Output change vs previous row
fp32 forward (training dtype)	8.53 GiB	—
bf16 forward	4.74 GiB	54.5 dB PSNR (visually identical)
+ in-place ops, early frees	4.28 GiB	bit-identical
+ split-concat up-block	3.51 GiB	fp reduction order only (43.6 dB)
+ chunked executor	2.50 GiB	bit-identical

1. bf16 execution. The model was trained in fp32, but bf16 matches it to ~54 dB PSNR while halving peak memory and roughly doubling forward speed. The sinusoidal timestep embedding stays fp32 and is cast to the projection weight dtype (ConditioningEncoder.forward), which is what makes reduced precision work at all. The pipeline default remains fp32; pass dtype=torch.bfloat16.
2. In-place ops. SiLU, FiLM scale/shift, and residual adds mutate their operands at inference (backward would need the pre-op values, hence the grad gate). Conditioning tensors are freed the moment their FiLM layer has consumed them instead of surviving the whole forward.
3. Split-concat up-block. The first block after each skip concat is the network's memory peak: GN/SiLU/conv on cat([x, skip]) at double width. Because the GroupNorm group count (32 over 768 channels) aligns with the concat boundary (channel 384 = group 16), the block decomposes exactly: GN and SiLU apply per half, and the 3x3/1x1 convs over the concat are the sums of per-half convs. No 768-channel tensor is ever materialized; the only numeric change is floating-point summation order.
4. Chunked inference executor — the main mechanism, engaged when the latent is at/above CHUNK_MIN_PIXELS (512x512, i.e. a 4096px output). Ops execute one at a time over resident full tensors, with only strip-sized working memory (CHUNK_STRIP_PIXELS, ~128k px per strip):
   • 3x3/1x1 convs run in row strips with a 1px halo into preallocated outputs (bit-exact vs the full conv); each block's final conv accumulates directly into its residual buffer.
   • FiLM applies strip-wise in place, never materializing the full-res 2x-channel scale/shift map (0.74 GiB at an 8K-UHD latent).
   • GroupNorm stays the native full-tensor op. Its statistics are global by construction — the seam/color-shift failure mode of naive spatial tiling cannot occur — and since its full-size output just fills a live-set slot the neighboring convs already require, keeping it native costs no memory. It also keeps the input pristine, so the input doubles as the residual for identity-skip blocks. The net effect: each block's live set is capped at 3-4 activation-sized tensors, and the chunked path's output is array_equal with the standard inference path — there is no quality trade-off to track.

Execution paths. Every module picks between three code paths at call time, so no configuration is needed:

Path	When	Ops
Training	grad enabled	original out-of-place ops, upstream-identical
Unchunked inference	grad off, latent < CHUNK_MIN_PIXELS	whole-tensor ops with optimizations 1-3 (in-place, split-concat)
Chunked inference	grad off, latent >= CHUNK_MIN_PIXELS (512x512, i.e. 4096px output)	optimization 4

The unchunked path exists because below the threshold the entire transient set is small (one activation at a 2048px output is 48 MB), so strip loops would add per-strip kernel-launch overhead with nothing to save. Since both inference paths are bit-identical, the threshold is purely a memory/latency trade-off — never a quality one. Measured at the crossover (2048 -> 4096px, the first size that chunks): 1.65 GiB chunked vs 1.93 GiB forced-unchunked at equal speed (~2.0 s end-to-end); a 1024 -> 2048px pass runs unchunked at 1.05 GiB total, dominated by the VAE ends rather than the UNet. At 8K the unchunked path would peak at 3.51 GiB (the split-concat row of the table above) versus 2.50 GiB chunked, with the gap widening quadratically at larger outputs.

Memory model. Peak VRAM ~= 4 x one full-res activation + ~1 GiB of non-scaling base (weights, low-res conditioning stages, noise latent, strip temporaries), where one activation is (H/8)(W/8) x 384 x 2 bytes in bf16. Chunking bounds the per-op overhead, not the total: GroupNorm's global stats force each op's full input and output to exist between sync points, so the resident term still scales with latent area. Measured end-to-end with the MageFlow VAE on the L4:

Output	UNet pass	VAE decode	Wall (total)
8K UHD (from 4K)	2.50 GiB	1.8 GiB	~6 s
16K UHD (from 8K)	7.63 GiB	5.5 GiB	~25 s

MAX_OUTPUT_SIDE = 16384 is calibrated from these numbers. If the upscaler shares a GPU with resident diffusion pipelines, setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True before torch initializes CUDA reduces allocator fragmentation (~0.6 GiB reserved-memory savings measured).

Two further tiers exist if flat, resolution-independent memory is ever needed, both unimplemented: accumulating GN statistics strip-wise and recomputing convs in a second pass (~2 resident tensors, ~2x conv compute), or streaming the resident tensors through pinned CPU memory (peak becomes tile-sized, ~2 s of PCIe traffic at 16K).

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%:

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:

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

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