(
*,
rgb_nchw: torch.Tensor,
vae: SupportedAutoencoder,
)
| 117 | |
| 118 | |
| 119 | def to_latent( |
| 120 | *, |
| 121 | rgb_nchw: torch.Tensor, |
| 122 | vae: SupportedAutoencoder, |
| 123 | ) -> torch.Tensor: |
| 124 | rgb_nchw = VaeImageProcessor.normalize(rgb_nchw) # type: ignore |
| 125 | encoding_nchw = vae.encode(typing.cast(torch.FloatTensor, rgb_nchw)) |
| 126 | if isinstance(encoding_nchw, AutoencoderKLOutput): |
| 127 | latent = encoding_nchw.latent_dist.sample() # type: ignore |
| 128 | assert isinstance(latent, torch.Tensor) |
| 129 | elif isinstance(encoding_nchw, AutoencoderTinyOutput): |
| 130 | latent = encoding_nchw.latents |
| 131 | do_internal_vae_scaling = False # Is this needed? |
| 132 | if do_internal_vae_scaling: |
| 133 | latent = vae.scale_latents(latent).mul(255).round().byte() # type: ignore |
| 134 | latent = vae.unscale_latents(latent / 255.0) # type: ignore |
| 135 | assert isinstance(latent, torch.Tensor) |
| 136 | else: |
| 137 | assert False, f"Unknown encoding type: {type(encoding_nchw)}" |
| 138 | return latent |
| 139 | |
| 140 | |
| 141 | def from_latent( |
no test coverage detected