From 7a13a3f4ba86dc44fcf7d9944b179018744862f5 Mon Sep 17 00:00:00 2001
From: Sakura-Luna <53183413+Sakura-Luna@users.noreply.github.com>
Date: Wed, 17 May 2023 17:39:07 +0800
Subject: [PATCH] TAESD fix

---
 modules/sd_samplers_common.py | 9 +++++----
 modules/sd_vae_taesd.py       | 2 +-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index ceda6a35..d99c933d 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -35,13 +35,14 @@ def single_sample_to_image(sample, approximation=None):
     elif approximation == 1:
         x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
     elif approximation == 3:
-        x_sample = sd_vae_taesd.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
-        x_sample = sd_vae_taesd.TAESD.unscale_latents(x_sample)  # returns value in [-2, 2]
-        x_sample = x_sample * 0.5
+        x_sample = sample * 1.5
+        x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
     else:
         x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0]
 
-    x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
+    if approximation != 3:
+        x_sample = (x_sample + 1.0) / 2.0
+    x_sample = torch.clamp(x_sample, min=0.0, max=1.0)
     x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
     x_sample = x_sample.astype(np.uint8)
 
diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py
index d23812ef..5e8496e8 100644
--- a/modules/sd_vae_taesd.py
+++ b/modules/sd_vae_taesd.py
@@ -45,7 +45,7 @@ def decoder():
 
 
 class TAESD(nn.Module):
-    latent_magnitude = 2
+    latent_magnitude = 3
     latent_shift = 0.5
 
     def __init__(self, decoder_path="taesd_decoder.pth"):