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/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": 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) diff --git a/comfy/model_management.py b/comfy/model_management.py index 762531d9f..a423bfc68 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: @@ -271,7 +264,6 @@ try: except: logging.warning("Could not pick default device.") -logging.info("VAE dtype: {}".format(VAE_DTYPE)) current_loaded_models = [] @@ -632,9 +624,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 f1e487132..cfbf8fa4d 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: @@ -189,7 +190,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() + 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 @@ -236,14 +238,15 @@ 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 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 @@ -264,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] @@ -294,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) @@ -321,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 diff --git a/comfy/taesd/taesd.py b/comfy/taesd/taesd.py index 8f96c54e5..ce36f1a84 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,12 +48,13 @@ 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)) + 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: @@ -69,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/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..06ea46216 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 @@ -668,8 +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 @@ -682,7 +695,7 @@ 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) 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) {