Add files via upload
This commit is contained in:
AI Lab
@@ -0,0 +1,164 @@
|
||||
# Based on code from ComfyUI-Florence2 by Jukka Seppänen (kijai), MIT License
|
||||
import json
|
||||
from dataclasses import dataclass, field, fields
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class VisionConfig:
|
||||
drop_path_rate: float = 0.1
|
||||
patch_size: List[int] = field(default_factory=lambda: [7, 3, 3, 3])
|
||||
patch_stride: List[int] = field(default_factory=lambda: [4, 2, 2, 2])
|
||||
patch_padding: List[int] = field(default_factory=lambda: [3, 1, 1, 1])
|
||||
patch_prenorm: List[bool] = field(default_factory=lambda: [False, True, True, True])
|
||||
dim_embed: List[int] = field(default_factory=lambda: [256, 512, 1024, 2048])
|
||||
num_heads: List[int] = field(default_factory=lambda: [8, 16, 32, 64])
|
||||
num_groups: List[int] = field(default_factory=lambda: [8, 16, 32, 64])
|
||||
depths: List[int] = field(default_factory=lambda: [1, 1, 9, 1])
|
||||
window_size: int = 12
|
||||
projection_dim: int = 1024
|
||||
visual_temporal_embedding: Optional[Dict] = None
|
||||
image_pos_embed: Optional[Dict] = None
|
||||
image_feature_source: List[str] = field(default_factory=lambda: ["spatial_avg_pool", "temporal_avg_pool"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class LanguageConfig:
|
||||
vocab_size: int = 51289
|
||||
max_position_embeddings: int = 1024
|
||||
encoder_layers: int = 12
|
||||
encoder_ffn_dim: int = 4096
|
||||
encoder_attention_heads: int = 16
|
||||
decoder_layers: int = 12
|
||||
decoder_ffn_dim: int = 4096
|
||||
decoder_attention_heads: int = 16
|
||||
activation_function: str = "gelu"
|
||||
d_model: int = 1024
|
||||
scale_embedding: bool = False
|
||||
pad_token_id: int = 1
|
||||
bos_token_id: int = 0
|
||||
eos_token_id: int = 2
|
||||
decoder_start_token_id: int = 2
|
||||
forced_bos_token_id: int = 0
|
||||
forced_eos_token_id: int = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class Florence2Config:
|
||||
vision_config: VisionConfig = field(default_factory=VisionConfig)
|
||||
text_config: LanguageConfig = field(default_factory=LanguageConfig)
|
||||
vocab_size: int = 51289
|
||||
projection_dim: int = 1024
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, path):
|
||||
with open(path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
vision_data = data.get("vision_config", {})
|
||||
text_data = data.get("text_config", {})
|
||||
|
||||
vision_fields = {f.name for f in fields(VisionConfig)}
|
||||
text_fields = {f.name for f in fields(LanguageConfig)}
|
||||
|
||||
vision_config = VisionConfig(**{k: v for k, v in vision_data.items() if k in vision_fields})
|
||||
text_config = LanguageConfig(**{k: v for k, v in text_data.items() if k in text_fields})
|
||||
|
||||
return cls(
|
||||
vision_config=vision_config,
|
||||
text_config=text_config,
|
||||
vocab_size=data.get("vocab_size", text_config.vocab_size),
|
||||
projection_dim=data.get("projection_dim", vision_config.projection_dim),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_state_dict(cls, sd):
|
||||
"""Detect base vs large from weight shapes."""
|
||||
q_key = "language_model.model.encoder.layers.0.self_attn.q_proj.weight"
|
||||
if q_key in sd:
|
||||
d_model = sd[q_key].shape[0]
|
||||
else:
|
||||
d_model = 1024
|
||||
|
||||
encoder_layer_keys = [k for k in sd if k.startswith("language_model.model.encoder.layers.")]
|
||||
if encoder_layer_keys:
|
||||
max_layer = max(int(k.split(".")[4]) for k in encoder_layer_keys)
|
||||
encoder_layers = max_layer + 1
|
||||
else:
|
||||
encoder_layers = 12
|
||||
|
||||
decoder_layer_keys = [k for k in sd if k.startswith("language_model.model.decoder.layers.")]
|
||||
if decoder_layer_keys:
|
||||
max_layer = max(int(k.split(".")[4]) for k in decoder_layer_keys)
|
||||
decoder_layers = max_layer + 1
|
||||
else:
|
||||
decoder_layers = 12
|
||||
|
||||
shared_key = "language_model.model.shared.weight"
|
||||
vocab_size = sd[shared_key].shape[0] if shared_key in sd else 51289
|
||||
|
||||
fc1_key = "language_model.model.encoder.layers.0.fc1.weight"
|
||||
encoder_ffn_dim = sd[fc1_key].shape[0] if fc1_key in sd else d_model * 4
|
||||
|
||||
dec_fc1_key = "language_model.model.decoder.layers.0.fc1.weight"
|
||||
decoder_ffn_dim = sd[dec_fc1_key].shape[0] if dec_fc1_key in sd else d_model * 4
|
||||
|
||||
# Detect vision config from state dict
|
||||
dim_embed = []
|
||||
for i in range(4):
|
||||
conv_key = f"vision_tower.convs.{i}.proj.weight"
|
||||
if conv_key in sd:
|
||||
dim_embed.append(sd[conv_key].shape[0])
|
||||
|
||||
if not dim_embed:
|
||||
dim_embed = [256, 512, 1024, 2048]
|
||||
|
||||
num_heads = [d // 32 for d in dim_embed] # head_dim=32 for DaViT spatial attention
|
||||
|
||||
# Count depths
|
||||
depths = []
|
||||
for stage in range(4):
|
||||
stage_keys = [k for k in sd if k.startswith(f"vision_tower.blocks.{stage}.")]
|
||||
if stage_keys:
|
||||
max_depth = max(int(k.split(".")[3]) for k in stage_keys) + 1
|
||||
depths.append(max_depth)
|
||||
else:
|
||||
depths.append(1)
|
||||
|
||||
if not any(k.startswith("vision_tower.") for k in sd):
|
||||
depths = [1, 1, 9, 1]
|
||||
|
||||
projection_dim = dim_embed[-1] if dim_embed else 1024
|
||||
proj_key = "image_projection"
|
||||
if proj_key in sd:
|
||||
projection_dim = sd[proj_key].shape[1]
|
||||
|
||||
vision_config = VisionConfig(
|
||||
dim_embed=dim_embed,
|
||||
num_heads=num_heads,
|
||||
num_groups=num_heads,
|
||||
depths=depths,
|
||||
projection_dim=projection_dim,
|
||||
image_pos_embed={"type": "learned_abs_2d", "max_pos_embeddings": 50},
|
||||
visual_temporal_embedding={"type": "COSINE", "max_temporal_embeddings": 100},
|
||||
)
|
||||
|
||||
text_config = LanguageConfig(
|
||||
vocab_size=vocab_size,
|
||||
d_model=d_model,
|
||||
encoder_layers=encoder_layers,
|
||||
decoder_layers=decoder_layers,
|
||||
encoder_attention_heads=d_model // 64,
|
||||
decoder_attention_heads=d_model // 64,
|
||||
encoder_ffn_dim=encoder_ffn_dim,
|
||||
decoder_ffn_dim=decoder_ffn_dim,
|
||||
max_position_embeddings=1024,
|
||||
scale_embedding=False,
|
||||
)
|
||||
|
||||
return cls(
|
||||
vision_config=vision_config,
|
||||
text_config=text_config,
|
||||
vocab_size=vocab_size,
|
||||
projection_dim=projection_dim,
|
||||
)
|
||||
@@ -0,0 +1,226 @@
|
||||
# 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
|
||||
@@ -0,0 +1,579 @@
|
||||
# Based on code from ComfyUI-Florence2 by Jukka Seppänen (kijai), MIT License
|
||||
"""BART encoder-decoder language model, Florence2 top-level model, and autoregressive generation."""
|
||||
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .davit import DaViT
|
||||
import comfy.ops
|
||||
from comfy.ldm.modules.attention import optimized_attention_for_device
|
||||
from comfy.utils import ProgressBar
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
_POSITION_OFFSET = 2 # BART positional id offset
|
||||
|
||||
|
||||
def make_attention_mask(mask_2d, dtype):
|
||||
"""[B, S] -> [B, 1, 1, S] with -inf for padding positions."""
|
||||
inverted = 1.0 - mask_2d[:, None, None, :].to(dtype)
|
||||
return inverted.masked_fill(inverted.bool(), torch.finfo(dtype).min)
|
||||
|
||||
|
||||
def make_causal_mask(shape, dtype, device, past_kv_len=0):
|
||||
"""Create [B, 1, T, T+past] causal mask with -inf above diagonal."""
|
||||
bsz, tgt_len = shape
|
||||
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device, dtype=dtype)
|
||||
cond = torch.arange(tgt_len, device=device)
|
||||
mask.masked_fill_(cond < (cond + 1).view(-1, 1), 0)
|
||||
if past_kv_len > 0:
|
||||
mask = torch.cat([torch.zeros(tgt_len, past_kv_len, dtype=dtype, device=device), mask], dim=-1)
|
||||
return mask[None, None].expand(bsz, 1, tgt_len, tgt_len + past_kv_len)
|
||||
|
||||
|
||||
def _make_position_ids(input_ids, past_key_values_length, device):
|
||||
bsz, seq_len = input_ids.shape[:2]
|
||||
positions = torch.arange(past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=device).expand(bsz, -1)
|
||||
return positions + _POSITION_OFFSET
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""Multi-headed attention with optional KV cache for decoder layers."""
|
||||
|
||||
def __init__(self, embed_dim, num_heads, is_decoder, is_causal, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = embed_dim // num_heads
|
||||
self.is_decoder = is_decoder
|
||||
self.is_causal = is_causal
|
||||
|
||||
self.q_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||
self.k_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||
self.v_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||
self.out_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||
|
||||
def _shape(self, tensor, seq_len, bsz):
|
||||
return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
|
||||
|
||||
def forward(self, hidden_states, key_value_states=None, past_key_value=None, attention_mask=None):
|
||||
is_cross_attention = key_value_states is not None
|
||||
bsz, tgt_len, _ = hidden_states.size()
|
||||
|
||||
query_states = self._shape(self.q_proj(hidden_states), tgt_len, bsz)
|
||||
|
||||
if is_cross_attention and past_key_value is not None and past_key_value[0].shape[2] == key_value_states.shape[1]:
|
||||
key_states = past_key_value[0]
|
||||
value_states = past_key_value[1]
|
||||
elif is_cross_attention:
|
||||
key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
|
||||
value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
|
||||
elif past_key_value is not None:
|
||||
key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
|
||||
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
|
||||
key_states = torch.cat([past_key_value[0], key_states], dim=2)
|
||||
value_states = torch.cat([past_key_value[1], value_states], dim=2)
|
||||
else:
|
||||
key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
|
||||
value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
|
||||
|
||||
if self.is_decoder:
|
||||
past_key_value = (key_states, value_states)
|
||||
|
||||
# Ensure matching dtypes (cross-attention cache may differ when using manual_cast ops)
|
||||
if key_states.dtype != query_states.dtype:
|
||||
key_states = key_states.to(query_states.dtype)
|
||||
value_states = value_states.to(query_states.dtype)
|
||||
|
||||
optimized_attention = optimized_attention_for_device(hidden_states.device, mask=attention_mask is not None, small_input=True)
|
||||
attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, mask=attention_mask, skip_reshape=True, skip_output_reshape=True)
|
||||
|
||||
attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.embed_dim)
|
||||
attn_output = self.out_proj(attn_output)
|
||||
return attn_output, past_key_value
|
||||
|
||||
|
||||
class EncoderLayer(nn.Module):
|
||||
def __init__(self, config, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.self_attn = Attention(d, config.encoder_attention_heads, is_decoder=False, is_causal=False, dtype=dtype, device=device, operations=operations)
|
||||
self.self_attn_layer_norm = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
self.fc1 = operations.Linear(d, config.encoder_ffn_dim, dtype=dtype, device=device)
|
||||
self.fc2 = operations.Linear(config.encoder_ffn_dim, d, dtype=dtype, device=device)
|
||||
self.final_layer_norm = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, hidden_states, attention_mask=None):
|
||||
residual = hidden_states
|
||||
hidden_states, _ = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = F.gelu(self.fc1(hidden_states))
|
||||
hidden_states = self.fc2(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states = self.final_layer_norm(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, config, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.self_attn = Attention(d, config.decoder_attention_heads, is_decoder=True, is_causal=True, dtype=dtype, device=device, operations=operations)
|
||||
self.self_attn_layer_norm = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
self.encoder_attn = Attention(d, config.decoder_attention_heads, is_decoder=True, is_causal=False, dtype=dtype, device=device, operations=operations)
|
||||
self.encoder_attn_layer_norm = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
self.fc1 = operations.Linear(d, config.decoder_ffn_dim, dtype=dtype, device=device)
|
||||
self.fc2 = operations.Linear(config.decoder_ffn_dim, d, dtype=dtype, device=device)
|
||||
self.final_layer_norm = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, hidden_states, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, past_key_value=None):
|
||||
# Self-attention
|
||||
self_attn_past = past_key_value[:2] if past_key_value is not None else None
|
||||
residual = hidden_states
|
||||
hidden_states, self_attn_cache = self.self_attn(hidden_states=hidden_states, past_key_value=self_attn_past, attention_mask=attention_mask)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||||
|
||||
# Cross-attention
|
||||
cross_attn_cache = None
|
||||
if encoder_hidden_states is not None:
|
||||
cross_attn_past = past_key_value[2:] if past_key_value is not None else None
|
||||
residual = hidden_states
|
||||
hidden_states, cross_attn_cache = self.encoder_attn(
|
||||
hidden_states=hidden_states, key_value_states=encoder_hidden_states,
|
||||
past_key_value=cross_attn_past, attention_mask=encoder_attention_mask,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states = self.encoder_attn_layer_norm(hidden_states)
|
||||
|
||||
# FFN
|
||||
residual = hidden_states
|
||||
hidden_states = F.gelu(self.fc1(hidden_states))
|
||||
hidden_states = self.fc2(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states = self.final_layer_norm(hidden_states)
|
||||
|
||||
present_key_value = None
|
||||
if self_attn_cache is not None:
|
||||
present_key_value = self_attn_cache
|
||||
if cross_attn_cache is not None:
|
||||
present_key_value = present_key_value + cross_attn_cache
|
||||
return hidden_states, present_key_value
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, config, embed_tokens=None, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.embed_scale = math.sqrt(d) if config.scale_embedding else 1.0
|
||||
self.embed_tokens = operations.Embedding(config.vocab_size, d, padding_idx=config.pad_token_id, dtype=dtype, device=device)
|
||||
if embed_tokens is not None:
|
||||
self.embed_tokens.weight = embed_tokens.weight
|
||||
self.embed_positions = operations.Embedding(config.max_position_embeddings + _POSITION_OFFSET, d, dtype=dtype, device=device)
|
||||
self.layers = nn.ModuleList([EncoderLayer(config, dtype=dtype, device=device, operations=operations) for _ in range(config.encoder_layers)])
|
||||
self.layernorm_embedding = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None):
|
||||
if input_ids is not None and inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
|
||||
|
||||
pos_input = input_ids if input_ids is not None else inputs_embeds[:, :, -1]
|
||||
position_ids = _make_position_ids(pos_input, 0, self.embed_positions.weight.device)
|
||||
embed_pos = self.embed_positions(position_ids).to(inputs_embeds.device)
|
||||
|
||||
hidden_states = self.layernorm_embedding(inputs_embeds + embed_pos)
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = make_attention_mask(attention_mask, hidden_states.dtype)
|
||||
|
||||
for layer in self.layers:
|
||||
hidden_states = layer(hidden_states, attention_mask=attention_mask)
|
||||
return (hidden_states,)
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, config, embed_tokens=None, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
d = config.d_model
|
||||
self.embed_scale = math.sqrt(d) if config.scale_embedding else 1.0
|
||||
self.embed_tokens = operations.Embedding(config.vocab_size, d, padding_idx=config.pad_token_id, dtype=dtype, device=device)
|
||||
if embed_tokens is not None:
|
||||
self.embed_tokens.weight = embed_tokens.weight
|
||||
self.embed_positions = operations.Embedding(config.max_position_embeddings + _POSITION_OFFSET, d, dtype=dtype, device=device)
|
||||
self.layers = nn.ModuleList([DecoderLayer(config, dtype=dtype, device=device, operations=operations) for _ in range(config.decoder_layers)])
|
||||
self.layernorm_embedding = operations.LayerNorm(d, dtype=dtype, device=device)
|
||||
|
||||
def forward(self, input_ids=None, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None,
|
||||
past_key_values=None, use_cache=True, inputs_embeds=None):
|
||||
if input_ids is not None and inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
|
||||
|
||||
if input_ids is not None:
|
||||
input_shape = input_ids.shape
|
||||
pos_input = input_ids
|
||||
else:
|
||||
input_shape = inputs_embeds.size()[:-1]
|
||||
pos_input = inputs_embeds[:, :, -1]
|
||||
|
||||
past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
|
||||
|
||||
position_ids = _make_position_ids(pos_input, past_key_values_length, self.embed_positions.weight.device)
|
||||
positions = self.embed_positions(position_ids).to(inputs_embeds.device)
|
||||
|
||||
hidden_states = self.layernorm_embedding(inputs_embeds + positions)
|
||||
|
||||
# When decoding with KV cache and single token, the causal mask is trivially all-attend
|
||||
# (a 1xN mask with no masking). Skip mask computation entirely for this fast path.
|
||||
tgt_len = input_shape[-1]
|
||||
if past_key_values is not None and tgt_len == 1:
|
||||
causal_mask = None
|
||||
else:
|
||||
causal_mask = make_causal_mask(input_shape, hidden_states.dtype, hidden_states.device, past_kv_len=past_key_values_length)
|
||||
if attention_mask is not None:
|
||||
causal_mask = causal_mask + make_attention_mask(attention_mask, hidden_states.dtype)
|
||||
|
||||
enc_attn_mask = None
|
||||
if encoder_hidden_states is not None and encoder_attention_mask is not None:
|
||||
if encoder_attention_mask.ndim == 4:
|
||||
enc_attn_mask = encoder_attention_mask # already expanded
|
||||
else:
|
||||
enc_attn_mask = make_attention_mask(encoder_attention_mask, hidden_states.dtype)
|
||||
|
||||
next_cache = () if use_cache else None
|
||||
for idx, layer in enumerate(self.layers):
|
||||
past_kv = past_key_values[idx] if past_key_values is not None else None
|
||||
hidden_states, present_kv = layer(
|
||||
hidden_states, attention_mask=causal_mask,
|
||||
encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=enc_attn_mask,
|
||||
past_key_value=past_kv,
|
||||
)
|
||||
if use_cache:
|
||||
next_cache = next_cache + (present_kv,)
|
||||
return (hidden_states, next_cache)
|
||||
|
||||
|
||||
class LanguageModel(nn.Module):
|
||||
def __init__(self, config, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.shared = operations.Embedding(config.vocab_size, config.d_model, dtype=dtype, device=device)
|
||||
self.encoder = Encoder(config, self.shared, dtype=dtype, device=device, operations=operations)
|
||||
self.decoder = Decoder(config, self.shared, dtype=dtype, device=device, operations=operations)
|
||||
|
||||
def forward(self, input_ids=None, attention_mask=None, decoder_input_ids=None, encoder_outputs=None, past_key_values=None, inputs_embeds=None, use_cache=True):
|
||||
if encoder_outputs is None:
|
||||
encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds)
|
||||
decoder_outputs = self.decoder(
|
||||
input_ids=decoder_input_ids, encoder_hidden_states=encoder_outputs[0],
|
||||
encoder_attention_mask=attention_mask, past_key_values=past_key_values, use_cache=use_cache,
|
||||
)
|
||||
return decoder_outputs + encoder_outputs
|
||||
|
||||
|
||||
class LanguageModelWithLMHead(nn.Module):
|
||||
def __init__(self, config, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.model = LanguageModel(config, dtype=dtype, device=device, operations=operations)
|
||||
self.lm_head = operations.Linear(config.d_model, config.vocab_size, bias=False, dtype=dtype, device=device)
|
||||
# Always zero — exists so the key loads from safetensors without warnings
|
||||
self.register_buffer("final_logits_bias", torch.zeros((1, config.vocab_size)))
|
||||
|
||||
def forward(self, input_ids=None, attention_mask=None, decoder_input_ids=None, encoder_outputs=None, past_key_values=None, inputs_embeds=None, use_cache=True):
|
||||
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_outputs, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache)
|
||||
return self.lm_head(outputs[0]), outputs
|
||||
|
||||
def tie_weights(self):
|
||||
self.model.encoder.embed_tokens.weight = self.model.shared.weight
|
||||
self.model.decoder.embed_tokens.weight = self.model.shared.weight
|
||||
self.lm_head.weight = self.model.shared.weight
|
||||
|
||||
|
||||
class PositionEmbedding2D(nn.Module):
|
||||
"""Learned absolute 2D positional embedding (row + column)."""
|
||||
def __init__(self, embedding_dim, num_pos, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.row_embeddings = operations.Embedding(num_pos, embedding_dim // 2, dtype=dtype, device=device)
|
||||
self.column_embeddings = operations.Embedding(num_pos, embedding_dim - (embedding_dim // 2), dtype=dtype, device=device)
|
||||
|
||||
def forward(self, pixel_values):
|
||||
"""pixel_values: [B, H, W, C] -> [B, H, W, embedding_dim]"""
|
||||
height, width = pixel_values.shape[1:3]
|
||||
x_emb = self.column_embeddings(torch.arange(width, device=pixel_values.device))
|
||||
y_emb = self.row_embeddings(torch.arange(height, device=pixel_values.device))
|
||||
pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)
|
||||
pos = pos.permute(2, 0, 1).unsqueeze(0).repeat(pixel_values.shape[0], 1, 1, 1).permute(0, 2, 3, 1)
|
||||
return pos
|
||||
|
||||
|
||||
class CosinePositionEmbedding1D(nn.Module):
|
||||
"""Fixed sinusoidal 1D positional encoding. Buffer key: pos_idx_to_embed"""
|
||||
def __init__(self, embed_dim, max_seq_len):
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.max_seq_len = max_seq_len
|
||||
factor = math.log(10000)
|
||||
denominator = torch.exp(-factor * torch.arange(0, embed_dim, 2) / embed_dim)
|
||||
frequencies = torch.arange(0, max_seq_len).reshape(max_seq_len, 1) * denominator
|
||||
pos_idx_to_embed = torch.zeros((max_seq_len, embed_dim))
|
||||
pos_idx_to_embed[:, 0::2] = torch.sin(frequencies)
|
||||
pos_idx_to_embed[:, 1::2] = torch.cos(frequencies)
|
||||
self.register_buffer("pos_idx_to_embed", pos_idx_to_embed)
|
||||
|
||||
def forward(self, seq_embeds):
|
||||
"""seq_embeds: [T, D] or [B, T, D] -> positional embeddings of matching shape."""
|
||||
shape_len = len(seq_embeds.shape)
|
||||
assert 2 <= shape_len <= 3
|
||||
pos_embeds = comfy.ops.cast_to_input(self.pos_idx_to_embed[:seq_embeds.size(-2), :], seq_embeds)
|
||||
if shape_len == 3:
|
||||
pos_embeds = pos_embeds.view(1, pos_embeds.size(0), pos_embeds.size(1))
|
||||
return pos_embeds
|
||||
|
||||
|
||||
class Florence2(nn.Module):
|
||||
"""Florence2 = DaViT vision encoder + image projection + BART LM."""
|
||||
|
||||
def __init__(self, config, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
vision_cfg = config.vision_config
|
||||
text_cfg = config.text_config
|
||||
|
||||
self.vision_tower = DaViT(vision_cfg, dtype=dtype, device=device, operations=operations)
|
||||
|
||||
image_dim_out = vision_cfg.dim_embed[-1]
|
||||
projection_dim = vision_cfg.projection_dim
|
||||
self.image_projection = nn.Parameter(torch.empty(image_dim_out, projection_dim, dtype=dtype, device=device))
|
||||
self.image_proj_norm = operations.LayerNorm(projection_dim, dtype=dtype, device=device)
|
||||
|
||||
image_pos_embed_config = vision_cfg.image_pos_embed
|
||||
if image_pos_embed_config is not None and image_pos_embed_config.get("type") == "learned_abs_2d":
|
||||
self.image_pos_embed = PositionEmbedding2D(embedding_dim=image_dim_out, num_pos=image_pos_embed_config["max_pos_embeddings"], dtype=dtype, device=device, operations=operations)
|
||||
else:
|
||||
self.image_pos_embed = None
|
||||
|
||||
visual_temporal_config = vision_cfg.visual_temporal_embedding
|
||||
if visual_temporal_config is not None and visual_temporal_config.get("type") == "COSINE":
|
||||
self.visual_temporal_embed = CosinePositionEmbedding1D(embed_dim=image_dim_out, max_seq_len=visual_temporal_config["max_temporal_embeddings"])
|
||||
else:
|
||||
self.visual_temporal_embed = None
|
||||
|
||||
self.image_feature_source = vision_cfg.image_feature_source
|
||||
self.language_model = LanguageModelWithLMHead(text_cfg, dtype=dtype, device=device, operations=operations)
|
||||
|
||||
def encode_image(self, pixel_values):
|
||||
"""Encode image through DaViT + projection. [B, 3, H, W] -> [B, num_tokens, projection_dim]"""
|
||||
batch_size = pixel_values.shape[0]
|
||||
T = 1
|
||||
x = self.vision_tower.forward_features_unpool(pixel_values)
|
||||
|
||||
if self.image_pos_embed is not None:
|
||||
x = x.view(batch_size * T, -1, x.shape[-1])
|
||||
num_tokens = x.shape[-2]
|
||||
h = w = int(num_tokens ** 0.5)
|
||||
assert h * w == num_tokens, "only square feature maps supported"
|
||||
x = x.view(batch_size * T, h, w, x.shape[-1])
|
||||
x = x + self.image_pos_embed(x)
|
||||
x = x.view(batch_size, T * h * w, x.shape[-1])
|
||||
|
||||
if self.visual_temporal_embed is not None:
|
||||
temporal = self.visual_temporal_embed(x.view(batch_size, T, -1, x.shape[-1])[:, :, 0])
|
||||
x = x.view(batch_size, T, -1, x.shape[-1]) + temporal.view(1, T, 1, x.shape[-1])
|
||||
|
||||
x_feat_dict = {
|
||||
"spatial_avg_pool": x.view(batch_size, T, -1, x.shape[-1]).mean(dim=2),
|
||||
"temporal_avg_pool": x.view(batch_size, T, -1, x.shape[-1]).mean(dim=1),
|
||||
"last_frame": x.view(batch_size, T, -1, x.shape[-1])[:, -1],
|
||||
}
|
||||
x = torch.cat([x_feat_dict[src] for src in self.image_feature_source], dim=1)
|
||||
|
||||
# Project — cast nn.Parameter to match x dtype (manual_cast doesn't cover raw params)
|
||||
x = x @ self.image_projection.to(dtype=x.dtype, device=x.device)
|
||||
x = self.image_proj_norm(x)
|
||||
return x
|
||||
|
||||
def _merge_input_ids_with_image_features(self, image_features, inputs_embeds):
|
||||
batch_size, image_token_length = image_features.size()[:-1]
|
||||
device = image_features.device
|
||||
image_attention_mask = torch.ones(batch_size, image_token_length, device=device)
|
||||
if inputs_embeds is None:
|
||||
return image_features, image_attention_mask
|
||||
text_attention_mask = torch.ones(batch_size, inputs_embeds.size(1), device=device)
|
||||
return torch.cat([image_features, inputs_embeds], dim=1), torch.cat([image_attention_mask, text_attention_mask], dim=1)
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, input_ids, pixel_values, max_new_tokens=1024, num_beams=3, do_sample=False, temperature=1.0):
|
||||
eos_token_id = self.config.text_config.eos_token_id
|
||||
decoder_start_token_id = self.config.text_config.decoder_start_token_id
|
||||
forced_bos_token_id = self.config.text_config.forced_bos_token_id
|
||||
|
||||
image_features = self.encode_image(pixel_values)
|
||||
|
||||
if input_ids is not None:
|
||||
embed_scale = self.language_model.model.encoder.embed_scale
|
||||
inputs_embeds = self.language_model.model.shared(input_ids) * embed_scale
|
||||
else:
|
||||
inputs_embeds = None
|
||||
|
||||
merged_embeds, attention_mask = self._merge_input_ids_with_image_features(image_features, inputs_embeds)
|
||||
attention_mask = attention_mask.to(merged_embeds.dtype)
|
||||
encoder_outputs = self.language_model.model.encoder(inputs_embeds=merged_embeds, attention_mask=attention_mask)
|
||||
|
||||
# Florence2 has no padding — the encoder attention mask is all-attend.
|
||||
# Passing None skips mask computation in cross-attention (faster).
|
||||
batch_size = merged_embeds.shape[0]
|
||||
args = dict(encoder_outputs=encoder_outputs, encoder_attention_mask=None, batch_size=batch_size,
|
||||
decoder_start_token_id=decoder_start_token_id, eos_token_id=eos_token_id,
|
||||
forced_bos_token_id=forced_bos_token_id, max_new_tokens=max_new_tokens)
|
||||
|
||||
if num_beams > 1:
|
||||
return self._beam_search(**args, num_beams=num_beams, do_sample=do_sample)
|
||||
else:
|
||||
return self._greedy_or_sample(**args, do_sample=do_sample, temperature=temperature)
|
||||
|
||||
def _greedy_or_sample(self, encoder_outputs, encoder_attention_mask, batch_size,
|
||||
decoder_start_token_id, eos_token_id, forced_bos_token_id,
|
||||
max_new_tokens, do_sample, temperature):
|
||||
device = encoder_outputs[0].device
|
||||
generated = torch.full((batch_size, 1), decoder_start_token_id, dtype=torch.long, device=device)
|
||||
past_key_values = None
|
||||
finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
|
||||
|
||||
pbar = ProgressBar(max_new_tokens)
|
||||
tbar = tqdm(total=max_new_tokens, desc="Florence2", unit="tok")
|
||||
for step in range(max_new_tokens):
|
||||
cur_input_ids = generated[:, -1:] if past_key_values is not None else generated
|
||||
decoder_outputs = self.language_model.model.decoder(
|
||||
input_ids=cur_input_ids, encoder_hidden_states=encoder_outputs[0],
|
||||
encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=True,
|
||||
)
|
||||
past_key_values = decoder_outputs[1]
|
||||
logits = self.language_model.lm_head(decoder_outputs[0][:, -1:, :])[:, -1, :]
|
||||
|
||||
if step == 0 and forced_bos_token_id is not None:
|
||||
next_token = torch.full((batch_size, 1), forced_bos_token_id, dtype=torch.long, device=device)
|
||||
else:
|
||||
if forced_bos_token_id is not None and step > 0:
|
||||
logits[:, forced_bos_token_id] = torch.finfo(logits.dtype).min
|
||||
if do_sample:
|
||||
next_token = torch.multinomial(F.softmax(logits / temperature, dim=-1), num_samples=1)
|
||||
else:
|
||||
next_token = logits.argmax(dim=-1, keepdim=True)
|
||||
|
||||
generated = torch.cat([generated, next_token], dim=-1)
|
||||
finished = finished | (next_token.squeeze(-1) == eos_token_id)
|
||||
pbar.update(1)
|
||||
tbar.update(1)
|
||||
if finished.all():
|
||||
break
|
||||
tbar.close()
|
||||
return generated
|
||||
|
||||
def _beam_search(self, encoder_outputs, encoder_attention_mask, batch_size,
|
||||
decoder_start_token_id, eos_token_id, forced_bos_token_id,
|
||||
max_new_tokens, num_beams, do_sample=False):
|
||||
device = encoder_outputs[0].device
|
||||
vocab_size = self.language_model.config.vocab_size
|
||||
|
||||
# Expand encoder outputs for beams
|
||||
encoder_hidden = encoder_outputs[0].unsqueeze(1).expand(-1, num_beams, -1, -1).reshape(batch_size * num_beams, -1, encoder_outputs[0].shape[-1])
|
||||
encoder_outputs_expanded = (encoder_hidden,)
|
||||
|
||||
generated = torch.full((batch_size * num_beams, 1), decoder_start_token_id, dtype=torch.long, device=device)
|
||||
past_key_values = None
|
||||
beam_scores = torch.zeros(batch_size, num_beams, device=device)
|
||||
beam_scores[:, 1:] = torch.finfo(beam_scores.dtype).min
|
||||
beam_scores = beam_scores.view(-1)
|
||||
|
||||
pbar = ProgressBar(max_new_tokens)
|
||||
tbar = tqdm(total=max_new_tokens, desc="Florence2 beam", unit="tok")
|
||||
for step in range(max_new_tokens):
|
||||
cur_input_ids = generated[:, -1:] if past_key_values is not None else generated
|
||||
decoder_outputs = self.language_model.model.decoder(
|
||||
input_ids=cur_input_ids, encoder_hidden_states=encoder_outputs_expanded[0],
|
||||
encoder_attention_mask=encoder_attention_mask, past_key_values=past_key_values, use_cache=True,
|
||||
)
|
||||
past_key_values = decoder_outputs[1]
|
||||
next_token_logits = self.language_model.lm_head(decoder_outputs[0][:, -1:, :])[:, -1, :]
|
||||
|
||||
# Force bos at step 0, suppress it after
|
||||
if step == 0 and forced_bos_token_id is not None:
|
||||
forced_scores = torch.full_like(next_token_logits, torch.finfo(next_token_logits.dtype).min)
|
||||
forced_scores[:, forced_bos_token_id] = 0
|
||||
next_token_scores = forced_scores
|
||||
else:
|
||||
if forced_bos_token_id is not None and step > 0:
|
||||
next_token_logits[:, forced_bos_token_id] = torch.finfo(next_token_logits.dtype).min
|
||||
next_token_scores = F.log_softmax(next_token_logits, dim=-1)
|
||||
|
||||
next_scores = (next_token_scores + beam_scores[:, None]).view(batch_size, num_beams * vocab_size)
|
||||
|
||||
# Step 0 with forced_bos is deterministic (only one valid token across all beams),
|
||||
# so skip sampling there — multinomial would fail with fewer nonzero entries than samples.
|
||||
sampling = do_sample and not (step == 0 and forced_bos_token_id is not None)
|
||||
if sampling:
|
||||
# HF beam-sample: multinomial without replacement over softmax(next_scores), then sort desc
|
||||
probs = F.softmax(next_scores, dim=-1)
|
||||
topk_indices = torch.multinomial(probs, num_samples=2 * num_beams)
|
||||
topk_scores = torch.gather(next_scores, -1, topk_indices)
|
||||
topk_scores, _sort_idx = torch.sort(topk_scores, descending=True, dim=-1)
|
||||
topk_indices = torch.gather(topk_indices, -1, _sort_idx)
|
||||
else:
|
||||
topk_scores, topk_indices = torch.topk(next_scores, 2 * num_beams, dim=-1, largest=True, sorted=True)
|
||||
|
||||
topk_beam_indices = topk_indices.div(vocab_size, rounding_mode='floor')
|
||||
topk_token_ids = topk_indices % vocab_size
|
||||
|
||||
next_beam_scores = topk_scores[:, :num_beams]
|
||||
next_beam_tokens = topk_token_ids[:, :num_beams]
|
||||
next_beam_flat_indices = topk_beam_indices[:, :num_beams]
|
||||
|
||||
beam_offset = torch.arange(batch_size, device=device).unsqueeze(1) * num_beams
|
||||
beam_idx_flat = (next_beam_flat_indices + beam_offset).view(-1)
|
||||
|
||||
beam_scores = next_beam_scores.view(-1)
|
||||
beam_tokens = next_beam_tokens.view(-1)
|
||||
|
||||
generated = torch.cat([generated[beam_idx_flat], beam_tokens.unsqueeze(-1)], dim=-1)
|
||||
past_key_values = self._reorder_cache(past_key_values, beam_idx_flat)
|
||||
|
||||
pbar.update(1)
|
||||
tbar.update(1)
|
||||
if (beam_tokens.view(batch_size, num_beams) == eos_token_id).all():
|
||||
break
|
||||
tbar.close()
|
||||
|
||||
# Select best beam, truncate at first eos
|
||||
beam_scores = beam_scores.view(batch_size, num_beams)
|
||||
best_indices = beam_scores.argmax(dim=-1) + torch.arange(batch_size, device=device) * num_beams
|
||||
generated = generated[best_indices]
|
||||
|
||||
results = []
|
||||
for b in range(batch_size):
|
||||
seq = generated[b]
|
||||
eos_positions = (seq == eos_token_id).nonzero(as_tuple=False)
|
||||
for pos in eos_positions:
|
||||
if pos.item() > 0:
|
||||
seq = seq[:pos.item() + 1]
|
||||
break
|
||||
results.append(seq)
|
||||
|
||||
max_len = max(r.shape[0] for r in results)
|
||||
pad_token_id = self.language_model.config.pad_token_id
|
||||
padded = torch.full((batch_size, max_len), pad_token_id, dtype=torch.long, device=device)
|
||||
for b, r in enumerate(results):
|
||||
padded[b, :r.shape[0]] = r
|
||||
return padded
|
||||
|
||||
@staticmethod
|
||||
def _reorder_cache(past_key_values, beam_idx):
|
||||
if hasattr(past_key_values, "reorder_cache"):
|
||||
past_key_values.reorder_cache(beam_idx)
|
||||
return past_key_values
|
||||
return tuple(
|
||||
tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
|
||||
for layer_past in past_key_values
|
||||
)
|
||||
@@ -0,0 +1,364 @@
|
||||
# Based on code from ComfyUI-Florence2 by Jukka Seppänen (kijai), MIT License
|
||||
"""Florence-2 processor: image preprocessing, tokenization, task prompts, and post-processing."""
|
||||
|
||||
import re
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def preprocess_image(image, size=768, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
|
||||
"""[B, C, H, W] or [C, H, W] float [0,1] -> [B, C, size, size] normalized."""
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
image = torch.nn.functional.interpolate(image, size=(size, size), mode='bicubic', align_corners=False).clamp(0, 1)
|
||||
mean_t = torch.tensor(mean, device=image.device, dtype=image.dtype).view(1, 3, 1, 1)
|
||||
std_t = torch.tensor(std, device=image.device, dtype=image.dtype).view(1, 3, 1, 1)
|
||||
return (image - mean_t) / std_t
|
||||
|
||||
|
||||
class BoxQuantizer:
|
||||
def __init__(self, mode, bins):
|
||||
self.mode = mode
|
||||
self.bins = bins
|
||||
|
||||
def dequantize(self, boxes, size):
|
||||
bins_w, bins_h = self.bins
|
||||
size_w, size_h = size
|
||||
xmin, ymin, xmax, ymax = boxes.split(1, dim=-1)
|
||||
return torch.cat((
|
||||
(xmin + 0.5) * size_w / bins_w, (ymin + 0.5) * size_h / bins_h,
|
||||
(xmax + 0.5) * size_w / bins_w, (ymax + 0.5) * size_h / bins_h,
|
||||
), dim=-1)
|
||||
|
||||
|
||||
class CoordinatesQuantizer:
|
||||
def __init__(self, mode, bins):
|
||||
self.mode = mode
|
||||
self.bins = bins
|
||||
|
||||
def dequantize(self, coordinates, size):
|
||||
bins_w, bins_h = self.bins
|
||||
size_w, size_h = size
|
||||
x, y = coordinates.split(1, dim=-1)
|
||||
return torch.cat(((x + 0.5) * size_w / bins_w, (y + 0.5) * size_h / bins_h), dim=-1)
|
||||
|
||||
|
||||
class PostProcessor:
|
||||
"""Regex-based parsing of Florence-2 text outputs into structured results."""
|
||||
|
||||
def __init__(self, tokenizer=None):
|
||||
config = {
|
||||
'NUM_BBOX_HEIGHT_BINS': 1000, 'NUM_BBOX_WIDTH_BINS': 1000, 'BOX_QUANTIZATION_MODE': 'floor',
|
||||
'COORDINATES_HEIGHT_BINS': 1000, 'COORDINATES_WIDTH_BINS': 1000, 'COORDINATES_QUANTIZATION_MODE': 'floor',
|
||||
'PARSE_TASKS': [
|
||||
{'TASK_NAME': 'od', 'PATTERN': r'([a-zA-Z0-9 ]+)<loc_(\\d+)><loc_(\\d+)><loc_(\\d+)><loc_(\\d+)>'},
|
||||
{'TASK_NAME': 'ocr', 'PATTERN': r'(.+?)<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>', 'AREA_THRESHOLD': 0.00},
|
||||
{'TASK_NAME': 'phrase_grounding', 'FILTER_BY_BLACK_LIST': True},
|
||||
{'TASK_NAME': 'pure_text'}, {'TASK_NAME': 'description_with_bboxes'},
|
||||
{'TASK_NAME': 'description_with_polygons'}, {'TASK_NAME': 'polygons'},
|
||||
{'TASK_NAME': 'bboxes'}, {'TASK_NAME': 'description_with_bboxes_or_polygons'},
|
||||
],
|
||||
}
|
||||
self.config = config
|
||||
self.parse_tasks = [t['TASK_NAME'] for t in config['PARSE_TASKS']]
|
||||
self.parse_tasks_configs = {t['TASK_NAME']: t for t in config['PARSE_TASKS']}
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
if tokenizer is not None:
|
||||
self.all_special_tokens = set(tokenizer.all_special_tokens)
|
||||
|
||||
self.box_quantizer = BoxQuantizer('floor', (1000, 1000))
|
||||
self.coordinates_quantizer = CoordinatesQuantizer('floor', (1000, 1000))
|
||||
|
||||
self.black_list_of_phrase_grounding = set()
|
||||
if 'phrase_grounding' in self.parse_tasks and self.parse_tasks_configs['phrase_grounding'].get('FILTER_BY_BLACK_LIST'):
|
||||
self.black_list_of_phrase_grounding = {
|
||||
'it', 'I', 'me', 'mine', 'you', 'your', 'yours', 'he', 'him', 'his',
|
||||
'she', 'her', 'hers', 'they', 'them', 'their', 'theirs', 'one', 'oneself',
|
||||
'we', 'us', 'our', 'ours', 'mine', 'yours', 'his', 'hers', 'its',
|
||||
'ours', 'yours', 'theirs', 'myself', 'yourself', 'himself', 'herself',
|
||||
'itself', 'ourselves', 'yourselves', 'themselves', 'this', 'that',
|
||||
'these', 'those', 'who', 'whom', 'whose', 'which', 'what', 'that',
|
||||
'all', 'another', 'any', 'anybody', 'anyone', 'anything',
|
||||
'each', 'everybody', 'everyone', 'everything',
|
||||
'few', 'many', 'nobody', 'none', 'one', 'several',
|
||||
'some', 'somebody', 'someone', 'something',
|
||||
'each other', 'one another',
|
||||
'the image', 'image', 'images', 'the', 'a', 'an', 'a group',
|
||||
'other objects', 'lots', 'a set',
|
||||
}
|
||||
|
||||
def parse_od_from_text_and_spans(self, text, pattern, image_size, phrase_centric=False):
|
||||
parsed = list(re.finditer(pattern, text))
|
||||
instances = []
|
||||
for m in parsed:
|
||||
if phrase_centric:
|
||||
bbox_bins = [int(m.group(j)) for j in range(2, 6)]
|
||||
cat_name = m.group(1).lower().strip()
|
||||
else:
|
||||
bbox_bins = [int(m.group(j)) for j in range(1, 5)]
|
||||
cat_name = m.group(5).lower().strip()
|
||||
instances.append({
|
||||
'bbox': self.box_quantizer.dequantize(torch.tensor(bbox_bins), image_size).tolist(),
|
||||
'cat_name': cat_name,
|
||||
})
|
||||
return instances
|
||||
|
||||
def parse_ocr_from_text_and_spans(self, text, pattern, image_size, area_threshold=-1.0):
|
||||
text = text.replace('<s>', '')
|
||||
parsed = re.findall(pattern, text)
|
||||
instances = []
|
||||
image_width, image_height = image_size
|
||||
for ocr_line in parsed:
|
||||
ocr_content = ocr_line[0]
|
||||
quad_box = [int(i) for i in ocr_line[1:]]
|
||||
quad_box = self.coordinates_quantizer.dequantize(torch.tensor(np.array(quad_box).reshape(-1, 2)), image_size).reshape(-1).tolist()
|
||||
if area_threshold > 0:
|
||||
x_coords, y_coords = quad_box[0::2], quad_box[1::2]
|
||||
area = 0.5 * abs(sum(x_coords[i] * y_coords[i + 1] - x_coords[i + 1] * y_coords[i] for i in range(3)))
|
||||
if area < (image_width * image_height) * area_threshold:
|
||||
continue
|
||||
instances.append({'quad_box': quad_box, 'text': ocr_content})
|
||||
return instances
|
||||
|
||||
def parse_phrase_grounding_from_text_and_spans(self, text, pattern, image_size):
|
||||
text = text.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
|
||||
pattern = r"([^<]+(?:<loc_\d+>){4,})"
|
||||
phrases = re.findall(pattern, text)
|
||||
phrase_pattern = r'^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_)'
|
||||
box_pattern = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
|
||||
|
||||
instances = []
|
||||
for pharse_text in phrases:
|
||||
phrase_text_strip = pharse_text.replace('<ground>', '', 1).replace('<obj>', '', 1)
|
||||
if phrase_text_strip == '':
|
||||
continue
|
||||
phrase = re.search(phrase_pattern, phrase_text_strip)
|
||||
if phrase is None:
|
||||
continue
|
||||
bboxes_parsed = list(re.finditer(box_pattern, pharse_text))
|
||||
if not bboxes_parsed:
|
||||
continue
|
||||
phrase = phrase.group().strip()
|
||||
if phrase in self.black_list_of_phrase_grounding:
|
||||
continue
|
||||
bbox_bins = [[int(b.group(j)) for j in range(1, 5)] for b in bboxes_parsed]
|
||||
phrase = phrase.encode('ascii', errors='ignore').decode('ascii')
|
||||
instances.append({
|
||||
'bbox': self.box_quantizer.dequantize(torch.tensor(bbox_bins), image_size).tolist(),
|
||||
'cat_name': phrase,
|
||||
})
|
||||
return instances
|
||||
|
||||
def parse_description_with_bboxes_from_text_and_spans(self, text, pattern, image_size, allow_empty_phrase=False):
|
||||
text = text.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
|
||||
if allow_empty_phrase:
|
||||
pattern = r"(?:(?:<loc_\d+>){{4,}})"
|
||||
else:
|
||||
pattern = r"([^<]+(?:<loc_\d+>){4,})"
|
||||
phrases = re.findall(pattern, text)
|
||||
phrase_pattern = r'^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_)'
|
||||
box_pattern = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
|
||||
|
||||
instances = []
|
||||
for pharse_text in phrases:
|
||||
phrase_text_strip = pharse_text.replace('<ground>', '', 1).replace('<obj>', '', 1)
|
||||
if phrase_text_strip == '' and not allow_empty_phrase:
|
||||
continue
|
||||
phrase = re.search(phrase_pattern, phrase_text_strip)
|
||||
if phrase is None:
|
||||
continue
|
||||
phrase = phrase.group().strip()
|
||||
bboxes_parsed = list(re.finditer(box_pattern, pharse_text))
|
||||
if not bboxes_parsed:
|
||||
continue
|
||||
bbox_bins = [[int(b.group(j)) for j in range(1, 5)] for b in bboxes_parsed]
|
||||
bboxes = self.box_quantizer.dequantize(torch.tensor(bbox_bins), image_size).tolist()
|
||||
phrase = phrase.encode('ascii', errors='ignore').decode('ascii')
|
||||
for bbox in bboxes:
|
||||
instances.append({'bbox': bbox, 'cat_name': phrase})
|
||||
return instances
|
||||
|
||||
def parse_description_with_polygons_from_text_and_spans(self, text, pattern, image_size,
|
||||
allow_empty_phrase=False, polygon_sep_token='<sep>',
|
||||
polygon_start_token='<poly>', polygon_end_token='</poly>',
|
||||
with_box_at_start=False):
|
||||
text = text.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
|
||||
sep, start, end = re.escape(polygon_sep_token), re.escape(polygon_start_token), re.escape(polygon_end_token)
|
||||
if allow_empty_phrase:
|
||||
pattern = rf"(?:(?:<loc_\d+>|{sep}|{start}|{end}){{4,}})"
|
||||
else:
|
||||
pattern = rf"([^<]+(?:<loc_\d+>|{sep}|{start}|{end}){{4,}})"
|
||||
phrases = re.findall(pattern, text)
|
||||
phrase_string_pattern = r'^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_|<poly>)'
|
||||
box_pattern = rf'((?:<loc_\d+>)+)(?:{sep}|$)'
|
||||
polygons_instance_pattern = rf'{start}(.*?){end}'
|
||||
|
||||
instances = []
|
||||
for phrase_text in phrases:
|
||||
phrase_text_strip = re.sub(r'^loc_\d+>', '', phrase_text, count=1)
|
||||
if phrase_text_strip == '' and not allow_empty_phrase:
|
||||
continue
|
||||
phrase = re.search(phrase_string_pattern, phrase_text_strip)
|
||||
if phrase is None:
|
||||
continue
|
||||
phrase = phrase.group().strip()
|
||||
|
||||
if polygon_start_token in phrase_text and polygon_end_token in phrase_text:
|
||||
poly_instances = list(re.finditer(polygons_instance_pattern, phrase_text))
|
||||
else:
|
||||
poly_instances = [phrase_text]
|
||||
|
||||
for pi in poly_instances:
|
||||
poly_text = pi.group(1) if not isinstance(pi, str) else pi
|
||||
polygons_parsed = list(re.finditer(box_pattern, poly_text))
|
||||
if not polygons_parsed:
|
||||
continue
|
||||
|
||||
bbox, polygons = [], []
|
||||
for pp in polygons_parsed:
|
||||
coords = [int(m.group(1)) for m in re.finditer(r'<loc_(\d+)>', pp.group(1))]
|
||||
if with_box_at_start and not bbox:
|
||||
if len(coords) > 4:
|
||||
bbox = coords[:4]
|
||||
coords = coords[4:]
|
||||
else:
|
||||
bbox = [0, 0, 0, 0]
|
||||
if len(coords) % 2 == 1:
|
||||
coords = coords[:-1]
|
||||
polygons.append(self.coordinates_quantizer.dequantize(
|
||||
torch.tensor(np.array(coords).reshape(-1, 2)), image_size,
|
||||
).reshape(-1).tolist())
|
||||
|
||||
instance = {'cat_name': phrase, 'polygons': polygons}
|
||||
if bbox:
|
||||
instance['bbox'] = self.box_quantizer.dequantize(torch.tensor([bbox]), image_size).tolist()[0]
|
||||
instances.append(instance)
|
||||
return instances
|
||||
|
||||
def __call__(self, text=None, image_size=None, parse_tasks=None):
|
||||
if parse_tasks is not None:
|
||||
if isinstance(parse_tasks, str):
|
||||
parse_tasks = [parse_tasks]
|
||||
for t in parse_tasks:
|
||||
assert t in self.parse_tasks, f'parse task {t} not supported'
|
||||
|
||||
assert text is not None, 'text should be provided'
|
||||
parsed_dict = {'text': text}
|
||||
|
||||
for task in self.parse_tasks:
|
||||
if parse_tasks is not None and task not in parse_tasks:
|
||||
continue
|
||||
pattern = self.parse_tasks_configs[task].get('PATTERN', None)
|
||||
|
||||
if task == 'ocr':
|
||||
parsed_dict['ocr'] = self.parse_ocr_from_text_and_spans(text, pattern, image_size, self.parse_tasks_configs[task].get('AREA_THRESHOLD', 0.0))
|
||||
elif task == 'phrase_grounding':
|
||||
parsed_dict['phrase_grounding'] = self.parse_phrase_grounding_from_text_and_spans(text, pattern, image_size)
|
||||
elif task == 'pure_text':
|
||||
parsed_dict['pure_text'] = text
|
||||
elif task == 'description_with_bboxes':
|
||||
parsed_dict['description_with_bboxes'] = self.parse_description_with_bboxes_from_text_and_spans(text, pattern, image_size)
|
||||
elif task == 'description_with_polygons':
|
||||
parsed_dict['description_with_polygons'] = self.parse_description_with_polygons_from_text_and_spans(text, pattern, image_size)
|
||||
elif task == 'polygons':
|
||||
parsed_dict['polygons'] = self.parse_description_with_polygons_from_text_and_spans(text, pattern, image_size, allow_empty_phrase=True)
|
||||
elif task == 'bboxes':
|
||||
parsed_dict['bboxes'] = self.parse_description_with_bboxes_from_text_and_spans(text, pattern, image_size, allow_empty_phrase=True)
|
||||
elif task == 'description_with_bboxes_or_polygons':
|
||||
if '<poly>' in text:
|
||||
parsed_dict['description_with_bboxes_or_polygons'] = self.parse_description_with_polygons_from_text_and_spans(text, pattern, image_size)
|
||||
else:
|
||||
parsed_dict['description_with_bboxes_or_polygons'] = self.parse_description_with_bboxes_from_text_and_spans(text, pattern, image_size)
|
||||
else:
|
||||
raise ValueError(f"task {task} is not supported")
|
||||
return parsed_dict
|
||||
|
||||
|
||||
class Processor:
|
||||
def __init__(self, model_path):
|
||||
from .tokenizer import Florence2Tokenizer
|
||||
self.tokenizer = Florence2Tokenizer(model_path)
|
||||
self.image_seq_length = 577
|
||||
|
||||
self.tasks_answer_post_processing_type = {
|
||||
'<OCR>': 'pure_text', '<OCR_WITH_REGION>': 'ocr',
|
||||
'<CAPTION>': 'pure_text', '<DETAILED_CAPTION>': 'pure_text', '<MORE_DETAILED_CAPTION>': 'pure_text',
|
||||
'<OD>': 'description_with_bboxes', '<DENSE_REGION_CAPTION>': 'description_with_bboxes',
|
||||
'<CAPTION_TO_PHRASE_GROUNDING>': 'phrase_grounding',
|
||||
'<REFERRING_EXPRESSION_SEGMENTATION>': 'polygons', '<REGION_TO_SEGMENTATION>': 'polygons',
|
||||
'<OPEN_VOCABULARY_DETECTION>': 'description_with_bboxes_or_polygons',
|
||||
'<REGION_TO_CATEGORY>': 'pure_text', '<REGION_TO_DESCRIPTION>': 'pure_text', '<REGION_TO_OCR>': 'pure_text',
|
||||
'<REGION_PROPOSAL>': 'bboxes',
|
||||
}
|
||||
self.task_prompts_without_inputs = {
|
||||
'<OCR>': 'What is the text in the image?',
|
||||
'<OCR_WITH_REGION>': 'What is the text in the image, with regions?',
|
||||
'<CAPTION>': 'What does the image describe?',
|
||||
'<DETAILED_CAPTION>': 'Describe in detail what is shown in the image.',
|
||||
'<MORE_DETAILED_CAPTION>': 'Describe with a paragraph what is shown in the image.',
|
||||
'<OD>': 'Locate the objects with category name in the image.',
|
||||
'<DENSE_REGION_CAPTION>': 'Locate the objects in the image, with their descriptions.',
|
||||
'<REGION_PROPOSAL>': 'Locate the region proposals in the image.',
|
||||
}
|
||||
self.task_prompts_with_input = {
|
||||
'<CAPTION_TO_PHRASE_GROUNDING>': 'Locate the phrases in the caption: {input}',
|
||||
'<REFERRING_EXPRESSION_SEGMENTATION>': 'Locate {input} in the image with mask',
|
||||
'<REGION_TO_SEGMENTATION>': 'What is the polygon mask of region {input}',
|
||||
'<OPEN_VOCABULARY_DETECTION>': 'Locate {input} in the image.',
|
||||
'<REGION_TO_CATEGORY>': 'What is the region {input}?',
|
||||
'<REGION_TO_DESCRIPTION>': 'What does the region {input} describe?',
|
||||
'<REGION_TO_OCR>': 'What text is in the region {input}?',
|
||||
}
|
||||
self.post_processor = PostProcessor(tokenizer=self.tokenizer)
|
||||
|
||||
def _construct_prompts(self, text):
|
||||
for task_token, task_prompt in self.task_prompts_without_inputs.items():
|
||||
if task_token in text:
|
||||
return task_prompt
|
||||
for task_token, task_prompt in self.task_prompts_with_input.items():
|
||||
if task_token in text:
|
||||
return task_prompt.format(input=text.replace(task_token, ''))
|
||||
return text
|
||||
|
||||
def __call__(self, text, images):
|
||||
prompt = self._construct_prompts(text)
|
||||
encoded = self.tokenizer.encode(prompt)
|
||||
pixel_values = preprocess_image(images)
|
||||
return {'input_ids': encoded['input_ids'], 'pixel_values': pixel_values}
|
||||
|
||||
def batch_decode(self, token_ids, skip_special_tokens=False):
|
||||
return self.tokenizer.batch_decode(token_ids, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
def post_process_generation(self, text, task, image_size):
|
||||
pp_type = self.tasks_answer_post_processing_type.get(task, 'pure_text')
|
||||
result = self.post_processor(text=text, image_size=image_size, parse_tasks=pp_type)[pp_type]
|
||||
|
||||
if pp_type == 'pure_text':
|
||||
final = result.replace('<s>', '').replace('</s>', '')
|
||||
elif pp_type in ('od', 'description_with_bboxes', 'bboxes'):
|
||||
final = {'bboxes': [i['bbox'] for i in result], 'labels': [str(i['cat_name']) for i in result]}
|
||||
elif pp_type == 'ocr':
|
||||
final = {'quad_boxes': [i['quad_box'] for i in result], 'labels': [str(i['text']) for i in result]}
|
||||
elif pp_type == 'phrase_grounding':
|
||||
bboxes, labels = [], []
|
||||
for phrase in result:
|
||||
for bbox in phrase['bbox']:
|
||||
bboxes.append(bbox)
|
||||
labels.append(phrase['cat_name'])
|
||||
final = {'bboxes': bboxes, 'labels': labels}
|
||||
elif pp_type in ('description_with_polygons', 'polygons'):
|
||||
final = {'polygons': [r['polygons'] for r in result], 'labels': [r['cat_name'] for r in result]}
|
||||
elif pp_type == 'description_with_bboxes_or_polygons':
|
||||
bboxes, bl, polygons, pl = [], [], [], []
|
||||
for r in result:
|
||||
if 'polygons' in r:
|
||||
polygons.append(r['polygons'])
|
||||
pl.append(r['cat_name'])
|
||||
else:
|
||||
bboxes.append(r['bbox'])
|
||||
bl.append(r['cat_name'])
|
||||
final = {'bboxes': bboxes, 'bboxes_labels': bl, 'polygons': polygons, 'polygons_labels': pl}
|
||||
else:
|
||||
raise ValueError(f'Unknown post processing type: {pp_type}')
|
||||
return {task: final}
|
||||
@@ -0,0 +1,67 @@
|
||||
# Based on code from ComfyUI-Florence2 by Jukka Seppänen (kijai), MIT License
|
||||
import os
|
||||
from tokenizers import Tokenizer, AddedToken
|
||||
|
||||
|
||||
class Florence2Tokenizer:
|
||||
def __init__(self, model_path):
|
||||
tokenizer_path = os.path.join(model_path, "tokenizer.json")
|
||||
self.tokenizer = Tokenizer.from_file(tokenizer_path)
|
||||
self._add_special_tokens()
|
||||
|
||||
# Standard token IDs
|
||||
self.pad_token_id = 1
|
||||
self.bos_token_id = 0
|
||||
self.eos_token_id = 2
|
||||
|
||||
def _add_special_tokens(self):
|
||||
"""Add Florence2-specific special tokens."""
|
||||
special_tokens = ['<od>', '</od>', '<ocr>', '</ocr>']
|
||||
special_tokens += [f'<loc_{x}>' for x in range(1000)]
|
||||
special_tokens += [
|
||||
'<cap>', '</cap>', '<ncap>', '</ncap>', '<dcap>', '</dcap>',
|
||||
'<grounding>', '</grounding>', '<seg>', '</seg>', '<sep>',
|
||||
'<region_cap>', '</region_cap>', '<region_to_desciption>',
|
||||
'</region_to_desciption>', '<proposal>', '</proposal>',
|
||||
'<poly>', '</poly>', '<and>'
|
||||
]
|
||||
added = [AddedToken(t, special=True) for t in special_tokens]
|
||||
self.tokenizer.add_special_tokens(added)
|
||||
|
||||
def encode(self, text):
|
||||
"""Encode text to token ids. Returns dict with 'input_ids' tensor."""
|
||||
import torch
|
||||
encoding = self.tokenizer.encode(text)
|
||||
return {"input_ids": torch.tensor([encoding.ids], dtype=torch.long)}
|
||||
|
||||
def decode(self, token_ids, skip_special_tokens=False):
|
||||
"""Decode token ids to text."""
|
||||
if hasattr(token_ids, 'tolist'):
|
||||
token_ids = token_ids.tolist()
|
||||
if isinstance(token_ids[0], list):
|
||||
token_ids = token_ids[0]
|
||||
return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
def batch_decode(self, token_ids_batch, skip_special_tokens=False):
|
||||
"""Decode a batch of token id sequences."""
|
||||
results = []
|
||||
if hasattr(token_ids_batch, 'tolist'):
|
||||
token_ids_batch = token_ids_batch.tolist()
|
||||
for ids in token_ids_batch:
|
||||
results.append(self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens))
|
||||
return results
|
||||
|
||||
@property
|
||||
def all_special_tokens(self):
|
||||
"""Return set of all special token strings. Needed by post-processor."""
|
||||
added = self.tokenizer.get_added_tokens_decoder()
|
||||
tokens = set()
|
||||
for token_obj in added.values():
|
||||
tokens.add(str(token_obj))
|
||||
# Also add the built-in ones
|
||||
tokens.update({'<s>', '</s>', '<pad>', '<unk>', '<mask>'})
|
||||
return tokens
|
||||
|
||||
@property
|
||||
def vocab_size(self):
|
||||
return self.tokenizer.get_vocab_size()
|
||||
Reference in New Issue
Block a user