562 lines
19 KiB
Python
562 lines
19 KiB
Python
import logging
|
|
import os
|
|
from types import SimpleNamespace
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from huggingface_hub import hf_hub_download
|
|
from safetensors import safe_open
|
|
from transformers import CLIPTokenizer
|
|
|
|
import comfy.model_management
|
|
import comfy.sample
|
|
import comfy.utils
|
|
|
|
|
|
class DecoderBlock(nn.Module):
|
|
def __init__(self, embed_dim: int, num_heads: int):
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(embed_dim)
|
|
self.self_attn = nn.MultiheadAttention(
|
|
embed_dim=embed_dim,
|
|
num_heads=num_heads,
|
|
batch_first=True,
|
|
)
|
|
|
|
self.norm2 = nn.LayerNorm(embed_dim)
|
|
self.cross_attn = nn.MultiheadAttention(
|
|
embed_dim=embed_dim,
|
|
num_heads=num_heads,
|
|
batch_first=True,
|
|
)
|
|
|
|
self.norm3 = nn.LayerNorm(embed_dim)
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(embed_dim, embed_dim * 4),
|
|
nn.GELU(),
|
|
nn.Identity(),
|
|
nn.Linear(embed_dim * 4, embed_dim),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
memory: torch.Tensor,
|
|
self_attn_mask: Optional[torch.Tensor] = None,
|
|
):
|
|
# self attention with mask
|
|
residual = x
|
|
x = self.norm1(x)
|
|
attn_output, _ = self.self_attn(
|
|
query=x,
|
|
key=x,
|
|
value=x,
|
|
attn_mask=self_attn_mask,
|
|
need_weights=False,
|
|
is_causal=True,
|
|
)
|
|
x = residual + attn_output
|
|
|
|
# cross attention
|
|
residual = x
|
|
x = self.norm2(x)
|
|
attn_output, _ = self.cross_attn(
|
|
query=x, key=memory, value=memory, need_weights=False
|
|
)
|
|
x = residual + attn_output
|
|
|
|
# FFN
|
|
residual = x
|
|
x = self.norm3(x)
|
|
x = residual + self.mlp(x)
|
|
|
|
return x
|
|
|
|
|
|
class Captioner(nn.Module):
|
|
def __init__(self, config, vision_embed_dim: int, vocab_size: int):
|
|
super().__init__()
|
|
|
|
self.embed_dim = vision_embed_dim
|
|
self.hidden_dim = config.hidden_dim
|
|
self.vocab_size = vocab_size
|
|
self.max_length = config.max_length
|
|
|
|
# projection from ViT dimension to decoder dimension
|
|
self.projection = nn.Linear(self.embed_dim, self.hidden_dim)
|
|
self.memory_pos_embedding = nn.Parameter(torch.zeros(1, 257, self.hidden_dim))
|
|
|
|
# decoder layers
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
DecoderBlock(config.hidden_dim, config.num_heads)
|
|
for _ in range(config.num_blocks)
|
|
]
|
|
)
|
|
|
|
causal_mask = nn.Transformer.generate_square_subsequent_mask(self.max_length)
|
|
self.register_buffer("causal_mask", causal_mask, persistent=False)
|
|
|
|
def generate(
|
|
self,
|
|
image_features: torch.Tensor,
|
|
temperature: float,
|
|
tokenizer: CLIPTokenizer,
|
|
output_projection: nn.Module,
|
|
token_embedding_: nn.Module,
|
|
pos_embedding_: nn.Module,
|
|
seed: Optional[int] = None,
|
|
ramble: bool = False,
|
|
):
|
|
# project and add positional embeddings to image features
|
|
memory = self.projection(image_features)
|
|
memory = memory + self.memory_pos_embedding
|
|
|
|
# start with BOS token
|
|
batch_size = image_features.size(0)
|
|
current_tokens = torch.full(
|
|
(batch_size, 1),
|
|
tokenizer.bos_token_id,
|
|
dtype=torch.long,
|
|
device=image_features.device,
|
|
)
|
|
generated_tokens = [current_tokens]
|
|
|
|
generator = torch.Generator(device=image_features.device)
|
|
if seed is not None:
|
|
generator.manual_seed(seed)
|
|
|
|
for _ in range(self.max_length - 2):
|
|
# embed current tokens
|
|
token_embeddings = token_embedding_(current_tokens)
|
|
positions = torch.arange(
|
|
current_tokens.size(1), device=current_tokens.device
|
|
)
|
|
pos_embeddings = pos_embedding_(positions)
|
|
x = token_embeddings + pos_embeddings
|
|
|
|
# get causal mask for self-attention
|
|
seq_len = x.size(1)
|
|
mask = self.causal_mask[:seq_len, :seq_len]
|
|
|
|
# pass through decoder layers
|
|
for layer in self.layers:
|
|
x = layer(x, memory, self_attn_mask=mask)
|
|
|
|
# get next token
|
|
logits = output_projection(x[:, -1:])
|
|
logits = logits / temperature
|
|
if ramble:
|
|
logits[:, :, tokenizer.eos_token_id] = -float("inf")
|
|
probs = F.softmax(logits, dim=-1)
|
|
next_token = torch.multinomial(probs.squeeze(1), 1, generator=generator)
|
|
|
|
generated_tokens.append(next_token)
|
|
current_tokens = torch.cat([current_tokens, next_token], dim=1)
|
|
|
|
# stop if all sequences have EOS token
|
|
if (next_token == tokenizer.eos_token_id).all():
|
|
break
|
|
|
|
return torch.cat(generated_tokens, dim=1)
|
|
|
|
|
|
class CLIPtionModel(nn.Module):
|
|
def __init__(self, config, clip, clip_vision):
|
|
super().__init__()
|
|
|
|
if not hasattr(clip, "cond_stage_model"):
|
|
raise ValueError("CLIP is missing from model checkpoint")
|
|
if not hasattr(clip.cond_stage_model, "clip_l"):
|
|
raise ValueError("Must use model which includes CLIP-L")
|
|
|
|
# store CLIP model references
|
|
self.clip_text = clip
|
|
self.clip_vision = clip_vision
|
|
self.tokenizer = clip.tokenizer.clip_l.tokenizer
|
|
self.text_model = clip.cond_stage_model.clip_l.transformer.text_model
|
|
self.vision_model = clip_vision.model.vision_model
|
|
|
|
# clip.cond_stage_model.clip_l.transformer.text_projection is empty
|
|
# so load a copy from the CLIPtion safetensors file instead
|
|
self.text_projection = nn.Linear(768, 768, bias=False)
|
|
|
|
# create caption decoder
|
|
self.captioner = Captioner(config, 1024, self.tokenizer.vocab_size)
|
|
|
|
# use CLIP's token embeddings for output projection
|
|
clip_embed_weight = self.text_model.embeddings.token_embedding.weight
|
|
self.output_projection = nn.Linear(
|
|
self.captioner.hidden_dim, self.tokenizer.vocab_size, bias=False
|
|
)
|
|
self.output_projection.weight = nn.Parameter(
|
|
clip_embed_weight.clone(), requires_grad=False
|
|
)
|
|
|
|
def generate(
|
|
self,
|
|
images: torch.Tensor,
|
|
seed: int = 42,
|
|
temperature: float = 0.7,
|
|
best_of: int = 1,
|
|
ramble: bool = False,
|
|
) -> List[str]:
|
|
device = comfy.model_management.get_torch_device()
|
|
|
|
image_outputs = self.clip_vision.encode_image(images)
|
|
image_features = image_outputs.last_hidden_state
|
|
|
|
image_embeds = image_outputs.image_embeds
|
|
image_embeds = image_embeds.to(device, dtype=torch.float16)
|
|
image_embeds /= image_embeds.norm(dim=-1, keepdim=True)
|
|
|
|
captions = []
|
|
for image_idx in range(image_features.size(0)):
|
|
candidates = []
|
|
for _ in range(best_of):
|
|
features = (
|
|
image_features[image_idx]
|
|
.unsqueeze(0)
|
|
.to(device, dtype=torch.float16)
|
|
)
|
|
tokens = self.captioner.generate(
|
|
features,
|
|
temperature,
|
|
self.tokenizer,
|
|
self.output_projection,
|
|
self.text_model.embeddings.token_embedding,
|
|
self.text_model.embeddings.position_embedding,
|
|
seed=seed,
|
|
ramble=ramble,
|
|
)
|
|
seed += 1
|
|
text = self.tokenizer.decode(
|
|
tokens.squeeze(),
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=True,
|
|
)
|
|
candidates.append(text)
|
|
|
|
if best_of == 1:
|
|
captions.append(candidates[0])
|
|
continue
|
|
|
|
# calculate CLIP similarity for each candidate
|
|
scored = []
|
|
image_embed = image_embeds[image_idx : image_idx + 1]
|
|
for text in candidates:
|
|
text_embeds = self._text_to_embed(text, device)
|
|
clip_sim = torch.sum(image_embed * text_embeds, dim=-1)[0]
|
|
scored.append((clip_sim, text))
|
|
|
|
# return the candidate with the highest similarity
|
|
scored.sort(key=lambda x: x[0], reverse=True)
|
|
for score, text in scored:
|
|
logging.debug(f"({score.item()}) {text}")
|
|
captions.append(scored[0][1])
|
|
|
|
return captions
|
|
|
|
def generate_beam(
|
|
self,
|
|
images: torch.Tensor,
|
|
temperature: float = 0.7,
|
|
beam_width: int = 4,
|
|
ramble: bool = False,
|
|
) -> List[str]:
|
|
device = comfy.model_management.get_torch_device()
|
|
|
|
# get image features and embeddings
|
|
image_outputs = self.clip_vision.encode_image(images)
|
|
image_features = image_outputs.last_hidden_state.to(device, dtype=torch.float16)
|
|
image_embeds = image_outputs.image_embeds.to(device, dtype=torch.float16)
|
|
image_embeds /= image_embeds.norm(dim=-1, keepdim=True)
|
|
|
|
captions = []
|
|
for image_idx in range(image_features.size(0)):
|
|
features = image_features[image_idx].unsqueeze(0)
|
|
candidates = self._beam_search(
|
|
features,
|
|
image_embeds[image_idx : image_idx + 1],
|
|
device,
|
|
beam_width=beam_width,
|
|
temperature=temperature,
|
|
ramble=ramble,
|
|
)
|
|
# pick highest scoring candidate
|
|
candidates.sort(key=lambda x: x[0], reverse=True)
|
|
for score, text in candidates:
|
|
logging.debug(f"({score:.3f}) {text}")
|
|
captions.append(candidates[0][1])
|
|
return captions
|
|
|
|
def get_tokenizer(self) -> CLIPTokenizer:
|
|
return self.tokenizer
|
|
|
|
def _beam_search(
|
|
self,
|
|
image_features: torch.Tensor,
|
|
image_embed: torch.Tensor,
|
|
device: torch.device,
|
|
beam_width: int = 5,
|
|
temperature: float = 0.7,
|
|
ramble: bool = False,
|
|
):
|
|
tokenizer = self.tokenizer
|
|
captioner = self.captioner
|
|
token_embedding = self.text_model.embeddings.token_embedding
|
|
pos_embedding = self.text_model.embeddings.position_embedding
|
|
vocab_size = tokenizer.vocab_size
|
|
|
|
# project image features
|
|
memory = captioner.projection(image_features)
|
|
memory = memory + captioner.memory_pos_embedding
|
|
|
|
# start with beam_width copies of BOS token
|
|
current_tokens = torch.full(
|
|
(beam_width, 1), tokenizer.bos_token_id, dtype=torch.long, device=device
|
|
)
|
|
|
|
# track sequence scores
|
|
scores = torch.zeros(beam_width, device=device)
|
|
|
|
for step in range(captioner.max_length - 2):
|
|
# embed current tokens
|
|
token_embeddings = token_embedding(current_tokens)
|
|
positions = torch.arange(current_tokens.size(1), device=device)
|
|
pos_embeddings = pos_embedding(positions)
|
|
x = token_embeddings + pos_embeddings
|
|
|
|
# run decoder layers
|
|
seq_len = x.size(1)
|
|
mask = captioner.causal_mask[:seq_len, :seq_len]
|
|
for layer in captioner.layers:
|
|
x = layer(x, memory.repeat(beam_width, 1, 1), self_attn_mask=mask)
|
|
|
|
# get next token logits and log probabilities
|
|
logits = self.output_projection(x[:, -1:])
|
|
logits = logits / temperature
|
|
if ramble:
|
|
logits[:, :, tokenizer.eos_token_id] = -float("inf")
|
|
log_probs = F.log_softmax(logits, dim=-1)
|
|
|
|
if step == 0:
|
|
# pick top-k tokens for first step
|
|
scores = log_probs.squeeze(1)[0]
|
|
scores, indices = scores.topk(beam_width)
|
|
current_tokens = torch.cat(
|
|
[current_tokens[0:1].repeat(beam_width, 1), indices.unsqueeze(1)],
|
|
dim=1,
|
|
)
|
|
else:
|
|
# find sequences that have already ended
|
|
finished_seqs = (current_tokens == tokenizer.eos_token_id).any(dim=1)
|
|
|
|
# calculate scores for next tokens
|
|
next_scores = scores.unsqueeze(1) + log_probs.squeeze(
|
|
1
|
|
) # [8 sequences x vocab_size options]
|
|
next_scores = next_scores.view(-1) # flatten to [8*vocab_size]
|
|
|
|
# mask out all continuations of finished sequences except first EOS token
|
|
if finished_seqs.any():
|
|
# for each finished sequence, only allow one EOS continuation
|
|
for seq_idx in finished_seqs.nonzero().squeeze(1):
|
|
vocab_start = seq_idx * vocab_size
|
|
# set all continuations to -inf except the EOS token
|
|
next_scores[vocab_start : vocab_start + vocab_size] = float(
|
|
"-inf"
|
|
)
|
|
next_scores[vocab_start + tokenizer.eos_token_id] = scores[
|
|
seq_idx
|
|
]
|
|
|
|
# pick top beam_width sequences
|
|
scores, indices = next_scores.topk(beam_width)
|
|
beam_indices = indices // vocab_size # which sequence each came from
|
|
token_indices = indices % vocab_size # which token to append
|
|
|
|
current_tokens = torch.cat(
|
|
[current_tokens[beam_indices], token_indices.unsqueeze(1)], dim=1
|
|
)
|
|
|
|
# check if all beams ended with EOS
|
|
if (current_tokens[:, -1] == tokenizer.eos_token_id).all():
|
|
break
|
|
|
|
# rank final candidates by CLIP similarity
|
|
candidates = []
|
|
for idx in range(beam_width):
|
|
text = tokenizer.decode(
|
|
current_tokens[idx],
|
|
skip_special_tokens=True,
|
|
clean_up_tokenization_spaces=True,
|
|
)
|
|
text_embeds = self._text_to_embed(text, device)
|
|
clip_sim = torch.sum(image_embed * text_embeds, dim=-1)[0]
|
|
candidates.append((clip_sim.item(), text))
|
|
return candidates
|
|
|
|
def _text_to_embed(self, text: str, device: torch.device) -> torch.Tensor:
|
|
tokens = self.clip_text.tokenize(text)
|
|
self.clip_text.load_model()
|
|
self.clip_text.cond_stage_model.reset_clip_options()
|
|
self.clip_text.cond_stage_model.set_clip_options({"projected_pooled": False})
|
|
clip_l = self.clip_text.cond_stage_model.clip_l
|
|
_, pooled = clip_l.encode_token_weights(tokens["l"])
|
|
text_embeds = self.text_projection(pooled.to(device, dtype=torch.float16))
|
|
text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
|
|
return text_embeds
|
|
|
|
|
|
class CLIPtionLoader:
|
|
CATEGORY = "pharmapsychotic"
|
|
FUNCTION = "load"
|
|
RETURN_TYPES = ("CLIPTION_MODEL",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"clip": (
|
|
"CLIP",
|
|
{"tooltip": "The CLIP model used for encoding the text."},
|
|
),
|
|
"clip_vision": ("CLIP_VISION",),
|
|
}
|
|
}
|
|
|
|
def load(self, clip, clip_vision):
|
|
state_dict = {}
|
|
file = "CLIPtion_20241219_fp16.safetensors"
|
|
base_path = os.path.dirname(os.path.abspath(__file__))
|
|
if os.path.exists(os.path.join(base_path, file)):
|
|
model_path = os.path.join(base_path, file)
|
|
else:
|
|
repo_id = "pharmapsychotic/CLIPtion"
|
|
revision = "15ee8cb77a902616478a033332011ff640e72277"
|
|
model_path = hf_hub_download(
|
|
repo_id=repo_id, filename=file, revision=revision
|
|
)
|
|
with safe_open(model_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
state_dict[key] = f.get_tensor(key)
|
|
tp_dict = {"weight": state_dict.pop("text_projection.weight")}
|
|
|
|
config = SimpleNamespace(
|
|
**{"hidden_dim": 768, "num_heads": 8, "num_blocks": 6, "max_length": 77}
|
|
)
|
|
model = CLIPtionModel(config, clip, clip_vision)
|
|
model.captioner.load_state_dict(state_dict)
|
|
model.text_projection.load_state_dict(tp_dict)
|
|
model.eval()
|
|
model.to(comfy.model_management.get_torch_device(), dtype=torch.float16)
|
|
return (model,)
|
|
|
|
|
|
class CLIPtionGenerate:
|
|
CATEGORY = "pharmapsychotic"
|
|
FUNCTION = "caption"
|
|
OUTPUT_IS_LIST = (True,)
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("CLIPTION_MODEL", {"tooltip": "The CLIPtion model."}),
|
|
"image": ("IMAGE",),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xFFFFFFFFFFFFFFFF,
|
|
"tooltip": "The random seed used for creating the caption.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"temperature": (
|
|
"FLOAT",
|
|
{"default": 0.7, "tooltip": "Temperature for sampling."},
|
|
),
|
|
"best_of": (
|
|
"INT",
|
|
{
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 64,
|
|
"tooltip": "Number of options to evaluate.",
|
|
},
|
|
),
|
|
"ramble": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
def caption(
|
|
self,
|
|
model: CLIPtionModel,
|
|
image: torch.Tensor,
|
|
seed: int,
|
|
temperature: float = 0.7,
|
|
best_of: int = 1,
|
|
ramble: bool = False,
|
|
):
|
|
with torch.inference_mode():
|
|
captions = model.generate(image, seed, temperature, best_of, ramble)
|
|
return (captions,)
|
|
|
|
|
|
class CLIPtionBeamSearch:
|
|
CATEGORY = "pharmapsychotic"
|
|
FUNCTION = "caption"
|
|
OUTPUT_IS_LIST = (True,)
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("CLIPTION_MODEL", {"tooltip": "The CLIPtion model."}),
|
|
"image": ("IMAGE",),
|
|
"beam_width": (
|
|
"INT",
|
|
{
|
|
"default": 4,
|
|
"min": 1,
|
|
"max": 64,
|
|
"tooltip": "Number of beams to maintain during search.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"temperature": (
|
|
"FLOAT",
|
|
{"default": 0.7, "tooltip": "Temperature for token sampling."},
|
|
),
|
|
"ramble": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
def caption(
|
|
self,
|
|
model: CLIPtionModel,
|
|
image: torch.Tensor,
|
|
beam_width: int = 4,
|
|
temperature: float = 0.7,
|
|
ramble: bool = False,
|
|
):
|
|
with torch.inference_mode():
|
|
captions = model.generate_beam(image, temperature, beam_width, ramble)
|
|
return (captions,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CLIPtionBeamSearch": CLIPtionBeamSearch,
|
|
"CLIPtionGenerate": CLIPtionGenerate,
|
|
"CLIPtionLoader": CLIPtionLoader,
|
|
}
|