Adapt to Ascend NPU

environment_npu.yml

@@ -0,0 +1,15 @@
+name: DiT
+channels:
+  - pytorch
+dependencies:
+  - python >= 3.8
+  - pytorch >= 1.13
+  - torchvision
+  - pip:
+    - timm
+    - diffusers
+    - accelerate
+    - protobuf
+    - decorator
+    - scipy
+    - attrs

sample.py

@@ -16,6 +16,9 @@
 from download import find_model
 from models import DiT_models
 import argparse
+from utils.device_utils import is_npu_available
+if is_npu_available():
+    from torch_npu.contrib import transfer_to_npu
 
 
 def main(args):

train.py

@@ -30,6 +30,9 @@
 from models import DiT_models
 from diffusion import create_diffusion
 from diffusers.models import AutoencoderKL
+from utils.device_utils import is_npu_available
+if is_npu_available():
+    from torch_npu.contrib import transfer_to_npu
 
 
 #################################################################################

utils/device_utils.py

@@ -0,0 +1,17 @@
+import torch
+import importlib
+
+
+def is_npu_available():
+    "Checks if `torch_npu` is installed and potentially if a NPU is in the environment"
+    if importlib.util.find_spec("torch") is None or importlib.util.find_spec("torch_npu") is None:
+        return False
+
+    import torch_npu
+
+    try:
+        # Will raise a RuntimeError if no NPU is found
+        _ = torch.npu.device_count()
+        return torch.npu.is_available()
+    except RuntimeError:
+        return False