update readme

.gitignore

@@ -5,6 +5,7 @@ _sc.py
 *.ckpt
 *.bin
 
+checkpoints
 .idea
 .idea/workspace.xml
 .DS_Store

README.md

@@ -4,9 +4,17 @@
 <a href="https://arxiv.org/pdf/2302.13848.pdf"><img src="https://img.shields.io/badge/arXiv-2302.13848-b31b1b.svg" height=22.5></a>
 <a href="https://huggingface.co/spaces/ELITE-library/ELITE"><img src="https://img.shields.io/static/v1?label=HuggingFace&message=gradio demo&color=darkgreen" height=22.5></a>
 
-## Getting Started
+![method](assets/results.png)
+
+
+## Method Details
+
+![method](assets/method.png)
+
+Given an image indicates the target concept (usually an object), we propose a learning-based encoder ELITE to encode the visual concept into the textual embeddings, which can be further flexibly composed into new scenes.  It consists of two modules: (a) a global mapping network is first trained to encode a concept image into multiple textual word embeddings, where one primary word (w0) for well-editable concept and other auxiliary words (w1···N) to exclude irrelevant disturbances. (b) A local mapping network is further trained, which projects the foreground object into textual feature space to provide local details.
+
 
-----
+## Getting Started
 
 ### Environment Setup
 
@@ -22,13 +30,21 @@ pip install -r requirements.txt
 
 We provide the pretrained checkpoints in [Google Drive](https://drive.google.com/drive/folders/1VkiVZzA_i9gbfuzvHaLH2VYh7kOTzE0x?usp=sharing). One can download them and save to the directory `checkpoints`.
 
-### Setting up Diffusers
+### Setting up HuggingFace
+
+Our code is built on the [diffusers](https://github.com/huggingface/diffusers/) version of Stable Diffusion, you need to accept the [model license](https://huggingface.co/CompVis/stable-diffusion-v1-4) before downloading or using the weights. In our experiments, we use model version v1-4.
+
+You have to be a registered user in Hugging Face Hub, and you'll also need to use an access token for the code to work. For more information on access tokens, please refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).
 
-Our code is built on the [diffusers](https://github.com/huggingface/diffusers/), and you can follow the guideline [here](https://github.com/huggingface/diffusers/tree/main/examples/textual_inversion#cat-toy-example) to set it.
+Run the following command to authenticate your token
+```shell
+huggingface-cli login
+```
+If you have already cloned the repo, then you won't need to go through these steps.
 
 ### Customized Generation
 
-We provide the testing dataset in [test_datasets](./test_datasets), which contains both images and object masks. For testing, you can run, 
+We provide some testing images in [test_datasets](./test_datasets), which contains both images and object masks. For testing, you can run, 
 ```
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
 export DATA_DIR='./test_datasets/'
@@ -42,16 +58,14 @@ CUDA_VISIBLE_DEVICES=0 python inference_local.py \
   --global_mapper_path="./checkpoints/global_mapper.pt" \
   --local_mapper_path="./checkpoints/local_mapper.pt"
 ```
-or you can use the shell script:
+or you can use the shell script,
 ```
 bash inference_local.sh
 ```
-If you want to test your customized dataset, you should align the image to ensure the object is at the center of image, and also provide the corresponding object mask. The object mask can be obtained by [image-matting-app](https://huggingface.co/spaces/SankarSrin/image-matting-app), or other image matting methods. 
+If you want to test your customized dataset, you should align the image to ensure the object is at the center of image, and also provide the corresponding object mask. The object mask can be obtained by [image-matting-app](https://huggingface.co/spaces/SankarSrin/image-matting-app), or other image matting methods.
 
 ## Training
 
-----
-
 ### Preparing Dataset
 
 We use the **test** dataset of Open-Images V6 to train our ELITE. You can prepare the dataset as follows:
@@ -87,7 +101,7 @@ datasets
 
 ### Training Global Mapping Network
 
-To train the global mapping network, run the following command:
+To train the global mapping network, you can run,
 
 ```Shell
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
@@ -106,14 +120,14 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_p
   --output_dir="./elite_experiments/global_mapping" \
   --save_steps 200
 ```
-or you can use the shell script:
+or you can use the shell script,
 ```shell
 bash train_global.sh
 ```
 
 ### Training Local Mapping Network
 
-After the global mapping is trained, you can train the local mapping by running the following command:
+After the global mapping network is trained, you can train the local mapping network by running,
 
 ```Shell
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
@@ -133,7 +147,7 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_p
   --output_dir="./elite_experiments/local_mapping" \
   --save_steps 200
 ```
-or you can use the shell script:
+or you can use the shell script,
 ```shell
 bash train_local.sh
 ```
@@ -152,4 +166,4 @@ bash train_local.sh
 
 ## Acknowledgements
 
-This code is built on [diffusers](https://github.com/huggingface/diffusers/). We thank the authors for sharing the codes.
+This code is built on [diffusers](https://github.com/huggingface/diffusers/) version of [Stable Diffusion](https://github.com/CompVis/stable-diffusion). We thank the authors for sharing the codes.

inference_global.sh

@@ -1,7 +1,7 @@
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
 export DATA_DIR='./test_datasets/'
 
-CUDA_VISIBLE_DEVICES=7 python inference_global.py \
+CUDA_VISIBLE_DEVICES=0 python inference_global.py \
   --pretrained_model_name_or_path=$MODEL_NAME \
   --test_data_dir=$DATA_DIR \
   --output_dir="./outputs/global_mapping"  \
@@ -9,5 +9,5 @@ CUDA_VISIBLE_DEVICES=7 python inference_global.py \
   --token_index="0" \
   --template="a photo of a S" \
   --global_mapper_path="./checkpoints/global_mapper.pt" \
-  --seed 42
+  --seed=42
 

inference_local.sh

@@ -1,6 +1,6 @@
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
 export DATA_DIR='./test_datasets/'
-CUDA_VISIBLE_DEVICES=7 python inference_local.py \
+CUDA_VISIBLE_DEVICES=0 python inference_local.py \
   --pretrained_model_name_or_path=$MODEL_NAME \
   --test_data_dir=$DATA_DIR \
   --output_dir="./outputs/local_mapping"  \
@@ -9,5 +9,5 @@ CUDA_VISIBLE_DEVICES=7 python inference_local.py \
   --llambda="0.8" \
   --global_mapper_path="./checkpoints/global_mapper.pt" \
   --local_mapper_path="./checkpoints/local_mapper.pt" \
-  --seed 42
+  --seed=42
 

train_global.sh

@@ -11,5 +11,5 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_p
   --learning_rate=1e-06 --scale_lr \
   --lr_scheduler="constant" \
   --lr_warmup_steps=0 \
-  --output_dir="./elite_experiments/global_mapping_new" \
+  --output_dir="./elite_experiments/global_mapping" \
   --save_steps 200