import torch import torch.nn as nn import math from .rope import get_1d_dype_yarn_pos_embed, get_1d_yarn_pos_embed, get_1d_ntk_pos_embed class DyPEBasePosEmbed(nn.Module): """ Base class for Dynamic Position Extrapolation. Handles the calculation of DyPE scaling factors and raw (cos, sin) components. Subclasses must implement `forward` to format the output for specific model architectures. """ def __init__(self, theta: int, axes_dim: list[int], method: str = 'yarn', yarn_alt_scaling: bool = False, dype: bool = True, dype_scale: float = 2.0, dype_exponent: float = 2.0, base_resolution: int = 1024, dype_start_sigma: float = 1.0, base_patch_grid: tuple[int, int] = None): super().__init__() self.theta = theta self.axes_dim = axes_dim self.method = method self.yarn_alt_scaling = yarn_alt_scaling self.dype = True if method == 'vision_yarn' else (dype if method != 'base' else False) self.dype_scale = dype_scale self.dype_exponent = dype_exponent self.base_resolution = base_resolution self.dype_start_sigma = max(0.001, min(1.0, dype_start_sigma)) # Clamp 0.001-1.0 self.current_timestep = 1.0 # Determine Base Patch Grid and Max Patches if base_patch_grid is None: # Default heuristic: 1024px -> 128 latent -> 64 patches (assuming patch_size=2) val = (self.base_resolution // 8) // 2 self.base_patch_grid = (val, val) elif isinstance(base_patch_grid, int): self.base_patch_grid = (base_patch_grid, base_patch_grid) else: self.base_patch_grid = base_patch_grid self.base_patches = max(self.base_patch_grid) def set_timestep(self, timestep: float): self.current_timestep = timestep @staticmethod def _axis_token_span(axis_pos: torch.Tensor) -> float: flat = axis_pos.float().reshape(-1) if flat.numel() <= 1: return 1.0 min_val, max_val = flat.min(), flat.max() span = max_val - min_val if span <= 0: return 1.0 unique_vals = torch.unique(flat) if unique_vals.numel() <= 1: return 1.0 step = torch.diff(unique_vals).min().item() if step <= 1e-6: return float(flat.numel()) return float((span / step) + 1.0) def _get_mscale(self, scale_global): mscale_start = 0.1 * math.log(scale_global) + 1.0 mscale_end = 1.0 t_effective = self.current_timestep t_norm = 1.0 if t_effective > self.dype_start_sigma else (t_effective / self.dype_start_sigma) return mscale_end + (mscale_start - mscale_end) * math.pow(t_norm, self.dype_exponent) def _calc_vision_yarn_components(self, pos: torch.Tensor, freqs_dtype: torch.dtype): n_axes = pos.shape[-1] components = [] if n_axes >= 3: h_span = self._axis_token_span(pos[..., 1]) w_span = self._axis_token_span(pos[..., 2]) scale_global = max(1.0, max(h_span/self.base_patch_grid[0], w_span/self.base_patch_grid[1])) else: max_current_patches = self._axis_token_span(pos) scale_global = max(1.0, max_current_patches / self.base_patches) current_mscale = self._get_mscale(scale_global) for i in range(n_axes): axis_pos = pos[..., i] axis_dim = self.axes_dim[i] current_patches = self._axis_token_span(axis_pos) common_kwargs = {'dim': axis_dim, 'pos': axis_pos, 'theta': self.theta, 'use_real': True, 'repeat_interleave_real': True, 'freqs_dtype': freqs_dtype} dype_kwargs = {'dype': self.dype, 'current_timestep': self.current_timestep, 'dype_scale': self.dype_scale, 'dype_exponent': self.dype_exponent, 'ntk_scale': scale_global, 'override_mscale': current_mscale} if i > 0: base_axis_len = self.base_patch_grid[i-1] if (n_axes >=3 and i-1 < len(self.base_patch_grid)) else self.base_patches scale_local = max(1.0, current_patches / base_axis_len) dype_kwargs['linear_scale'] = scale_local if scale_global > 1.0: cos, sin = get_1d_dype_yarn_pos_embed(**common_kwargs, ori_max_pe_len=base_axis_len, **dype_kwargs) else: cos, sin = get_1d_ntk_pos_embed(**common_kwargs, ntk_factor=1.0) else: cos, sin = get_1d_ntk_pos_embed(**common_kwargs, ntk_factor=1.0) components.append((cos, sin)) return components def _calc_yarn_components(self, pos: torch.Tensor, freqs_dtype: torch.dtype): n_axes = pos.shape[-1] components = [] if n_axes >= 3: h_span = self._axis_token_span(pos[..., 1]) w_span = self._axis_token_span(pos[..., 2]) max_current_patches = max(h_span, w_span) else: max_current_patches = self._axis_token_span(pos) needs_extrapolation = (max_current_patches > self.base_patches) if needs_extrapolation and self.yarn_alt_scaling: for i in range(n_axes): axis_pos = pos[..., i] axis_dim = self.axes_dim[i] common_kwargs = {'dim': axis_dim, 'pos': axis_pos, 'theta': self.theta, 'use_real': True, 'repeat_interleave_real': True, 'freqs_dtype': freqs_dtype} dype_kwargs = {'dype': self.dype, 'current_timestep': self.current_timestep, 'dype_scale': self.dype_scale, 'dype_exponent': self.dype_exponent} current_patches = self._axis_token_span(axis_pos) base_axis_len = self.base_patch_grid[i-1] if (n_axes >=3 and i > 0 and i-1 < len(self.base_patch_grid)) else self.base_patches if i > 0 and current_patches > base_axis_len: max_pe_len = torch.tensor(current_patches, dtype=freqs_dtype, device=pos.device) cos, sin = get_1d_yarn_pos_embed(**common_kwargs, max_pe_len=max_pe_len, ori_max_pe_len=base_axis_len, **dype_kwargs, use_aggressive_mscale=True) else: cos, sin = get_1d_ntk_pos_embed(**common_kwargs, ntk_factor=1.0) components.append((cos, sin)) else: cos_full_spatial, sin_full_spatial = None, None if needs_extrapolation: spatial_axis_dim = self.axes_dim[1] square_pos = torch.arange(0, max_current_patches, device=pos.device).float() max_pe_len = torch.tensor(max_current_patches, dtype=freqs_dtype, device=pos.device) common_kwargs_spatial = {'dim': spatial_axis_dim, 'theta': self.theta, 'use_real': True, 'repeat_interleave_real': True, 'freqs_dtype': freqs_dtype} dype_kwargs = {'dype': self.dype, 'current_timestep': self.current_timestep, 'dype_scale': self.dype_scale, 'dype_exponent': self.dype_exponent} cos_full_spatial, sin_full_spatial = get_1d_yarn_pos_embed( **common_kwargs_spatial, pos=square_pos, max_pe_len=max_pe_len, ori_max_pe_len=self.base_patches, **dype_kwargs, use_aggressive_mscale=False ) for i in range(n_axes): axis_pos = pos[..., i] axis_dim = self.axes_dim[i] if i > 0 and needs_extrapolation: offset_indices = axis_pos.long() - axis_pos.long().min() pos_indices = offset_indices.view(-1) pos_indices = torch.clamp(pos_indices, max=cos_full_spatial.shape[0]-1) cos = cos_full_spatial[pos_indices].view(*axis_pos.shape, -1) sin = sin_full_spatial[pos_indices].view(*axis_pos.shape, -1) else: common_kwargs = {'dim': axis_dim, 'pos': axis_pos, 'theta': self.theta, 'use_real': True, 'repeat_interleave_real': True, 'freqs_dtype': freqs_dtype} cos, sin = get_1d_ntk_pos_embed(**common_kwargs, ntk_factor=1.0) components.append((cos, sin)) return components def _calc_ntk_components(self, pos: torch.Tensor, freqs_dtype: torch.dtype): n_axes = pos.shape[-1] components = [] if n_axes >= 3: h_span = self._axis_token_span(pos[..., 1]) w_span = self._axis_token_span(pos[..., 2]) scale_global = max(1.0, max(h_span/self.base_patch_grid[0], w_span/self.base_patch_grid[1])) else: max_current_patches = self._axis_token_span(pos) scale_global = max(1.0, max_current_patches / self.base_patches) for i in range(n_axes): axis_pos = pos[..., i] axis_dim = self.axes_dim[i] common_kwargs = {'dim': axis_dim, 'pos': axis_pos, 'theta': self.theta, 'use_real': True, 'repeat_interleave_real': True, 'freqs_dtype': freqs_dtype} ntk_factor = 1.0 if i > 0 and scale_global > 1.0: base_ntk = scale_global ** (axis_dim / (axis_dim - 2)) if self.dype: k_t = self.dype_scale * (self.current_timestep ** self.dype_exponent) ntk_factor = base_ntk ** k_t else: ntk_factor = base_ntk ntk_factor = max(1.0, ntk_factor) cos, sin = get_1d_ntk_pos_embed(**common_kwargs, ntk_factor=ntk_factor) components.append((cos, sin)) return components # Public Interface def get_components(self, pos: torch.Tensor, freqs_dtype: torch.dtype): if self.method == 'vision_yarn': return self._calc_vision_yarn_components(pos, freqs_dtype) elif self.method == 'yarn': return self._calc_yarn_components(pos, freqs_dtype) else: return self._calc_ntk_components(pos, freqs_dtype) def forward(self, ids: torch.Tensor) -> torch.Tensor: raise NotImplementedError("Base class does not implement forward. Use a specific model subclass.")