diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 6a57bca1c..f2fe854c5 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -224,6 +224,7 @@ class Flux2(LatentFormat): self.latent_rgb_factors_bias = [-0.0329, -0.0718, -0.0851] self.latent_rgb_factors_reshape = lambda t: t.reshape(t.shape[0], 32, 2, 2, t.shape[-2], t.shape[-1]).permute(0, 1, 4, 2, 5, 3).reshape(t.shape[0], 32, t.shape[-2] * 2, t.shape[-1] * 2) + self.taesd_decoder_name = "taef2_decoder" def process_in(self, latent): return latent diff --git a/comfy/taesd/taesd.py b/comfy/taesd/taesd.py index ce36f1a84..80121a637 100644 --- a/comfy/taesd/taesd.py +++ b/comfy/taesd/taesd.py @@ -17,32 +17,79 @@ class Clamp(nn.Module): return torch.tanh(x / 3) * 3 class Block(nn.Module): - def __init__(self, n_in, n_out): + def __init__(self, n_in: int, n_out: int, use_midblock_gn: bool = False): super().__init__() self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out)) self.skip = comfy.ops.disable_weight_init.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() self.fuse = nn.ReLU() - def forward(self, x): + if not use_midblock_gn: + self.pool = None + return + n_gn = n_in * 4 + self.pool = nn.Sequential( + comfy.ops.disable_weight_init.Conv2d(n_in, n_gn, 1, bias=False), + comfy.ops.disable_weight_init.GroupNorm(4, n_gn), + nn.ReLU(inplace=True), + comfy.ops.disable_weight_init.Conv2d(n_gn, n_in, 1, bias=False), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.pool is not None: + x = x + self.pool(x) return self.fuse(self.conv(x) + self.skip(x)) -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, latent_channels), - ) +class Encoder(nn.Sequential): + def __init__(self, latent_channels: int = 4, use_gn: bool = False): + super().__init__( + 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, use_gn), Block(64, 64, use_gn), Block(64, 64, use_gn), + conv(64, latent_channels), + ) +class Decoder(nn.Sequential): + def __init__(self, latent_channels: int = 4, use_gn: bool = False): + super().__init__( + Clamp(), conv(latent_channels, 64), nn.ReLU(), + Block(64, 64, use_gn), Block(64, 64, use_gn), Block(64, 64, use_gn), 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), + Block(64, 64), conv(64, 3), + ) + +class DecoderFlux2(Decoder): + def __init__(self, latent_channels: int = 128, use_gn: bool = True): + if latent_channels != 128 or not use_gn: + raise ValueError("Unexpected parameters for Flux2 TAE module") + super().__init__(latent_channels=32, use_gn=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, C, H, W = x.shape + x = ( + x + .reshape(B, 32, 2, 2, H, W) + .permute(0, 1, 4, 2, 5, 3) + .reshape(B, 32, H * 2, W * 2) + ) + return super().forward(x) + +class EncoderFlux2(Decoder): + def __init__(self, latent_channels: int = 128, use_gn: bool = True): + if latent_channels != 128 or not use_gn: + raise ValueError("Unexpected parameters for Flux2 TAE module") + super().__init__(latent_channels=32, use_gn=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + result = super().forward(x) + B, C, H, W = result.shape + return ( + result + .reshape(B, C, H // 2, 2, W // 2, 2) + .permute(0, 1, 3, 5, 2, 4) + .reshape(B, 128, H // 2, W // 2) + ) -def Decoder(latent_channels=4): - return nn.Sequential( - 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), - Block(64, 64), conv(64, 3), - ) class TAESD(nn.Module): latent_magnitude = 3 @@ -51,8 +98,15 @@ class TAESD(nn.Module): 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(latent_channels=latent_channels) - self.taesd_decoder = Decoder(latent_channels=latent_channels) + if latent_channels == 128: + encoder_class = EncoderFlux2 + decoder_class = DecoderFlux2 + else: + encoder_class = Encoder + decoder_class = Decoder + self.taesd_encoder = encoder_class(latent_channels=latent_channels) + self.taesd_decoder = decoder_class(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: @@ -61,19 +115,19 @@ class TAESD(nn.Module): self.taesd_decoder.load_state_dict(comfy.utils.load_torch_file(decoder_path, safe_load=True)) @staticmethod - def scale_latents(x): + def scale_latents(x: torch.Tensor) -> torch.Tensor: """raw latents -> [0, 1]""" return x.div(2 * TAESD.latent_magnitude).add(TAESD.latent_shift).clamp(0, 1) @staticmethod - def unscale_latents(x): + def unscale_latents(x: torch.Tensor) -> torch.Tensor: """[0, 1] -> raw latents""" return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude) - def decode(self, x): + def decode(self, x: torch.Tensor) -> torch.Tensor: 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): + def encode(self, x: torch.Tensor) -> torch.Tensor: return (self.taesd_encoder(x * 0.5 + 0.5) / self.vae_scale) + self.vae_shift