Foreign-Mirrors/Fooocus
1
0
Fork 0
mirror of https://github.com/lllyasviel/Fooocus.git synced 2026-08-16 13:13:16 +02:00
Code Issues Packages Projects Releases Wiki Activity

maintain

Browse Source
This commit is contained in:
lllyasviel
2023-11-23 13:46:50 -08:00
parent bd4d40203c
commit 3bc9ac88fd
13 changed files with 189 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/headless/fcbh/taesd/taesd.py
+14 -5
View File
@@ -46,15 +46,16 @@ class TAESD(nn.Module):
latent_magnitude = 3
latent_shift = 0.5
def __init__(self, encoder_path="taesd_encoder.pth", decoder_path="taesd_decoder.pth"):
def __init__(self, encoder_path=None, decoder_path=None):
"""Initialize pretrained TAESD on the given device from the given checkpoints."""
super().__init__()
self.encoder = Encoder()
self.decoder = Decoder()
self.taesd_encoder = Encoder()
self.taesd_decoder = Decoder()
self.vae_scale = torch.nn.Parameter(torch.tensor(1.0))
if encoder_path is not None:
self.encoder.load_state_dict(fcbh.utils.load_torch_file(encoder_path, safe_load=True))
self.taesd_encoder.load_state_dict(fcbh.utils.load_torch_file(encoder_path, safe_load=True))
if decoder_path is not None:
self.decoder.load_state_dict(fcbh.utils.load_torch_file(decoder_path, safe_load=True))
self.taesd_decoder.load_state_dict(fcbh.utils.load_torch_file(decoder_path, safe_load=True))
@staticmethod
def scale_latents(x):
@@ -65,3 +66,11 @@ class TAESD(nn.Module):
def unscale_latents(x):
"""[0, 1] -> raw latents"""
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 = 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