Use torchruntime's device utilities

Author: cmdr2

sdkit/__init__.py

@@ -58,9 +58,9 @@ def __init__(self) -> None:
         https://github.com/sczhou/CodeFormer/blob/master/LICENSE
         """
 
-        from sdkit.utils import get_torch_platform
+        from torchruntime.utils import get_installed_torch_platform
 
-        self.device = get_torch_platform()[0]
+        self.device = get_installed_torch_platform()[0]
 
     # hacky approach, but we need to enforce full precision for some devices
     # we also need to force full precision for these devices (haven't implemented this yet):
@@ -73,7 +73,7 @@ def device(self):
     def device(self, d):
         self._device = d
 
-        from sdkit.utils import get_device
+        from torchruntime.utils import get_device
 
         if d.split(":")[0] in ("cpu", "mps"):
             from sdkit.utils import log

sdkit/utils/__init__.py

@@ -44,10 +44,6 @@
 )
 from .device_utils import (
     has_amd_gpu,
-    get_torch_platform,
-    get_device_count,
-    get_device_name,
-    get_device,
     mem_get_info,
     memory_allocated,
     memory_stats,

sdkit/utils/memory_utils.py

@@ -20,8 +20,8 @@ def gc(context: Context):
 
 
 def get_device_usage(device, log_info=False, process_usage_only=True, log_prefix=""):
-    import torch
-    from sdkit.utils import log, is_cpu_device, get_device, mem_get_info, memory_allocated, memory_stats
+    from torchruntime.utils import get_device
+    from sdkit.utils import log, is_cpu_device, mem_get_info, memory_allocated, memory_stats
 
     if isinstance(device, str):
         device = get_device(device)
@@ -120,7 +120,7 @@ def get_tensors_in_memory(device):
     prevent garbage-collection of all the tensors in memory.**
     """
     import torch
-    from .device_utils import get_device
+    from torchruntime.utils import get_device
 
     if isinstance(device, str):
         device = get_device(device)

setup.py

@@ -17,5 +17,6 @@
         "controlnet-aux==0.0.6",
         "invisible-watermark==0.2.0",  # required for SD XL
         "huggingface_hub==0.21.4",
+        "torchruntime>=1.7.0",
     ],
 )