nodes/city96-ComfyUI_ExtraModels
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Consistency decoder full

Browse Source 
With model definitions from https://gist.github.com/madebyollin/865fa6a18d9099351ddbdfbe7299ccbf
This commit is contained in:
City
2023-11-07 21:26:54 +01:00
parent 4fb1292c50
commit a3e1804def
3 changed files with 284 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
VAE/loader.py
+15 -16
View File
@@ -19,28 +19,27 @@ class EXVAE(comfy.sd.VAE):
self.offload_device = model_management.vae_offload_device()
self.vae_dtype = vae_dtype_dict.get(dtype, "auto")
sd = None
sd = comfy.utils.load_torch_file(model_path)
model = None
if model_conf["type"] == "AutoencoderKL":
from .models.kl import AutoencoderKL
model = AutoencoderKL(config=model_conf)
sd = comfy.utils.load_torch_file(model_path)
if model_conf["type"] == "VQModel":
from .models.vq import VQModel
model = VQModel(config=model_conf)
sd = comfy.utils.load_torch_file(model_path)
if model_conf["type"] == "ConsistencyDecoder":
from .models.consistencydecoder import ConsistencyDecoder
model = ConsistencyDecoder(model_path, self.device)
sd = model.ckpt.state_dict()
if sd:
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys():
sd = diffusers_convert.convert_vae_state_dict(sd)
self.first_stage_model = model.eval()
m, u = self.first_stage_model.load_state_dict(sd, strict=False)
if len(m) > 0: print("Missing VAE keys", m)
if len(u) > 0: print("Leftover VAE keys", u)
elif model_conf["type"] == "VQModel":
from .models.vq import VQModel
model = VQModel(config=model_conf)
elif model_conf["type"] == "ConsistencyDecoder":
from .models.consistencydecoder import ConsistencyDecoder
model = ConsistencyDecoder()
sd = {f"model.{k}":v for k,v in sd.items()}
else:
raise NotImplementedError(f"Unknown VAE type '{model_conf['type']}'")
self.first_stage_model = model.eval()
m, u = self.first_stage_model.load_state_dict(sd, strict=False)
if len(m) > 0: print("Missing VAE keys", m)
if len(u) > 0: print("Leftover VAE keys", u)
self.first_stage_model.to(self.vae_dtype).to(self.offload_device)
VAE/models/consistencydecoder.py
+265 -5
View File
@@ -1,5 +1,11 @@
import math
import torch
import torch.nn.functional as F
import torch.nn as nn
"""
Code below ported from https://github.com/openai/consistencydecoder
"""
def _extract_into_tensor(arr, timesteps, broadcast_shape):
# from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L895 """
@@ -18,9 +24,9 @@ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
class ConsistencyDecoder(torch.nn.Module):
# From https://github.com/openai/consistencydecoder
def __init__(self, model_path, device):
def __init__(self):
super().__init__()
self.ckpt = torch.jit.load(model_path, map_location=device)
self.model = ConvUNetVAE()
self.n_distilled_steps = 64
sigma_data = 0.5
@@ -91,8 +97,8 @@ class ConsistencyDecoder(torch.nn.Module):
schedule_timesteps = [int((1024 - 1) * s) for s in schedule]
for i in schedule_timesteps:
t = ts[i].item()
t_ = torch.tensor([t] * features.shape[0]).to(features.device)
noise = torch.randn_like(x_start)
t_ = torch.tensor([t] * features.shape[0], device=features.device)
noise = torch.randn_like(x_start, device=features.device)
x_start = (
_extract_into_tensor(self.sqrt_alphas_cumprod, t_, x_start.shape)
* x_start
@@ -102,7 +108,7 @@ class ConsistencyDecoder(torch.nn.Module):
* noise
)
c_in = _extract_into_tensor(self.c_in, t_, x_start.shape)
model_output = self.ckpt(c_in * x_start, t_, features=features)
model_output = self.model((c_in * x_start).to(features.dtype), t_, features=features)
B, C = x_start.shape[:2]
model_output, _ = torch.split(model_output, C, dim=1)
pred_xstart = (
@@ -114,3 +120,257 @@ class ConsistencyDecoder(torch.nn.Module):
def encode(self, *args, **kwargs):
raise NotImplementedError("ConsistencyDecoder can't be used for encoding!")
"""
Model definitions ported from:
https://gist.github.com/madebyollin/865fa6a18d9099351ddbdfbe7299ccbf
https://gist.github.com/mrsteyk/74ad3ec2f6f823111ae4c90e168505ac.
"""
class TimestepEmbedding(nn.Module):
def __init__(self, n_time=1024, n_emb=320, n_out=1280) -> None:
super().__init__()
self.emb = nn.Embedding(n_time, n_emb)
self.f_1 = nn.Linear(n_emb, n_out)
self.f_2 = nn.Linear(n_out, n_out)
def forward(self, x) -> torch.Tensor:
x = self.emb(x)
x = self.f_1(x)
x = F.silu(x)
return self.f_2(x)
class ImageEmbedding(nn.Module):
def __init__(self, in_channels=7, out_channels=320) -> None:
super().__init__()
self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
def forward(self, x) -> torch.Tensor:
return self.f(x)
class ImageUnembedding(nn.Module):
def __init__(self, in_channels=320, out_channels=6) -> None:
super().__init__()
self.gn = nn.GroupNorm(32, in_channels)
self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
def forward(self, x) -> torch.Tensor:
return self.f(F.silu(self.gn(x)))
class ConvResblock(nn.Module):
def __init__(self, in_features=320, out_features=320) -> None:
super().__init__()
self.f_t = nn.Linear(1280, out_features * 2)
self.gn_1 = nn.GroupNorm(32, in_features)
self.f_1 = nn.Conv2d(in_features, out_features, kernel_size=3, padding=1)
self.gn_2 = nn.GroupNorm(32, out_features)
self.f_2 = nn.Conv2d(out_features, out_features, kernel_size=3, padding=1)
skip_conv = in_features != out_features
self.f_s = (
nn.Conv2d(in_features, out_features, kernel_size=1, padding=0)
if skip_conv
else nn.Identity()
)
def forward(self, x, t):
x_skip = x
t = self.f_t(F.silu(t))
t = t.chunk(2, dim=1)
t_1 = t[0].unsqueeze(dim=2).unsqueeze(dim=3) + 1
t_2 = t[1].unsqueeze(dim=2).unsqueeze(dim=3)
gn_1 = F.silu(self.gn_1(x))
f_1 = self.f_1(gn_1)
gn_2 = self.gn_2(f_1)
return self.f_s(x_skip) + self.f_2(F.silu(gn_2 * t_1 + t_2))
# Also ConvResblock
class Downsample(nn.Module):
def __init__(self, in_channels=320) -> None:
super().__init__()
self.f_t = nn.Linear(1280, in_channels * 2)
self.gn_1 = nn.GroupNorm(32, in_channels)
self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
self.gn_2 = nn.GroupNorm(32, in_channels)
self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
def forward(self, x, t) -> torch.Tensor:
x_skip = x
t = self.f_t(F.silu(t))
t_1, t_2 = t.chunk(2, dim=1)
t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1
t_2 = t_2.unsqueeze(2).unsqueeze(3)
gn_1 = F.silu(self.gn_1(x))
avg_pool2d = F.avg_pool2d(gn_1, kernel_size=(2, 2), stride=None)
f_1 = self.f_1(avg_pool2d)
gn_2 = self.gn_2(f_1)
f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2)))
return f_2 + F.avg_pool2d(x_skip, kernel_size=(2, 2), stride=None)
# Also ConvResblock
class Upsample(nn.Module):
def __init__(self, in_channels=1024) -> None:
super().__init__()
self.f_t = nn.Linear(1280, in_channels * 2)
self.gn_1 = nn.GroupNorm(32, in_channels)
self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
self.gn_2 = nn.GroupNorm(32, in_channels)
self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
def forward(self, x, t) -> torch.Tensor:
x_skip = x
t = self.f_t(F.silu(t))
t_1, t_2 = t.chunk(2, dim=1)
t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1
t_2 = t_2.unsqueeze(2).unsqueeze(3)
gn_1 = F.silu(self.gn_1(x))
upsample = F.interpolate(gn_1.float(), scale_factor=2, mode="nearest").to(gn_1.dtype)
f_1 = self.f_1(upsample)
gn_2 = self.gn_2(f_1)
f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2)))
return f_2 + F.interpolate(x_skip.float(), scale_factor=2, mode="nearest").to(x_skip.dtype)
class ConvUNetVAE(nn.Module):
def __init__(self) -> None:
super().__init__()
self.embed_image = ImageEmbedding()
self.embed_time = TimestepEmbedding()
down_0 = nn.ModuleList(
[
ConvResblock(320, 320),
ConvResblock(320, 320),
ConvResblock(320, 320),
Downsample(320),
]
)
down_1 = nn.ModuleList(
[
ConvResblock(320, 640),
ConvResblock(640, 640),
ConvResblock(640, 640),
Downsample(640),
]
)
down_2 = nn.ModuleList(
[
ConvResblock(640, 1024),
ConvResblock(1024, 1024),
ConvResblock(1024, 1024),
Downsample(1024),
]
)
down_3 = nn.ModuleList(
[
ConvResblock(1024, 1024),
ConvResblock(1024, 1024),
ConvResblock(1024, 1024),
]
)
self.down = nn.ModuleList(
[
down_0,
down_1,
down_2,
down_3,
]
)
self.mid = nn.ModuleList(
[
ConvResblock(1024, 1024),
ConvResblock(1024, 1024),
]
)
up_3 = nn.ModuleList(
[
ConvResblock(1024 * 2, 1024),
ConvResblock(1024 * 2, 1024),
ConvResblock(1024 * 2, 1024),
ConvResblock(1024 * 2, 1024),
Upsample(1024),
]
)
up_2 = nn.ModuleList(
[
ConvResblock(1024 * 2, 1024),
ConvResblock(1024 * 2, 1024),
ConvResblock(1024 * 2, 1024),
ConvResblock(1024 + 640, 1024),
Upsample(1024),
]
)
up_1 = nn.ModuleList(
[
ConvResblock(1024 + 640, 640),
ConvResblock(640 * 2, 640),
ConvResblock(640 * 2, 640),
ConvResblock(320 + 640, 640),
Upsample(640),
]
)
up_0 = nn.ModuleList(
[
ConvResblock(320 + 640, 320),
ConvResblock(320 * 2, 320),
ConvResblock(320 * 2, 320),
ConvResblock(320 * 2, 320),
]
)
self.up = nn.ModuleList(
[
up_0,
up_1,
up_2,
up_3,
]
)
self.output = ImageUnembedding()
def forward(self, x, t, features) -> torch.Tensor:
x = torch.cat([x, F.interpolate(features.float(),scale_factor=8,mode="nearest").to(features.dtype)], dim=1)
t = self.embed_time(t)
x = self.embed_image(x)
skips = [x]
for down in self.down:
for block in down:
x = block(x, t)
skips.append(x)
for i in range(2):
x = self.mid[i](x, t)
for up in self.up[::-1]:
for block in up:
if isinstance(block, ConvResblock):
x = torch.concat([x, skips.pop()], dim=1)
x = block(x, t)
return self.output(x)