move

Author: akaitsuki-ii

diffsynth_engine/utils/fp8_linear.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from contextlib import contextmanager
-from diffsynth_engine.utils.constants import DTYPE_FP8
+from diffsynth_engine.utils.platform import DTYPE_FP8
 
 
 def enable_fp8_linear(module: nn.Module):

diffsynth_engine/utils/platform.py

@@ -1,7 +1,15 @@
+# cross-platform definitions and utilities
 import torch
 import gc
 
-# 存放跨平台的工具类
+
+# data type
+# AMD only supports float8_e4m3fnuz
+# https://onnx.ai/onnx/technical/float8.html
+if torch.version.hip and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName:
+    DTYPE_FP8 = torch.float8_e4m3fnuz
+else:
+    DTYPE_FP8 = torch.float8_e4m3fn
 
 
 def empty_cache():