Update app.py

app.py

@@ -120,8 +120,11 @@ def infer(
     pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir)
     pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, text_encoder_sub_dir)
 
-    pipe.unet.load_state_dict({k: lora_scale*v for k, v in pipe.unet.state_dict().items()})
-    pipe.text_encoder.load_state_dict({k: lora_scale*v for k, v in pipe.text_encoder.state_dict().items()})
+    pipe.unet.add_weighted_adapter(['default'], lora_scale, 'lora')
+    pipe.text_encoder.add_weighted_adapter(['default'], lora_scale, 'lora')
+
+    # pipe.unet.load_state_dict({k: lora_scale*v for k, v in pipe.unet.state_dict().items()})
+    # pipe.text_encoder.load_state_dict({k: lora_scale*v for k, v in pipe.text_encoder.state_dict().items()})
 
     if tiny_vae:
         pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesd", torch_dtype=torch_dtype)