fix(scheduler): device mismatch from upstream b114620 - move rk and b to device before torch.stack

Author: terarachang

src/diffusers/schedulers/scheduling_unipc_multistep.py

@@ -900,7 +900,7 @@ def multistep_uni_p_bh_update(
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
             lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
             rk = (lambda_si - lambda_s0) / h
-            rks.append(rk)
+            rks.append(rk.to(device))
             D1s.append((mi - m0) / rk)
 
         rks.append(torch.ones((), device=device))
@@ -924,7 +924,7 @@ def multistep_uni_p_bh_update(
 
         for i in range(1, order + 1):
             R.append(torch.pow(rks, i - 1))
-            b.append(h_phi_k * factorial_i / B_h)
+            b.append((h_phi_k * factorial_i / B_h).to(device))
             factorial_i *= i + 1
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
@@ -1035,7 +1035,7 @@ def multistep_uni_c_bh_update(
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
             lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
             rk = (lambda_si - lambda_s0) / h
-            rks.append(rk)
+            rks.append(rk.to(device))
             D1s.append((mi - m0) / rk)
 
         rks.append(torch.ones((), device=device))
@@ -1059,7 +1059,7 @@ def multistep_uni_c_bh_update(
 
         for i in range(1, order + 1):
             R.append(torch.pow(rks, i - 1))
-            b.append(h_phi_k * factorial_i / B_h)
+            b.append((h_phi_k * factorial_i / B_h).to(device))
             factorial_i *= i + 1
             h_phi_k = h_phi_k / hh - 1 / factorial_i