| 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 |
|
|
