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：每个图层单独编码 / 解码，再按原文件的混合规则合成。
线稿图层：细线、半透明笔触是逐层 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。
密集文字：纯色背景上的文字，以及透明图层上的文字（棋盘格表示透明）。
照片：h3、g3 与 Qwen-Image-2.1 VAE 的重建对比（第二行为局部放大）。
HDR：同一窗口在不同曝光档下的样子；Qwen-Image-2.1 是先做 PU21 编码再过它的 VAE。
用法
pip install torch safetensors huggingface_hub
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 档）：
g3 的细节看起来真实，但不等于原图；需要逐层精确还原的编辑请用 h3。
训练数据含第三方图像。
许可证
权重（h3/、g3/）与推理代码均以 Apache License 2.0 发布。
文件
h3/、g3/：model.safetensors（float32）和config.jsonpsd_layer_vae/：推理代码（模型定义与编解码接口），不含训练代码；输出与训练代码逐位一致LICENSE：Apache License 2.0 全文SHA256SUMS：本仓库全部文件的校验和