Function TemporalTransformer_forward

motionctrl/lvdm_modified_modules.py:19–69

(self, x, context=None, is_imgbatch=False)

def TemporalTransformer_forward(self, x, context=None, is_imgbatch=False):
    b, c, t, h, w = x.shape
    x_in = x
    x = self.norm(x)
    x = rearrange(x, 'b c t h w -> (b h w) c t').contiguous()
    if not self.use_linear:
        x = self.proj_in(x)
    x = rearrange(x, 'bhw c t -> bhw t c').contiguous()
    if self.use_linear:
        x = self.proj_in(x)

    temp_mask = None
    if self.causal_attention:
        temp_mask = torch.tril(torch.ones([1, t, t]))
    if is_imgbatch:
        temp_mask = torch.eye(t).unsqueeze(0)
    if temp_mask is not None:
        mask = temp_mask.to(x.device)
        mask = repeat(mask, 'l i j -> (l bhw) i j', bhw=b*h*w)
    else:
        mask = None

    if self.only_self_att:
        ## note: if no context is given, cross-attention defaults to self-attention
        for i, block in enumerate(self.transformer_blocks):
            x = block(x, context=context, mask=mask)
        x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous()
    else:
        x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous()
        context = rearrange(context, '(b t) l con -> b t l con', t=t).contiguous()
        for i, block in enumerate(self.transformer_blocks):
            # calculate each batch one by one (since number in shape could not greater then 65,535 for some package)
            for j in range(b):
                unit_context = context[j][0:1]
                context_j = repeat(unit_context, 't l con -> (t r) l con', r=(h * w)).contiguous()
                ## note: causal mask will not applied in cross-attention case
                x[j] = block(x[j], context=context_j)
    
    if self.use_linear:
        x = self.proj_out(x)
        x = rearrange(x, 'b (h w) t c -> b c t h w', h=h, w=w).contiguous()
    if not self.use_linear:
        x = rearrange(x, 'b hw t c -> (b hw) c t').contiguous()
        x = self.proj_out(x)
        x = rearrange(x, '(b h w) c t -> b c t h w', b=b, h=h, w=w).contiguous()

    if self.use_image_dataset:
        x = 0.0 * x + x_in
    else:
        x = x + x_in
    return x

def selfattn_forward_unet(self, x, timesteps, context=None, y=None, features_adapter=None, is_imgbatch=False, T=None,  **kwargs):
        b,_,t,_,_ = x.shape