Extensions/flux2-resolution-guard
ComfyUI Extension

flux2-resolution-guard

FLUX.2 high-resolution drift correction and ComfyUI node pack

By xmarre·Created 5 months ago·Updated 5 months ago· 2
xmarre/FLUX.2-Resolution-Guard
Nodes—
On cloudLocal install
Stars2
Updated5 months ago
Readme

FLUX.2 Resolution Guard

A research-first correction module and ComfyUI node pack for FLUX.2 high-resolution drift control.

This repository is built around one practical observation:

  • FLUX.2 edits and generations can become less geometry-stable and less color-stable as effective resolution climbs.
  • The symptoms often travel together: slight outward "breathing"/expansion, low-frequency relighting, and washed-out chroma/contrast.
  • The problem is not exclusive to inpainting or crop-detailing. It can affect whole-image FLUX.2 outputs as well.
  • A useful correction strategy is to treat a lower-resolution FLUX.2 result as a more stable anchor manifold, then pull the high-resolution output back toward that manifold in a controlled way.

This repo implements that strategy in two forms:

  1. A trainable correction model called SMIC
    Stable-Manifold Inward Compander

  2. A training-free analytic fallback usable immediately in ComfyUI before trained weights exist.

The code is deliberately written so the same machinery can be used for:

  • whole-image FLUX.2 outputs
  • masked regions / FaceDetailer style crops
  • anchor-conditioned correction
  • anchorless correction using a learned prior or analytic companding

What this repo contains

Core package

src/flux2_resolution_guard/

  • models/smic.py
    Small correction network with:

    • low-rank warp head
    • RGB residual head
    • gate/confidence head

  • data/synthetic.py
    Synthetic pretraining dataset that creates:

    • outward radial drift
    • anisotropic expansion
    • low-frequency washout
    • mild seam stress
    • whole-image or masked perturbations

  • data/triplets.py
    Dataset loader for FLUX.2 self-distillation triplets:

    • original
    • anchor (stable lower-resolution FLUX.2 pass)
    • highres (problematic FLUX.2 pass)
    • mask
    • metadata

  • training/engine.py
    Complete training loop with:

    • pixel loss
    • low-frequency loss
    • seam loss
    • OKLab low-frequency loss
    • identity/no-op loss
    • warp smoothness regularization

  • inference/image.py
    High-level inference API:

    • load checkpoint
    • correct whole image
    • correct masked image
    • analytic fallback correction

ComfyUI node pack

Repo root contains a working custom-node package:

  • __init__.py
  • nodes.py

Nodes included:

  • FLUX2 RG Load Model
  • FLUX2 RG Apply Correction
  • FLUX2 RG Analytic Compand

These nodes work on standard ComfyUI IMAGE and MASK types.

Scripts

  • scripts/train_synthetic.py
  • scripts/train_triplets.py

Tests

  • tests/test_model.py
  • tests/test_inference.py
  • tests/test_datasets.py

Design goals

This project is aimed at FLUX.2 specifically, but not limited to one workflow shape.

It is designed to address resolution-linked drift in:

  • full-frame FLUX.2 generations
  • edited full-frame FLUX.2 outputs
  • inpainted/masked regions
  • FaceDetailer / crop-and-stitch workflows

The implementation starts in the image domain on purpose.

Why?

  • It is robust and immediately usable.
  • It avoids hard-coding brittle assumptions about specific FLUX.2 runtime internals.
  • It still supports FLUX.2-specific training by using FLUX.2 triplets as supervision.
  • It can later be extended to latent-side correction while preserving the same public interfaces.

Stable-Manifold Inward Companding

The core idea is simple:

  1. Use a lower-resolution FLUX.2 result as a more stable low-frequency anchor.
  2. Use the high-resolution FLUX.2 result for detail.
  3. Predict a smooth inward/companding warp and a residual correction.
  4. Restore low-frequency geometry and color stability without erasing high-frequency detail.

Mathematically, the model predicts:

  • a low-rank warp field
  • an RGB residual
  • a confidence gate

The output is:

corrected = warped(edit) + gate * residual

with strong regularization to keep the warp broad and controlled.

Installation

As a Python package

git clone https://github.com/yourname/flux2-resolution-guard.git
cd flux2-resolution-guard
pip install -e .

For ComfyUI

Clone or copy this repo into your ComfyUI custom nodes directory:

cd ComfyUI/custom_nodes
git clone https://github.com/yourname/flux2-resolution-guard.git

Restart ComfyUI.

Minimal Python inference example

from PIL import Image
from flux2_resolution_guard.inference.image import correct_image_with_checkpoint

edited = Image.open("highres.png").convert("RGB")
anchor = Image.open("anchor_1mp.png").convert("RGB")

corrected = correct_image_with_checkpoint(
    image=edited,
    checkpoint_path="checkpoints/smic_best.pt",
    anchor_image=anchor,
    mp_ratio=1.8,
    strength=1.0,
)
corrected.save("corrected.png")

Analytic fallback example

from PIL import Image
from flux2_resolution_guard.inference.image import analytic_compand_correction

edited = Image.open("highres.png").convert("RGB")
anchor = Image.open("anchor_1mp.png").convert("RGB")

corrected = analytic_compand_correction(
    image=edited,
    anchor_image=anchor,
    mp_ratio=1.8,
    strength=0.7,
)
corrected.save("analytic_corrected.png")

Training

1. Synthetic pretraining

Builds a correction prior from synthetic perturbations.

python scripts/train_synthetic.py \
  --image-dir /path/to/images \
  --output-dir runs/synthetic \
  --epochs 20 \
  --batch-size 4

2. FLUX.2 triplet training

Train on real FLUX.2 captures with a manifest file.

python scripts/train_triplets.py \
  --manifest data/flux2_triplets.json \
  --output-dir runs/flux2_triplets \
  --epochs 10 \
  --batch-size 2

Manifest format:

[
  {
    "original": "data/originals/example.png",
    "anchor": "data/anchors/example_anchor.png",
    "highres": "data/highres/example_highres.png",
    "mask": "data/masks/example_mask.png",
    "mp_ratio": 1.85
  }
]

ComfyUI nodes

FLUX2 RG Load Model

Inputs:

  • checkpoint_path
  • device

Returns:

  • RG_MODEL

FLUX2 RG Apply Correction

Inputs:

  • model
  • image
  • anchor_image (optional in practice; pass same image if unavailable)
  • mask (optional; defaults to full image)
  • mp_ratio
  • strength

Returns:

  • corrected IMAGE

FLUX2 RG Analytic Compand

Inputs:

  • image
  • anchor_image
  • mask
  • mp_ratio
  • strength
  • blur_sigma

Returns:

  • corrected IMAGE

Notes on current scope

This repository gives you:

  • a real trainable codebase
  • a real analytic fallback
  • a real ComfyUI node pack
  • tests and working interfaces

It does not ship with pretrained weights.

For best real-world performance you should train on your own FLUX.2 captures, especially:

  • whole-frame pairs across 1MP → highres
  • face crops
  • portrait work
  • the exact denoise/settings you care about

Suggested next experiments

  • Add latent-domain backend while keeping current API stable
  • Export trained checkpoints to .safetensors
  • Add face-landmark alignment losses
  • Add optional CLIP/identity preservation losses
  • Add whole-image low-frequency anchor blending directly inside the model

License

MIT