Upload model

config.json

@@ -0,0 +1,17 @@
+{
+  "architectures": [
+    "CSDModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.CSDConfig",
+    "AutoModel": "model.CSDModel"
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.2",
+  "vit_heads": 16,
+  "vit_input_resolution": 224,
+  "vit_layers": 24,
+  "vit_output_dim": 768,
+  "vit_patch_size": 14,
+  "vit_width": 1024
+}

config.py

@@ -0,0 +1,22 @@
+from transformers import PretrainedConfig
+
+
+class CSDConfig(PretrainedConfig):
+    def __init__(
+        self,
+        vit_input_resolution: int = 224,
+        vit_patch_size: int = 14,
+        vit_width: int = 1024,
+        vit_layers: int = 24,
+        vit_heads: int = 16,
+        vit_output_dim: int = 768,
+        **kwargs
+    ) -> None:
+        super(CSDConfig, self).__init__(**kwargs)
+
+        self.vit_input_resolution = vit_input_resolution
+        self.vit_patch_size = vit_patch_size
+        self.vit_width = vit_width
+        self.vit_layers = vit_layers
+        self.vit_heads = vit_heads
+        self.vit_output_dim = vit_output_dim

csd.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from copy import deepcopy
+from clip.model import VisionTransformer
+from typing import Tuple
+
+
+class CSD(nn.Module):
+    def __init__(
+        self,
+        vit_input_resolution: int = 224,
+        vit_patch_size: int = 14,
+        vit_width: int = 1024,
+        vit_layers: int = 768,
+        vit_heads: int = 16,
+        vit_output_dim: int = 768,
+    ) -> None:
+        super(CSD, self).__init__()
+
+        self.backbone = VisionTransformer(
+            input_resolution=vit_input_resolution,
+            patch_size=vit_patch_size,
+            width=vit_width,
+            layers=vit_layers,
+            heads=vit_heads,
+            output_dim=vit_output_dim,
+        )
+
+        self.last_layer_style = deepcopy(self.backbone.proj)
+        self.last_layer_content = deepcopy(self.backbone.proj)
+        self.backbone.proj = None
+
+    def forward(self, pixel_values: torch.Tensor) -> Tuple[torch.Tensor]:
+        features = self.backbone(pixel_values)
+
+        style_output = features @ self.last_layer_style
+        style_output = F.normalize(style_output, dim=1, p=2)
+
+        content_output = features @ self.last_layer_content
+        content_output = F.normalize(content_output, dim=1, p=2)
+
+        return features, content_output, style_output

model.py

@@ -0,0 +1,36 @@
+import torch
+
+from typing import Tuple
+from dataclasses import dataclass
+from transformers import PretrainedConfig, PreTrainedModel
+
+from .csd import CSD
+from .config import CSDConfig
+
+
+@dataclass
+class CSDOutput:
+    image_embeds: torch.Tensor
+    style_embeds: torch.Tensor
+    content_embeds: torch.Tensor
+
+
+class CSDModel(PreTrainedModel):
+    config_class = CSDConfig
+
+    def __init__(self, config: CSDConfig) -> None:
+        super(CSDModel, self).__init__(config)
+
+        self.model = CSD(
+            vit_input_resolution=config.vit_input_resolution,
+            vit_patch_size=config.vit_patch_size,
+            vit_width=config.vit_width,
+            vit_layers=config.vit_layers,
+            vit_heads=config.vit_heads,
+            vit_output_dim=config.vit_output_dim,
+        )
+
+    @torch.inference_mode()
+    def forward(self, pixel_values: torch.Tensor) -> CSDOutput:
+        image_embeds, style_embeds, content_embeds = self.model(pixel_values)
+        return CSDOutput(image_embeds=image_embeds, style_embeds=style_embeds, content_embeds=content_embeds)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4edeb72ee261d99700b654ec40d89484ed3ff02c49a277a63668897a9261914
+size 1219048024