import torch import torch.nn.functional as F import torch.nn as nn from torch import Tensor, FloatTensor from typing import Optional, Callable, Tuple, Dict, List, Any, Union import einops from einops import rearrange import copy import comfy from .latents import gaussian_blur_2d, median_blur_2d # WIP... not yet in use... class StyleTransfer: def __init__(self, style_method = "WCT", embedder_method = None, patch_size = 1, pinv_dtype = torch.float64, dtype = torch.float64, ): self.style_method = style_method self.embedder_method = None self.unembedder_method = None if embedder_method is not None: self.set_embedder_method(embedder_method) self.patch_size = patch_size #if embedder_type == "conv2d": # self.unembedder = self.invert_conv2d self.pinv_dtype = pinv_dtype self.dtype = dtype self.patchify = None self.unpatchify = None self.orig_shape = None self.grid_sizes = None #self.x_embed_ndim = 0 def set_patchify_method(self, patchify_method=None): self.patchify_method = patchify_method def set_unpatchify_method(self, unpatchify_method=None): self.unpatchify_method = unpatchify_method def set_embedder_method(self, embedder_method): self.embedder_method = copy.deepcopy(embedder_method).to(self.pinv_dtype) self.W = self.embedder_method.weight self.B = self.embedder_method.bias if isinstance(embedder_method, nn.Linear): self.unembedder_method = self.invert_linear elif isinstance(embedder_method, nn.Conv2d): self.unembedder_method = self.invert_conv2d elif isinstance(embedder_method, nn.Conv3d): self.unembedder_method = self.invert_conv3d def set_patch_size(self, patch_size): self.patch_size = patch_size def unpatchify(self, x: Tensor) -> List[Tensor]: x_arr = [] for i, img_size in enumerate(self.img_sizes): # [[64,64]] , img_sizes: List[Tuple[int, int]] pH, pW = img_size x_arr.append( einops.rearrange(x[i, :pH*pW].reshape(1, pH, pW, -1), 'B H W (p1 p2 C) -> B C (H p1) (W p2)', p1=self.patch_size, p2=self.patch_size) ) x = torch.cat(x_arr, dim=0) return x def patchify(self, x: Tensor): x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size)) pH, pW = x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size self.img_sizes = [[pH, pW]] * x.shape[0] x = einops.rearrange(x, 'B C (H p1) (W p2) -> B (H W) (p1 p2 C)', p1=self.patch_size, p2=self.patch_size) return x def embedder(self, x): if isinstance(self.embedder_method, nn.Linear): x = self.patchify(x) self.orig_shape = x.shape x = self.embedder_method(x) self.grid_sizes = x.shape[2:] #self.x_embed_ndim = x.ndim #if x.ndim > 3: # x = einops.rearrange(x, "B C H W -> B (H W) C") return x def unembedder(self, x): #if self.x_embed_ndim > 3: # x = einops.rearrange(x, "B (H W) C -> B C H W", W=self.orig_shape[-1]) x = self.unembedder_method(x) return x def invert_linear(self, x : torch.Tensor,) -> torch.Tensor: x = x.to(self.pinv_dtype) #x = (x - self.B.to(self.dtype)) @ torch.linalg.pinv(self.W.to(self.pinv_dtype)).T.to(self.dtype) x = (x - self.B) @ torch.linalg.pinv(self.W).T return x.to(self.dtype) def invert_conv2d(self, z: torch.Tensor,) -> torch.Tensor: z = z.to(self.pinv_dtype) conv = self.embedder_method B, C_in, H, W = self.orig_shape C_out, _, kH, kW = conv.weight.shape stride_h, stride_w = conv.stride pad_h, pad_w = conv.padding b = conv.bias.view(1, C_out, 1, 1).to(z) z_nobias = z - b W_flat = conv.weight.view(C_out, -1).to(z) W_pinv = torch.linalg.pinv(W_flat) Bz, Co, Hp, Wp = z_nobias.shape z_flat = z_nobias.reshape(Bz, Co, -1) x_patches = W_pinv @ z_flat x_sum = F.fold( x_patches, output_size=(H + 2*pad_h, W + 2*pad_w), kernel_size=(kH, kW), stride=(stride_h, stride_w), ) ones = torch.ones_like(x_patches) count = F.fold( ones, output_size=(H + 2*pad_h, W + 2*pad_w), kernel_size=(kH, kW), stride=(stride_h, stride_w), ) x_recon = x_sum / count.clamp(min=1e-6) if pad_h > 0 or pad_w > 0: x_recon = x_recon[..., pad_h:pad_h+H, pad_w:pad_w+W] return x_recon.to(self.dtype) def invert_conv3d(self, z: torch.Tensor, ) -> torch.Tensor: z = z.to(self.pinv_dtype) conv = self.embedder_method grid_sizes = self.grid_sizes B, C_in, D, H, W = self.orig_shape pD, pH, pW = self.patch_size sD, sH, sW = pD, pH, pW if z.ndim == 3: # [B, S, C_out] -> reshape to [B, C_out, D', H', W'] S = z.shape[1] if grid_sizes is None: Dp = D // pD Hp = H // pH # getting actual patchified dims Wp = W // pW else: Dp, Hp, Wp = grid_sizes C_out = z.shape[2] z = z.transpose(1, 2).reshape(B, C_out, Dp, Hp, Wp) else: B2, C_out, Dp, Hp, Wp = z.shape assert B2 == B, "Batch size mismatch... ya sharked it." b = conv.bias.view(1, C_out, 1, 1, 1) # need to kncokout bias to invert via weight z_nobias = z - b # 2D filter -> pinv w3 = conv.weight # [C_out, C_in, 1, pH, pW] w2 = w3.squeeze(2) # [C_out, C_in, pH, pW] out_ch, in_ch, kH, kW = w2.shape W_flat = w2.view(out_ch, -1) # [C_out, in_ch*pH*pW] W_pinv = torch.linalg.pinv(W_flat) # [in_ch*pH*pW, C_out] # merge depth for 2D unfold wackiness z2 = z_nobias.permute(0,2,1,3,4).reshape(B*Dp, C_out, Hp, Wp) # apply pinv ... get patch vectors z_flat = z2.reshape(B*Dp, C_out, -1) # [B*Dp, C_out, L] x_patches = W_pinv @ z_flat # [B*Dp, in_ch*pH*pW, L] # fold -> restore spatial frames x2 = F.fold( x_patches, output_size=(H, W), kernel_size=(pH, pW), stride=(sH, sW) ) # → [B*Dp, C_in, H, W] # unmerge depth (de-depth charge) x2 = x2.reshape(B, Dp, in_ch, H, W) # [B, Dp, C_in, H, W] x_recon = x2.permute(0,2,1,3,4).contiguous() # [B, C_in, D, H, W] return x_recon.to(self.dtype) def adain_seq_inplace(self, content: torch.Tensor, style: torch.Tensor, eps: float = 1e-7) -> torch.Tensor: mean_c = content.mean(1, keepdim=True) std_c = content.std (1, keepdim=True).add_(eps) mean_s = style.mean (1, keepdim=True) std_s = style.std (1, keepdim=True).add_(eps) content.sub_(mean_c).div_(std_c).mul_(std_s).add_(mean_s) return content class StyleWCT: def __init__(self, dtype=torch.float64, use_svd=False,): self.dtype = dtype self.use_svd = use_svd self.y0_adain_embed = None self.mu_s = None self.y0_color = None self.spatial_shape = None def whiten(self, f_s_centered: torch.Tensor, set=False): cov = (f_s_centered.T.double() @ f_s_centered.double()) / (f_s_centered.size(0) - 1) if self.use_svd: U_svd, S_svd, Vh_svd = torch.linalg.svd(cov + 1e-5 * torch.eye(cov.size(0), dtype=cov.dtype, device=cov.device)) S_eig = S_svd U_eig = U_svd else: S_eig, U_eig = torch.linalg.eigh(cov + 1e-5 * torch.eye(cov.size(0), dtype=cov.dtype, device=cov.device)) if set: S_eig_root = S_eig.clamp(min=0).sqrt() # eigenvalues -> singular values else: S_eig_root = S_eig.clamp(min=0).rsqrt() # inverse square root whiten = U_eig @ torch.diag(S_eig_root) @ U_eig.T return whiten.to(f_s_centered) def set(self, y0_adain_embed: torch.Tensor, spatial_shape=None): if self.y0_adain_embed is None or self.y0_adain_embed.shape != y0_adain_embed.shape or torch.norm(self.y0_adain_embed - y0_adain_embed) > 0: self.y0_adain_embed = y0_adain_embed.clone() if spatial_shape is not None: self.spatial_shape = spatial_shape f_s = y0_adain_embed[0] # if y0_adain_embed.ndim > 4 else y0_adain_embed self.mu_s = f_s.mean(dim=0, keepdim=True) f_s_centered = f_s - self.mu_s self.y0_color = self.whiten(f_s_centered, set=True) def get(self, denoised_embed: torch.Tensor): for wct_i in range(denoised_embed.shape[0]): f_c = denoised_embed[wct_i] mu_c = f_c.mean(dim=0, keepdim=True) f_c_centered = f_c - mu_c whiten = self.whiten(f_c_centered) f_c_whitened = f_c_centered @ whiten.T f_cs = f_c_whitened @ self.y0_color.T + self.mu_s denoised_embed[wct_i] = f_cs return denoised_embed class WaveletStyleWCT(StyleWCT): def set(self, y0_adain_embed: torch.Tensor, h_len, w_len): if self.y0_adain_embed is None or self.y0_adain_embed.shape != y0_adain_embed.shape or torch.norm(self.y0_adain_embed - y0_adain_embed) > 0: self.y0_adain_embed = y0_adain_embed.clone() B, HW, C = y0_adain_embed.shape LL, _, _, _ = haar_wavelet_decompose(y0_adain_embed.contiguous().view(B, C, h_len, w_len)) B_LL, C_LL, H_LL, W_LL = LL.shape #flat = rearrange(LL, 'b c h w -> b (h w) c') flat = LL.contiguous().view(B_LL, H_LL * W_LL, C_LL) f_s = flat[0] # assuming batch size 1 or using only the first self.mu_s = f_s.mean(dim=0, keepdim=True) f_s_centered = f_s - self.mu_s self.y0_color = self.whiten(f_s_centered, set=True) #self.y0_adain_embed = flat # cache if needed def get(self, denoised_embed: torch.Tensor, h_len, w_len, stylize_highfreq=False): B, HW, C = denoised_embed.shape denoised_embed = denoised_embed.contiguous().view(B, C, h_len, w_len) for i in range(B): x = denoised_embed[i:i+1] # [1, C, H, W] LL, LH, HL, HH = haar_wavelet_decompose(x) def process_band(band): Bc, Cc, Hc, Wc = band.shape flat = band.contiguous().view(Bc, Hc * Wc, Cc) styled = super(WaveletStyleWCT, self).get(flat) return styled.contiguous().view(Bc, Cc, Hc, Wc) LL_styled = process_band(LL) if stylize_highfreq: LH_styled = process_band(LH) HL_styled = process_band(HL) HH_styled = process_band(HH) else: LH_styled, HL_styled, HH_styled = LH, HL, HH recon = haar_wavelet_reconstruct(LL_styled, LH_styled, HL_styled, HH_styled) denoised_embed[i] = recon.squeeze(0) return denoised_embed.view(B, HW, C) def haar_wavelet_decompose(x): """ Orthonormal Haar decomposition. Input: [B, C, H, W] Output: LL, LH, HL, HH with shape [B, C, H//2, W//2] """ if x.dtype != torch.float32: x = x.float() B, C, H, W = x.shape assert H % 2 == 0 and W % 2 == 0, "Input must have even H, W" # Precompute norm = 1 / 2**0.5 x00 = x[:, :, 0::2, 0::2] x01 = x[:, :, 0::2, 1::2] x10 = x[:, :, 1::2, 0::2] x11 = x[:, :, 1::2, 1::2] LL = (x00 + x01 + x10 + x11) * norm * 0.5 LH = (x00 - x01 + x10 - x11) * norm * 0.5 HL = (x00 + x01 - x10 - x11) * norm * 0.5 HH = (x00 - x01 - x10 + x11) * norm * 0.5 return LL, LH, HL, HH def haar_wavelet_reconstruct(LL, LH, HL, HH): """ Orthonormal inverse Haar reconstruction. Input: LL, LH, HL, HH [B, C, H, W] Output: Reconstructed [B, C, H*2, W*2] """ norm = 1 / 2**0.5 B, C, H, W = LL.shape x00 = (LL + LH + HL + HH) * norm x01 = (LL - LH + HL - HH) * norm x10 = (LL + LH - HL - HH) * norm x11 = (LL - LH - HL + HH) * norm out = torch.zeros(B, C, H * 2, W * 2, device=LL.device, dtype=LL.dtype) out[:, :, 0::2, 0::2] = x00 out[:, :, 0::2, 1::2] = x01 out[:, :, 1::2, 0::2] = x10 out[:, :, 1::2, 1::2] = x11 return out """ class StyleFeatures: def __init__(self, dtype=torch.float64,): self.dtype = dtype def set(self, y0_adain_embed: torch.Tensor): def get(self, denoised_embed: torch.Tensor): return "Norpity McNerp" """ class Retrojector: def __init__(self, proj=None, patch_size=2, pinv_dtype=torch.float64, dtype=torch.float64, ENDO=False): self.proj = proj self.patch_size = patch_size self.pinv_dtype = pinv_dtype self.dtype = dtype self.LINEAR = isinstance(proj, nn.Linear) self.CONV2D = isinstance(proj, nn.Conv2d) self.CONV3D = isinstance(proj, nn.Conv3d) self.ENDO = ENDO self.W = proj.weight.data.to(dtype=pinv_dtype).cuda() if self.LINEAR: self.W_inv = torch.linalg.pinv(self.W.cuda()) elif self.CONV2D: C_out, _, kH, kW = proj.weight.shape W_flat = proj.weight.view(C_out, -1).to(dtype=pinv_dtype) self.W_inv = torch.linalg.pinv(W_flat.cuda()) if proj.bias is None: if self.LINEAR: bias_size = proj.out_features else: bias_size = proj.out_channels self.b = torch.zeros(bias_size, dtype=pinv_dtype, device=self.W_inv.device) else: self.b = proj.bias.data.to(dtype=pinv_dtype).to(self.W_inv.device) def embed(self, img: torch.Tensor): self.h = img.shape[-2] // self.patch_size self.w = img.shape[-1] // self.patch_size img = comfy.ldm.common_dit.pad_to_patch_size(img, (self.patch_size, self.patch_size)) if self.CONV2D: self.orig_shape = img.shape # for unembed img_embed = F.conv2d( img.to(self.W), weight=self.W, bias=self.b, stride=self.proj.stride, padding=self.proj.padding ) #img_embed = rearrange(img_embed, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=self.patch_size, pw=self.patch_size) img_embed = rearrange(img_embed, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=1, pw=1) elif self.LINEAR: if img.ndim == 4: img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=self.patch_size, pw=self.patch_size) if self.ENDO: img_embed = F.linear(img.to(self.b) - self.b, self.W_inv) else: img_embed = F.linear(img.to(self.W), self.W, self.b) return img_embed.to(img) def unembed(self, img_embed: torch.Tensor): if self.CONV2D: #img_embed = rearrange(img_embed, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=self.h, w=self.w, ph=self.patch_size, pw=self.patch_size) img_embed = rearrange(img_embed, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=self.h, w=self.w, ph=1, pw=1) img = self.invert_conv2d(img_embed) elif self.LINEAR: if self.ENDO: img = F.linear(img_embed.to(self.W), self.W, self.b) else: img = F.linear(img_embed.to(self.b) - self.b, self.W_inv) if img.ndim == 3: img = rearrange(img, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=self.h, w=self.w, ph=self.patch_size, pw=self.patch_size) return img.to(img_embed) def invert_conv2d(self, z: torch.Tensor,) -> torch.Tensor: z_dtype = z.dtype z = z.to(self.pinv_dtype) conv = self.proj B, C_in, H, W = self.orig_shape C_out, _, kH, kW = conv.weight.shape stride_h, stride_w = conv.stride pad_h, pad_w = conv.padding b = conv.bias.view(1, C_out, 1, 1).to(z) z_nobias = z - b #W_flat = conv.weight.view(C_out, -1).to(z) #W_pinv = torch.linalg.pinv(W_flat) Bz, Co, Hp, Wp = z_nobias.shape z_flat = z_nobias.reshape(Bz, Co, -1) x_patches = self.W_inv @ z_flat x_sum = F.fold( x_patches, output_size=(H + 2*pad_h, W+ 2*pad_w), kernel_size=(kH, kW), stride=(stride_h, stride_w), ) ones = torch.ones_like(x_patches) count = F.fold( ones, output_size=(H + 2*pad_h, W + 2*pad_w), kernel_size=(kH, kW), stride=(stride_h, stride_w), ) x_recon = x_sum / count.clamp(min=1e-6) if pad_h > 0 or pad_w > 0: x_recon = x_recon[..., pad_h:pad_h+H, pad_w:pad_w+W] return x_recon.to(z_dtype) def invert_patch_embedding(self, z: torch.Tensor, original_shape: torch.Size, grid_sizes: Optional[Tuple[int,int,int]] = None) -> torch.Tensor: B, C_in, D, H, W = original_shape pD, pH, pW = self.patch_size sD, sH, sW = pD, pH, pW if z.ndim == 3: # [B, S, C_out] -> reshape to [B, C_out, D', H', W'] S = z.shape[1] if grid_sizes is None: Dp = D // pD Hp = H // pH Wp = W // pW else: Dp, Hp, Wp = grid_sizes C_out = z.shape[2] z = z.transpose(1, 2).reshape(B, C_out, Dp, Hp, Wp) else: B2, C_out, Dp, Hp, Wp = z.shape assert B2 == B, "Batch size mismatch... ya sharked it." # kncokout bias b = self.patch_embedding.bias.view(1, C_out, 1, 1, 1) z_nobias = z - b # 2D filter -> pinv w3 = self.patch_embedding.weight # [C_out, C_in, 1, pH, pW] w2 = w3.squeeze(2) # [C_out, C_in, pH, pW] out_ch, in_ch, kH, kW = w2.shape W_flat = w2.view(out_ch, -1) # [C_out, in_ch*pH*pW] W_pinv = torch.linalg.pinv(W_flat) # [in_ch*pH*pW, C_out] # merge depth for 2D unfold wackiness z2 = z_nobias.permute(0,2,1,3,4).reshape(B*Dp, C_out, Hp, Wp) # apply pinv ... get patch vectors z_flat = z2.reshape(B*Dp, C_out, -1) # [B*Dp, C_out, L] x_patches = W_pinv @ z_flat # [B*Dp, in_ch*pH*pW, L] # fold -> spatial frames x2 = F.fold( x_patches, output_size=(H, W), kernel_size=(pH, pW), stride=(sH, sW) ) # → [B*Dp, C_in, H, W] # un-merge depth x2 = x2.reshape(B, Dp, in_ch, H, W) # [B, Dp, C_in, H, W] x_recon = x2.permute(0,2,1,3,4).contiguous() # [B, C_in, D, H, W] return x_recon def invert_conv2d( conv: torch.nn.Conv2d, z: torch.Tensor, original_shape: torch.Size, ) -> torch.Tensor: import torch.nn.functional as F B, C_in, H, W = original_shape C_out, _, kH, kW = conv.weight.shape stride_h, stride_w = conv.stride pad_h, pad_w = conv.padding if conv.bias is not None: b = conv.bias.view(1, C_out, 1, 1).to(z) z_nobias = z - b else: z_nobias = z W_flat = conv.weight.view(C_out, -1).to(z) W_pinv = torch.linalg.pinv(W_flat) Bz, Co, Hp, Wp = z_nobias.shape z_flat = z_nobias.reshape(Bz, Co, -1) x_patches = W_pinv @ z_flat x_sum = F.fold( x_patches, output_size=(H + 2*pad_h, W + 2*pad_w), kernel_size=(kH, kW), stride=(stride_h, stride_w), ) ones = torch.ones_like(x_patches) count = F.fold( ones, output_size=(H + 2*pad_h, W + 2*pad_w), kernel_size=(kH, kW), stride=(stride_h, stride_w), ) x_recon = x_sum / count.clamp(min=1e-6) if pad_h > 0 or pad_w > 0: x_recon = x_recon[..., pad_h:pad_h+H, pad_w:pad_w+W] return x_recon def adain_seq_inplace(content: torch.Tensor, style: torch.Tensor, dim=1, eps: float = 1e-7) -> torch.Tensor: mean_c = content.mean(dim, keepdim=True) std_c = content.std (dim, keepdim=True).add_(eps) # in-place add mean_s = style.mean (dim, keepdim=True) std_s = style.std (dim, keepdim=True).add_(eps) content.sub_(mean_c).div_(std_c).mul_(std_s).add_(mean_s) # in-place chain return content def adain_seq(content: torch.Tensor, style: torch.Tensor, eps: float = 1e-7) -> torch.Tensor: return ((content - content.mean(1, keepdim=True)) / (content.std(1, keepdim=True) + eps)) * (style.std(1, keepdim=True) + eps) + style.mean(1, keepdim=True) def apply_scattersort_tiled( denoised_spatial : torch.Tensor, y0_adain_spatial : torch.Tensor, tile_h : int, tile_w : int, pad : int, ): """ Apply spatial scattersort between denoised_spatial and y0_adain_spatial using local tile-wise sorted value matching. Args: denoised_spatial (Tensor): (B, C, H, W) tensor. y0_adain_spatial (Tensor): (B, C, H, W) reference tensor. tile_h (int): tile height. tile_w (int): tile width. pad (int): padding size to apply around tiles. Returns: denoised_embed (Tensor): (B, H*W, C) tensor after sortmatch. """ denoised_padded = F.pad(denoised_spatial, (pad, pad, pad, pad), mode='reflect') y0_padded = F.pad(y0_adain_spatial, (pad, pad, pad, pad), mode='reflect') denoised_padded_out = denoised_padded.clone() _, _, h_len, w_len = denoised_spatial.shape for ix in range(pad, h_len, tile_h): for jx in range(pad, w_len, tile_w): tile = denoised_padded[:, :, ix - pad:ix + tile_h + pad, jx - pad:jx + tile_w + pad] y0_tile = y0_padded[:, :, ix - pad:ix + tile_h + pad, jx - pad:jx + tile_w + pad] tile = rearrange(tile, "b c h w -> b c (h w)", h=tile_h + pad * 2, w=tile_w + pad * 2) y0_tile = rearrange(y0_tile, "b c h w -> b c (h w)", h=tile_h + pad * 2, w=tile_w + pad * 2) src_sorted, src_idx = tile.sort(dim=-1) ref_sorted, ref_idx = y0_tile.sort(dim=-1) new_tile = tile.scatter(dim=-1, index=src_idx, src=ref_sorted.expand(src_sorted.shape)) new_tile = rearrange(new_tile, "b c (h w) -> b c h w", h=tile_h + pad * 2, w=tile_w + pad * 2) denoised_padded_out[:, :, ix:ix + tile_h, jx:jx + tile_w] = ( new_tile if pad == 0 else new_tile[:, :, pad:-pad, pad:-pad] ) denoised_padded_out = denoised_padded_out if pad == 0 else denoised_padded_out[:, :, pad:-pad, pad:-pad] return denoised_padded_out def apply_scattersort_masked( denoised_embed : torch.Tensor, y0_adain_embed : torch.Tensor, y0_style_pos_mask : torch.Tensor | None, y0_style_pos_mask_edge : torch.Tensor | None, h_len : int, w_len : int ): if y0_style_pos_mask is None: flatmask = torch.ones((1,1,h_len,w_len)).bool().flatten().bool() else: flatmask = F.interpolate(y0_style_pos_mask, size=(h_len, w_len)).bool().flatten().cpu() flatunmask = ~flatmask if y0_style_pos_mask_edge is not None: edgemask = F.interpolate( y0_style_pos_mask_edge.unsqueeze(0), size=(h_len, w_len) ).bool().flatten() flatmask = flatmask & (~edgemask) flatunmask = flatunmask & (~edgemask) denoised_masked = denoised_embed[:, flatmask, :].clone() y0_adain_masked = y0_adain_embed[:, flatmask, :].clone() src_sorted, src_idx = denoised_masked.sort(dim=-2) ref_sorted, ref_idx = y0_adain_masked.sort(dim=-2) denoised_embed[:, flatmask, :] = src_sorted.scatter(dim=-2, index=src_idx, src=ref_sorted.expand(src_sorted.shape)) if (flatunmask == True).any(): denoised_unmasked = denoised_embed[:, flatunmask, :].clone() y0_adain_unmasked = y0_adain_embed[:, flatunmask, :].clone() src_sorted, src_idx = denoised_unmasked.sort(dim=-2) ref_sorted, ref_idx = y0_adain_unmasked.sort(dim=-2) denoised_embed[:, flatunmask, :] = src_sorted.scatter(dim=-2, index=src_idx, src=ref_sorted.expand(src_sorted.shape)) if y0_style_pos_mask_edge is not None: denoised_edgemasked = denoised_embed[:, edgemask, :].clone() y0_adain_edgemasked = y0_adain_embed[:, edgemask, :].clone() src_sorted, src_idx = denoised_edgemasked.sort(dim=-2) ref_sorted, ref_idx = y0_adain_edgemasked.sort(dim=-2) denoised_embed[:, edgemask, :] = src_sorted.scatter(dim=-2, index=src_idx, src=ref_sorted.expand(src_sorted.shape)) return denoised_embed def apply_scattersort( denoised_embed : torch.Tensor, y0_adain_embed : torch.Tensor, ): #src_sorted, src_idx = denoised_embed.cpu().sort(dim=-2) src_idx = denoised_embed.argsort(dim=-2) ref_sorted = y0_adain_embed.sort(dim=-2)[0] denoised_embed.scatter_(dim=-2, index=src_idx, src=ref_sorted.expand(ref_sorted.shape)) return denoised_embed def apply_scattersort_spatial( denoised_spatial : torch.Tensor, y0_adain_spatial : torch.Tensor, ): denoised_embed = rearrange(denoised_spatial, "b c h w -> b (h w) c") y0_adain_embed = rearrange(y0_adain_spatial, "b c h w -> b (h w) c") src_sorted, src_idx = denoised_embed.sort(dim=-2) ref_sorted, ref_idx = y0_adain_embed.sort(dim=-2) denoised_embed = src_sorted.scatter(dim=-2, index=src_idx, src=ref_sorted.expand(src_sorted.shape)) return rearrange(denoised_embed, "b (h w) c -> b c h w", h=denoised_spatial.shape[-2], w=denoised_spatial.shape[-1]) def apply_scattersort_spatial( x_spatial : torch.Tensor, y_spatial : torch.Tensor, ): x_emb = rearrange(x_spatial, "b c h w -> b (h w) c") y_emb = rearrange(y_spatial, "b c h w -> b (h w) c") x_sorted, x_idx = x_emb.sort(dim=-2) y_sorted, y_idx = y_emb.sort(dim=-2) x_emb = x_sorted.scatter(dim=-2, index=x_idx, src=y_sorted.expand(x_sorted.shape)) return rearrange(x_emb, "b (h w) c -> b c h w", h=x_spatial.shape[-2], w=x_spatial.shape[-1]) def apply_adain_spatial( x_spatial : torch.Tensor, y_spatial : torch.Tensor, ): x_emb = rearrange(x_spatial, "b c h w -> b (h w) c") y_emb = rearrange(y_spatial, "b c h w -> b (h w) c") x_mean = x_emb.mean(-2, keepdim=True) x_std = x_emb.std (-2, keepdim=True) y_mean = y_emb.mean(-2, keepdim=True) y_std = y_emb.std (-2, keepdim=True) assert (x_std == 0).any() == 0, "Target tensor has no variance!" assert (y_std == 0).any() == 0, "Reference tensor has no variance!" x_emb_adain = (x_emb - x_mean) / x_std x_emb_adain = (x_emb_adain * y_std) + y_mean return x_emb_adain.reshape_as(x_spatial) def adain_patchwise(content: torch.Tensor, style: torch.Tensor, sigma: float = 1.0, kernel_size: int = None, eps: float = 1e-5) -> torch.Tensor: # this one is really slow B, C, H, W = content.shape device = content.device dtype = content.dtype if kernel_size is None: kernel_size = int(2 * math.ceil(3 * sigma) + 1) if kernel_size % 2 == 0: kernel_size += 1 pad = kernel_size // 2 coords = torch.arange(kernel_size, dtype=torch.float64, device=device) - pad gauss = torch.exp(-0.5 * (coords / sigma) ** 2) gauss /= gauss.sum() kernel_2d = (gauss[:, None] * gauss[None, :]).to(dtype=dtype) weight = kernel_2d.view(1, 1, kernel_size, kernel_size) content_padded = F.pad(content, (pad, pad, pad, pad), mode='reflect') style_padded = F.pad(style, (pad, pad, pad, pad), mode='reflect') result = torch.zeros_like(content) for i in range(H): for j in range(W): c_patch = content_padded[:, :, i:i + kernel_size, j:j + kernel_size] s_patch = style_padded[:, :, i:i + kernel_size, j:j + kernel_size] w = weight.expand_as(c_patch) c_mean = (c_patch * w).sum(dim=(-1, -2), keepdim=True) c_std = ((c_patch - c_mean)**2 * w).sum(dim=(-1, -2), keepdim=True).sqrt() + eps s_mean = (s_patch * w).sum(dim=(-1, -2), keepdim=True) s_std = ((s_patch - s_mean)**2 * w).sum(dim=(-1, -2), keepdim=True).sqrt() + eps normed = (c_patch[:, :, pad:pad+1, pad:pad+1] - c_mean) / c_std stylized = normed * s_std + s_mean result[:, :, i, j] = stylized.squeeze(-1).squeeze(-1) return result def adain_patchwise_row_batch(content: torch.Tensor, style: torch.Tensor, sigma: float = 1.0, kernel_size: int = None, eps: float = 1e-5) -> torch.Tensor: B, C, H, W = content.shape device, dtype = content.device, content.dtype if kernel_size is None: kernel_size = int(2 * math.ceil(3 * sigma) + 1) if kernel_size % 2 == 0: kernel_size += 1 pad = kernel_size // 2 coords = torch.arange(kernel_size, dtype=torch.float64, device=device) - pad gauss = torch.exp(-0.5 * (coords / sigma) ** 2) gauss = (gauss / gauss.sum()).to(dtype) kernel_2d = (gauss[:, None] * gauss[None, :]) weight = kernel_2d.view(1, 1, kernel_size, kernel_size) content_padded = F.pad(content, (pad, pad, pad, pad), mode='reflect') style_padded = F.pad(style, (pad, pad, pad, pad), mode='reflect') result = torch.zeros_like(content) for i in range(H): c_row_patches = torch.stack([ content_padded[:, :, i:i+kernel_size, j:j+kernel_size] for j in range(W) ], dim=0) # [W, B, C, k, k] s_row_patches = torch.stack([ style_padded[:, :, i:i+kernel_size, j:j+kernel_size] for j in range(W) ], dim=0) w = weight.expand_as(c_row_patches[0]) c_mean = (c_row_patches * w).sum(dim=(-1, -2), keepdim=True) c_std = ((c_row_patches - c_mean) ** 2 * w).sum(dim=(-1, -2), keepdim=True).sqrt() + eps s_mean = (s_row_patches * w).sum(dim=(-1, -2), keepdim=True) s_std = ((s_row_patches - s_mean) ** 2 * w).sum(dim=(-1, -2), keepdim=True).sqrt() + eps center = kernel_size // 2 central = c_row_patches[:, :, :, center:center+1, center:center+1] normed = (central - c_mean) / c_std stylized = normed * s_std + s_mean result[:, :, i, :] = stylized.squeeze(-1).squeeze(-1).permute(1, 2, 0) # [B,C,W] return result def adain_patchwise_row_batch_med(content: torch.Tensor, style: torch.Tensor, sigma: float = 1.0, kernel_size: int = None, eps: float = 1e-5, mask: torch.Tensor = None, use_median_blur: bool = False, lowpass_weight=1.0, highpass_weight=1.0) -> torch.Tensor: B, C, H, W = content.shape device, dtype = content.device, content.dtype if kernel_size is None: kernel_size = int(2 * math.ceil(3 * abs(sigma)) + 1) if kernel_size % 2 == 0: kernel_size += 1 pad = kernel_size // 2 content_padded = F.pad(content, (pad, pad, pad, pad), mode='reflect') style_padded = F.pad(style, (pad, pad, pad, pad), mode='reflect') result = torch.zeros_like(content) scaling = torch.ones((B, 1, H, W), device=device, dtype=dtype) sigma_scale = torch.ones((H, W), device=device, dtype=torch.float32) if mask is not None: with torch.no_grad(): padded_mask = F.pad(mask.float(), (pad, pad, pad, pad), mode="reflect") blurred_mask = F.avg_pool2d(padded_mask, kernel_size=kernel_size, stride=1, padding=pad) blurred_mask = blurred_mask[..., pad:-pad, pad:-pad] edge_proximity = blurred_mask * (1.0 - blurred_mask) scaling = 1.0 - (edge_proximity / 0.25).clamp(0.0, 1.0) sigma_scale = scaling[0, 0] # assuming single-channel mask broadcasted across B, C if not use_median_blur: coords = torch.arange(kernel_size, dtype=torch.float64, device=device) - pad base_gauss = torch.exp(-0.5 * (coords / sigma) ** 2) base_gauss = (base_gauss / base_gauss.sum()).to(dtype) gaussian_table = {} for s in sigma_scale.unique(): sig = float((sigma * s + eps).clamp(min=1e-3)) gauss_local = torch.exp(-0.5 * (coords / sig) ** 2) gauss_local = (gauss_local / gauss_local.sum()).to(dtype) kernel_2d = gauss_local[:, None] * gauss_local[None, :] gaussian_table[s.item()] = kernel_2d for i in range(H): row_result = torch.zeros(B, C, W, dtype=dtype, device=device) for j in range(W): c_patch = content_padded[:, :, i:i+kernel_size, j:j+kernel_size] s_patch = style_padded[:, :, i:i+kernel_size, j:j+kernel_size] if use_median_blur: # Median blur with residual restoration unfolded_c = c_patch.reshape(B, C, -1) unfolded_s = s_patch.reshape(B, C, -1) c_median = unfolded_c.median(dim=-1, keepdim=True).values s_median = unfolded_s.median(dim=-1, keepdim=True).values center = kernel_size // 2 central = c_patch[:, :, center, center].view(B, C, 1) residual = central - c_median stylized = lowpass_weight * s_median + residual * highpass_weight else: k = gaussian_table[float(sigma_scale[i, j].item())] local_weight = k.view(1, 1, kernel_size, kernel_size).expand(B, C, kernel_size, kernel_size) c_mean = (c_patch * local_weight).sum(dim=(-1, -2), keepdim=True) c_std = ((c_patch - c_mean) ** 2 * local_weight).sum(dim=(-1, -2), keepdim=True).sqrt() + eps s_mean = (s_patch * local_weight).sum(dim=(-1, -2), keepdim=True) s_std = ((s_patch - s_mean) ** 2 * local_weight).sum(dim=(-1, -2), keepdim=True).sqrt() + eps center = kernel_size // 2 central = c_patch[:, :, center:center+1, center:center+1] normed = (central - c_mean) / c_std stylized = normed * s_std + s_mean local_scaling = scaling[:, :, i, j].view(B, 1, 1) stylized = central * (1 - local_scaling) + stylized * local_scaling row_result[:, :, j] = stylized.squeeze(-1) result[:, :, i, :] = row_result return result def weighted_mix_n(tensor_list, weight_list, dim=-1, offset=0): assert all(t.shape == tensor_list[0].shape for t in tensor_list) assert len(tensor_list) == len(weight_list) total_weight = sum(weight_list) ratios = [w / total_weight for w in weight_list] length = tensor_list[0].shape[dim] idx = torch.arange(length) # Create a bin index tensor based on weighted slots float_bins = (idx + offset) * len(ratios) / length bin_idx = torch.floor(float_bins).long() % len(ratios) # Allocate slots based on ratio using a cyclic pattern counters = [0.0 for _ in ratios] slots = torch.empty_like(idx) for i in range(length): # Assign to the group that's most under-allocated expected = [r * (i + 1) for r in ratios] errors = [expected[j] - counters[j] for j in range(len(ratios))] k = max(range(len(errors)), key=lambda j: errors[j]) slots[i] = k counters[k] += 1 # Create mask for each tensor out = tensor_list[0].clone() for i, tensor in enumerate(tensor_list): mask = slots == i while mask.dim() < tensor.dim(): mask = mask.unsqueeze(0) mask = mask.expand_as(tensor) out = torch.where(mask, tensor, out) return out from torch import vmap BLOCK_NAMES = {"double_blocks", "single_blocks", "up_blocks", "middle_blocks", "down_blocks", "input_blocks", "output_blocks"} DEFAULT_BLOCK_WEIGHTS_MMDIT = { "attn_norm" : 0.0, "attn_norm_mod": 0.0, "attn" : 1.0, "attn_gated" : 0.0, "attn_res" : 1.0, "ff_norm" : 0.0, "ff_norm_mod" : 0.0, "ff" : 1.0, "ff_gated" : 0.0, "ff_res" : 1.0, "h_tile" : 8, "w_tile" : 8, } DEFAULT_ATTN_WEIGHTS_MMDIT = { "q_proj": 0.0, "k_proj": 0.0, "v_proj": 1.0, "q_norm": 0.0, "k_norm": 0.0, "out" : 1.0, "h_tile": 8, "w_tile": 8, } DEFAULT_BASE_WEIGHTS_MMDIT = { "proj_in" : 1.0, "proj_out": 1.0, "h_tile" : 8, "w_tile" : 8, } class Stylizer: buffer = {} CLS_WCT = StyleWCT() CLS_WCT2 = WaveletStyleWCT() def __init__(self, dtype=torch.float64, device=torch.device("cuda")): self.dtype = dtype self.device = device self.mask = [None] self.apply_to = [""] self.method = ["passthrough"] self.h_tile = [-1] self.w_tile = [-1] self.w_len = 0 self.h_len = 0 self.img_len = 0 self.IMG_1ST = True self.HEADS = 0 self.KONTEXT = 0 def set_mode(self, mode): self.method = [mode] #[getattr(self, mode)] def set_weights(self, **kwargs): for k, v in kwargs.items(): if hasattr(self, k): setattr(self, k, [v]) def set_weights_recursive(self, **kwargs): for name, val in kwargs.items(): if hasattr(self, name): setattr(self, name, [val]) for attr_name, attr_val in vars(self).items(): if isinstance(attr_val, Stylizer): attr_val.set_weights_recursive(**kwargs) for list_name in BLOCK_NAMES: lst = getattr(self, list_name, None) if isinstance(lst, list): for element in lst: if isinstance(element, Stylizer): element.set_weights_recursive(**kwargs) def merge_weights(self, other): def recursive_merge(a, b, path): if isinstance(a, list) and isinstance(b, list): if path in BLOCK_NAMES: out = [] for i in range(max(len(a), len(b))): if i < len(a) and i < len(b): out.append(recursive_merge(a[i], b[i], path=None)) elif i < len(a): out.append(a[i]) else: out.append(b[i]) return out return a + b if isinstance(a, dict) and isinstance(b, dict): merged = dict(a) for k, v_b in b.items(): if k in merged: merged[k] = recursive_merge(merged[k], v_b, path=None) else: merged[k] = v_b return merged if hasattr(a, "__dict__") and hasattr(b, "__dict__"): for attr, val_b in vars(b).items(): val_a = getattr(a, attr, None) if val_a is not None: setattr(a, attr, recursive_merge(val_a, val_b, path=attr)) else: setattr(a, attr, val_b) return a return b for attr in vars(self): if attr in BLOCK_NAMES: merged = recursive_merge(getattr(self, attr), getattr(other, attr, []), path=attr) elif hasattr(other, attr): merged = recursive_merge(getattr(self, attr), getattr(other, attr), path=attr) else: continue setattr(self, attr, merged) def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): self.h_len = h_len self.w_len = w_len self.img_slice = img_slice self.txt_slice = txt_slice self.img_len = h_len * w_len self.HEADS = HEADS @staticmethod def middle_slice(length, weight): """ Returns a slice object that selects the middle `weight` fraction of a dimension. Example: weight=1.0 → full slice; weight=0.5 → middle 50% """ if weight >= 1.0: return slice(None) wr = int((length * (1 - weight)) // 2) return slice(wr, -wr if wr > 0 else None) @staticmethod def get_outer_slice(x, weight): if weight >= 0.0: return x length = x.shape[-2] wr = int((length * (1 - (-weight))) // 2) return torch.cat([x[...,:wr,:], x[...,-wr:,:]], dim=-2) @staticmethod def restore_outer_slice(x, x_outer, weight): if weight >= 0.0: return x length = x.shape[-2] wr = int((length * (1 - (-weight))) // 2) x[...,:wr,:] = x_outer[...,:wr,:] x[...,-wr:,:] = x_outer[...,-wr:,:] return x def __call__(self, x, attr): if x.shape[0] == 1 and not self.KONTEXT: return x weight_list = getattr(self, attr) weights_all_zero = all(weight == 0.0 for weight in weight_list) if weights_all_zero: return x #self.HEADS=24 #x_ndim = x.ndim #if x_ndim == 3: # B, HW, C = x.shape # if x.shape[-2] != self.HEADS and self.HEADS != 0: # x = x.reshape(B,self.HEADS,HW,-1) HEAD_DIM = x.shape[1] if HEAD_DIM == self.HEADS: B, HEAD_DIM, HW, C = x.shape x = x.reshape(B, HW, C*HEAD_DIM) if hasattr(self, "KONTEXT") and self.KONTEXT == 1: x = x.reshape(2, x.shape[1] // 2, x.shape[2]) txt_slice, img_slice, ktx_slice = self.txt_slice, self.img_slice, None if hasattr(self, "KONTEXT") and self.KONTEXT == 2: ktx_slice = self.img_slice # slice(2 * self.img_slice.start, None) img_slice = slice(2 * self.img_slice.start, self.img_slice.start) txt_slice = slice(None, 2 * self.txt_slice.stop) weights_all_one = all(weight == 1.0 for weight in weight_list) methods_all_scattersort = all(name == "scattersort" for name in self.method) masks_all_none = all(mask is None for mask in self.mask) if weights_all_one and methods_all_scattersort and len(weight_list) > 1 and masks_all_none: buf = Stylizer.buffer buf['src_idx'] = x[0:1].argsort(dim=-2) buf['ref_sorted'], buf['ref_idx'] = x[1:].reshape(1, -1, x.shape[-1]).sort(dim=-2) buf['src'] = buf['ref_sorted'][:,::len(weight_list)].expand_as(buf['src_idx']) # interleave_stride = len(weight_list) x[0:1] = x[0:1].scatter_(dim=-2, index=buf['src_idx'], src=buf['src'],) else: for i, (weight, mask) in enumerate(zip(weight_list, self.mask)): if mask is not None: x01 = x[0:1].clone() slc = Stylizer.middle_slice(x.shape[-2], weight) #slc = slice(None) txt_method_name = self.method[i].removeprefix("tiled_") txt_method = getattr(self, txt_method_name) method_name = self.method[i].removeprefix("tiled_") if self.img_len > x.shape[-2] or self.h_len < 0 else self.method[i] method = getattr(self, method_name) apply_to = self.apply_to[i] if weight == 0.0: continue else: # if weight == 1.0: if weight > 0 and weight < 1: x_clone = x.clone() if self.img_len == x.shape[-2] or apply_to == "img+txt" or self.h_len < 0: x = method(x, idx=i+1, slc=slc) elif self.img_len < x.shape[-2]: if "img" in apply_to: x[...,img_slice,:] = method(x[...,img_slice,:], idx=i+1, slc=slc) #if ktx_slice is not None: # x[...,ktx_slice,:] = method(x[...,ktx_slice,:], idx=i+1) #x[:,:self.img_len,:] = method(x[:,:self.img_len,:], idx=i+1) if "txt" in apply_to: x[...,txt_slice,:] = txt_method(x[...,txt_slice,:], idx=i+1, slc=slc) #x[:,self.img_len:,:] = method(x[:,self.img_len:,:], idx=i+1) if not "img" in apply_to and not "txt" in apply_to: pass else: x = method(x, idx=i+1, slc=slc) if weight > 0 and weight < 1 and txt_method_name != "scattersort": x = torch.lerp(x_clone, x, weight) #else: # x = torch.lerp(x, method(x.clone(), idx=i+1), weight) if mask is not None: x[0:1,...,img_slice,:] = torch.lerp(x01[...,img_slice,:], x[0:1,...,img_slice,:], mask.view(1, -1, 1)) if ktx_slice is not None: x[0:1,...,ktx_slice,:] = torch.lerp(x01[...,ktx_slice,:], x[0:1,...,ktx_slice,:], mask.view(1, -1, 1)) #x[0:1,:self.img_len] = torch.lerp(x01[:,:self.img_len], x[0:1,:self.img_len], mask.view(1, -1, 1)) #if x_ndim == 3: # return x.view(B,HW,C) if hasattr(self, "KONTEXT") and self.KONTEXT == 1: x = x.reshape(1, x.shape[1] * 2, x.shape[2]) if HEAD_DIM == self.HEADS: return x.reshape(B, HEAD_DIM, HW, C) else: return x def WCT(self, x, idx=1): Stylizer.CLS_WCT.set(x[idx:idx+1]) x[0:1] = Stylizer.CLS_WCT.get(x[0:1]) return x def WCT2(self, x, idx=1): Stylizer.CLS_WCT2.set(x[idx:idx+1], self.h_len, self.w_len) x[0:1] = Stylizer.CLS_WCT2.get(x[0:1], self.h_len, self.w_len) return x @staticmethod def AdaIN_(x, y, eps: float = 1e-7) -> torch.Tensor: mean_c = x.mean(-2, keepdim=True) std_c = x.std (-2, keepdim=True).add_(eps) # in-place add mean_s = y.mean (-2, keepdim=True) std_s = y.std (-2, keepdim=True).add_(eps) x.sub_(mean_c).div_(std_c).mul_(std_s).add_(mean_s) # in-place chain return x def AdaIN(self, x, idx=1, eps: float = 1e-7) -> torch.Tensor: mean_c = x[0:1].mean(-2, keepdim=True) std_c = x[0:1].std (-2, keepdim=True).add_(eps) # in-place add mean_s = x[idx:idx+1].mean (-2, keepdim=True) std_s = x[idx:idx+1].std (-2, keepdim=True).add_(eps) x[0:1].sub_(mean_c).div_(std_c).mul_(std_s).add_(mean_s) # in-place chain return x def injection(self, x:torch.Tensor, idx=1) -> torch.Tensor: x[0:1] = x[idx:idx+1] return x @staticmethod def injection_(x:torch.Tensor, y:torch.Tensor) -> torch.Tensor: return y @staticmethod def passthrough(x:torch.Tensor, idx=1) -> torch.Tensor: return x @staticmethod def decompose_magnitude_direction(x, dim=-1, eps=1e-8): magnitude = x.norm(p=2, dim=dim, keepdim=True) direction = x / (magnitude + eps) return magnitude, direction @staticmethod def scattersort_dir_(x, y, dim=-2): #buf = Stylizer.buffer #buf['src_sorted'], buf['src_idx'] = x.sort(dim=-2) #buf['ref_sorted'], buf['ref_idx'] = y.sort(dim=-2) #mag, _ = Stylizer.decompose_magnitude_direction(buf['src_sorted'], dim) #_, dir = Stylizer.decompose_magnitude_direction(buf['ref_sorted'], dim) mag, _ = Stylizer.decompose_magnitude_direction(x.to(torch.float64), dim) buf = Stylizer.buffer buf['src_idx'] = x.argsort(dim=-2) buf['ref_sorted'], buf['ref_idx'] = y .sort(dim=-2) x.scatter_(dim=-2, index=buf['src_idx'], src=buf['ref_sorted'].expand_as(buf['src_idx'])) _, dir = Stylizer.decompose_magnitude_direction(x.to(torch.float64), dim) return (mag * dir).to(x) @staticmethod def scattersort_dir2_(x, y, dim=-2): #buf = Stylizer.buffer #buf['src_sorted'], buf['src_idx'] = x.sort(dim=-2) #buf['ref_sorted'], buf['ref_idx'] = y.sort(dim=-2) #mag, _ = Stylizer.decompose_magnitude_direction(buf['src_sorted'], dim) #_, dir = Stylizer.decompose_magnitude_direction(buf['ref_sorted'], dim) buf = Stylizer.buffer buf['src_sorted'], buf['src_idx'] = x.sort(dim=dim) buf['ref_sorted'], buf['ref_idx'] = y.sort(dim=dim) buf['x_sub'], buf['x_sub_idx'] = buf['src_sorted'].sort(dim=-1) buf['y_sub'], buf['y_sub_idx'] = buf['ref_sorted'].sort(dim=-1) mag, _ = Stylizer.decompose_magnitude_direction(buf['x_sub'].to(torch.float64), -1) _, dir = Stylizer.decompose_magnitude_direction(buf['y_sub'].to(torch.float64), -1) buf['y_sub'] = (mag * dir).to(x) buf['ref_sorted'].scatter_(dim=-1, index=buf['y_sub_idx'], src=buf['y_sub'].expand_as(buf['y_sub_idx'])) mag, _ = Stylizer.decompose_magnitude_direction(buf['src_sorted'].to(torch.float64), dim) _, dir = Stylizer.decompose_magnitude_direction(buf['ref_sorted'].to(torch.float64), dim) buf['ref_sorted'] = (mag * dir).to(x) x.scatter_(dim=dim, index=buf['src_idx'], src=buf['ref_sorted'].expand_as(buf['src_idx'])) return x @staticmethod def scattersort_dir(x, idx=1): x[0:1] = Stylizer.scattersort_dir_(x[0:1], x[idx:idx+1]) return x @staticmethod def scattersort_dir2(x, idx=1): x[0:1] = Stylizer.scattersort_dir2_(x[0:1], x[idx:idx+1]) return x @staticmethod def scattersort_(x, y, slc=slice(None)): buf = Stylizer.buffer buf['src_idx'] = x.argsort(dim=-2) buf['ref_sorted'], buf['ref_idx'] = y .sort(dim=-2) return x.scatter_(dim=-2, index=buf['src_idx'][...,slc,:], src=buf['ref_sorted'][...,slc,:].expand_as(buf['src_idx'][...,slc,:])) @staticmethod def scattersort_double(x, y): buf = Stylizer.buffer buf['src_sorted'], buf['src_idx'] = x.sort(dim=-2) buf['ref_sorted'], buf['ref_idx'] = y.sort(dim=-2) buf['x_sub_idx'] = buf['src_sorted'].argsort(dim=-1) buf['y_sub'], buf['y_sub_idx'] = buf['ref_sorted'].sort(dim=-1) x.scatter_(dim=-1, index=buf['x_sub_idx'], src=buf['y_sub'].expand_as(buf['x_sub_idx'])) return x.scatter_(dim=-2, index=buf['src_idx'], src=buf['ref_sorted'].expand_as(buf['src_idx'])) def scattersort_aoeu(self, x, idx=1, slc=slice(None)): x[0:1] = Stylizer.scattersort_(x[0:1], x[idx:idx+1], slc) return x def scattersort(self, x, idx=1, slc=slice(None)): if x.shape[0] != 2: x[0:1] = Stylizer.scattersort_(x[0:1], x[idx:idx+1], slc) return x buf = Stylizer.buffer buf['sorted'], buf['idx'] = x.sort(dim=-2) return x.scatter_(dim=-2, index=buf['idx'][0:1][...,slc,:], src=buf['sorted'][1:2][...,slc,:].expand_as(buf['idx'][0:1][...,slc,:])) def tiled_scattersort(self, x, idx=1): #, h_tile=None, w_tile=None): #if HDModel.RECON_MODE: # return denoised_embed #den = x[0:1] [:,:self.img_len,:].view(-1, 2560, self.h_len, self.w_len) #style = x[idx:idx+1][:,:self.img_len,:].view(-1, 2560, self.h_len, self.w_len) #h_tile = self.h_tile[idx-1] if h_tile is None else h_tile #w_tile = self.w_tile[idx-1] if w_tile is None else w_tile C = x.shape[-1] den = x[0:1] [:,self.img_slice,:].reshape(-1, C, self.h_len, self.w_len) style = x[idx:idx+1][:,self.img_slice,:].reshape(-1, C, self.h_len, self.w_len) tiles = Stylizer.get_tiles_as_strided(den, self.h_tile[idx-1], self.w_tile[idx-1]) ref_tile = Stylizer.get_tiles_as_strided(style, self.h_tile[idx-1], self.w_tile[idx-1]) # rearrange for vmap to run on (nH, nW) ( as outer axes) tiles_v = tiles .permute(2, 3, 0, 1, 4, 5) # (nH, nW, B, C, tile_h, tile_w) ref_tile_v = ref_tile.permute(2, 3, 0, 1, 4, 5) # (nH, nW, B, C, tile_h, tile_w) # vmap over spatial dimms (nH, nW)... num of tiles high, num tiles wide vmap2 = torch.vmap(torch.vmap(Stylizer.apply_scattersort_per_tile, in_dims=0), in_dims=0) result = vmap2(tiles_v, ref_tile_v) # (nH, nW, B, C, tile_h, tile_w) # --> (B, C, nH, nW, tile_h, tile_w) result = result.permute(2, 3, 0, 1, 4, 5) #( B, C, nH, nW, tile_h, tile_w) # in-place copy, werx if result has same shape/strides as tiles... overwrites same mem location "content" is using tiles.copy_(result) return x def tiled_AdaIN(self, x, idx=1): #if HDModel.RECON_MODE: # return denoised_embed #den = x[0:1] [:,:self.img_len,:].view(-1, 2560, self.h_len, self.w_len) #style = x[idx:idx+1][:,:self.img_len,:].view(-1, 2560, self.h_len, self.w_len) C = x.shape[-1] den = x[0:1] [:,self.img_slice,:].reshape(-1, C, self.h_len, self.w_len) style = x[idx:idx+1][:,self.img_slice,:].reshape(-1, C, self.h_len, self.w_len) tiles = Stylizer.get_tiles_as_strided(den, self.h_tile[idx-1], self.w_tile[idx-1]) ref_tile = Stylizer.get_tiles_as_strided(style, self.h_tile[idx-1], self.w_tile[idx-1]) # rearrange for vmap to run on (nH, nW) ( as outer axes) tiles_v = tiles .permute(2, 3, 0, 1, 4, 5) # (nH, nW, B, C, tile_h, tile_w) ref_tile_v = ref_tile.permute(2, 3, 0, 1, 4, 5) # (nH, nW, B, C, tile_h, tile_w) # vmap over spatial dimms (nH, nW)... num of tiles high, num tiles wide vmap2 = torch.vmap(torch.vmap(Stylizer.apply_AdaIN_per_tile, in_dims=0), in_dims=0) result = vmap2(tiles_v, ref_tile_v) # (nH, nW, B, C, tile_h, tile_w) # --> (B, C, nH, nW, tile_h, tile_w) result = result.permute(2, 3, 0, 1, 4, 5) #( B, C, nH, nW, tile_h, tile_w) # in-place copy, werx if result has same shape/strides as tiles... overwrites same mem location "content" is using tiles.copy_(result) return x @staticmethod def get_tiles_as_strided(x, tile_h, tile_w): B, C, H, W = x.shape stride = x.stride() nH = H // tile_h nW = W // tile_w tiles = x.as_strided( size=(B, C, nH, nW, tile_h, tile_w), stride=(stride[0], stride[1], stride[2] * tile_h, stride[3] * tile_w, stride[2], stride[3]) ) return tiles # shape: (B, C, nH, nW, tile_h, tile_w) @staticmethod def apply_scattersort_per_tile(tile, ref_tile): flat = tile .flatten(-2, -1) ref_flat = ref_tile.flatten(-2, -1) sorted_ref, _ = ref_flat .sort(dim=-1) src_sorted, src_idx = flat.sort(dim=-1) out = flat.scatter(dim=-1, index=src_idx, src=sorted_ref) return out.view_as(tile) @staticmethod def apply_AdaIN_per_tile(tile, ref_tile, eps: float = 1e-7): mean_c = tile.mean(-2, keepdim=True) std_c = tile.std (-2, keepdim=True).add_(eps) # in-place add mean_s = ref_tile.mean (-2, keepdim=True) std_s = ref_tile.std (-2, keepdim=True).add_(eps) tile.sub_(mean_c).div_(std_c).mul_(std_s).add_(mean_s) # in-place chain return tile class StyleMMDiT_Attn(Stylizer): def __init__(self, mode): super().__init__() self.q_proj = [0.0] self.k_proj = [0.0] self.v_proj = [0.0] self.q_norm = [0.0] self.k_norm = [0.0] self.out = [0.0] class StyleMMDiT_FF(Stylizer): # these hit img or joint only, never txt def __init__(self, mode): super().__init__() self.ff_1 = [0.0] self.ff_1_silu = [0.0] self.ff_3 = [0.0] self.ff_13 = [0.0] self.ff_2 = [0.0] class StyleMMDiT_MoE(Stylizer): # these hit img or joint only, never txt def __init__(self, mode): super().__init__() self.FF_SHARED = StyleMMDiT_FF(mode) self.FF_SEPARATE = StyleMMDiT_FF(mode) self.shared = [0.0] self.gate = [False] self.topk_weight = [0.0] self.separate = [0.0] self.sum = [0.0] self.out = [0.0] class StyleMMDiT_SubBlock(Stylizer): def __init__(self, mode): super().__init__() self.ATTN = StyleMMDiT_Attn(mode) # options for attn itself: qkv proj, qk norm, attn out self.attn_norm = [0.0] self.attn_norm_mod = [0.0] self.attn = [0.0] self.attn_gated = [0.0] self.attn_res = [0.0] self.ff_norm = [0.0] self.ff_norm_mod = [0.0] self.ff = [0.0] self.ff_gated = [0.0] self.ff_res = [0.0] self.mask = [None] def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.ATTN.set_len(h_len, w_len, img_slice, txt_slice, HEADS) class StyleMMDiT_IMG_Block(StyleMMDiT_SubBlock): # img or joint def __init__(self, mode): super().__init__(mode) self.FF = StyleMMDiT_MoE(mode) # options for MoE if img or joint def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.FF.set_len(h_len, w_len, img_slice, txt_slice, HEADS) class StyleMMDiT_TXT_Block(StyleMMDiT_SubBlock): # txt only def __init__(self, mode): super().__init__(mode) self.FF = StyleMMDiT_FF(mode) # options for FF within MoE for img or joint; or for txt alone def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.FF.set_len(h_len, w_len, img_slice, txt_slice, HEADS) class StyleMMDiT_BaseBlock: def __init__(self, mode="passthrough"): self.img = StyleMMDiT_IMG_Block(mode) self.txt = StyleMMDiT_TXT_Block(mode) self.mask = [None] self.attn_mask = [None] def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): self.h_len = h_len self.w_len = w_len self.img_len = h_len * w_len self.img_slice = img_slice self.txt_slice = txt_slice self.HEADS = HEADS self.img.set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.txt.set_len(-1, -1, img_slice, txt_slice, HEADS) for i, mask in enumerate(self.mask): if mask is not None and mask.ndim > 1: self.mask[i] = F.interpolate(mask.unsqueeze(0), size=(h_len, w_len)).flatten().to(torch.bfloat16).cuda() self.img.mask = self.mask for i, mask in enumerate(self.attn_mask): if mask is not None and mask.ndim > 1: self.attn_mask[i] = F.interpolate(mask.unsqueeze(0), size=(h_len, w_len)).flatten().to(torch.bfloat16).cuda() self.img.ATTN.mask = self.attn_mask class StyleMMDiT_DoubleBlock(StyleMMDiT_BaseBlock): def __init__(self, mode="passthrough"): super().__init__(mode) self.txt = StyleMMDiT_TXT_Block(mode) def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.txt.set_len(-1, -1, img_slice, txt_slice, HEADS) class StyleMMDiT_SingleBlock(StyleMMDiT_BaseBlock): def __init__(self, mode="passthrough"): super().__init__(mode) class StyleUNet_Resample(Stylizer): def __init__(self, mode): super().__init__() self.conv = [0.0] class StyleUNet_Attn(Stylizer): def __init__(self, mode): super().__init__() self.q_proj = [0.0] self.k_proj = [0.0] self.v_proj = [0.0] self.out = [0.0] class StyleUNet_FF(Stylizer): def __init__(self, mode): super().__init__() self.proj = [0.0] self.geglu = [0.0] self.linear = [0.0] class StyleUNet_TransformerBlock(Stylizer): def __init__(self, mode): super().__init__() self.ATTN1 = StyleUNet_Attn(mode) # self-attn self.FF = StyleUNet_FF (mode) self.ATTN2 = StyleUNet_Attn(mode) # cross-attn self.self_attn = [0.0] self.ff = [0.0] self.cross_attn = [0.0] self.self_attn_res = [0.0] self.cross_attn_res = [0.0] self.ff_res = [0.0] self.norm1 = [0.0] self.norm2 = [0.0] self.norm3 = [0.0] def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.ATTN1.set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.ATTN2.set_len(h_len, w_len, img_slice, txt_slice, HEADS) class StyleUNet_SpatialTransformer(Stylizer): def __init__(self, mode): super().__init__() self.TFMR = StyleUNet_TransformerBlock(mode) self.spatial_norm_in = [0.0] self.spatial_proj_in = [0.0] self.spatial_transformer_block = [0.0] self.spatial_transformer = [0.0] self.spatial_proj_out = [0.0] self.spatial_res = [0.0] def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.TFMR.set_len(h_len, w_len, img_slice, txt_slice, HEADS) class StyleUNet_ResBlock(Stylizer): def __init__(self, mode): super().__init__() self.in_norm = [0.0] self.in_silu = [0.0] self.in_conv = [0.0] self.emb_silu = [0.0] self.emb_linear = [0.0] self.emb_res = [0.0] self.out_norm = [0.0] self.out_silu = [0.0] self.out_conv = [0.0] self.residual = [0.0] class StyleUNet_BaseBlock(Stylizer): def __init__(self, mode="passthrough"): self.resample_block = StyleUNet_Resample(mode) self.res_block = StyleUNet_ResBlock(mode) self.spatial_block = StyleUNet_SpatialTransformer(mode) self.resample = [0.0] self.res = [0.0] self.spatial = [0.0] self.mask = [None] self.attn_mask = [None] self.KONTEXT = 0 def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): self.h_len = h_len self.w_len = w_len self.img_len = h_len * w_len self.img_slice = img_slice self.txt_slice = txt_slice self.HEADS = HEADS self.resample_block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.res_block .set_len(h_len, w_len, img_slice, txt_slice, HEADS) self.spatial_block .set_len(h_len, w_len, img_slice, txt_slice, HEADS) for i, mask in enumerate(self.mask): if mask is not None and mask.ndim > 1: self.mask[i] = F.interpolate(mask.unsqueeze(0), size=(h_len, w_len)).flatten().to(torch.bfloat16).cuda() self.resample_block.mask = self.mask self.res_block.mask = self.mask self.spatial_block.mask = self.mask self.spatial_block.TFMR.mask = self.mask for i, mask in enumerate(self.attn_mask): if mask is not None and mask.ndim > 1: self.attn_mask[i] = F.interpolate(mask.unsqueeze(0), size=(h_len, w_len)).flatten().to(torch.bfloat16).cuda() self.spatial_block.TFMR.ATTN1.mask = self.attn_mask def __call__(self, x, attr): B, C, H, W = x.shape x = super().__call__(x.reshape(B, H*W, C), attr) return x.reshape(B,C,H,W) class StyleUNet_InputBlock(StyleUNet_BaseBlock): def __init__(self, mode="passthrough"): super().__init__(mode) class StyleUNet_MiddleBlock(StyleUNet_BaseBlock): def __init__(self, mode="passthrough"): super().__init__(mode) class StyleUNet_OutputBlock(StyleUNet_BaseBlock): def __init__(self, mode="passthrough"): super().__init__(mode) class Style_Model(Stylizer): def __init__(self, dtype=torch.float64, device=torch.device("cuda")): super().__init__(dtype, device) self.guides = [] self.GUIDES_INITIALIZED = False #self.double_blocks = [StyleMMDiT_DoubleBlock() for _ in range(100)] #self.single_blocks = [StyleMMDiT_SingleBlock() for _ in range(100)] self.h_len = -1 self.w_len = -1 self.img_len = -1 self.h_tile = [-1] self.w_tile = [-1] self.proj_in = [0.0] # these are for img only! not sliced self.proj_out = [0.0] self.cond_pos = [None] self.cond_neg = [None] self.noise_mode = "update" self.recon_lure = "none" self.data_shock = "none" self.data_shock_start_step = 0 self.data_shock_end_step = 0 self.Retrojector = None self.Endojector = None self.IMG_1ST = True self.HEADS = 0 self.KONTEXT = 0 def __call__(self, x, attr): if x.shape[0] == 1 and not self.KONTEXT: return x weight_list = getattr(self, attr) weights_all_zero = all(weight == 0.0 for weight in weight_list) if weights_all_zero: return x """x_ndim = x.ndim if x_ndim == 4: B, HEAD, HW, C = x.shape if x_ndim == 3: B, HW, C = x.shape if x.shape[-2] != self.HEADS and self.HEADS != 0: x = x.reshape(B,self.HEADS,HW,-1)""" HEAD_DIM = x.shape[1] if HEAD_DIM == self.HEADS: B, HEAD_DIM, HW, C = x.shape x = x.reshape(B, HW, C*HEAD_DIM) if self.KONTEXT == 1: x = x.reshape(2, x.shape[1] // 2, x.shape[2]) weights_all_one = all(weight == 1.0 for weight in weight_list) methods_all_scattersort = all(name == "scattersort" for name in self.method) masks_all_none = all(mask is None for mask in self.mask) if weights_all_one and methods_all_scattersort and len(weight_list) > 1 and masks_all_none: buf = Stylizer.buffer buf['src_idx'] = x[0:1].argsort(dim=-2) buf['ref_sorted'], buf['ref_idx'] = x[1:].reshape(1, -1, x.shape[-1]).sort(dim=-2) buf['src'] = buf['ref_sorted'][:,::len(weight_list)].expand_as(buf['src_idx']) # interleave_stride = len(weight_list) x[0:1] = x[0:1].scatter_(dim=-2, index=buf['src_idx'], src=buf['src'],) else: for i, (weight, mask) in enumerate(zip(weight_list, self.mask)): if weight > 0 and weight < 1: x_clone = x.clone() if mask is not None: x01 = x[0:1].clone() slc = Stylizer.middle_slice(x.shape[-2], weight) method = getattr(self, self.method[i]) if weight == 0.0: continue elif weight == 1.0: x = method(x, idx=i+1) else: x = method(x, idx=i+1, slc=slc) if weight > 0 and weight < 1 and self.method[i] != "scattersort": x = torch.lerp(x_clone, x, weight) #else: # x = torch.lerp(x, method(x.clone(), idx=i), weight) if mask is not None: x[0:1] = torch.lerp(x01, x[0:1], mask.view(1, -1, 1)) #if x_ndim == 3: # return x.view(B,HW,C) if self.KONTEXT == 1: x = x.reshape(1, x.shape[1] * 2, x.shape[2]) if HEAD_DIM == self.HEADS: return x.reshape(B, HEAD_DIM, HW, C) else: return x def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): self.h_len = h_len self.w_len = w_len self.img_len = h_len * w_len self.img_slice = img_slice self.txt_slice = txt_slice self.HEADS = HEADS #for block in self.double_blocks: # block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) #for block in self.single_blocks: # block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) for i, mask in enumerate(self.mask): if mask is not None and mask.ndim > 1: self.mask[i] = F.interpolate(mask.unsqueeze(0), size=(h_len, w_len)).flatten().to(torch.bfloat16).cuda() def init_guides(self, model): if not self.GUIDES_INITIALIZED: if self.guides == []: self.guides = None elif self.guides is not None: for i, latent in enumerate(self.guides): if type(latent) is dict: latent = model.inner_model.inner_model.process_latent_in(latent['samples']).to(dtype=self.dtype, device=self.device) elif type(latent) is torch.Tensor: latent = latent.to(dtype=self.dtype, device=self.device) else: latent = None #raise ValueError(f"Invalid latent type: {type(latent)}") #if self.VIDEO and latent.shape[2] == 1: # latent = latent.repeat(1, 1, x.shape[2], 1, 1) self.guides[i] = latent if any(g is None for g in self.guides): self.guides = None print("Style guide nonetype set for Kontext.") else: self.guides = torch.cat(self.guides, dim=0) self.GUIDES_INITIALIZED = True def set_conditioning(self, positive, negative): self.cond_pos = [positive] self.cond_neg = [negative] def apply_style_conditioning(self, UNCOND, base_context, base_y=None, base_llama3=None): def get_max_token_lengths(style_conditioning, base_context, base_y=None, base_llama3=None): context_max_len = base_context.shape[-2] llama3_max_len = base_llama3.shape[-2] if base_llama3 is not None else -1 y_max_len = base_y.shape[-1] if base_y is not None else -1 for style_cond in style_conditioning: if style_cond is None: continue context_max_len = max(context_max_len, style_cond[0][0].shape[-2]) if base_llama3 is not None: llama3_max_len = max(llama3_max_len, style_cond[0][1]['conditioning_llama3'].shape[-2]) if base_y is not None: y_max_len = max(y_max_len, style_cond[0][1]['pooled_output'].shape[-1]) return context_max_len, llama3_max_len, y_max_len def pad_to_len(x, target_len, pad_value=0.0, dim=1): if target_len < 0: return x cur_len = x.shape[dim] if cur_len == target_len: return x return F.pad(x, (0, 0, 0, target_len - cur_len), value=pad_value) style_conditioning = self.cond_pos if not UNCOND else self.cond_neg context_max_len, llama3_max_len, y_max_len = get_max_token_lengths( style_conditioning = style_conditioning, base_context = base_context, base_y = base_y, base_llama3 = base_llama3, ) bsz_style = len(style_conditioning) context = base_context.repeat(bsz_style + 1, 1, 1) y = base_y.repeat(bsz_style + 1, 1) if base_y is not None else None llama3 = base_llama3.repeat(bsz_style + 1, 1, 1, 1) if base_llama3 is not None else None context = pad_to_len(context, context_max_len, dim=-2) llama3 = pad_to_len(llama3, llama3_max_len, dim=-2) if base_llama3 is not None else None y = pad_to_len(y, y_max_len, dim=-1) if base_y is not None else None for ci, style_cond in enumerate(style_conditioning): if style_cond is None: continue context[ci+1:ci+2] = pad_to_len(style_cond[0][0], context_max_len, dim=-2).to(context) if llama3 is not None: llama3 [ci+1:ci+2] = pad_to_len(style_cond[0][1]['conditioning_llama3'], llama3_max_len, dim=-2).to(llama3) if y is not None: y [ci+1:ci+2] = pad_to_len(style_cond[0][1]['pooled_output'], y_max_len, dim=-1).to(y) return context, y, llama3 def WCT_data(self, denoised_embed, y0_style_embed): Stylizer.CLS_WCT.set(y0_style_embed.to(denoised_embed)) return Stylizer.CLS_WCT.get(denoised_embed) def WCT2_data(self, denoised_embed, y0_style_embed): Stylizer.CLS_WCT2.set(y0_style_embed.to(denoised_embed)) return Stylizer.CLS_WCT2.get(denoised_embed) def apply_to_data(self, denoised, y0_style=None, mode="none"): if mode == "none": return denoised y0_style = self.guides if y0_style is None else y0_style y0_style_embed = self.Retrojector.embed(y0_style) denoised_embed = self.Retrojector.embed(denoised) B,HW,C = y0_style_embed.shape embed = torch.cat([denoised_embed, y0_style_embed.view(1,B*HW,C)[:,::B,:]], dim=0) method = getattr(self, mode) if mode == "scattersort": slc = Stylizer.middle_slice(embed.shape[-2], self.data_shock_weight) embed = method(embed, slc=slc) else: embed = method(embed) return self.Retrojector.unembed(embed[0:1]) def apply_recon_lure(self, denoised, y0_style): if self.recon_lure == "none": return denoised for i in range(denoised.shape[0]): denoised[i:i+1] = self.apply_to_data(denoised[i:i+1], y0_style, self.recon_lure) return denoised def apply_data_shock(self, denoised): if self.data_shock == "none": return denoised datashock_ref = getattr(self, "datashock_ref", None) if self.data_shock == "scattersort": return self.apply_to_data(denoised, datashock_ref, self.data_shock) else: return torch.lerp(denoised, self.apply_to_data(denoised, datashock_ref, self.data_shock), torch.Tensor([self.data_shock_weight]).double().cuda()) class StyleMMDiT_Model(Style_Model): def __init__(self, dtype=torch.float64, device=torch.device("cuda")): super().__init__(dtype, device) self.double_blocks = [StyleMMDiT_DoubleBlock() for _ in range(100)] self.single_blocks = [StyleMMDiT_SingleBlock() for _ in range(100)] def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) for block in self.double_blocks: block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) for block in self.single_blocks: block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) class StyleUNet_Model(Style_Model): def __init__(self, dtype=torch.float64, device=torch.device("cuda")): super().__init__(dtype, device) self.input_blocks = [StyleUNet_InputBlock() for _ in range(100)] self.middle_blocks = [StyleUNet_MiddleBlock() for _ in range(100)] self.output_blocks = [StyleUNet_OutputBlock() for _ in range(100)] def set_len(self, h_len, w_len, img_slice, txt_slice, HEADS): super().set_len(h_len, w_len, img_slice, txt_slice, HEADS) for block in self.input_blocks: block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) for block in self.middle_blocks: block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) for block in self.output_blocks: block.set_len(h_len, w_len, img_slice, txt_slice, HEADS) def __call__(self, x, attr): B, C, H, W = x.shape x = super().__call__(x.reshape(B, H*W, C), attr) return x.reshape(B,C,H,W)