nodes/spacepxl-ComfyUI-VAE-Utils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add pad channel value and fix channel detection

Browse Source
This commit is contained in:
spacepxl
2025-12-20 00:43:41 -05:00
parent 0504284791
commit 40b839f15c
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/sd.py
+8 -2
View File
@@ -41,11 +41,13 @@ class CustomVAE(VAE):
self.latent_dim = 2
self.input_channels = 3
self.output_channels = 3
self.pad_channel_value = None
self.process_input = lambda image: image * 2.0 - 1.0
self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
self.working_dtypes = [torch.bfloat16, torch.float32]
self.disable_offload = False
self.not_video = False
self.size = None
self.downscale_index_formula = None
self.upscale_index_formula = None
@@ -247,13 +249,17 @@ class CustomVAE(VAE):
self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
self.downscale_index_formula = (4, 8, 8)
self.input_channels = sd["encoder.conv1.weight"].shape[1]
self.output_channels = sd["decoder.head.2.weight"].shape[0]
# self.output_channels = sd["decoder.head.2.weight"].shape[0]
self.output_channels = self.input_channels
raw_output_channels = sd["decoder.head.2.weight"].shape[0]
self.latent_dim = 3
self.latent_channels = 16
self.pad_channel_value = 1.0
ddconfig = {
"in_channels": self.input_channels,
"out_channels": self.output_channels,
# "out_channels": self.output_channels,
"out_channels": raw_output_channels,
"dim": 96,
"z_dim": self.latent_channels,
"dim_mult": [1, 2, 4, 4],