""" Vantage DyPE Patch Node for ComfyUI ----------------------------------- Refactored, renamed, and optimized for clarity and maintainability. This node patches Flux models to enable DyPE (Dynamic Positional Encoding) for extended context and higher-resolution image generation. It also outputs an empty latent tensor for sampling in ComfyUI. Features: - DyPE positional embedding (Yarn/NTK/Base) - Dynamic σ-shift scaling based on token count - 64-pixel grid alignment for Flux - Adaptive precision switching (FP16/BF16) - Quantization awareness (GGUF/INT8/FP8/NF4) - Optional half-precision performance mode """ import math import types import torch import torch.nn as nn import numpy as np from typing import List from comfy.model_patcher import ModelPatcher from comfy import model_sampling, model_management # ===================================================================== # Utility Functions # ===================================================================== def compute_correction_factor(rotations: float, dim: int, base: float, max_positions: int) -> float: """Return the frequency correction factor for scaling rotary embeddings.""" return (dim * math.log(max_positions / (rotations * 2 * math.pi))) / (2 * math.log(base)) def compute_correction_range(low_ratio: float, high_ratio: float, dim: int, base: float, max_positions: int): """Return index range for smooth correction transition.""" low = np.floor(compute_correction_factor(low_ratio, dim, base, max_positions)) high = np.ceil(compute_correction_factor(high_ratio, dim, base, max_positions)) return max(low, 0), min(high, dim - 1) def create_linear_mask(start: float, end: float, dim: int) -> torch.Tensor: """Return a linear ramp mask for gradual frequency blending.""" if start == end: end += 1e-3 linear = (torch.arange(dim, dtype=torch.float32) - start) / (end - start) return torch.clamp(linear, 0, 1) def compute_ntk_base(dim: int, base: float, scale: float) -> float: """Return NTK-adjusted base frequency scaling.""" return base * (scale ** (dim / (dim - 2))) def detect_quantized_model(model) -> bool: """Return True if model is quantized (GGUF, FP8, INT8, NF4, etc.).""" try: dm = model.model.diffusion_model if getattr(dm, "is_quantized", False) or hasattr(dm, "gguf"): return True for module in dm.modules(): name = str(type(module)).lower() if any(keyword in name for keyword in ["gguf", "quant", "fp8", "int8", "nf4", "k_m"]): return True except Exception: pass return False # ===================================================================== # Rotary & DyPE Frequency Generation # ===================================================================== def generate_rotary_frequencies( dim: int, pos: torch.Tensor, theta: float = 10000.0, linear_factor: float = 1.0, ntk_factor: float = 1.0, repeat_real: bool = True, use_real: bool = True, freqs_dtype: torch.dtype = torch.float32, yarn_mode: bool = False, max_len: torch.Tensor = None, base_len: int = 64, dype_enabled: bool = False, timestep_norm: float = 1.0, exponent: float = 2.0, ): """ Generate rotary positional embedding frequencies (cos/sin) with optional DyPE and Yarn/NTK extrapolation. """ assert dim % 2 == 0, "Embedding dimension must be even." device = pos.device freqs_dtype = freqs_dtype or torch.float32 # --- Yarn extrapolation for extended context --- if yarn_mode and max_len is not None and max_len > base_len: if not isinstance(max_len, torch.Tensor): max_len = torch.tensor(max_len, dtype=freqs_dtype, device=device) scale = torch.clamp_min(max_len / base_len, 1.0) # Constants for blending regions beta_0, beta_1, gamma_0, gamma_1 = 1.25, 0.75, 16, 2 # Apply DyPE time scaling if dype_enabled: scale_factor = exponent * (timestep_norm ** exponent) beta_0 **= scale_factor beta_1 **= scale_factor gamma_0 **= scale_factor gamma_1 **= scale_factor # Base frequency arrays base_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=device) / dim)) linear_freq = 1.0 / torch.einsum( "... , f -> ... f", scale, (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=device) / dim)) ) ntk_base = compute_ntk_base(dim, theta, scale) ntk_freq = 1.0 / torch.pow(ntk_base, (torch.arange(0, dim, 2, dtype=freqs_dtype, device=device) / dim)) # Blend between NTK and linear low, high = compute_correction_range(beta_0, beta_1, dim, theta, base_len) mask = (1 - create_linear_mask(low, high, dim // 2).to(device).to(freqs_dtype)) freqs = linear_freq * (1 - mask) + ntk_freq * mask # Blend between base and Yarn frequencies low, high = compute_correction_range(gamma_0, gamma_1, dim, theta, base_len) mask = (1 - create_linear_mask(low, high, dim // 2).to(device).to(freqs_dtype)) freqs = freqs * (1 - mask) + base_freq * mask else: # Standard NTK extrapolation theta_scaled = theta * ntk_factor if dype_enabled and ntk_factor > 1.0: theta_scaled *= (ntk_factor ** (exponent * (timestep_norm ** exponent))) freqs = 1.0 / (theta_scaled ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=device) / dim)) / linear_factor freqs = torch.einsum("...s,d->...sd", pos, freqs) if use_real and repeat_real: cos_vals = freqs.cos().repeat_interleave(2, dim=-1).float() sin_vals = freqs.sin().repeat_interleave(2, dim=-1).float() if yarn_mode and max_len is not None and max_len > base_len: mscale = torch.where(scale <= 1., torch.tensor(1.0), 0.1 * torch.log(scale) + 1.0).to(scale) cos_vals *= mscale sin_vals *= mscale return cos_vals, sin_vals return freqs.cos().float(), freqs.sin().float() # ===================================================================== # DyPE Positional Embedding # ===================================================================== class FluxPosEmbed(nn.Module): """ Flux positional embedding layer enhanced with DyPE (Dynamic Positional Encoding) for high-resolution, long-context adaptability. """ def __init__(self, theta: int, axes_dims: List[int], mode: str = "yarn", enable_dype: bool = True, exponent: float = 2.0): super().__init__() self.theta = theta self.axes_dims = axes_dims self.mode = mode self.enable_dype = enable_dype if mode != "base" else False self.exponent = exponent self.base_resolution = 1024 self.patch_size = 16 self.base_patch_count = self.base_resolution // self.patch_size self.timestep_norm = 1.0 print(f"[DyPE] Initialized PosEmbed | base={self.base_resolution}px | patches={self.base_patch_count}") def set_timestep(self, timestep: float): """Store current normalized timestep (0.0–1.0) for DyPE scaling.""" self.timestep_norm = timestep def forward(self, pos_ids: torch.Tensor) -> torch.Tensor: """ Compute positional embedding tensor compatible with ComfyUI Flux attention. Returns: Tensor [B, 1, num_axes*dim, 2, 2] """ n_axes = pos_ids.shape[-1] embeddings = [] pos = pos_ids.float() freqs_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 for i in range(n_axes): axis_pos = pos[..., i] dim = self.axes_dims[i] base_args = dict( dim=dim, pos=axis_pos, theta=self.theta, repeat_real=True, use_real=True, freqs_dtype=freqs_dtype ) # Determine extrapolation mode if i > 0: current_patches = int(axis_pos.max().item() + 1) if self.mode == "yarn" and current_patches > self.base_patch_count: max_len = torch.tensor(current_patches, dtype=freqs_dtype, device=pos.device) cos, sin = generate_rotary_frequencies( **base_args, yarn_mode=True, max_len=max_len, base_len=self.base_patch_count, dype_enabled=self.enable_dype, timestep_norm=self.timestep_norm, exponent=self.exponent ) elif self.mode == "ntk" and current_patches > self.base_patch_count: base_ntk = (current_patches / self.base_patch_count) ** (dim / (dim - 2)) ntk_factor = base_ntk ** (2.0 * (self.timestep_norm ** self.exponent)) if self.enable_dype else base_ntk ntk_factor = max(1.0, ntk_factor) cos, sin = generate_rotary_frequencies(**base_args, ntk_factor=ntk_factor) else: cos, sin = generate_rotary_frequencies(**base_args) else: cos, sin = generate_rotary_frequencies(**base_args) # Match expected Comfy Flux shape cos_r = cos.view(*cos.shape[:-1], -1, 2)[..., :1] sin_r = sin.view(*sin.shape[:-1], -1, 2)[..., :1] row1 = torch.cat([cos_r, -sin_r], dim=-1) row2 = torch.cat([sin_r, cos_r], dim=-1) embeddings.append(torch.stack([row1, row2], dim=-2)) combined = torch.cat(embeddings, dim=-3) return combined.unsqueeze(1).to(pos.device) # ===================================================================== # DyPE Model Patching Logic # ===================================================================== def apply_dype_patch(model: ModelPatcher, width: int, height: int, mode: str, enable_dype: bool, exponent: float, base_shift: float, max_shift: float, performance_mode: bool = True, adaptive_precision: bool = False, base_resolution: int = 1024) -> ModelPatcher: """ Patch a Flux model to enable DyPE with dynamic σ-shift and positional embedding replacement. """ patched_model = model.clone() sampler = patched_model.model.model_sampling # Align to 64-pixel grid height, width = map(lambda v: int(round(v / 64)) * 64, (height, width)) print(f"[DyPE] 🧮 Using {height}x{width} (aligned ×64)") # Token grid calculations latent_h, latent_w = height // 8, width // 8 seq_len = latent_h * latent_w slope = (max_shift - base_shift) / (4096 - 256) sigma_shift = seq_len * slope + (base_shift - slope * 256) # σ-shift patch if isinstance(sampler, model_sampling.ModelSamplingFlux): def dynamic_sigma(self, timestep): return model_sampling.flux_time_shift(getattr(self, "_dype_shift", sigma_shift), 1.0, timestep) sampler.sigma = types.MethodType(dynamic_sigma, sampler) sampler._dype_shift = sigma_shift sampler._vantage_dype = True print(f"[DyPE] ✅ Patched σ-shift | seq_len={seq_len} | shift={sigma_shift:.4f}") # Positional embedding patch try: orig_embed = patched_model.model.diffusion_model.pe_embedder theta, axes_dims = orig_embed.theta, orig_embed.axes_dim except AttributeError: raise ValueError("❌ Incompatible model: missing pe_embedder.") new_embed = FluxPosEmbed(theta, axes_dims, mode, enable_dype, exponent) patched_model.add_object_patch("diffusion_model.pe_embedder", new_embed) sigma_max = patched_model.model.model_sampling.sigma_max.item() # Wrapper for DyPE time/precision handling def wrapped_model_fn(inner_fn, args_dict): if enable_dype: timestep = args_dict.get("timestep") if timestep is not None and timestep.numel() > 0: norm_t = min(max(timestep.item() / sigma_max, 0.0), 1.0) new_embed.set_timestep(norm_t) if adaptive_precision: try: if not hasattr(patched_model, "_quant_checked"): patched_model._is_quant = detect_quantized_model(patched_model) patched_model._quant_checked = True print(f"[DyPE] 🔍 Quantized={patched_model._is_quant}") if not patched_model._is_quant: dtype = torch.float16 if norm_t < 0.5 else torch.bfloat16 patched_model.model.diffusion_model.to(dtype=dtype) except Exception as e: print(f"[DyPE] ⚠️ Precision switch skipped: {e}") return inner_fn(args_dict["input"], args_dict.get("timestep"), **args_dict.get("c", {})) patched_model.set_model_unet_function_wrapper(wrapped_model_fn) if performance_mode: patched_model.model.diffusion_model.half() print("[DyPE] ⚡ Performance mode enabled") return patched_model # ===================================================================== # ComfyUI Node Definition # ===================================================================== class VantageFluxDyPEPatch: """ComfyUI node that applies DyPE patch and returns an empty latent tensor.""" @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "width": ("INT", {"default": 4096, "min": 512, "max": 8192, "step": 8}), "height": ("INT", {"default": 4096, "min": 512, "max": 8192, "step": 8}), "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}), "method": (["yarn", "ntk", "base"], {"default": "yarn"}), "enable_dype": ("BOOLEAN", {"default": True}), "dype_exponent": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 4.0, "step": 0.1}), "base_shift": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}), "max_shift": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 2.0, "step": 0.01}), "performance_mode": ("BOOLEAN", {"default": False}), "adaptive_precision": ("BOOLEAN", {"default": False}), } } RETURN_TYPES = ("MODEL", "LATENT") FUNCTION = "apply_patch_and_create_latent" CATEGORY = "Model/Patches" def apply_patch_and_create_latent( self, model, width, height, batch_size, method, enable_dype, dype_exponent, base_shift, max_shift, performance_mode, adaptive_precision, ): """Apply DyPE patch and create a latent tensor for sampling.""" is_quantized = detect_quantized_model(model) if is_quantized: print("[DyPE] Quantized model detected — disabling performance optimizations.") performance_mode = adaptive_precision = False patched = apply_dype_patch( model, width, height, method, enable_dype, dype_exponent, base_shift, max_shift, performance_mode, adaptive_precision ) latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=model_management.intermediate_device()) print(f"[DyPE] ✅ Patched ({method}) | Latent: {latent.shape}") return (patched, {"samples": latent})