Parts2Whole

[Arxiv 2024] 从部分到整体：可控人像生成的统一参考框架

🔥 更新

[2024-06-21] 发布训练代码。参见训练。 <br/> [2024-05-26] 数据集已在此处发布。参见数据集。 <br/> [2024-05-06] 🔥🔥🔥 代码已发布。尽情享受人体部件组合吧！

🏠 <a href="https://huanngzh.github.io/Parts2Whole/" target="_blank">项目主页</a> | <a href="https://arxiv.org/abs/2404.15267" target="_blank">论文</a> | <a href="https://huggingface.co/huanngzh/Parts2Whole" target="_blank">模型</a>

img:teaser

摘要：我们提出了Parts2Whole，这是一个新颖的框架，旨在从多个参考图像生成定制人像，包括姿势图像和人类外观的各个方面。首先，我们开发了一个语义感知外观编码器，以保留不同人体部位的细节，该编码器根据文本标签将每个图像处理为一系列多尺度特征图，而不是一个图像标记，从而保留了图像维度。其次，我们的框架通过在扩散过程中在参考和目标特征之间操作的共享自注意力机制支持多图像条件生成。我们通过结合参考人体图像的掩码信息来增强原始注意力机制，允许精确选择任何部位。

🔨 方法概述

img:pipeline

⚒️ 安装

克隆我们的仓库，并安装requirements.txt中的包。我们在80G A800 GPU上使用CUDA 11.8和PyTorch 2.0.1测试了我们的模型。但在较小的GPU上进行推理也是可能的。

conda create -n parts2whole
conda activate parts2whole
pip install -r requirements.txt

在<a href="https://huggingface.co/huanngzh/Parts2Whole" target="_blank">此处</a>下载检查点到pretrained_weights/parts2whole目录。我们还提供了一个简单的下载脚本，使用：

python download_weights.py

🎨 推理

查看inference.py。根据需要修改检查点路径和输入，然后运行命令：

python inference.py

您可能需要在inference.py脚本中修改以下代码：

### 定义配置 ###
device = "cuda"
torch_dtype = torch.float16
seed = 42
model_dir = "pretrained_weights/parts2whole"  # 本地机器上的检查点路径
use_decoupled_cross_attn = True
decoupled_cross_attn_path = "pretrained_weights/parts2whole/decoupled_attn.pth" # 包含在model_dir中

### 定义输入数据 ###
height, width = 768, 512
prompt = "这个人穿着短袖衬衫。" # 输入提示
input_dict = {
    "appearance": {
        "face": "testset/face_man1.jpg",
        "whole body clothes": "testset/clothes_man1.jpg",
    },
    "mask": {
        "face": "testset/face_man1_mask.jpg",
        "whole body clothes": "testset/clothes_man1_mask.jpg",
    },
    "structure": {"densepose": "testset/densepose_man1.jpg"},
}

⭐️⭐️⭐️ 值得注意的是，input_dict应包含键appearance、mask和structure。前两个表示指定多个参考图像中部分的外观，structure表示姿势，如densepose。

⭐️⭐️⭐️ 这三个部分中的键也有解释。appearance和mask中的键应该相同。选项包括"upper body clothes"、"lower body clothes"、"whole body clothes"、"hair or headwear"、"face"、"shoes"。structure的键应为"densepose"。（openpose模型尚未发布。）

🔨🔨🔨 为了方便获取每个参考图像的掩码，我们还提供了相应的工具，并在工具中解释了如何使用它们。首先，您可以使用Real-ESRGAN提高参考图像的分辨率，并使用segformer获取人体各个部位的掩码。

📊 数据集

data_sample

我们的数据集已在此处发布。我们在download_dataset.py中提供了下载和解压脚本，请使用以下命令：

python download_dataset.py

它将在data/DeepFashion-MultiModal-Parts2Whole文件夹中准备数据集，这样您就可以运行我们的配置来训练模型或运行我们的数据集文件parts2whole/data/ref_trg.py来检查我们的数据集。

确保数据集按如下方式组织：

DeepFashion-MultiModal-Parts2Whole
# 结构信号
|-- densepose
|-- openpose
# 外观条件
|-- face
|-- hair_headwear
|-- lower_body_clothes
|-- upper_body_clothes
|-- whole_body_clothes
|-- shoes
# 目标图像
|-- images
# 说明文件
|-- train.jsonl
`-- test.jsonl

这个人像数据集包含约41,500对参考-目标对。该数据集中的每对都包括多个参考图像，包括姿势图、人类外观的各个方面（如头发、脸部、衣服、鞋子），以及一张featuring相同个体（ID）的目标图像，还有文本说明。有关数据集的详细信息，请参阅我们的数据集仓库。

我们的数据集是从DeepFashion-Multimodal数据集后处理而来的。

🏋️ 训练

如果在单个设备上训练我们的parts2whole，请使用以下命令：

python train.py --config configs/train-sd15.yaml

如果在 DDP 环境中训练（假设有 8 个设备），请运行以下命令：

accelerate launch \
--mixed_precision=fp16 \
--num_processes=$((8*$WORLD_SIZE)) \  # 8 是设备数量
--num_machines=$WORLD_SIZE \
--multi_gpu \
--machine_rank=$RANK \
train.py --config configs/train-sd15.yaml

在我们的配置文件中，每个设备的批量大小设置为 8（这是针对 80G 内存设备的推荐值）。如果你在内存较小的设备上训练，需要减小这个值。

😊 评估

进行评估之前，请先安装以下额外的包：

pip install git+https://github.com/openai/CLIP.git # 用于 clip
pip install dreamsim # 用于 dreamsim
pip install lpips # 用于 lpips

我们在 scripts/evals 文件夹中提供了易于使用的评估脚本。这些脚本接收统一格式的数据，以两个图像列表作为输入。根据需要修改加载图像的代码。更多细节请查看我们的脚本。

🔨 工具

Real-ESRGAN

要使用 Real-ESRGAN 来恢复图像，请先下载 RealESRGAN_x4plus.pth 到 ./pretrained_weights/Real-ESRGAN 目录。然后运行以下命令：

python -m scripts.real_esrgan -n RealESRGAN_x4plus -i /path/to/dir -o /path/to/dir --face_enhance

SegFormer

要使用 segformer 对人物图像进行分割并获取帽子、头发、面部、衣服等部分，请运行以下命令：

python scripts.segformer_b2_clothes.py --image-path /path/to/image --output-dir /path/to/dir

标签：0: "背景", 1: "帽子", 2: "头发", 3: "太阳镜", 4: "上衣", 5: "裙子", 6: "裤子", 7: "连衣裙", 8: "腰带", 9: "左鞋", 10: "右鞋", 11: "脸", 12: "左腿", 13: "右腿", 14: "左臂", 15: "右臂", 16: "包", 17: "围巾"

😭 局限性

目前，训练数据的泛化性一般，女性数量相对较多，因此模型的泛化能力需要改进，比如风格化等方面。我们正在努力提高模型的鲁棒性和能力，同时也期待并欢迎社区的贡献和拉取请求。

🤝 致谢

我们感谢以下项目的开源：

diffusers magic-animate Moore-AnimateAnyone DeepFashion-MultiModal Real-ESRGAN

📎 引用

如果您觉得这个仓库有用，请考虑引用：

@misc{huang2024parts2whole,
  title={From Parts to Whole: A Unified Reference Framework for Controllable Human Image Generation},
  author={Huang, Zehuan and Fan, Hongxing and Wang, Lipeng and Sheng, Lu},
  journal={arXiv preprint arXiv:2404.15267},
  year={2024}
}