birder-clip 0.0.2.dev6__tar.gz → 0.0.2.dev7__tar.gz

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder_clip
-Version: 0.0.2.dev6
+Version: 0.0.2.dev7
 Summary: A Birder extension for CLIP-style image-text modeling and multimodal computer vision workflows.
 Author: Ofer Hasson
 License-Expression: Apache-2.0
@@ -38,7 +38,7 @@ Requires-Dist: bandit~=1.9.4; extra == "dev"
 Requires-Dist: black~=26.5.0; extra == "dev"
 Requires-Dist: build~=1.5.0; extra == "dev"
 Requires-Dist: bumpver~=2026.1132; extra == "dev"
-Requires-Dist: coverage~=7.14.1; extra == "dev"
+Requires-Dist: coverage~=7.14.2; extra == "dev"
 Requires-Dist: debugpy; extra == "dev"
 Requires-Dist: flake8-pep585~=0.1.7; extra == "dev"
 Requires-Dist: flake8~=7.3.0; extra == "dev"

birder_clip-0.0.2.dev7/birder_clip/inference/data_parallel.py

@@ -0,0 +1,118 @@
+"""
+Inference-optimized multi-GPU parallelization for image-text models
+
+This module provides ZeroShotInferenceDataParallel, a CLIP-style zero-shot
+specialization of Birder's InferenceDataParallel.
+"""
+
+from typing import Optional
+
+import torch
+from birder.inference.data_parallel import InferenceDataParallel
+
+from birder_clip.inference.zero_shot import ZeroShotInference
+from birder_clip.net.base import BaseNet
+
+
+class ZeroShotInferenceDataParallel(InferenceDataParallel):
+    """
+    Distributes zero-shot image inference batches across multiple GPUs
+
+    This wrapper scatters the image batch across devices and keeps a replicated
+    copy of the zero-shot text embeddings on each device. Each replica computes
+    image embeddings and zero-shot logits locally before outputs are gathered.
+
+    Important
+    ---------
+    This class assumes the model is already configured for inference mode
+    (i.e., loaded with inference=True in load_model or manually set to eval mode
+    with requires_grad=False on all parameters).
+    """
+
+    def __init__(
+        self,
+        module: BaseNet,
+        text_embeddings: torch.Tensor,
+        device_ids: Optional[list[int]] = None,
+        output_device: Optional[int | str | torch.device] = None,
+        compile_replicas: bool = False,
+        compile_methods: Optional[list[str]] = None,
+        compile_mode: Optional[str] = None,
+    ) -> None:
+        if compile_methods is None:
+            compile_methods = ["encode_image", "forward_logits"]
+
+        super().__init__(
+            module,
+            device_ids=device_ids,
+            output_device=output_device,
+            compile_replicas=compile_replicas,
+            compile_methods=compile_methods,
+            compile_mode=compile_mode,
+        )
+        self.set_text_embeddings(text_embeddings)
+
+    def set_text_embeddings(self, text_embeddings: torch.Tensor) -> None:
+        if text_embeddings.ndim != 2:
+            raise ValueError(
+                f"text_embeddings must be a 2D tensor of shape (num_classes, embedding_size), "
+                f"got shape {text_embeddings.size()}"
+            )
+
+        self.text_embeddings = [
+            text_embeddings.to(f"cuda:{device_id}", non_blocking=True) for device_id in self.device_ids
+        ]
+        self.inference_modules = [
+            ZeroShotInference(replica, embeddings) for replica, embeddings in zip(self.replicas, self.text_embeddings)
+        ]
+
+    def forward(  # type: ignore[override] # pylint: disable=arguments-differ
+        self, inputs: torch.Tensor, *, tta: bool = False, return_logits: bool = False
+    ) -> torch.Tensor:
+        """
+        Run zero-shot inference distributed across GPUs
+
+        Parameters
+        ----------
+        inputs
+            Input image batch to process.
+        tta
+            Run inference with oversampling.
+        return_logits
+            If True, return raw logits instead of probabilities after softmax.
+        """
+
+        if len(self.device_ids) == 1:
+            output = self.inference_modules[0](
+                inputs,
+                tta=tta,
+                return_logits=return_logits,
+            )
+            return self._gather([output])
+
+        scattered = self._scatter(inputs, {})
+
+        outputs = []
+        for inference, (input_chunk, _), device_id in zip(self.inference_modules, scattered, self.device_ids):
+            if input_chunk is not None and input_chunk.size(0) > 0:
+                with torch.cuda.device(device_id):
+                    output = inference(
+                        input_chunk,
+                        tta=tta,
+                        return_logits=return_logits,
+                    )
+                    outputs.append(output)
+            else:
+                outputs.append(None)
+
+        return self._gather(outputs)
+
+    def __repr__(self) -> str:
+        return (
+            f"ZeroShotInferenceDataParallel(\n"
+            f"  devices={self.device_ids},\n"
+            f"  output_device={self.output_device},\n"
+            f"  src_device={self.src_device},\n"
+            f"  text_embeddings_shape={tuple(self.text_embeddings[0].shape)}\n"
+            f")"
+        )

