Quantized ImageBind

.gitignore

@@ -5,4 +5,5 @@
 build/
 imagebind.egg-info
 .DS_Store
-venv/
+venv/
+.checkpoints

imagebind/models/helpers.py

@@ -18,6 +18,9 @@ def __init__(self, dim: int) -> None:
         self.dim = dim
 
     def forward(self, x):
+        if x.is_quantized:
+            x = x.dequantize()
+
         return torch.nn.functional.normalize(x, dim=self.dim, p=2)
 
 
@@ -42,8 +45,10 @@ def forward(self, x):
         return torch.clip(self.log_logit_scale.exp(), max=self.max_logit_scale) * x
 
     def extra_repr(self):
-        st = f"logit_scale_init={self.logit_scale_init},learnable={self.learnable}," \
-             f" max_logit_scale={self.max_logit_scale}"
+        st = (
+            f"logit_scale_init={self.logit_scale_init},learnable={self.learnable},"
+            f" max_logit_scale={self.max_logit_scale}"
+        )
         return st