From df7db0e0279f58d9f2f3f33ddb60bb238b6d0dc8 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" <128333288+ltdrdata@users.noreply.github.com> Date: Sun, 16 Jun 2024 15:03:53 +0900 Subject: [PATCH 1/9] support TAESD3 (#3738) --- comfy/latent_formats.py | 1 + comfy/sd.py | 6 +++--- comfy/taesd/taesd.py | 15 ++++++++------- latent_preview.py | 2 +- nodes.py | 17 +++++++++++++++-- 5 files changed, 28 insertions(+), 13 deletions(-) diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 92f39d5c2..4b4a9eda2 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -129,6 +129,7 @@ class SD3(LatentFormat): [-0.0749, -0.0634, -0.0456], [-0.1418, -0.1457, -0.1259] ] + self.taesd_decoder_name = "taesd3_decoder" def process_in(self, latent): return (latent - self.shift_factor) * self.scale_factor diff --git a/comfy/sd.py b/comfy/sd.py index f1e487132..d2720ec13 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -166,7 +166,7 @@ class CLIP: return self.patcher.get_key_patches() class VAE: - def __init__(self, sd=None, device=None, config=None, dtype=None): + def __init__(self, sd=None, device=None, config=None, dtype=None, latent_channels=4): if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format sd = diffusers_convert.convert_vae_state_dict(sd) @@ -174,7 +174,7 @@ class VAE: self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype) self.downscale_ratio = 8 self.upscale_ratio = 8 - self.latent_channels = 4 + self.latent_channels = latent_channels self.output_channels = 3 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) @@ -189,7 +189,7 @@ class VAE: encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config}, decoder_config={'target': "comfy.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config}) elif "taesd_decoder.1.weight" in sd: - self.first_stage_model = comfy.taesd.taesd.TAESD() + self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=self.latent_channels) elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade self.first_stage_model = StageA() self.downscale_ratio = 4 diff --git a/comfy/taesd/taesd.py b/comfy/taesd/taesd.py index 8f96c54e5..74031c60d 100644 --- a/comfy/taesd/taesd.py +++ b/comfy/taesd/taesd.py @@ -25,18 +25,19 @@ class Block(nn.Module): def forward(self, x): return self.fuse(self.conv(x) + self.skip(x)) -def Encoder(): +def Encoder(latent_channels=4): return nn.Sequential( conv(3, 64), Block(64, 64), conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), conv(64, 64, stride=2, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), - conv(64, 4), + conv(64, latent_channels), ) -def Decoder(): + +def Decoder(latent_channels=4): return nn.Sequential( - Clamp(), conv(4, 64), nn.ReLU(), + Clamp(), conv(latent_channels, 64), nn.ReLU(), Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False), @@ -47,11 +48,11 @@ class TAESD(nn.Module): latent_magnitude = 3 latent_shift = 0.5 - def __init__(self, encoder_path=None, decoder_path=None): + def __init__(self, encoder_path=None, decoder_path=None, latent_channels=4): """Initialize pretrained TAESD on the given device from the given checkpoints.""" super().__init__() - self.taesd_encoder = Encoder() - self.taesd_decoder = Decoder() + self.taesd_encoder = Encoder(latent_channels=latent_channels) + self.taesd_decoder = Decoder(latent_channels=latent_channels) self.vae_scale = torch.nn.Parameter(torch.tensor(1.0)) if encoder_path is not None: self.taesd_encoder.load_state_dict(comfy.utils.load_torch_file(encoder_path, safe_load=True)) diff --git a/latent_preview.py b/latent_preview.py index 54aa233f2..ae6c106e4 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -64,7 +64,7 @@ def get_previewer(device, latent_format): if method == LatentPreviewMethod.TAESD: if taesd_decoder_path: - taesd = TAESD(None, taesd_decoder_path).to(device) + taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device) previewer = TAESDPreviewerImpl(taesd) else: logging.warning("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name)) diff --git a/nodes.py b/nodes.py index 0b2a96f73..ca10ca32a 100644 --- a/nodes.py +++ b/nodes.py @@ -634,6 +634,8 @@ class VAELoader: sdxl_taesd_dec = False sd1_taesd_enc = False sd1_taesd_dec = False + sd3_taesd_enc = False + sd3_taesd_dec = False for v in approx_vaes: if v.startswith("taesd_decoder."): @@ -644,10 +646,16 @@ class VAELoader: sdxl_taesd_dec = True elif v.startswith("taesdxl_encoder."): sdxl_taesd_enc = True + elif v.startswith("taesd3_decoder."): + sd3_taesd_dec = True + elif v.startswith("taesd3_encoder."): + sd3_taesd_enc = True if sd1_taesd_dec and sd1_taesd_enc: vaes.append("taesd") if sdxl_taesd_dec and sdxl_taesd_enc: vaes.append("taesdxl") + if sd3_taesd_dec and sd3_taesd_enc: + vaes.append("taesd3") return vaes @staticmethod @@ -670,6 +678,8 @@ class VAELoader: sd["vae_scale"] = torch.tensor(0.18215) elif name == "taesdxl": sd["vae_scale"] = torch.tensor(0.13025) + elif name == "taesd3": + sd["vae_scale"] = torch.tensor(1.5305) return sd @classmethod @@ -682,12 +692,15 @@ class VAELoader: #TODO: scale factor? def load_vae(self, vae_name): - if vae_name in ["taesd", "taesdxl"]: + if vae_name in ["taesd", "taesdxl", "taesd3"]: sd = self.load_taesd(vae_name) else: vae_path = folder_paths.get_full_path("vae", vae_name) sd = comfy.utils.load_torch_file(vae_path) - vae = comfy.sd.VAE(sd=sd) + + latent_channels = 16 if vae_name == 'taesd3' else 4 + + vae = comfy.sd.VAE(sd=sd, latent_channels=latent_channels) return (vae,) class ControlNetLoader: From 04e8798c37d958d74ea6bda506b86f51356d6caf Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 16 Jun 2024 02:04:24 -0400 Subject: [PATCH 2/9] Improvements to the TAESD3 implementation. --- comfy/sd.py | 6 +++--- comfy/taesd/taesd.py | 5 +++-- nodes.py | 8 ++++---- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index d2720ec13..8d90dd063 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -166,7 +166,7 @@ class CLIP: return self.patcher.get_key_patches() class VAE: - def __init__(self, sd=None, device=None, config=None, dtype=None, latent_channels=4): + def __init__(self, sd=None, device=None, config=None, dtype=None): if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format sd = diffusers_convert.convert_vae_state_dict(sd) @@ -174,7 +174,7 @@ class VAE: self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype) self.downscale_ratio = 8 self.upscale_ratio = 8 - self.latent_channels = latent_channels + self.latent_channels = 4 self.output_channels = 3 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) @@ -189,7 +189,7 @@ class VAE: encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config}, decoder_config={'target': "comfy.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config}) elif "taesd_decoder.1.weight" in sd: - self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=self.latent_channels) + self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=sd["taesd_decoder.1.weight"].shape[1]) elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade self.first_stage_model = StageA() self.downscale_ratio = 4 diff --git a/comfy/taesd/taesd.py b/comfy/taesd/taesd.py index 74031c60d..ce36f1a84 100644 --- a/comfy/taesd/taesd.py +++ b/comfy/taesd/taesd.py @@ -54,6 +54,7 @@ class TAESD(nn.Module): self.taesd_encoder = Encoder(latent_channels=latent_channels) self.taesd_decoder = Decoder(latent_channels=latent_channels) self.vae_scale = torch.nn.Parameter(torch.tensor(1.0)) + self.vae_shift = torch.nn.Parameter(torch.tensor(0.0)) if encoder_path is not None: self.taesd_encoder.load_state_dict(comfy.utils.load_torch_file(encoder_path, safe_load=True)) if decoder_path is not None: @@ -70,9 +71,9 @@ class TAESD(nn.Module): return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude) def decode(self, x): - x_sample = self.taesd_decoder(x * self.vae_scale) + x_sample = self.taesd_decoder((x - self.vae_shift) * self.vae_scale) x_sample = x_sample.sub(0.5).mul(2) return x_sample def encode(self, x): - return self.taesd_encoder(x * 0.5 + 0.5) / self.vae_scale + return (self.taesd_encoder(x * 0.5 + 0.5) / self.vae_scale) + self.vae_shift diff --git a/nodes.py b/nodes.py index ca10ca32a..06ea46216 100644 --- a/nodes.py +++ b/nodes.py @@ -676,10 +676,13 @@ class VAELoader: if name == "taesd": sd["vae_scale"] = torch.tensor(0.18215) + sd["vae_shift"] = torch.tensor(0.0) elif name == "taesdxl": sd["vae_scale"] = torch.tensor(0.13025) + sd["vae_shift"] = torch.tensor(0.0) elif name == "taesd3": sd["vae_scale"] = torch.tensor(1.5305) + sd["vae_shift"] = torch.tensor(0.0609) return sd @classmethod @@ -697,10 +700,7 @@ class VAELoader: else: vae_path = folder_paths.get_full_path("vae", vae_name) sd = comfy.utils.load_torch_file(vae_path) - - latent_channels = 16 if vae_name == 'taesd3' else 4 - - vae = comfy.sd.VAE(sd=sd, latent_channels=latent_channels) + vae = comfy.sd.VAE(sd=sd) return (vae,) class ControlNetLoader: From 746a0410d43b02dcaf91ac5abc2c1803e420203d Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 16 Jun 2024 03:10:04 -0400 Subject: [PATCH 3/9] Fix VAEEncode with taesd3. --- comfy/sd.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index 8d90dd063..82f9aeab8 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -189,7 +189,8 @@ class VAE: encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config}, decoder_config={'target': "comfy.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config}) elif "taesd_decoder.1.weight" in sd: - self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=sd["taesd_decoder.1.weight"].shape[1]) + self.latent_channels = sd["taesd_decoder.1.weight"].shape[1] + self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=self.latent_channels) elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade self.first_stage_model = StageA() self.downscale_ratio = 4 From ca9d300a804fd1bfc67b3de9200c2d09b78899d0 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 16 Jun 2024 11:47:32 -0400 Subject: [PATCH 4/9] Better estimation for memory usage during audio VAE encoding/decoding. --- comfy/sd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 82f9aeab8..e16cd8e53 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -237,8 +237,8 @@ class VAE: decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig}) elif "decoder.layers.0.weight_v" in sd: self.first_stage_model = AudioOobleckVAE() - self.memory_used_encode = lambda shape, dtype: (1767 * shape[2]) * model_management.dtype_size(dtype) #TODO: tweak for the audio VAE - self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * 64) * model_management.dtype_size(dtype) + self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype) + self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * 2048) * model_management.dtype_size(dtype) self.latent_channels = 64 self.output_channels = 2 self.upscale_ratio = 2048 From 8ddc151a4ce7f59de75b24ea8349f0e406ba0da5 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 16 Jun 2024 13:06:23 -0400 Subject: [PATCH 5/9] Squash depreciation warning on new pytorch. --- comfy/ldm/audio/autoencoder.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/audio/autoencoder.py b/comfy/ldm/audio/autoencoder.py index 7363131e0..8123e66a5 100644 --- a/comfy/ldm/audio/autoencoder.py +++ b/comfy/ldm/audio/autoencoder.py @@ -75,10 +75,16 @@ class SnakeBeta(nn.Module): return x def WNConv1d(*args, **kwargs): - return torch.nn.utils.weight_norm(ops.Conv1d(*args, **kwargs)) + try: + return torch.nn.utils.parametrizations.weight_norm(ops.Conv1d(*args, **kwargs)) + except: + return torch.nn.utils.weight_norm(ops.Conv1d(*args, **kwargs)) #support pytorch 2.1 and older def WNConvTranspose1d(*args, **kwargs): - return torch.nn.utils.weight_norm(ops.ConvTranspose1d(*args, **kwargs)) + try: + return torch.nn.utils.parametrizations.weight_norm(ops.ConvTranspose1d(*args, **kwargs)) + except: + return torch.nn.utils.weight_norm(ops.ConvTranspose1d(*args, **kwargs)) #support pytorch 2.1 and older def get_activation(activation: Literal["elu", "snake", "none"], antialias=False, channels=None) -> nn.Module: if activation == "elu": From 6425252c4f2f6acd8f4ad59a2135f5bdae3452e4 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 16 Jun 2024 13:12:54 -0400 Subject: [PATCH 6/9] Use fp16 as the default vae dtype for the audio VAE. --- comfy/model_management.py | 35 ++++++++++++++++++++--------------- comfy/sd.py | 5 ++++- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 8b8d3ff06..047193290 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -167,7 +167,7 @@ if args.use_pytorch_cross_attention: ENABLE_PYTORCH_ATTENTION = True XFORMERS_IS_AVAILABLE = False -VAE_DTYPE = torch.float32 +VAE_DTYPES = [torch.float32] try: if is_nvidia(): @@ -176,7 +176,7 @@ try: if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False: ENABLE_PYTORCH_ATTENTION = True if torch.cuda.is_bf16_supported() and torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8: - VAE_DTYPE = torch.bfloat16 + VAE_DTYPES = [torch.bfloat16] + VAE_DTYPES if is_intel_xpu(): if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: ENABLE_PYTORCH_ATTENTION = True @@ -184,17 +184,10 @@ except: pass if is_intel_xpu(): - VAE_DTYPE = torch.bfloat16 + VAE_DTYPES = [torch.bfloat16] + VAE_DTYPES if args.cpu_vae: - VAE_DTYPE = torch.float32 - -if args.fp16_vae: - VAE_DTYPE = torch.float16 -elif args.bf16_vae: - VAE_DTYPE = torch.bfloat16 -elif args.fp32_vae: - VAE_DTYPE = torch.float32 + VAE_DTYPES = [torch.float32] if ENABLE_PYTORCH_ATTENTION: @@ -258,7 +251,6 @@ try: except: logging.warning("Could not pick default device.") -logging.info("VAE dtype: {}".format(VAE_DTYPE)) current_loaded_models = [] @@ -619,9 +611,22 @@ def vae_offload_device(): else: return torch.device("cpu") -def vae_dtype(): - global VAE_DTYPE - return VAE_DTYPE +def vae_dtype(device=None, allowed_dtypes=[]): + global VAE_DTYPES + if args.fp16_vae: + return torch.float16 + elif args.bf16_vae: + return torch.bfloat16 + elif args.fp32_vae: + return torch.float32 + + for d in allowed_dtypes: + if d == torch.float16 and should_use_fp16(device, prioritize_performance=False): + return d + if d in VAE_DTYPES: + return d + + return VAE_DTYPES[0] def get_autocast_device(dev): if hasattr(dev, 'type'): diff --git a/comfy/sd.py b/comfy/sd.py index e16cd8e53..58a858aa0 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -178,6 +178,7 @@ class VAE: self.output_channels = 3 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] if config is None: if "decoder.mid.block_1.mix_factor" in sd: @@ -245,6 +246,7 @@ class VAE: self.downscale_ratio = 2048 self.process_output = lambda audio: audio self.process_input = lambda audio: audio + self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32] else: logging.warning("WARNING: No VAE weights detected, VAE not initalized.") self.first_stage_model = None @@ -265,12 +267,13 @@ class VAE: self.device = device offload_device = model_management.vae_offload_device() if dtype is None: - dtype = model_management.vae_dtype() + dtype = model_management.vae_dtype(self.device, self.working_dtypes) self.vae_dtype = dtype self.first_stage_model.to(self.vae_dtype) self.output_device = model_management.intermediate_device() self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device) + logging.debug("VAE load device: {}, offload device: {}, dtype: {}".format(self.device, offload_device, self.vae_dtype)) def vae_encode_crop_pixels(self, pixels): dims = pixels.shape[1:-1] From b7c473d1abda165e1ff6f90dfacc6f589c3d1c38 Mon Sep 17 00:00:00 2001 From: Janek Mann Date: Mon, 17 Jun 2024 12:55:06 +0100 Subject: [PATCH 7/9] Fix lora keys for SimpleTuner (#3759) --- comfy/lora.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/comfy/lora.py b/comfy/lora.py index 37254b03f..082a8b3cb 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -30,6 +30,7 @@ def load_lora(lora, to_load): regular_lora = "{}.lora_up.weight".format(x) diffusers_lora = "{}_lora.up.weight".format(x) diffusers2_lora = "{}.lora_B.weight".format(x) + diffusers3_lora = "{}.lora.up.weight".format(x) transformers_lora = "{}.lora_linear_layer.up.weight".format(x) A_name = None @@ -45,6 +46,10 @@ def load_lora(lora, to_load): A_name = diffusers2_lora B_name = "{}.lora_A.weight".format(x) mid_name = None + elif diffusers3_lora in lora.keys(): + A_name = diffusers3_lora + B_name = "{}.lora.down.weight".format(x) + mid_name = None elif transformers_lora in lora.keys(): A_name = transformers_lora B_name ="{}.lora_linear_layer.down.weight".format(x) From 379ff92e9e872a588ad2ceeabdd08b909590b25b Mon Sep 17 00:00:00 2001 From: Juanjuan <472185568@qq.com> Date: Mon, 17 Jun 2024 19:56:53 +0800 Subject: [PATCH 8/9] fix app.js no graph defined (#3754) * local test * fix "graph" not found * fix --------- Co-authored-by: Xiujuan Li --- web/scripts/app.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/scripts/app.js b/web/scripts/app.js index f96d197a8..58fb765c4 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -2239,7 +2239,7 @@ export class ComfyApp { const node = LiteGraph.createNode(data.class_type); node.id = isNaN(+id) ? id : +id; node.title = data._meta?.title ?? node.title - graph.add(node); + app.graph.add(node); } for (const id of ids) { From a45df695706e8b80a66ff6815789582d88571326 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 17 Jun 2024 22:48:23 -0400 Subject: [PATCH 9/9] Basic tiled decoding for audio VAE. --- comfy/sd.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/comfy/sd.py b/comfy/sd.py index 58a858aa0..cfbf8fa4d 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -298,6 +298,17 @@ class VAE: / 3.0) return output + def decode_tiled_1d(self, samples, tile_x=128, overlap=64): + output = torch.empty((samples.shape[0], self.output_channels) + tuple(map(lambda a: a * self.upscale_ratio, samples.shape[2:])), device=self.output_device) + + for j in range(samples.shape[0]): + for i in range(0, samples.shape[-1], tile_x - overlap): + f = i + t = i + tile_x + output[j:j+1,:,f * self.upscale_ratio:t * self.upscale_ratio] = self.first_stage_model.decode(samples[j:j+1,:,f:t].to(self.vae_dtype).to(self.device)).float() + + return output + def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64): steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap) steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap) @@ -325,7 +336,10 @@ class VAE: pixel_samples[x:x+batch_number] = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float()) except model_management.OOM_EXCEPTION as e: logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.") - pixel_samples = self.decode_tiled_(samples_in) + if len(samples_in.shape) == 3: + pixel_samples = self.decode_tiled_1d(samples_in) + else: + pixel_samples = self.decode_tiled_(samples_in) pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1) return pixel_samples