add add mps support

2025-05-04 20:03:57 +08:00
parent 6cb7c1d894
commit d77c5f42f1
6 changed files with 56 additions and 57 deletions
@@ -17,7 +17,10 @@ def project(
    dims=[-1, -2],
 ):
    dtype = v0.dtype
-    v0, v1 = v0.double(), v1.double()
+    if v0.device.type == "mps":
+        v0, v1 = v0.float(), v1.float()
+    else:
+        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
@@ -53,6 +56,7 @@ def apg_forward(
 def cfg_forward(cond_output, uncond_output, cfg_strength):
    return uncond_output + cfg_strength * (cond_output - uncond_output)

+
 def cfg_double_condition_forward(
    cond_output,
    uncond_output,