Skip to content

Community Pipeline: Add HRDiT (training-free high-resolution FLUX.1-dev) - #26

Open
smellslikeml wants to merge 19 commits into
mainfrom
flux-hrdit-community-pipeline
Open

Community Pipeline: Add HRDiT (training-free high-resolution FLUX.1-dev)#26
smellslikeml wants to merge 19 commits into
mainfrom
flux-hrdit-community-pipeline

Conversation

@smellslikeml

@smellslikeml smellslikeml commented Aug 14, 2026

Copy link
Copy Markdown
Owner

Staged on this fork for review; targets huggingface/diffusersexamples/community.

What

Adds examples/community/pipeline_flux_hrdit.pytraining-free high-resolution image generation for off-the-shelf diffusion transformers, implementing HRDiT (arXiv:2608.07003, official repo).

It generates natively at high resolution (up to 4096×4096) straight from a text prompt on stock FLUX.1-devno fine-tuning, no new weights, no separate upscaler model, and no src/diffusers changes. Users of the base checkpoint get 4K generation directly, instead of the usual generate-at-1K-then-upscale (or train-a-high-res-model) workaround.

This is high-resolution generation, not super-resolution of an existing image — it complements, rather than replaces, image upscalers.

Validation (FLUX.1-dev · A100-80GB · torch 2.11)

HRDiT high-res

Both 2048² and 4096² are coherent — no tiling, no washout. A 4096² generation runs in ~2 min at ~26 GB peak.

Runnable notebook: https://colab.research.google.com/drive/1AU6QNOGDMCPpGdocVyIyKpXXGt3hoWdo?usp=sharing

Method (on top of the stock FluxPipeline denoise loop)

  • NTK-aware RoPE scaling — the primary high-res mechanism; per-stage scaling of the RoPE base brings out-of-range positions back into the trained band, applied on every upscale step.
  • Spatial Position Alignment (SPA) — leading-steps-only nudge; a monotonic bundle coarsening of position ids (no wrapping → no periodic tiling), averaged over sliding variants inside attention.
  • Structure-guided progressive ladder (1024 → 2048 → 4096) — each stage decodes / upscales / re-encodes the previous latent as a structural prior, then injects its low-frequency band every step (Butterworth FFT split, weight alpha) with a velocity-momentum term (beta) to hold the coarse structure stable at the highest stage.

Out of scope (not needed for the results above; noted in the pipeline docstring): HAP attention pruning, swin_pachify, and DWT guidance.

Scope & generality

HRDiT is a DiT-general technique, not FLUX-specific. SPA operates on the token indices (i, j) fed to the model's positional function — which, per the paper, "though different mainstream DiT models may differ in their specific implementations… all share a common structure: f_pe takes token indices as inputs." So the architecture-agnostic parts of this pipeline — build_bundle_id_variants (SPA on token indices) and the structure-guided progressive ladder — carry over unchanged to other DiTs; only the per-model position scaling differs.

  • Applies to diffusion transformers. The paper's two evaluation targets are FLUX (this PR) and SD3; an SD3 instantiation would reuse this same SPA + structure-guidance core, swapping RoPE-base scaling for SD3's positional-embedding interpolation.
  • Not for U-Net Stable Diffusion (1.5 / 2.x / SDXL) — that space is already served by the DemoFusion / FreeScale community pipelines.
  • Prior-art note. The structure-guidance ladder follows a DemoFusion-style skip-residual; the DiT-specific contributions here are the NTK-aware position scaling and SPA.

Files

  • examples/community/pipeline_flux_hrdit.py — the pipeline
  • examples/community/README.md — usage entry
  • tests/others/test_community_pipeline_flux_hrdit.py — CPU-only unit tests (SPA variants, NTK RoPE, structure-guidance helpers; 13 passing)
  • benchmarks/benchmarking_flux_hrdit.py — naive FluxPipeline vs HRDiT timing / peak memory

Disclosure

Training-free / adds no weights. Initially drafted by Outrider (remyx) and then corrected and completed with AI assistance (Claude) against the reference implementation, with each iteration validated on GPU before merge.

github-actions Bot and others added 12 commits August 14, 2026 16:30
…ion (OOM: autograd graph + eager full-score materialization)
…ttention variant averaging + proportional scale

