import logging import os import re from types import SimpleNamespace from typing import Optional, Tuple 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 import comfy.model_management import folder_paths from comfy_api.latest import ComfyExtension, io # register comfyui/models/cliption as a known model folder so the # downloaded decoder weights are cached alongside ComfyUI's other models _CLIPTION_MODELS_DIR = os.path.join(folder_paths.models_dir, "cliption") os.makedirs(_CLIPTION_MODELS_DIR, exist_ok=True) folder_paths.add_model_folder_path("cliption", _CLIPTION_MODELS_DIR) def _fix_punctuation_spacing(text: str) -> str: """Remove stray whitespace before commas, periods, and apostrophes left over from CLIPTokenizer's BPE decoding (e.g. "cats , dogs . it ' s" -> "cats, dogs. it's").""" return re.sub(r"\s+([,.'])", r"\1", text) 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: torch.Tensor): # 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.hidden_dim = config.hidden_dim self.max_length = config.max_length # projection from ViT dimension to decoder dimension self.projection = nn.Linear(vision_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) class CLIPtionModel(nn.Module): def __init__(self, config, clip, clip_vision, device=None): 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 # use specified device, or fall back to ComfyUI default at runtime self.inference_device = device self.tokenizer = clip.tokenizer.clip_l.tokenizer self.text_model = clip.cond_stage_model.clip_l.transformer.text_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()) def generate_beam(self, images: torch.Tensor, beam_width: int = 4) -> list: device = self.inference_device or comfy.model_management.get_torch_device() image_features, image_embeds = self._images_to_embeds(images, device) 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 ) # pick highest scoring candidate candidates.sort(key=lambda x: x[0]) for score, text in candidates: logging.debug(f"({score:.3f}) {text}") captions.append(candidates[-1][1]) return captions def _beam_search( self, image_features: torch.Tensor, image_embed: torch.Tensor, device: torch.device, beam_width: int, ): 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 # determine dtype from captioner weights to handle FP32 models correctly model_dtype = next(captioner.parameters()).dtype memory = captioner.projection(image_features.to(dtype=model_dtype)) 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 ) 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) # cast to model dtype to handle FP32/FP16 mismatch x = (token_embeddings + pos_embeddings).to(dtype=model_dtype) # 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 log probabilities logits = self.output_projection(x[:, -1:]) 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: # calculate scores for next tokens [beam_width x vocab_size] next_scores = scores.unsqueeze(1) + log_probs.squeeze(1) # force sequences to continue EOS after first one prev_is_eos = current_tokens[:, -1] == tokenizer.eos_token_id vocab_mask = torch.zeros_like(next_scores) vocab_mask[prev_is_eos] = float("-inf") vocab_mask[prev_is_eos, tokenizer.eos_token_id] = 0 next_scores = next_scores + vocab_mask # pick top beam_width sequences next_scores = next_scores.view(-1) 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 # add final EOS token current_tokens = torch.cat( [current_tokens, torch.full((beam_width, 1), tokenizer.eos_token_id, device=device)], dim=1, ) # rank final candidates by CLIP similarity candidates = [] for idx in range(beam_width): tokens = current_tokens[idx] # trim everything after the first EOS token (inclusive) before decoding eos_positions = (tokens == tokenizer.eos_token_id).nonzero(as_tuple=True)[0] if len(eos_positions) > 0: tokens = tokens[: eos_positions[0]] text = tokenizer.decode(tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True) text = _fix_punctuation_spacing(text) 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 _images_to_embeds(self, images: torch.Tensor, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: if images.size(-1) == 1: images = images.repeat(1, 1, 1, 3) elif images.size(-1) == 4: images = images[..., :3] outputs = self.clip_vision.encode_image(images) # features go into the FP16 CLIPtion decoder, so cast to FP16 features = outputs.last_hidden_state.to(device, dtype=torch.float16) if features.size(2) != 1024: raise ValueError( f"Expected image features to have 1024 dimensions but got {features.size(2)}. Please ensure you are using CLIP L." ) # embeds are used only for CLIP similarity scoring, preserve original dtype (FP32) for accuracy embeds = outputs.image_embeds.to(device) embeds = embeds / embeds.norm(dim=-1, keepdim=True) return features, embeds def _text_to_embed(self, text: str, device: torch.device) -> torch.Tensor: # load CLIP model and disable final projection since that's missing from comfy checkpoints 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}) # calculate text embedding tokens = self.clip_text.tokenize(text) clip_l = self.clip_text.cond_stage_model.clip_l _, pooled = clip_l.encode_token_weights(tokens["l"]) # preserve original dtype (FP32) for accurate CLIP similarity scoring text_embeds = self.text_projection(pooled.to(device)) text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True) return text_embeds class _DecoderCache: model: Optional["CLIPtionModel"] = None @classmethod def is_loaded(cls) -> bool: return cls.model is not None @classmethod def set(cls, model: "CLIPtionModel") -> None: cls.model = model @classmethod def clear(cls) -> None: cls.model = None class CLIPtionBeamSearchIntegrated(io.ComfyNode): # CLIPtion decoder config (fixed for the released checkpoint) _CAPTIONER_CONFIG = SimpleNamespace(hidden_dim=768, num_heads=8, num_blocks=6, max_length=77) _SAFETENSORS_FILE = "CLIPtion_20241219_fp16.safetensors" _HF_REPO_ID = "easygoing0114/ComfyUI-use-models" _HF_REVISION = "4158578ad9da6f1a54338f00e4310af7f66b5eb7" @classmethod def define_schema(cls) -> io.Schema: return io.Schema( node_id="CLIPtionBeamSearchIntegrated", display_name="CLIPtion Beam Search (Integrated)", category="pharmapsychotic", description="Loads the CLIPtion decoder on demand and runs beam search to caption image(s), " "ranking candidates by CLIP similarity to the input image.", inputs=[ io.Clip.Input("clip", tooltip="CLIP text encoder (must include CLIP-L)."), io.Custom("CLIP_VISION").Input( "clip_vision", tooltip="CLIP vision encoder (must be CLIP-L)." ), io.Image.Input("image"), io.Int.Input( "beam_width", default=9, min=1, max=64, tooltip="Number of beams to maintain during search.", ), io.Boolean.Input( "unload_after_run", default=True, tooltip="Unload CLIPtion decoder from VRAM after captioning.", ), io.Boolean.Input( "force_cpu", default=False, optional=True, advanced=True, tooltip="Run the CLIPtion decoder on CPU instead of the default ComfyUI device.", ), ], outputs=[ io.String.Output(is_output_list=True), ], ) @classmethod def execute( cls, clip, clip_vision, image: torch.Tensor, beam_width: int = 4, unload_after_run: bool = False, force_cpu: bool = False, ) -> io.NodeOutput: try: cls._load_model(clip, clip_vision, force_cpu) with torch.inference_mode(): captions = _DecoderCache.model.generate_beam(image, beam_width) finally: if unload_after_run: cls._unload_model() return io.NodeOutput(captions) @classmethod def _load_model(cls, clip, clip_vision, force_cpu: bool): """Load CLIPtion decoder from disk and move to target device.""" if _DecoderCache.is_loaded(): return # 1) allow a manually placed copy anywhere under comfyui/models/cliption # (covers extra_model_paths / multiple registered search paths too) existing_paths = folder_paths.get_filename_list("cliption") if cls._SAFETENSORS_FILE in existing_paths: model_path = folder_paths.get_full_path("cliption", cls._SAFETENSORS_FILE) else: # 2) fall back to a copy sitting next to this node (legacy location) base_path = os.path.dirname(os.path.abspath(__file__)) local_path = os.path.join(base_path, cls._SAFETENSORS_FILE) if os.path.exists(local_path): model_path = local_path else: # 3) download into comfyui/models/cliption and cache it there model_path = hf_hub_download( repo_id=cls._HF_REPO_ID, filename=cls._SAFETENSORS_FILE, revision=cls._HF_REVISION, local_dir=_CLIPTION_MODELS_DIR, ) state_dict = {} 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")} inference_device = torch.device("cpu") if force_cpu else None model = CLIPtionModel(cls._CAPTIONER_CONFIG, clip, clip_vision, device=inference_device) model.captioner.load_state_dict(state_dict) model.text_projection.load_state_dict(tp_dict) model.eval() load_device = inference_device or comfy.model_management.get_torch_device() model.to(load_device, dtype=torch.float16) # keep text_projection in FP32 to preserve CLIP similarity scoring accuracy model.text_projection.to(dtype=torch.float32) _DecoderCache.set(model) logging.info(f"{cls.__name__}: decoder loaded on {load_device}") @classmethod def _unload_model(cls): """Remove the CLIPtion decoder from VRAM and CPU RAM completely.""" if not _DecoderCache.is_loaded(): return _DecoderCache.clear() # works across CUDA/XPU/MPS, unlike calling torch.cuda directly comfy.model_management.soft_empty_cache() logging.info(f"{cls.__name__}: decoder unloaded") class CLIPtionIntegratedExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: return [CLIPtionBeamSearchIntegrated] async def comfy_entrypoint() -> CLIPtionIntegratedExtension: return CLIPtionIntegratedExtension()