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

Update custom node to v3.1

Browse Source
This commit is contained in:
City
2023-10-11 06:32:14 +02:00
parent c8ed3d1128
commit 7df4f8a004
1 changed files with 62 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfy_latent_interposer.py
+62 -15
View File
@@ -1,3 +1,4 @@
import os
import torch
import torch.nn as nn
from safetensors.torch import load_file
@@ -9,21 +10,54 @@ class Interposer(nn.Module):
Basic NN layout, ported from:
https://github.com/city96/SD-Latent-Interposer/blob/main/interposer.py
"""
version = 1.1 # network revision
version = 3.1 # network revision
def __init__(self):
super().__init__()
module_list = [
nn.Conv2d(4, 32, kernel_size=5, padding=2),
self.chan = 4
self.hid = 128
self.head_join = nn.ReLU()
self.head_short = nn.Conv2d(self.chan, self.hid, kernel_size=3, stride=1, padding=1)
self.head_long = nn.Sequential(
nn.Conv2d(self.chan, self.hid, kernel_size=3, stride=1, padding=1),
nn.LeakyReLU(0.1),
nn.Conv2d(self.hid, self.hid, kernel_size=3, stride=1, padding=1),
nn.LeakyReLU(0.1),
nn.Conv2d(self.hid, self.hid, kernel_size=3, stride=1, padding=1),
)
self.core = nn.Sequential(
Block(self.hid),
Block(self.hid),
Block(self.hid),
)
self.tail = nn.Sequential(
nn.ReLU(),
nn.Conv2d(32, 128, kernel_size=7, padding=3),
nn.ReLU(),
nn.Conv2d(128, 32, kernel_size=7, padding=3),
nn.ReLU(),
nn.Conv2d(32, 4, kernel_size=5, padding=2),
]
self.sequential = nn.Sequential(*module_list)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.sequential(x)
nn.Conv2d(self.hid, self.chan, kernel_size=3, stride=1, padding=1)
)
def forward(self, x):
y = self.head_join(
self.head_long(x)+
self.head_short(x)
)
z = self.core(y)
return self.tail(z)
class Block(nn.Module):
def __init__(self, size):
super().__init__()
self.join = nn.ReLU()
self.long = nn.Sequential(
nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),
nn.LeakyReLU(0.1),
nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),
nn.LeakyReLU(0.1),
nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),
)
def forward(self, x):
y = self.long(x)
z = self.join(y + x)
return z
class LatentInterposer:
@@ -48,10 +82,23 @@ class LatentInterposer:
if latent_src == latent_dst:
return (samples,)
model = Interposer()
weights = str(hf_hub_download(
repo_id="city96/SD-Latent-Interposer",
filename=f"{latent_src}-to-{latent_dst}_interposer-v{model.version}.safetensors")
model.eval()
filename = f"{latent_src}-to-{latent_dst}_interposer-v{model.version}.safetensors"
local = os.path.join(
os.path.join(os.path.dirname(os.path.realpath(__file__)),"models"),
filename
)
if os.path.isfile(local):
print("LatentInterposer: Using local model")
weights = local
else:
print("LatentInterposer: Using HF Hub model")
weights = str(hf_hub_download(
repo_id="city96/SD-Latent-Interposer",
filename=filename)
)
model.load_state_dict(load_file(weights))
lt = samples["samples"]
lt = model(lt)