The first draft mis-ported SPA two ways, producing periodic-tiled mush worse than
naive high-res generation (caught on GPU validation @2048):
  - position map used modulo wrapping ((y-shift)%H)%bundle -> periodic tiling;
    the paper uses a monotonic bundle coarsening phi(x)=ceil((x+1-n1)/size).
  - averaging happened over 4 full model forwards at the output; the paper
    averages per-RoPE-variant attention *outputs* inside each layer (identity
    mean(softmax(A_n))@v == mean(softmax(A_n)@v)), one forward per step.
Also adds the reference's proportional attention scale and corrects group_num
default to 80. HAP pruning, NTK RoPE scaling and per-step SPA scheduling are
documented as follow-ups (not yet ported). Ref: https://github.com/zylwithxy/HRDiT

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…-res mechanism)

GPU validation of the SPA-only port: 2048 coherent but cross-hatched, 4096 washed
out. Root cause vs reference inference.py: SPA is a *leading-few-steps* nudge
(spa_steps=[3,0] -> 3 steps at 2048, none at 4096), while NTK-aware RoPE scaling
(theta*=ntk_factor, [4,10] per stage) applied on *every* step is the primary
high-res mechanism. Running SPA every step over-averaged -> washout.

  - flux_rope(): NTK RoPE (== diffusers FluxPosEmbed at factor 1).
  - _SPAState: base NTK rope on every step; SPA variants only while spa_active.
  - per-stage ntk_factor / spa_steps / guidance_scale_highres, gated per step.

Remaining (documented): frequency-domain structure guidance + HAP pruning.
Ref: https://github.com/zylwithxy/HRDiT (inference.py, hrdit/transformer.py)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… for the top stage

Adds the DemoFusion-style structure guidance the reference relies on at high res,
so 4096 stops drifting to a washed-out mean:
  - shared flow-match schedule across stages (pred_x0 references align per timestep)
  - per-stage prior: decode -> bicubic upscale -> sharpen -> re-encode
  - custom flowmatch step: per-step low-frequency injection from the upsampled
    previous-stage pred_x0 (alpha, Butterworth FFT split) + velocity momentum (beta)
Keeps validated NTK RoPE (every step) + gated SPA (leading steps). Reference
hyperparameters wired: ntk=[4,10], spa_steps=[3,0], steps=[17,10], guidance=[4.5,6],
alphas=[1.0,0.25], betas=[0.5,0.5], filter_ratio=0.2.

Component-verified offline vs diffusers: flux_rope==FluxPosEmbed, pack/unpack dims,
scheduler tail-index alignment, bundle variants in-range. Deferred: HAP, swin, DWT.
Ref: https://github.com/zylwithxy/HRDiT (pipeline.py flowmatch_step, inference.py)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…e NTK, structure-guidance helpers; drop HAP)
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
smellslikeml and others added 4 commits August 14, 2026 15:18
…perators

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
…C408)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
…lack Forest Labs + HuggingFace Team)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
@smellslikeml
smellslikeml marked this pull request as ready for review August 14, 2026 22:39
@smellslikeml smellslikeml added community-examples Community pipeline / example contribution pipelines labels Aug 14, 2026
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
smellslikeml and others added 2 commits August 14, 2026 17:03
…e scheduler API + dead code

Self-review fixes against the diffusers .ai rubric:
- _flowmatch_step now hands the Euler update to self.scheduler.step() instead of
  inlining prev = sample + (sigma_next-sigma)*model_output and managing
  scheduler._init_step_index/_step_index by hand (pipelines.md gotcha #3). Reads
  the current sigma via scheduler.index_for_timestep(); structure guidance still
  only adjusts the velocity before the step. Behaviour-preserving (same sigma /
  formula / dtype).
- removed the always-true _SPAState.proportional flag + its dead branch, the unused
  self._current_timestep bookkeeping, and the torchvision try/except fallback
  (hard-import, matching the reference) per AGENTS.md 'no defensive/just-in-case code'.
- momentum state (_mo_high/_mo_ref) now reset per stage by the caller.
13/13 unit tests pass; ruff check + format clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-examples Community pipeline / example contribution examples tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant