基于扩散模型的图像压缩：创建基于Stable Diffusion的有损压缩编解码器

Stable Diffusion是最近在图像生成领域大火的模型，在对他研究的时候我发现它可以作为非常强大的有损图像压缩编解码器。在介绍之前，我们先看看它与高压缩率下的JPG和WebP相比的而结果，所有都是512x512像素的分辨率为单位：

上面这图的细节如下：

这是JPG压缩结果大小为5.66kB

这是WebP ，大小为6.74kB：

最后是Stable Diffusion 大小4.97kB

还有一些其他示例，请参考我们最后提供的源代码地址

这些示例非常明显，与JPG和WebP相比，用Stable Diffusion压缩这些图像在较小的文件尺寸下会产生极好的图像质量。这么看来这是一种非常有前途的有损图像压缩的选择，但是Stable Diffusion还会带来一些问题，我们在后面介绍。

Stable Diffusion的潜在空间

Stable Diffusion使用三个常用的神经网络进行组合：

1、VAE

2、U-net

3、文本编码器

Variational Auto Encoder（VAE）编码和解码器将图像转换为一些潜在空间表示。

VAE如何将图像编码到这个潜在空间中呢，它在训练过程中自行学习，随着模型的进一步训练，模型的不同版本的潜在空间表示可能会有所不同，但是Stable Diffusion v1.4的表示形式看起来像这样（重新映射为4通道颜色图像）：

在重新缩放和解释潜在的颜色值（带有alpha通道）后，图像的主要特征仍然是肉眼可见，VAE将较高的分辨率特征编码到这些像素值中。一个vae的编码/解码的过程如下：

左：512x512@24bpp的真实图像，中：64x64@128bpp的潜在空间，右：潜在空间解码生成的512x512@24bpp图像

通过潜在空间还原的过程并不是无损的。可以看到解码后，我们的羊驼头上蓝色带子包含的名字就不太清晰了。这是因为Stable Diffusion v1.4的VAE通常不太擅长表示小文本以及人脸。

而我们知道Stable Diffusion的主要算法从短文本描述中生成新图像，它从潜在空间表示中的随机噪声开始，然后通过使用训练过的U-Net对潜在空间图像进行迭代去噪，简单地说，U-Net输出它认为在噪声中“看到”的东西的预测。当生成图像时，这个迭代的去噪步骤由第三个ML模型(文本编码器)进行指导，该模型向U-Net提供关于它应该在噪声中尝试看到什么内容的信息。所以对于我们这个压缩的方案来说，不需要文本编码器，也就是希望没有文本来进行引导。所以我们这里只创建一个空字符串的一次性编码，这样相当于是告诉U-Net在图像重建期间进行无引导去噪。

压缩方法

为了使用Stable Diffusion作为图像压缩编解码器，我们还需要研究了如何有效地压缩由VAE生成的潜在表示。在这个实验中，进行下采样或者应用现有的有损图像压缩方法都会大大降低重构图像的质量。但是我发现VAE的解码似乎对潜在的量化非常稳健。

对潜在空间进行从浮点到8位无符号整数的量化，通过scaling，clamping，然后remapping ，只有非常小的可见重构误差:

为了量化由VAE生成的空间，我首先将它们按1 / 0.18215进行缩放，这个数字是在Stable Diffusion源代码中出现的。用这个数字可以很好地映射到[- 1,1]的范围。

通过将潜在项量化到8位，图像表示的数据大小现在是646448位= 16 kB(真实图像为51251238位= 768 kB)

将潜在对象量化到8位以下并没有产生好的结果，但通过重整色调和抖动来进一步量化它们却出乎意料地好。使用256个48位向量和Floyd-Steinberg抖动的潜在色调创建了一个色调表示。使用256的色调用单个8位索引表示每个潜在向量，使数据大小达到64648+2564*8位= 5 kB。

色调的抖动会引入了噪声使解码结果失真。但由于扩散模型是基于潜波的去噪，所以可以使用U-Net来去除抖动带来的噪声。经过4次迭代，重构结果在视觉上非常接近于未量化的版本:

左：抖动后，中:4个去噪步骤，右:真实图像

数据大小的极大减少了(压缩系数为155倍)，结果非常好，但是我们也看到心脏符号上新增了光泽阴影，这是在压缩之前不存在。这种压缩方案会引入新的内容从而影响图像表示而不是图像质量，这也说明扩散模型是在创造而不仅仅是还原。

结果评估

