composition-stable-diffusion's Introduction

Image Composition via Stable Diffusion

We achieve image composition via Stable Diffusion Model. Application includes Virtual Clothes / Funiture Try-on.

Demo 1: Virtual Clothes Try-on

Demo 2: Virtual Furniture Try-on

Installation

Requirements

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch 
pip install -r requirements.txt

export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION="python"

Initialize an 🤗Accelerate environment with
```
accelerate config default
```
Run the following command to authenticate your token
```
huggingface-cli login
```

0. Prepare Images

Please provide at least one images in .jpg format and instance prompt. For example, images in ./data/sofa

1. Set Environment

export MODEL_NAME="runwayml/stable-diffusion-inpainting"
export INSTANCE_DIR="data/sofa"
export TEST_DIR="data/sofa_test"
export MODEL_DIR="logs/sofa"
export OUT_DIR="out/sofa"
export INSTANCE_PROMPT="sofa"

2. Preprocess Images

Please provide at least one images in .jpg format and instance prompt. The preprocess.py script will generate captions and instance masks.

python preprocess.py --instance_data_dir $INSTANCE_DIR \
                     --instance_prompt $INSTANCE_PROMPT

3. Finetune

We then embed the instance images and prompt into stable diffusion model.

accelerate launch --num_processes 1 train.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$MODEL_DIR \
  --instance_prompt=$INSTANCE_PROMPT \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1000

4. Image Composition

Finally, you can provide new images to achieve image composition.

python inference.py --image_path $TEST_DIR \
                    --model_path $MODEL_DIR \
                    --out_path $OUT_DIR \
                    --instance_prompt $INSTANCE_PROMPT