pin memory bug fix (#141)

Author: tenderness-git

diffsynth_engine/utils/offload.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from typing import Dict
-
+import platform
 
 def enable_sequential_cpu_offload(module: nn.Module, device: str = "cuda"):
     module = module.to("cpu")
@@ -26,13 +26,14 @@ def _forward_pre_hook(module: nn.Module, input_):
             for name, buffer in module.named_buffers(recurse=recurse):
                 buffer.data = buffer.data.to(device=device)
             return tuple(x.to(device=device) if isinstance(x, torch.Tensor) else x for x in input_)
-
-        for name, param in module.named_parameters(recurse=recurse):
-            param.data = param.data.pin_memory()
+        for name, param in module.named_parameters(recurse=recurse):            
+            if platform.system()  == 'Linux':
+                param.data = param.data.pin_memory()
             offload_param_dict[name] = param.data
             param.data = param.data.to(device=device)
         for name, buffer in module.named_buffers(recurse=recurse):
-            buffer.data = buffer.data.pin_memory()
+            if platform.system()  == 'Linux':            
+                buffer.data = buffer.data.pin_memory()
             offload_param_dict[name] = buffer.data
             buffer.data = buffer.data.to(device=device)
         setattr(module, "_offload_param_dict", offload_param_dict)
@@ -58,10 +59,12 @@ def offload_model_to_dict(module: nn.Module) -> Dict[str, torch.Tensor]:
     module = module.to("cpu")
     offload_param_dict = {}
     for name, param in module.named_parameters(recurse=True):
-        param.data = param.data.pin_memory()
+        if platform.system()  == 'Linux':        
+            param.data = param.data.pin_memory()
         offload_param_dict[name] = param.data
     for name, buffer in module.named_buffers(recurse=True):
-        buffer.data = buffer.data.pin_memory()
+        if platform.system()  == 'Linux':        
+            buffer.data = buffer.data.pin_memory()
         offload_param_dict[name] = buffer.data
     return offload_param_dict