Upload AestheticsPredictorV1

config.json

@@ -4,6 +4,9 @@
     "AestheticsPredictorV1"
   ],
   "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_v1.AestheticsPredictorV1"
+  },
   "dropout": 0.0,
   "hidden_act": "quick_gelu",
   "hidden_size": 1024,

configuration_predictor.py

@@ -0,0 +1,39 @@
+from transformers.models.clip.configuration_clip import CLIPVisionConfig
+
+
+class AestheticsPredictorConfig(CLIPVisionConfig):
+    model_type = "aesthetics_predictor"
+
+    def __init__(
+        self,
+        hidden_size: int = 768,
+        intermediate_size: int = 3072,
+        projection_dim: int = 512,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 12,
+        num_channels: int = 3,
+        image_size: int = 224,
+        patch_size: int = 32,
+        hidden_act: str = "quick_gelu",
+        layer_norm_eps: float = 0.00001,
+        attention_dropout: float = 0,
+        initializer_range: float = 0.02,
+        initializer_factor: float = 1,
+        **kwargs,
+    ):
+        super().__init__(
+            hidden_size,
+            intermediate_size,
+            projection_dim,
+            num_hidden_layers,
+            num_attention_heads,
+            num_channels,
+            image_size,
+            patch_size,
+            hidden_act,
+            layer_norm_eps,
+            attention_dropout,
+            initializer_range,
+            initializer_factor,
+            **kwargs,
+        )

modeling_v1.py

@@ -0,0 +1,63 @@
+from typing import Dict, Final, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from transformers import CLIPVisionModelWithProjection, logging
+from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
+
+from .configuration_predictor import AestheticsPredictorConfig
+
+logging.set_verbosity_error()
+
+URLS: Final[Dict[str, str]] = {
+    "openai/clip-vit-base-patch16": "https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_b_16_linear.pth",
+    "openai/clip-vit-base-patch32": "https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_b_32_linear.pth",
+    "openai/clip-vit-large-patch14": "https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_l_14_linear.pth",
+}
+
+
+class AestheticsPredictorV1(CLIPVisionModelWithProjection):
+    def __init__(self, config: AestheticsPredictorConfig) -> None:
+        super().__init__(config)
+        self.predictor = nn.Linear(config.projection_dim, 1)
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = super().forward(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        image_embeds = outputs[0]  # image_embeds
+        image_embeds /= image_embeds.norm(dim=-1, keepdim=True)
+
+        prediction = self.predictor(image_embeds)
+
+        if not return_dict:
+            return (None, prediction, image_embeds)
+
+        return ImageClassifierOutputWithNoAttention(
+            loss=None,
+            logits=prediction,
+            hidden_states=image_embeds,
+        )
+
+
+def convert_from_openai_clip(openai_model_name: str) -> AestheticsPredictorV1:
+    model = AestheticsPredictorV1.from_pretrained(openai_model_name)
+    state_dict = torch.hub.load_state_dict_from_url(URLS[openai_model_name])
+    model.predictor.load_state_dict(state_dict)
+    model.eval()
+
+    return model