为了评估此压缩编解码器，所以没有使用任何公开的图像，这样以确保图像不会包含在扩散模型的训练集中（这样的图像可能会获得不公平的压缩优势，因为它们的一部分数据可能已经在训练时出现了）。为了使比较尽可能公平，使用了Python图像库的JPG和WebP压缩的最高质量设置，并且还使用MozJPEG库对压缩JPG数据进行了无损压缩。

虽然看起来Stable Diffusion结果比JPG和WebP压缩图像好得多，但就标准测量指标（如PSNR或SSIM）而言，它们并没有明显好（但也没更差）。这可能是因为Stable Diffusion创造了一些不太引人注意的新的内容，这些内容不会会影响图像评分，但是对人类而言感觉却变好了，例如重构内容可能会受到新生成内容的影响，即使它看起来非常清晰（原图没有，但是还原图被生成出来的一些细节欺骗了）。例如这张旧金山测试图像中的一个细节:

左:JPG压缩图，中:真实图片 -右:Stable Diffusion压缩

可以非常清晰的看到，建筑物的形状的改变。

一些限制

有些特征目前并没有被保留得很好;特别是小文本和(参考 Stable Diffusion v1.4的Limitations and Biases节)。下面是一个简单的例子

这是真实图像：

vae潜在空间 32bit

8bit的解码和去噪

可以非常清晰的看到当前的Stable Diffusion1.4模型在潜在空间中不能很好地保存小的文本，所以这个压缩方案也有这些限制。但是据说从1.5以后似乎在人脸处理方面做得更好，期待stability.ai尽快发布更好的版本。

这种方法与基于ml的图像恢复有何不同?

深度学习非常成功地用于恢复降级的图像和视频或将其提高到更高的分辨率。这种方法与恢复图像压缩的图像有何不同？

这里重要的区别是是要理解Stable Diffusion不是用来恢复由于压缩而降级的图像。它是将有损的压缩应用于对图像的内部理解，然后试图“修复”有损压缩对内部表示造成的损害(这与修复降级图像不同)。

例如，传统的图像恢复保留相机纹理(定性)以及任何其他定性退化，而AI恢复一个高度压缩的jpg将无法恢复相机纹理，因为任何关于它的信息已经从数据中丢失了。

一下这个比喻可能会更恰当一点：

假设我们有一个拥有过目不忘的高技能艺术家。我们给他们看一张图像，然后让他们重新创造它，艺术家可以根据自己的记忆创造一个近乎完美的副本。这位艺术家的摄影记忆就是是Stable Diffusion中VAE的潜在空间。

而传统的图像恢复，我们会给这位艺术家看一张被图像压缩后严重降级的图像，然后让他们根据记忆重建这张图像。

然对于我们的量化操作，可以认为是我们向艺术家展示原始的、未压缩的图像并要求他们尽可能地记住它。然后我们对艺术家进行脑部手术，通过对其进行一些有损压缩来缩小他们记忆中的数据（量化或者上面提到的压缩），这种压缩去除了看起来不重要的细微差别，并用相同的变化替换了记忆图像中非常非常相似的概念和方面的变化。

手术后，我们还是要求艺术家根据自己的记忆，创造一个完美的图像重建。他们仍然会记住图像的所有重要方面，例如，从内容到相机纹理的定性属性，以及他们看到的每座建筑的位置和大致外观，虽然可能每个点的确切位置不再相同，并且在他们记忆中还有一些奇怪的缺失，但是这些小的缺失对于整理来说影响并不大。

最后，艺术家通过这些手术后的记忆（压缩后）来从新绘画，细节变少了也就是被压缩了，但是整体来说还是不变的。

所以在这种情况下，艺术家被要求根据他们的记忆和经验画出他们认为应该是什么样子的东西，但由于这里的压缩应用于他们对图像的记忆表示，图像的视觉质量和概念内容都已降低，所以在某些部分可能是因为其他记忆产生的所以艺术家会在图画中发明创造一些新的东西。这也就是我们刚才说到的创造。

总结

作为有损图像压缩方案，Stable Diffusion似乎非常有前途。我们可以通过更多的实验来完整这个方案并且还可以专门设计和训练用于图像压缩的VAE，或者其他更好的方法都是我们在未来研究的方向，如果你感兴趣，请看看下本文的源代码并加以改进。

https://avoid.overfit.cn/post/3b8ce036928f431d93061e5a307f8067

作者：Matthias Bühlmann