[core] Add DyPE for training-free high-resolution generation

[smellslikeml]

What does this PR do?

Adds DyPE (Dynamic Position Extrapolation) as a hooks utility for training-free ultra-high-resolution generation on RoPE DiTs such as FluxTransformer2DModel — no fine-tuning and no extra sampling cost.

apply_dype(pipe.transformer) swaps the positional embedding for a timestep-aware YaRN / NTK-by-parts schedule (κ = t²) that engages only above the trained resolution (1024² for Flux) and is otherwise a bit-for-bit no-op.

importtorchfromdiffusersimportFluxPipeline, apply_dypepipe=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
apply_dype(pipe.transformer)
# Above the trained resolution, also flatten the flow-matching shift schedule (see notes).pipe.scheduler.register_to_config(base_shift=1.15, max_shift=1.15)
image=pipe(
"a sunlit alpine meadow, snow-capped peaks, clear blue sky",
height=4096,
width=4096,
guidance_scale=4.5,
num_inference_steps=28,
).images[0]

Notes

• High-resolution shift (required above ~2K). FLUX's default flow-matching shift grows mu with the image sequence length, which collapses the sigma schedule at 4K (the sampler never leaves the noise regime). Pinning base_shift == max_shift holds mu constant. This is a pipeline-level setting and is intentionally kept out of the hook's scope.
• Timestep feed. The normalized timestep is delivered to the embedding through a native register_forward_pre_hook, so it keeps working when the transformer's forward is re-wrapped by enable_model_cpu_offload.
• Known limitation. As with DyPE itself, 4K outputs can carry mild high-frequency speckle (a training-free-extrapolation artifact). The complementary SEGA method targets this and is left for a follow-up.

Validation

Reproduces the reference implementation (guyyariv/DyPE, MIT) on FLUX.1-Krea-dev at 4096² (guidance 4.5, 28 steps); the positional-embedding output matches the reference to Δ = 0. Unit tests cover the schedule helpers, the no-op-at-trained-resolution guarantee, the timestep schedule, and the hook lifecycle.

The schedule is adapted from https://github.com/guyyariv/DyPE (MIT-licensed).

---

🤖 Generated with Claude Code