birder_clip-0.0.2.dev7/birder_clip/inference/image_embeddings.py

@@ -0,0 +1,63 @@
+import sys
+from collections.abc import Iterator
+from typing import Optional
+
+import numpy as np
+import numpy.typing as npt
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from birder_clip.net.base import BaseNet
+
+DataloaderInferenceResult = tuple[list[str], npt.NDArray[np.float32]]
+
+
+def infer_dataloader_iter(
+    device: torch.device,
+    net: BaseNet,
+    dataloader: DataLoader,
+    model_dtype: torch.dtype = torch.float32,
+    amp: bool = False,
+    amp_dtype: Optional[torch.dtype] = None,
+    num_samples: Optional[int] = None,
+    chunk_size: Optional[float] = None,
+) -> Iterator[DataloaderInferenceResult]:
+    if chunk_size is None:
+        chunk_size = float("inf")
+
+    net.to(device, dtype=model_dtype)
+    embeddings_list: list[npt.NDArray[np.float32]] = []
+    sample_paths: list[str] = []
+    sample_count = 0
+    with tqdm(total=num_samples, initial=0, unit="images", unit_scale=True, leave=False) as progress:
+        for file_paths, inputs, _targets in dataloader:
+            batch_size = inputs.size(0)
+
+            # Inference
+            inputs = inputs.to(device, dtype=model_dtype)
+            with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
+                embeddings = net.encode_image(inputs, normalize=True)
+                embeddings = embeddings.cpu().float().numpy()
+
+            embeddings_list.append(embeddings)
+
+            # Set sample list
+            sample_paths.extend(file_paths)
+
+            # Update progress bar
+            progress.update(n=batch_size)
+
+            # Yield results when we reach chunk_size
+            sample_count += batch_size
+            if sample_count >= chunk_size:
+                with tqdm.external_write_mode(file=sys.stderr):
+                    yield (sample_paths, np.concatenate(embeddings_list, axis=0))
+
+                # Reset for next chunk
+                embeddings_list = []
+                sample_paths = []
+                sample_count = 0
+
+    if len(embeddings_list) > 0:
+        yield (sample_paths, np.concatenate(embeddings_list, axis=0))

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev7}/birder_clip/inference/zero_shot.py

@@ -13,6 +13,7 @@ from collections.abc import Callable
 from collections.abc import Iterator
 from collections.abc import Sequence
 from typing import Optional
+from typing import Protocol
 
 import numpy as np
 import numpy.typing as npt
@@ -27,6 +28,42 @@ from birder_clip.net.base import BaseNet
 from birder_clip.tokenizers.base import Tokenizer
 
 
