diff --git a/comfy_latent_interposer.py b/comfy_latent_interposer.py index 252dde6..6306187 100644 --- a/comfy_latent_interposer.py +++ b/comfy_latent_interposer.py @@ -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)