fix

.gitignore

@@ -0,0 +1,2 @@
+.venv
+.mypy_cache

.python-version

@@ -0,0 +1 @@
+3.11

app.py

@@ -1,13 +1,14 @@
 import gradio as gr
 import numpy as np
-from transformers import TimmWrapper 
+from PIL import Image
+from transformers import TimmWrapperModel
 
 import torch
-import torchvision.transform.v2 as T
+import torchvision.transforms.v2 as T
 
 
 MODEL_MAP = {
-    "hf_hub:p1atdev/style_250412.vit_base_patch16_siglip_384.v2_webli": {
+    "p1atdev/style_250412.vit_base_patch16_siglip_384.v2_webli": {
         "mean": [0, 0, 0],
         "std": [1.0, 1.0, 1.0],
         "image_size": 384,
@@ -15,9 +16,11 @@ MODEL_MAP = {
     }
 }
 
+
 def config_to_processor(config: dict):
     return T.Compose(
         [
+            T.PILToTensor(),
             T.Resize(
                 size=None,
                 max_size=config["image_size"],
@@ -25,55 +28,60 @@ def config_to_processor(config: dict):
             ),
             T.Pad(
                 padding=config["image_size"] // 2,
-                fill=config["background]",  # black
+                fill=config["background"],
             ),
             T.CenterCrop(
                 size=(config["image_size"], config["image_size"]),
             ),
-            T.PILToTensor(),
-            T.ToDtype(dtype=torch.float32, scale=True), # 0~255 -> 0~1
+            T.ToDtype(dtype=torch.float32, scale=True),  # 0~255 -> 0~1
             T.Normalize(mean=config["mean"], std=config["std"]),
         ]
     )
 
+
 def load_model(name: str):
-    return TimmWrapper.from_pretrained(name).eval().requires_grad_False)
+    return TimmWrapperModel.from_pretrained(name).eval().requires_grad_(False)
+
 
 MODELS = {
     name: {
         "model": load_model(name),
         "processor": config_to_processor(config),
     }
-    for name, config in MODEL_NAMES.items()
+    for name, config in MODEL_MAP.items()
 }
 
 
 @torch.inference_mode()
-def calculate_similarity(model:_name str, image_1: Image.Image, image_2: Image.Image):
+def calculate_similarity(model_name: str, image_1: Image.Image, image_2: Image.Image):
     model = MODELS[model_name]["model"]
     processor = MODELS[model_name]["processor"]
-    
-    pixel_values = torch.cat([
-        processor(image) for image in [image_1, image_2]
-    ])
-    embeddings = model(pixel_values)
+
+    pixel_values = torch.stack([processor(image) for image in [image_1, image_2]])
+
+    embeddings = model(pixel_values).pooler_output
     embeddings /= embeddings.norm(p=2, dim=-1, keepdim=True)
 
-    similarity = (embeddings[0] * embeddings[1]).item()
+    similarity = (embeddings[0] @ embeddings[1].T).item()
+
     return similarity
-    
+
 
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            image_1 = gr.Image("Image 1", type="pil")
-            image_2 = gr.Image("Image 2", type="pil")
+            image_1 = gr.Image(label="Image 1", type="pil")
+            image_2 = gr.Image(label="Image 2", type="pil")
 
-            model_name = gr.Dropdwon("Model", choices=list(MODELS.keys())
+            model_name = gr.Dropdown(
+                label="Model",
+                choices=list(MODELS.keys()),
+                value=list(MODELS.keys())[0],
+            )
             submit_btn = gr.Button("Submit", variant="primary")
-        
+
         with gr.Column():
-            similarity = gr.Text("Similarity")
+            similarity = gr.Label(label="Similarity")
 
     gr.on(
         triggers=[submit_btn.click],
@@ -83,7 +91,7 @@ with gr.Blocks() as demo:
             image_1,
             image_2,
         ],
-        outputs=[image_2],
+        outputs=[similarity],
     )
 
 if __name__ == "__main__":

pyproject.toml

@@ -0,0 +1,20 @@
+[project]
+name = "style-demo"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "hf-xet>=1.0.3",
+    "safetensors>=0.5.3",
+    "timm>=1.0.15",
+    "torch>=2.6.0",
+    "torchvision>=0.21.0",
+    "transformers>=4.51.2",
+]
+
+[dependency-groups]
+dev = [
+    "gradio>=5.25.0",
+    "ruff>=0.11.5",
+]

requirements.txt

@@ -1,4 +1,5 @@
 torch
+torchvision
 transformers
 timm
 safetensors