--- 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: diffusers base_model: - TensorForger/FlowUpscaler - black-forest-labs/FLUX.2-klein-4B 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`) 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. ```python 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](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. 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%: ![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)