diff --git a/models/florence2/__init__.py b/models/florence2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/florence2/config.py b/models/florence2/config.py new file mode 100644 index 0000000..c849c04 --- /dev/null +++ b/models/florence2/config.py @@ -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, + ) diff --git a/models/florence2/davit.py b/models/florence2/davit.py new file mode 100644 index 0000000..3094eea --- /dev/null +++ b/models/florence2/davit.py @@ -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 diff --git a/models/florence2/model.py b/models/florence2/model.py new file mode 100644 index 0000000..f3f9c9b --- /dev/null +++ b/models/florence2/model.py @@ -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 + ) diff --git a/models/florence2/processing.py b/models/florence2/processing.py new file mode 100644 index 0000000..bcbd3a1 --- /dev/null +++ b/models/florence2/processing.py @@ -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 ]+)'}, + {'TASK_NAME': 'ocr', 'PATTERN': r'(.+?)', '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('', '') + 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('', '').replace('', '').replace('', '') + pattern = r"([^<]+(?:){4,})" + phrases = re.findall(pattern, text) + phrase_pattern = r'^\s*(.*?)(?=||||||' + + instances = [] + for pharse_text in phrases: + phrase_text_strip = pharse_text.replace('', '', 1).replace('', '', 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('', '').replace('', '').replace('', '') + if allow_empty_phrase: + pattern = r"(?:(?:){{4,}})" + else: + pattern = r"([^<]+(?:){4,})" + phrases = re.findall(pattern, text) + phrase_pattern = r'^\s*(.*?)(?=||||||' + + instances = [] + for pharse_text in phrases: + phrase_text_strip = pharse_text.replace('', '', 1).replace('', '', 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='', + polygon_start_token='', polygon_end_token='', + with_box_at_start=False): + text = text.replace('', '').replace('', '').replace('', '') + sep, start, end = re.escape(polygon_sep_token), re.escape(polygon_start_token), re.escape(polygon_end_token) + if allow_empty_phrase: + pattern = rf"(?:(?:|{sep}|{start}|{end}){{4,}})" + else: + pattern = rf"([^<]+(?:|{sep}|{start}|{end}){{4,}})" + phrases = re.findall(pattern, text) + phrase_string_pattern = r'^\s*(.*?)(?=||||||)' + box_pattern = rf'((?:)+)(?:{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'', 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 '' 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 = { + '': 'pure_text', '': 'ocr', + '': 'pure_text', '': 'pure_text', '': 'pure_text', + '': 'description_with_bboxes', '': 'description_with_bboxes', + '': 'phrase_grounding', + '': 'polygons', '': 'polygons', + '': 'description_with_bboxes_or_polygons', + '': 'pure_text', '': 'pure_text', '': 'pure_text', + '': 'bboxes', + } + self.task_prompts_without_inputs = { + '': 'What is the text in the image?', + '': 'What is the text in the image, with regions?', + '': 'What does the image describe?', + '': 'Describe in detail what is shown in the image.', + '': 'Describe with a paragraph what is shown in the image.', + '': 'Locate the objects with category name in the image.', + '': 'Locate the objects in the image, with their descriptions.', + '': 'Locate the region proposals in the image.', + } + self.task_prompts_with_input = { + '': 'Locate the phrases in the caption: {input}', + '': 'Locate {input} in the image with mask', + '': 'What is the polygon mask of region {input}', + '': 'Locate {input} in the image.', + '': 'What is the region {input}?', + '': 'What does the region {input} describe?', + '': '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('', '').replace('', '') + 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} diff --git a/models/florence2/tokenizer.py b/models/florence2/tokenizer.py new file mode 100644 index 0000000..d1a1dee --- /dev/null +++ b/models/florence2/tokenizer.py @@ -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 = ['', '', '', ''] + special_tokens += [f'' for x in range(1000)] + special_tokens += [ + '', '', '', '', '', '', + '', '', '', '', '', + '', '', '', + '', '', '', + '', '', '' + ] + 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({'', '', '', '', ''}) + return tokens + + @property + def vocab_size(self): + return self.tokenizer.get_vocab_size()