Approximate Simple for SDXL

modules/sd_samplers_common.py

@@ -33,8 +33,6 @@ def single_sample_to_image(sample, approximation=None):
             x_sample = x_sample[[2,1,0],:,:] # BGR to RGB
     elif approximation == 2:
         x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
-        if shared.sd_model_type == "sdxl":
-            x_sample = x_sample[[2,1,0],:,:] # BGR to RGB
     elif approximation == 3:
         # x_sample = sample * 1.5
         # x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()

modules/sd_vae_approx.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from torch import nn
-from modules import devices, paths
+from modules import devices, paths, shared
 
 
 sd_vae_approx_model = None
@@ -49,14 +49,21 @@ def model():
 
 def cheap_approximation(sample):
     # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/2
-    coefs = torch.tensor([
-        [0.298, 0.207, 0.208],
-        [0.187, 0.286, 0.173],
-        [-0.158, 0.189, 0.264],
-        [-0.184, -0.271, -0.473],
-    ]).to(sample.device)
+    if shared.sd_model_type == "sdxl":
+        weight = torch.tensor([
+            [0.4543,-0.2868, 0.1566,-0.4748],
+            [0.5008, 0.0952, 0.2155,-0.3268],
+            [0.5294, 0.1625,-0.0624,-0.3793]
+        ]).reshape(3, 4, 1, 1).to(sample.device)
+        bias = torch.tensor([0.1375, 0.0144, -0.0675]).to(sample.device)
+    else:
+        weight = torch.tensor([
+            [0.298, 0.187,-0.158,-0.184],
+            [0.207, 0.286, 0.189,-0.271],
+            [0.208, 0.173, 0.264,-0.473],
+        ]).reshape(3, 4, 1, 1).to(sample.device)
+        bias = None
     try:
-        x_sample = torch.einsum("lxy,lr -> rxy", sample, coefs)
-        return x_sample
+        return nn.functional.conv2d(sample, weight, bias)
     except Exception:
         return sample