Function make_attn

diff2flow/kl_autoencoder.py:231–244

(in_channels, attn_type="vanilla", natten_kernel_size=-1, use_null_attention=False)

def make_attn(in_channels, attn_type="vanilla", natten_kernel_size=-1, use_null_attention=False):
    if attn_type == "vanilla":
        assert not (use_null_attention and natten_kernel_size > -1), "use_null_attention and natten_kernel_size > -1 are mutually exclusive"
        if natten_kernel_size > -1 and not NATTEN_IS_AVAILBLE:
            raise ValueError("natten_kernel_size > -1 but natten is not available")
        if use_null_attention:
            print(f"Using null attention to save memory and compute...")
        return MemoryEfficientAttnBlock(in_channels, natten_kernel_size=natten_kernel_size, use_null_attention=use_null_attention)
    
    elif attn_type == "none":
        return nn.Identity(in_channels)
    
    else:
        raise NotImplementedError(f"attn_type {attn_type} not implemented")


""" Encoder and Decoder """