Update README.md

b73067f verified 2 days ago

8.12 kB

---
license: apache-2.0
library_name: pytorch
tags:
- vae
- autoencoder
- rgba
- layered
- psd
- hdr
---
# psd-layer-vae

**面向可编辑分层 PSD 的 RGBA 自编码器。** 每个图层各自编码成一个 latent（1/8 分辨率、32 通道，即 f8c32），解码时同一作品的各图层之间做跨图层注意力，输出直通（straight）RGBA。另有 HDR 通道，可以还原高于 SDR 白的线性光。模型从零训练，约 82M 参数（编码器 34.8M，解码器 47.2M）。

| 子目录 | 适用 | 特点 |
|---|---|---|
| `h3/` | **默认，用于编辑和逐层还原** | 像素保真最好，所有评测门槛都通过 |
| `g3/` | 生成时需要更锐利的解码 | rFID 低一半，像素保真略低（约 0.3 dB） |

h3 和 g3 的 latent 不同，编码器和解码器不能混用。

## 效果

**分层 PSD**：每个图层单独编码 / 解码，再按原文件的混合规则合成。

![layers](assets/layers.png)

**线稿图层**：细线、半透明笔触是逐层 RGBA 最难的部分。下图选的是"h3 与 Qwen 的差距最接近平均差距"的两个线稿层（不是挑最好的例子）；每组依次是图层本身、alpha、alpha 放大 64 倍。

| 72 个留出线稿层的平均值 | h3 | Qwen-Image-2.1 逐层往返 |
|---|---|---|
| alpha soft IoU | **.958** | .578 |
| 只看笔画（alpha < 0.02 置零后） | **.960** | .879 |
| 透明区域的平均残留 alpha | **0.02/255** | 0.80/255 |

笔画本身两者差距不大（.960 对 .879）；主要差别在透明区域：Qwen 在整个图层上留下一层很淡的 alpha 残影，笔画边缘还有光晕（放大 64 倍那一行可以看到），它的 soft IoU 损失有八成以上来自这些本该为 0 的像素。单层看不出来，但几十个图层叠在一起就会成为一层灰纱，存成 8 位 PSD 时这些残影大多会变成 1/255，图层的包围盒也就不再紧贴内容。h3 在透明区域基本输出 0。

![lines](assets/lines.png)

**密集文字**：纯色背景上的文字，以及透明图层上的文字（棋盘格表示透明）。

![text](assets/text.png)

**照片**：h3、g3 与 Qwen-Image-2.1 VAE 的重建对比（第二行为局部放大）。

![photo](assets/photo.png)

**HDR**：同一窗口在不同曝光档下的样子；Qwen-Image-2.1 是先做 PU21 编码再过它的 VAE。

![hdr](assets/hdr.png)

## 用法

```bash
pip install torch safetensors huggingface_hub
```

```python
import torch
import psd_layer_vae as P          # 本仓库的 psd_layer_vae/ 目录
model = P.load('Yidhar/psd-layer-vae', 'h3')       # 或本地目录；'g3' 为感知版
# 一个作品的 L 个图层，自下而上，同一画布窗口，H、W 为 8 的倍数（可用 P.pad_to_8 补齐）
rgba = torch.rand(5, 4, 256, 256, device='cuda')   # 直通 RGBA，sRGB 编码，[0, 1]
latent, box = P.encode(model, rgba)                # latent: (5, 32, 32, 32)；box: 各图层的包围盒
out = P.decode(model, latent, box)                 # out['rgba']: (5, 4, 256, 256)
image = P.composite(out['rgba'])                   # 普通混合模式合成（混合模式 / 剪贴蒙版 / 不透明度需自行处理）
# HDR：颜色为相对线性光（1.0 = SDR 白，可远大于 1），alpha 照常在 [0, 1]
hdr = torch.cat([linear_rgb, alpha], 1)             # (L, 4, H, W)
lin_rgb, alpha = P.reconstruct(model, hdr, linear=True)
```

