Fix memory leak (#103)

* set start method when init ParaModel

* fast safetensors load tensor to device & fix unload lora

Author: qzzz95

diffsynth_engine/models/basic/lora.py

@@ -131,8 +131,7 @@ def clear(self):
         self._lora_dict.clear()
         self._frozen_lora_list = []
         if self._original_weight is not None:
-            self.weight.data = self._original_weight
-            self._original_weight = None
+            self.weight.data.copy_(self._original_weight)
 
     def forward(self, x):
         w_x = super().forward(x)

diffsynth_engine/pipelines/base.py

@@ -62,7 +62,7 @@ def from_state_dict(
     def load_loras(self, lora_list: List[Tuple[str, float]], fused: bool = True, save_original_weight: bool = False):
         for lora_path, lora_scale in lora_list:
             logger.info(f"loading lora from {lora_path} with scale {lora_scale}")
-            state_dict = load_file(lora_path, device="cpu")
+            state_dict = load_file(lora_path, device=self.device)
             lora_state_dict = self.lora_converter.convert(state_dict)
             for model_name, state_dict in lora_state_dict.items():
                 model = getattr(self, model_name)

diffsynth_engine/utils/loader.py

@@ -24,7 +24,7 @@ def load_file(path: str | os.PathLike, device: str = "cpu"):
             direct_io=(os.environ.get("FAST_SAFETENSORS_DIRECT_IO", "False").upper() == "TRUE"),
         )
         logger.info(f"FastSafetensors Load Model End. Time: {time.time() - start_time:.2f}s")
-        return result
+        return {k: v.to(device) for k, v in result.items()}
     else:
         logger.info(f"Safetensors load model from {path}")
         start_time = time.time()

diffsynth_engine/utils/parallel.py

@@ -329,6 +329,12 @@ def __init__(
         master_port: int = 29500,
         device: str = "cuda",
     ):
+        current_method = mp.get_start_method(allow_none=True)
+        if current_method is None or current_method != 'spawn':
+            try:
+                mp.set_start_method('spawn')
+            except RuntimeError as e:
+                raise RuntimeError("Failed to set start method to spawn:", e)
         super().__init__()
         self.world_size = cfg_degree * sp_ulysses_degree * sp_ring_degree * tp_degree
         self.queue_in = mp.Queue()