Implementation for sgm_uniform branch

2024-03-19 20:05:54 +08:00
parent c4a00affc5
commit a6b5a513f9
3 changed files with 21 additions and 2 deletions
@@ -0,0 +1,12 @@
+import torch
+
+
+def sgm_uniform(n, sigma_min, sigma_max, inner_model, device):
+    start = inner_model.sigma_to_t(torch.tensor(sigma_max))
+    end = inner_model.sigma_to_t(torch.tensor(sigma_min))
+    sigs = [
+        inner_model.t_to_sigma(ts)
+        for ts in torch.linspace(start, end, n)[:-1]
+    ]
+    sigs += [0.0]
+    return torch.FloatTensor(sigs).to(device)