feat: 修改调用方法以支持传递网络和缩放参数

app.py

@@ -62,6 +62,20 @@ def merge_lora_networks(networks):
                 base_network.state_dict()[name] = param.clone()
     return base_network
 
+# 修改 call 方法以支持传递 networks 参数
+def rw_sd_call(self, *args, networks=None, scales=None, **kwargs):
+    if networks is not None and scales is not None:
+        for network, scale in zip(networks, scales):
+            for name, param in network.named_parameters():
+                if name in self.unet.state_dict():
+                    self.unet.state_dict()[name].add_(param * scale)
+                else:
+                    self.unet.state_dict()[name] = param.clone() * scale
+    return self.__original_call__(*args, **kwargs)
+
+StableDiffusionXLPipeline.__original_call__ = StableDiffusionXLPipeline.__call__
+StableDiffusionXLPipeline.__call__ = rw_sd_call
+
 class Demo:
 
     def __init__(self) -> None:
@@ -382,7 +396,7 @@ class Demo:
                     alpha=alpha,
                     train_method=train_method,
                 ).to(self.device, dtype=self.weight_dtype)
-            network.load_state_dict(torch.load(model_path))
+            network.load_state_dict(torch.load(model_path, weights_only=True))
             networks.append(network)
 
         # 设置种子