feat: add code

app.py

@@ -0,0 +1,120 @@
+import numpy as np
+import torch
+
+from transformers import (
+    AutoProcessor,
+)
+
+from PIL import Image
+import gradio as gr
+
+from modeling_siglip import SiglipForImageClassification
+
+MODEL_NAME = "p1atdev/siglip-tagger-test-3"
+PROCESSOR_NAME = "google/siglip-so400m-patch14-384"
+
+model = SiglipForImageClassification.from_pretrained(
+    MODEL_NAME,
+)
+# model = torch.compile(model)
+processor = AutoProcessor.from_pretrained(PROCESSOR_NAME)
+
+
+def compose_text(results: dict[str, float], threshold: float = 0.3):
+    return ", ".join(
+        [
+            key
+            for key, value in sorted(results.items(), key=lambda x: x[1], reverse=True)
+            if value > threshold
+        ]
+    )
+
+
+@torch.no_grad()
+def predict_tags(image: Image.Image, threshold: float):
+    inputs = processor(images=image, return_tensors="pt")
+
+    logits = model(**inputs.to(model.device, model.dtype)).logits.detach().cpu()
+
+    logits = np.clip(logits, 0.0, 1.0)
+
+    results = {}
+
+    for prediction in logits:
+        for i, prob in enumerate(prediction):
+            if prob.item() > 0:
+                results[model.config.id2label[i]] = prob.item()
+
+    return compose_text(results, threshold), results
+
+
+css = """\
+.sticky {
+  position: sticky;
+  top: 16px;
+}
+
+.gradio-container {
+  overflow: clip;
+}
+"""
+
+
+def demo():
+    with gr.Blocks(css=css) as ui:
+        gr.Markdown(
+            """\
+## SigLIP Tagger Test 3
+An experimental model for tagging danbooru tags of images using SigLIP.
+
+Models:
+- (soon)
+
+Example images by NovelAI and niji・journey.
+
+"""
+        )
+
+        with gr.Row():
+            with gr.Column():
+                with gr.Row(elem_classes="sticky"):
+                    with gr.Column():
+                        input_img = gr.Image(
+                            label="Input image", type="pil", height=480
+                        )
+
+                        with gr.Group():
+                            tag_threshold_slider = gr.Slider(
+                                label="Tags threshold",
+                                minimum=0.0,
+                                maximum=1.0,
+                                value=0.3,
+                                step=0.01,
+                            )
+
+                        start_btn = gr.Button(value="Start", variant="primary")
+
+                        gr.Examples(
+                            examples=[["./sample.jpg"], ["./sample2.webp"]],
+                            inputs=[input_img],
+                            cache_examples=False,
+                        )
+
+            with gr.Column():
+                output_tags = gr.Text(label="Output text", interactive=False)
+                output_label = gr.Label(label="Output tags")
+
+        start_btn.click(
+            fn=predict_tags,
+            inputs=[input_img, tag_threshold_slider],
+            outputs=[output_tags, output_label],
+        )
+
+    ui.launch(
+        debug=True,
+        # share=True
+    )
+
+
+if __name__ == "__main__":
+    demo()

modeling_siglip.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+
+from transformers import SiglipVisionModel, SiglipPreTrainedModel, SiglipVisionConfig
+from transformers.utils import ModelOutput
+
+
+@dataclass
+class SiglipForImageClassifierOutput(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    pooler_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+class SiglipForImageClassification(SiglipPreTrainedModel):
+    config_class = SiglipVisionConfig
+    main_input_name = "pixel_values"
+
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.siglip = SiglipVisionModel(config)
+
+        # Classifier head
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels)
+            if config.num_labels > 0
+            else nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self, pixel_values: torch.FloatTensor, labels: torch.LongTensor | None = None
+    ):
+        outputs = self.siglip(pixel_values)
+        pooler_output = outputs.pooler_output
+        logits = self.classifier(pooler_output)
+
+        loss = None
+
+        return SiglipForImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            pooler_output=outputs.pooler_output,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

requirements.txt

@@ -0,0 +1,3 @@
+torch
+accelerate
+transformers==4.37.2