xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-12 22:57:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

remove guess_latent_channels and add latent_channels info directly to flux model

Browse Source
This commit is contained in:
Matthew Turnshek 2024-08-16 12:06:33 -04:00
parent e792dec45f
commit 1c14fc89db
2 changed files with 2 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/latent_formats.py
View File

@ -141,6 +141,7 @@ class StableAudio1(LatentFormat):
latent_channels = 64 latent_channels = 64
class Flux(SD3): class Flux(SD3):
latent_channels = 16
def __init__(self): def __init__(self):
self.scale_factor = 0.3611 self.scale_factor = 0.3611
self.shift_factor = 0.1159 self.shift_factor = 0.1159

12
comfy/taesd/taesd.py
View File

@ -48,11 +48,9 @@ class TAESD(nn.Module):
latent_magnitude = 3 latent_magnitude = 3
latent_shift = 0.5 latent_shift = 0.5
def __init__(self, encoder_path=None, decoder_path=None, latent_channels=None): def __init__(self, encoder_path=None, decoder_path=None, latent_channels=4):
"""Initialize pretrained TAESD on the given device from the given checkpoints.""" """Initialize pretrained TAESD on the given device from the given checkpoints."""
super().__init__() super().__init__()
if latent_channels is None:
latent_channels = self.guess_latent_channels(str(encoder_path))
self.taesd_encoder = Encoder(latent_channels=latent_channels) self.taesd_encoder = Encoder(latent_channels=latent_channels)
self.taesd_decoder = Decoder(latent_channels=latent_channels) self.taesd_decoder = Decoder(latent_channels=latent_channels)
self.vae_scale = torch.nn.Parameter(torch.tensor(1.0)) self.vae_scale = torch.nn.Parameter(torch.tensor(1.0))
@ -62,14 +60,6 @@ class TAESD(nn.Module):
if decoder_path is not None: if decoder_path is not None:
self.taesd_decoder.load_state_dict(comfy.utils.load_torch_file(decoder_path, safe_load=True)) self.taesd_decoder.load_state_dict(comfy.utils.load_torch_file(decoder_path, safe_load=True))
def guess_latent_channels(self, encoder_path):
"""guess latent channel count based on encoder filename"""
if "taef1" in encoder_path:
return 16
if "taesd3" in encoder_path:
return 16
return 4
@staticmethod @staticmethod
def scale_latents(x): def scale_latents(x):
"""raw latents -> [0, 1]""" """raw latents -> [0, 1]"""