+class ZeroShotInferenceModule(Protocol):
+    def __call__(self, inputs: torch.Tensor, *, tta: bool = False, return_logits: bool = False) -> torch.Tensor: ...
+
+
+class ZeroShotInference:
+    def __init__(self, net: BaseNet, text_embeddings: torch.Tensor) -> None:
+        self.net = net
+        self.text_embeddings = text_embeddings
+
+    def __call__(self, inputs: torch.Tensor, *, tta: bool = False, return_logits: bool = False) -> torch.Tensor:
+        inputs = inputs.to(self.text_embeddings.device, non_blocking=True)
+        if tta is True:
+            _, _, H, W = inputs.size()
+            crop_h = int(H * 0.8)
+            crop_w = int(W * 0.8)
+            tta_inputs = five_crop(inputs, size=[crop_h, crop_w])
+            t = v2.Resize((H, W), interpolation=v2.InterpolationMode.BICUBIC, antialias=True)
+            outs = []
+            for tta_input in tta_inputs:
+                image_embeddings = self.net.encode_image(t(tta_input), normalize=True)
+                logits = self.net.forward_logits(image_embeddings, self.text_embeddings)
+                if return_logits is True:
+                    outs.append(logits)
+                else:
+                    outs.append(F.softmax(logits, dim=-1))
+
+            return torch.stack(outs).mean(dim=0)
+
+        image_embeddings = self.net.encode_image(inputs, normalize=True)
+        logits = self.net.forward_logits(image_embeddings, self.text_embeddings)
+        if return_logits is True:
+            return logits
+
+        return F.softmax(logits, dim=-1)
+
+
 def render_prompts(class_names: Sequence[str], templates: Sequence[str]) -> list[str]:
     return [template.format(class_name) for class_name in class_names for template in templates]
 
