stable-diffusion-webui

1
"""RNG imitiating torch cuda randn on CPU. You are welcome.
2

3
Usage:
4

5
```
6
g = Generator(seed=0)
7
print(g.randn(shape=(3, 4)))
8
```
9

10
Expected output:
11
```
12
[[-0.92466259 -0.42534415 -2.6438457   0.14518388]
13
 [-0.12086647 -0.57972564 -0.62285122 -0.32838709]
14
 [-1.07454231 -0.36314407 -1.67105067  2.26550497]]
15
```
16
"""
17

18
import numpy as np
19

20
philox_m = [0xD2511F53, 0xCD9E8D57]
21
philox_w = [0x9E3779B9, 0xBB67AE85]
22

23
two_pow32_inv = np.array([2.3283064e-10], dtype=np.float32)
24
two_pow32_inv_2pi = np.array([2.3283064e-10 * 6.2831855], dtype=np.float32)
25

26

27
def uint32(x):
28
    """Converts (N,) np.uint64 array into (2, N) np.unit32 array."""
29
    return x.view(np.uint32).reshape(-1, 2).transpose(1, 0)
30

31

32
def philox4_round(counter, key):
33
    """A single round of the Philox 4x32 random number generator."""
34

35
    v1 = uint32(counter[0].astype(np.uint64) * philox_m[0])
36
    v2 = uint32(counter[2].astype(np.uint64) * philox_m[1])
37

38
    counter[0] = v2[1] ^ counter[1] ^ key[0]
39
    counter[1] = v2[0]
40
    counter[2] = v1[1] ^ counter[3] ^ key[1]
41
    counter[3] = v1[0]
42

43

44
def philox4_32(counter, key, rounds=10):
45
    """Generates 32-bit random numbers using the Philox 4x32 random number generator.
46

47
    Parameters:
48
        counter (numpy.ndarray): A 4xN array of 32-bit integers representing the counter values (offset into generation).
49
        key (numpy.ndarray): A 2xN array of 32-bit integers representing the key values (seed).
50
        rounds (int): The number of rounds to perform.
51

52
    Returns:
53
        numpy.ndarray: A 4xN array of 32-bit integers containing the generated random numbers.
54
    """
55

56
    for _ in range(rounds - 1):
57
        philox4_round(counter, key)
58

59
        key[0] = key[0] + philox_w[0]
60
        key[1] = key[1] + philox_w[1]
61

62
    philox4_round(counter, key)
63
    return counter
64

65

66
def box_muller(x, y):
67
    """Returns just the first out of two numbers generated by Box–Muller transform algorithm."""
68
    u = x * two_pow32_inv + two_pow32_inv / 2
69
    v = y * two_pow32_inv_2pi + two_pow32_inv_2pi / 2
70

71
    s = np.sqrt(-2.0 * np.log(u))
72

73
    r1 = s * np.sin(v)
74
    return r1.astype(np.float32)
75

76

77
class Generator:
78
    """RNG that produces same outputs as torch.randn(..., device='cuda') on CPU"""
79

80
    def __init__(self, seed):
81
        self.seed = seed
82
        self.offset = 0
83

84
    def randn(self, shape):
85
        """Generate a sequence of n standard normal random variables using the Philox 4x32 random number generator and the Box-Muller transform."""
86

87
        n = 1
88
        for x in shape:
89
            n *= x
90

91
        counter = np.zeros((4, n), dtype=np.uint32)
92
        counter[0] = self.offset
93
        counter[2] = np.arange(n, dtype=np.uint32)  # up to 2^32 numbers can be generated - if you want more you'd need to spill into counter[3]
94
        self.offset += 1
95

96
        key = np.empty(n, dtype=np.uint64)
97
        key.fill(self.seed)
98
        key = uint32(key)
99

100
        g = philox4_32(counter, key)
101

102
        return box_muller(g[0], g[1]).reshape(shape)  # discard g[2] and g[3]
103