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/hooks/dype.py b/src/diffusers/hooks/dype.py index fa4fcb39f663..8a27014c5c25 100644 --- a/src/diffusers/hooks/dype.py +++ b/src/diffusers/hooks/dype.py @@ -14,7 +14,6 @@ import math -import numpy as np import torch from ..utils import get_logger @@ -28,8 +27,12 @@ _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). +# 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): @@ -41,8 +44,8 @@ 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)) + 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 @@ -62,6 +65,164 @@ def find_newbase_ntk(dim, base, scale): 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, @@ -76,10 +237,12 @@ def _dype_rotary_pos_embed( 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. + 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`): @@ -110,6 +273,9 @@ def _dype_rotary_pos_embed( 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 @@ -180,7 +346,12 @@ def _dype_rotary_pos_embed( 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: + 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) @@ -205,6 +376,9 @@ class _DyPEPosEmbed(torch.nn.Module): 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__( @@ -213,6 +387,10 @@ def __init__( 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 @@ -220,14 +398,64 @@ def __init__( 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 = [] @@ -248,7 +476,7 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor: max_pos = pos[:, i].max().item() current_patches = max_pos + 1 - if self.method == "yarn" and current_patches > self.base_patches: + 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, @@ -258,7 +486,16 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor: 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: @@ -276,19 +513,32 @@ 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. + 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 timestep is fed through a native `torch.nn.Module` forward pre-hook rather than `ModelHook.pre_forward`. This - keeps the timestep 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`. + 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) -> None: + 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 @@ -306,20 +556,33 @@ def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module: 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, ) - def _feed_timestep(mod, args, kwargs): + 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)) - self._timestep_hook_handle = module.register_forward_pre_hook(_feed_timestep, with_kwargs=True) + 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 @@ -335,10 +598,41 @@ def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module: return module -def apply_dype(module: torch.nn.Module, method: str = "yarn", dype: bool = True) -> None: +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. + ultra-high-resolution generation, optionally with [SEGA](https://huggingface.co/papers/2605.22668) spectral + attention. Args: module (`torch.nn.Module`): @@ -346,11 +640,21 @@ def apply_dype(module: torch.nn.Module, method: str = "yarn", dype: bool = True) `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. + 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 position extrapolation schedule by the diffusion timestep (`kappa = t^2`). If - `False`, the extrapolation schedule is static across timesteps. + 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 @@ -360,25 +664,27 @@ def apply_dype(module: torch.nn.Module, method: str = "yarn", dype: bool = True) >>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=torch.bfloat16) >>> pipe.enable_model_cpu_offload() - >>> apply_dype(pipe.transformer) + >>> apply_dype(pipe.transformer, method="spectral") # or method="yarn" for plain DyPE - >>> # Above the trained resolution, also flatten the flow-matching shift schedule. The default FLUX schedule - >>> # grows the timestep shift `mu` with the image sequence length, which collapses the sigma schedule well - >>> # before 4K (the sampler spends nearly every step at pure noise). Pinning `base_shift == max_shift` keeps - >>> # `mu` constant so the sampler denoises normally. This is a pipeline-level setting; the DyPE hook only - >>> # governs the positional embedding. + >>> # 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) - >>> prompt = "a photo of a cat" - >>> image = pipe(prompt, height=4096, width=4096, guidance_scale=4.5, num_inference_steps=28).images[0] + >>> image = pipe("a photo of a cat", height=4096, width=4096, guidance_scale=4.5, num_inference_steps=28).images[0] ``` """ - if method != "yarn": - raise ValueError(f'`method` must be "yarn", but got {method!r}. Other methods are not supported yet.') + 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) + 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/tests/hooks/test_dype.py b/tests/hooks/test_dype.py index 41508fce8021..cb7376e86d10 100644 --- a/tests/hooks/test_dype.py +++ b/tests/hooks/test_dype.py @@ -20,6 +20,10 @@ 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, @@ -267,7 +271,7 @@ def test_timestep_fed_via_native_forward_pre_hook(self): assert len(model._forward_pre_hooks) == 0 # torn down on removal def test_apply_dype_validation(self): - with pytest.raises(ValueError, match='must be "yarn"'): + with pytest.raises(ValueError, match="must be one of"): apply_dype(DummyFluxLikeTransformer(), method="ntk") class NoPosEmbedModel(torch.nn.Module): @@ -276,3 +280,110 @@ def forward(self, 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