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

birder_clip-0.0.2.dev6/PKG-INFO

@@ -0,0 +1,151 @@
+Metadata-Version: 2.4
+Name: birder_clip
+Version: 0.0.2.dev6
+Summary: A Birder extension for CLIP-style image-text modeling and multimodal computer vision workflows.
+Author: Ofer Hasson
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://gitlab.com/birder/birder-clip
+Project-URL: Issues, https://gitlab.com/birder/birder-clip/-/issues
+Project-URL: Changelog, https://gitlab.com/birder/birder-clip/-/blob/main/CHANGELOG.md
+Keywords: computer-vision,clip,image-text,pytorch,deep-learning,artificial intelligence
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: birder>=0.6.0
+Requires-Dist: ftfy>=6.3.1
+Requires-Dist: regex>=2025.7.29
+Requires-Dist: tqdm>=4.67.0
+Requires-Dist: webdataset>=0.2.111
+Requires-Dist: huggingface_hub
+Requires-Dist: transformers
+Requires-Dist: torch>=2.10.0
+Requires-Dist: torchvision
+Provides-Extra: dev
+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: debugpy; extra == "dev"
+Requires-Dist: flake8-pep585~=0.1.7; extra == "dev"
+Requires-Dist: flake8~=7.3.0; extra == "dev"
+Requires-Dist: invoke~=3.0.3; extra == "dev"
+Requires-Dist: ipython; extra == "dev"
+Requires-Dist: isort~=8.0.1; extra == "dev"
+Requires-Dist: mkdocs~=1.6.1; extra == "dev"
+Requires-Dist: mkdocs-exclude~=1.0.2; extra == "dev"
+Requires-Dist: mypy~=2.1.0; extra == "dev"
+Requires-Dist: parameterized~=0.9.0; extra == "dev"
+Requires-Dist: pylint~=4.0.6; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: requests~=2.34.2; extra == "dev"
+Requires-Dist: safetensors~=0.7.0; extra == "dev"
+Requires-Dist: setuptools; extra == "dev"
+Requires-Dist: twine~=6.2.0; extra == "dev"
+Requires-Dist: types-requests~=2.33.0; extra == "dev"
+Requires-Dist: urllib3~=2.7.0; extra == "dev"
+Requires-Dist: wheel; extra == "dev"
+Dynamic: license-file
+
+# Birder CLIP
+
+Birder CLIP is an early-stage Birder extension for CLIP-style image-text models, focused on practical inference and fine-tuning workflows.
+
+- [Introduction](#introduction)
+- [Setup](#setup)
+- [Getting Started](#getting-started)
+- [Training](#training)
+- [Project Status and Contributions](#project-status-and-contributions)
+- [Licenses](#licenses)
+- [Acknowledgments](#acknowledgments)
+
+## Introduction
+
+Birder CLIP extends [Birder](https://gitlab.com/birder/birder) with image-text models for zero-shot classification, image-text retrieval style workflows, caption generation and related multimodal computer vision tasks.
+
+The project is aimed at image-text modeling rather than general vision-language model (VLM) chat or instruction-following systems.
+It currently includes CLIP-style components, tokenizers, model registry utilities, inference scripts and training code.
+Full training is supported, but for large-scale CLIP pretraining you are probably better served by [OpenCLIP](https://github.com/mlfoundations/open_clip).
+
+## Setup
+
+1. Ensure your environment meets the minimum requirements:
+   - Python 3.11 or newer
+   - PyTorch 2.10 or newer (installed for your hardware/driver stack)
+   - Birder 0.6.0 or newer
+
+1. Install the latest Birder CLIP version:
+
+```sh
+pip install birder-clip
+```
+
+## Getting Started
+
+List available image-text models:
+
+```sh
+python -m birder_clip.tools list-models --image-text
+```
+
+List available pretrained weights:
+
+```sh
+python -m birder_clip.tools list-models --pretrained --verbose
+```
+
+Run zero-shot classification on a directory of images:
+
+```sh
+python -m birder_clip.scripts.zero_shot -n laion_clip_vit_l14 --classes eagle hawk falcon --template-set default --gpu data/images
+```
+
+For detailed options, run:
+
+```sh
+python -m birder_clip.scripts.zero_shot --help
+python -m birder_clip.tools --help
+```
+
+## Training
+
+Birder CLIP includes training support for image-text datasets in CSV and WebDataset formats, including CLIP, CoCa and LiT-style workflows.
+
+## Project Status and Contributions
+
+Birder CLIP is an early alpha project. APIs, model names, checkpoints, training recipes and command-line options may change without notice.
+
+This is currently a personal project in active development. Suggestions, bug reports and feedback are welcome through the project's issue tracker, but the project is not yet stable enough for broad external contributions.
+
+## Licenses
+
+The code in this project is primarily licensed under Apache 2.0. See [LICENSE](LICENSE) for details.
+
+Some model implementations, pretrained weights, tokenizers and converted artifacts may be derived from or depend on projects and datasets with their own licenses and usage restrictions.
+
+**You are responsible for ensuring compliance with all licenses and conditions of any dependent licenses.**
+
+### Disclaimer
+
+If you intend to use Birder CLIP, its pretrained weights, or any associated datasets in a commercial product, we strongly recommend seeking legal advice to ensure compliance with all relevant licenses and terms of use.
+
+It's the user's responsibility to ensure that their use of this project, including any pretrained weights or datasets, complies with all applicable licenses and legal requirements.
+
+## Acknowledgments
+
+Birder CLIP owes much to the work of others in computer vision, image-text representation learning and open-source machine learning.
+
+Special thanks to the [OpenCLIP](https://github.com/mlfoundations/open_clip) project, which serves as the main reference implementation and inspiration for much of the CLIP-style modeling and training work here. The same principle as in Birder applies: this project stands on the shoulders of many open-source projects, papers and datasets. If an attribution is missing, please open an issue.

birder_clip-0.0.2.dev6/README.md

@@ -0,0 +1,89 @@
+# Birder CLIP
+
+Birder CLIP is an early-stage Birder extension for CLIP-style image-text models, focused on practical inference and fine-tuning workflows.
+
+- [Introduction](#introduction)
+- [Setup](#setup)
+- [Getting Started](#getting-started)
+- [Training](#training)
+- [Project Status and Contributions](#project-status-and-contributions)
+- [Licenses](#licenses)
+- [Acknowledgments](#acknowledgments)
+
+## Introduction
+
+Birder CLIP extends [Birder](https://gitlab.com/birder/birder) with image-text models for zero-shot classification, image-text retrieval style workflows, caption generation and related multimodal computer vision tasks.
+
+The project is aimed at image-text modeling rather than general vision-language model (VLM) chat or instruction-following systems.
+It currently includes CLIP-style components, tokenizers, model registry utilities, inference scripts and training code.
+Full training is supported, but for large-scale CLIP pretraining you are probably better served by [OpenCLIP](https://github.com/mlfoundations/open_clip).
+
+## Setup
+
+1. Ensure your environment meets the minimum requirements:
+   - Python 3.11 or newer
+   - PyTorch 2.10 or newer (installed for your hardware/driver stack)
+   - Birder 0.6.0 or newer
+
+1. Install the latest Birder CLIP version:
+
+```sh
+pip install birder-clip
+```
+
+## Getting Started
+
+List available image-text models:
+
+```sh
+python -m birder_clip.tools list-models --image-text
+```
+
+List available pretrained weights:
+
+```sh
+python -m birder_clip.tools list-models --pretrained --verbose
+```
+
+Run zero-shot classification on a directory of images:
+
+```sh
+python -m birder_clip.scripts.zero_shot -n laion_clip_vit_l14 --classes eagle hawk falcon --template-set default --gpu data/images
+```
+
+For detailed options, run:
+
+```sh
+python -m birder_clip.scripts.zero_shot --help
+python -m birder_clip.tools --help
+```
+
+## Training
+
+Birder CLIP includes training support for image-text datasets in CSV and WebDataset formats, including CLIP, CoCa and LiT-style workflows.
+
+## Project Status and Contributions
+
+Birder CLIP is an early alpha project. APIs, model names, checkpoints, training recipes and command-line options may change without notice.
+
+This is currently a personal project in active development. Suggestions, bug reports and feedback are welcome through the project's issue tracker, but the project is not yet stable enough for broad external contributions.
+
+## Licenses
+
+The code in this project is primarily licensed under Apache 2.0. See [LICENSE](LICENSE) for details.
+
+Some model implementations, pretrained weights, tokenizers and converted artifacts may be derived from or depend on projects and datasets with their own licenses and usage restrictions.
+
+**You are responsible for ensuring compliance with all licenses and conditions of any dependent licenses.**
+
+### Disclaimer
+
+If you intend to use Birder CLIP, its pretrained weights, or any associated datasets in a commercial product, we strongly recommend seeking legal advice to ensure compliance with all relevant licenses and terms of use.
+
+It's the user's responsibility to ensure that their use of this project, including any pretrained weights or datasets, complies with all applicable licenses and legal requirements.
+
+## Acknowledgments
+
+Birder CLIP owes much to the work of others in computer vision, image-text representation learning and open-source machine learning.
+
+Special thanks to the [OpenCLIP](https://github.com/mlfoundations/open_clip) project, which serves as the main reference implementation and inspiration for much of the CLIP-style modeling and training work here. The same principle as in Birder applies: this project stands on the shoulders of many open-source projects, papers and datasets. If an attribution is missing, please open an issue.

{birder_clip-0.0.2.dev5 → birder_clip-0.0.2.dev6}/birder_clip/common/lib.py

@@ -43,7 +43,7 @@ def get_image_text_network_name(
     parts = [network]
     if image_encoder is not None:
         parts.append(image_encoder)
-    if text_encoder is not None and text_encoder != "text_transformer":
+    if text_encoder is not None and text_encoder != "transformer_encoder":
         parts.append(text_encoder)
 
     if registry.exists(network) is True:

{birder_clip-0.0.2.dev5 → birder_clip-0.0.2.dev6}/birder_clip/common/training_cli.py

@@ -21,6 +21,11 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("-n", "--network", type=str, help="the image-text network to train")
     parser.add_argument("-t", "--tag", type=str, help="add model tag")
     parser.add_argument("--image-encoder", type=str, help="the image encoder to use")
+    parser.add_argument(
+        "--image-encoder-pretrained",
+        type=str,
+        help="pretrained Birder image model weights path to load into the image encoder",
+    )
     parser.add_argument("--text-encoder", type=str, help="the text encoder to use")
     parser.add_argument("--embed-dim", type=int, metavar="N", help="shared image-text embedding dimension")
     parser.add_argument("--tokenizer", type=str, help="the tokenizer to use")
@@ -43,7 +48,23 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
 
 def add_loss_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Loss parameters")
-    group.add_argument("--loss", type=str, choices=["clip"], default="clip", help="loss function to use")
+    group.add_argument("--loss", type=str, choices=["clip", "coca"], default="clip", help="loss function to use")
+    group.add_argument(
+        "--coca-caption-loss-weight", type=float, default=1.0, help="weight assigned to CoCa caption loss"
+    )
+    group.add_argument(
+        "--coca-contrastive-loss-weight", type=float, default=1.0, help="weight assigned to CoCa contrastive loss"
+    )
+
+
+def add_freeze_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Freeze parameters")
+    group.add_argument(
+        "--freeze-image-encoder",
+        default=False,
+        action="store_true",
+        help="freeze image encoder body, leaving the projection head trainable",
+    )
 
 
 def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: int = 32) -> None:
@@ -66,7 +87,12 @@ def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: i
         metavar="N",
         help="number of iterations to accumulate gradients per optimizer step",
     )
-    # NOTE: Add flag for negative sample caching in grad accum mode
+    group.add_argument(
+        "--grad-accum-cache-negatives",
+        default=False,
+        action="store_true",
+        help="cache features so CLIP loss uses all accumulated microbatches as negatives",
+    )
 
 
 def add_lr_wd_args(parser: argparse.ArgumentParser) -> None:
@@ -250,6 +276,8 @@ def add_data_aug_args(
     )
     group.add_argument("--ra-magnitude", type=int, default=9, help="magnitude for all the RandAugment transformations")
     group.add_argument("--augmix-severity", type=int, default=3, help="severity of AugMix policy")
+    group.add_argument("--clip-color-jitter-prob", type=float, default=0.0, help="CLIP color jitter probability")
+    group.add_argument("--clip-gray-prob", type=float, default=0.0, help="CLIP grayscale probability")
     group.add_argument("--resize-min-scale", type=float, default=default_min_scale, help="random resize min scale")
     group.add_argument(
         "--re-prob",
@@ -356,6 +384,29 @@ def add_precision_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
+def add_grad_checkpointing_args(parser: argparse.ArgumentParser) -> None:
+    group = parser.add_argument_group("Gradient checkpointing parameters")
+    group.add_argument(
+        "--grad-checkpointing",
+        default=False,
+        action="store_true",
+        help="enable gradient checkpointing for supported models",
+    )
+    group.add_argument(
+        "--grad-checkpointing-segments",
+        type=int,
+        metavar="N",
+        help="number of checkpoint segments to request from supported models",
+    )
+    group.add_argument(
+        "--no-grad-checkpointing-preserve-rng-state",
+        dest="grad_checkpointing_preserve_rng_state",
+        default=True,
+        action="store_false",
+        help="disable RNG state preservation during gradient checkpointing recomputation",
+    )
+
+
 def add_compile_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Compilation parameters")
     group.add_argument("--compile", default=False, action="store_true", help="enable compilation")
@@ -578,8 +629,12 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--load-states requires --resume-epoch to be set")
     if args.load_scheduler is True and args.resume_epoch is None:
         raise cli.ValidationError("--load-scheduler requires --resume-epoch to be set")
-    if hasattr(args, "pretrained") is True and args.pretrained is True and args.resume_epoch is not None:
+    if args.pretrained is True and args.resume_epoch is not None:
         raise cli.ValidationError("--pretrained cannot be used with --resume-epoch")
+    if args.image_encoder_pretrained is not None and args.resume_epoch is not None:
+        raise cli.ValidationError("--image-encoder-pretrained cannot be used with --resume-epoch")
+    if args.pretrained is True and args.image_encoder_pretrained is not None:
+        raise cli.ValidationError("--image-encoder-pretrained cannot be used with --pretrained")
 
     if args.freeze_bn is True and args.sync_bn is True:
         raise cli.ValidationError("--freeze-bn cannot be used with --sync-bn")
@@ -605,10 +660,29 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--context-length must be positive")
     if args.grad_accum_steps < 1:
         raise cli.ValidationError("--grad-accum-steps must be >= 1")
+    if args.grad_accum_cache_negatives is True and args.grad_accum_steps == 1:
+        raise cli.ValidationError("--grad-accum-cache-negatives requires --grad-accum-steps greater than 1")
+    if args.grad_accum_cache_negatives is True and args.loss == "coca":
+        raise cli.ValidationError("--grad-accum-cache-negatives is only supported with --loss clip")
+
+    if args.coca_caption_loss_weight < 0.0:
+        raise cli.ValidationError("--coca-caption-loss-weight must be non-negative")
+    if args.coca_contrastive_loss_weight < 0.0:
+        raise cli.ValidationError("--coca-contrastive-loss-weight must be non-negative")
+
+    # EMA
     if args.model_ema_steps < 1:
         raise cli.ValidationError("--model-ema-steps must be >= 1")
 
+    # Gradient checkpointing args
+    if args.grad_checkpointing_segments is not None and args.grad_checkpointing_segments < 1:
+        raise cli.ValidationError("--grad-checkpointing-segments must be >= 1")
+    if args.grad_checkpointing_segments is not None and args.grad_checkpointing is False:
+        raise cli.ValidationError("--grad-checkpointing-segments requires --grad-checkpointing")
+
     if args.distributed_mode == "fsdp":
+        if args.grad_checkpointing is True:
+            raise cli.ValidationError("--grad-checkpointing cannot be used with --distributed-mode fsdp")
         if args.sync_bn is True:
             raise cli.ValidationError("--sync-bn cannot be used with --distributed-mode fsdp")
         if args.find_unused_parameters is True:

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

@@ -19,6 +19,8 @@ import numpy.typing as npt
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
+from torchvision.transforms import v2
+from torchvision.transforms.v2.functional import five_crop
 from tqdm import tqdm
 
 from birder_clip.net.base import BaseNet
@@ -64,11 +66,40 @@ 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 | torch.ScriptModule,
+    net: BaseNet,
     dataloader: DataLoader,
     text_embeddings: torch.Tensor,
+    tta: bool = False,
     return_logits: bool = False,
     model_dtype: torch.dtype = torch.float32,
     amp: bool = False,
@@ -92,12 +123,8 @@ def infer_dataloader_iter(
             # Inference
             inputs = inputs.to(device, dtype=model_dtype)
             with torch.amp.autocast(device.type, enabled=amp, dtype=amp_dtype):
-                image_embeddings = net.encode_image(inputs, normalize=True)
-                logits = net.forward_logits(image_embeddings, text_embeddings)
-                if return_logits is True:
-                    out = logits.cpu().float().numpy()
-                else:
-                    out = F.softmax(logits, dim=-1).cpu().float().numpy()
+                out = infer_batch(net, inputs, text_embeddings, return_logits=return_logits, tta=tta)
+                out = out.cpu().float().numpy()
 
             out_list.append(out)
 

{birder_clip-0.0.2.dev5 → birder_clip-0.0.2.dev6}/birder_clip/loss/__init__.py

@@ -1,5 +1,7 @@
+from birder_clip.loss.coca import CoCaLoss
 from birder_clip.loss.contrastive import CLIPLoss
 
 __all__ = [
+    "CoCaLoss",
     "CLIPLoss",
 ]

birder_clip-0.0.2.dev6/birder_clip/loss/coca.py

@@ -0,0 +1,66 @@
+"""
+CoCa loss, adapted from
+https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/loss.py
+
+Paper "CoCa: Contrastive Captioners are Image-Text Foundation Models",
+https://arxiv.org/abs/2205.01917
+
+Generated by gpt-5.5 xhigh.
+"""
+
+# Reference license: MIT
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from birder_clip.loss.contrastive import CLIPLoss
+
+
+class CoCaLoss(torch.nn.Module):
+    """
+    CoCa contrastive and captioning loss
+
+    Combines the symmetric CLIP contrastive objective with autoregressive
+    captioning cross entropy over decoder logits.
+    """
+
+    def __init__(
+        self, *, caption_loss_weight: float = 1.0, clip_loss_weight: float = 1.0, pad_token_id: int = 0
+    ) -> None:
+        super().__init__()
+        self.caption_loss_weight = caption_loss_weight
+        self.clip_loss_weight = clip_loss_weight
+        self.pad_token_id = pad_token_id
+        self.clip_loss = CLIPLoss()
+
+    def forward(
+        self,
+        image_features: torch.Tensor,
+        text_features: torch.Tensor,
+        logits: torch.Tensor,
+        texts: torch.Tensor,
+        logit_scale: torch.Tensor,
+        logit_bias: Optional[torch.Tensor] = None,
+    ) -> dict[str, torch.Tensor]:
+        if self.clip_loss_weight != 0.0:
+            losses = self.clip_loss(image_features, text_features, logit_scale, logit_bias=logit_bias)
+            contrastive_loss = losses["contrastive_loss"] * self.clip_loss_weight
+        else:
+            contrastive_loss = logits.new_zeros(())
+
+        if self.caption_loss_weight != 0.0:
+            caption_loss = F.cross_entropy(
+                logits[:, :-1].permute(0, 2, 1),
+                texts[:, 1:],
+                ignore_index=self.pad_token_id,
+            )
+            caption_loss = caption_loss * self.caption_loss_weight
+        else:
+            caption_loss = logits.new_zeros(())
+
+        return {
+            "contrastive_loss": contrastive_loss,
+            "caption_loss": caption_loss,
+        }

{birder_clip-0.0.2.dev5 → birder_clip-0.0.2.dev6}/birder_clip/model_registry/model_registry.py

@@ -15,38 +15,44 @@ from birder_clip.tokenizers.registry import get_tokenizer_info
 
 if TYPE_CHECKING is True:
     from birder_clip.net.base import BaseNet  # pylint: disable=cyclic-import
-    from birder_clip.net.text.base import TextBaseNet  # pylint: disable=cyclic-import
+    from birder_clip.net.text.base import TextDecoderBaseNet  # pylint: disable=cyclic-import
+    from birder_clip.net.text.base import TextEncoderBaseNet  # pylint: disable=cyclic-import
 
-    NetType: TypeAlias = type[BaseNet] | type[TextBaseNet]
+    BaseNetObjType: TypeAlias = BaseNet | TextEncoderBaseNet | TextDecoderBaseNet
+    BaseNetType: TypeAlias = type[BaseNet] | type[TextEncoderBaseNet] | type[TextDecoderBaseNet]
 
 
 class Task(StrEnum):
     IMAGE_TEXT = "image_text"
-    TEXT = "text"
+    TEXT_ENCODER = "text_encoder"
+    TEXT_DECODER = "text_decoder"
 
 
 class ModelRegistry:
     def __init__(self) -> None:
-        self.registered_configs: dict[str, "NetType"] = {}
+        self.registered_configs: dict[str, "BaseNetType"] = {}
         self._image_text_nets: dict[str, type["BaseNet"]] = {}
-        self._text_nets: dict[str, type[TextBaseNet]] = {}
+        self._text_encoder_nets: dict[str, type[TextEncoderBaseNet]] = {}
+        self._text_decoder_nets: dict[str, type[TextDecoderBaseNet]] = {}
         self._pretrained_nets = manifest.REGISTRY_MANIFEST
 
     @property
-    def all_nets(self) -> dict[str, "NetType"]:
-        return {**self._image_text_nets, **self._text_nets}
+    def all_nets(self) -> dict[str, "BaseNetType"]:
+        return {**self._image_text_nets, **self._text_encoder_nets, **self._text_decoder_nets}
 
-    def _get_models_for_task(self, task: Task) -> dict[str, "NetType"]:
+    def _get_models_for_task(self, task: Task) -> dict[str, "BaseNetType"]:
         if task == Task.IMAGE_TEXT:
-            nets: dict[str, "NetType"] = self._image_text_nets
-        elif task == Task.TEXT:
-            nets = self._text_nets
+            nets: dict[str, "BaseNetType"] = self._image_text_nets
+        elif task == Task.TEXT_ENCODER:
+            nets = self._text_encoder_nets
+        elif task == Task.TEXT_DECODER:
+            nets = self._text_decoder_nets
         else:
             raise ValueError(f"Unsupported model task: {task}")
 
         return nets
 
-    def _register_model(self, name: str, net_type: "NetType") -> None:
+    def _register_model(self, name: str, net_type: "BaseNetType") -> None:
         name_key = name.lower()
         task = Task(net_type.task)
         nets = self._get_models_for_task(task)
@@ -57,7 +63,9 @@ class ModelRegistry:
 
         nets[name_key] = net_type
 
-    def register_model_config(self, name: str, net_type: "NetType", *, config: Optional[dict[str, Any]] = None) -> None:
+    def register_model_config(
+        self, name: str, net_type: "BaseNetType", *, config: Optional[dict[str, Any]] = None
+    ) -> None:
         name_key = name.lower()
         registered_net_type = type(name, (net_type,), {"config": config})
         self._register_model(name_key, registered_net_type)
@@ -180,17 +188,30 @@ class ModelRegistry:
 
         return metadata
 
-    def text_factory(
+    def text_encoder_factory(
         self, name: str, *, config: Optional[dict[str, Any]] = None, context_length: Optional[int] = None
-    ) -> "TextBaseNet":
+    ) -> "TextEncoderBaseNet":
         name_key = name.lower()
-        return self._text_nets[name_key](config=config, context_length=context_length)
+        if name_key in self._text_encoder_nets:
+            return self._text_encoder_nets[name_key](config=config, context_length=context_length)
+
+        if name_key.startswith("hf:"):
+            config = {"source": name[len("hf:") :], **({} if config is None else config)}
+            name_key = "hf_text_encoder"
+
+        return self._text_encoder_nets[name_key](config=config, context_length=context_length)
+
+    def text_decoder_factory(
+        self, name: str, *, config: Optional[dict[str, Any]] = None, context_length: Optional[int] = None
+    ) -> "TextDecoderBaseNet":
+        name_key = name.lower()
+        return self._text_decoder_nets[name_key](config=config, context_length=context_length)
 
     def net_factory(self, name: str, *, config: Optional[dict[str, Any]] = None) -> "BaseNet":
         name_key = name.lower()
         return self._image_text_nets[name_key](config=config)
 
-    def _metadata_type_name(self, model: "BaseNet | TextBaseNet") -> str:
+    def _metadata_type_name(self, model: "BaseNetObjType") -> str:
         cls = model.__class__
         bases = cls.__bases__
         if len(bases) > 1 and bases[0].__name__ == "FSDPModule":
@@ -198,14 +219,14 @@ class ModelRegistry:
 
         return cls.__name__.lower()
 
-    def get_model_base_name(self, model: "BaseNet | TextBaseNet") -> str:
+    def get_model_base_name(self, model: "BaseNetObjType") -> str:
         type_name = self._metadata_type_name(model)
         if type_name in self.registered_configs:
             type_name = self.registered_configs[type_name].__bases__[0].__name__.lower()
 
         return type_name
 
-    def get_registered_name(self, model: "BaseNet | TextBaseNet") -> Optional[str]:
+    def get_registered_name(self, model: "BaseNetObjType") -> Optional[str]:
         type_name = self._metadata_type_name(model)
         if type_name in self.registered_configs:
             return type_name

{birder_clip-0.0.2.dev5 → birder_clip-0.0.2.dev6}/birder_clip/net/__init__.py

@@ -1,5 +1,7 @@
 from birder_clip.net.clip import CLIP
+from birder_clip.net.coca import CoCa
 
 __all__ = [
     "CLIP",
+    "CoCa",
 ]