psd-layer-vae

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

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

效果

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

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

笔画本身两者差距不大（.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 为逐层顺序调用 / 各层放进一个批次）

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

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

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

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

评测

图像集合各取 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.json
• psd_layer_vae/：推理代码（模型定义与编解码接口）
• LICENSE：Apache License 2.0 全文
• SHA256SUMS：本仓库全部文件的校验和