Update diffusers_helper/memory.py

diffusers_helper/memory.py

@@ -1,13 +1,36 @@
 # By lllyasviel
+# WindowsとHugging Face Space環境の両方に対応した DynamicSwap + zeroGPU 対応バージョン
 
-
+import os
 import torch
 
+# Hugging Face Space環境で実行されているかどうかを確認
+IN_HF_SPACE = os.environ.get('SPACE_ID') is not None
 
+# CPU デバイスを設定
 cpu = torch.device('cpu')
-gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
-gpu_complete_modules = []
 
+# ステートレスGPU環境では、メインプロセスでCUDAを初期化しない
+def get_gpu_device():
+    if IN_HF_SPACE:
+        # Spacesではデバイスの初期化を遅延させる
+        return 'cuda'
+    try:
+        if torch.cuda.is_available():
+            return torch.device(f'cuda:{torch.cuda.current_device()}')
+        else:
+            print("CUDAが利用できません。デフォルトデバイスとしてCPUを使用します")
+            return torch.device('cpu')
+    except Exception as e:
+        print(f"CUDAデバイスの初期化中にエラーが発生しました: {e}")
+        print("CPUデバイスにフォールバックします")
+        return torch.device('cpu')
+
+# GPUデバイスを取得（文字列または実際のデバイスオブジェクト）
+gpu = get_gpu_device()
+
+# 完全にGPUにロードされたモジュールのリスト
+gpu_complete_modules = []
 
 class DynamicSwapInstaller:
     @staticmethod
@@ -32,37 +55,36 @@ class DynamicSwapInstaller:
                     return _buffers[name].to(**kwargs)
             return super(original_class, self).__getattr__(name)
 
-        module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
-            '__getattr__': hacked_get_attr,
-        })
-
-        return
+        module.__class__ = type(
+            'DynamicSwap_' + original_class.__name__,
+            (original_class,),
+            {'__getattr__': hacked_get_attr}
+        )
 
     @staticmethod
     def _uninstall_module(module: torch.nn.Module):
         if 'forge_backup_original_class' in module.__dict__:
             module.__class__ = module.__dict__.pop('forge_backup_original_class')
-        return
 
     @staticmethod
     def install_model(model: torch.nn.Module, **kwargs):
         for m in model.modules():
             DynamicSwapInstaller._install_module(m, **kwargs)
-        return
 
     @staticmethod
     def uninstall_model(model: torch.nn.Module):
         for m in model.modules():
             DynamicSwapInstaller._uninstall_module(m)
-        return
 
 
-def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
+def fake_diffusers_current_device(model: torch.nn.Module, target_device):
+    # 文字列デバイスをtorch.deviceに変換
+    if isinstance(target_device, str):
+        target_device = torch.device(target_device)
     if hasattr(model, 'scale_shift_table'):
         model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
         return
-
-    for k, p in model.named_modules():
+    for _, p in model.named_modules():
         if hasattr(p, 'weight'):
             p.to(target_device)
             return
@@ -71,64 +93,80 @@ def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.d
 def get_cuda_free_memory_gb(device=None):
     if device is None:
         device = gpu
-
-    memory_stats = torch.cuda.memory_stats(device)
-    bytes_active = memory_stats['active_bytes.all.current']
-    bytes_reserved = memory_stats['reserved_bytes.all.current']
-    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
-    bytes_inactive_reserved = bytes_reserved - bytes_active
-    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
-    return bytes_total_available / (1024 ** 3)
+    if isinstance(device, str):
+        device = torch.device(device)
+    if device.type != 'cuda':
+        # CUDAでない場合はデフォルト値
+        return 6.0
+    try:
+        stats = torch.cuda.memory_stats(device)
+        active = stats['active_bytes.all.current']
+        reserved = stats['reserved_bytes.all.current']
+        free_cuda, _ = torch.cuda.mem_get_info(device)
+        inactive = reserved - active
+        available = free_cuda + inactive
+        return available / (1024 ** 3)
+    except Exception as e:
+        print(f"CUDAメモリ情報取得エラー: {e}")
+        return 6.0
 
 
 def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
-    print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
-
+    print(f"{model.__class__.__name__} を {target_device} に移動します。保持メモリ: {preserved_memory_gb} GB")
+    if isinstance(target_device, str):
+        target_device = torch.device(target_device)
+    # CPUまたはGPU未使用時は直接移動
+    if target_device.type == 'cpu':
+        model.to(device=target_device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return
     for m in model.modules():
         if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
             torch.cuda.empty_cache()
             return
-
         if hasattr(m, 'weight'):
             m.to(device=target_device)
-
     model.to(device=target_device)
     torch.cuda.empty_cache()
-    return
 
 
 def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
-    print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
-
+    print(f"メモリ保持のため {model.__class__.__name__} を {target_device} からオフロードします: {preserved_memory_gb} GB")
+    if isinstance(target_device, str):
+        target_device = torch.device(target_device)
+    if target_device.type == 'cpu':
+        model.to(device=cpu)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return
     for m in model.modules():
         if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
             torch.cuda.empty_cache()
             return
-
         if hasattr(m, 'weight'):
             m.to(device=cpu)
-
     model.to(device=cpu)
     torch.cuda.empty_cache()
-    return
 
 
 def unload_complete_models(*args):
     for m in gpu_complete_modules + list(args):
+        if m is None:
+            continue
         m.to(device=cpu)
-        print(f'Unloaded {m.__class__.__name__} as complete.')
-
+        print(f"{m.__class__.__name__} を完全にアンロードしました")
     gpu_complete_modules.clear()
-    torch.cuda.empty_cache()
-    return
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
 
 def load_model_as_complete(model, target_device, unload=True):
+    if isinstance(target_device, str):
+        target_device = torch.device(target_device)
     if unload:
         unload_complete_models()
-
     model.to(device=target_device)
-    print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
-
+    print(f"{model.__class__.__name__} を {target_device} に完全にロードしました")
     gpu_complete_modules.append(model)
-    return
+