Repository: huggingface/diffusers

Issue:
TAESD-encoded latents are too dark
### Describe the bug

AutoencodeTiny (TAESD) _decoder_ seems to work fine. _encoding_ on the other hand is producing poor results, and an encode-decode round-trip turns out poorly:

input:
![benz](https://github.com/huggingface/diffusers/assets/83819/df80128c-b0c9-448f-a246-940402f2f99c)

output:
![sad benz](https://cdn.discordapp.com/attachments/862888301081133099/1142296650426105896/sad-benz.webp)


### Reproduction

see https://gist.github.com/keturn/b0a10a3b388e1e49cdf38567b76eb30c

```py
import diffusers, torch
from PIL.Image import Image, open as image_open

device = torch.device("cuda:0")

with torch.inference_mode():
    taesd = diffusers.AutoencoderTiny.from_pretrained("madebyollin/taesd", torch_dtype=torch.float16).to(device=device)
    vaesd = diffusers.AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae", variant="fp16", torch_dtype=torch.float16).to(device=device)

from diffusers.utils.testing_utils import load_image

image = load_image(
    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"
)

from diffusers.image_processor import VaeImageProcessor
vae_processor = VaeImageProcessor()
image_tensor: torch.FloatTensor = vae_processor.preprocess(image).to(dtype=torch.float16, device=device)

print(f"image tensor range: {image_tensor.min()} < {image_tensor.mean()} < {image_tensor.max()})")


with torch.inference_mode():
    taesd_latents = taesd.encode(image_tensor).latents
    print(f"taesd-encoded latent range: {taesd_latents.min()} < {taesd_latents.mean()} (σ={taesd_latents.std()}) < {taesd_latents.max()})")

    vaesd_latents = vaesd.encode(image_tensor).latent_dist.sample()
    print(f"vaesd-encoded latent range: {vaesd_latents.min()} < {vaesd_latents.mean()} (σ={vaesd_latents.std()}) < {vaesd_latents.max()})")


with torch.inference_mode():
    redecoded_tensor = taesd.decode(taesd_latents).sample


redecoded_image = vae_processor.postprocess(redecoded_tensor)
display(image, redecoded_image[0])


from diffusers.commands import env
env.EnvironmentCommand().run()
```

### System Info

  - `diffusers` version: 0.20.0
  - Platform: Linux-5.15.0-79-generic-x86_64-with-glibc2.35
  - Python version: 3.11.4
  - PyTorch version (GPU?): 2.0.1+cu118 (True)
  - Huggingface_hub version: 0.16.4
  - Transformers version: 4.31.0
  - Accelerate version: 0.21.0
  - xFormers version: 0.0.21
  - Using GPU in script?: yes
  - Using distributed or parallel set-up in script?: no


### Who can help?

_No response_


Discussion:
Oops, It's probably [this](https://github.com/huggingface/diffusers/pull/4384#discussion_r1279430003) (@sayakpaul 😅)
yeah I was side-eyeing that suspiciously
> Oops, It's probably https://github.com/huggingface/diffusers/pull/4384#discussion_r1279430003

I am not sure what's meant by that comment. 

Maybe the encode-decode roundtrip is best done by referring to @madebyollin's original notebook here:
https://github.com/madebyollin/taesd/blob/main/examples/Previewing_During_Image_Generation.ipynb
Is there a round-trip in that notebook? I don't see any encoding.
@sayakpaul
Per the [README](https://github.com/madebyollin/taesd/blob/main/README.md?plain=1#L44), TAESD's raw model assumes a [0, 1] scaling convention for input / output images.

I think either @keturn's sample code needs to use a special preprocessor for `AutoencoderTiny`, or (to match the `AutoencoderKL` behavior), `AutoencoderTiny.encode` could add rescaling of the inputs, inverse of the [rescaling](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoder_tiny.py#L317) in `AutoencoderTiny.decode`.

Right now the encoder is getting an image in [-1, 1], encoding the values that are in [0, 1] (above 50% brightness), and clamping the rest of the values (everything below 50% brightness) to black - which is why the decoder decodes a darkened image.

(My fault for not following up on the review comment - sorry!)

@keturn Since a few people have now been (understandably 😅) tripped up by the way TAESD bakes in the SD-VAE input / output scale-shift transforms, I've added an [example "Encoding / Decoding" notebook](https://github.com/madebyollin/taesd/blob/main/examples/Encoding_and_Decoding.ipynb) to hopefully clear things up. I also added (hopefully) [clearer language](https://github.com/madebyollin/taesd/commit/8de0ac9b3aeac0df0783f21f31284ee2f4dff8f6?short_path=b335630#diff-b335630551682c19a781afebcf4d07bf978fb1f8ac04c6bf87428ed5106870f5) to the README.

![enc_dec_example](https://github.com/huggingface/diffusers/assets/8733615/24cbe7fb-3449-4d32-9ffa-512d3b56d981)

It seems that it's avaliable for `VaeImageProcessor` to normalize/denormalize the image in preprocess/postprocess.

Add `do_normalize=False` to `VaeImageProcessor` and `do_denormalize` to `vae_processor.postprocess` could create normal image.

But it seems that the workflow would be a little bit complicated if it needs switching between `vae`/`tae`. Maybe we can move `do_normalize` from processor to the `preprocess` method for flexibility?

## Test code
```python
import diffusers, torch
from PIL.Image import Image, open as image_open

device = torch.device("cuda:0")

with torch.inference_mode():
    taesd = diffusers.AutoencoderTiny.from_pretrained("madebyollin/taesd", torch_dtype=torch.float16).to(device=device)
    vaesd = diffusers.AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae", variant="fp16", torch_dtype=torch.float16).to(device=device)

from diffusers.utils.testing_utils import load_image

image = load_image(
    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"
)

from diffusers.image_processor import VaeImageProcessor
vae_processor = VaeImageProcessor(do_normalize=False)
image_tensor: torch.FloatTensor = vae_processor.preprocess(image).to(dtype=torch.float16, device=device)

print(f"image tensor range: {image_tensor.min()} < {image_tensor.mean()} < {image_tensor.max()})")


with torch.inference_mode():
    taesd_latents = taesd.encode(image_tensor).latents
    print(f"taesd-encoded latent range: {taesd_latents.min()} < {taesd_latents.mean()} (σ={taesd_latents.std()}) < {taesd_latents.max()})")

    vaesd_latents = vaesd.encode(image_tensor).latent_dist.sample()
    print(f"vaesd-encoded latent range: {vaesd_latents.min()} < {vaesd_latents.mean()} (σ={vaesd_latents.std()}) < {vaesd_latents.max()})")


with torch.inference_mode():
    redecoded_tensor = taesd.decode(taesd_latents).sample


redecoded_image = vae_processor.postprocess(redecoded_tensor, do_denormalize=[True])
display(image, redecoded_image[0])
```

## Outputs
![benz](https://github.com/huggingface/diffusers/assets/41363108/44505aa5-f338-4463-b7d9-e63e8c15b745)

`encode` and `decode` should be inverses of one another, such that `x == encode(decode(x))` and `y == decode(encode(y))`, give or take bit due to the lossy nature of the encoding. So the domain of the input of `encode` should be consistent with the range of the output of `decode` — it would be quite unexpected for one of them to be [0..1) and the other to be [-1, 1].

> But it seems that the workflow would be a little bit complicated if it needs switching between `vae`/`tae`.

Yes. Given that TAESD was _explicitly designed_ as a drop-in replacement for the Stable Diffusion VAE, and the diffusers library implements them both, it would be very much appreciated if the library offered an interface that's consistent to both.
Even though it was developed as a drop-in replacement, I think its main usefulness lies in speedy decoding. For now, almost all (if not all) encoders we have in the library expect the value range to be in [-1, 1]. So, we won't likely be changing it. But happy to review any PRs in this regard. 

Resolving patch:
diff --git a/src/diffusers/models/autoencoder_tiny.py b/src/diffusers/models/autoencoder_tiny.py
--- a/src/diffusers/models/autoencoder_tiny.py
+++ b/src/diffusers/models/autoencoder_tiny.py
@@ -312,9 +312,6 @@ def decode(self, x: torch.FloatTensor, return_dict: bool = True) -> Union[Decode
             output = torch.cat(output)
         else:
             output = self._tiled_decode(x) if self.use_tiling else self.decoder(x)
-        # Refer to the following discussion to know why this is needed.
-        # https://github.com/huggingface/diffusers/pull/4384#discussion_r1279401854
-        output = output.mul_(2).sub_(1)
 
         if not return_dict:
             return (output,)
@@ -333,8 +330,15 @@ def forward(
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
         """
         enc = self.encode(sample).latents
+
+        # scale latents to be in [0, 1], then quantize latents to a byte tensor,
+        # as if we were storing the latents in an RGBA uint8 image.
         scaled_enc = self.scale_latents(enc).mul_(255).round_().byte()
-        unscaled_enc = self.unscale_latents(scaled_enc)
+
+        # unquantize latents back into [0, 1], then unscale latents back to their original range,
+        # as if we were loading the latents from an RGBA uint8 image.
+        unscaled_enc = self.unscale_latents(scaled_enc / 255.0)
+
         dec = self.decode(unscaled_enc)
 
         if not return_dict:
diff --git a/src/diffusers/models/vae.py b/src/diffusers/models/vae.py
--- a/src/diffusers/models/vae.py
+++ b/src/diffusers/models/vae.py
@@ -732,7 +732,8 @@ def custom_forward(*inputs):
                 x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.layers), x)
 
         else:
-            x = self.layers(x)
+            # scale image from [-1, 1] to [0, 1] to match TAESD convention
+            x = self.layers(x.add(1).div(2))
 
         return x
 
@@ -790,4 +791,5 @@ def custom_forward(*inputs):
         else:
             x = self.layers(x)
 
-        return x
+        # scale image from [0, 1] to [-1, 1] to match diffusers convention
+        return x.mul(2).sub(1)


Test patch:
diff --git a/tests/models/test_models_vae.py b/tests/models/test_models_vae.py
--- a/tests/models/test_models_vae.py
+++ b/tests/models/test_models_vae.py
@@ -312,10 +312,32 @@ def test_stable_diffusion(self):
         assert sample.shape == image.shape
 
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
-        expected_output_slice = torch.tensor([0.9858, 0.9262, 0.8629, 1.0974, -0.091, -0.2485, 0.0936, 0.0604])
+        expected_output_slice = torch.tensor([0.0093, 0.6385, -0.1274, 0.1631, -0.1762, 0.5232, -0.3108, -0.0382])
 
         assert torch_all_close(output_slice, expected_output_slice, atol=3e-3)
 
+    @parameterized.expand([(True,), (False,)])
+    def test_tae_roundtrip(self, enable_tiling):
+        # load the autoencoder
+        model = self.get_sd_vae_model()
+        if enable_tiling:
+            model.enable_tiling()
+
+        # make a black image with a white square in the middle,
+        # which is large enough to split across multiple tiles
+        image = -torch.ones(1, 3, 1024, 1024, device=torch_device)
+        image[..., 256:768, 256:768] = 1.0
+
+        # round-trip the image through the autoencoder
+        with torch.no_grad():
+            sample = model(image).sample
+
+        # the autoencoder reconstruction should match original image, sorta
+        def downscale(x):
+            return torch.nn.functional.avg_pool2d(x, model.spatial_scale_factor)
+
+        assert torch_all_close(downscale(sample), downscale(image), atol=0.125)
+
 
 @slow
 class AutoencoderKLIntegrationTests(unittest.TestCase):