File size: 660 Bytes
59630ba | 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 | from typing import Callable
import torch
from einops import rearrange
def video_to_image(func: Callable) -> Callable:
def wrapper(self, x: torch.Tensor, *args, **kwargs):
if x.dim() == 5:
t = x.shape[2]
x = rearrange(x, "b c t h w -> (b t) c h w")
x = func(self, x, *args, **kwargs)
x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
else:
x = func(self, x, *args, **kwargs)
return x
return wrapper
def nonlinearity(x):
return x * torch.sigmoid(x)
def cast_tuple(t, length=1):
return t if isinstance(t, tuple) or isinstance(t, list) else ((t,) * length)
|