File size: 4,428 Bytes
2e1a430 | 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 | import torch
from einops import rearrange
LATENT_SHIFT = (
0.01984364, 0.10149707, 0.29689495, 0.27188619, -0.21445648, -0.15979549,
0.05021099, -0.15083604, -0.15360136, -0.20131799, 0.01922352, 0.0622626,
0.10140969, -0.06739428, 0.3758261, -0.233712, 0.35164491, -0.02590912,
-0.0271935, -0.10833897, -0.1476848, -0.01130957, -0.2298372, 0.23526423,
-0.10893522, 0.11957631, 0.04047799, 0.3134589, -0.17225064, -0.18646109,
-0.34691978, -0.03571246, 0.02583857, 0.10190072, 0.28402294, 0.26952152,
-0.21634675, -0.17938656, 0.04358909, -0.15007621, -0.1548502, -0.18971131,
0.02710861, 0.05609494, 0.10697846, -0.06854968, 0.38167698, -0.24269937,
0.35705471, -0.03063305, -0.02946109, -0.11244286, -0.14336038, -0.01362137,
-0.21863696, 0.23228983, -0.11739769, 0.11693044, 0.02563311, 0.31356594,
-0.17420591, -0.19006285, -0.34905377, -0.04025005, 0.01924137, 0.07652984,
0.2995608, 0.2628057, -0.22011674, -0.12715361, 0.04879879, -0.14075719,
-0.15935895, -0.2123584, 0.01974813, 0.05523547, 0.10011992, -0.06428964,
0.37781868, -0.21491644, 0.34254215, -0.03153528, -0.0310082, -0.10761415,
-0.14730405, -0.02475182, -0.2285588, 0.2515081, -0.10445128, 0.12446,
0.07062869, 0.30880162, -0.18016875, -0.18869164, -0.34533499, -0.0129177,
0.02578168, 0.07993659, 0.28642181, 0.26038408, -0.22459419, -0.14820155,
0.04059549, -0.14043529, -0.16111187, -0.2020305, 0.02602069, 0.04852717,
0.10432153, -0.06309942, 0.38402443, -0.22397003, 0.34814481, -0.03774432,
-0.03381438, -0.11245691, -0.14128767, -0.02853208, -0.21752016, 0.24872463,
-0.11399775, 0.1222687, 0.05620835, 0.309178, -0.18065738, -0.19401479,
-0.34495114, -0.01760592,
)
LATENT_SCALE = (
1.63933691, 1.70204478, 1.73642566, 1.90004803, 1.6675316, 1.69059584,
1.56853198, 1.62314944, 1.89106626, 1.58086668, 1.60822129, 1.60962993,
1.63322129, 1.56074359, 1.73419528, 1.7919265, 1.64040632, 1.66802808,
1.60390303, 1.75480492, 1.63187587, 1.64334594, 1.61722884, 1.60146046,
1.63459219, 1.55291476, 1.68771497, 1.68415657, 1.78966054, 1.66631641,
1.65626686, 1.65976433, 1.63487607, 1.69513249, 1.72933756, 1.91310663,
1.67035057, 1.72286863, 1.56719251, 1.61934825, 1.88628859, 1.56911539,
1.59455129, 1.60829869, 1.62470611, 1.56052853, 1.73677003, 1.77563606,
1.63732541, 1.66370527, 1.59508952, 1.75153949, 1.63029275, 1.64517667,
1.61659342, 1.59722044, 1.64103121, 1.5408531, 1.68610394, 1.67772755,
1.78998563, 1.66621713, 1.65458955, 1.66041308, 1.64710857, 1.68163503,
1.74000294, 1.92784786, 1.67411194, 1.67395548, 1.57406532, 1.62199356,
1.87618195, 1.5584375, 1.57438785, 1.61711053, 1.63094305, 1.55644029,
1.73124302, 1.80666627, 1.6463621, 1.65932006, 1.60816188, 1.75682671,
1.64695873, 1.63121722, 1.61380832, 1.60478651, 1.63396035, 1.53505068,
1.65534289, 1.67132281, 1.80317197, 1.6767314, 1.65700938, 1.68426259,
1.65339716, 1.67540638, 1.73298504, 1.94067348, 1.67893609, 1.70635117,
1.5730906, 1.61928553, 1.87148809, 1.56244866, 1.56697152, 1.61584394,
1.62759496, 1.55480378, 1.73484107, 1.79055143, 1.64688773, 1.66121492,
1.60135887, 1.75254572, 1.64798332, 1.62989921, 1.61381592, 1.60792883,
1.63939668, 1.53075757, 1.65371318, 1.66801185, 1.80029087, 1.67591476,
1.65655173, 1.68533454,
)
def get_latent_norm(device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
shift = torch.tensor(LATENT_SHIFT, dtype=torch.float32, device=device)
scale = torch.tensor(LATENT_SCALE, dtype=torch.float32, device=device)
return shift, scale
def decode(vae, latents, height, width, torch_dtype):
latent_shift, latent_scale = get_latent_norm(latents.device)
latents = latents.float() * latent_scale + latent_shift
latents = rearrange(latents, "B (H W) (P Q C) -> B C (H P) (W Q)", P=2, Q=2, H=height//16, W=width//16).to(torch.bfloat16)
latents = latents.to(torch_dtype)
image = vae._decode(latents)
return image
def encode(vae, image, height, width, torch_dtype):
latents = vae._encode(image)[:, :32]
latent_shift, latent_scale = get_latent_norm(latents.device)
latents = rearrange(latents, "B C (H P) (W Q) -> B (H W) (P Q C)", P=2, Q=2, H=height//16, W=width//16).to(torch.bfloat16)
latents = (latents.float() - latent_shift) / latent_scale
latents = latents.to(torch_dtype)
return latents
|