File size: 7,849 Bytes
62d3300 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
"""Fourier embeddings for given noise levels.
We supply fixed weights and biases for the Fourier embeddings. These were
initially generated by the following code, but we make them into constants
to future proof against changes in jax rng generation:
```
dim = 256
w_key, b_key = jax.random.split(jax.random.PRNGKey(42))
weight = jax.random.normal(w_key, shape=[dim])
bias = jax.random.uniform(b_key, shape=[dim])
```
"""
import jax.numpy as jnp
# pyformat: disable
# pylint: disable=bad-whitespace
# pylint: disable=bad-continuation
_WEIGHT = [
0.45873642, 0.06516238, -0.07278306, -0.26992258, 0.64292115,
-0.40763968, 3.60116863, 0.54461384, -0.32644904, 2.10888267,
1.30805349, 1.19838560, -1.37745857, 1.99475312, -1.64120293,
1.07823789, -0.02288206, 0.88305283, 0.48099944, 0.17655374,
0.30281949, 0.80646873, 0.62605333, -0.23965347, -1.02609432,
0.75006109, -0.19913037, 0.07466396, 0.66431236, -0.60990530,
-0.69709194, -0.44453633, -1.77656078, 0.02299878, 0.04095552,
0.35485864, -0.47602659, -0.98820388, -0.24106771, -1.07254291,
-0.99741757, 0.22697604, 1.41390419, 1.54984057, -0.12237291,
0.20156337, 0.61767143, 0.23959029, 0.92454034, 1.84082258,
0.89030224, 0.39598912, -1.52224910, 0.29669049, 1.52356744,
-0.33968377, 0.24155144, -0.52308381, -0.23622665, 0.92825454,
-0.63864607, -0.62169307, 0.78623551, -0.80352145, -0.45496067,
1.30877995, -0.06686528, 1.00248849, -0.63593471, 0.16372502,
-1.46133232, 1.10562658, -0.01693927, 0.28684548, -0.72843230,
0.66133535, -1.92225552, 0.70241231, -0.96868867, -0.47309339,
-1.66894221, 0.46018723, -0.56806105, 0.32694784, -0.46529883,
1.02299964, 0.84688205, 1.19581807, -1.82454145, 0.05999713,
-0.59530073, 1.44862521, -0.34933713, -0.46564487, -0.55005538,
-1.61170268, 0.17502306, 0.38670063, -1.12133658, -0.29343036,
-0.52527446, -1.26285112, 1.07982683, 0.51215219, 1.48963666,
1.09847653, -0.01563358, 0.32574457, 1.94779706, -1.29198587,
1.06249654, -0.86965990, 0.22975266, -0.27182648, -0.21130897,
-0.41773933, -0.02329035, 1.31049252, 0.05579265, -1.23127055,
-0.99691105, 0.27058721, -0.72509319, -0.14421797, -1.48605061,
1.35041201, 1.29619241, -1.01022530, -0.79787987, -0.16166858,
0.87210685, 1.69248152, 1.42469788, -0.72325104, -1.24823737,
0.07051118, 0.71332991, -0.07360429, -0.91955227, -2.68856549,
-0.44033936, 0.35482934, -0.57933813, 0.97468042, -0.31050494,
-0.88454425, -2.08785224, 0.47322822, -0.02400172, 0.26644820,
-0.19147627, -2.10538960, -1.27962470, -1.35999286, 2.09867334,
0.65099514, 0.21604492, -0.45951018, 0.15994427, -0.31420693,
-0.65202618, -0.61077976, -1.06100249, -1.47254968, 1.18165290,
-0.78656220, 1.28182006, 1.80323684, 1.09196901, 0.26118696,
-0.30168581, 0.39749333, 0.26812574, -1.51995814, -0.46909946,
0.03874255, -1.36774313, 2.30143976, 2.06959820, -0.41647521,
1.85624206, 0.49019700, -0.06726539, 0.00457313, 0.23915423,
-1.84971249, -0.20482327, -0.34097880, -0.57933033, -1.10541213,
-0.30269983, -0.16430426, -0.82371718, 0.10345812, 1.78753936,
0.04786763, 1.86778629, -0.65214992, 0.81544143, -0.28214937,
0.31187257, 0.57661986, 1.21938801, -1.56046617, 0.38046429,
-0.18235965, 0.81794524, -0.40474343, 0.46538028, -1.15558851,
0.59625793, -1.07801270, 0.07310858, 0.61526084, 0.55518496,
-0.49787554, 0.92703879, -1.27780271, -0.83373469, -0.43015575,
0.41877759, -1.03987372, -1.46055734, 0.61282396, 0.15590595,
-0.34269521, 0.56509072, -1.17904210, 0.11374855, -1.83310866,
0.38734794, -0.58623004, 0.77931106, 1.53930688, -0.70299625,
-0.11389336, -1.14818096, -0.44400632, 1.21887410, 0.64066756,
-0.70249403, -0.27244881, 0.38586098, -1.07925785, 0.12448707,
-1.28286278, 0.37827531, 0.68812364, 1.65695465, 0.12440517,
-0.03689830, 1.10224664, -0.28323629, -0.47939169, 0.70120829,
-0.67204583
]
_BIAS = [
0.00465965, 0.21738243, 0.22277749, 0.68463874, 0.84596848, 0.17337036,
0.39573753, 0.78153563, 0.86311185, 0.21782327, 0.24377882, 0.42310703,
0.19887352, 0.10486019, 0.48707581, 0.22205460, 0.97263455, 0.29714966,
0.11244559, 0.53020525, 0.36796236, 0.37294638, 0.80261672, 0.04669094,
0.86319661, 0.75907171, 0.77297020, 0.01114798, 0.55850804, 0.91799915,
0.23032320, 0.12154722, 0.26701927, 0.42934716, 0.47951782, 0.96782577,
0.86785042, 0.61985648, 0.05743814, 0.41800117, 0.68881893, 0.60575199,
0.21058667, 0.64412105, 0.63958526, 0.89390790, 0.69755554, 0.89345169,
0.53330755, 0.56985939, 0.30724049, 0.00984561, 0.91407037, 0.92118979,
0.94153070, 0.81097460, 0.70537627, 0.32810748, 0.47227263, 0.11821401,
0.44983089, 0.30767226, 0.31756389, 0.62969446, 0.69892538, 0.16949117,
0.06207097, 0.46717727, 0.95348179, 0.62363589, 0.49018729, 0.06920040,
0.39333904, 0.41299903, 0.52514863, 0.61197245, 0.56871891, 0.65053988,
0.22203422, 0.46748531, 0.86931503, 0.87050021, 0.40208721, 0.32084906,
0.55084610, 0.94584596, 0.76279902, 0.36250532, 0.74272907, 0.66682065,
0.96452832, 0.64768302, 0.88070846, 0.56995463, 0.06395614, 0.69499350,
0.44494808, 0.39775658, 0.20280898, 0.33363521, 0.05999005, 0.44414878,
0.65227020, 0.01199079, 0.71995056, 0.19045687, 0.48342144, 0.25127733,
0.66515994, 0.22465158, 0.22313106, 0.06302810, 0.55783665, 0.93625581,
0.58800840, 0.72525370, 0.52879298, 0.77195418, 0.15548682, 0.01028740,
0.39325142, 0.45401239, 0.71494079, 0.33011997, 0.05050695, 0.26381660,
0.63064706, 0.47604024, 0.08593416, 0.00383425, 0.06352687, 0.05510247,
0.03552997, 0.35810637, 0.56094289, 0.60922170, 0.88599777, 0.45419788,
0.40486634, 0.71297824, 0.34976673, 0.97825217, 0.12915993, 0.09566259,
0.64318919, 0.16717327, 0.82308614, 0.32672071, 0.81688786, 0.84857118,
0.99922776, 0.07551706, 0.18766022, 0.13051236, 0.39136350, 0.08768725,
0.92048228, 0.87185788, 0.39158428, 0.79224777, 0.17492688, 0.68902445,
0.81980729, 0.70458186, 0.59489477, 0.93324888, 0.49986637, 0.40705478,
0.89202917, 0.20673239, 0.39339757, 0.20996964, 0.02923799, 0.53992438,
0.40119815, 0.10366607, 0.08044600, 0.95551598, 0.20518017, 0.68826210,
0.90159297, 0.69008791, 0.86880815, 0.16246438, 0.89628279, 0.11481643,
0.61353648, 0.41545081, 0.92478311, 0.78212476, 0.48292696, 0.79621077,
0.11947489, 0.01747024, 0.22928023, 0.87387264, 0.86349785, 0.89526737,
0.58904779, 0.13896775, 0.68194926, 0.55824125, 0.44428205, 0.55422378,
0.28189969, 0.27923775, 0.09979951, 0.66994715, 0.45943546, 0.71207762,
0.17300689, 0.83434916, 0.02573085, 0.45858085, 0.55934799, 0.30676675,
0.52219367, 0.34544575, 0.19280875, 0.26937950, 0.07147646, 0.06295013,
0.76382887, 0.38737607, 0.58825982, 0.17423475, 0.05509448, 0.97228825,
0.94380617, 0.91664016, 0.18800116, 0.41771865, 0.59420645, 0.77371931,
0.64687788, 0.27284670, 0.22310913, 0.15663862, 0.45573199, 0.50386798,
0.66712272, 0.71649647, 0.28475654, 0.83415413, 0.75261366, 0.61517799,
0.93544555, 0.76141870, 0.85474241, 0.74766934, 0.33459592, 0.78477907,
0.07250881, 0.10174239, 0.95332730, 0.80793905
]
# pyformat: enable
# pylint: enable=bad-whitespace
# pylint: enable=bad-continuation
def noise_embeddings(sigma_scaled_noise_level: jnp.ndarray) -> jnp.ndarray:
"""Returns Fourier noise level embeddings for diffusion model."""
transformed_noise_level = (1 / 4) * jnp.log(sigma_scaled_noise_level)
weight = jnp.array(_WEIGHT, dtype=jnp.float32)
bias = jnp.array(_BIAS, dtype=jnp.float32)
embeddings = transformed_noise_level[..., None] * weight + bias
return jnp.cos(2 * jnp.pi * embeddings)
|