nyanko7 commited on
Commit
4508ef4
1 Parent(s): 1d9ec1d

Update modules/lora.py

Browse files
Files changed (1) hide show
  1. modules/lora.py +5 -1
modules/lora.py CHANGED
@@ -7,6 +7,7 @@
7
  import math
8
  import os
9
  import torch
 
10
  import modules.safe as _
11
  from safetensors.torch import load_file
12
 
@@ -115,7 +116,10 @@ class LoRANetwork(torch.nn.Module):
115
  else:
116
  self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
117
  print(f"Create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
118
-
 
 
 
119
  self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
120
  print(f"Create LoRA for U-Net: {len(self.unet_loras)} modules.")
121
 
 
7
  import math
8
  import os
9
  import torch
10
+ import diffusers
11
  import modules.safe as _
12
  from safetensors.torch import load_file
13
 
 
116
  else:
117
  self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
118
  print(f"Create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
119
+
120
+ if diffusers.__version__ >= "0.15.0":
121
+ LoRANetwork.UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
122
+
123
  self.unet_loras = create_modules(LoRANetwork.LORA_PREFIX_UNET, unet, LoRANetwork.UNET_TARGET_REPLACE_MODULE)
124
  print(f"Create LoRA for U-Net: {len(self.unet_loras)} modules.")
125