feat: 修改调用方法以支持多个网络和缩放参数的处理

app.py

@@ -63,18 +63,18 @@ def merge_lora_networks(networks):
     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
+# 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:
 

utils.py

@@ -66,9 +66,11 @@ def call(
         negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
         negative_target_size: Optional[Tuple[int, int]] = None,
     
-        network=None, 
+        network=None,
+        networks=None,
         start_noise=None,
         scale=None,
+        scales=None,
         unet=None,
     ):
         r"""
@@ -318,12 +320,21 @@ def call(
             num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
             timesteps = timesteps[:num_inference_steps]
         latents = latents.to(unet.dtype)
+        
+        # 统一处理 network，scale | 处理成 list
+        if network is not None:
+            networks = [network]
+        if scale is not None:
+            scales = [scale]
+
         with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if t>start_noise:
-                    network.set_lora_slider(scale=0)
-                else:
-                    network.set_lora_slider(scale=scale)
+            for i, t in enumerate(timesteps):               
+                # 遍历所有网络，设置 scale
+                if networks is not None and scales is not None:
+                    for _network, _scale in zip(networks, scales):
+                        with _network:
+                            _network.set_lora_slider(scale=0 if t > start_noise else _scale)
+                            
                 # expand the latents if we are doing classifier free guidance
                 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
 
@@ -331,15 +342,18 @@ def call(
 
                 # predict the noise residual
                 added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
-                with network:
-                    noise_pred = unet(
-                        latent_model_input,
-                        t,
-                        encoder_hidden_states=prompt_embeds,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                        added_cond_kwargs=added_cond_kwargs,
-                        return_dict=False,
-                    )[0]
+                # 应用多个 LoRA 网络
+                if networks is not None and scales is not None:
+                    for _network, _scale in zip(networks, scales):
+                        with _network:
+                            noise_pred = self.unet(
+                                latent_model_input,
+                                t,
+                                encoder_hidden_states=prompt_embeds,
+                                cross_attention_kwargs=cross_attention_kwargs,
+                                added_cond_kwargs=added_cond_kwargs,
+                                return_dict=False,
+                            )[0]
 
                 # perform guidance
                 if do_classifier_free_guidance: