Update custom node to v3.1

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()
+		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=f"{latent_src}-to-{latent_dst}_interposer-v{model.version}.safetensors")
+				filename=filename)
 			)
+
 		model.load_state_dict(load_file(weights))
 		lt = samples["samples"]
 		lt = model(lt)