diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index b1b7ffebb780..70e2172ba607 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -174,6 +174,8 @@ title: bitsandbytes - local: quantization/gguf title: gguf + - local: quantization/nunchaku + title: Nunchaku Lite - local: quantization/torchao title: torchao - local: quantization/quanto @@ -182,6 +184,8 @@ title: NVIDIA ModelOpt - local: quantization/autoround title: AutoRound + - local: quantization/sdnq + title: SDNQ title: Quantization - isExpanded: false sections: @@ -218,6 +222,8 @@ title: Task recipes - local: using-diffusers/write_own_pipeline title: Understanding pipelines, models and schedulers + - local: using-diffusers/cli + title: Command line interface - local: community_projects title: Projects built with Diffusers - local: conceptual/philosophy @@ -228,8 +234,6 @@ title: How to contribute? - local: conceptual/ethical_guidelines title: Diffusers' Ethical Guidelines - - local: conceptual/evaluation - title: Evaluating Diffusion Models title: Resources - isExpanded: false sections: @@ -353,8 +357,12 @@ title: HunyuanVideoTransformer3DModel - local: api/models/ideogram4_transformer2d title: Ideogram4Transformer2DModel + - local: api/models/transformer_joyimage_edit_plus + title: JoyImageEditPlusTransformer3DModel - local: api/models/transformer_joyimage title: JoyImageEditTransformer3DModel + - local: api/models/krea2_transformer2d + title: Krea2Transformer2DModel - local: api/models/latte_transformer3d title: LatteTransformer3DModel - local: api/models/longcat_image_transformer2d @@ -367,6 +375,10 @@ title: Lumina2Transformer2DModel - local: api/models/lumina_nextdit2d title: LuminaNextDiT2DModel + - local: api/models/minimax_h3_transformer3d + title: MiniMaxH3Transformer3DModel + - local: api/models/minimax_music3_transformer + title: MiniMaxMusic3Transformer1DModel - local: api/models/mochi_transformer3d title: MochiTransformer3DModel - local: api/models/motif_video_transformer_3d @@ -395,6 +407,8 @@ title: Transformer2DModel - local: api/models/transformer_temporal title: TransformerTemporalModel + - local: api/models/wan_animate_2_transformer_3d + title: WanAnimate2Transformer3DModel - local: api/models/wan_animate_transformer_3d title: WanAnimateTransformer3DModel - local: api/models/wan_transformer_3d @@ -451,6 +465,10 @@ title: AutoencoderKLLTXVideo - local: api/models/autoencoderkl_magvit title: AutoencoderKLMagvit + - local: api/models/autoencoderkl_minimax_h3 + title: AutoencoderKLMiniMaxH3 + - local: api/models/autoencoderkl_minimax_h3_audio + title: AutoencoderKLMiniMaxH3Audio - local: api/models/autoencoderkl_mochi title: AutoencoderKLMochi - local: api/models/autoencoderkl_qwenimage @@ -461,6 +479,8 @@ title: AutoencoderRAE - local: api/models/consistency_decoder_vae title: ConsistencyDecoderVAE + - local: api/models/ltx2_diffusion_decoder + title: LTX2VideoDiffusionDecoderModel - local: api/models/autoencoder_oobleck title: Oobleck AutoEncoder - local: api/models/autoencoder_tiny @@ -553,6 +573,8 @@ title: InstructPix2Pix - local: api/pipelines/joyimage_edit title: JoyImage Edit + - local: api/pipelines/joyimage_edit_plus + title: JoyImage Edit Plus - local: api/pipelines/kandinsky title: Kandinsky 2.1 - local: api/pipelines/kandinsky_v22 @@ -563,6 +585,8 @@ title: Kandinsky 5.0 Image - local: api/pipelines/kolors title: Kolors + - local: api/pipelines/krea2 + title: Krea 2 - local: api/pipelines/latent_consistency_models title: Latent Consistency Models - local: api/pipelines/latent_diffusion @@ -591,6 +615,8 @@ title: PixArt-Σ - local: api/pipelines/prx title: PRX + - local: api/pipelines/prx_pixel + title: PRX Pixel - local: api/pipelines/qwenimage title: QwenImage - local: api/pipelines/sana @@ -641,6 +667,8 @@ title: Z-Image title: Image - sections: + - local: api/pipelines/diffusion_gemma + title: DiffusionGemma - local: api/pipelines/llada2 title: LLaDA2 title: Text @@ -675,6 +703,10 @@ title: LTX-2 - local: api/pipelines/ltx_video title: LTXVideo + - local: api/pipelines/minimax_music3 + title: MiniMax Music 3 + - local: api/pipelines/minimax_h3 + title: MiniMax-H3 - local: api/pipelines/mochi title: Mochi - local: api/pipelines/motif_video @@ -685,6 +717,8 @@ title: Stable Video Diffusion - local: api/pipelines/wan title: Wan + - local: api/pipelines/wan_animate_2 + title: Wan-Animate-2 title: Video title: Pipelines - sections: @@ -710,6 +744,8 @@ title: DDPMScheduler - local: api/schedulers/deis title: DEISMultistepScheduler + - local: api/schedulers/discrete_ddim + title: DiscreteDDIMScheduler - local: api/schedulers/multistep_dpm_solver_inverse title: DPMSolverMultistepInverse - local: api/schedulers/multistep_dpm_solver @@ -722,6 +758,8 @@ title: EDMDPMSolverMultistepScheduler - local: api/schedulers/edm_euler title: EDMEulerScheduler + - local: api/schedulers/entropy_bound + title: EntropyBoundScheduler - local: api/schedulers/euler_ancestral title: EulerAncestralDiscreteScheduler - local: api/schedulers/euler @@ -750,6 +788,8 @@ title: LCMScheduler - local: api/schedulers/lms_discrete title: LMSDiscreteScheduler + - local: api/schedulers/minimax_h3 + title: MiniMaxH3Scheduler - local: api/schedulers/pndm title: PNDMScheduler - local: api/schedulers/repaint @@ -774,6 +814,8 @@ title: Custom activation functions - local: api/cache title: Caching methods + - local: api/dype + title: Resolution extrapolation - local: api/normalization title: Custom normalization layers - local: api/utilities diff --git a/docs/source/en/api/dype.md b/docs/source/en/api/dype.md new file mode 100644 index 000000000000..525d5a188a21 --- /dev/null +++ b/docs/source/en/api/dype.md @@ -0,0 +1,46 @@ + + +# Resolution extrapolation + +Training-free methods that let RoPE-based diffusion transformers such as [`FluxTransformer2DModel`] generate above their trained resolution (for example 4096x4096 from a model trained at 1024x1024), with no fine-tuning and no extra sampling cost. + +[DyPE](https://huggingface.co/papers/2510.20766) (Dynamic Position Extrapolation) swaps the transformer's rotary positional embedding for a timestep-aware YaRN / NTK-by-parts schedule that only engages above the trained resolution. [SEGA](https://huggingface.co/papers/2605.22668) (Spectral-Energy Guided Attention) additionally applies a per-frequency, content-aware attention temperature derived from the latent's spectrum, which suppresses the high-frequency speckle a scalar temperature can leave in flat regions at very high resolutions. Both are enabled through [`apply_dype`]. + +```python +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() + +# method="yarn" (default) is plain DyPE; method="spectral" adds SEGA spectral attention. +apply_dype(pipe.transformer, method="spectral") + +# 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( + "a sunlit alpine meadow, snow-capped peaks, clear blue sky", + height=4096, + width=4096, + guidance_scale=4.5, + num_inference_steps=28, +).images[0] +``` + +> [!TIP] +> `apply_dype` is a no-op at or below the trained resolution (1024x1024 for Flux), so the hook can stay applied for standard-resolution generation. + +## apply_dype + +[[autodoc]] apply_dype diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index da77fa67df52..7712f566f471 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -175,6 +175,7 @@ ) _import_structure["hooks"].extend( [ + "DyPEHook", "FasterCacheConfig", "FirstBlockCacheConfig", "HookRegistry", @@ -184,6 +185,7 @@ "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "TextKVCacheConfig", + "apply_dype", "apply_faster_cache", "apply_first_block_cache", "apply_layer_skip", @@ -1036,6 +1038,7 @@ TangentialClassifierFreeGuidance, ) from .hooks import ( + DyPEHook, FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, @@ -1045,6 +1048,7 @@ SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, TextKVCacheConfig, + apply_dype, apply_faster_cache, apply_first_block_cache, apply_layer_skip, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..1adc3b514a07 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -17,6 +17,7 @@ if is_torch_available(): from .context_parallel import apply_context_parallel + from .dype import DyPEHook, apply_dype from .faster_cache import FasterCacheConfig, apply_faster_cache from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache from .group_offloading import apply_group_offloading diff --git a/src/diffusers/hooks/dype.py b/src/diffusers/hooks/dype.py new file mode 100644 index 000000000000..8a27014c5c25 --- /dev/null +++ b/src/diffusers/hooks/dype.py @@ -0,0 +1,690 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch + +from ..utils import get_logger +from ..utils.torch_utils import maybe_adjust_dtype_for_device +from .hooks import HookRegistry, ModelHook + + +logger = get_logger(__name__) # pylint: disable=invalid-name + + +_DYPE_HOOK = "dype_hook" + + +# DyPE schedule adapted from https://github.com/guyyariv/DyPE (MIT). DyPE: "Dynamic Position Extrapolation for Ultra +# High Resolution Diffusion" (https://arxiv.org/abs/2510.20766). +# +# SEGA spectral mscale (method="spectral") adapted from https://github.com/wildminder/ComfyUI-DyPE (Apache-2.0). SEGA: +# "Spectral-Energy Guided Attention for Resolution Extrapolation in Diffusion Transformers" +# (https://arxiv.org/abs/2605.22668). + + +def find_correction_factor(num_rotations, dim, base, max_position_embeddings): + # Inverse dim formula to find the dimension index of a given number of rotations + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + +def find_correction_range(low_ratio, high_ratio, dim, base, ori_max_pe_len): + """ + Find the correction range for NTK-by-parts interpolation. + """ + low = math.floor(find_correction_factor(low_ratio, dim, base, ori_max_pe_len)) + high = math.ceil(find_correction_factor(high_ratio, dim, base, ori_max_pe_len)) + return max(low, 0), min(high, dim - 1) # Clamp values just in case + + +def linear_ramp_mask(min, max, dim): + if min == max: + max += 0.001 # Prevent singularity + + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + +def find_newbase_ntk(dim, base, scale): + """ + Calculate the new base for NTK-aware scaling. + """ + return base * (scale ** (dim / (dim - 2))) + + +# --------------------------------------------------------------------------- +# SEGA: spectral-energy helpers (method="spectral") +# --------------------------------------------------------------------------- + + +def compute_base_mscale(target_res: int, training_res: int, coefficient: float = 0.08) -> float: + r"""Reference magnitude ``m_ref = (R_target / R_train) ** kappa`` (>= 1.0).""" + s = max(float(target_res) / float(training_res), 1.0) + return s**coefficient + + +@torch.no_grad() +def compute_spectral_energy_profile( + hidden_states: torch.Tensor, height: int, width: int, n_bins: int +) -> torch.Tensor: + r""" + Radial (isotropic) spectral energy profile ``E_iso``. Reshapes the leading ``height * width`` tokens into an + ``H x W`` spatial map, averages over batch and channels, mean-centres, computes the 2D FFT power spectrum, and + bins the power into ``n_bins`` concentric rings. + """ + if hidden_states.dim() == 3: + B, S, C = hidden_states.shape + n_spatial = min(S, height * width) + spatial = hidden_states[:, :n_spatial].reshape(B, height, width, C) + elif hidden_states.dim() == 4: + spatial = hidden_states + else: + raise ValueError(f"hidden_states must be 3-D or 4-D, got {hidden_states.dim()}-D") + + spatial_map = spatial.float().mean(dim=(0, -1)) # (H, W) + spatial_map = spatial_map - spatial_map.mean() + + power = torch.fft.fftshift(torch.fft.fft2(spatial_map)).abs().pow(2) + + cy, cx = height / 2.0, width / 2.0 + y = torch.arange(height, device=power.device, dtype=torch.float32) - cy + x = torch.arange(width, device=power.device, dtype=torch.float32) - cx + yy, xx = torch.meshgrid(y, x, indexing="ij") + radius = torch.sqrt(yy**2 + xx**2) + radius = radius / (radius.max() + 1e-8) + + bin_idx = (radius * n_bins).long().clamp(0, n_bins - 1).flatten() + flat_pw = power.flatten() + energy_sum = torch.zeros(n_bins, device=power.device, dtype=torch.float32) + energy_cnt = torch.zeros(n_bins, device=power.device, dtype=torch.float32) + energy_sum.scatter_add_(0, bin_idx, flat_pw) + energy_cnt.scatter_add_(0, bin_idx, torch.ones_like(flat_pw)) + return energy_sum / (energy_cnt + 1e-8) + + +@torch.no_grad() +def compute_axis_spectral_profiles( + hidden_states: torch.Tensor, height: int, width: int, n_bins_h: int, n_bins_w: int +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Per-axis (height, width) 1-D spectral energy profiles, so horizontal and vertical RoPE dimensions can be adjusted + independently. + """ + if hidden_states.dim() == 3: + B, S, C = hidden_states.shape + n_spatial = min(S, height * width) + spatial = hidden_states[:, :n_spatial].reshape(B, height, width, C) + elif hidden_states.dim() == 4: + spatial = hidden_states + else: + raise ValueError(f"hidden_states must be 3-D or 4-D, got {hidden_states.dim()}-D") + + sm = spatial.float().mean(dim=(0, -1)) # (H, W) + sm = sm - sm.mean() + + def _axis_profile(sm: torch.Tensor, axis: int, n_bins: int, length: int) -> torch.Tensor: + fft = torch.fft.fft(sm, dim=axis) + power = fft.abs().pow(2).mean(dim=1 - axis) + half = length // 2 + 1 + power = power[:half] + freq_norm = torch.linspace(0.0, 1.0, half, device=power.device) + bin_idx = (freq_norm * n_bins).long().clamp(0, n_bins - 1) + energy_sum = torch.zeros(n_bins, device=power.device, dtype=torch.float32) + energy_cnt = torch.zeros(n_bins, device=power.device, dtype=torch.float32) + energy_sum.scatter_add_(0, bin_idx, power.float()) + energy_cnt.scatter_add_(0, bin_idx, torch.ones_like(power, dtype=torch.float32)) + return energy_sum / (energy_cnt + 1e-8) + + return ( + _axis_profile(sm, axis=0, n_bins=n_bins_h, length=height), + _axis_profile(sm, axis=1, n_bins=n_bins_w, length=width), + ) + + +@torch.no_grad() +def compute_dynamic_spread( + energy_profile: torch.Tensor, spread_min: float = 0.0, spread_max: float = 1.0, alpha: float = 1.5 +) -> float: + r""" + Spectral-flatness-driven spread in ``[spread_min, spread_max]``. Flat (noise-like) spectrum -> ``spread_min``; + concentrated (structured) spectrum -> ``spread_max``. + """ + eps = 1e-8 + energy = energy_profile.clamp(min=eps) + geo_mean = torch.exp(torch.log(energy).mean()) + arith_mean = energy.mean() + flatness = (geo_mean / (arith_mean + eps)).clamp(0.0, 1.0) + concentration = 1.0 - flatness.item() + return spread_min + (spread_max - spread_min) * (1.0 - (1.0 - concentration) ** alpha) + + +@torch.no_grad() +def compute_spectral_allocation( + energy_profile: torch.Tensor, + freqs: torch.Tensor, + base_mscale: float, + spread: float, + alpha: float = 0.15, + beta: float = 1.5, + min_mscale: float = 1.0, +) -> torch.Tensor: + r""" + Per-RoPE-dimension mscale from a spectral energy profile (SEGA). High spectral-energy dimensions get a *lower* + ``m_k`` (sharpness-biased) while low-energy dimensions get a *higher* ``m_k``; the zero-mean constraint on the + correction redistributes rather than shifts the magnitude. + """ + D_half = freqs.shape[0] + eps = 1e-8 + + # Degenerate cases -> uniform base_mscale + if spread <= 0.0 or alpha <= 0.0: + return torch.full((D_half,), float(base_mscale), device=freqs.device, dtype=torch.float32) + + # Map each RoPE dim to its FFT bin via log-period + periods = 2.0 * math.pi / freqs.clamp(min=eps) + log_periods = torch.log(periods) + min_lp, max_lp = log_periods.min(), log_periods.max() + if (max_lp - min_lp).item() > 1e-6: + lp_norm = (log_periods - min_lp) / (max_lp - min_lp) # 0=high-freq, 1=low-freq + else: + lp_norm = torch.zeros_like(log_periods) + + n_bins = energy_profile.shape[0] + bin_pos = (1.0 - lp_norm) * (n_bins - 1) + j_low = bin_pos.floor().long().clamp(0, n_bins - 1) + j_high = (j_low + 1).clamp(0, n_bins - 1) + frac = (bin_pos - j_low.to(bin_pos.dtype)).clamp(0.0, 1.0) + + E = energy_profile.to(freqs.device).clamp(min=eps) + log_E = torch.log(E) + raw = log_E[j_low] * (1.0 - frac) + log_E[j_high] * frac + + # Standardise + tanh + re-centre (zero-sum property) + z = raw - raw.mean() + z = z / z.std().clamp(min=eps) + s = torch.tanh(float(beta) * z) + s = s - s.mean() + + # Final per-dim mscale + m = float(base_mscale) * (1.0 - float(alpha) * float(spread) * s) + return m.clamp(min=float(min_mscale)).to(torch.float32) + + +def _dype_rotary_pos_embed( + dim: int, + pos: torch.Tensor, + theta: float = 10000.0, + use_real=False, + linear_factor=1.0, + ntk_factor=1.0, + repeat_interleave_real=True, + freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux) + yarn=False, + max_pe_len=None, + ori_max_pe_len=64, + dype=False, + current_timestep=1.0, + spectral_mscale=None, +): + r""" + Precompute the frequency tensor for complex exponentials (cis) with RoPE. Supports YaRN interpolation (optionally + modulated by the DyPE timestep schedule) and, via `spectral_mscale`, SEGA's per-dimension spectral attention + temperature. + + Args: + dim (`int`): + Dimension of the frequency tensor. + pos (`torch.Tensor`): + Position indices for the frequency tensor. [S] or scalar. + theta (`float`, *optional*, defaults to `10000.0`): + Scaling factor for frequency computation. + use_real (`bool`, *optional*, defaults to `False`): + If True, return real part and imaginary part separately. Otherwise, return complex numbers. + linear_factor (`float`, *optional*, defaults to `1.0`): + Scaling factor for linear interpolation. + ntk_factor (`float`, *optional*, defaults to `1.0`): + Scaling factor for NTK-Aware RoPE. + repeat_interleave_real (`bool`, *optional*, defaults to `True`): + If True and use_real, real and imaginary parts are interleaved with themselves to reach `dim`. Otherwise, + they are concatenated. + freqs_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + Data type of the frequency tensor. `torch.float64` is used by models such as Flux. + yarn (`bool`, *optional*, defaults to `False`): + If True, use YaRN interpolation combining NTK, linear, and base methods. + max_pe_len (`int` or `torch.Tensor`, *optional*): + Maximum position encoding length (current patches per axis for vision models). + ori_max_pe_len (`int`, *optional*, defaults to `64`): + Original maximum position encoding length (base patches per axis, 1024 // 16 = 64 for Flux). + dype (`bool`, *optional*, defaults to `False`): + If True, enable DyPE (Dynamic Position Extrapolation) with timestep-aware scaling of the correction + ranges (`kappa = current_timestep**2`). + current_timestep (`float`, *optional*, defaults to `1.0`): + Current timestep for DyPE, normalized to [0, 1] where 1 is pure noise. + spectral_mscale (`torch.Tensor`, *optional*): + Per-dimension SEGA attention temperature of shape `[dim // 2]`. When provided, it replaces the scalar YaRN + attention temperature and is applied to the returned `cos`/`sin`. + + Returns: + `torch.Tensor`: Precomputed frequency tensor for complex exponentials. [S, D/2]. If `use_real=True`, returns a + tuple of `(cos, sin)` tensors. + """ + assert dim % 2 == 0 + + device = pos.device + + if yarn and max_pe_len is not None and max_pe_len > ori_max_pe_len: + if not isinstance(max_pe_len, torch.Tensor): + max_pe_len = torch.tensor(max_pe_len, dtype=freqs_dtype, device=device) + + scale = torch.clamp_min(max_pe_len / ori_max_pe_len, 1.0) + + beta_0 = 1.25 + beta_1 = 0.75 + gamma_0 = 16 + gamma_1 = 2 + + exponents = torch.arange(0, dim, 2, dtype=freqs_dtype, device=device) / dim + freqs_base = 1.0 / (theta**exponents) + # Position interpolation (PI) frequencies + freqs_linear = 1.0 / (scale * theta**exponents) + + new_base = find_newbase_ntk(dim, theta, scale) + if new_base.dim() > 0: + new_base = new_base.view(-1, 1) + freqs_ntk = 1.0 / torch.pow(new_base, exponents) + if freqs_ntk.dim() > 1: + freqs_ntk = freqs_ntk.squeeze() + + if dype: + kappa = current_timestep**2.0 # kappa(t) = t^lambda_t, with lambda_t = 2 + beta_0 = beta_0 * kappa + beta_1 = beta_1 * kappa + + low, high = find_correction_range(beta_0, beta_1, dim, theta, ori_max_pe_len) + low = max(0, low) + high = min(dim // 2, high) + + freqs_mask = 1 - linear_ramp_mask(low, high, dim // 2).to(device).to(freqs_dtype) + freqs = freqs_linear * (1 - freqs_mask) + freqs_ntk * freqs_mask + + if dype: + gamma_0 = gamma_0 * kappa + gamma_1 = gamma_1 * kappa + + low, high = find_correction_range(gamma_0, gamma_1, dim, theta, ori_max_pe_len) + low = max(0, low) + high = min(dim // 2, high) + + freqs_mask = 1 - linear_ramp_mask(low, high, dim // 2).to(device).to(freqs_dtype) + freqs = freqs * (1 - freqs_mask) + freqs_base * freqs_mask + else: + theta_ntk = theta * ntk_factor + exponents = torch.arange(0, dim, 2, dtype=freqs_dtype, device=device) / dim + freqs = 1.0 / (theta_ntk**exponents) / linear_factor + + freqs = torch.outer(pos, freqs) + + is_npu = freqs.device.type == "npu" + if is_npu: + freqs = freqs.float() + + if use_real and repeat_interleave_real: + # flux, hunyuan-dit, cogvideox + freqs_cos = freqs.cos().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1, output_size=freqs.shape[1] * 2).float() # [S, D] + + if spectral_mscale is not None: + # SEGA per-RoPE-dimension attention temperature. Replaces the scalar YaRN temperature below. + ms = spectral_mscale.to(device=freqs_cos.device, dtype=freqs_cos.dtype).repeat_interleave(2) # [D] + freqs_cos = freqs_cos * ms + freqs_sin = freqs_sin * ms + elif yarn and max_pe_len is not None and max_pe_len > ori_max_pe_len: + # YaRN attention temperature. `torch.ones_like` is used instead of a plain `torch.tensor(1.0)` so the + # constant is materialized on the same device/dtype as `scale`. + mscale = torch.where(scale <= 1.0, torch.ones_like(scale), 0.1 * torch.log(scale) + 1.0) + freqs_cos = freqs_cos * mscale + freqs_sin = freqs_sin * mscale + + return freqs_cos, freqs_sin + elif use_real: + # stable audio, allegro + freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D] + freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +class _DyPEPosEmbed(torch.nn.Module): + r""" + Drop-in replacement for the positional embedding of `FluxTransformer2DModel` that applies the DyPE schedule to the + spatial axes. The first axis (text positions) always uses plain RoPE, and the scheduled path only engages on a + spatial axis when the number of patches on that axis (`max_pos + 1`) exceeds the number of patches at the trained + resolution (`base_resolution // patch_size = 1024 // 16 = 64`). As a result, generation at or below the trained + resolution is a no-op compared to the stock positional embedding. + + With `method="spectral"` the spatial axes use NTK-scaled frequencies together with SEGA's per-dimension spectral + attention temperature, computed from the latent's Fourier spectrum (set per step via `set_spectral_data`). + """ + + def __init__( + self, + theta: int, + axes_dim: list[int], + method: str = "yarn", + dype: bool = True, + spectral_alpha: float = 0.15, + spectral_beta: float = 1.5, + spectral_kappa: float = 0.08, + spectral_min_mscale: float = 1.0, + ): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + self.base_resolution = 1024 + self.patch_size = 16 + self.base_patches = self.base_resolution // self.patch_size + self.training_res_pixels = self.base_resolution + self.method = method + self.dype = dype if method != "base" else False + # SEGA parameters + self.spectral_alpha = spectral_alpha + self.spectral_beta = spectral_beta + self.spectral_kappa = spectral_kappa + self.spectral_min_mscale = spectral_min_mscale + # DyPE runtime state + self.current_timestep = 1.0 + # SEGA runtime state (set per step by the hook when method == "spectral") + self._energy_profile_h = None + self._energy_profile_w = None + self._dynamic_spread = 0.0 + self._target_res_h = 0 + self._target_res_w = 0 + + def set_timestep(self, timestep: float): + """Set current timestep for DyPE. Timestep normalized to [0, 1] where 1 is pure noise.""" + self.current_timestep = timestep + + def set_spectral_data(self, energy_profile_h, energy_profile_w, dynamic_spread, target_res_h=0, target_res_w=0): + """Set the per-step SEGA spectral data (called by the hook before each forward when `method == 'spectral'`).""" + self._energy_profile_h = energy_profile_h + self._energy_profile_w = energy_profile_w + self._dynamic_spread = dynamic_spread + self._target_res_h = target_res_h + self._target_res_w = target_res_w + + def _compute_spectral_mscale(self, axis_idx, axis_dim, scale, device): + # Per-dimension SEGA attention temperature for a spatial axis, or None to fall back to plain RoPE. + energy_profile = self._energy_profile_h if axis_idx == 1 else self._energy_profile_w + target_res = self._target_res_h if axis_idx == 1 else self._target_res_w + + m_ref = compute_base_mscale( + target_res if target_res > 0 else 2 * self.training_res_pixels, + self.training_res_pixels, + coefficient=self.spectral_kappa, + ) + if m_ref <= 1.0 + 1e-8: + return None + if energy_profile is None or self._dynamic_spread <= 0.0: + # No spectral data available yet -> uniform reference magnitude. + return torch.full((axis_dim // 2,), float(m_ref), device=device, dtype=torch.float32) + + exponents = torch.arange(0, axis_dim, 2, dtype=torch.float32, device=device) / axis_dim + theta_ntk = self.theta * (scale ** (axis_dim / (axis_dim - 2))) + freqs = 1.0 / (theta_ntk**exponents) + return compute_spectral_allocation( + energy_profile=energy_profile, + freqs=freqs, + base_mscale=m_ref, + spread=self._dynamic_spread, + alpha=self.spectral_alpha, + beta=self.spectral_beta, + min_mscale=self.spectral_min_mscale, + ) + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + n_axes = ids.shape[-1] + cos_out = [] + sin_out = [] + pos = ids.float() + freqs_dtype = maybe_adjust_dtype_for_device(torch.float64, ids.device) + for i in range(n_axes): + common_kwargs = { + "dim": self.axes_dim[i], + "pos": pos[:, i], + "theta": self.theta, + "repeat_interleave_real": True, + "use_real": True, + "freqs_dtype": freqs_dtype, + } + + if i > 0: + max_pos = pos[:, i].max().item() + current_patches = max_pos + 1 + + if current_patches > self.base_patches and self.method == "yarn": + max_pe_len = torch.tensor(current_patches, dtype=freqs_dtype, device=pos.device) + cos, sin = _dype_rotary_pos_embed( + **common_kwargs, + yarn=True, + max_pe_len=max_pe_len, + ori_max_pe_len=self.base_patches, + dype=self.dype, + current_timestep=self.current_timestep, + ) + elif current_patches > self.base_patches and self.method == "spectral": + scale = current_patches / self.base_patches + ntk_factor = scale ** (self.axes_dim[i] / (self.axes_dim[i] - 2)) + spectral_mscale = self._compute_spectral_mscale(i, self.axes_dim[i], scale, pos.device) + cos, sin = _dype_rotary_pos_embed( + **common_kwargs, + yarn=False, + ntk_factor=ntk_factor, + spectral_mscale=spectral_mscale, + ) + else: + cos, sin = _dype_rotary_pos_embed(**common_kwargs) + else: + cos, sin = _dype_rotary_pos_embed(**common_kwargs) + + cos_out.append(cos) + sin_out.append(sin) + + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +class DyPEHook(ModelHook): + r""" + A hook that swaps the positional embedding of a Flux-like transformer for a `_DyPEPosEmbed` and feeds it the + current (normalized) diffusion timestep at every forward pass, enabling training-free ultra-high-resolution + generation. With `method="spectral"` it additionally computes the latent's spectral energy each step and feeds it to + the embedding for SEGA's per-dimension attention temperature. + + The runtime state (timestep, and the SEGA spectral profiles) is delivered through a native `torch.nn.Module` + forward pre-hook rather than `ModelHook.pre_forward`. This keeps it reaching the positional embedding even when + the transformer's `forward` is additionally wrapped by another mechanism -- most importantly accelerate's + `enable_model_cpu_offload`, which re-wraps `forward` and would otherwise bypass `pre_forward`. + """ + + def __init__( + self, + method: str = "yarn", + dype: bool = True, + spectral_alpha: float = 0.15, + spectral_beta: float = 1.5, + spectral_kappa: float = 0.08, + spectral_min_mscale: float = 1.0, + ) -> None: + super().__init__() + + self.method = method + self.dype = dype + self.spectral_alpha = spectral_alpha + self.spectral_beta = spectral_beta + self.spectral_kappa = spectral_kappa + self.spectral_min_mscale = spectral_min_mscale + self._original_pos_embed = None + self._timestep_hook_handle = None + + def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + pos_embed = getattr(module, "pos_embed", None) + if pos_embed is None: + raise ValueError( + "DyPE requires the module to have a `pos_embed` attribute with `theta` and `axes_dim` attributes, as " + "found on `FluxTransformer2DModel`. Please apply the hook to a compatible transformer." + ) + + self._original_pos_embed = pos_embed + module.pos_embed = _DyPEPosEmbed( + theta=pos_embed.theta, + axes_dim=pos_embed.axes_dim, + method=self.method, + dype=self.dype, + spectral_alpha=self.spectral_alpha, + spectral_beta=self.spectral_beta, + spectral_kappa=self.spectral_kappa, + spectral_min_mscale=self.spectral_min_mscale, + ) + + use_spectral = self.method == "spectral" + + def _feed_runtime_state(mod, args, kwargs): + timestep = kwargs.get("timestep", None) + hidden_states = kwargs.get("hidden_states", None) + img_ids = kwargs.get("img_ids", None) + if timestep is None and len(args) > 3: + # `FluxTransformer2DModel.forward` receives (hidden_states, encoder_hidden_states, + # pooled_projections, timestep, ...) positionally when timestep is not passed as a kwarg. + hidden_states = args[0] if hidden_states is None and len(args) > 0 else hidden_states + timestep = args[3] + + if timestep is not None: + if torch.is_tensor(timestep): + timestep = timestep.flatten()[0] + mod.pos_embed.set_timestep(float(timestep)) + + if use_spectral and hidden_states is not None and img_ids is not None: + _feed_spectral_data(mod.pos_embed, hidden_states, img_ids) + + self._timestep_hook_handle = module.register_forward_pre_hook(_feed_runtime_state, with_kwargs=True) + return module + + # NOTE: the method name intentionally matches `ModelHook.deinitalize_hook` / `HookRegistry.remove_hook`, which + # are spelled without the second "i" in diffusers' hooks framework. Renaming to the correct spelling would stop + # this override from being called on removal. + def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + if self._timestep_hook_handle is not None: + self._timestep_hook_handle.remove() + self._timestep_hook_handle = None + if self._original_pos_embed is not None: + module.pos_embed = self._original_pos_embed + self._original_pos_embed = None + return module + + +def _feed_spectral_data(pos_embed, hidden_states, img_ids): + # Compute SEGA's per-step spectral profiles from the packed latent and hand them to the positional embedding. + if not torch.is_tensor(hidden_states) or hidden_states.dim() != 3 or not torch.is_tensor(img_ids): + return + h = int(img_ids[:, 1].max().item()) + 1 + w = int(img_ids[:, 2].max().item()) + 1 + if h * w > hidden_states.shape[1] or h < 2 or w < 2: + return + n_bins_h = max(h // 2, 8) + n_bins_w = max(w // 2, 8) + energy_h, energy_w = compute_axis_spectral_profiles(hidden_states, h, w, n_bins_h, n_bins_w) + iso = compute_spectral_energy_profile(hidden_states, h, w, max(h, w) // 2) + spread = compute_dynamic_spread(iso) + pos_embed.set_spectral_data( + energy_h, + energy_w, + spread, + target_res_h=h * pos_embed.patch_size, + target_res_w=w * pos_embed.patch_size, + ) + + +def apply_dype( + module: torch.nn.Module, + method: str = "yarn", + dype: bool = True, + spectral_alpha: float = 0.15, + spectral_beta: float = 1.5, + spectral_kappa: float = 0.08, + spectral_min_mscale: float = 1.0, +) -> None: + r""" + Applies [DyPE](https://huggingface.co/papers/2510.20766) to a given transformer to enable training-free + ultra-high-resolution generation, optionally with [SEGA](https://huggingface.co/papers/2605.22668) spectral + attention. + + Args: + module (`torch.nn.Module`): + The transformer to apply DyPE to. This should be a RoPE-based DiT with a `pos_embed` attribute exposing + `theta` and `axes_dim`, such as the stock `FluxTransformer2DModel`. At or below the trained resolution + (1024x1024 for Flux), the hook is a no-op. + method (`str`, defaults to `"yarn"`): + The position extrapolation method. `"yarn"` is DyPE's YaRN / NTK-by-parts schedule. `"spectral"` uses + NTK-scaled frequencies together with SEGA's per-dimension spectral attention temperature, which reduces + the high-frequency speckle that a scalar temperature (`"yarn"`) can leave at very high resolutions. + dype (`bool`, defaults to `True`): + Whether to modulate the YaRN extrapolation schedule by the diffusion timestep (`kappa = t^2`). Only + affects `method="yarn"`. + spectral_alpha (`float`, defaults to `0.15`): + SEGA correction amplitude (only used when `method="spectral"`). + spectral_beta (`float`, defaults to `1.5`): + SEGA tanh sharpness (only used when `method="spectral"`). + spectral_kappa (`float`, defaults to `0.08`): + Exponent of the SEGA reference magnitude `m_ref = (R_target / R_train) ** kappa` (only used when + `method="spectral"`). + spectral_min_mscale (`float`, defaults to `1.0`): + Floor for the per-dimension SEGA temperature (only used when `method="spectral"`). + + Example: + ```python + >>> 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") # or method="yarn" for plain DyPE + + >>> # 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 photo of a cat", height=4096, width=4096, guidance_scale=4.5, num_inference_steps=28).images[0] + ``` + """ + + if method not in ("yarn", "spectral", "base"): + raise ValueError(f'`method` must be one of "yarn", "spectral", "base", but got {method!r}.') + + logger.debug(f"Enabling DyPE (method={method}, dype={dype}) on {module.__class__.__name__}") + + hook = DyPEHook( + method=method, + dype=dype, + spectral_alpha=spectral_alpha, + spectral_beta=spectral_beta, + spectral_kappa=spectral_kappa, + spectral_min_mscale=spectral_min_mscale, + ) + registry = HookRegistry.check_if_exists_or_initialize(module) + registry.register_hook(hook, _DYPE_HOOK) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 8439a2b93371..18d910b6334a 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -167,6 +167,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class DyPEHook(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class FasterCacheConfig(metaclass=DummyObject): _backends = ["torch"] @@ -302,6 +317,10 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +def apply_dype(*args, **kwargs): + requires_backends(apply_dype, ["torch"]) + + def apply_faster_cache(*args, **kwargs): requires_backends(apply_faster_cache, ["torch"]) diff --git a/tests/hooks/test_dype.py b/tests/hooks/test_dype.py new file mode 100644 index 000000000000..cb7376e86d10 --- /dev/null +++ b/tests/hooks/test_dype.py @@ -0,0 +1,389 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.hooks import HookRegistry, apply_dype +from diffusers.hooks.dype import ( + DyPEHook, + _DyPEPosEmbed, + _dype_rotary_pos_embed, + compute_axis_spectral_profiles, + compute_base_mscale, + compute_dynamic_spread, + compute_spectral_allocation, + find_correction_factor, + find_correction_range, + find_newbase_ntk, + linear_ramp_mask, +) +from diffusers.models.embeddings import get_1d_rotary_pos_embed +from diffusers.models.transformers.transformer_flux import FluxPosEmbed + + +# Flux trains at 1024x1024 with patch size 16, i.e. 64 patches per spatial axis. +BASE_PATCHES = 64 +AXES_DIM = [16, 56, 56] +THETA = 10000 + + +def build_flux_style_ids(num_txt_tokens: int, patch_grid: int) -> torch.Tensor: + # Mirrors the ids fed to `FluxTransformer2DModel.pos_embed`: `txt_ids` are all-zero and `img_ids` carry the (h, w) + # patch grid positions on the two spatial axes. + txt_ids = torch.zeros(num_txt_tokens, 3) + h, w = torch.meshgrid(torch.arange(patch_grid), torch.arange(patch_grid), indexing="ij") + img_ids = torch.cat([torch.zeros(patch_grid * patch_grid, 1), torch.stack((h, w), dim=-1).reshape(-1, 2)], dim=1) + return torch.cat((txt_ids, img_ids), dim=0).float() + + +class TestDypeScheduleHelpers: + def test_find_correction_factor(self): + # (dim * ln(max_pe / (n_rot * 2pi))) / (2 * ln(base)) + assert find_correction_factor(1.25, 56, THETA, BASE_PATCHES) == pytest.approx(6.37763064832401) + assert find_correction_factor(0.75, 56, THETA, BASE_PATCHES) == pytest.approx(7.9305718956385) + assert find_correction_factor(16, 56, THETA, BASE_PATCHES) == pytest.approx(-1.3728391392110684) + + def test_find_correction_range(self): + # Floor/ceil of the correction factors, clamped to [0, dim - 1]. At kappa=1 (timestep 1.0), the beta ramp + # covers rotations in (1.25, 0.75) and the gamma ramp in (16, 2). + assert find_correction_range(1.25, 0.75, 56, THETA, BASE_PATCHES) == (6.0, 8.0) + assert find_correction_range(16, 2, 56, THETA, BASE_PATCHES) == (0.0, 5.0) + + def test_linear_ramp_mask(self): + torch.testing.assert_close( + linear_ramp_mask(2, 6, 8), torch.tensor([0.0, 0.0, 0.0, 0.25, 0.5, 0.75, 1.0, 1.0]) + ) + # A degenerate (min == max) range is guarded against singularity. + torch.testing.assert_close(linear_ramp_mask(3, 3, 4), torch.zeros(4)) + + def test_find_newbase_ntk(self): + assert find_newbase_ntk(56, THETA, 4.0) == pytest.approx(42107.40810555822) + # `scale` may also be a 0-dim tensor, as in the yarn path. + assert find_newbase_ntk(56, THETA, torch.tensor(4.0, dtype=torch.float64)).item() == pytest.approx( + 42107.40810555822 + ) + + +class TestDypeRotaryPosEmbed: + def test_noop_at_trained_resolution(self): + # With max_pe_len <= ori_max_pe_len the yarn branch must not engage, and the plain branch must be bitwise + # identical to the stock rotary embedding used by FluxPosEmbed. + pos = torch.arange(BASE_PATCHES) + common_kwargs = { + "theta": THETA, + "use_real": True, + "repeat_interleave_real": True, + "freqs_dtype": torch.float64, + } + for max_pe_len in (BASE_PATCHES, BASE_PATCHES // 2): + cos, sin = _dype_rotary_pos_embed( + 56, pos, yarn=True, max_pe_len=max_pe_len, ori_max_pe_len=BASE_PATCHES, dype=True, **common_kwargs + ) + expected_cos, expected_sin = get_1d_rotary_pos_embed(56, pos, **common_kwargs) + assert torch.equal(cos, expected_cos) + assert torch.equal(sin, expected_sin) + + def test_yarn_engages_above_trained_resolution(self): + pos = torch.arange(128) + cos, sin = _dype_rotary_pos_embed( + 56, + pos, + yarn=True, + max_pe_len=128, + ori_max_pe_len=BASE_PATCHES, + dype=True, + current_timestep=1.0, + theta=THETA, + use_real=True, + repeat_interleave_real=True, + freqs_dtype=torch.float64, + ) + expected_cos, expected_sin = get_1d_rotary_pos_embed( + 56, pos, theta=THETA, use_real=True, repeat_interleave_real=True, freqs_dtype=torch.float64 + ) + + assert cos.shape == expected_cos.shape == (128, 56) + assert not torch.equal(cos, expected_cos) + assert not torch.equal(sin, expected_sin) + + # The yarn attention temperature (mscale = 0.1 * ln(scale) + 1.0, scale = 128 / 64 = 2) scales cos/sin, so + # position 0 evaluates to exactly mscale (cos(0) = 1) and 0 (sin(0) = 0). + assert cos[0, 0].item() == pytest.approx(1.0693147180559945) + assert sin[0, 0].item() == pytest.approx(0.0, abs=1e-7) + + def test_yarn_follows_the_timestep_schedule(self): + pos = torch.arange(128) + common_kwargs = { + "yarn": True, + "max_pe_len": 128, + "ori_max_pe_len": BASE_PATCHES, + "theta": THETA, + "use_real": True, + "repeat_interleave_real": True, + "freqs_dtype": torch.float64, + } + cos_noise, sin_noise = _dype_rotary_pos_embed(56, pos, dype=True, current_timestep=1.0, **common_kwargs) + cos_mid, sin_mid = _dype_rotary_pos_embed(56, pos, dype=True, current_timestep=0.5, **common_kwargs) + cos_static, sin_static = _dype_rotary_pos_embed(56, pos, dype=False, **common_kwargs) + + # kappa = t^2 shifts the beta/gamma correction ranges, so the embedding must change across timesteps. + assert not torch.equal(cos_noise, cos_mid) + assert not torch.equal(sin_noise, sin_mid) + # At t = 1.0, kappa = 1 and the DyPE schedule degenerates to static yarn. + assert torch.equal(cos_noise, cos_static) + assert torch.equal(sin_noise, sin_static) + + +class TestDypePosEmbed: + def test_matches_stock_flux_pos_embed_at_trained_resolution(self): + pos_embed = _DyPEPosEmbed(THETA, AXES_DIM) + stock_pos_embed = FluxPosEmbed(THETA, AXES_DIM) + + # Byte-identical output at 1024x1024 (64 patches) and below. + for patch_grid in (8, 32, BASE_PATCHES): + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=patch_grid) + cos, sin = pos_embed(ids) + expected_cos, expected_sin = stock_pos_embed(ids) + assert torch.equal(cos, expected_cos) + assert torch.equal(sin, expected_sin) + + def test_engages_above_trained_resolution(self): + pos_embed = _DyPEPosEmbed(THETA, AXES_DIM) + stock_pos_embed = FluxPosEmbed(THETA, AXES_DIM) + pos_embed.set_timestep(0.5) + + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=128) + cos, sin = pos_embed(ids) + expected_cos, expected_sin = stock_pos_embed(ids) + + assert cos.shape == expected_cos.shape + assert not torch.equal(cos, expected_cos) + assert not torch.equal(sin, expected_sin) + + # The first (text) axis is plain RoPE, so its slice must stay untouched. + assert torch.equal(cos[:, : AXES_DIM[0]], expected_cos[:, : AXES_DIM[0]]) + assert torch.equal(sin[:, : AXES_DIM[0]], expected_sin[:, : AXES_DIM[0]]) + + def test_base_method_disables_dype(self): + pos_embed = _DyPEPosEmbed(THETA, AXES_DIM, method="base", dype=True) + assert pos_embed.dype is False + + +class DummyFluxLikeTransformer(torch.nn.Module): + # Minimal stand-in for `FluxTransformer2DModel` with the same forward signature prefix, so that the hook's + # positional fallback for `timestep` can be exercised. + def __init__(self): + super().__init__() + self.pos_embed = FluxPosEmbed(THETA, AXES_DIM) + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + pooled_projections=None, + timestep=None, + img_ids=None, + txt_ids=None, + ): + ids = torch.cat((txt_ids, img_ids), dim=0) + cos, sin = self.pos_embed(ids) + return cos, sin, timestep + + +class TestDypeHook: + def test_swaps_and_restores_pos_embed(self): + model = DummyFluxLikeTransformer() + original_pos_embed = model.pos_embed + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=8) + hidden_states = torch.randn(8, 4) + + apply_dype(model) + + registry = HookRegistry.check_if_exists_or_initialize(model) + assert isinstance(registry.get_hook("dype_hook"), DyPEHook) + assert isinstance(model.pos_embed, _DyPEPosEmbed) + assert model.pos_embed.theta == original_pos_embed.theta + assert model.pos_embed.axes_dim == original_pos_embed.axes_dim + assert model.pos_embed.method == "yarn" + assert model.pos_embed.dype is True + + registry.remove_hook("dype_hook") + + assert model.pos_embed is original_pos_embed + cos, sin, _ = model(hidden_states, img_ids=ids[8:], txt_ids=ids[:8], timestep=torch.tensor(1.0)) + assert cos.shape == (8 + 64, sum(AXES_DIM)) + + def test_timestep_is_read_from_kwargs(self): + model = DummyFluxLikeTransformer() + apply_dype(model) + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=8) + hidden_states = torch.randn(8, 4) + + # Stock Flux passes the timestep already normalized to [0, 1], where 1 is pure noise. + model(hidden_states, img_ids=ids[8:], txt_ids=ids[:8], timestep=torch.tensor([0.75])) + assert model.pos_embed.current_timestep == 0.75 + + model(hidden_states, img_ids=ids[8:], txt_ids=ids[:8], timestep=0.25) + assert model.pos_embed.current_timestep == 0.25 + + def test_timestep_is_read_from_args(self): + model = DummyFluxLikeTransformer() + apply_dype(model) + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=8) + hidden_states = torch.randn(8, 4) + + model(hidden_states, None, None, torch.tensor(0.5), ids[8:], ids[:8]) + assert model.pos_embed.current_timestep == 0.5 + + def test_timestep_fed_via_native_forward_pre_hook(self): + # Regression: accelerate's `enable_model_cpu_offload` re-wraps the transformer's `forward`, which bypasses + # `ModelHook.pre_forward`. DyPE therefore feeds the timestep with a native forward pre-hook, which + # `nn.Module._call_impl` runs before `forward` regardless of how `forward` is subsequently wrapped. + model = DummyFluxLikeTransformer() + assert len(model._forward_pre_hooks) == 0 + + apply_dype(model) + assert len(model._forward_pre_hooks) == 1 # native pre-hook installed by initialize_hook + + # The timestep must still reach the embedding when `forward` is replaced by an external wrapper. + import functools + + inner_forward = model.forward + model.forward = functools.update_wrapper(lambda *a, **k: inner_forward(*a, **k), inner_forward) + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=8) + model(torch.randn(8, 4), img_ids=ids[8:], txt_ids=ids[:8], timestep=torch.tensor([0.4])) + assert model.pos_embed.current_timestep == pytest.approx(0.4) + + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.remove_hook("dype_hook") + assert len(model._forward_pre_hooks) == 0 # torn down on removal + + def test_apply_dype_validation(self): + with pytest.raises(ValueError, match="must be one of"): + apply_dype(DummyFluxLikeTransformer(), method="ntk") + + class NoPosEmbedModel(torch.nn.Module): + def forward(self, hidden_states): + return hidden_states + + with pytest.raises(ValueError, match="`pos_embed`"): + apply_dype(NoPosEmbedModel()) + + +class DummyFluxLikeTransformerWithLatent(torch.nn.Module): + # Like DummyFluxLikeTransformer but exposes `hidden_states`/`img_ids` so the SEGA spectral path can be exercised. + def __init__(self): + super().__init__() + self.pos_embed = FluxPosEmbed(THETA, AXES_DIM) + + def forward(self, hidden_states, encoder_hidden_states=None, pooled_projections=None, timestep=None, img_ids=None, txt_ids=None): + ids = torch.cat((txt_ids, img_ids), dim=0) + cos, sin = self.pos_embed(ids) + return cos, sin + + +class TestSpectralHelpers: + def test_compute_base_mscale(self): + # m_ref = (target / train) ** kappa, clamped so the ratio is >= 1. + assert compute_base_mscale(4096, 1024, coefficient=0.08) == pytest.approx(4.0**0.08) + assert compute_base_mscale(2048, 1024, coefficient=0.08) == pytest.approx(2.0**0.08) + # At or below the trained resolution the ratio clamps to 1 -> m_ref == 1. + assert compute_base_mscale(512, 1024, coefficient=0.08) == pytest.approx(1.0) + + def test_compute_dynamic_spread_endpoints(self): + # A perfectly flat spectrum is maximally noise-like -> spread == spread_min. + flat = torch.ones(32) + assert compute_dynamic_spread(flat, spread_min=0.0, spread_max=1.0) == pytest.approx(0.0, abs=1e-5) + # A sharply concentrated spectrum is highly structured -> spread near spread_max. + peaked = torch.full((32,), 1e-6) + peaked[0] = 1.0 + assert compute_dynamic_spread(peaked, spread_min=0.0, spread_max=1.0) > 0.9 + + def test_compute_spectral_allocation_zero_sum_and_shape(self): + # Non-flat profile -> non-uniform per-dim mscale; with min_mscale=0 the redistribution is zero-mean so the + # average temperature stays at the reference magnitude. + energy = torch.linspace(1.0, 10.0, 64) + freqs = 1.0 / (THETA ** (torch.arange(0, 56, 2).float() / 56)) + m = compute_spectral_allocation(energy, freqs, base_mscale=1.12, spread=1.0, alpha=0.15, beta=1.5, min_mscale=0.0) + assert m.shape == (28,) + assert (m.max() - m.min()).item() > 1e-3 # non-uniform + assert m.mean().item() == pytest.approx(1.12, abs=1e-3) # zero-sum redistribution + + def test_compute_spectral_allocation_degenerate_is_uniform(self): + energy = torch.linspace(1.0, 10.0, 64) + freqs = 1.0 / (THETA ** (torch.arange(0, 56, 2).float() / 56)) + m = compute_spectral_allocation(energy, freqs, base_mscale=1.12, spread=0.0, alpha=0.15) + assert torch.allclose(m, torch.full((28,), 1.12), atol=1e-6) + + def test_axis_profiles_shape(self): + hs = torch.randn(1, 128 * 128, 8) + e_h, e_w = compute_axis_spectral_profiles(hs, 128, 128, n_bins_h=64, n_bins_w=64) + assert e_h.shape == (64,) and e_w.shape == (64,) + + +class TestSpectralPosEmbed: + def test_noop_at_trained_resolution(self): + # SEGA must be a no-op at/below 1024x1024, identical to plain rope ("base"). + spectral_pe = _DyPEPosEmbed(THETA, AXES_DIM, method="spectral") + base = _DyPEPosEmbed(THETA, AXES_DIM, method="base") + for patch_grid in (32, BASE_PATCHES): + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=patch_grid) + cs, _ = spectral_pe(ids) + cb, _ = base(ids) + assert torch.equal(cs, cb) + + def test_engages_above_trained_resolution(self): + spectral_pe = _DyPEPosEmbed(THETA, AXES_DIM, method="spectral") + yarn = _DyPEPosEmbed(THETA, AXES_DIM, method="yarn") + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=128) + + # Without spectral data, SEGA falls back to a uniform reference magnitude but still differs from YaRN. + cs, _ = spectral_pe(ids) + cy, _ = yarn(ids) + assert cs.shape == cy.shape + assert not torch.equal(cs, cy) + # Text axis (plain rope) is untouched. + assert torch.equal(cs[:, : AXES_DIM[0]], cy[:, : AXES_DIM[0]]) + + def test_per_dim_mscale_is_non_uniform_with_spectral_data(self): + spectral_pe = _DyPEPosEmbed(THETA, AXES_DIM, method="spectral") + energy = torch.linspace(1.0, 10.0, 64) + spectral_pe.set_spectral_data(energy, energy, dynamic_spread=1.0, target_res_h=2048, target_res_w=2048) + m = spectral_pe._compute_spectral_mscale(1, AXES_DIM[1], scale=128 / BASE_PATCHES, device=torch.device("cpu")) + assert m.shape == (AXES_DIM[1] // 2,) + assert (m.max() - m.min()).item() > 1e-3 + + +class TestSpectralHook: + def test_spectral_reads_latent_and_sets_spectral_data(self): + model = DummyFluxLikeTransformerWithLatent() + apply_dype(model, method="spectral") + assert len(model._forward_pre_hooks) == 1 + + G = 96 # > 64 base patches so SEGA engages + ids = build_flux_style_ids(num_txt_tokens=8, patch_grid=G) + img_ids, txt_ids = ids[8:], ids[:8] + hidden_states = torch.randn(1, G * G, 8) + + model(hidden_states=hidden_states, timestep=torch.tensor([0.5]), img_ids=img_ids, txt_ids=txt_ids) + pe = model.pos_embed + assert pe.current_timestep == pytest.approx(0.5) + assert pe._energy_profile_h is not None and pe._energy_profile_w is not None + assert pe._target_res_h == G * pe.patch_size + + registry = HookRegistry.check_if_exists_or_initialize(model) + registry.remove_hook("dype_hook") + assert model.pos_embed.__class__ is FluxPosEmbed + assert len(model._forward_pre_hooks) == 0