[core] Add SEGA spectral attention mode to the DyPE hook

[smellslikeml]

What does this PR do?

Adds an optional SEGA (Spectral-Energy Guided Attention) mode to the DyPE hook. It removes the high-frequency speckle that training-free resolution extrapolation leaves in flat regions at 4K, while preserving fine detail.

apply_dype(pipe.transformer, method="spectral") keeps DyPE's NTK positions but replaces the single YaRN attention-temperature scalar with a per-RoPE-dimension, content-aware temperature derived from the latent's spectral energy at each denoising step — sharpening low-energy frequency bands and damping high-energy ones. method="yarn" (plain DyPE) remains the default and is unchanged.

Stacked on #28 (DyPE); this PR is the method="spectral" delta.

Motivation

DyPE's YaRN temperature is a single scalar (0.1·ln(scale)+1) applied to every RoPE frequency, and it is the source of the residual speckle: dialing it down removes the speckle but also flattens genuine detail — a scalar cannot tell noise from wanted high-frequency content. SEGA does the scaling per-frequency, guided by the image's own spectrum.

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

Flat-sky high-frequency energy (Laplacian variance) drops ~6× with SEGA, with detail retained:

	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, coherent composition:

(SEGA modulates attention per step, so at a fixed seed it yields a different — still coherent — composition than plain DyPE. It is an alternate mode, not a reproduction of the yarn image.)

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
pipe.scheduler.register_to_config(base_shift=1.15, max_shift=1.15)  # high-res shift, see #28
image = pipe(prompt, height=4096, width=4096, guidance_scale=4.5, num_inference_steps=28).images[0]

Implementation

• method="spectral": NTK-scaled RoPE frequencies + per-dim spectral mscale m_ref·(1 − α·σ·s_d) (κ=0.08, α=0.15, β=1.5), exposed as spectral_* kwargs.
• The existing native forward pre-hook additionally reads the packed latent + img_ids each step, computes axis + radial FFT energy profiles, and feeds them via set_spectral_data. One 2D FFT per step; offload-robust (same plumbing as the timestep feed).
• DyPE method="yarn" is byte-identical (parity preserved); SEGA is a no-op at/below the trained resolution.

AI assistance & self-review notes

Written with Claude Code. Self-review of the diff:

• DyPE yarn path unchanged — verified byte-identical positional-embedding output (Δ = 0) vs the base branch.
• Unit tests cover the schedule math, spread endpoints, allocation zero-sum/degenerate cases, the no-op-at-trained-resolution guarantee, and the hook lifecycle (swap → per-step spectral feed → teardown).
• Deferred / by design: (a) SEGA uses NTK base frequencies (faithful to the reference), not layered on DyPE's timestep-modulated YaRN schedule — composing the two is a follow-up; (b) under classifier-free guidance the spectral FFT averages over the batch; (c) the residual sky energy (12.7) is slightly above the fully-flat floor and is tunable via spectral_alpha; (d) deinitalize_hook intentionally matches diffusers' framework spelling (commented inline) — renaming would stop the override from being called on removal.

Attribution

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

Before submitting

• Did you use an AI agent (Claude Code) to help with this PR? Yes.
  • Read the Coding with AI agents guide.
  • Ran the repo self-review skill on the diff — one blocking finding (NumPy in the forward path) fixed; full notes in a comment below.
  • Shared the self-review notes (see comment below).
• Read the contributor guideline.
• Discussed/approved via a GitHub issue — not yet; happy to open a tracking issue.
• Updated documentation — docs/source/en/api/dype.md + _toctree.yml entry.
• 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

[smellslikeml]

Self-review notes

Ran the repo self-review rubric (.ai/review-rules.md + models.md / testing.md / AGENTS.md) over the whole main...dype-sega diff.

Fixed before submitting

• NumPy in the forward path — find_correction_range() used np.floor/np.ceil and runs every forward on the YaRN path. Per models.md ("do not insert NumPy operations in forward implementations", torch.compile graph breaks) this is replaced with math.floor/math.ceil and the numpy import is dropped. DyPE parity re-verified byte-identical after the change.
• Docs — added docs/source/en/api/dype.md + _toctree.yml entry.

Left for the reviewer (flagged, not guessed at)

• .item() device syncs. _DyPEPosEmbed.forward reads pos[:, i].max().item() per spatial axis (and the hook reads img_ids[...].max().item()) to decide whether extrapolation engages. This is data-dependent control flow the stock FluxPosEmbed doesn't have; it forces a CPU sync / graph break under torch.compile. Not obviously removable without changing the engagement contract — happy to take direction.
• Tests use a hand-rolled DummyFluxLikeTransformer rather than a tiny real FluxTransformer2DModel. Per testing.md a real-class fixture is preferred; the mock covers the hook mechanics but not the real pos_embed/forward contract. Can switch to a tiny real transformer if you'd like it in this PR.
• High-resolution shift. Above ~2K the flow-matching mu grows with sequence length and stalls the sampler, so usage requires pipe.scheduler.register_to_config(base_shift=1.15, max_shift=1.15). It's documented, but a pos-embed hook can't reach the pipeline schedule — open to a cleaner surface for it.

Advisory (dead code, inherited from the generic rotary signature)

• In _dype_rotary_pos_embed, the use_real=False / non-repeat_interleave_real branches and the linear_factor argument are unreachable under the Flux pos_embed call path (always use_real=True, repeat_interleave_real=True). Left in to mirror the stock rotary helper; happy to trim to the Flux path if preferred.

Verdict after the fixes: ready apart from the flagged items above, which I'd rather raise than guess at.