Community Pipeline: Add HRDiT (training-free high-resolution FLUX.1-dev) - #26
Open
smellslikeml wants to merge 19 commits into
Open
Community Pipeline: Add HRDiT (training-free high-resolution FLUX.1-dev)#26smellslikeml wants to merge 19 commits into
smellslikeml wants to merge 19 commits into
Conversation
…'t crash on import
…e-assign at inference)
… takes no desc kwarg)
…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)
…ion (drop removed use_hap/group_num=1)
…dance method and simplified API
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>
…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
marked this pull request as ready for review
August 14, 2026 22:39
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>
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds
examples/community/pipeline_flux_hrdit.py— training-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-dev — no fine-tuning, no new weights, no separate upscaler model, and no
src/diffuserschanges. 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)
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
FluxPipelinedenoise loop)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_petakes 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.Files
examples/community/pipeline_flux_hrdit.py— the pipelineexamples/community/README.md— usage entrytests/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— naiveFluxPipelinevs HRDiT timing / peak memoryDisclosure
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.