update

trellis2/models/sparse_structure_flow.py

@@ -4,8 +4,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from ..trainers.utils import str_to_dtype
-from ..modules.utils import convert_module_to, manual_cast
+from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
 from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
 from ..modules.attention import RotaryPositionEmbedder
 

trellis2/models/structured_latent_flow.py

@@ -4,8 +4,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from ..trainers.utils import str_to_dtype
-from ..modules.utils import convert_module_to, manual_cast
+from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
 from ..modules.transformer import AbsolutePositionEmbedder
 from ..modules import sparse as sp
 from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock

trellis2/modules/utils.py

@@ -72,3 +72,16 @@ def manual_cast(tensor, dtype):
     if not torch.is_autocast_enabled():
         return tensor.type(dtype)
     return tensor
+
+
+def str_to_dtype(dtype_str: str):
+    return {
+        'f16': torch.float16,
+        'fp16': torch.float16,
+        'float16': torch.float16,
+        'bf16': torch.bfloat16,
+        'bfloat16': torch.bfloat16,
+        'f32': torch.float32,
+        'fp32': torch.float32,
+        'float32': torch.float32,
+    }[dtype_str]