readme and new requirements

.gitignore

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

2_gpu.json

@@ -1,11 +0,0 @@
-{
-  "compute_environment": "LOCAL_MACHINE",
-  "distributed_type": "MULTI_GPU",
-  "fp16": false,
-  "machine_rank": 0,
-  "main_process_ip": null,
-  "main_process_port": null,
-  "main_training_function": "main",
-  "num_machines": 1,
-  "num_processes": 2
-}

3_gpu.json

@@ -1,11 +0,0 @@
-{
-  "compute_environment": "LOCAL_MACHINE",
-  "distributed_type": "MULTI_GPU",
-  "fp16": false,
-  "machine_rank": 0,
-  "main_process_ip": null,
-  "main_process_port": null,
-  "main_training_function": "main",
-  "num_machines": 1,
-  "num_processes": 3
-}

README.md

@@ -1,3 +1,155 @@
-# ELITE
+# ELITE: Encoding Visual Concepts into Textual Embeddings for Customized Text-to-Image Generation
 
-The detailed README is coming soom.
+
+<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
+
+----
+
+### Environment Setup
+
+```shell
+git clone https://github.com/csyxwei/ELITE.git
+cd ELITE
+conda create -n elite python=3.9
+conda activate elite
+pip install -r requirements.txt
+```
+
+### Pretrained Models
+
+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
+
+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.
+
+### Customized Generation
+
+We provide the testing dataset 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/'
+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"  \
+  --suffix="object" \
+  --template="a photo of a S" \
+  --llambda="0.8" \
+  --global_mapper_path="./checkpoints/global_mapper.pt" \
+  --local_mapper_path="./checkpoints/local_mapper.pt"
+```
+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. 
+
+## Training
+
+----
+
+### Preparing Dataset
+
+We use the **test** dataset of Open-Images V6 to train our ELITE. You can prepare the dataset as follows:
+
+- Download Open-Images test dataset from [CVDF's site](https://github.com/cvdfoundation/open-images-dataset#download-images-with-bounding-boxes-annotations) and unzip it to the directory `datasets/Open_Images/images/test`.
+- Download attribute names file `oidv6-attributes-description.csv` of Open-Images test dataset from [Open-Images official site](https://storage.googleapis.com/openimages/web/download_v7.html#download-manually) and save it to the directory `datasets/Open_Images/annotations/`.
+- Download bbox annotations file `test-annotations-bbox.csv` of Open-Images test dataset from [Open-Images official site](https://storage.googleapis.com/openimages/web/download_v7.html#download-manually) and save it to the directory `datasets/Open_Images/annotations/`.
+- Download segmentation annotations of Open-Images test dataset from [Open-Images official site](https://storage.googleapis.com/openimages/web/download_v7.html#download-manually) and unzip them to the directory `datasets/Open_Images/segs/test`. And put the `test-annotations-object-segmentation.csv` into `datasets/Open_Images/annotations/`. 
+- Obtain the mask bbox by running the following command:
+    ```shell
+    python data_scripts/cal_bbox_by_seg.py
+    ```
+
+The final data structure is like this:
+
+```
+datasets
+├── Open_Images
+│  ├── annotations
+│  │  ├── oidv6-class-descriptions.csv
+│  │  ├── test-annotations-object-segmentation.csv
+│  │  ├── test-annotations-bbox.csv
+│  ├── images
+│  │  ├── test
+│  │  │  ├── xxx.jpg
+│  │  │  ├── ...
+│  ├── segs
+│  │  ├── test
+│  │  │  ├── xxx.png
+│  │  │  ├── ...
+│  │  ├── test_bbox_dict.npy
+```
+
+### Training Global Mapping Network
+
+To train the global mapping network, run the following command:
+
+```Shell
+export MODEL_NAME="CompVis/stable-diffusion-v1-4"
+export DATA_DIR='./datasets/Open_Images/'
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_process_port 25656 train_global.py \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATA_DIR \
+  --placeholder_token="S" \
+  --resolution=512 \
+  --train_batch_size=4 \
+  --gradient_accumulation_steps=4 \
+  --max_train_steps=200000 \
+  --learning_rate=1e-06 --scale_lr \
+  --lr_scheduler="constant" \
+  --lr_warmup_steps=0 \
+  --output_dir="./elite_experiments/global_mapping" \
+  --save_steps 200
+```
+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:
+
+```Shell
+export MODEL_NAME="CompVis/stable-diffusion-v1-4"
+export DATA_DIR='/home/weiyuxiang/datasets/Open_Images/'
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_process_port 25657 train_local.py \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATA_DIR \
+  --placeholder_token="S" \
+  --resolution=512 \
+  --train_batch_size=2 \
+  --gradient_accumulation_steps=4 \
+  --max_train_steps=200000 \
+  --learning_rate=1e-5 --scale_lr \
+  --lr_scheduler="constant" \
+  --lr_warmup_steps=0 \
+  --global_mapper_path "./elite_experiments/global_mapping/mapper_070000.pt" \
+  --output_dir="./elite_experiments/local_mapping" \
+  --save_steps 200
+```
+or you can use the shell script:
+```shell
+bash train_local.sh
+```
+
+
+## Citation
+
+```
+@article{wei2023elite,
+  title={ELITE: Encoding Visual Concepts into Textual Embeddings for Customized Text-to-Image Generation},
+  author={Wei, Yuxiang and Zhang, Yabo and Ji, Zhilong and Bai, Jinfeng and Zhang, Lei and Zuo, Wangmeng},
+  journal={arXiv preprint arXiv:2302.13848},
+  year={2023}
+}
+```
+
+## Acknowledgements
+
+This code is built on [diffusers](https://github.com/huggingface/diffusers/). We thank the authors for sharing the codes.

data_scripts/cal_bbox_by_seg.py

@@ -0,0 +1,59 @@
+import cv2
+from os.path import join
+import os
+import numpy as np
+from tqdm import tqdm
+
+dir = './datasets/Open_Images/'
+mode = 'test'
+
+image_dir = join(dir, 'images', mode)
+seg_dir = join(dir, 'segs', mode)
+
+files = os.listdir(seg_dir)
+
+data_dict = {}
+
+for file in tqdm(files):
+    seg_path = join(seg_dir, file)
+    image_path = join(image_dir, file.split('_')[0] + '.jpg')
+    seg = cv2.imread(seg_path)
+    image = cv2.imread(image_path)
+    seg = cv2.resize(seg, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+
+    seg = seg[:, :, 0]
+
+    # obtain contours point set： contours
+    contours = cv2.findContours(seg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    contours = contours[0] if len(contours) == 2 else contours[1]
+
+    if len(contours) > 1:
+        cntr = np.vstack(contours)
+    elif len(contours) == 1:
+        cntr = contours[0]
+    else:
+        continue
+
+    if len(cntr) < 2:
+        continue
+
+    hs, he = np.min(cntr[:, :, 1]), np.max(cntr[:, :, 1])
+    ws, we = np.min(cntr[:, :, 0]), np.max(cntr[:, :, 0])
+
+    h, w = seg.shape
+
+    if (he - hs) % 2 == 1 and (he + 1) <= h:
+        he = he + 1
+    if (he - hs) % 2 == 1 and (hs - 1) >= 0:
+        hs = hs - 1
+    if (we - ws) % 2 == 1 and (we + 1) <= w:
+        we = we + 1
+    if (we - ws) % 2 == 1 and (ws - 1) >= 0:
+        ws = ws - 1
+
+    if he - hs < 2 or we - ws < 2:
+        continue
+
+    data_dict[file] = [cntr, hs, he, ws, we]
+
+np.save(join(dir, 'segs', f'{mode}_bbox_dict.npy'), data_dict)

datasets.py

@@ -141,7 +141,9 @@ class CustomDatasetWithBG(Dataset):
         image = Image.open(self.image_paths[i % self.num_images])
 
         mask_path = self.image_paths[i % self.num_images].replace('.jpeg', '.png').replace('.jpg', '.png').replace('.JPEG', '.png')[:-4] + '_bg.png'
-        mask = np.array(Image.open(mask_path)) / 255.0
+        mask = np.array(Image.open(mask_path))
+
+        mask = np.where(mask > 0, 1, 0)
 
         if not image.mode == "RGB":
             image = image.convert("RGB")

elite.yaml

@@ -1,147 +0,0 @@
-name: elite
-channels:
-  - defaults
-dependencies:
-  - _libgcc_mutex=0.1=main
-  - ca-certificates=2022.10.11=h06a4308_0
-  - certifi=2022.9.24=py39h06a4308_0
-  - ld_impl_linux-64=2.38=h1181459_1
-  - libffi=3.3=he6710b0_2
-  - libgcc-ng=9.1.0=hdf63c60_0
-  - libstdcxx-ng=9.1.0=hdf63c60_0
-  - ncurses=6.3=h7f8727e_2
-  - openssl=1.1.1s=h7f8727e_0
-  - pip=22.2.2=py39h06a4308_0
-  - python=3.9.12=h12debd9_1
-  - readline=8.1.2=h7f8727e_1
-  - sqlite=3.38.5=hc218d9a_0
-  - tk=8.6.12=h1ccaba5_0
-  - wheel=0.37.1=pyhd3eb1b0_0
-  - xz=5.2.5=h7f8727e_1
-  - zlib=1.2.12=h7f8727e_2
-  - pip:
-    - absl-py==1.3.0
-    - accelerate==0.15.0
-    - aiohttp==3.8.3
-    - aiosignal==1.3.1
-    - albumentations==1.1.0
-    - altair==4.2.0
-    - antlr4-python3-runtime==4.8
-    - async-timeout==4.0.2
-    - attrs==22.1.0
-    - blinker==1.5
-    - cachetools==5.2.0
-    - charset-normalizer==2.1.1
-    - click==8.1.3
-    - commonmark==0.9.1
-    - contourpy==1.0.6
-    - cycler==0.11.0
-    - cython==0.29.33
-    - decorator==5.1.1
-    - diffusers==0.11.1
-    - einops==0.4.1
-    - emoji==2.2.0
-    - entrypoints==0.4
-    - faiss-gpu==1.7.2
-    - filelock==3.8.0
-    - fonttools==4.38.0
-    - frozenlist==1.3.3
-    - fsspec==2022.11.0
-    - ftfy==6.1.1
-    - future==0.18.2
-    - gitdb==4.0.9
-    - gitpython==3.1.29
-    - google-auth==2.14.1
-    - google-auth-oauthlib==0.4.6
-    - grpcio==1.50.0
-    - huggingface-hub==0.11.0
-    - idna==3.4
-    - imageio==2.14.1
-    - imageio-ffmpeg==0.4.7
-    - importlib-metadata==5.0.0
-    - jinja2==3.1.2
-    - joblib==1.2.0
-    - jsonschema==4.17.0
-    - kiwisolver==1.4.4
-    - kornia==0.6.0
-    - markdown==3.4.1
-    - markupsafe==2.1.1
-    - matplotlib==3.6.2
-    - multidict==6.0.2
-    - networkx==2.8.8
-    - nltk==3.7
-    - numpy==1.23.4
-    - oauthlib==3.2.2
-    - omegaconf==2.1.1
-    - opencv-python==4.6.0.66
-    - opencv-python-headless==4.6.0.66
-    - packaging==21.3
-    - pandas==1.5.1
-    - pillow==9.0.1
-    - protobuf==3.20.1
-    - psutil==5.9.4
-    - pudb==2019.2
-    - pyarrow==10.0.0
-    - pyasn1==0.4.8
-    - pyasn1-modules==0.2.8
-    - pycocotools==2.0.6
-    - pydeck==0.8.0
-    - pydensecrf==1.0rc2
-    - pydeprecate==0.3.2
-    - pygments==2.13.0
-    - pympler==1.0.1
-    - pyparsing==3.0.9
-    - pyrsistent==0.19.2
-    - python-dateutil==2.8.2
-    - python-dotenv==0.21.0
-    - pytorch-lightning==1.6.5
-    - pytz==2022.6
-    - pytz-deprecation-shim==0.1.0.post0
-    - pywavelets==1.4.1
-    - pyyaml==6.0
-    - qudida==0.0.4
-    - regex==2022.10.31
-    - requests==2.28.1
-    - requests-oauthlib==1.3.1
-    - rich==12.6.0
-    - rsa==4.9
-    - sacremoses==0.0.53
-    - scikit-image==0.19.3
-    - scikit-learn==1.1.3
-    - scipy==1.9.3
-    - semver==2.13.0
-    - setuptools==59.5.0
-    - six==1.16.0
-    - smmap==5.0.0
-    - stanza==1.4.2
-    - streamlit==1.15.0
-    - tensorboard==2.11.0
-    - tensorboard-data-server==0.6.1
-    - tensorboard-plugin-wit==1.8.1
-    - test-tube==0.7.5
-    - threadpoolctl==3.1.0
-    - tifffile==2022.10.10
-    - timm==0.6.12
-    - tokenizers==0.12.1
-    - toml==0.10.2
-    - toolz==0.12.0
-    - torch==1.12.1+cu116
-    - torch-fidelity==0.3.0
-    - torchaudio==0.12.1+cu116
-    - torchmetrics==0.6.0
-    - torchvision==0.13.1+cu116
-    - tornado==6.2
-    - tqdm==4.64.1
-    - transformers==4.25.1
-    - typing-extensions==4.4.0
-    - tzdata==2022.6
-    - tzlocal==4.2
-    - urllib3==1.26.12
-    - urwid==2.1.2
-    - validators==0.20.0
-    - watchdog==2.1.9
-    - wcwidth==0.2.5
-    - werkzeug==2.2.2
-    - yarl==1.8.1
-    - zipp==3.10.0
-

inference_global.py

@@ -170,6 +170,12 @@ def parse_args():
         help="Data index. -1 for all.",
     )
 
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="A seed for testing.",
+    )
     args = parser.parse_args()
     return args
 
@@ -204,11 +210,7 @@ if __name__ == "__main__":
         batch["input_ids"] = batch["input_ids"].to("cuda:0")
         batch["index"] = batch["index"].to("cuda:0").long()
         print(step, batch['text'])
-        seeds = [0, 42, 10086, 777, 555, 222, 111, 999, 327, 283, 190, 218, 2371, 9329, 2938, 2073, 27367, 293,
-                 8269, 87367, 29379, 4658, 39, 598]
-        seeds = sorted(seeds)
-        for seed in seeds:
-            syn_images = validation(batch, tokenizer, image_encoder, text_encoder, unet, mapper, vae, batch["pixel_values_clip"].device, 5,
-                                    token_index=args.token_index, seed=seed)
-            concat = np.concatenate((np.array(syn_images[0]), th2image(batch["pixel_values"][0])), axis=1)
-            Image.fromarray(concat).save(os.path.join(save_dir, f'{str(step).zfill(5)}_{str(seed).zfill(5)}.jpg'))
+        syn_images = validation(batch, tokenizer, image_encoder, text_encoder, unet, mapper, vae, batch["pixel_values_clip"].device, 5,
+                                token_index=args.token_index, seed=args.seed)
+        concat = np.concatenate((np.array(syn_images[0]), th2image(batch["pixel_values"][0])), axis=1)
+        Image.fromarray(concat).save(os.path.join(save_dir, f'{str(step).zfill(5)}_{str(args.seed).zfill(5)}.jpg'))

inference_global.sh

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

inference_local.py

@@ -199,6 +199,13 @@ def parse_args():
         help="Lambda for fuse the global and local feature.",
     )
 
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="A seed for testing.",
+    )
+
     args = parser.parse_args()
     return args
 
@@ -236,12 +243,8 @@ if __name__ == "__main__":
         batch["input_ids"] = batch["input_ids"].to("cuda:0")
         batch["index"] = batch["index"].to("cuda:0").long()
         print(step, batch['text'])
-        seeds = [0, 42, 10086, 777, 555, 222, 111, 999, 327, 283, 190, 218, 2371, 9329, 2938, 2073, 27367, 293,
-                 8269, 87367, 29379, 4658, 39, 598]
-        seeds = sorted(seeds)
-        for seed in seeds:
-            syn_images = validation(batch, tokenizer, image_encoder, text_encoder, unet, mapper, mapper_local, vae,
-                                    batch["pixel_values_clip"].device, 5,
-                                    seed=seed, llambda=float(args.llambda))
-            concat = np.concatenate((np.array(syn_images[0]), th2image(batch["pixel_values"][0])), axis=1)
-            Image.fromarray(concat).save(os.path.join(save_dir, f'{str(step).zfill(5)}_{str(seed).zfill(5)}.jpg'))
+        syn_images = validation(batch, tokenizer, image_encoder, text_encoder, unet, mapper, mapper_local, vae,
+                                batch["pixel_values_clip"].device, 5,
+                                seed=args.seed, llambda=float(args.llambda))
+        concat = np.concatenate((np.array(syn_images[0]), th2image(batch["pixel_values"][0])), axis=1)
+        Image.fromarray(concat).save(os.path.join(save_dir, f'{str(step).zfill(5)}_{str(args.seed).zfill(5)}.jpg'))

inference_local.sh

@@ -5,8 +5,9 @@ CUDA_VISIBLE_DEVICES=7 python inference_local.py \
   --test_data_dir=$DATA_DIR \
   --output_dir="./outputs/local_mapping"  \
   --suffix="object" \
-  --template="a photo of a {}" \
+  --template="a photo of a S" \
   --llambda="0.8" \
   --global_mapper_path="./checkpoints/global_mapper.pt" \
-  --local_mapper_path="./checkpoints/local_mapper.pt"
+  --local_mapper_path="./checkpoints/local_mapper.pt" \
+  --seed 42
 

requirements.txt

@@ -0,0 +1,11 @@
+accelerate==0.16.0
+albumentations==1.3.0
+diffusers==0.11.1
+gradio==3.20.1
+huggingface-hub==0.13.0
+opencv-python-headless==4.7.0.68
+Pillow==9.4.0
+torch==1.13.1
+torchvision==0.14.1
+tqdm==4.65.0
+transformers==4.26.1

train_global.py

@@ -11,7 +11,6 @@ import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch.utils.data import Dataset
 
-import PIL
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
@@ -31,7 +30,7 @@ from PIL import Image
 from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel
 
-from typing import Any, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 from datasets import OpenImagesDataset
 
 
@@ -362,7 +361,6 @@ def get_full_repo_name(model_id: str, organization: Optional[str] = None, token:
     else:
         return f"{organization}/{model_id}"
 
-
 def freeze_params(params):
     for param in params:
         param.requires_grad = False

train_global.sh

@@ -1,6 +1,6 @@
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
-export DATA_DIR='/home/weiyuxiang/datasets/Open_Images/'
-CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch --config_file 4_gpu.json --main_process_port 25656 train_global.py \
+export DATA_DIR='./datasets/Open_Images/'
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_process_port 25656 train_global.py \
   --pretrained_model_name_or_path=$MODEL_NAME \
   --train_data_dir=$DATA_DIR \
   --placeholder_token="S" \
@@ -11,5 +11,5 @@ CUDA_VISIBLE_DEVICES=4,5,6,7 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" \
+  --output_dir="./elite_experiments/global_mapping_new" \
   --save_steps 200

train_local.py

@@ -1,4 +1,3 @@
-
 import argparse
 import itertools
 import math
@@ -16,15 +15,13 @@ import PIL
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
-from diffusers import AutoencoderKL, DDPMScheduler, PNDMScheduler, StableDiffusionPipeline, UNet2DConditionModel, LMSDiscreteScheduler
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, LMSDiscreteScheduler
 from diffusers.optimization import get_scheduler
-from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from huggingface_hub import HfFolder, Repository, whoami
 
-# TODO: remove and import from diffusers.utils when the new version of diffusers is released
 from PIL import Image
 from tqdm.auto import tqdm
-from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer, CLIPVisionModel
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel
 
 
 from typing import Optional

train_local.sh

@@ -1,6 +1,6 @@
 export MODEL_NAME="CompVis/stable-diffusion-v1-4"
-export DATA_DIR='/home/weiyuxiang/datasets/Open_Images/'
-CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch --config_file 4_gpu.json --main_process_port 25657 train_local.py \
+export DATA_DIR='./datasets/Open_Images/'
+CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch --config_file 4_gpu.json --main_process_port 25657 train_local.py \
   --pretrained_model_name_or_path=$MODEL_NAME \
   --train_data_dir=$DATA_DIR \
   --placeholder_token="S" \