flyingshutter
9 months ago
32 changed files with 1858 additions and 142 deletions
@ -0,0 +1,161 @@
|
||||
""" |
||||
This file is part of ComfyUI. |
||||
Copyright (C) 2024 Stability AI |
||||
|
||||
This program is free software: you can redistribute it and/or modify |
||||
it under the terms of the GNU General Public License as published by |
||||
the Free Software Foundation, either version 3 of the License, or |
||||
(at your option) any later version. |
||||
|
||||
This program is distributed in the hope that it will be useful, |
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
||||
GNU General Public License for more details. |
||||
|
||||
You should have received a copy of the GNU General Public License |
||||
along with this program. If not, see <https://www.gnu.org/licenses/>. |
||||
""" |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
from comfy.ldm.modules.attention import optimized_attention |
||||
|
||||
class Linear(torch.nn.Linear): |
||||
def reset_parameters(self): |
||||
return None |
||||
|
||||
class Conv2d(torch.nn.Conv2d): |
||||
def reset_parameters(self): |
||||
return None |
||||
|
||||
class OptimizedAttention(nn.Module): |
||||
def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.heads = nhead |
||||
|
||||
self.to_q = operations.Linear(c, c, bias=True, dtype=dtype, device=device) |
||||
self.to_k = operations.Linear(c, c, bias=True, dtype=dtype, device=device) |
||||
self.to_v = operations.Linear(c, c, bias=True, dtype=dtype, device=device) |
||||
|
||||
self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device) |
||||
|
||||
def forward(self, q, k, v): |
||||
q = self.to_q(q) |
||||
k = self.to_k(k) |
||||
v = self.to_v(v) |
||||
|
||||
out = optimized_attention(q, k, v, self.heads) |
||||
|
||||
return self.out_proj(out) |
||||
|
||||
class Attention2D(nn.Module): |
||||
def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.attn = OptimizedAttention(c, nhead, dtype=dtype, device=device, operations=operations) |
||||
# self.attn = nn.MultiheadAttention(c, nhead, dropout=dropout, bias=True, batch_first=True, dtype=dtype, device=device) |
||||
|
||||
def forward(self, x, kv, self_attn=False): |
||||
orig_shape = x.shape |
||||
x = x.view(x.size(0), x.size(1), -1).permute(0, 2, 1) # Bx4xHxW -> Bx(HxW)x4 |
||||
if self_attn: |
||||
kv = torch.cat([x, kv], dim=1) |
||||
# x = self.attn(x, kv, kv, need_weights=False)[0] |
||||
x = self.attn(x, kv, kv) |
||||
x = x.permute(0, 2, 1).view(*orig_shape) |
||||
return x |
||||
|
||||
|
||||
def LayerNorm2d_op(operations): |
||||
class LayerNorm2d(operations.LayerNorm): |
||||
def __init__(self, *args, **kwargs): |
||||
super().__init__(*args, **kwargs) |
||||
|
||||
def forward(self, x): |
||||
return super().forward(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) |
||||
return LayerNorm2d |
||||
|
||||
class GlobalResponseNorm(nn.Module): |
||||
"from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105" |
||||
def __init__(self, dim, dtype=None, device=None): |
||||
super().__init__() |
||||
self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim, dtype=dtype, device=device)) |
||||
self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim, dtype=dtype, device=device)) |
||||
|
||||
def forward(self, x): |
||||
Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) |
||||
Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) |
||||
return self.gamma.to(device=x.device, dtype=x.dtype) * (x * Nx) + self.beta.to(device=x.device, dtype=x.dtype) + x |
||||
|
||||
|
||||
class ResBlock(nn.Module): |
||||
def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0, dtype=None, device=None, operations=None): # , num_heads=4, expansion=2): |
||||
super().__init__() |
||||
self.depthwise = operations.Conv2d(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c, dtype=dtype, device=device) |
||||
# self.depthwise = SAMBlock(c, num_heads, expansion) |
||||
self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
self.channelwise = nn.Sequential( |
||||
operations.Linear(c + c_skip, c * 4, dtype=dtype, device=device), |
||||
nn.GELU(), |
||||
GlobalResponseNorm(c * 4, dtype=dtype, device=device), |
||||
nn.Dropout(dropout), |
||||
operations.Linear(c * 4, c, dtype=dtype, device=device) |
||||
) |
||||
|
||||
def forward(self, x, x_skip=None): |
||||
x_res = x |
||||
x = self.norm(self.depthwise(x)) |
||||
if x_skip is not None: |
||||
x = torch.cat([x, x_skip], dim=1) |
||||
x = self.channelwise(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) |
||||
return x + x_res |
||||
|
||||
|
||||
class AttnBlock(nn.Module): |
||||
def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0, dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.self_attn = self_attn |
||||
self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
self.attention = Attention2D(c, nhead, dropout, dtype=dtype, device=device, operations=operations) |
||||
self.kv_mapper = nn.Sequential( |
||||
nn.SiLU(), |
||||
operations.Linear(c_cond, c, dtype=dtype, device=device) |
||||
) |
||||
|
||||
def forward(self, x, kv): |
||||
kv = self.kv_mapper(kv) |
||||
x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn) |
||||
return x |
||||
|
||||
|
||||
class FeedForwardBlock(nn.Module): |
||||
def __init__(self, c, dropout=0.0, dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
self.channelwise = nn.Sequential( |
||||
operations.Linear(c, c * 4, dtype=dtype, device=device), |
||||
nn.GELU(), |
||||
GlobalResponseNorm(c * 4, dtype=dtype, device=device), |
||||
nn.Dropout(dropout), |
||||
operations.Linear(c * 4, c, dtype=dtype, device=device) |
||||
) |
||||
|
||||
def forward(self, x): |
||||
x = x + self.channelwise(self.norm(x).permute(0, 2, 3, 1)).permute(0, 3, 1, 2) |
||||
return x |
||||
|
||||
|
||||
class TimestepBlock(nn.Module): |
||||
def __init__(self, c, c_timestep, conds=['sca'], dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.mapper = operations.Linear(c_timestep, c * 2, dtype=dtype, device=device) |
||||
self.conds = conds |
||||
for cname in conds: |
||||
setattr(self, f"mapper_{cname}", operations.Linear(c_timestep, c * 2, dtype=dtype, device=device)) |
||||
|
||||
def forward(self, x, t): |
||||
t = t.chunk(len(self.conds) + 1, dim=1) |
||||
a, b = self.mapper(t[0])[:, :, None, None].chunk(2, dim=1) |
||||
for i, c in enumerate(self.conds): |
||||
ac, bc = getattr(self, f"mapper_{c}")(t[i + 1])[:, :, None, None].chunk(2, dim=1) |
||||
a, b = a + ac, b + bc |
||||
return x * (1 + a) + b |
||||
@ -0,0 +1,258 @@
|
||||
""" |
||||
This file is part of ComfyUI. |
||||
Copyright (C) 2024 Stability AI |
||||
|
||||
This program is free software: you can redistribute it and/or modify |
||||
it under the terms of the GNU General Public License as published by |
||||
the Free Software Foundation, either version 3 of the License, or |
||||
(at your option) any later version. |
||||
|
||||
This program is distributed in the hope that it will be useful, |
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
||||
GNU General Public License for more details. |
||||
|
||||
You should have received a copy of the GNU General Public License |
||||
along with this program. If not, see <https://www.gnu.org/licenses/>. |
||||
""" |
||||
|
||||
import torch |
||||
from torch import nn |
||||
from torch.autograd import Function |
||||
|
||||
class vector_quantize(Function): |
||||
@staticmethod |
||||
def forward(ctx, x, codebook): |
||||
with torch.no_grad(): |
||||
codebook_sqr = torch.sum(codebook ** 2, dim=1) |
||||
x_sqr = torch.sum(x ** 2, dim=1, keepdim=True) |
||||
|
||||
dist = torch.addmm(codebook_sqr + x_sqr, x, codebook.t(), alpha=-2.0, beta=1.0) |
||||
_, indices = dist.min(dim=1) |
||||
|
||||
ctx.save_for_backward(indices, codebook) |
||||
ctx.mark_non_differentiable(indices) |
||||
|
||||
nn = torch.index_select(codebook, 0, indices) |
||||
return nn, indices |
||||
|
||||
@staticmethod |
||||
def backward(ctx, grad_output, grad_indices): |
||||
grad_inputs, grad_codebook = None, None |
||||
|
||||
if ctx.needs_input_grad[0]: |
||||
grad_inputs = grad_output.clone() |
||||
if ctx.needs_input_grad[1]: |
||||
# Gradient wrt. the codebook |
||||
indices, codebook = ctx.saved_tensors |
||||
|
||||
grad_codebook = torch.zeros_like(codebook) |
||||
grad_codebook.index_add_(0, indices, grad_output) |
||||
|
||||
return (grad_inputs, grad_codebook) |
||||
|
||||
|
||||
class VectorQuantize(nn.Module): |
||||
def __init__(self, embedding_size, k, ema_decay=0.99, ema_loss=False): |
||||
""" |
||||
Takes an input of variable size (as long as the last dimension matches the embedding size). |
||||
Returns one tensor containing the nearest neigbour embeddings to each of the inputs, |
||||
with the same size as the input, vq and commitment components for the loss as a touple |
||||
in the second output and the indices of the quantized vectors in the third: |
||||
quantized, (vq_loss, commit_loss), indices |
||||
""" |
||||
super(VectorQuantize, self).__init__() |
||||
|
||||
self.codebook = nn.Embedding(k, embedding_size) |
||||
self.codebook.weight.data.uniform_(-1./k, 1./k) |
||||
self.vq = vector_quantize.apply |
||||
|
||||
self.ema_decay = ema_decay |
||||
self.ema_loss = ema_loss |
||||
if ema_loss: |
||||
self.register_buffer('ema_element_count', torch.ones(k)) |
||||
self.register_buffer('ema_weight_sum', torch.zeros_like(self.codebook.weight)) |
||||
|
||||
def _laplace_smoothing(self, x, epsilon): |
||||
n = torch.sum(x) |
||||
return ((x + epsilon) / (n + x.size(0) * epsilon) * n) |
||||
|
||||
def _updateEMA(self, z_e_x, indices): |
||||
mask = nn.functional.one_hot(indices, self.ema_element_count.size(0)).float() |
||||
elem_count = mask.sum(dim=0) |
||||
weight_sum = torch.mm(mask.t(), z_e_x) |
||||
|
||||
self.ema_element_count = (self.ema_decay * self.ema_element_count) + ((1-self.ema_decay) * elem_count) |
||||
self.ema_element_count = self._laplace_smoothing(self.ema_element_count, 1e-5) |
||||
self.ema_weight_sum = (self.ema_decay * self.ema_weight_sum) + ((1-self.ema_decay) * weight_sum) |
||||
|
||||
self.codebook.weight.data = self.ema_weight_sum / self.ema_element_count.unsqueeze(-1) |
||||
|
||||
def idx2vq(self, idx, dim=-1): |
||||
q_idx = self.codebook(idx) |
||||
if dim != -1: |
||||
q_idx = q_idx.movedim(-1, dim) |
||||
return q_idx |
||||
|
||||
def forward(self, x, get_losses=True, dim=-1): |
||||
if dim != -1: |
||||
x = x.movedim(dim, -1) |
||||
z_e_x = x.contiguous().view(-1, x.size(-1)) if len(x.shape) > 2 else x |
||||
z_q_x, indices = self.vq(z_e_x, self.codebook.weight.detach()) |
||||
vq_loss, commit_loss = None, None |
||||
if self.ema_loss and self.training: |
||||
self._updateEMA(z_e_x.detach(), indices.detach()) |
||||
# pick the graded embeddings after updating the codebook in order to have a more accurate commitment loss |
||||
z_q_x_grd = torch.index_select(self.codebook.weight, dim=0, index=indices) |
||||
if get_losses: |
||||
vq_loss = (z_q_x_grd - z_e_x.detach()).pow(2).mean() |
||||
commit_loss = (z_e_x - z_q_x_grd.detach()).pow(2).mean() |
||||
|
||||
z_q_x = z_q_x.view(x.shape) |
||||
if dim != -1: |
||||
z_q_x = z_q_x.movedim(-1, dim) |
||||
return z_q_x, (vq_loss, commit_loss), indices.view(x.shape[:-1]) |
||||
|
||||
|
||||
class ResBlock(nn.Module): |
||||
def __init__(self, c, c_hidden): |
||||
super().__init__() |
||||
# depthwise/attention |
||||
self.norm1 = nn.LayerNorm(c, elementwise_affine=False, eps=1e-6) |
||||
self.depthwise = nn.Sequential( |
||||
nn.ReplicationPad2d(1), |
||||
nn.Conv2d(c, c, kernel_size=3, groups=c) |
||||
) |
||||
|
||||
# channelwise |
||||
self.norm2 = nn.LayerNorm(c, elementwise_affine=False, eps=1e-6) |
||||
self.channelwise = nn.Sequential( |
||||
nn.Linear(c, c_hidden), |
||||
nn.GELU(), |
||||
nn.Linear(c_hidden, c), |
||||
) |
||||
|
||||
self.gammas = nn.Parameter(torch.zeros(6), requires_grad=True) |
||||
|
||||
# Init weights |
||||
def _basic_init(module): |
||||
if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d): |
||||
torch.nn.init.xavier_uniform_(module.weight) |
||||
if module.bias is not None: |
||||
nn.init.constant_(module.bias, 0) |
||||
|
||||
self.apply(_basic_init) |
||||
|
||||
def _norm(self, x, norm): |
||||
return norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) |
||||
|
||||
def forward(self, x): |
||||
mods = self.gammas |
||||
|
||||
x_temp = self._norm(x, self.norm1) * (1 + mods[0]) + mods[1] |
||||
try: |
||||
x = x + self.depthwise(x_temp) * mods[2] |
||||
except: #operation not implemented for bf16 |
||||
x_temp = self.depthwise[0](x_temp.float()).to(x.dtype) |
||||
x = x + self.depthwise[1](x_temp) * mods[2] |
||||
|
||||
x_temp = self._norm(x, self.norm2) * (1 + mods[3]) + mods[4] |
||||
x = x + self.channelwise(x_temp.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) * mods[5] |
||||
|
||||
return x |
||||
|
||||
|
||||
class StageA(nn.Module): |
||||
def __init__(self, levels=2, bottleneck_blocks=12, c_hidden=384, c_latent=4, codebook_size=8192, |
||||
scale_factor=0.43): # 0.3764 |
||||
super().__init__() |
||||
self.c_latent = c_latent |
||||
self.scale_factor = scale_factor |
||||
c_levels = [c_hidden // (2 ** i) for i in reversed(range(levels))] |
||||
|
||||
# Encoder blocks |
||||
self.in_block = nn.Sequential( |
||||
nn.PixelUnshuffle(2), |
||||
nn.Conv2d(3 * 4, c_levels[0], kernel_size=1) |
||||
) |
||||
down_blocks = [] |
||||
for i in range(levels): |
||||
if i > 0: |
||||
down_blocks.append(nn.Conv2d(c_levels[i - 1], c_levels[i], kernel_size=4, stride=2, padding=1)) |
||||
block = ResBlock(c_levels[i], c_levels[i] * 4) |
||||
down_blocks.append(block) |
||||
down_blocks.append(nn.Sequential( |
||||
nn.Conv2d(c_levels[-1], c_latent, kernel_size=1, bias=False), |
||||
nn.BatchNorm2d(c_latent), # then normalize them to have mean 0 and std 1 |
||||
)) |
||||
self.down_blocks = nn.Sequential(*down_blocks) |
||||
self.down_blocks[0] |
||||
|
||||
self.codebook_size = codebook_size |
||||
self.vquantizer = VectorQuantize(c_latent, k=codebook_size) |
||||
|
||||
# Decoder blocks |
||||
up_blocks = [nn.Sequential( |
||||
nn.Conv2d(c_latent, c_levels[-1], kernel_size=1) |
||||
)] |
||||
for i in range(levels): |
||||
for j in range(bottleneck_blocks if i == 0 else 1): |
||||
block = ResBlock(c_levels[levels - 1 - i], c_levels[levels - 1 - i] * 4) |
||||
up_blocks.append(block) |
||||
if i < levels - 1: |
||||
up_blocks.append( |
||||
nn.ConvTranspose2d(c_levels[levels - 1 - i], c_levels[levels - 2 - i], kernel_size=4, stride=2, |
||||
padding=1)) |
||||
self.up_blocks = nn.Sequential(*up_blocks) |
||||
self.out_block = nn.Sequential( |
||||
nn.Conv2d(c_levels[0], 3 * 4, kernel_size=1), |
||||
nn.PixelShuffle(2), |
||||
) |
||||
|
||||
def encode(self, x, quantize=False): |
||||
x = self.in_block(x) |
||||
x = self.down_blocks(x) |
||||
if quantize: |
||||
qe, (vq_loss, commit_loss), indices = self.vquantizer.forward(x, dim=1) |
||||
return qe / self.scale_factor, x / self.scale_factor, indices, vq_loss + commit_loss * 0.25 |
||||
else: |
||||
return x / self.scale_factor |
||||
|
||||
def decode(self, x): |
||||
x = x * self.scale_factor |
||||
x = self.up_blocks(x) |
||||
x = self.out_block(x) |
||||
return x |
||||
|
||||
def forward(self, x, quantize=False): |
||||
qe, x, _, vq_loss = self.encode(x, quantize) |
||||
x = self.decode(qe) |
||||
return x, vq_loss |
||||
|
||||
|
||||
class Discriminator(nn.Module): |
||||
def __init__(self, c_in=3, c_cond=0, c_hidden=512, depth=6): |
||||
super().__init__() |
||||
d = max(depth - 3, 3) |
||||
layers = [ |
||||
nn.utils.spectral_norm(nn.Conv2d(c_in, c_hidden // (2 ** d), kernel_size=3, stride=2, padding=1)), |
||||
nn.LeakyReLU(0.2), |
||||
] |
||||
for i in range(depth - 1): |
||||
c_in = c_hidden // (2 ** max((d - i), 0)) |
||||
c_out = c_hidden // (2 ** max((d - 1 - i), 0)) |
||||
layers.append(nn.utils.spectral_norm(nn.Conv2d(c_in, c_out, kernel_size=3, stride=2, padding=1))) |
||||
layers.append(nn.InstanceNorm2d(c_out)) |
||||
layers.append(nn.LeakyReLU(0.2)) |
||||
self.encoder = nn.Sequential(*layers) |
||||
self.shuffle = nn.Conv2d((c_hidden + c_cond) if c_cond > 0 else c_hidden, 1, kernel_size=1) |
||||
self.logits = nn.Sigmoid() |
||||
|
||||
def forward(self, x, cond=None): |
||||
x = self.encoder(x) |
||||
if cond is not None: |
||||
cond = cond.view(cond.size(0), cond.size(1), 1, 1, ).expand(-1, -1, x.size(-2), x.size(-1)) |
||||
x = torch.cat([x, cond], dim=1) |
||||
x = self.shuffle(x) |
||||
x = self.logits(x) |
||||
return x |
||||
@ -0,0 +1,257 @@
|
||||
""" |
||||
This file is part of ComfyUI. |
||||
Copyright (C) 2024 Stability AI |
||||
|
||||
This program is free software: you can redistribute it and/or modify |
||||
it under the terms of the GNU General Public License as published by |
||||
the Free Software Foundation, either version 3 of the License, or |
||||
(at your option) any later version. |
||||
|
||||
This program is distributed in the hope that it will be useful, |
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
||||
GNU General Public License for more details. |
||||
|
||||
You should have received a copy of the GNU General Public License |
||||
along with this program. If not, see <https://www.gnu.org/licenses/>. |
||||
""" |
||||
|
||||
import math |
||||
import numpy as np |
||||
import torch |
||||
from torch import nn |
||||
from .common import AttnBlock, LayerNorm2d_op, ResBlock, FeedForwardBlock, TimestepBlock |
||||
|
||||
class StageB(nn.Module): |
||||
def __init__(self, c_in=4, c_out=4, c_r=64, patch_size=2, c_cond=1280, c_hidden=[320, 640, 1280, 1280], |
||||
nhead=[-1, -1, 20, 20], blocks=[[2, 6, 28, 6], [6, 28, 6, 2]], |
||||
block_repeat=[[1, 1, 1, 1], [3, 3, 2, 2]], level_config=['CT', 'CT', 'CTA', 'CTA'], c_clip=1280, |
||||
c_clip_seq=4, c_effnet=16, c_pixels=3, kernel_size=3, dropout=[0, 0, 0.0, 0.0], self_attn=True, |
||||
t_conds=['sca'], stable_cascade_stage=None, dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.dtype = dtype |
||||
self.c_r = c_r |
||||
self.t_conds = t_conds |
||||
self.c_clip_seq = c_clip_seq |
||||
if not isinstance(dropout, list): |
||||
dropout = [dropout] * len(c_hidden) |
||||
if not isinstance(self_attn, list): |
||||
self_attn = [self_attn] * len(c_hidden) |
||||
|
||||
# CONDITIONING |
||||
self.effnet_mapper = nn.Sequential( |
||||
operations.Conv2d(c_effnet, c_hidden[0] * 4, kernel_size=1, dtype=dtype, device=device), |
||||
nn.GELU(), |
||||
operations.Conv2d(c_hidden[0] * 4, c_hidden[0], kernel_size=1, dtype=dtype, device=device), |
||||
LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
) |
||||
self.pixels_mapper = nn.Sequential( |
||||
operations.Conv2d(c_pixels, c_hidden[0] * 4, kernel_size=1, dtype=dtype, device=device), |
||||
nn.GELU(), |
||||
operations.Conv2d(c_hidden[0] * 4, c_hidden[0], kernel_size=1, dtype=dtype, device=device), |
||||
LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
) |
||||
self.clip_mapper = operations.Linear(c_clip, c_cond * c_clip_seq, dtype=dtype, device=device) |
||||
self.clip_norm = operations.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
|
||||
self.embedding = nn.Sequential( |
||||
nn.PixelUnshuffle(patch_size), |
||||
operations.Conv2d(c_in * (patch_size ** 2), c_hidden[0], kernel_size=1, dtype=dtype, device=device), |
||||
LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
) |
||||
|
||||
def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0, self_attn=True): |
||||
if block_type == 'C': |
||||
return ResBlock(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout, dtype=dtype, device=device, operations=operations) |
||||
elif block_type == 'A': |
||||
return AttnBlock(c_hidden, c_cond, nhead, self_attn=self_attn, dropout=dropout, dtype=dtype, device=device, operations=operations) |
||||
elif block_type == 'F': |
||||
return FeedForwardBlock(c_hidden, dropout=dropout, dtype=dtype, device=device, operations=operations) |
||||
elif block_type == 'T': |
||||
return TimestepBlock(c_hidden, c_r, conds=t_conds, dtype=dtype, device=device, operations=operations) |
||||
else: |
||||
raise Exception(f'Block type {block_type} not supported') |
||||
|
||||
# BLOCKS |
||||
# -- down blocks |
||||
self.down_blocks = nn.ModuleList() |
||||
self.down_downscalers = nn.ModuleList() |
||||
self.down_repeat_mappers = nn.ModuleList() |
||||
for i in range(len(c_hidden)): |
||||
if i > 0: |
||||
self.down_downscalers.append(nn.Sequential( |
||||
LayerNorm2d_op(operations)(c_hidden[i - 1], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device), |
||||
operations.Conv2d(c_hidden[i - 1], c_hidden[i], kernel_size=2, stride=2, dtype=dtype, device=device), |
||||
)) |
||||
else: |
||||
self.down_downscalers.append(nn.Identity()) |
||||
down_block = nn.ModuleList() |
||||
for _ in range(blocks[0][i]): |
||||
for block_type in level_config[i]: |
||||
block = get_block(block_type, c_hidden[i], nhead[i], dropout=dropout[i], self_attn=self_attn[i]) |
||||
down_block.append(block) |
||||
self.down_blocks.append(down_block) |
||||
if block_repeat is not None: |
||||
block_repeat_mappers = nn.ModuleList() |
||||
for _ in range(block_repeat[0][i] - 1): |
||||
block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device)) |
||||
self.down_repeat_mappers.append(block_repeat_mappers) |
||||
|
||||
# -- up blocks |
||||
self.up_blocks = nn.ModuleList() |
||||
self.up_upscalers = nn.ModuleList() |
||||
self.up_repeat_mappers = nn.ModuleList() |
||||
for i in reversed(range(len(c_hidden))): |
||||
if i > 0: |
||||
self.up_upscalers.append(nn.Sequential( |
||||
LayerNorm2d_op(operations)(c_hidden[i], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device), |
||||
operations.ConvTranspose2d(c_hidden[i], c_hidden[i - 1], kernel_size=2, stride=2, dtype=dtype, device=device), |
||||
)) |
||||
else: |
||||
self.up_upscalers.append(nn.Identity()) |
||||
up_block = nn.ModuleList() |
||||
for j in range(blocks[1][::-1][i]): |
||||
for k, block_type in enumerate(level_config[i]): |
||||
c_skip = c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0 |
||||
block = get_block(block_type, c_hidden[i], nhead[i], c_skip=c_skip, dropout=dropout[i], |
||||
self_attn=self_attn[i]) |
||||
up_block.append(block) |
||||
self.up_blocks.append(up_block) |
||||
if block_repeat is not None: |
||||
block_repeat_mappers = nn.ModuleList() |
||||
for _ in range(block_repeat[1][::-1][i] - 1): |
||||
block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device)) |
||||
self.up_repeat_mappers.append(block_repeat_mappers) |
||||
|
||||
# OUTPUT |
||||
self.clf = nn.Sequential( |
||||
LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device), |
||||
operations.Conv2d(c_hidden[0], c_out * (patch_size ** 2), kernel_size=1, dtype=dtype, device=device), |
||||
nn.PixelShuffle(patch_size), |
||||
) |
||||
|
||||
# --- WEIGHT INIT --- |
||||
# self.apply(self._init_weights) # General init |
||||
# nn.init.normal_(self.clip_mapper.weight, std=0.02) # conditionings |
||||
# nn.init.normal_(self.effnet_mapper[0].weight, std=0.02) # conditionings |
||||
# nn.init.normal_(self.effnet_mapper[2].weight, std=0.02) # conditionings |
||||
# nn.init.normal_(self.pixels_mapper[0].weight, std=0.02) # conditionings |
||||
# nn.init.normal_(self.pixels_mapper[2].weight, std=0.02) # conditionings |
||||
# torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02) # inputs |
||||
# nn.init.constant_(self.clf[1].weight, 0) # outputs |
||||
# |
||||
# # blocks |
||||
# for level_block in self.down_blocks + self.up_blocks: |
||||
# for block in level_block: |
||||
# if isinstance(block, ResBlock) or isinstance(block, FeedForwardBlock): |
||||
# block.channelwise[-1].weight.data *= np.sqrt(1 / sum(blocks[0])) |
||||
# elif isinstance(block, TimestepBlock): |
||||
# for layer in block.modules(): |
||||
# if isinstance(layer, nn.Linear): |
||||
# nn.init.constant_(layer.weight, 0) |
||||
# |
||||
# def _init_weights(self, m): |
||||
# if isinstance(m, (nn.Conv2d, nn.Linear)): |
||||
# torch.nn.init.xavier_uniform_(m.weight) |
||||
# if m.bias is not None: |
||||
# nn.init.constant_(m.bias, 0) |
||||
|
||||
def gen_r_embedding(self, r, max_positions=10000): |
||||
r = r * max_positions |
||||
half_dim = self.c_r // 2 |
||||
emb = math.log(max_positions) / (half_dim - 1) |
||||
emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp() |
||||
emb = r[:, None] * emb[None, :] |
||||
emb = torch.cat([emb.sin(), emb.cos()], dim=1) |
||||
if self.c_r % 2 == 1: # zero pad |
||||
emb = nn.functional.pad(emb, (0, 1), mode='constant') |
||||
return emb |
||||
|
||||
def gen_c_embeddings(self, clip): |
||||
if len(clip.shape) == 2: |
||||
clip = clip.unsqueeze(1) |
||||
clip = self.clip_mapper(clip).view(clip.size(0), clip.size(1) * self.c_clip_seq, -1) |
||||
clip = self.clip_norm(clip) |
||||
return clip |
||||
|
||||
def _down_encode(self, x, r_embed, clip): |
||||
level_outputs = [] |
||||
block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers) |
||||
for down_block, downscaler, repmap in block_group: |
||||
x = downscaler(x) |
||||
for i in range(len(repmap) + 1): |
||||
for block in down_block: |
||||
if isinstance(block, ResBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
ResBlock)): |
||||
x = block(x) |
||||
elif isinstance(block, AttnBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
AttnBlock)): |
||||
x = block(x, clip) |
||||
elif isinstance(block, TimestepBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
TimestepBlock)): |
||||
x = block(x, r_embed) |
||||
else: |
||||
x = block(x) |
||||
if i < len(repmap): |
||||
x = repmap[i](x) |
||||
level_outputs.insert(0, x) |
||||
return level_outputs |
||||
|
||||
def _up_decode(self, level_outputs, r_embed, clip): |
||||
x = level_outputs[0] |
||||
block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers) |
||||
for i, (up_block, upscaler, repmap) in enumerate(block_group): |
||||
for j in range(len(repmap) + 1): |
||||
for k, block in enumerate(up_block): |
||||
if isinstance(block, ResBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
ResBlock)): |
||||
skip = level_outputs[i] if k == 0 and i > 0 else None |
||||
if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)): |
||||
x = torch.nn.functional.interpolate(x, skip.shape[-2:], mode='bilinear', |
||||
align_corners=True) |
||||
x = block(x, skip) |
||||
elif isinstance(block, AttnBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
AttnBlock)): |
||||
x = block(x, clip) |
||||
elif isinstance(block, TimestepBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
TimestepBlock)): |
||||
x = block(x, r_embed) |
||||
else: |
||||
x = block(x) |
||||
if j < len(repmap): |
||||
x = repmap[j](x) |
||||
x = upscaler(x) |
||||
return x |
||||
|
||||
def forward(self, x, r, effnet, clip, pixels=None, **kwargs): |
||||
if pixels is None: |
||||
pixels = x.new_zeros(x.size(0), 3, 8, 8) |
||||
|
||||
# Process the conditioning embeddings |
||||
r_embed = self.gen_r_embedding(r).to(dtype=x.dtype) |
||||
for c in self.t_conds: |
||||
t_cond = kwargs.get(c, torch.zeros_like(r)) |
||||
r_embed = torch.cat([r_embed, self.gen_r_embedding(t_cond).to(dtype=x.dtype)], dim=1) |
||||
clip = self.gen_c_embeddings(clip) |
||||
|
||||
# Model Blocks |
||||
x = self.embedding(x) |
||||
x = x + self.effnet_mapper( |
||||
nn.functional.interpolate(effnet, size=x.shape[-2:], mode='bilinear', align_corners=True)) |
||||
x = x + nn.functional.interpolate(self.pixels_mapper(pixels), size=x.shape[-2:], mode='bilinear', |
||||
align_corners=True) |
||||
level_outputs = self._down_encode(x, r_embed, clip) |
||||
x = self._up_decode(level_outputs, r_embed, clip) |
||||
return self.clf(x) |
||||
|
||||
def update_weights_ema(self, src_model, beta=0.999): |
||||
for self_params, src_params in zip(self.parameters(), src_model.parameters()): |
||||
self_params.data = self_params.data * beta + src_params.data.clone().to(self_params.device) * (1 - beta) |
||||
for self_buffers, src_buffers in zip(self.buffers(), src_model.buffers()): |
||||
self_buffers.data = self_buffers.data * beta + src_buffers.data.clone().to(self_buffers.device) * (1 - beta) |
||||
@ -0,0 +1,271 @@
|
||||
""" |
||||
This file is part of ComfyUI. |
||||
Copyright (C) 2024 Stability AI |
||||
|
||||
This program is free software: you can redistribute it and/or modify |
||||
it under the terms of the GNU General Public License as published by |
||||
the Free Software Foundation, either version 3 of the License, or |
||||
(at your option) any later version. |
||||
|
||||
This program is distributed in the hope that it will be useful, |
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
||||
GNU General Public License for more details. |
||||
|
||||
You should have received a copy of the GNU General Public License |
||||
along with this program. If not, see <https://www.gnu.org/licenses/>. |
||||
""" |
||||
|
||||
import torch |
||||
from torch import nn |
||||
import numpy as np |
||||
import math |
||||
from .common import AttnBlock, LayerNorm2d_op, ResBlock, FeedForwardBlock, TimestepBlock |
||||
# from .controlnet import ControlNetDeliverer |
||||
|
||||
class UpDownBlock2d(nn.Module): |
||||
def __init__(self, c_in, c_out, mode, enabled=True, dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
assert mode in ['up', 'down'] |
||||
interpolation = nn.Upsample(scale_factor=2 if mode == 'up' else 0.5, mode='bilinear', |
||||
align_corners=True) if enabled else nn.Identity() |
||||
mapping = operations.Conv2d(c_in, c_out, kernel_size=1, dtype=dtype, device=device) |
||||
self.blocks = nn.ModuleList([interpolation, mapping] if mode == 'up' else [mapping, interpolation]) |
||||
|
||||
def forward(self, x): |
||||
for block in self.blocks: |
||||
x = block(x) |
||||
return x |
||||
|
||||
|
||||
class StageC(nn.Module): |
||||
def __init__(self, c_in=16, c_out=16, c_r=64, patch_size=1, c_cond=2048, c_hidden=[2048, 2048], nhead=[32, 32], |
||||
blocks=[[8, 24], [24, 8]], block_repeat=[[1, 1], [1, 1]], level_config=['CTA', 'CTA'], |
||||
c_clip_text=1280, c_clip_text_pooled=1280, c_clip_img=768, c_clip_seq=4, kernel_size=3, |
||||
dropout=[0.0, 0.0], self_attn=True, t_conds=['sca', 'crp'], switch_level=[False], stable_cascade_stage=None, |
||||
dtype=None, device=None, operations=None): |
||||
super().__init__() |
||||
self.dtype = dtype |
||||
self.c_r = c_r |
||||
self.t_conds = t_conds |
||||
self.c_clip_seq = c_clip_seq |
||||
if not isinstance(dropout, list): |
||||
dropout = [dropout] * len(c_hidden) |
||||
if not isinstance(self_attn, list): |
||||
self_attn = [self_attn] * len(c_hidden) |
||||
|
||||
# CONDITIONING |
||||
self.clip_txt_mapper = operations.Linear(c_clip_text, c_cond, dtype=dtype, device=device) |
||||
self.clip_txt_pooled_mapper = operations.Linear(c_clip_text_pooled, c_cond * c_clip_seq, dtype=dtype, device=device) |
||||
self.clip_img_mapper = operations.Linear(c_clip_img, c_cond * c_clip_seq, dtype=dtype, device=device) |
||||
self.clip_norm = operations.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device) |
||||
|
||||
self.embedding = nn.Sequential( |
||||
nn.PixelUnshuffle(patch_size), |
||||
operations.Conv2d(c_in * (patch_size ** 2), c_hidden[0], kernel_size=1, dtype=dtype, device=device), |
||||
LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6) |
||||
) |
||||
|
||||
def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0, self_attn=True): |
||||
if block_type == 'C': |
||||
return ResBlock(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout, dtype=dtype, device=device, operations=operations) |
||||
elif block_type == 'A': |
||||
return AttnBlock(c_hidden, c_cond, nhead, self_attn=self_attn, dropout=dropout, dtype=dtype, device=device, operations=operations) |
||||
elif block_type == 'F': |
||||
return FeedForwardBlock(c_hidden, dropout=dropout, dtype=dtype, device=device, operations=operations) |
||||
elif block_type == 'T': |
||||
return TimestepBlock(c_hidden, c_r, conds=t_conds, dtype=dtype, device=device, operations=operations) |
||||
else: |
||||
raise Exception(f'Block type {block_type} not supported') |
||||
|
||||
# BLOCKS |
||||
# -- down blocks |
||||
self.down_blocks = nn.ModuleList() |
||||
self.down_downscalers = nn.ModuleList() |
||||
self.down_repeat_mappers = nn.ModuleList() |
||||
for i in range(len(c_hidden)): |
||||
if i > 0: |
||||
self.down_downscalers.append(nn.Sequential( |
||||
LayerNorm2d_op(operations)(c_hidden[i - 1], elementwise_affine=False, eps=1e-6), |
||||
UpDownBlock2d(c_hidden[i - 1], c_hidden[i], mode='down', enabled=switch_level[i - 1], dtype=dtype, device=device, operations=operations) |
||||
)) |
||||
else: |
||||
self.down_downscalers.append(nn.Identity()) |
||||
down_block = nn.ModuleList() |
||||
for _ in range(blocks[0][i]): |
||||
for block_type in level_config[i]: |
||||
block = get_block(block_type, c_hidden[i], nhead[i], dropout=dropout[i], self_attn=self_attn[i]) |
||||
down_block.append(block) |
||||
self.down_blocks.append(down_block) |
||||
if block_repeat is not None: |
||||
block_repeat_mappers = nn.ModuleList() |
||||
for _ in range(block_repeat[0][i] - 1): |
||||
block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device)) |
||||
self.down_repeat_mappers.append(block_repeat_mappers) |
||||
|
||||
# -- up blocks |
||||
self.up_blocks = nn.ModuleList() |
||||
self.up_upscalers = nn.ModuleList() |
||||
self.up_repeat_mappers = nn.ModuleList() |
||||
for i in reversed(range(len(c_hidden))): |
||||
if i > 0: |
||||
self.up_upscalers.append(nn.Sequential( |
||||
LayerNorm2d_op(operations)(c_hidden[i], elementwise_affine=False, eps=1e-6), |
||||
UpDownBlock2d(c_hidden[i], c_hidden[i - 1], mode='up', enabled=switch_level[i - 1], dtype=dtype, device=device, operations=operations) |
||||
)) |
||||
else: |
||||
self.up_upscalers.append(nn.Identity()) |
||||
up_block = nn.ModuleList() |
||||
for j in range(blocks[1][::-1][i]): |
||||
for k, block_type in enumerate(level_config[i]): |
||||
c_skip = c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0 |
||||
block = get_block(block_type, c_hidden[i], nhead[i], c_skip=c_skip, dropout=dropout[i], |
||||
self_attn=self_attn[i]) |
||||
up_block.append(block) |
||||
self.up_blocks.append(up_block) |
||||
if block_repeat is not None: |
||||
block_repeat_mappers = nn.ModuleList() |
||||
for _ in range(block_repeat[1][::-1][i] - 1): |
||||
block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device)) |
||||
self.up_repeat_mappers.append(block_repeat_mappers) |
||||
|
||||
# OUTPUT |
||||
self.clf = nn.Sequential( |
||||
LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device), |
||||
operations.Conv2d(c_hidden[0], c_out * (patch_size ** 2), kernel_size=1, dtype=dtype, device=device), |
||||
nn.PixelShuffle(patch_size), |
||||
) |
||||
|
||||
# --- WEIGHT INIT --- |
||||
# self.apply(self._init_weights) # General init |
||||
# nn.init.normal_(self.clip_txt_mapper.weight, std=0.02) # conditionings |
||||
# nn.init.normal_(self.clip_txt_pooled_mapper.weight, std=0.02) # conditionings |
||||
# nn.init.normal_(self.clip_img_mapper.weight, std=0.02) # conditionings |
||||
# torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02) # inputs |
||||
# nn.init.constant_(self.clf[1].weight, 0) # outputs |
||||
# |
||||
# # blocks |
||||
# for level_block in self.down_blocks + self.up_blocks: |
||||
# for block in level_block: |
||||
# if isinstance(block, ResBlock) or isinstance(block, FeedForwardBlock): |
||||
# block.channelwise[-1].weight.data *= np.sqrt(1 / sum(blocks[0])) |
||||
# elif isinstance(block, TimestepBlock): |
||||
# for layer in block.modules(): |
||||
# if isinstance(layer, nn.Linear): |
||||
# nn.init.constant_(layer.weight, 0) |
||||
# |
||||
# def _init_weights(self, m): |
||||
# if isinstance(m, (nn.Conv2d, nn.Linear)): |
||||
# torch.nn.init.xavier_uniform_(m.weight) |
||||
# if m.bias is not None: |
||||
# nn.init.constant_(m.bias, 0) |
||||
|
||||
def gen_r_embedding(self, r, max_positions=10000): |
||||
r = r * max_positions |
||||
half_dim = self.c_r // 2 |
||||
emb = math.log(max_positions) / (half_dim - 1) |
||||
emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp() |
||||
emb = r[:, None] * emb[None, :] |
||||
emb = torch.cat([emb.sin(), emb.cos()], dim=1) |
||||
if self.c_r % 2 == 1: # zero pad |
||||
emb = nn.functional.pad(emb, (0, 1), mode='constant') |
||||
return emb |
||||
|
||||
def gen_c_embeddings(self, clip_txt, clip_txt_pooled, clip_img): |
||||
clip_txt = self.clip_txt_mapper(clip_txt) |
||||
if len(clip_txt_pooled.shape) == 2: |
||||
clip_txt_pooled = clip_txt_pooled.unsqueeze(1) |
||||
if len(clip_img.shape) == 2: |
||||
clip_img = clip_img.unsqueeze(1) |
||||
clip_txt_pool = self.clip_txt_pooled_mapper(clip_txt_pooled).view(clip_txt_pooled.size(0), clip_txt_pooled.size(1) * self.c_clip_seq, -1) |
||||
clip_img = self.clip_img_mapper(clip_img).view(clip_img.size(0), clip_img.size(1) * self.c_clip_seq, -1) |
||||
clip = torch.cat([clip_txt, clip_txt_pool, clip_img], dim=1) |
||||
clip = self.clip_norm(clip) |
||||
return clip |
||||
|
||||
def _down_encode(self, x, r_embed, clip, cnet=None): |
||||
level_outputs = [] |
||||
block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers) |
||||
for down_block, downscaler, repmap in block_group: |
||||
x = downscaler(x) |
||||
for i in range(len(repmap) + 1): |
||||
for block in down_block: |
||||
if isinstance(block, ResBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
ResBlock)): |
||||
if cnet is not None: |
||||
next_cnet = cnet() |
||||
if next_cnet is not None: |
||||
x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear', |
||||
align_corners=True) |
||||
x = block(x) |
||||
elif isinstance(block, AttnBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
AttnBlock)): |
||||
x = block(x, clip) |
||||
elif isinstance(block, TimestepBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
TimestepBlock)): |
||||
x = block(x, r_embed) |
||||
else: |
||||
x = block(x) |
||||
if i < len(repmap): |
||||
x = repmap[i](x) |
||||
level_outputs.insert(0, x) |
||||
return level_outputs |
||||
|
||||
def _up_decode(self, level_outputs, r_embed, clip, cnet=None): |
||||
x = level_outputs[0] |
||||
block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers) |
||||
for i, (up_block, upscaler, repmap) in enumerate(block_group): |
||||
for j in range(len(repmap) + 1): |
||||
for k, block in enumerate(up_block): |
||||
if isinstance(block, ResBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
ResBlock)): |
||||
skip = level_outputs[i] if k == 0 and i > 0 else None |
||||
if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)): |
||||
x = torch.nn.functional.interpolate(x, skip.shape[-2:], mode='bilinear', |
||||
align_corners=True) |
||||
if cnet is not None: |
||||
next_cnet = cnet() |
||||
if next_cnet is not None: |
||||
x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear', |
||||
align_corners=True) |
||||
x = block(x, skip) |
||||
elif isinstance(block, AttnBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
AttnBlock)): |
||||
x = block(x, clip) |
||||
elif isinstance(block, TimestepBlock) or ( |
||||
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module, |
||||
TimestepBlock)): |
||||
x = block(x, r_embed) |
||||
else: |
||||
x = block(x) |
||||
if j < len(repmap): |
||||
x = repmap[j](x) |
||||
x = upscaler(x) |
||||
return x |
||||
|
||||
def forward(self, x, r, clip_text, clip_text_pooled, clip_img, cnet=None, **kwargs): |
||||
# Process the conditioning embeddings |
||||
r_embed = self.gen_r_embedding(r).to(dtype=x.dtype) |
||||
for c in self.t_conds: |
||||
t_cond = kwargs.get(c, torch.zeros_like(r)) |
||||
r_embed = torch.cat([r_embed, self.gen_r_embedding(t_cond).to(dtype=x.dtype)], dim=1) |
||||
clip = self.gen_c_embeddings(clip_text, clip_text_pooled, clip_img) |
||||
|
||||
# Model Blocks |
||||
x = self.embedding(x) |
||||
if cnet is not None: |
||||
cnet = ControlNetDeliverer(cnet) |
||||
level_outputs = self._down_encode(x, r_embed, clip, cnet) |
||||
x = self._up_decode(level_outputs, r_embed, clip, cnet) |
||||
return self.clf(x) |
||||
|
||||
def update_weights_ema(self, src_model, beta=0.999): |
||||
for self_params, src_params in zip(self.parameters(), src_model.parameters()): |
||||
self_params.data = self_params.data * beta + src_params.data.clone().to(self_params.device) * (1 - beta) |
||||
for self_buffers, src_buffers in zip(self.buffers(), src_model.buffers()): |
||||
self_buffers.data = self_buffers.data * beta + src_buffers.data.clone().to(self_buffers.device) * (1 - beta) |
||||
@ -0,0 +1,95 @@
|
||||
""" |
||||
This file is part of ComfyUI. |
||||
Copyright (C) 2024 Stability AI |
||||
|
||||
This program is free software: you can redistribute it and/or modify |
||||
it under the terms of the GNU General Public License as published by |
||||
the Free Software Foundation, either version 3 of the License, or |
||||
(at your option) any later version. |
||||
|
||||
This program is distributed in the hope that it will be useful, |
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
||||
GNU General Public License for more details. |
||||
|
||||
You should have received a copy of the GNU General Public License |
||||
along with this program. If not, see <https://www.gnu.org/licenses/>. |
||||
""" |
||||
import torch |
||||
import torchvision |
||||
from torch import nn |
||||
|
||||
|
||||
# EfficientNet |
||||
class EfficientNetEncoder(nn.Module): |
||||
def __init__(self, c_latent=16): |
||||
super().__init__() |
||||
self.backbone = torchvision.models.efficientnet_v2_s().features.eval() |
||||
self.mapper = nn.Sequential( |
||||
nn.Conv2d(1280, c_latent, kernel_size=1, bias=False), |
||||
nn.BatchNorm2d(c_latent, affine=False), # then normalize them to have mean 0 and std 1 |
||||
) |
||||
self.mean = nn.Parameter(torch.tensor([0.485, 0.456, 0.406])) |
||||
self.std = nn.Parameter(torch.tensor([0.229, 0.224, 0.225])) |
||||
|
||||
def forward(self, x): |
||||
x = x * 0.5 + 0.5 |
||||
x = (x - self.mean.view([3,1,1])) / self.std.view([3,1,1]) |
||||
o = self.mapper(self.backbone(x)) |
||||
return o |
||||
|
||||
|
||||
# Fast Decoder for Stage C latents. E.g. 16 x 24 x 24 -> 3 x 192 x 192 |
||||
class Previewer(nn.Module): |
||||
def __init__(self, c_in=16, c_hidden=512, c_out=3): |
||||
super().__init__() |
||||
self.blocks = nn.Sequential( |
||||
nn.Conv2d(c_in, c_hidden, kernel_size=1), # 16 channels to 512 channels |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden), |
||||
|
||||
nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1), |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden), |
||||
|
||||
nn.ConvTranspose2d(c_hidden, c_hidden // 2, kernel_size=2, stride=2), # 16 -> 32 |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden // 2), |
||||
|
||||
nn.Conv2d(c_hidden // 2, c_hidden // 2, kernel_size=3, padding=1), |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden // 2), |
||||
|
||||
nn.ConvTranspose2d(c_hidden // 2, c_hidden // 4, kernel_size=2, stride=2), # 32 -> 64 |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden // 4), |
||||
|
||||
nn.Conv2d(c_hidden // 4, c_hidden // 4, kernel_size=3, padding=1), |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden // 4), |
||||
|
||||
nn.ConvTranspose2d(c_hidden // 4, c_hidden // 4, kernel_size=2, stride=2), # 64 -> 128 |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden // 4), |
||||
|
||||
nn.Conv2d(c_hidden // 4, c_hidden // 4, kernel_size=3, padding=1), |
||||
nn.GELU(), |
||||
nn.BatchNorm2d(c_hidden // 4), |
||||
|
||||
nn.Conv2d(c_hidden // 4, c_out, kernel_size=1), |
||||
) |
||||
|
||||
def forward(self, x): |
||||
return (self.blocks(x) - 0.5) * 2.0 |
||||
|
||||
class StageC_coder(nn.Module): |
||||
def __init__(self): |
||||
super().__init__() |
||||
self.previewer = Previewer() |
||||
self.encoder = EfficientNetEncoder() |
||||
|
||||
def encode(self, x): |
||||
return self.encoder(x) |
||||
|
||||
def decode(self, x): |
||||
return self.previewer(x) |
||||
@ -0,0 +1,109 @@
|
||||
""" |
||||
This file is part of ComfyUI. |
||||
Copyright (C) 2024 Stability AI |
||||
|
||||
This program is free software: you can redistribute it and/or modify |
||||
it under the terms of the GNU General Public License as published by |
||||
the Free Software Foundation, either version 3 of the License, or |
||||
(at your option) any later version. |
||||
|
||||
This program is distributed in the hope that it will be useful, |
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of |
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
||||
GNU General Public License for more details. |
||||
|
||||
You should have received a copy of the GNU General Public License |
||||
along with this program. If not, see <https://www.gnu.org/licenses/>. |
||||
""" |
||||
|
||||
import torch |
||||
import nodes |
||||
import comfy.utils |
||||
|
||||
|
||||
class StableCascade_EmptyLatentImage: |
||||
def __init__(self, device="cpu"): |
||||
self.device = device |
||||
|
||||
@classmethod |
||||
def INPUT_TYPES(s): |
||||
return {"required": { |
||||
"width": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), |
||||
"height": ("INT", {"default": 1024, "min": 256, "max": nodes.MAX_RESOLUTION, "step": 8}), |
||||
"compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), |
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}) |
||||
}} |
||||
RETURN_TYPES = ("LATENT", "LATENT") |
||||
RETURN_NAMES = ("stage_c", "stage_b") |
||||
FUNCTION = "generate" |
||||
|
||||
CATEGORY = "_for_testing/stable_cascade" |
||||
|
||||
def generate(self, width, height, compression, batch_size=1): |
||||
c_latent = torch.zeros([batch_size, 16, height // compression, width // compression]) |
||||
b_latent = torch.zeros([batch_size, 4, height // 4, width // 4]) |
||||
return ({ |
||||
"samples": c_latent, |
||||
}, { |
||||
"samples": b_latent, |
||||
}) |
||||
|
||||
class StableCascade_StageC_VAEEncode: |
||||
def __init__(self, device="cpu"): |
||||
self.device = device |
||||
|
||||
@classmethod |
||||
def INPUT_TYPES(s): |
||||
return {"required": { |
||||
"image": ("IMAGE",), |
||||
"vae": ("VAE", ), |
||||
"compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}), |
||||
}} |
||||
RETURN_TYPES = ("LATENT", "LATENT") |
||||
RETURN_NAMES = ("stage_c", "stage_b") |
||||
FUNCTION = "generate" |
||||
|
||||
CATEGORY = "_for_testing/stable_cascade" |
||||
|
||||
def generate(self, image, vae, compression): |
||||
width = image.shape[-2] |
||||
height = image.shape[-3] |
||||
out_width = (width // compression) * vae.downscale_ratio |
||||
out_height = (height // compression) * vae.downscale_ratio |
||||
|
||||
s = comfy.utils.common_upscale(image.movedim(-1,1), out_width, out_height, "bicubic", "center").movedim(1,-1) |
||||
|
||||
c_latent = vae.encode(s[:,:,:,:3]) |
||||
b_latent = torch.zeros([c_latent.shape[0], 4, height // 4, width // 4]) |
||||
return ({ |
||||
"samples": c_latent, |
||||
}, { |
||||
"samples": b_latent, |
||||
}) |
||||
|
||||
class StableCascade_StageB_Conditioning: |
||||
@classmethod |
||||
def INPUT_TYPES(s): |
||||
return {"required": { "conditioning": ("CONDITIONING",), |
||||
"stage_c": ("LATENT",), |
||||
}} |
||||
RETURN_TYPES = ("CONDITIONING",) |
||||
|
||||
FUNCTION = "set_prior" |
||||
|
||||
CATEGORY = "_for_testing/stable_cascade" |
||||
|
||||
def set_prior(self, conditioning, stage_c): |
||||
c = [] |
||||
for t in conditioning: |
||||
d = t[1].copy() |
||||
d['stable_cascade_prior'] = stage_c['samples'] |
||||
n = [t[0], d] |
||||
c.append(n) |
||||
return (c, ) |
||||
|
||||
NODE_CLASS_MAPPINGS = { |
||||
"StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage, |
||||
"StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning, |
||||
"StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode, |
||||
} |
||||
@ -0,0 +1,49 @@
|
||||
from PIL import Image, ImageOps |
||||
from io import BytesIO |
||||
import numpy as np |
||||
import struct |
||||
import comfy.utils |
||||
import time |
||||
|
||||
#You can use this node to save full size images through the websocket, the |
||||
#images will be sent in exactly the same format as the image previews: as |
||||
#binary images on the websocket with a 8 byte header indicating the type |
||||
#of binary message (first 4 bytes) and the image format (next 4 bytes). |
||||
|
||||
#The reason this node is disabled by default is because there is a small |
||||
#issue when using it with the default ComfyUI web interface: When generating |
||||
#batches only the last image will be shown in the UI. |
||||
|
||||
#Note that no metadata will be put in the images saved with this node. |
||||
|
||||
class SaveImageWebsocket: |
||||
@classmethod |
||||
def INPUT_TYPES(s): |
||||
return {"required": |
||||
{"images": ("IMAGE", ),} |
||||
} |
||||
|
||||
RETURN_TYPES = () |
||||
FUNCTION = "save_images" |
||||
|
||||
OUTPUT_NODE = True |
||||
|
||||
CATEGORY = "image" |
||||
|
||||
def save_images(self, images): |
||||
pbar = comfy.utils.ProgressBar(images.shape[0]) |
||||
step = 0 |
||||
for image in images: |
||||
i = 255. * image.cpu().numpy() |
||||
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) |
||||
pbar.update_absolute(step, images.shape[0], ("PNG", img, None)) |
||||
step += 1 |
||||
|
||||
return {} |
||||
|
||||
def IS_CHANGED(s, images): |
||||
return time.time() |
||||
|
||||
NODE_CLASS_MAPPINGS = { |
||||
"SaveImageWebsocket": SaveImageWebsocket, |
||||
} |
||||
Loading…
Reference in new issue