DEVICE

mmaudio/ext/rotary_embeddings.py

@@ -7,7 +7,7 @@ from torch import Tensor
 # Ref: https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py
 # Ref: https://github.com/lucidrains/rotary-embedding-torch
 
-
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 def compute_rope_rotations(length: int,
                            dim: int,
                            theta: int,
@@ -17,7 +17,7 @@ def compute_rope_rotations(length: int,
     assert dim % 2 == 0
 
     # with torch.amp.autocast(device_type='cuda', enabled=False):
-    with torch.amp.autocast(device_type=device, enabled=False):
+    with torch.amp.autocast(device_type=DEVICE, enabled=False):
         pos = torch.arange(length, dtype=torch.float32, device=device)
         freqs = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
         freqs *= freq_scaling
@@ -28,9 +28,9 @@ def compute_rope_rotations(length: int,
         return rot
 
 
-def apply_rope(x: Tensor, rot: Tensor, device: Union[torch.device, str] = 'cpu') -> tuple[Tensor, Tensor]:
+def apply_rope(x: Tensor, rot: Tensor) -> tuple[Tensor, Tensor]:
     # with torch.amp.autocast(device_type='cuda', enabled=False):
-    with torch.amp.autocast(device_type=device, enabled=False):
+    with torch.amp.autocast(device_type=DEVICE, enabled=False):
         _x = x.float()
         _x = _x.view(*_x.shape[:-1], -1, 1, 2)
         x_out = rot[..., 0] * _x[..., 0] + rot[..., 1] * _x[..., 1]