xiaozhiob/diffusion-tryon-trainer

Diffusion-Tryon-Trainer

This is a unoffical training code for OOTDiffusion

This repository contains the training code for the OOTDiffusion project. We extend our gratitude to the contributions of OOTDiffusion and have built upon this foundation by utilizing Huggingface's Diffusors library to implement training on the VTON dataset for virtual try-on. Our project aims to enhance the accuracy and realism of virtual try-ons through cutting-edge diffusion model technology, providing users with a more authentic try-on experience.

Some results

Install requirements

Python env

conda env create -f environment.yaml
conda activate groot

Data Preparation

VITON-HD

1. Download VITON-HD dataset
2. Download pre-warped cloth image/mask from Google Driver or Baidu Cloud and put it under your VITON-HD dataset
3. Download cloth captions from train Google Driver test Google Driver

After these, the folder structure should look like this (the unpaired-cloth* only included in test directory):

├── VITON-HD
|   ├── test_pairs.txt
|   ├── train_pairs.txt
│   ├── [train | test]
|   |   ├── image
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── cloth
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── cloth-mask
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── cloth-warp
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── cloth-warp-mask
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── unpaired-cloth-warp
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
│   │   ├── unpaired-cloth-warp-mask
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
|   |   ├── cloth_caption
│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]

How to run training

bash train_ootd.sh

Inference over CP dataset

Weight preparation

Train the weights or download a pretrained weight from Huggingface The weights need to be put under checkpoints dir.

Run inference

sh inference_test_dataset.sh

Acknowledgements

Our unet code is directly from OOTDiffusion. We also thank DCI-VTON-Virtual-Try-On, our dataset module depends on it.

This code is only for study and research.

Contributors

nftblackmagic 💻

Yu (Brian) Yao 💻

Stevada 💻

Dingkang Wang 💻

xiaoweilu 💻