prettier u

2025-10-01 18:52:40 +09:00
parent 29eb04d1a4
commit a601dc6095
3 changed files with 341 additions and 7 deletions
--- a/src/rf.py
+++ b/src/rf.py
@@ -1,5 +1,3 @@
-import math
-
 import lpips
 import torch

@@ -32,8 +30,9 @@ class RF:
        b = x0.size(0)
        if self.ushaped:
            a = 4.0  # HYPERPARMS
-            u = torch.rand((b,)).to(x0.device)
-            t = torch.asinh((2 * u - 1) * math.sinh(a)) / a
+            u = torch.rand((b,), device=x0.device)
+            t = torch.acosh(1 + (torch.cosh(torch.tensor(a)) - 1) * u) / a
+            t = t * (torch.randint(0, 2, (b,), device=x0.device) * 2 - 1) * 0.5 + 0.5
        elif self.ln:
            nt = torch.randn((b,)).to(x0.device)
            t = torch.sigmoid(nt)