import comfy from comfy.model_patcher import ModelPatcher from comfy.patcher_extension import WrapperExecutor, WrappersMP from comfy_api.latest import io from comfy.ldm.cosmos.predict2 import MiniTrainDIT from comfy.ldm.cosmos.position_embedding import VideoRopePosition3DEmb import functools import math from typing import Optional, Union import torch from torchvision import transforms from einops import rearrange, repeat File_x_SEGA_Anima_key = "File_x_SEGA_Anima_key" File_x_SEGA_Anima_state = "File_x_SEGA_Anima_state" def compute_base_mscale( target_res: int, training_res: int, formula: str = "power_res", coefficient: Optional[float] = None, ) -> float: s = max(float(target_res) / float(training_res), 1.0) if formula == "power_res": c = 0.1 if coefficient is None else float(coefficient) return s ** c if formula == "log_res": c = 0.1 if coefficient is None else float(coefficient) return 1.0 + c * math.log(s) raise ValueError(f"Unknown base_mscale formula: {formula!r}. Use 'power_res' or 'log_res'.") @torch.no_grad() def compute_sega_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: """ Spectral-Energy Guided Attention (SEGA): per-dim RoPE mscale (asymmetric band). z_k = standardise( log E[bin(k)] ) s_k = tanh( beta * z_k ), re-centred m_k = base_mscale * ( 1 - alpha * spread * s_k ) High spectral-energy dims get lower m_k (sharpness-biased). """ D_half = freqs.shape[0] eps = 1e-8 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 z = raw - raw.mean() z = z / z.std().clamp(min=eps) s = torch.tanh(float(beta) * z) s = s - s.mean() # direction = -1 => subtract m = float(base_mscale) * (1.0 - float(alpha) * float(spread) * s) return m.clamp(min=float(min_mscale)).to(torch.float32) class File_x_SEGA_Anima_RoPE(torch.nn.Module): def __init__( self, axes_dims_rope: tuple[int, int, int] = (44, 42, 42), theta: float = 10000.0, base_mscale_formula: str = "power_res", base_mscale_coefficient: Optional[float] = None, training_resolution_h: float = 1024, training_resolution_w: float = 1024, mscale_alpha: float = 0.15, mscale_beta: float = 1.5, mscale_min: float = 1.0, debug: int = 0b0, *args, **kwargs ): super().__init__(*args, **kwargs) self.axes_dim = axes_dims_rope self.theta = theta self.base_mscale_formula = base_mscale_formula self.base_mscale_coefficient = base_mscale_coefficient self.training_resolution_h = training_resolution_h self.training_resolution_w = training_resolution_w self.mscale_alpha = mscale_alpha self.mscale_beta = mscale_beta self.mscale_min = mscale_min self.frame = 0 self.height = 64 self.width = 64 self.build_freqs_with_size(self.frame, self.height, self.width, 1.0, 1.0, 1.0, debug) @staticmethod def rope_params(index, dim, theta=10000, ntk_factor=1.0, debug=0b0): assert dim % 2 == 0 scaled_theta = theta * ntk_factor freqs = torch.outer(index, 1.0 / torch.pow(scaled_theta, torch.arange(0, dim, 2).to(torch.float32).div(dim))) # freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs def build_freqs_with_size(self, frame, height, width, ntk_t: float, ntk_h: float, ntk_w: float, debug:int): """(Re)build pos/neg frequency tensors for the given per-axis NTK factors.""" pos_index_frame = torch.arange(frame, dtype=torch.float) pos_index_height = torch.arange(height, dtype=torch.float) pos_index_width = torch.arange(width, dtype=torch.float) neg_index_frame = torch.arange(frame, dtype=torch.float).flip(0) * -1 - 1 neg_index_height = torch.arange(height, dtype=torch.float).flip(0) * -1 - 1 neg_index_width = torch.arange(width, dtype=torch.float).flip(0) * -1 - 1 self.pos_freqs_t = self.rope_params(pos_index_frame, self.axes_dim[0], self.theta, ntk_factor=ntk_t, debug=debug) self.pos_freqs_h = self.rope_params(pos_index_height, self.axes_dim[1], self.theta, ntk_factor=ntk_h, debug=debug) self.pos_freqs_w = self.rope_params(pos_index_width, self.axes_dim[2], self.theta, ntk_factor=ntk_w, debug=debug) self.neg_freqs_t = self.rope_params(neg_index_frame, self.axes_dim[0], self.theta, ntk_factor=ntk_t, debug=debug) self.neg_freqs_h = self.rope_params(neg_index_height, self.axes_dim[1], self.theta, ntk_factor=ntk_h, debug=debug) self.neg_freqs_w = self.rope_params(neg_index_width, self.axes_dim[2], self.theta, ntk_factor=ntk_w, debug=debug) if ntk_t > 1.0: scaled_theta_t = self.theta * ntk_t temporal_dim = self.axes_dim[0] self._inv_freqs_t = 1.0 / torch.pow( scaled_theta_t, torch.arange(0, temporal_dim, 2).float().div(temporal_dim) ) if ntk_h > 1.0: scaled_theta_h = self.theta * ntk_h spatial_dim = self.axes_dim[1] self._inv_freqs_h = 1.0 / torch.pow( scaled_theta_h, torch.arange(0, spatial_dim, 2).float().div(spatial_dim) ) if ntk_w > 1.0: scaled_theta_w = self.theta * ntk_w spatial_dim = self.axes_dim[2] self._inv_freqs_w = 1.0 / torch.pow( scaled_theta_w, torch.arange(0, spatial_dim, 2).float().div(spatial_dim) ) self.ntk_t = ntk_t self.ntk_h = ntk_h self.ntk_w = ntk_w if hasattr(self, 'compute_condition_freqs'): self.compute_condition_freqs.cache_clear() def ensure_ntk_factor(self, frame: int, height: int, width: int, ntk_t: float, ntk_h: float, ntk_w: float, debug: int): """Rebuild frequencies if per-axis NTK factors changed.""" if (frame != self.frame or height != self.height or width != self.width or abs(ntk_t - self.ntk_t) > 1e-8 or abs(ntk_h - self.ntk_h) > 1e-8 or abs(ntk_w - self.ntk_w) > 1e-8): self.build_freqs_with_size(frame, height, width, ntk_t, ntk_h, ntk_w, debug) @torch.no_grad() def compute_mscale(self, ntk_factor_t, ntk_factor_h, ntk_factor_w, mscale_spread, energy_profile_t=None, energy_profile_h=None, energy_profile_w=None, target_res_h=None, target_res_w=None, device=None, debug=0b0): """ Compute per-frequency-dimension mscale with SEGA (Spectral-Energy Guided Attention). Returns None if no scaling needed, otherwise [D_total/2] real tensor. """ if (ntk_factor_h <= 1.0 and ntk_factor_w <= 1.0) or mscale_spread is None: return None # temporal axis SEGA if ntk_factor_t > 1.0 and hasattr(self, '_inv_freqs_t'): base_ms_t = compute_base_mscale( 1.0, 1.0, self.base_mscale_formula, self.base_mscale_coefficient, ) dev = device or self._inv_freqs_t.device mscale_t = torch.full((self.axes_dim[0] // 2,), base_ms_t, device=dev, dtype=torch.float32) else: mscale_t = None # Height axis SEGA if ntk_factor_h > 1.0 and hasattr(self, '_inv_freqs_h'): base_ms_h = compute_base_mscale( target_res_h or self.training_resolution_h, self.training_resolution_h, self.base_mscale_formula, self.base_mscale_coefficient, ) use_sega_h = ( energy_profile_h is not None and float(mscale_spread) > 0.0 and base_ms_h > 1.0 + 1e-8 ) if use_sega_h: inv_f = self._inv_freqs_h.to(device) if device is not None else self._inv_freqs_h mscale_h = compute_sega_allocation( energy_profile_h, inv_f, base_ms_h, float(mscale_spread), float(self.mscale_alpha), float(self.mscale_beta), float(self.mscale_min), ) else: dev = device or self._inv_freqs_h.device mscale_h = torch.full((self.axes_dim[1] // 2,), base_ms_h, device=dev, dtype=torch.float32) else: mscale_h = None # Width axis SEGA if ntk_factor_w > 1.0 and hasattr(self, '_inv_freqs_w'): base_ms_w = compute_base_mscale( target_res_w or self.training_resolution_w, self.training_resolution_w, self.base_mscale_formula, self.base_mscale_coefficient, ) use_sega_w = ( energy_profile_w is not None and float(mscale_spread) > 0.0 and base_ms_w > 1.0 + 1e-8 ) if use_sega_w: inv_f = self._inv_freqs_w.to(device) if device is not None else self._inv_freqs_w mscale_w = compute_sega_allocation( energy_profile_w, inv_f, base_ms_w, float(mscale_spread), float(self.mscale_alpha), float(self.mscale_beta), float(self.mscale_min), ) else: dev = device or self._inv_freqs_w.device mscale_w = torch.full((self.axes_dim[2] // 2,), base_ms_w, device=dev, dtype=torch.float32) else: mscale_w = None if mscale_h is None and mscale_w is None: return None dev = device or (mscale_h.device if mscale_h is not None else mscale_w.device) if mscale_t is None: mscale_t = torch.ones(self.axes_dim[0] // 2, device=dev) if mscale_h is None: mscale_h = torch.ones(self.axes_dim[1] // 2, device=dev) if mscale_w is None: mscale_w = torch.ones(self.axes_dim[2] // 2, device=dev) mscale_full = torch.cat([ mscale_t, mscale_h, mscale_w, ]) # vals = [f"{v:.4f}" for v in mscale_full.tolist()] # print( # f"[SEGA] mscale (per-dim, {mscale_full.shape[0]} dims): " # f"min={mscale_full.min().item():.4f} " # f"mean={mscale_full.mean().item():.4f} " # f"max={mscale_full.max().item():.4f}\n" # f"[SEGA] mscale values: {vals}" # ) return mscale_full @functools.lru_cache(maxsize=None) def compute_condition_freqs(self, frame: int, height: int, width: int, device: torch.device, debug: int): seq_lens = frame * height * width pos_freqs_t = self.pos_freqs_t.to(device) if device is not None else self.pos_freqs_t pos_freqs_h = self.pos_freqs_h.to(device) if device is not None else self.pos_freqs_h pos_freqs_w = self.pos_freqs_w.to(device) if device is not None else self.pos_freqs_w neg_freqs_t = self.neg_freqs_t.to(device) if device is not None else self.neg_freqs_t neg_freqs_h = self.neg_freqs_h.to(device) if device is not None else self.neg_freqs_h neg_freqs_w = self.neg_freqs_w.to(device) if device is not None else self.neg_freqs_w freqs_frame = pos_freqs_t[:frame] freqs_height = pos_freqs_h[:height] freqs_width = pos_freqs_w[:width] freqs_frame = torch.stack([freqs_frame.cos(), -freqs_frame.sin(), freqs_frame.sin(), freqs_frame.cos()], dim=-1) freqs_height = torch.stack([freqs_height.cos(), -freqs_height.sin(), freqs_height.sin(), freqs_height.cos()], dim=-1) freqs_width = torch.stack([freqs_width.cos(), -freqs_width.sin(), freqs_width.sin(), freqs_width.cos()], dim=-1) em_T_H_W_D = torch.cat( [ repeat(freqs_frame, "t d x -> t h w d x", h=height, w=width), repeat(freqs_height, "h d x -> t h w d x", t=frame, w=width), repeat(freqs_width, "w d x -> t h w d x", t=frame, h=height), ] , dim=-2, ) return rearrange(em_T_H_W_D, "t h w d (i j) -> (t h w) d i j", i=2, j=2).float() def forward( self, x_B_T_H_W_D: torch.Tensor, ntk_t: float = 1.0, ntk_h: float = 1.0, ntk_w: float = 1.0, mscale_spread: Optional[float] = None, energy_profile_t: Optional[torch.Tensor] = None, energy_profile_h: Optional[torch.Tensor] = None, energy_profile_w: Optional[torch.Tensor] = None, target_resolution_h: Optional[int] = None, target_resolution_w: Optional[int] = None, debug: int = 0b0 ) -> torch.Tensor: device = x_B_T_H_W_D.device _, frame, height, width, _ = x_B_T_H_W_D.shape self.ensure_ntk_factor(frame, height, width, ntk_t, ntk_h, ntk_w, debug) video_freq = self.compute_condition_freqs(frame, height, width, device, debug) mscale = self.compute_mscale( ntk_factor_t=ntk_t, ntk_factor_h=ntk_h, ntk_factor_w=ntk_w, mscale_spread=mscale_spread, energy_profile_t=energy_profile_t, energy_profile_h=energy_profile_h, energy_profile_w=energy_profile_w, target_res_h=target_resolution_h, target_res_w=target_resolution_w, device=device, debug=debug ) if mscale is not None: return video_freq * mscale.unsqueeze(0).unsqueeze(-1).unsqueeze(-1) else: return video_freq @torch.no_grad() def compute_axis_spectral_profiles(hidden_states, height, width, n_bins_h, n_bins_w): reordered_hidden_states = rearrange(hidden_states, "b c t h w -> b (t h w) c") B, S, C = reordered_hidden_states.shape # x_B_C_T_H_W n_spatial = min(S, height * width) spatial = reordered_hidden_states[:, :n_spatial].reshape(B, height, width, C) sm = spatial.float().mean(dim=(0, -1)) sm = sm - sm.mean() def _axis_profile(sm, axis, n_bins, length): 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_spectral_energy_profile(hidden_states, height, width, n_bins): reordered_hidden_states = rearrange(hidden_states, "b c t h w -> b (t h w) c") B, S, C = reordered_hidden_states.shape n_spatial = min(S, height * width) spatial = reordered_hidden_states[:, :n_spatial].reshape(B, height, width, C) spatial_map = spatial.float().mean(dim=(0, -1)) 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_norm = torch.sqrt(yy ** 2 + xx ** 2) radius_norm = radius_norm / (radius_norm.max() + 1e-8) bin_idx = (radius_norm * 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_dynamic_spread(energy_profile, spread_min=0.0, spread_max=1.0, alpha=1.5): 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) def prepare_embedded_sequence( class_obj: MiniTrainDIT, x_B_C_T_H_W: torch.Tensor, fps: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None, transformer_options: Optional[dict] = None ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: """ Prepares an embedded sequence tensor by applying positional embeddings and handling padding masks. Args: x_B_C_T_H_W (torch.Tensor): video fps (Optional[torch.Tensor]): Frames per second tensor to be used for positional embedding when required. If None, a default value (`class_obj.base_fps`) will be used. padding_mask (Optional[torch.Tensor]): current it is not used Returns: Tuple[torch.Tensor, Optional[torch.Tensor]]: - A tensor of shape (B, T, H, W, D) with the embedded sequence. - An optional positional embedding tensor, returned only if the positional embedding class (`class_obj.pos_emb_cls`) includes 'rope'. Otherwise, None. Notes: - If `class_obj.concat_padding_mask` is True, a padding mask channel is concatenated to the input tensor. - The method of applying positional embeddings depends on the value of `class_obj.pos_emb_cls`. - If 'rope' is in `class_obj.pos_emb_cls` (case insensitive), the positional embeddings are generated using the `class_obj.pos_embedder` with the shape [T, H, W]. - If "fps_aware" is in `class_obj.pos_emb_cls`, the positional embeddings are generated using the `class_obj.pos_embedder` with the fps tensor. - Otherwise, the positional embeddings are generated without considering fps. """ if class_obj.concat_padding_mask: if padding_mask is None: padding_mask = torch.zeros(x_B_C_T_H_W.shape[0], 1, x_B_C_T_H_W.shape[3], x_B_C_T_H_W.shape[4], dtype=x_B_C_T_H_W.dtype, device=x_B_C_T_H_W.device) else: padding_mask = transforms.functional.resize( padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST ) x_B_C_T_H_W = torch.cat( [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1 ) x_B_T_H_W_D = class_obj.x_embedder(x_B_C_T_H_W) if class_obj.extra_per_block_abs_pos_emb: extra_pos_emb = class_obj.extra_pos_embedder(x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device, dtype=x_B_C_T_H_W.dtype) else: extra_pos_emb = None if "rope" in class_obj.pos_emb_cls.lower(): state: dict = transformer_options.get(File_x_SEGA_Anima_state, None) if state is None: raise RuntimeError("Failed to pass parameters to internal function.") training_resolution_h: int = state.get("training_resolution_h", 1024) training_resolution_w: int = state.get("training_resolution_w", 1024) debug: int = state.get("debug", 0) d_mul: float = state.get("d_mul", 2.0) log_s_mul: float = state.get("log_s_mul", 0.1) axes_dims_rope: tuple[int, int, int] = state.get("axes_dims_rope", (44, 42, 42)) rope_embedder: File_x_SEGA_Anima_RoPE = state.get("RoPE_Embedder", None) if rope_embedder is None: raise RuntimeError("Class File_x_SEGA_Anima_RoPE is not properly transfered.") img_h = x_B_C_T_H_W.shape[3] img_w = x_B_C_T_H_W.shape[4] target_resolution_h = img_h * 8 target_resolution_w = img_w * 8 s_h = target_resolution_h / training_resolution_h s_w = target_resolution_w / training_resolution_w d_h = axes_dims_rope[1] d_w = axes_dims_rope[2] ntk_h = s_h ** (d_mul * d_h / (d_h - 2)) / (1 + log_s_mul * math.log(s_h)) ntk_w = s_w ** (d_mul * d_w / (d_w - 2)) / (1 + log_s_mul * math.log(s_w)) big = max(float(s_h), float(s_w)) small = min(float(s_h), float(s_w)) if big > 1.0 and small > 0.0 and abs(s_h - s_w) > 1e-8: ntk_big = max(float(ntk_h), float(ntk_w)) ntk_small = 1.0 + (small / big) * max(0.0, (ntk_big - 1.0)) if s_h < s_w: ntk_h = ntk_small ntk_w = ntk_big else: ntk_h = ntk_big ntk_w = ntk_small if max(ntk_h, ntk_w) > 1.0 and img_h > 0 and img_w > 0: n_bins_div = 1 if debug & 0b100 == 0 else 2 n_bins_h = max(img_h // n_bins_div, 8) n_bins_w = max(img_w // n_bins_div, 8) energy_profile_h, energy_profile_w = compute_axis_spectral_profiles(x_B_C_T_H_W, img_h, img_w, n_bins_h, n_bins_w) iso_profile = compute_spectral_energy_profile(x_B_C_T_H_W, img_h, img_w, max(img_h, img_w) // 2) dynamic_spread = compute_dynamic_spread(iso_profile, 0.0, 1.0, 1.5) else: energy_profile_h = None energy_profile_w = None dynamic_spread = None rope_emb_L_1_1_D = rope_embedder( x_B_T_H_W_D, ntk_h=ntk_h, ntk_w=ntk_w, mscale_spread=dynamic_spread, energy_profile_h=energy_profile_h, energy_profile_w=energy_profile_w, target_resolution_h=target_resolution_h, target_resolution_w=target_resolution_w, debug=debug ) return x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb x_B_T_H_W_D = x_B_T_H_W_D + class_obj.pos_embedder(x_B_T_H_W_D, device=x_B_C_T_H_W.device) # [B, T, H, W, D] return x_B_T_H_W_D, None, extra_pos_emb def File_x_SEGA_Anima_wrapper_forward( class_obj: MiniTrainDIT, x: torch.Tensor, timesteps: torch.Tensor, context: torch.Tensor, fps: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None, **kwargs, ): orig_shape = list(x.shape) x = comfy.ldm.common_dit.pad_to_patch_size(x, (class_obj.patch_temporal, class_obj.patch_spatial, class_obj.patch_spatial)) x_B_C_T_H_W = x timesteps_B_T = timesteps crossattn_emb = context transformer_options: dict = kwargs.get("transformer_options", {}) """ Args: x: (B, C, T, H, W) tensor of spatial-temp inputs timesteps: (B, ) tensor of timesteps crossattn_emb: (B, N, D) tensor of cross-attention embeddings """ x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = prepare_embedded_sequence( class_obj, x_B_C_T_H_W, fps=fps, padding_mask=padding_mask, transformer_options=transformer_options ) if timesteps_B_T.ndim == 1: timesteps_B_T = timesteps_B_T.unsqueeze(1) t_embedding_B_T_D, adaln_lora_B_T_3D = class_obj.t_embedder[1](class_obj.t_embedder[0](timesteps_B_T).to(x_B_T_H_W_D.dtype)) t_embedding_B_T_D = class_obj.t_embedding_norm(t_embedding_B_T_D) # for logging purpose affline_scale_log_info = {} affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach() class_obj.affline_scale_log_info = affline_scale_log_info class_obj.affline_emb = t_embedding_B_T_D class_obj.crossattn_emb = crossattn_emb if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None: assert ( x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape ), f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}" block_kwargs = { "rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0), "adaln_lora_B_T_3D": adaln_lora_B_T_3D, "extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D, "transformer_options": transformer_options, } # The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream # in fp32, but run attention and MLP modules in fp16. # An alternate method that clamps fp16 values "works" in the sense that it makes coherent images, but there is noticeable # quality degradation and visual artifacts. if x_B_T_H_W_D.dtype == torch.float16: x_B_T_H_W_D = x_B_T_H_W_D.float() for block in class_obj.blocks: x_B_T_H_W_D = block( x_B_T_H_W_D, t_embedding_B_T_D, crossattn_emb, **block_kwargs, ) x_B_T_H_W_O = class_obj.final_layer(x_B_T_H_W_D.to(crossattn_emb.dtype), t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D) x_B_C_Tt_Hp_Wp = class_obj.unpatchify(x_B_T_H_W_O)[:, :, :orig_shape[-3], :orig_shape[-2], :orig_shape[-1]] return x_B_C_Tt_Hp_Wp def File_x_SEGA_Anima_wrapper(executor: WrapperExecutor, x, timesteps, context, fps, padding_mask, **kwargs): return File_x_SEGA_Anima_wrapper_forward(executor.class_obj, x, timesteps, context, fps, padding_mask, **kwargs) class File_x_SEGA_Anima_(io.ComfyNode): @classmethod def define_schema(cls): return io.Schema(node_id="File_x SEGA Anima", display_name="SEGA for Anima", category="model/patch/anima", inputs=[io.Model.Input(id="model"), io.Int.Input(id="training_resolution", default=1024, min=1, max=65536, step=1), io.Float.Input(id="theta", default=10000.0, min=10.0, max=1000000.0, step=10.0), io.Combo.Input(id="base_mscale_formula", options=["power_res", "log_res"], default="power_res"), io.Float.Input(id="base_mscale_coefficient", default=0.08, min=-100.000, max=100.000, step=0.001), io.Float.Input(id="mscale_alpha", default=0.15, min=-100.000, max=100.000, step=0.001), io.Float.Input(id="mscale_beta", default=1.5, min=-100.000, max=100.000, step=0.001), io.Float.Input(id="mscale_min", default=1.0, min=-100.000, max=100.000, step=0.001), io.Float.Input(id="d_mul", default=2.0, min=-100.000, max=100.000, step=0.001), io.Float.Input(id="log_s_mul", default=0.1, min=-100.000, max=100.000, step=0.001), # io.Int.Input(id="debug", default=0b0, min=0, max=0b1111_1111_1111_1111, step=1, advanced=True), ], outputs=[io.Model.Output(id=None),],) @classmethod def execute(cls, model: ModelPatcher, training_resolution, theta, base_mscale_formula, base_mscale_coefficient, mscale_alpha, mscale_beta, mscale_min, d_mul, log_s_mul): base_model = getattr(model, "model", None) if base_model is None: raise RuntimeError("Model input has no model object.") image_model = None model_config = getattr(base_model, "model_config", None) if model_config is not None: image_model = model_config.unet_config.get("image_model", None) if image_model != "anima": raise RuntimeError("Model input is not Anima model.") diffusion_model = getattr(base_model, "diffusion_model", None) if diffusion_model is None: raise RuntimeError("Model input has invalid Anima model.") modified_model = model.clone() head_dim = modified_model.model.diffusion_model.model_channels // modified_model.model.diffusion_model.num_heads dim_h = head_dim // 6 * 2 dim_w = dim_h dim_t = head_dim - 2 * dim_h assert head_dim == dim_h + dim_w + dim_t, f"bad dim: {head_dim} != {dim_h} + {dim_w} + {dim_t}" axes_dims_rope = (dim_t, dim_h, dim_w) modified_model.remove_wrappers_with_key(WrappersMP.DIFFUSION_MODEL, File_x_SEGA_Anima_key) transformer_options = modified_model.model_options.setdefault("transformer_options", {}) transformer_options[File_x_SEGA_Anima_state] = { "training_resolution_h": training_resolution, "training_resolution_w": training_resolution, "debug": 0b0, "d_mul": d_mul, "log_s_mul": log_s_mul, "axes_dims_rope": axes_dims_rope, "RoPE_Embedder": File_x_SEGA_Anima_RoPE(axes_dims_rope, theta, base_mscale_formula, base_mscale_coefficient, training_resolution, training_resolution, mscale_alpha, mscale_beta, mscale_min, 0b0) } modified_model.add_wrapper_with_key(WrappersMP.DIFFUSION_MODEL, File_x_SEGA_Anima_key, File_x_SEGA_Anima_wrapper) return io.NodeOutput(modified_model)