@@ -66,39 +103,10 @@ def build_class_text_embeddings(
 DataloaderInferenceResult = tuple[list[str], npt.NDArray[np.float32], npt.NDArray[np.int64]]
 
 
-def infer_batch(
-    net: BaseNet, inputs: torch.Tensor, text_embeddings: torch.Tensor, tta: bool = False, return_logits: bool = False
-) -> torch.Tensor:
-    if tta is True:
-        _, _, H, W = inputs.size()
-        crop_h = int(H * 0.8)
-        crop_w = int(W * 0.8)
-        tta_inputs = five_crop(inputs, size=[crop_h, crop_w])
-        t = v2.Resize((H, W), interpolation=v2.InterpolationMode.BICUBIC, antialias=True)
-        outs = []
-        for tta_input in tta_inputs:
-            image_embeddings = net.encode_image(t(tta_input), normalize=True)
-            logits = net.forward_logits(image_embeddings, text_embeddings)
-            if return_logits is True:
-                outs.append(logits)
-            else:
-                outs.append(F.softmax(logits, dim=-1))
-
-        return torch.stack(outs).mean(dim=0)
-
-    image_embeddings = net.encode_image(inputs, normalize=True)
-    logits = net.forward_logits(image_embeddings, text_embeddings)
-    if return_logits is True:
-        return logits
-
-    return F.softmax(logits, dim=-1)
-
-
 def infer_dataloader_iter(
     device: torch.device,
-    net: BaseNet,
+    net: ZeroShotInferenceModule,
     dataloader: DataLoader,
-    text_embeddings: torch.Tensor,
     tta: bool = False,
     return_logits: bool = False,
     model_dtype: torch.dtype = torch.float32,
@@ -111,7 +119,6 @@ def infer_dataloader_iter(
     if chunk_size is None:
         chunk_size = float("inf")
 
-    net.to(device, dtype=model_dtype)
     out_list: list[npt.NDArray[np.float32]] = []
     labels_list: list[npt.NDArray[np.int64]] = []
     sample_paths: list[str] = []
@@ -121,9 +128,9 @@ def infer_dataloader_iter(
             batch_size = inputs.size(0)
 
             # Inference
-            inputs = inputs.to(device, dtype=model_dtype)
+            inputs = inputs.to(dtype=model_dtype)
             with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
-                out = infer_batch(net, inputs, text_embeddings, return_logits=return_logits, tta=tta)
+                out = net(inputs, return_logits=return_logits, tta=tta)
                 out = out.cpu().float().numpy()
 
             out_list.append(out)

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev7}/birder_clip/net/clip.py

@@ -557,6 +557,29 @@ registry.register_model_config(
     },
 )
 
+# EVA02 CLIP - https://arxiv.org/abs/2303.15389
+registry.register_model_config(
+    "eva02_clip_l14",
+    CLIP,
+    config={
+        "image": {
+            "network": "rope_i_vit_l14_nf_swiglu_c1",
+            "config": {"drop_path_rate": 0.0},
+            "size": (336, 336),
+        },
+        "text": {
+            "network": "transformer_encoder",
+            "config": {
+                "hidden_dim": 768,
+                "num_heads": 12,
+                "output_dim": 768,
+            },
+        },
+        "embed_dim": 768,
+        "tokenizer": "openai_clip_bpe",
+    },
+)
+
 # MetaCLIP - https://arxiv.org/abs/2309.16671
 registry.register_model_config(
     "metaclip_v1_fullcc25b_b16",

{birder_clip-0.0.2.dev6 → birder_clip-0.0.2.dev7}/birder_clip/net/coca.py

@@ -329,6 +329,38 @@ registry.register_model_config(
     },
 )
 
+registry.register_model_config(
+    "laion_coca_vit_b32",
+    CoCa,
+    config={
+        "image": {
+            "network": "vit_b32_pn",
+            "size": (224, 224),
+        },
+        "text": {
+            "network": "transformer_encoder",
+            "config": {
+                "hidden_dim": 512,
+                "num_heads": 8,
+                "output_dim": 512,
+                "class_token": True,
+                "norm_after_pool": True,
+                "pool_type": "last",
+            },
+            "context_length": 76,
+        },
+        "decoder": {
+            "network": "conditioned_decoder",
+            "config": {
+                "num_heads": 8,
+            },
+            "context_length": 76,
+        },
+        "visual_pool": {},
+        "embed_dim": 512,
+        "tokenizer": "openai_clip_bpe",
+    },
+)
 registry.register_model_config(
     "laion_coca_vit_l14",
     CoCa,
@@ -362,6 +394,40 @@ registry.register_model_config(
     },
 )
 
+# LAION CoCa - laion/CoCa-ViT-B-32-laion2B-s13B-b90k
+registry.register_model_config(
+    "laion_coca_vit_b32_openclip_legacy",
+    OpenCLIPLegacyCoCa,
+    config={
+        "image": {
+            "network": "vit_b32_pn",
+            "size": (224, 224),
+        },
+        "text": {
+            "network": "transformer_encoder",
+            "config": {
+                "hidden_dim": 512,
+                "num_heads": 8,
+                "output_dim": 512,
+                "class_token": True,
+                "norm_after_pool": True,
+                "pool_type": "last",
+            },
+            "context_length": 76,
+        },
+        "decoder": {
+            "network": "conditioned_decoder",
+            "config": {
+                "num_heads": 8,
+            },
+            "context_length": 76,
+        },
+        "visual_pool": {},
+        "embed_dim": 512,
+        "tokenizer": "openai_clip_bpe",
+    },
+)
+
 # LAION CoCa - laion/CoCa-ViT-L-14-laion2B-s13B-b90k
 registry.register_model_config(
     "laion_coca_vit_l14_openclip_legacy",

birder_clip-0.0.2.dev7/birder_clip/scripts/__main__.py

@@ -0,0 +1,25 @@
+import pkgutil
+
+import birder_clip.scripts
+
+
+def list_scripts() -> list[str]:
+    scripts = []
+    for _, name, is_pkg in pkgutil.iter_modules(birder_clip.scripts.__path__):
+        if name.startswith("_") is True:  # Skip private modules
+            continue
+
+        if is_pkg is False:
+            scripts.append(name)
+
+    return scripts
+
+
+def main() -> None:
+    print("Available birder scripts:")
+    for script in list_scripts():
+        print(f"    birder_clip.scripts.{script}")
+
+
+if __name__ == "__main__":
+    main()