227 lines
10 KiB
Python
227 lines
10 KiB
Python
# Based on code from ComfyUI-Florence2 by Jukka Seppänen (kijai), MIT License
|
|
"""DaViT (Dual-Attention Vision Transformer) for Florence2. ComfyUI ops pattern."""
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from collections import OrderedDict
|
|
from comfy.ldm.modules.attention import optimized_attention_for_device
|
|
|
|
|
|
class MySequential(nn.Sequential):
|
|
def forward(self, *inputs):
|
|
for module in self._modules.values():
|
|
if isinstance(inputs, tuple):
|
|
inputs = module(*inputs)
|
|
else:
|
|
inputs = module(inputs)
|
|
return inputs
|
|
|
|
|
|
class PreNorm(nn.Module):
|
|
"""Pre-norm wrapper with residual. Attrs named norm/fn to match weight keys."""
|
|
def __init__(self, norm, fn):
|
|
super().__init__()
|
|
self.norm = norm
|
|
self.fn = fn
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
shortcut = x
|
|
if self.norm is not None:
|
|
x, size = self.fn(self.norm(x), *args, **kwargs)
|
|
else:
|
|
x, size = self.fn(x, *args, **kwargs)
|
|
return shortcut + x, size
|
|
|
|
|
|
class Mlp(nn.Module):
|
|
"""MLP with net.fc1/net.fc2 weight keys via OrderedDict Sequential."""
|
|
def __init__(self, in_features, hidden_features=None, out_features=None, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
out_features = out_features or in_features
|
|
hidden_features = hidden_features or in_features
|
|
self.net = nn.Sequential(OrderedDict([
|
|
("fc1", operations.Linear(in_features, hidden_features, dtype=dtype, device=device)),
|
|
("act", nn.GELU()),
|
|
("fc2", operations.Linear(hidden_features, out_features, dtype=dtype, device=device)),
|
|
]))
|
|
|
|
def forward(self, x, size):
|
|
return self.net(x), size
|
|
|
|
|
|
class DepthWiseConv2d(nn.Module):
|
|
def __init__(self, dim_in, kernel_size, padding, stride, bias=True, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
self.dw = operations.Conv2d(dim_in, dim_in, kernel_size=kernel_size, padding=padding, groups=dim_in, stride=stride, bias=bias, dtype=dtype, device=device)
|
|
|
|
def forward(self, x, size):
|
|
B, N, C = x.shape
|
|
H, W = size
|
|
x = self.dw(x.transpose(1, 2).view(B, C, H, W))
|
|
size = (x.size(-2), x.size(-1))
|
|
return x.flatten(2).transpose(1, 2), size
|
|
|
|
|
|
class ConvEmbed(nn.Module):
|
|
"""Patch embedding via Conv2d + optional LayerNorm."""
|
|
def __init__(self, patch_size=7, in_chans=3, embed_dim=64, stride=4, padding=2, norm_layer_on=True, pre_norm=True, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.pre_norm = pre_norm
|
|
self.proj = operations.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride, padding=padding, dtype=dtype, device=device)
|
|
if norm_layer_on:
|
|
self.norm = operations.LayerNorm(in_chans if pre_norm else embed_dim, dtype=dtype, device=device)
|
|
else:
|
|
self.norm = None
|
|
|
|
def forward(self, x, size):
|
|
H, W = size
|
|
if len(x.size()) == 3:
|
|
if self.norm and self.pre_norm:
|
|
x = self.norm(x)
|
|
B, _, C = x.shape
|
|
x = x.view(B, H, W, C).permute(0, 3, 1, 2)
|
|
x = self.proj(x)
|
|
_, _, H, W = x.shape
|
|
x = x.flatten(2).transpose(1, 2)
|
|
if self.norm and not self.pre_norm:
|
|
x = self.norm(x)
|
|
return x, (H, W)
|
|
|
|
|
|
def window_partition(x, window_size):
|
|
B, H, W, C = x.shape
|
|
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
|
|
return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
|
|
|
|
|
|
def window_reverse(windows, batch_size, window_size, H, W):
|
|
x = windows.view(batch_size, H // window_size, W // window_size, window_size, window_size, -1)
|
|
return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(batch_size, H, W, -1)
|
|
|
|
|
|
class WindowAttention(nn.Module):
|
|
def __init__(self, dim, num_heads, window_size, qkv_bias=True, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.window_size = window_size
|
|
self.num_heads = num_heads
|
|
self.qkv = operations.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
|
|
self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
|
|
|
|
def forward(self, x, size):
|
|
H, W = size
|
|
B, L, C = x.shape
|
|
x = x.view(B, H, W, C)
|
|
|
|
pad_r = (self.window_size - W % self.window_size) % self.window_size
|
|
pad_b = (self.window_size - H % self.window_size) % self.window_size
|
|
x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
|
|
_, Hp, Wp, _ = x.shape
|
|
|
|
x = window_partition(x, self.window_size).view(-1, self.window_size * self.window_size, C)
|
|
|
|
B_, N, C = x.shape
|
|
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv[0], qkv[1], qkv[2]
|
|
optimized_attention = optimized_attention_for_device(x.device, small_input=True)
|
|
x = optimized_attention(q, k, v, self.num_heads, skip_reshape=True)
|
|
x = self.proj(x)
|
|
|
|
x = window_reverse(x.view(-1, self.window_size, self.window_size, C), B, self.window_size, Hp, Wp)
|
|
if pad_r > 0 or pad_b > 0:
|
|
x = x[:, :H, :W, :].contiguous()
|
|
return x.view(B, H * W, C), size
|
|
|
|
|
|
class ChannelAttention(nn.Module):
|
|
def __init__(self, dim, groups=8, qkv_bias=True, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
self.groups = groups
|
|
self.qkv = operations.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
|
|
self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
|
|
|
|
def forward(self, x, size):
|
|
B, N, C = x.shape
|
|
qkv = self.qkv(x).reshape(B, N, 3, self.groups, C // self.groups).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv[0], qkv[1], qkv[2]
|
|
q = q * (float(N) ** -0.5)
|
|
attn = (q.transpose(-1, -2) @ k).softmax(dim=-1)
|
|
x = (attn @ v.transpose(-1, -2)).transpose(-1, -2).transpose(1, 2).reshape(B, N, C)
|
|
return self.proj(x), size
|
|
|
|
|
|
class SpatialBlock(nn.Module):
|
|
def __init__(self, dim, num_heads, window_size, mlp_ratio=4., qkv_bias=True, conv_at_attn=True, conv_at_ffn=True, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1, dtype=dtype, device=device, operations=operations)) if conv_at_attn else None
|
|
self.window_attn = PreNorm(operations.LayerNorm(dim, dtype=dtype, device=device), WindowAttention(dim, num_heads, window_size, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations))
|
|
self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1, dtype=dtype, device=device, operations=operations)) if conv_at_ffn else None
|
|
self.ffn = PreNorm(operations.LayerNorm(dim, dtype=dtype, device=device), Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), dtype=dtype, device=device, operations=operations))
|
|
|
|
def forward(self, x, size):
|
|
if self.conv1:
|
|
x, size = self.conv1(x, size)
|
|
x, size = self.window_attn(x, size)
|
|
if self.conv2:
|
|
x, size = self.conv2(x, size)
|
|
x, size = self.ffn(x, size)
|
|
return x, size
|
|
|
|
|
|
class ChannelBlock(nn.Module):
|
|
def __init__(self, dim, groups, mlp_ratio=4., qkv_bias=True, conv_at_attn=True, conv_at_ffn=True, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1, dtype=dtype, device=device, operations=operations)) if conv_at_attn else None
|
|
self.channel_attn = PreNorm(operations.LayerNorm(dim, dtype=dtype, device=device), ChannelAttention(dim, groups=groups, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations))
|
|
self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1, dtype=dtype, device=device, operations=operations)) if conv_at_ffn else None
|
|
self.ffn = PreNorm(operations.LayerNorm(dim, dtype=dtype, device=device), Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), dtype=dtype, device=device, operations=operations))
|
|
|
|
def forward(self, x, size):
|
|
if self.conv1:
|
|
x, size = self.conv1(x, size)
|
|
x, size = self.channel_attn(x, size)
|
|
if self.conv2:
|
|
x, size = self.conv2(x, size)
|
|
x, size = self.ffn(x, size)
|
|
return x, size
|
|
|
|
|
|
class DaViT(nn.Module):
|
|
"""DaViT: Dual-Attention Vision Transformer. [B, 3, H, W] -> [B, N, C]"""
|
|
def __init__(self, config, dtype=None, device=None, operations=None):
|
|
super().__init__()
|
|
embed_dims = config.dim_embed
|
|
num_heads = config.num_heads
|
|
num_groups = config.num_groups
|
|
depths = config.depths
|
|
self.embed_dims = embed_dims
|
|
|
|
convs, blocks = [], []
|
|
for i in range(len(embed_dims)):
|
|
convs.append(ConvEmbed(
|
|
patch_size=config.patch_size[i], in_chans=3 if i == 0 else embed_dims[i - 1],
|
|
embed_dim=embed_dims[i], stride=config.patch_stride[i], padding=config.patch_padding[i],
|
|
pre_norm=config.patch_prenorm[i], dtype=dtype, device=device, operations=operations,
|
|
))
|
|
blocks.append(MySequential(*[
|
|
MySequential(OrderedDict([
|
|
('spatial_block', SpatialBlock(embed_dims[i], num_heads[i], config.window_size, dtype=dtype, device=device, operations=operations)),
|
|
('channel_block', ChannelBlock(embed_dims[i], num_groups[i], dtype=dtype, device=device, operations=operations)),
|
|
])) for j in range(depths[i])
|
|
]))
|
|
self.convs = nn.ModuleList(convs)
|
|
self.blocks = nn.ModuleList(blocks)
|
|
|
|
@property
|
|
def dim_out(self):
|
|
return self.embed_dims[-1]
|
|
|
|
def forward_features_unpool(self, x):
|
|
"""[B, 3, H, W] -> [B, N, C]"""
|
|
input_size = (x.size(2), x.size(3))
|
|
for conv, block in zip(self.convs, self.blocks):
|
|
x, input_size = conv(x, input_size)
|
|
x, input_size = block(x, input_size)
|
|
return x
|