Update CLIP-style model

Moreover:
- add .gitignore
- clean up .gitattributes
- add SigLIP-like model code
- clean up a bit the Artoria example

.gitattributes

@@ -34,8 +34,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.index filter=lfs diff=lfs merge=lfs -text
-
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class

.gitignore

@@ -0,0 +1,4 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class

Models/CLIP.py

@@ -28,12 +28,18 @@ class Text(flax.linen.Module):
 
 class CLIP(flax.linen.Module):
     out_units: int = 1024
-    logit_scale: float = 1.0
 
     def setup(self):
         self.image_enc = Image(self.out_units)
         self.text_enc = Text(self.out_units)
 
+        self.logit_scale = self.variable(
+            "params",
+            "logit_scale",
+            lambda x: jnp.log(10) * jnp.ones((1,)),
+            None,
+        ).value
+
     @flax.linen.compact
     def __call__(self, image, text, training=False):
         image_emb = self.image_enc(image, training=training)
@@ -43,8 +49,8 @@ class CLIP(flax.linen.Module):
         image_emb = image_emb / jnp.linalg.norm(image_emb, axis=-1, keepdims=True)
         text_emb = text_emb / jnp.linalg.norm(text_emb, axis=-1, keepdims=True)
 
-        image_sim = self.logit_scale * image_emb @ text_emb.T
-        text_sim = self.logit_scale * text_emb @ image_emb.T
+        image_sim = jnp.exp(self.logit_scale) * image_emb @ text_emb.T
+        text_sim = jnp.exp(self.logit_scale) * text_emb @ image_emb.T
         return image_sim, text_sim
 
     def encode_text(self, text):

Models/SigLIP.py

@@ -0,0 +1,63 @@
+import flax
+import jax.numpy as jnp
+
+
+class Image(flax.linen.Module):
+    out_units: int = 1024
+
+    @flax.linen.compact
+    def __call__(self, x, training=False):
+        x = flax.linen.Dropout(0.1)(x, deterministic=not training)
+        return x
+
+
+class Text(flax.linen.Module):
+    out_units: int = 1024
+
+    @flax.linen.compact
+    def __call__(self, x, training=False):
+        x = flax.linen.Dense(features=self.out_units)(x)
+
+        res = flax.linen.silu(x)
+        res = flax.linen.Dense(features=self.out_units)(res)
+        res = flax.linen.Dropout(0.1)(res, deterministic=not training)
+
+        x = x + res
+        return x
+
+
+class CLIP(flax.linen.Module):
+    out_units: int = 1024
+
+    def setup(self):
+        self.image_enc = Image(self.out_units)
+        self.text_enc = Text(self.out_units)
+
+        self.logit_scale = self.variable(
+            "params",
+            "logit_scale",
+            lambda x: jnp.log(10) * jnp.ones((1,)),
+            None,
+        ).value
+        self.logit_bias = self.variable(
+            "params",
+            "logit_bias",
+            lambda x: -10 * jnp.ones((1,)),
+            None,
+        ).value
+
+    @flax.linen.compact
+    def __call__(self, image, text, training=False):
+        image_emb = self.image_enc(image, training=training)
+        text_emb = self.text_enc(text, training=training)
+
+        # Normalize
+        image_emb = image_emb / jnp.linalg.norm(image_emb, axis=-1, keepdims=True)
+        text_emb = text_emb / jnp.linalg.norm(text_emb, axis=-1, keepdims=True)
+
+        logits = image_emb @ text_emb.T * jnp.exp(self.logit_scale) + self.logit_bias
+        return logits
+
+    def encode_text(self, text):
+        text_emb = self.text_enc(text, training=False)
+        return text_emb

app.py

@@ -243,7 +243,7 @@ def main():
                     None,
                     None,
                     "artoria_pendragon_(fate),solo",
-                    "excalibur_(fate/stay_night),green_eyes,monochrome,blonde_hair",
+                    "green_eyes",
                     "CLIP",
                     ["General", "Sensitive"],
                     5,

data/wd-v1-4-convnext-tagger-v2/clip.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0e32b62f6bee5e8db4b17a05d605435dcfa24dc99d0eb26582078f2181567031
-size 48689306
+oid sha256:699734f96af810340b397f1a4e30a2ec40771f4288fc5fb861e7437000984e6d
+size 48689338