367 lines
15 KiB
Python
367 lines
15 KiB
Python
"""
|
||
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})
|
||
|