mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 06:10:50 +08:00
179 lines
7.5 KiB
Python
179 lines
7.5 KiB
Python
import json
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from . import clip_model
|
|
from . import model_management
|
|
from . import model_patcher
|
|
from . import ops
|
|
from .image_encoders import dino2
|
|
from .component_model import files
|
|
from .model_management import load_models_gpu
|
|
from .utils import load_torch_file, transformers_convert, state_dict_prefix_replace
|
|
|
|
|
|
class Output:
|
|
def __getitem__(self, key):
|
|
return getattr(self, key)
|
|
|
|
def __setitem__(self, key, item):
|
|
setattr(self, key, item)
|
|
|
|
|
|
def clip_preprocess(image, size=224, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711], crop=True):
|
|
image = image[:, :, :, :3] if image.shape[3] > 3 else image
|
|
mean = torch.tensor(mean, device=image.device, dtype=image.dtype)
|
|
std = torch.tensor(std, device=image.device, dtype=image.dtype)
|
|
image = image.movedim(-1, 1)
|
|
if not (image.shape[2] == size and image.shape[3] == size):
|
|
if crop:
|
|
scale = (size / min(image.shape[2], image.shape[3]))
|
|
scale_size = (round(scale * image.shape[2]), round(scale * image.shape[3]))
|
|
else:
|
|
scale_size = (size, size)
|
|
|
|
image = torch.nn.functional.interpolate(image, size=scale_size, mode="bicubic", antialias=True)
|
|
h = (image.shape[2] - size) // 2
|
|
w = (image.shape[3] - size) // 2
|
|
image = image[:, :, h:h + size, w:w + size]
|
|
image = torch.clip((255. * image), 0, 255).round() / 255.0
|
|
return (image - mean.view([3, 1, 1])) / std.view([3, 1, 1])
|
|
|
|
|
|
IMAGE_ENCODERS = {
|
|
"clip_vision_model": clip_model.CLIPVisionModelProjection,
|
|
"siglip_vision_model": clip_model.CLIPVisionModelProjection,
|
|
"dinov2": dino2.Dinov2Model,
|
|
}
|
|
|
|
|
|
class ClipVisionModel():
|
|
def __init__(self, json_config: dict | str):
|
|
if isinstance(json_config, dict):
|
|
config = json_config
|
|
elif json_config is not None and isinstance(json_config, str):
|
|
if json_config.startswith("{"):
|
|
config = json.loads(json_config)
|
|
else:
|
|
with open(json_config) as f:
|
|
config = json.load(f)
|
|
else:
|
|
raise ValueError(f"json_config had invalid value={json_config}")
|
|
|
|
self.image_size = config.get("image_size", 224)
|
|
self.image_mean = config.get("image_mean", [0.48145466, 0.4578275, 0.40821073])
|
|
self.image_std = config.get("image_std", [0.26862954, 0.26130258, 0.27577711])
|
|
model_type = config.get("model_type", "clip_vision_model")
|
|
model_class = IMAGE_ENCODERS.get(model_type)
|
|
if model_type == "siglip_vision_model":
|
|
self.return_all_hidden_states = True
|
|
else:
|
|
self.return_all_hidden_states = False
|
|
|
|
self.load_device = model_management.text_encoder_device()
|
|
offload_device = model_management.text_encoder_offload_device()
|
|
self.dtype = model_management.text_encoder_dtype(self.load_device)
|
|
self.model = model_class(config, self.dtype, offload_device, ops.manual_cast)
|
|
self.model.eval()
|
|
|
|
self.patcher = model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
|
|
|
|
def load_sd(self, sd):
|
|
return self.model.load_state_dict(sd, strict=False)
|
|
|
|
def get_sd(self):
|
|
return self.model.state_dict()
|
|
|
|
def encode_image(self, image, crop=True):
|
|
load_models_gpu([self.patcher])
|
|
pixel_values = clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=crop).float()
|
|
out = self.model(pixel_values=pixel_values, intermediate_output='all' if self.return_all_hidden_states else -2)
|
|
|
|
outputs = Output()
|
|
outputs["last_hidden_state"] = out[0].to(model_management.intermediate_device())
|
|
outputs["image_embeds"] = out[2].to(model_management.intermediate_device())
|
|
if self.return_all_hidden_states:
|
|
all_hs = out[1].to(model_management.intermediate_device())
|
|
outputs["penultimate_hidden_states"] = all_hs[:, -2]
|
|
outputs["all_hidden_states"] = all_hs
|
|
else:
|
|
outputs["penultimate_hidden_states"] = out[1].to(model_management.intermediate_device())
|
|
outputs["mm_projected"] = out[3]
|
|
return outputs
|
|
|
|
|
|
def convert_to_transformers(sd, prefix):
|
|
sd_k = sd.keys()
|
|
if "{}transformer.resblocks.0.attn.in_proj_weight".format(prefix) in sd_k:
|
|
keys_to_replace = {
|
|
"{}class_embedding".format(prefix): "vision_model.embeddings.class_embedding",
|
|
"{}conv1.weight".format(prefix): "vision_model.embeddings.patch_embedding.weight",
|
|
"{}positional_embedding".format(prefix): "vision_model.embeddings.position_embedding.weight",
|
|
"{}ln_post.bias".format(prefix): "vision_model.post_layernorm.bias",
|
|
"{}ln_post.weight".format(prefix): "vision_model.post_layernorm.weight",
|
|
"{}ln_pre.bias".format(prefix): "vision_model.pre_layrnorm.bias",
|
|
"{}ln_pre.weight".format(prefix): "vision_model.pre_layrnorm.weight",
|
|
}
|
|
|
|
for x in keys_to_replace:
|
|
if x in sd_k:
|
|
sd[keys_to_replace[x]] = sd.pop(x)
|
|
|
|
if "{}proj".format(prefix) in sd_k:
|
|
sd['visual_projection.weight'] = sd.pop("{}proj".format(prefix)).transpose(0, 1)
|
|
|
|
sd = transformers_convert(sd, prefix, "vision_model.", 48)
|
|
else:
|
|
replace_prefix = {prefix: ""}
|
|
sd = state_dict_prefix_replace(sd, replace_prefix)
|
|
return sd
|
|
|
|
|
|
def load_clipvision_from_sd(sd, prefix="", convert_keys=False) -> Optional[ClipVisionModel]:
|
|
json_config: dict = {}
|
|
if convert_keys:
|
|
sd = convert_to_transformers(sd, prefix)
|
|
if "vision_model.encoder.layers.47.layer_norm1.weight" in sd:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_config_g.json")
|
|
elif "vision_model.encoder.layers.30.layer_norm1.weight" in sd:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_config_h.json")
|
|
elif "vision_model.encoder.layers.22.layer_norm1.weight" in sd:
|
|
embed_shape = sd["vision_model.embeddings.position_embedding.weight"].shape[0]
|
|
if sd["vision_model.encoder.layers.0.layer_norm1.weight"].shape[0] == 1152:
|
|
if embed_shape == 729:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_siglip_384.json")
|
|
elif embed_shape == 1024:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_siglip_512.json")
|
|
elif embed_shape == 577:
|
|
if "multi_modal_projector.linear_1.bias" in sd:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_config_vitl_336_llava.json")
|
|
else:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_config_vitl_336.json")
|
|
else:
|
|
json_config = files.get_path_as_dict(None, "clip_vision_config_vitl.json")
|
|
elif "embeddings.patch_embeddings.projection.weight" in sd:
|
|
json_config = files.get_path_as_dict(None, "dino2_giant.json", package="comfy.image_encoders")
|
|
else:
|
|
return None
|
|
|
|
clip = ClipVisionModel(json_config)
|
|
m, u = clip.load_sd(sd)
|
|
if len(m) > 0:
|
|
logging.warning("missing clip vision: {}".format(m))
|
|
u = set(u)
|
|
keys = list(sd.keys())
|
|
for k in keys:
|
|
if k not in u:
|
|
sd.pop(k)
|
|
return clip
|
|
|
|
|
|
def load(ckpt_path):
|
|
sd = load_torch_file(ckpt_path)
|
|
if "visual.transformer.resblocks.0.attn.in_proj_weight" in sd:
|
|
return load_clipvision_from_sd(sd, prefix="visual.", convert_keys=True)
|
|
else:
|
|
return load_clipvision_from_sd(sd)
|