zero gpu

apg_guidance.py

@@ -17,14 +17,15 @@ def project(
     dims=[-1, -2],
 ):
     dtype = v0.dtype
-    if v0.device.type == "mps":
-        v0, v1 = v0.float(), v1.float()
-    else:
-        v0, v1 = v0.double(), v1.double()
+    device_type = v0.device.type
+    if device_type == "mps":
+        v0, v1 = v0.cpu(), v1.cpu()
+
+    v0, v1 = v0.double(), v1.double()
     v1 = torch.nn.functional.normalize(v1, dim=dims)
     v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
     v0_orthogonal = v0 - v0_parallel
-    return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)
 
 
 def apg_forward(

pipeline_ace_step.py

@@ -955,6 +955,7 @@ class ACEStepPipeline:
         latents, _ = self.music_dcae.encode(input_audio, sr=sr)
         return latents
 
+    @spaces.GPU
     def __call__(
         self,
         audio_duration: float = 60.0,