165 lines
6.2 KiB
Python
165 lines
6.2 KiB
Python
# 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,
|
|
)
|