Refactor previews into one command line argument.

Clean up a few things.
This commit is contained in:
comfyanonymous
2023-06-06 01:26:52 -04:00
parent 081134f5c8
commit a3a713b6c5
8 changed files with 107 additions and 100 deletions

View File

@@ -50,9 +50,9 @@ class TAESD(nn.Module):
self.encoder = Encoder()
self.decoder = Decoder()
if encoder_path is not None:
self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu"))
self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu", weights_only=True))
if decoder_path is not None:
self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu"))
self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu", weights_only=True))
@staticmethod
def scale_latents(x):