[core] DyPE + SEGA hook for training-free high-resolution generation (upstream candidate)

[smellslikeml]

Consolidated upstream candidate. This is the squashed feature branch (dype-upstream) shown against main — a single commit with the exact diff intended for huggingface/diffusers. The incremental development PRs are #28 (DyPE) and #29 (SEGA); this one exists to share the whole PR body + diff for coordination.

Coordination: huggingface/diffusers#14520 — placement (core hook vs community pipeline) is being settled with maintainers before this is filed upstream.

What does this PR do?

Adds apply_dype, a training-free hook that lets RoPE-based diffusion transformers such as [FluxTransformer2DModel] generate well above their trained resolution (e.g. 4096×4096 from a 1024×1024-trained model), with no fine-tuning and no extra sampling cost. Two modes:

• DyPE (method="yarn") — a timestep-aware YaRN / NTK-by-parts RoPE schedule that engages only above the trained resolution and is otherwise a no-op.
• SEGA (method="spectral") — DyPE positions plus a per-frequency, spectral-energy-guided attention temperature that removes the residual high-frequency speckle a scalar temperature leaves in flat regions at very high resolutions.

The hook swaps the transformer's positional embedding and feeds it the normalized timestep — and, for SEGA, the latent's per-step spectral energy — through a native register_forward_pre_hook, so it keeps working when forward is re-wrapped by enable_model_cpu_offload.

Motivation

DyPE reaches 4K on Flux training-free, but its YaRN attention temperature is a single scalar applied to every RoPE frequency, which leaves a fine multicolored speckle in flat regions. Turning the scalar down removes the speckle but also flattens genuine detail — a scalar cannot separate noise from wanted high-frequency content. SEGA does the scaling per-frequency, guided by the image's own spectrum, so it suppresses the speckle while keeping detail.

Results — FLUX.1-Krea-dev, 4096², guidance 4.5, 28 steps

• DyPE (method="yarn") reproduces the reference implementation (guyyariv/DyPE, MIT) bit-for-bit: positional-embedding output matches to Δ = 0.
• SEGA (method="spectral") drops flat-sky high-frequency energy (Laplacian variance) ~6× while keeping detail crisp:

	sky HF energy
DyPE (method="yarn")	76.5
DyPE + SEGA (method="spectral")	12.7

Full method="spectral" frame — clean sky, crisp daisies and snow/ridge detail:

Usage

import torch
from diffusers import FluxPipeline, apply_dype

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()

apply_dype(pipe.transformer, method="spectral")  # method="yarn" (default) = plain DyPE

# Above the trained resolution, also flatten the flow-matching shift schedule so the sampler
# does not stall near pure noise (the default shift `mu` grows with the image sequence length).
pipe.scheduler.register_to_config(base_shift=1.15, max_shift=1.15)

image = pipe(prompt, height=4096, width=4096, guidance_scale=4.5, num_inference_steps=28).images[0]

Implementation

• pos_embed is swapped for a _DyPEPosEmbed; a no-op at/below the trained resolution (1024² for Flux).
• Offload-robust native forward pre-hook feeds the timestep (and, for SEGA, the packed latent + img_ids → axis/radial FFT energy profiles) each step. One 2D FFT per step for SEGA; negligible cost.
• method="yarn": YaRN / NTK-by-parts frequencies with DyPE timestep modulation κ = t². method="spectral": NTK frequencies + per-dim spectral mscale m_ref·(1 − α·σ·s_d) (κ=0.08, α=0.15, β=1.5), exposed as spectral_* kwargs.
• High-resolution shift. Above ~2K, FLUX's default flow-matching mu grows with the sequence length and collapses the sigma schedule; pinning base_shift == max_shift holds mu constant. This is a pipeline-level setting and is intentionally kept out of the hook's scope (see usage).

Documentation & tests

• docs/source/en/api/dype.md + _toctree.yml entry ("Resolution extrapolation").
• tests/hooks/test_dype.py — schedule helpers, the no-op-at-trained-resolution guarantee, the DyPE timestep schedule, SEGA spectral allocation, and the hook lifecycle (swap → per-step feed → teardown).

AI assistance & self-review notes

Written with Claude Code, reviewed with the repo self-review rubric.

• Fixed: NumPy in the forward path (np.floor/np.ceil → math) per models.md's torch.compile rule; DyPE parity re-verified byte-identical afterward.
• Left for the reviewer (flagged, not guessed at): (a) .item() device syncs in the pos-embed forward — data-dependent extrapolation-engagement control flow the stock embedding doesn't have; (b) tests use a hand-rolled DummyFluxLikeTransformer rather than a tiny real FluxTransformer2DModel (testing.md prefers real classes); (c) the high-res shift is documented usage, not code — a pos-embed hook can't reach the pipeline schedule; (d) minor dead branches in _dype_rotary_pos_embed inherited from the generic rotary helper.

Attribution

DyPE schedule from https://github.com/guyyariv/DyPE (MIT). SEGA spectral mscale adapted from https://github.com/wildminder/ComfyUI-DyPE (Apache-2.0); method from SEGA (arXiv:2605.22668).

Before submitting

• Did you use an AI agent (Claude Code)? Yes — see the self-review notes above.
  • Read the Coding with AI agents guide.
  • Ran the repo self-review skill on the diff.
  • Shared the self-review notes (above).
• Read the contributor guideline.
• Discussed/approved via a GitHub issue — huggingface/diffusers#14520 (placement discussion).
• Updated documentation (docs/source/en/api/dype.md + toctree).
• Wrote new tests (tests/hooks/test_dype.py).

Who can review?

General functionalities / hooks: yiyixuxu, DN6, sayakpaul (live @-tags to be added when filed against huggingface/diffusers).

---

🤖 Generated with Claude Code