overwrite scheduling_unipc_multistep.py with main's version

Author: Ting-Yun Chang

src/diffusers/schedulers/scheduling_unipc_multistep.py

@@ -900,11 +900,11 @@ 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.to(device))
+            rks.append(rk)
             D1s.append((mi - m0) / rk)
 
-        rks.append(torch.ones((), device=device))
-        rks = torch.stack(rks)
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
 
         R = []
         b = []
@@ -924,18 +924,18 @@ 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).to(device))
+            b.append(h_phi_k * factorial_i / B_h)
             factorial_i *= i + 1
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
         R = torch.stack(R)
-        b = torch.stack(b) if len(b) > 0 else torch.tensor(b, device=device)
+        b = torch.tensor(b, device=device)
 
         if len(D1s) > 0:
             D1s = torch.stack(D1s, dim=1)  # (B, K)
             # for order 2, we use a simplified version
             if order == 2:
-                rhos_p = torch.ones(1, dtype=x.dtype, device=device) * 0.5
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
             else:
                 rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
         else:
@@ -1035,11 +1035,11 @@ 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.to(device))
+            rks.append(rk)
             D1s.append((mi - m0) / rk)
 
-        rks.append(torch.ones((), device=device))
-        rks = torch.stack(rks)
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
 
         R = []
         b = []
@@ -1059,12 +1059,12 @@ 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).to(device))
+            b.append(h_phi_k * factorial_i / B_h)
             factorial_i *= i + 1
             h_phi_k = h_phi_k / hh - 1 / factorial_i
 
         R = torch.stack(R)
-        b = torch.stack(b) if len(b) > 0 else torch.tensor(b, device=device)
+        b = torch.tensor(b, device=device)
 
         if len(D1s) > 0:
             D1s = torch.stack(D1s, dim=1)
@@ -1073,7 +1073,7 @@ def multistep_uni_c_bh_update(
 
         # for order 1, we use a simplified version
         if order == 1:
-            rhos_c = torch.ones(1, dtype=x.dtype, device=device) * 0.5
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
         else:
             rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)