asagi4
55d26c10c0
The model produces different output on different devices for some reason, but this matches the behaviour of the original code.
238 lines
8.2 KiB
Python
238 lines
8.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import functional as F
|
|
from timm import create_model
|
|
import einops
|
|
|
|
from diffusers.models.normalization import AdaGroupNorm
|
|
|
|
from timm.layers import use_fused_attn
|
|
|
|
|
|
from comfy.utils import common_upscale
|
|
|
|
|
|
class Attention(nn.Module):
|
|
fused_attn = True
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int = 8,
|
|
qkv_bias: bool = False,
|
|
qk_norm: bool = False,
|
|
attn_drop: float = 0.0,
|
|
proj_drop: float = 0.0,
|
|
norm_layer: nn.Module = nn.LayerNorm,
|
|
) -> None:
|
|
super().__init__()
|
|
assert dim % num_heads == 0, "dim should be divisible by num_heads"
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.scale = self.head_dim**-0.5
|
|
self.fused_attn = use_fused_attn()
|
|
|
|
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
|
self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
|
self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
|
self.attn_drop = nn.Dropout(attn_drop)
|
|
self.proj = nn.Linear(dim, dim)
|
|
self.proj_drop = nn.Dropout(proj_drop)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
B, N, C = x.shape
|
|
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv.unbind(0)
|
|
q, k = self.q_norm(q), self.k_norm(k)
|
|
|
|
if self.fused_attn:
|
|
x = F.scaled_dot_product_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
dropout_p=self.attn_drop.p if self.training else 0.0,
|
|
)
|
|
else:
|
|
q = q * self.scale
|
|
attn = q @ k.transpose(-2, -1)
|
|
attn = attn.softmax(dim=-1)
|
|
attn = self.attn_drop(attn)
|
|
x = attn @ v
|
|
|
|
x = x.transpose(1, 2).reshape(B, N, C)
|
|
x = self.proj(x)
|
|
x = self.proj_drop(x)
|
|
return x
|
|
|
|
|
|
class SVDNoiseUnet(nn.Module):
|
|
def __init__(self, in_channels=4, out_channels=4, resolution=128): # resolution = size // 8
|
|
super(SVDNoiseUnet, self).__init__()
|
|
|
|
_in = int(resolution * in_channels // 2)
|
|
_out = int(resolution * out_channels // 2)
|
|
self.mlp1 = nn.Sequential(
|
|
nn.Linear(_in, 64),
|
|
nn.ReLU(inplace=True),
|
|
nn.Linear(64, _out),
|
|
)
|
|
self.mlp2 = nn.Sequential(
|
|
nn.Linear(_in, 64),
|
|
nn.ReLU(inplace=True),
|
|
nn.Linear(64, _out),
|
|
)
|
|
|
|
self.mlp3 = nn.Sequential(
|
|
nn.Linear(_in, _out),
|
|
)
|
|
|
|
self.attention = Attention(_out)
|
|
|
|
self.bn = nn.BatchNorm2d(_out)
|
|
|
|
self.mlp4 = nn.Sequential(
|
|
nn.Linear(_out, 1024),
|
|
nn.ReLU(inplace=True),
|
|
nn.Linear(1024, _out),
|
|
)
|
|
|
|
def forward(self, x, residual=False):
|
|
b, c, h, w = x.shape
|
|
x = einops.rearrange(x, "b (a c)h w ->b (a h)(c w)", a=2, c=2) # x -> [1, 256, 256]
|
|
U, s, V = torch.linalg.svd(x) # U->[b 256 256], s-> [b 256], V->[b 256 256]
|
|
U_T = U.permute(0, 2, 1)
|
|
out = self.mlp1(U_T) + self.mlp2(V) + self.mlp3(s).unsqueeze(1) # s -> [b, 1, 256] => [b, 256, 256]
|
|
out = self.attention(out).mean(1)
|
|
out = self.mlp4(out) + s
|
|
pred = U @ torch.diag_embed(out) @ V
|
|
return einops.rearrange(pred, "b (a h)(c w) -> b (a c) h w", a=2, c=2)
|
|
|
|
|
|
class SVDNoiseUnet_Concise(nn.Module):
|
|
def __init__(self, in_channels=4, out_channels=4, resolution=128):
|
|
super(SVDNoiseUnet_Concise, self).__init__()
|
|
|
|
|
|
class NoiseTransformer(nn.Module):
|
|
def __init__(self, resolution=128):
|
|
super().__init__()
|
|
self.upsample = lambda x: F.interpolate(x, [224, 224])
|
|
self.downsample = lambda x: F.interpolate(x, [resolution, resolution])
|
|
self.upconv = nn.Conv2d(7, 4, (1, 1), (1, 1), (0, 0))
|
|
self.downconv = nn.Conv2d(4, 3, (1, 1), (1, 1), (0, 0))
|
|
# self.upconv = nn.Conv2d(7,4,(1,1),(1,1),(0,0))
|
|
self.swin = create_model("swin_tiny_patch4_window7_224", pretrained=True)
|
|
|
|
def forward(self, x, residual=False):
|
|
if residual:
|
|
x = self.upconv(self.downsample(self.swin.forward_features(self.downconv(self.upsample(x))))) + x
|
|
else:
|
|
x = self.upconv(self.downsample(self.swin.forward_features(self.downconv(self.upsample(x)))))
|
|
|
|
return x
|
|
|
|
|
|
class NPNet(nn.Module):
|
|
def __init__(self, model_id, pretrained_path, device="cuda") -> None:
|
|
super().__init__()
|
|
|
|
assert model_id in ["SDXL", "DreamShaper", "DiT"]
|
|
self.device = device
|
|
self.model_id = model_id
|
|
self.pretrained_path = pretrained_path
|
|
self.unet_embedding = NoiseTransformer(resolution=128)
|
|
self.unet_svd = SVDNoiseUnet(resolution=128)
|
|
if self.model_id == "DiT":
|
|
self.text_embedding = AdaGroupNorm(1024 * 77, 4, 1, eps=1e-6)
|
|
else:
|
|
self.text_embedding = AdaGroupNorm(2048 * 77, 4, 1, eps=1e-6)
|
|
sd = torch.load(self.pretrained_path, weights_only=True, map_location=device)
|
|
self.unet_embedding.load_state_dict(sd.pop("unet_embedding"))
|
|
self.unet_svd.load_state_dict(sd.pop("unet_svd"))
|
|
self.text_embedding.load_state_dict(sd.pop("embeeding"))
|
|
self.alpha = sd["alpha"]
|
|
self.beta = sd["beta"]
|
|
self.to(dtype=torch.float32, device=device)
|
|
|
|
def to(self, *args, **kwargs):
|
|
super().to(*args, **kwargs)
|
|
self.unet_embedding.to(*args, **kwargs)
|
|
self.unet_svd.to(*args, **kwargs)
|
|
self.text_embedding.to(*args, **kwargs)
|
|
self.alpha = self.alpha.to(*args, **kwargs)
|
|
self.beta = self.beta.to(*args, **kwargs)
|
|
self.device = self.alpha.device
|
|
return self
|
|
|
|
def forward(self, initial_noise, prompt_embeds):
|
|
prompt_embeds = prompt_embeds.float().view(prompt_embeds.shape[0], -1)
|
|
text_emb = self.text_embedding(initial_noise.float(), prompt_embeds)
|
|
|
|
encoder_hidden_states_svd = initial_noise
|
|
encoder_hidden_states_embedding = initial_noise + text_emb
|
|
|
|
golden_embedding = self.unet_embedding(encoder_hidden_states_embedding.float())
|
|
|
|
golden_noise = (
|
|
self.unet_svd(encoder_hidden_states_svd.float())
|
|
+ (2 * torch.sigmoid(self.alpha) - 1) * text_emb
|
|
+ self.beta * golden_embedding
|
|
)
|
|
|
|
return golden_noise
|
|
|
|
|
|
class NPNetGoldenNoise:
|
|
npnet = None
|
|
noise = None
|
|
cond = None
|
|
seed = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"noise": ("NOISE",),
|
|
"prompt": ("CONDITIONING",),
|
|
"model_path": ("STRING", {"default": "/path/to/sdxl.pth"}),
|
|
"model_type": (["SDXL", "DreamShaper", "DiT"],),
|
|
"device": (["cuda", "cpu"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("NOISE",)
|
|
CATEGORY = "_for_testing/golden_noise"
|
|
|
|
FUNCTION = "doit"
|
|
|
|
def generate_noise(self, input_latent):
|
|
self.seed = self.noise.seed
|
|
orig_shape = input_latent["samples"].shape
|
|
if orig_shape[-2] != 128 or orig_shape[-1] != 128:
|
|
input_latent = input_latent.copy()
|
|
print("Latent must be 128x128 for the NPNet model to work; generating square noise and reshaping...")
|
|
input_latent["samples"] = common_upscale(input_latent["samples"], 128, 128, "nearest-exact", "disabled")
|
|
init_noise = self.noise.generate_noise(input_latent).to(self.npnet.device)
|
|
cond = self.cond[0].clone().to(self.npnet.device)
|
|
if cond.shape[1] != 77:
|
|
print("NPNet can't handle conds >77 tokens, truncating...")
|
|
cond = cond[:, :77, :]
|
|
print("Applying NPNet to noise")
|
|
r = self.npnet(init_noise, cond).to("cpu")
|
|
if orig_shape[-2] != 128 or orig_shape[-1] != 128:
|
|
r = common_upscale(r, orig_shape[-1], orig_shape[-2], "nearest-exact", "disabled")
|
|
return r
|
|
|
|
def doit(self, noise, prompt, model_path, model_type, device):
|
|
if self.npnet is None:
|
|
print("Loading NPNet")
|
|
self.npnet = NPNet(model_type, model_path, device=device)
|
|
self.npnet.to(device)
|
|
self.noise = noise
|
|
self.cond = prompt[0]
|
|
|
|
return (self,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"NPNetGoldenNoise": NPNetGoldenNoise}
|