diff --git a/comfy/clip_vision.py b/comfy/clip_vision.py
index 297b3bca..87d32a66 100644
--- a/comfy/clip_vision.py
+++ b/comfy/clip_vision.py
@@ -9,6 +9,7 @@ import comfy.model_patcher
import comfy.model_management
import comfy.utils
import comfy.clip_model
+import comfy.image_encoders.dino2
class Output:
def __getitem__(self, key):
@@ -34,6 +35,12 @@ def clip_preprocess(image, size=224, mean=[0.48145466, 0.4578275, 0.40821073], s
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": comfy.clip_model.CLIPVisionModelProjection,
+ "siglip_vision_model": comfy.clip_model.CLIPVisionModelProjection,
+ "dinov2": comfy.image_encoders.dino2.Dinov2Model,
+}
+
class ClipVisionModel():
def __init__(self, json_config):
with open(json_config) as f:
@@ -42,10 +49,11 @@ class ClipVisionModel():
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_class = IMAGE_ENCODERS.get(config.get("model_type", "clip_vision_model"))
self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
- self.model = comfy.clip_model.CLIPVisionModelProjection(config, self.dtype, offload_device, comfy.ops.manual_cast)
+ self.model = model_class(config, self.dtype, offload_device, comfy.ops.manual_cast)
self.model.eval()
self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
@@ -111,6 +119,8 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336.json")
else:
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
+ elif "embeddings.patch_embeddings.projection.weight" in sd:
+ json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
else:
return None
diff --git a/comfy/image_encoders/dino2.py b/comfy/image_encoders/dino2.py
new file mode 100644
index 00000000..130ed6fd
--- /dev/null
+++ b/comfy/image_encoders/dino2.py
@@ -0,0 +1,141 @@
+import torch
+from comfy.text_encoders.bert import BertAttention
+import comfy.model_management
+from comfy.ldm.modules.attention import optimized_attention_for_device
+
+
+class Dino2AttentionOutput(torch.nn.Module):
+ def __init__(self, input_dim, output_dim, layer_norm_eps, dtype, device, operations):
+ super().__init__()
+ self.dense = operations.Linear(input_dim, output_dim, dtype=dtype, device=device)
+
+ def forward(self, x):
+ return self.dense(x)
+
+
+class Dino2AttentionBlock(torch.nn.Module):
+ def __init__(self, embed_dim, heads, layer_norm_eps, dtype, device, operations):
+ super().__init__()
+ self.attention = BertAttention(embed_dim, heads, dtype, device, operations)
+ self.output = Dino2AttentionOutput(embed_dim, embed_dim, layer_norm_eps, dtype, device, operations)
+
+ def forward(self, x, mask, optimized_attention):
+ return self.output(self.attention(x, mask, optimized_attention))
+
+
+class LayerScale(torch.nn.Module):
+ def __init__(self, dim, dtype, device, operations):
+ super().__init__()
+ self.lambda1 = torch.nn.Parameter(torch.empty(dim, device=device, dtype=dtype))
+
+ def forward(self, x):
+ return x * comfy.model_management.cast_to_device(self.lambda1, x.device, x.dtype)
+
+
+class SwiGLUFFN(torch.nn.Module):
+ def __init__(self, dim, dtype, device, operations):
+ super().__init__()
+ in_features = out_features = dim
+ hidden_features = int(dim * 4)
+ hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+
+ self.weights_in = operations.Linear(in_features, 2 * hidden_features, bias=True, device=device, dtype=dtype)
+ self.weights_out = operations.Linear(hidden_features, out_features, bias=True, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x = self.weights_in(x)
+ x1, x2 = x.chunk(2, dim=-1)
+ x = torch.nn.functional.silu(x1) * x2
+ return self.weights_out(x)
+
+
+class Dino2Block(torch.nn.Module):
+ def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations):
+ super().__init__()
+ self.attention = Dino2AttentionBlock(dim, num_heads, layer_norm_eps, dtype, device, operations)
+ self.layer_scale1 = LayerScale(dim, dtype, device, operations)
+ self.layer_scale2 = LayerScale(dim, dtype, device, operations)
+ self.mlp = SwiGLUFFN(dim, dtype, device, operations)
+ self.norm1 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
+ self.norm2 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
+
+ def forward(self, x, optimized_attention):
+ x = x + self.layer_scale1(self.attention(self.norm1(x), None, optimized_attention))
+ x = x + self.layer_scale2(self.mlp(self.norm2(x)))
+ return x
+
+
+class Dino2Encoder(torch.nn.Module):
+ def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations):
+ super().__init__()
+ self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations) for _ in range(num_layers)])
+
+ def forward(self, x, intermediate_output=None):
+ optimized_attention = optimized_attention_for_device(x.device, False, small_input=True)
+
+ if intermediate_output is not None:
+ if intermediate_output < 0:
+ intermediate_output = len(self.layer) + intermediate_output
+
+ intermediate = None
+ for i, l in enumerate(self.layer):
+ x = l(x, optimized_attention)
+ if i == intermediate_output:
+ intermediate = x.clone()
+ return x, intermediate
+
+
+class Dino2PatchEmbeddings(torch.nn.Module):
+ def __init__(self, dim, num_channels=3, patch_size=14, image_size=518, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.projection = operations.Conv2d(
+ in_channels=num_channels,
+ out_channels=dim,
+ kernel_size=patch_size,
+ stride=patch_size,
+ bias=True,
+ dtype=dtype,
+ device=device
+ )
+
+ def forward(self, pixel_values):
+ return self.projection(pixel_values).flatten(2).transpose(1, 2)
+
+
+class Dino2Embeddings(torch.nn.Module):
+ def __init__(self, dim, dtype, device, operations):
+ super().__init__()
+ patch_size = 14
+ image_size = 518
+
+ self.patch_embeddings = Dino2PatchEmbeddings(dim, patch_size=patch_size, image_size=image_size, dtype=dtype, device=device, operations=operations)
+ self.position_embeddings = torch.nn.Parameter(torch.empty(1, (image_size // patch_size) ** 2 + 1, dim, dtype=dtype, device=device))
+ self.cls_token = torch.nn.Parameter(torch.empty(1, 1, dim, dtype=dtype, device=device))
+ self.mask_token = torch.nn.Parameter(torch.empty(1, dim, dtype=dtype, device=device))
+
+ def forward(self, pixel_values):
+ x = self.patch_embeddings(pixel_values)
+ # TODO: mask_token?
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+ x = x + comfy.model_management.cast_to_device(self.position_embeddings, x.device, x.dtype)
+ return x
+
+
+class Dinov2Model(torch.nn.Module):
+ def __init__(self, config_dict, dtype, device, operations):
+ super().__init__()
+ num_layers = config_dict["num_hidden_layers"]
+ dim = config_dict["hidden_size"]
+ heads = config_dict["num_attention_heads"]
+ layer_norm_eps = config_dict["layer_norm_eps"]
+
+ self.embeddings = Dino2Embeddings(dim, dtype, device, operations)
+ self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations)
+ self.layernorm = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
+
+ def forward(self, pixel_values, attention_mask=None, intermediate_output=None):
+ x = self.embeddings(pixel_values)
+ x, i = self.encoder(x, intermediate_output=intermediate_output)
+ x = self.layernorm(x)
+ pooled_output = x[:, 0, :]
+ return x, i, pooled_output, None
diff --git a/comfy/image_encoders/dino2_giant.json b/comfy/image_encoders/dino2_giant.json
new file mode 100644
index 00000000..f6076a4d
--- /dev/null
+++ b/comfy/image_encoders/dino2_giant.json
@@ -0,0 +1,21 @@
+{
+ "attention_probs_dropout_prob": 0.0,
+ "drop_path_rate": 0.0,
+ "hidden_act": "gelu",
+ "hidden_dropout_prob": 0.0,
+ "hidden_size": 1536,
+ "image_size": 518,
+ "initializer_range": 0.02,
+ "layer_norm_eps": 1e-06,
+ "layerscale_value": 1.0,
+ "mlp_ratio": 4,
+ "model_type": "dinov2",
+ "num_attention_heads": 24,
+ "num_channels": 3,
+ "num_hidden_layers": 40,
+ "patch_size": 14,
+ "qkv_bias": true,
+ "use_swiglu_ffn": true,
+ "image_mean": [0.485, 0.456, 0.406],
+ "image_std": [0.229, 0.224, 0.225]
+}