23 lines
771 B
Python
23 lines
771 B
Python
|
|
def patchify_latents(latents):
|
|
batch_size, num_channels_latents, height, width = latents.shape
|
|
latents = latents.view(
|
|
batch_size, num_channels_latents, height // 2, 2, width // 2, 2
|
|
)
|
|
latents = latents.permute(0, 1, 3, 5, 2, 4)
|
|
latents = latents.reshape(
|
|
batch_size, num_channels_latents * 4, height // 2, width // 2
|
|
)
|
|
return latents
|
|
|
|
|
|
def unpatchify_latents(latents):
|
|
batch_size, num_channels_latents, height, width = latents.shape
|
|
latents = latents.reshape(
|
|
batch_size, num_channels_latents // (2 * 2), 2, 2, height, width
|
|
)
|
|
latents = latents.permute(0, 1, 4, 2, 5, 3)
|
|
latents = latents.reshape(
|
|
batch_size, num_channels_latents // (2 * 2), height * 2, width * 2
|
|
)
|
|
return latents |