Update the code(now can be used with single node single gpu, no mix precision)

Author: SH

src/diffusers_training.py

@@ -246,7 +246,7 @@
 from src.diffusers_data_pipeline import CustomDiffusionDataset, PromptDataset, collate_fn
 from src import retrieve
 
-check_min_version("0.21.4")
+# check_min_version("0.21.4")
 
 logger = get_logger(__name__)
 

src/diffusers_training_sdxl.py

@@ -12,7 +12,6 @@
 #                               Apache License
 #                            Version 2.0, January 2004
 #                         http://www.apache.org/licenses/
-
 #    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
 
 #    1. Definitions.
@@ -220,6 +219,7 @@
 import math
 import os
 import shutil
+import json
 import warnings
 from pathlib import Path
 
@@ -259,7 +259,7 @@
 from src import retrieve
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.21.4")
+# check_min_version("0.21.4")
 
 logger = get_logger(__name__)
 
@@ -269,13 +269,13 @@ def create_custom_diffusion(unet, freeze_model):
         if freeze_model == 'crossattn':
             if 'attn2' in name:
                 params.requires_grad = True
-                print(name)
+                # print(name)
             else:
                 params.requires_grad = False
         elif freeze_model == "crossattn_kv":
             if 'attn2.to_k' in name or 'attn2.to_v' in name:
                 params.requires_grad = True
-                print(name)
+                # print(name)
             else:
                 params.requires_grad = False
         else:
@@ -830,7 +830,6 @@ def main(args):
                         sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
                     ):
                         images = pipeline(example["prompt"]).images
-
                         for i, image in enumerate(images):
                             hash_image = hashlib.sha1(image.tobytes()).hexdigest()
                             image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg"

train.sh

@@ -0,0 +1,37 @@
+## launch training script (2 GPUs recommended, increase --max_train_steps to 500 if 1 GPU)
+export CUDA_VISIBLE_DEVICES=0
+export MODEL_NAME="/data/home/chensh/data/huggingface_model/stable-diffusion-xl-base-1.0"
+
+export INSTANCE_DIR="./data/cat"
+export INSTANCE_PROMPT="photo of a <new1> cat"
+export CLASS_DIR="./sample_reg/samples_cat/"
+export CLASS_PROMPT="cat"
+export OUTPUT_DIR="./logs/cat"
+export modifier_token="<new1>"
+
+#export INSTANCE_DIR="./data/wooden_pot"
+#export INSTANCE_PROMPT="photo of a <new2> wooden pot"
+#export CLASS_DIR="./data/prior_woodenpot/"
+#export CLASS_PROMPT="wooden pot"
+#export OUTPUT_DIR="./logs/wooden_pot"
+#export modifier_token="<new2>"
+
+accelerate launch src/diffusers_training_sdxl.py \
+          --pretrained_model_name_or_path=$MODEL_NAME  \
+          --instance_data_dir=$INSTANCE_DIR \
+          --class_data_dir=$CLASS_DIR \
+          --output_dir=$OUTPUT_DIR  \
+          --with_prior_preservation --prior_loss_weight=1.0 \
+          --instance_prompt="${INSTANCE_PROMPT}"  \
+          --class_prompt="${CLASS_PROMPT}" \
+          --resolution=1024  \
+          --train_batch_size=1  \
+          --learning_rate=1e-5  \
+          --lr_warmup_steps=0 \
+          --max_train_steps=1000 \
+          --num_class_images=200 \
+          --scale_lr --hflip  \
+          --modifier_token="${modifier_token}"
+
+### sample
+#python src/diffusers_sample.py --delta_ckpt logs/cat/delta.bin --ckpt "CompVis/stable-diffusion-v1-4" --prompt "<new1> cat playing with a ball"