- `decode` 需要编码时用的 `box`（各图层的 ROI），以及图层顺序 `z`（默认 0..L-1，自下而上）。解码器只在同一 8×8 格内、被多个图层覆盖的位置做跨图层注意力。
- 在 CUDA 上默认用 bf16 autocast，下面的评测也是这样测的；传 `bf16=False` 可改用 fp32。

## 多层解码：与逐层标准 VAE 对比

对照方式：用 Qwen-Image-2.1 VAE（RGBA，f16c64）把每个图层单独编码、解码；本模型把一个作品的所有图层一次联合解码。H200，bf16。

**速度**（解码，毫秒；Qwen 为逐层顺序调用 / 各层放进一个批次）

| 窗口 | 图层数 | 本模型 | Qwen 逐层 | Qwen 批量 | 显存峰值 本模型 / Qwen 批量 |
|---|---|---|---|---|---|
| 512² | 1 | 13 | 31 | 31 | 2.8 / 3.9 GB |
| 512² | 16 | 130 | 479 | 395 | 9.7 / 27.4 GB |
| 512² | 32 | 259 | 966 | 795 | 17.0 / 52.6 GB |
| 1024² | 16 | 685 | 1758 | 1653 | 42.4 / 103 GB |
| 1024² | 32 | 1368 | 3536 | 显存不足 | 82.5 / — |

解码快约 2.5–3.7 倍，显存约为 Qwen 批量解码的 1/3；时间随图层数线性增长，跨图层注意力本身的开销很小。编码速度在 512² 与 Qwen 相当，在 1024² 慢约 1.5 倍。速度优势主要来自解码器更轻（高分辨率阶段为深度可分离卷积），不是联合解码本身。

**质量**（20 个留出的 PSD 窗口，平均每个 15.7 层）

| 指标 | 本模型 | Qwen 逐层 |
|---|---|---|
| PSD 合成 PSNR | **37.00** | 34.97 |
| 逐层 PSNR（预乘 RGBA） | **53.96** | 41.74 |
| 线稿层 alpha soft IoU（按窗口平均） | **.959** | .627 |
| 其他图层 alpha soft IoU | **.991** | .939 |

**联合解码的代价：跨图层串扰。** 解码器在图层之间交换信息，所以删掉或修改一个图层后重新解码，其余图层也会轻微变化：删掉每个窗口里最大的图层，其余图层平均变化 0.12/255，个别像素最多 7–18/255，对其余图层的重建质量几乎无影响（−0.005 dB）。逐层解码的标准 VAE 没有这个问题；需要严格不变时，可以只对被编辑的图层重新解码、其余图层沿用之前的解码结果。

**说明**：上面的质量优势来自本模型针对 RGBA 图层的训练（alpha、细线、合成损失）和更大的 latent（f8c32），不能归功于跨图层注意力：训练早期的消融里，去掉跨图层注意力对单层指标没有可测的影响，只在 latent 带噪声时有小幅优势（噪声 0.3 倍时 +0.33 dB）。

## 评测

| 项目 | h3 | g3 | Qwen-Image-2.1 | FLUX.2 |
|---|---|---|---|---|
| PSD 窗口合成 PSNR | **37.00** | 36.62 | 34.97 | —（仅 RGB） |
| 不透明 PNG PSNR | **33.41** | 33.22 | 31.43 | 30.71 |
| 纯色文字页 / 透明文字页 PSNR | **34.30 / 42.06** | 33.89 / 41.63 | 30.66 / 29.82 | 29.82 / — |
| 细线图层 soft IoU / 透明文字 alpha soft IoU | **.958 / .945** | .956 / .941 | .578 / .745 | — |
| HDR PU21-PSNR（340 个窗口） | **35.51** | 35.09 | 35.43 | — |
| ImageNet val 5 万张 rFID | .622 | .291 | 1.14 | **.16** |
| ImageNet val 5 万张 PSNR / SSIM | **33.94 / .958** | 33.54 / .953 | 31.44 / .911 | 32.61 / .942 |
| LAION 500 张 PSNR / LPIPS | **36.00** / .0445 | 35.40 / .0485 | 33.09 / .0413 | 34.58 / **.0275** |
| 照片（FFHQ/DIV2K/LSDIR）PSNR / LPIPS | **35.00** / .0553 | 34.48 / .0583 | 32.19 / .0494 | 33.39 / **.0351** |
| 插画 PNG PSNR / LPIPS | **42.12** / .0232 | 41.68 / .0247 | 39.52 / .0222 | 38.98 / **.0217** |
| danbooru PSNR / LPIPS | **39.60** / .0203 | 39.01 / .0224 | 37.85 / .0166 | 37.69 / **.0158** |

