chore: sync with upstream

app.py

@@ -94,12 +94,17 @@ pipe = StableDiffusionPipeline(
 )
 
 unet.set_attn_processor(CrossAttnProcessor)
+pipe.setup_text_encoder(clip_skip, text_encoder)
 if torch.cuda.is_available():
     pipe = pipe.to("cuda")
 
 def get_model_list():
     return models + alt_models
 
+te_cache = {
+    base_name: text_encoder
+}
+
 unet_cache = {
     base_name: unet
 }
@@ -108,35 +113,50 @@ lora_cache = {
     base_name: LoRANetwork(text_encoder, unet)
 }
 
-def get_model(name):
-    local_models = models + alt_models
-    keys = [k[0] for k in local_models]
+def setup_model(name, lora_state=None, lora_scale=1.0):
+    global pipe
+
+    keys = [k[0] for k in models]
     if name not in unet_cache:
         if name not in keys:
             raise ValueError(name)
         else:
+
+            text_encoder = CLIPTextModel.from_pretrained(
+                models[keys.index(name)][1],
+                subfolder="text_encoder",
+                torch_dtype=torch.float16,
+            )
             unet = UNet2DConditionModel.from_pretrained(
-                local_models[keys.index(name)][1],
+                models[keys.index(name)][1],
                 subfolder="unet",
                 torch_dtype=torch.float16,
             )
-            unet.to("cuda")
+
+            if torch.cuda.is_available():
+                unet.to("cuda")
+                text_encoder.to("cuda")
+
             unet_cache[name] = unet
-            lora_cache[name] = LoRANetwork(lora_cache[base_name].text_encoder_loras, unet)
-
-    g_unet = unet_cache[name]
-    g_lora = lora_cache[name]
-    g_unet.set_attn_processor(CrossAttnProcessor())
-    g_lora.reset()
-    clip_skip = local_models[keys.index(name)][2]
-    if torch.cuda.is_available():
-        g_unet.to("cuda")
-        g_lora.to("cuda")
-    return g_unet, g_lora, clip_skip
+            te_cache[name] = text_encoder
+            lora_cache[name] = LoRANetwork(text_encoder, unet)
+
+    local_te, local_unet, local_lora, = te_cache[name], unet_cache[name], lora_cache[name]
+    local_unet.set_attn_processor(CrossAttnProcessor())
+    local_lora.reset()
+    clip_skip = models[keys.index(name)][2]
+
+    if lora_state is not None and lora_state != "":
+        local_lora.load(lora_state, lora_scale)
+        local_lora.to(local_unet.device, dtype=local_unet.dtype)
+
+    pipe.setup_unet(local_unet)
+    pipe.setup_text_encoder(clip_skip, local_te)
+    return pipe
 
 # precache on huggingface
 for model in models:
-    get_model(model[0])
+    setup_model(model[0])
 
 def error_str(error, title="Error"):
     return (
@@ -218,13 +238,7 @@ def inference(
     restore_all()
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
-    local_unet, local_lora, clip_skip = get_model(model)
-    pipe.set_clip_skip(clip_skip)
-    if lora_state is not None and lora_state != "":
-        local_lora.load(lora_state, lora_scale)
-        local_lora.to(local_unet.device, dtype=local_unet.dtype)
-
-    pipe.setup_unet(local_unet)
+    setup_model(model, lora_state, lora_scale)
     sampler_name, sampler_opt = None, None
     for label, funcname, options in samplers_k_diffusion:
         if label == sampler:

modules/lora.py

@@ -55,8 +55,9 @@ class LoRAModule(torch.nn.Module):
         self.org_module = org_module  # remove in applying
         self.enable = False
 
-    def resize(self, rank, alpha):
+    def resize(self, rank, alpha, multiplier):
         self.alpha = torch.tensor(alpha)
+        self.multiplier = multiplier
         self.scale = alpha / rank
         if self.lora_down.__class__.__name__ == "Conv2d":
             in_dim = self.lora_down.in_channels
@@ -172,10 +173,11 @@ class LoRANetwork(torch.nn.Module):
             weights_to_modify += self.unet_loras
 
         for lora in self.text_encoder_loras + self.unet_loras:
-            lora.resize(network_dim, network_alpha)
+            lora.resize(network_dim, network_alpha, scale)
             if lora in weights_to_modify:
                 lora.enable = True
 
         info = self.load_state_dict(weights, False)
-        print(f"Weights are loaded. Unexpect keys={info.unexpected_keys}")
-        
+        if len(info.unexpected_keys) > 0:
+            print(f"Weights are loaded. Unexpected keys={info.unexpected_keys}")
+            

modules/model.py

@@ -185,11 +185,13 @@ class StableDiffusionPipeline(DiffusionPipeline):
             scheduler=scheduler,
         )
         self.setup_unet(self.unet)
-        self.prompt_parser = FrozenCLIPEmbedderWithCustomWords(
-            self.tokenizer, self.text_encoder
-        )
-        
-    def set_clip_skip(self, n):
+        self.setup_text_encoder()
+
+    def setup_text_encoder(self, n=1, new_encoder=None):
+        if new_encoder is not None:
+            self.text_encoder = new_encoder
+
+        self.prompt_parser = FrozenCLIPEmbedderWithCustomWords(self.tokenizer, self.text_encoder)
         self.prompt_parser.CLIP_stop_at_last_layers = n
 
     def setup_unet(self, unet):