27 lines
640 B
Python
27 lines
640 B
Python
from einops import rearrange
|
|
|
|
|
|
def transform_to_2d(x):
|
|
b = x.shape[0]
|
|
return rearrange(x, 'b c f h w -> (b f) c h w'), b
|
|
|
|
|
|
def transformed_to_3d(x, b):
|
|
return rearrange(x, '(b f) c h w -> b c f h w', b=b)
|
|
|
|
|
|
def apply_unet_patch3d(h, hs, transformer_options, patch):
|
|
h, bh = transform_to_2d(h)
|
|
hs, bhs = transform_to_2d(hs)
|
|
h, hs = patch(h, hs, transformer_options)
|
|
h = transformed_to_3d(h, bh)
|
|
hs = transformed_to_3d(hs, bhs)
|
|
return h, hs
|
|
|
|
|
|
def apply_patch3d(x, transformer_options, patch):
|
|
x, b = transform_to_2d(x)
|
|
x = patch(x, transformer_options)
|
|
x = transformed_to_3d(x, b)
|
|
return x
|