fix(schedulers): guard Helios schedulers against float64 on MPS

.github/workflows/issue_labeler.yml

@@ -31,6 +31,11 @@ jobs:
           ISSUE_NUMBER: ${{ github.event.issue.number }}
           LABELS: ${{ steps.get-labels.outputs.labels }}
         run: |
+          existing=$(gh label list --limit 100 --json name | python -c "import json,sys; print('\n'.join(l['name'] for l in json.load(sys.stdin)))")
           for label in $(echo "$LABELS" | python -c "import json,sys; print('\n'.join(json.load(sys.stdin)))"); do
-            gh issue edit "$ISSUE_NUMBER" --add-label "$label"
+            if echo "$existing" | grep -Fqx "$label"; then
+              gh issue edit "$ISSUE_NUMBER" --add-label "$label"
+            else
+              echo "::warning::Issue labeler produced '$label', which does not exist in this repo; skipping."
+            fi
           done

src/diffusers/schedulers/scheduling_helios.py

@@ -235,8 +235,13 @@ def set_timesteps(
             ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps)
             sigmas = torch.from_numpy(ratios)
 
-        self.timesteps = torch.from_numpy(timesteps).to(device=device)
-        self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
+        if device is not None and torch.device(device).type == "mps":
+            # mps does not support float64
+            self.timesteps = torch.from_numpy(timesteps.astype(np.float32)).to(device=device)
+            self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device, dtype=torch.float32)
+        else:
+            self.timesteps = torch.from_numpy(timesteps).to(device=device)
+            self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
 
         self._step_index = None
         self.reset_scheduler_history()

src/diffusers/schedulers/scheduling_helios_dmd.py

@@ -213,8 +213,13 @@ def set_timesteps(
             ratios = np.linspace(stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps)
             sigmas = torch.from_numpy(ratios)
 
-        self.timesteps = torch.from_numpy(timesteps).to(device=device)
-        self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
+        if device is not None and torch.device(device).type == "mps":
+            # mps does not support float64
+            self.timesteps = torch.from_numpy(timesteps.astype(np.float32)).to(device=device)
+            self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device, dtype=torch.float32)
+        else:
+            self.timesteps = torch.from_numpy(timesteps).to(device=device)
+            self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
 
         self._step_index = None
         self.reset_scheduler_history()
@@ -275,7 +280,10 @@ def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps):
         # use higher precision for calculations
         original_dtype = flow_pred.dtype
         device = flow_pred.device
-        flow_pred, xt, sigmas, timesteps = (x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps))
+        target_dtype = torch.float32 if device.type == "mps" else torch.float64
+        flow_pred, xt, sigmas, timesteps = (
+            x.to(device=device, dtype=target_dtype) for x in (flow_pred, xt, sigmas, timesteps)
+        )
 
         timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
         sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)