Landing-zone spec: DyPE as a diffusers hook

[github-actions]

Test results

ℹ️ Tests could not run in CI — the runner lacks this repo's dependencies (a collection/import error, not a code failure). Run the suite locally to validate.

License & code availability

• 🟢 guyyariv/DyPE — MIT (permissive)

Drafted from a design brief supplied at dispatch time. No arXiv anchor — the brief IS the spec.

Design brief

Drafted by an autonomous coding loop — Remyx's Outrider action receives a design brief via workflow_dispatch, produces an implementation against this repo, and opens this PR for review. The brief is the sole input; there is no ranker-picked paper behind this PR.

Implementation by: Claude Code as autonomous agent

Brief

Landing-zone spec: DyPE as a diffusers hook

Goal

Add a diffusers-native DyPEHook that runs DyPE (Dynamic Position Extrapolation) on the
stock FluxTransformer2DModel via the existing src/diffusers/hooks/ framework — training-free
ultra-high-resolution text-to-image (up to 4096×4096) on off-the-shelf FLUX, no fork of the transformer,
no new weights, no src/diffusers/models/** changes. The upstream DyPE release ships a vendored fork of
diffusers' Flux; the missing piece this PR provides is the clean, model-general hook form.

Scope this PR to FLUX + the yarn (YaRN/NTK-by-parts) path — DyPE's default. Mention (in the PR
narrative only) that the technique generalizes to other RoPE DiTs; do not build speculative
multi-model abstraction or extra config surface for callers that don't exist yet.

References (read these)

• Paper: DyPE, "Dynamic Position Extrapolation for Ultra High Resolution Diffusion", ICML 2026 — arXiv:2510.20766
• Official implementation (MIT): https://github.com/guyyariv/DyPE — read flux/transformer_flux.py:
  get_1d_rotary_pos_embed (the YaRN/NTK-by-parts + DyPE κ=t² schedule), the FluxPosEmbed class
  (set_timestep, per-axis forward), and the transformer forward line self.pos_embed.set_timestep(timestep.item()).
• diffusers hooks framework: src/diffusers/hooks/hooks.py (ModelHook: initialize_hook/pre_forward/
  deinitialize_hook; HookRegistry.register_hook). Mirror the style of src/diffusers/hooks/faster_cache.py
  and layer_skip.py (an apply_* entrypoint that registers the hook).
• diffusers Flux (do NOT modify): src/diffusers/models/transformers/transformer_flux.py — stock FluxPosEmbed,
  the transformer forward(..., timestep=...), and image_rotary_emb = self.pos_embed(ids).

Landing zone (where the code goes)

• New file: src/diffusers/hooks/dype.py
• Export DyPEHook and apply_dype from src/diffusers/hooks/__init__.py
• New test: tests/hooks/test_dype.py (CPU-only unit tests — schedule helpers + no-op-at-trained-res)
• Do not touch src/diffusers/models/**.

The schedule (reproduce exactly from the MIT reference)

Port these verbatim (add # Adapted from https://github.com/guyyariv/DyPE (MIT)):

• Helpers: find_correction_factor(n_rot,dim,base,max_pe) = dim*ln(max_pe/(n_rot*2π))/(2*ln base);
  find_correction_range (floor/ceil, clamp to [0, dim-1]); linear_ramp_mask = clamp((arange-min)/(max-min),0,1);
  find_newbase_ntk(dim,base,scale) = base*scale**(dim/(dim-2)).
• _dype_rotary_pos_embed(...) — YaRN path when max_pe_len > ori_max_pe_len (=64):
  scale = clamp_min(max_pe_len/64, 1.0); build freqs_base, freqs_linear (PI, 1/(scale*θ**...)),
  freqs_ntk (1/find_newbase_ntk(...)**...). DyPE modulation: κ = current_timestep**2; scale
  β0=1.25, β1=0.75 by κ → correction range → blend freqs_linear↔freqs_ntk; then scale
  γ0=16, γ1=2 by κ → blend result↔freqs_base. YaRN attention temperature
  mscale = where(scale<=1, 1.0, 0.1*log(scale)+1.0) applied to cos/sin. Non-yarn path: θ_ntk = θ*ntk_factor.
• _DyPEPosEmbed(nn.Module) — mirrors stock FluxPosEmbed.forward per axis: base_patches = 1024//16 = 64;
  axis 0 = plain RoPE; spatial axes only, and the scheduled path engages only when
  current_patches (= max_pos+1) > base_patches. Holds current_timestep (default 1.0), method='yarn', dype=True.

Integration shape

• DyPEHook(ModelHook):
  • initialize_hook(transformer): construct _DyPEPosEmbed copying the transformer's existing
    pos_embed.theta / .axes_dim; swap transformer.pos_embed to it; stash the original.
  • pre_forward(module, *args, **kwargs): read timestep from kwargs/args (stock FLUX passes it already
    normalized /1000 ∈ [0,1], 1=noise); set module.pos_embed.current_timestep = float(timestep.flatten()[0]).
    Return (args, kwargs) unchanged.
  • deinitialize_hook(transformer): restore the original pos_embed.
• apply_dype(transformer, method="yarn", dype=True): register DyPEHook via HookRegistry.
• No-op below the trained resolution: at ≤1024² (current_patches ≤ 64) output must be byte-identical to stock.

Constraints

• Training-free; no new weights; no src/diffusers/models/** edits.
• Faithful port of the MIT reference schedule; keep the attribution comment.
• Match diffusers hook conventions (faster_cache.py); make style && make quality must pass; add the unit test.
• Disclose AI-assisted authorship in the PR body and commit message.

Validation (out of scope for the draft; maintainer runs on GPU)

On black-forest-labs/FLUX.1-Krea-dev (DyPE's validated checkpoint), guidance_scale=4.5, 4096²:
apply_dype(pipe.transformer) on the stock diffusers pipeline must reproduce DyPE's official-fork output;
naive (no hook) shows extrapolation collapse; at 1024² the hook is a no-op. (Note: plain FLUX.1-dev is NOT
DyPE-validated — use Krea-dev.)

Opened by the Remyx Recommendation orchestrator.

Co-Authored-By: remyx-ai[bot] <289541483+remyx-ai[bot]@users.noreply.github.com>

[smellslikeml]

Superseded by #28 — same change on a conventionally-named branch (dype-hook) with a rewritten description. Closing.