Upload folder using huggingface_hub

.gitattributes

@@ -2,7 +2,6 @@
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
@@ -22,10 +21,8 @@
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +30,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+diffusion_pytorch_model.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,83 @@
+---
+license: mit
+tags:
+- stable-diffusion
+- stable-diffusion-diffusers
+inference: false
+---
+# Improved Autoencoders
+
+## Utilizing
+These weights are intended to be used with the [🧨 diffusers library](https://github.com/huggingface/diffusers). If you are looking for the model to use with the original [CompVis Stable Diffusion codebase](https://github.com/CompVis/stable-diffusion), [come here](https://huggingface.co/stabilityai/sd-vae-ft-mse-original).
+
+#### How to use with 🧨 diffusers
+You can integrate this fine-tuned VAE decoder to your existing `diffusers` workflows, by including a `vae` argument to the `StableDiffusionPipeline`
+```py
+from diffusers.models import AutoencoderKL
+from diffusers import StableDiffusionPipeline
+
+model = "CompVis/stable-diffusion-v1-4"
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+pipe = StableDiffusionPipeline.from_pretrained(model, vae=vae)
+```
+
+## Decoder Finetuning
+We publish two kl-f8 autoencoder versions, finetuned from the original [kl-f8 autoencoder](https://github.com/CompVis/latent-diffusion#pretrained-autoencoding-models) on a 1:1 ratio of [LAION-Aesthetics](https://laion.ai/blog/laion-aesthetics/) and LAION-Humans, an unreleased subset containing only SFW images of humans. The intent was to fine-tune on the Stable Diffusion training set (the autoencoder was originally trained on OpenImages) but also enrich the dataset with images of humans to improve the reconstruction of faces.
+The first, _ft-EMA_, was resumed from the original checkpoint, trained for 313198 steps and uses EMA weights. It uses the same loss configuration as the original checkpoint (L1 + LPIPS).
+The second, _ft-MSE_, was resumed from _ft-EMA_ and uses EMA weights and was trained for another 280k steps using a different loss, with more emphasis 
+on MSE reconstruction (MSE + 0.1 * LPIPS). It produces somewhat ``smoother'' outputs. The batch size for both versions was 192 (16 A100s, batch size 12 per GPU).
+To keep compatibility with existing models, only the decoder part was finetuned; the checkpoints can be used as a drop-in replacement for the existing autoencoder.
+
+_Original kl-f8 VAE vs f8-ft-EMA vs f8-ft-MSE_
+
+## Evaluation 
+### COCO 2017 (256x256, val, 5000 images)
+| Model    | train steps | rFID | PSNR         | SSIM          | PSIM          | Link                                                                              | Comments                                                                                        
+|----------|---------|------|--------------|---------------|---------------|-----------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------|
+|          |         |      |              |               |               |                                                                                   |                                                                                                 |
+| original | 246803        | 4.99 | 23.4 +/- 3.8 | 0.69 +/- 0.14 | 1.01 +/- 0.28 | https://ommer-lab.com/files/latent-diffusion/kl-f8.zip                            | as used in SD                                                                                   |
+| ft-EMA   | 560001        | 4.42 | 23.8 +/- 3.9 | 0.69 +/- 0.13 | 0.96 +/- 0.27 | https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.ckpt | slightly better overall, with EMA                                                               |
+| ft-MSE   | 840001        | 4.70 | 24.5 +/- 3.7 | 0.71 +/- 0.13 | 0.92 +/- 0.27 | https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt | resumed with EMA from ft-EMA, emphasis on MSE (rec. loss = MSE + 0.1 * LPIPS), smoother outputs |
+
+
+### LAION-Aesthetics 5+ (256x256, subset, 10000 images)
+| Model    | train steps | rFID | PSNR         | SSIM          | PSIM          | Link                                                                              | Comments                                                                                        
+|----------|-----------|------|--------------|---------------|---------------|-----------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------|
+|          |           |      |              |               |               |                                                                                   |                                                                                                 |
+| original | 246803         | 2.61 | 26.0 +/- 4.4 | 0.81 +/- 0.12 | 0.75 +/- 0.36 | https://ommer-lab.com/files/latent-diffusion/kl-f8.zip                            | as used in SD                                                                                   |
+| ft-EMA   | 560001          | 1.77 | 26.7 +/- 4.8 | 0.82 +/- 0.12 | 0.67 +/- 0.34 | https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.ckpt | slightly better overall, with EMA                                                               |
+| ft-MSE   | 840001          | 1.88 | 27.3 +/- 4.7 | 0.83 +/- 0.11 | 0.65 +/- 0.34 | https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt | resumed with EMA from ft-EMA, emphasis on MSE (rec. loss = MSE + 0.1 * LPIPS), smoother outputs |
+
+
+### Visual
+_Visualization of reconstructions on  256x256 images from the COCO2017 validation dataset._ 
+
+<p align="center">
+  <br>
+  <b>
+256x256: ft-EMA (left), ft-MSE (middle), original (right)</b>
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00025_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00011_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00037_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00043_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00053_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00029_merged.png />
+</p>

config.json

@@ -0,0 +1,29 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.4.2",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b4889b6b1d4ce7ae320a02dedaeff1780ad77d415ea0d744b476155c6377ddc
+size 334707217

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815
+size 334643276

handler.py

@@ -0,0 +1,95 @@
+from typing import Dict, List, Any
+import torch
+from base64 import b64decode
+from huggingface_hub import model_info
+from diffusers import AutoencoderKL
+from diffusers.image_processor import VaeImageProcessor
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        kwargs = {"torch_dtype": torch.float16}
+        model_data = model_info(path)
+        has_model_index = any(
+            file.rfilename == "model_index.json" for file in model_data.siblings
+        )
+
+        if has_model_index:
+            kwargs["subfolder"] = "vae"
+
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = kwargs["torch_dtype"]
+        self.vae = AutoencoderKL.from_pretrained(path, **kwargs).to(self.device).eval()
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    @torch.no_grad()
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        """
+        tensor = data["inputs"]
+        tensor = b64decode(tensor.encode("utf-8"))
+        parameters = data.get("parameters", {})
+        if "shape" not in parameters:
+            raise ValueError("Expected `shape` in parameters.")
+        if "dtype" not in parameters:
+            raise ValueError("Expected `dtype` in parameters.")
+
+        DTYPE_MAP = {
+            "float16": torch.float16,
+            "float32": torch.float32,
+            "bfloat16": torch.bfloat16,
+        }
+
+        shape = parameters.get("shape")
+        dtype = DTYPE_MAP.get(parameters.get("dtype"))
+        tensor = torch.frombuffer(bytearray(tensor), dtype=dtype).reshape(shape)
+
+        needs_upcasting = (
+            self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+        )
+        if needs_upcasting:
+            self.vae = self.vae.to(torch.float32)
+            tensor = tensor.to(self.device, torch.float32)
+        else:
+            tensor = tensor.to(self.device, self.dtype)
+
+        # unscale/denormalize the latents
+        # denormalize with the mean and std if available and not None
+        has_latents_mean = (
+            hasattr(self.vae.config, "latents_mean")
+            and self.vae.config.latents_mean is not None
+        )
+        has_latents_std = (
+            hasattr(self.vae.config, "latents_std")
+            and self.vae.config.latents_std is not None
+        )
+        if has_latents_mean and has_latents_std:
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, 4, 1, 1)
+                .to(tensor.device, tensor.dtype)
+            )
+            latents_std = (
+                torch.tensor(self.vae.config.latents_std)
+                .view(1, 4, 1, 1)
+                .to(tensor.device, tensor.dtype)
+            )
+            tensor = (
+                tensor * latents_std / self.vae.config.scaling_factor + latents_mean
+            )
+        else:
+            tensor = tensor / self.vae.config.scaling_factor
+
+        with torch.no_grad():
+            image = self.vae.decode(tensor, return_dict=False)[0]
+
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float16)
+
+        image = self.image_processor.postprocess(image, output_type="pil")
+
+        return image[0]

requirements.txt

@@ -0,0 +1,3 @@
+huggingface_hub
+diffusers
+