图像集合各取 500 张，原分辨率中心 512 裁剪。Qwen-Image-2.1 用其 RGBA VAE，FLUX.2 只有 RGB，涉及 alpha 的项目不适用；HDR 时 Qwen-Image-2.1 先做 PU21 编码。所有数字用同一套评测工具和数据测得。

**局限**
- PSNR / SSIM 在所有图像集合上都领先 Qwen-Image-2.1 和 FLUX.2，但 LPIPS 仍落后两者；rFID 方面 g3 优于 Qwen-Image-2.1，仍落后 FLUX.2。
- h3 的 HDR 只比 Qwen 高 0.08 dB，单个种子；高光达 10 档以上的极亮场景仍有一部分窗口明显落后，例如下面这个电焊窗口（高光 14.5 档）：

  ![hdr failure](assets/hdr_failure.png)
- g3 的细节看起来真实，但不等于原图；需要逐层精确还原的编辑请用 h3。
- 训练数据含第三方图像。

## 许可证

权重（`h3/`、`g3/`）与推理代码均以 [Apache License 2.0](LICENSE) 发布。

## 文件

- `h3/`、`g3/`：`model.safetensors`（float32）和 `config.json`
- `psd_layer_vae/`：推理代码（模型定义与编解码接口）
- `LICENSE`：Apache License 2.0 全文
- `SHA256SUMS`：本仓库全部文件的校验和

	---
	license: apache-2.0
	library_name: pytorch
	tags:
	- vae
	- autoencoder
	- rgba
	- layered
	- psd
	- hdr
	---

	# psd-layer-vae

	面向可编辑分层 PSD 的 RGBA 自编码器。每个图层各自编码成一个 latent（1/8 分辨率、32 通道，即 f8c32），解码时同一作品的各图层之间做跨图层注意力，输出直通（straight）RGBA。另有 HDR 通道，可以还原高于 SDR 白的线性光。模型从零训练，约 82M 参数（编码器 34.8M，解码器 47.2M）。

	\| 子目录 \| 适用 \| 特点 \|
	\|---\|---\|---\|
	\| `h3/` \| 默认，用于编辑和逐层还原 \| 像素保真最好，所有评测门槛都通过 \|
	\| `g3/` \| 生成时需要更锐利的解码 \| rFID 低一半，像素保真略低（约 0.3 dB） \|

	h3 和 g3 的 latent 不同，编码器和解码器不能混用。

	## 效果

	分层 PSD：每个图层单独编码 / 解码，再按原文件的混合规则合成。

	![layers](assets/layers.png)

	线稿图层：细线、半透明笔触是逐层 RGBA 最难的部分。下图选的是"h3 与 Qwen 的差距最接近平均差距"的两个线稿层（不是挑最好的例子）；每组依次是图层本身、alpha、alpha 放大 64 倍。

	\| 72 个留出线稿层的平均值 \| h3 \| Qwen-Image-2.1 逐层往返 \|
	\|---\|---\|---\|
	\| alpha soft IoU \| .958 \| .578 \|
	\| 只看笔画（alpha < 0.02 置零后） \| .960 \| .879 \|
	\| 透明区域的平均残留 alpha \| 0.02/255 \| 0.80/255 \|

	笔画本身两者差距不大（.960 对 .879）；主要差别在透明区域：Qwen 在整个图层上留下一层很淡的 alpha 残影，笔画边缘还有光晕（放大 64 倍那一行可以看到），它的 soft IoU 损失有八成以上来自这些本该为 0 的像素。单层看不出来，但几十个图层叠在一起就会成为一层灰纱，存成 8 位 PSD 时这些残影大多会变成 1/255，图层的包围盒也就不再紧贴内容。h3 在透明区域基本输出 0。

	![lines](assets/lines.png)

	密集文字：纯色背景上的文字，以及透明图层上的文字（棋盘格表示透明）。

	![text](assets/text.png)

	照片：h3、g3 与 Qwen-Image-2.1 VAE 的重建对比（第二行为局部放大）。

	![photo](assets/photo.png)

	HDR：同一窗口在不同曝光档下的样子；Qwen-Image-2.1 是先做 PU21 编码再过它的 VAE。

	![hdr](assets/hdr.png)

	## 用法

	```bash
	pip install torch safetensors huggingface_hub
	```

	```python
	import torch
	import psd_layer_vae as P # 本仓库的 psd_layer_vae/ 目录

	model = P.load('Yidhar/psd-layer-vae', 'h3') # 或本地目录；'g3' 为感知版

	# 一个作品的 L 个图层，自下而上，同一画布窗口，H、W 为 8 的倍数（可用 P.pad_to_8 补齐）
	rgba = torch.rand(5, 4, 256, 256, device='cuda') # 直通 RGBA，sRGB 编码，[0, 1]

	latent, box = P.encode(model, rgba) # latent: (5, 32, 32, 32)；box: 各图层的包围盒
	out = P.decode(model, latent, box) # out['rgba']: (5, 4, 256, 256)
	image = P.composite(out['rgba']) # 普通混合模式合成（混合模式 / 剪贴蒙版 / 不透明度需自行处理）

	# HDR：颜色为相对线性光（1.0 = SDR 白，可远大于 1），alpha 照常在 [0, 1]
	hdr = torch.cat([linear_rgb, alpha], 1) # (L, 4, H, W)
	lin_rgb, alpha = P.reconstruct(model, hdr, linear=True)
	```

	- `decode` 需要编码时用的 `box`（各图层的 ROI），以及图层顺序 `z`（默认 0..L-1，自下而上）。解码器只在同一 8×8 格内、被多个图层覆盖的位置做跨图层注意力。
	- 在 CUDA 上默认用 bf16 autocast，下面的评测也是这样测的；传 `bf16=False` 可改用 fp32。

	## 多层解码：与逐层标准 VAE 对比

	对照方式：用 Qwen-Image-2.1 VAE（RGBA，f16c64）把每个图层单独编码、解码；本模型把一个作品的所有图层一次联合解码。H200，bf16。

	速度（解码，毫秒；Qwen 为逐层顺序调用 / 各层放进一个批次）

	\| 窗口 \| 图层数 \| 本模型 \| Qwen 逐层 \| Qwen 批量 \| 显存峰值本模型 / Qwen 批量 \|
	\|---\|---\|---\|---\|---\|---\|
	\| 512² \| 1 \| 13 \| 31 \| 31 \| 2.8 / 3.9 GB \|
	\| 512² \| 16 \| 130 \| 479 \| 395 \| 9.7 / 27.4 GB \|
	\| 512² \| 32 \| 259 \| 966 \| 795 \| 17.0 / 52.6 GB \|
	\| 1024² \| 16 \| 685 \| 1758 \| 1653 \| 42.4 / 103 GB \|
	\| 1024² \| 32 \| 1368 \| 3536 \| 显存不足 \| 82.5 / — \|

	解码快约 2.5–3.7 倍，显存约为 Qwen 批量解码的 1/3；时间随图层数线性增长，跨图层注意力本身的开销很小。编码速度在 512² 与 Qwen 相当，在 1024² 慢约 1.5 倍。速度优势主要来自解码器更轻（高分辨率阶段为深度可分离卷积），不是联合解码本身。

	质量（20 个留出的 PSD 窗口，平均每个 15.7 层）

	\| 指标 \| 本模型 \| Qwen 逐层 \|
	\|---\|---\|---\|
	\| PSD 合成 PSNR \| 37.00 \| 34.97 \|
	\| 逐层 PSNR（预乘 RGBA） \| 53.96 \| 41.74 \|
	\| 线稿层 alpha soft IoU（按窗口平均） \| .959 \| .627 \|
	\| 其他图层 alpha soft IoU \| .991 \| .939 \|

	联合解码的代价：跨图层串扰。解码器在图层之间交换信息，所以删掉或修改一个图层后重新解码，其余图层也会轻微变化：删掉每个窗口里最大的图层，其余图层平均变化 0.12/255，个别像素最多 7–18/255，对其余图层的重建质量几乎无影响（−0.005 dB）。逐层解码的标准 VAE 没有这个问题；需要严格不变时，可以只对被编辑的图层重新解码、其余图层沿用之前的解码结果。

	说明：上面的质量优势来自本模型针对 RGBA 图层的训练（alpha、细线、合成损失）和更大的 latent（f8c32），不能归功于跨图层注意力：训练早期的消融里，去掉跨图层注意力对单层指标没有可测的影响，只在 latent 带噪声时有小幅优势（噪声 0.3 倍时 +0.33 dB）。

	## 评测

	\| 项目 \| h3 \| g3 \| Qwen-Image-2.1 \| FLUX.2 \|
	\|---\|---\|---\|---\|---\|
	\| PSD 窗口合成 PSNR \| 37.00 \| 36.62 \| 34.97 \| —（仅 RGB） \|
	\| 不透明 PNG PSNR \| 33.41 \| 33.22 \| 31.43 \| 30.71 \|
	\| 纯色文字页 / 透明文字页 PSNR \| 34.30 / 42.06 \| 33.89 / 41.63 \| 30.66 / 29.82 \| 29.82 / — \|
	\| 细线图层 soft IoU / 透明文字 alpha soft IoU \| .958 / .945 \| .956 / .941 \| .578 / .745 \| — \|
	\| HDR PU21-PSNR（340 个窗口） \| 35.51 \| 35.09 \| 35.43 \| — \|
	\| ImageNet val 5 万张 rFID \| .622 \| .291 \| 1.14 \| .16 \|
	\| ImageNet val 5 万张 PSNR / SSIM \| 33.94 / .958 \| 33.54 / .953 \| 31.44 / .911 \| 32.61 / .942 \|
	\| LAION 500 张 PSNR / LPIPS \| 36.00 / .0445 \| 35.40 / .0485 \| 33.09 / .0413 \| 34.58 / .0275 \|
	\| 照片（FFHQ/DIV2K/LSDIR）PSNR / LPIPS \| 35.00 / .0553 \| 34.48 / .0583 \| 32.19 / .0494 \| 33.39 / .0351 \|
	\| 插画 PNG PSNR / LPIPS \| 42.12 / .0232 \| 41.68 / .0247 \| 39.52 / .0222 \| 38.98 / .0217 \|
	\| danbooru PSNR / LPIPS \| 39.60 / .0203 \| 39.01 / .0224 \| 37.85 / .0166 \| 37.69 / .0158 \|

	图像集合各取 500 张，原分辨率中心 512 裁剪。Qwen-Image-2.1 用其 RGBA VAE，FLUX.2 只有 RGB，涉及 alpha 的项目不适用；HDR 时 Qwen-Image-2.1 先做 PU21 编码。所有数字用同一套评测工具和数据测得。

	局限
	- PSNR / SSIM 在所有图像集合上都领先 Qwen-Image-2.1 和 FLUX.2，但 LPIPS 仍落后两者；rFID 方面 g3 优于 Qwen-Image-2.1，仍落后 FLUX.2。
	- h3 的 HDR 只比 Qwen 高 0.08 dB，单个种子；高光达 10 档以上的极亮场景仍有一部分窗口明显落后，例如下面这个电焊窗口（高光 14.5 档）：

	![hdr failure](assets/hdr_failure.png)
	- g3 的细节看起来真实，但不等于原图；需要逐层精确还原的编辑请用 h3。
	- 训练数据含第三方图像。

	## 许可证

	权重（`h3/`、`g3/`）与推理代码均以 [Apache License 2.0](LICENSE) 发布。

	## 文件

	- `h3/`、`g3/`：`model.safetensors`（float32）和 `config.json`
	- `psd_layer_vae/`：推理代码（模型定义与编解码接口）
	- `LICENSE`：Apache License 2.0 全文
	- `SHA256SUMS`：本仓库全部文件的校验和