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..31f2ac0c8482 --- /dev/null +++ b/src/diffusers/hooks/dype.py @@ -0,0 +1,363 @@ +# 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 numpy as np +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" + + +# Adapted from https://github.com/guyyariv/DyPE (MIT). DyPE: "Dynamic Position Extrapolation for Ultra High +# Resolution Diffusion" (https://arxiv.org/abs/2510.20766). + + +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 = np.floor(find_correction_factor(low_ratio, dim, base, ori_max_pe_len)) + high = np.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))) + + +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, +): + r""" + Precompute the frequency tensor for complex exponentials (cis) with RoPE. Supports YaRN interpolation, optionally + modulated by the DyPE timestep schedule. + + 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. + + 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 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. + """ + + def __init__( + self, + theta: int, + axes_dim: list[int], + method: str = "yarn", + dype: bool = True, + ): + 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.method = method + self.dype = dype if method != "base" else False + self.current_timestep = 1.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 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 self.method == "yarn" and current_patches > self.base_patches: + 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, + ) + + 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) timestep at every forward pass, enabling training-free ultra-high-resolution generation. + """ + + def __init__(self, method: str = "yarn", dype: bool = True) -> None: + super().__init__() + + self.method = method + self.dype = dype + self._original_pos_embed = 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, + ) + return module + + def pre_forward(self, module: torch.nn.Module, *args, **kwargs) -> tuple[tuple, dict]: + timestep = kwargs.get("timestep", None) + if timestep is None and len(args) > 3: + # Stock `FluxTransformer2DModel.forward` receives (hidden_states, encoder_hidden_states, + # pooled_projections, timestep, ...) positionally when not passed as a kwarg. + timestep = args[3] + + if timestep is not None: + if torch.is_tensor(timestep): + timestep = timestep.flatten()[0] + module.pos_embed.set_timestep(float(timestep)) + + return args, kwargs + + def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module: + if self._original_pos_embed is not None: + module.pos_embed = self._original_pos_embed + self._original_pos_embed = None + return module + + +def apply_dype(module: torch.nn.Module, method: str = "yarn", dype: bool = True) -> None: + r""" + Applies [DyPE](https://huggingface.co/papers/2510.20766) to a given transformer to enable training-free + ultra-high-resolution generation. + + 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 to use. Only `"yarn"` (YaRN / NTK-by-parts, DyPE's default) is currently + supported. + dype (`bool`, defaults to `True`): + Whether to modulate the position extrapolation schedule by the diffusion timestep (`kappa = t^2`). If + `False`, the extrapolation schedule is static across timesteps. + + 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.to("cuda") + + >>> apply_dype(pipe.transformer) + >>> image = pipe("a photo of a cat", height=4096, width=4096, guidance_scale=4.5).images[0] + ``` + """ + + if method != "yarn": + raise ValueError(f'`method` must be "yarn", but got {method!r}. Other methods are not supported yet.') + + logger.debug(f"Enabling DyPE (method={method}, dype={dype}) on {module.__class__.__name__}") + + hook = DyPEHook(method=method, dype=dype) + 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..c10e4af3eb6c --- /dev/null +++ b/tests/hooks/test_dype.py @@ -0,0 +1,255 @@ +# 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, + 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_apply_dype_validation(self): + with pytest.raises(ValueError, match='must be "yarn"'): + 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())