1 files changed, 3 insertions, 11 deletions

diff --git a/train_lora.py b/train_lora.py
index 737af58..dea58cf 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -15,7 +15,7 @@ import hidet
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
-from peft import LoraConfig, LoraModel
+from peft import LoraConfig, get_peft_model
 # from diffusers.models.attention_processor import AttnProcessor
 import transformers
 
@@ -731,7 +731,7 @@ def main():
         lora_dropout=args.lora_dropout,
         bias=args.lora_bias,
     )
-    unet = LoraModel(unet_config, unet)
+    unet = get_peft_model(unet, unet_config)
 
     text_encoder_config = LoraConfig(
         r=args.lora_text_encoder_r,
@@ -740,7 +740,7 @@ def main():
         lora_dropout=args.lora_text_encoder_dropout,
         bias=args.lora_text_encoder_bias,
     )
-    text_encoder = LoraModel(text_encoder_config, text_encoder)
+    text_encoder = get_peft_model(text_encoder, text_encoder_config)
 
     vae.enable_slicing()
 
@@ -1167,14 +1167,6 @@ def main():
         group_labels.append("unet")
 
         if training_iter < args.train_text_encoder_cycles:
-            # if len(placeholder_tokens) != 0:
-            #     params_to_optimize.append({
-            #         "params": text_encoder.text_model.embeddings.token_embedding.parameters(),
-            #         "lr": learning_rate_emb,
-            #         "weight_decay": 0,
-            #     })
-            #     group_labels.append("emb")
-
             params_to_optimize.append({
                 "params": (
                     param