# CatVTON

## 论文

`CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models`

* https://arxiv.org/pdf/2407.15886

## 模型结构

该模型基于`stable diffusion`结构，移除了`ReferenceNet`或其他图像（衣物）编码器，同时移除`cross-attention`及文本编码器。

![alt text](readme_imgs/model_arch.png)

## 算法原理

该算法基于`stable diffusion`，去除了多余的网络结构，直接将控制条件作为`Unet`的输入。

![alt text](readme_imgs/alg.png)

## 环境配置

### Docker（方法一）
    
    docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu22.04-dtk24.04.2-py3.10

    docker run --shm-size 50g --network=host --name=catvton --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash

    pip install -r requirements.txt

### Dockerfile（方法二）

    docker build -t <IMAGE_NAME>:<TAG> .

    docker run --shm-size 50g --network=host --name=catvton --privileged --device=/dev/kfd --device=/dev/dri --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v 项目地址(绝对路径):/home/ -v /opt/hyhal:/opt/hyhal:ro -it <your IMAGE ID> bash

    pip install -r requirements.txt

### Anaconda (方法三)

1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装：
https://developer.hpccube.com/tool/

    DTK驱动：dtk24.04.2
    python：python3.10
    torch: 2.1.0
    torchvision: 0.16.0

Tips：以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应

2、其它非特殊库参照requirements.txt安装

    pip install -r requirements.txt

## 数据集

完整数据集可通过官网链接下载，本项目提供测试数据，位于`datasets`中。

|DressCode|VITON-HD|

### 数据处理

VITON-HD本身包含已处理好的agnostic-mask，仅对DressCode处理即可。

```bash
HIP_VISIBLE_DEVICES=0 python preprocess_agnostic_mask.py \
--data_root_path <your_path_to_DressCode> 
```

```bash
├── VITON-HD
|   ├── test_pairs_unpaired.txt
│   ├── test
|   |   ├── image
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── cloth
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── agnostic-mask
│   │   │   ├── [000006_00.png | 000008_00.png | ...]
...
```

```bash
├── DressCode
|   ├── test_pairs_paired.txt
|   ├── test_pairs_unpaired.txt
│   ├── [dresses | lower_body | upper_body]
|   |   ├── test_pairs_paired.txt
|   |   ├── test_pairs_unpaired.txt
│   │   ├── images
│   │   │   ├── [013563_0.jpg | 013563_1.jpg | 013564_0.jpg | 013564_1.jpg | ...]
│   │   ├── agnostic_masks
│   │   │   ├── [013563_0.png| 013564_0.png | ...]
...
```

## 训练

无

## 推理

```bash
export HF_ENDPOINT=https://hf-mirror.com
```

### 命令行


```bash
HIP_VISIBLE_DEVICES=0 python inference.py \
--dataset [ dresscode | vitonhd ] \
--data_root_path <path/to/datasets> \
--output_dir temp_output \
--dataloader_num_workers 8 \
--batch_size 2 \
--repaint \
--eval_pair \
--base_model_path <path/to/stable-diffusion-inpainting> \
--resume_path <path/to/CatVTON/Models>
```

注意：请使用fp16或fp32进行推理。

### webui

```bash
HIP_VISIBLE_DEVICES=0 python app.py \
--output_dir="resource/demo/output" \
--mixed_precision="fp16" \
--base_model_path <path/to/stable-diffusion-inpainting> \
--resume_path <path/to/CatVTON/Models>
```

### 指标计算

```bash
HIP_VISIBLE_DEVICES=0 python eval.py \
--gt_folder <your_path_to_gt_image_folder> \
--pred_folder <your_path_to_predicted_image_folder> \
--paired \
--batch_size=2 \
--num_workers=2 
```

## result

|数据源|输入1|输入2|输出|
|:---:|:---:|:---:|:---:|
|VITON-HD|![alt text](readme_imgs/input1_vh.jpg)|<img src="readme_imgs/input2_vh.jpg" style="zoom:50%;">|![alt text](readme_imgs/output_vh.jpg)|
|DressCode|<img src="readme_imgs/input1_dr.jpg" style="zoom:50%;">|<img src="readme_imgs/input2_dr.jpg" style="zoom:50%;">|![alt text](readme_imgs/output_dr.jpg)|



### 精度

无

## 应用场景

### 算法类别

`AIGC`

### 热点应用行业

`电商,绘画,广媒`

## 预训练权重

stable-diffusion-inpainting: [huggingface](https://hf-mirror.com/booksforcharlie/stable-diffusion-inpainting/tree/main) 

CatVTON: [huggingface](https://hf-mirror.com/zhengchong/CatVTON)

注意：可将下载的权重文件放入`Models`文件夹（需自行创建）。

## 源码仓库及问题反馈

* https://developer.sourcefind.cn/codes/aibear/catvton_pytorch

## 参考资料

* https://github.com/Zheng-Chong/CatVTON

