Added Model.from_pretrained method to load a pretrained model

bonsai/models/__init__.py

@@ -6,6 +6,7 @@
 from bonsai.models.llada.modeling import LLaDAModel, ModelConfig as LLaDAModelConfig
 from bonsai.models.mamba2.modeling import Mamba2ForCausalLM, Mamba2Forecaster, Mamba2Model, ModelConfig as Mamba2Config
 from bonsai.models.qwen3.modeling import Qwen3, ModelConfig as Qwen3Config
+from bonsai.models.qwen3_vl.modeling import Qwen3VLForConditionalGeneration, ModelConfig as Qwen3VLConfig
 from bonsai.models.resnet.modeling import ResNet, ModelConfig as ResNetConfig
 from bonsai.models.sam2.modeling import SAM2Base, SAM2ImagePredictor, ModelConfig as SAM2Config
 from bonsai.models.umt5.modeling import UMT5Model, ModelConfig as UMT5Config
@@ -36,6 +37,8 @@
     "Mamba2Model",
     "Qwen3",
     "Qwen3Config",
+    "Qwen3VLConfig",
+    "Qwen3VLForConditionalGeneration",
     "ResNet",
     "ResNetConfig",
     "SAM2Base",

bonsai/models/convnext/modeling.py

@@ -119,6 +119,7 @@ def __init__(
         *,
         rngs: nnx.Rngs,
     ):
+        self.config = cfg
         self.embedding_layer = nnx.Sequential(
             nnx.Conv(cfg.in_channels, cfg.stage_dims[0], cfg.patch_size, cfg.patch_size, rngs=rngs),
             nnx.LayerNorm(cfg.stage_dims[0], epsilon=cfg.layernorm_eps, rngs=rngs),
@@ -141,6 +142,28 @@ def __call__(self, x: jax.Array, *, rngs: jax.Array, train: bool = False):
         x = self.norm(x)
         return self.head(x)
 
+    @classmethod
+    def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
+        """model_name the *model id* of a pretrained model hosted inside
+        a model repo on huggingface.co. For example, "facebook/convnext-small-224"
+        """
+        from huggingface_hub import snapshot_download
+        from bonsai.models.convnext import params
+
+        if config is None:
+            config_map = {
+                "facebook/convnext-tiny-224": ModelConfig.convnext_tiny_224,
+                "facebook/convnext-small-224": ModelConfig.convnext_small_224,
+                "facebook/convnext-base-224": ModelConfig.convnext_base_224,
+                "facebook/convnext-large-224": ModelConfig.convnext_large_224,
+            }
+            if model_name not in config_map:
+                raise ValueError(f"Model name '{model_name}' is unknown, please provide config argument")
+            config = config_map[model_name]()
+
+        model_ckpt_path = snapshot_download(repo_id=model_name, allow_patterns="*.h5")
+        return params.create_convnext_from_pretrained(model_ckpt_path, config)
+
 
 @partial(jax.jit, static_argnames=["graph_def", "train"])
 def forward(

bonsai/models/convnext/tests/run_model.py

@@ -3,31 +3,24 @@
 import jax
 import jax.numpy as jnp
 from flax import nnx
-from huggingface_hub import snapshot_download
 
 from bonsai.models.convnext import modeling as model_lib
-from bonsai.models.convnext import params
 
 
 def run_model():
-    # 1. Download Model Weights
+    # 1. Load Pretrained Model
     model_name = "facebook/convnext-small-224"
-    model_ckpt_path = snapshot_download(repo_id=model_name, allow_patterns="*.h5")
-
-    # 2. Load Pretrained Model
-    config = model_lib.ModelConfig.convnext_small_224()
-    model = params.create_convnext_from_pretrained(model_ckpt_path, config)
-
+    model = model_lib.ConvNeXt.from_pretrained(model_name)
     graphdef, state = nnx.split(model)
 
-    # 3. Prepare dummy input
+    # 2. Prepare dummy input
     batch_size, channels, image_size = 8, 3, 224
     dummy_input = jnp.ones((batch_size, image_size, image_size, channels), dtype=jnp.float32)
 
     key = jax.random.key(0)
     key, warmup_key, prof_key, time_key = jax.random.split(key, 4)
 
-    # 4. Warmup + profiling
+    # 3. Warmup + profiling
 
     _ = model_lib.forward(graphdef, state, dummy_input, rngs=warmup_key, train=False).block_until_ready()
 
@@ -41,7 +34,7 @@ def run_model():
         jax.block_until_ready(logits)
     jax.profiler.stop_trace()
 
-    # 5. Timed execution
+    # 4. Timed execution
 
     time_keys = jax.random.split(time_key, 10)
 
@@ -53,7 +46,7 @@ def run_model():
     print(f"Step time: {step_time:.4f} s")
     print(f"Throughput: {batch_size / step_time:.2f} images/s")
 
-    # 6. Show Top-1 Predicted Class
+    # 5. Show Top-1 Predicted Class
 
     pred = jnp.argmax(logits, axis=-1)
     print("Predicted classes (batch):", pred)

bonsai/models/convnext/tests/test_outputs_ConvNext.py

@@ -3,11 +3,9 @@
 import numpy as np
 import torch
 from absl.testing import absltest, parameterized
-from huggingface_hub import snapshot_download
 from transformers import ConvNextForImageClassification
 
 from bonsai.models.convnext import modeling as model_lib
-from bonsai.models.convnext import params
 
 
 class TestModuleForwardPasses(absltest.TestCase):
@@ -17,10 +15,9 @@ class TestModuleForwardPasses(absltest.TestCase):
     def setUpClass(cls):
         super().setUpClass()
         model_name = "facebook/convnext-large-224"
-        model_ckpt_path = snapshot_download(model_name)
 
-        cls.bonsai_config = model_lib.ModelConfig.convnext_large_224()
-        cls.bonsai_model = params.create_convnext_from_pretrained(model_ckpt_path, cls.bonsai_config)
+        cls.bonsai_model = model_lib.ConvNeXt.from_pretrained(model_name)
+        cls.bonsai_config = cls.bonsai_model.config
         cls.baseline_model = ConvNextForImageClassification.from_pretrained(model_name)
 
         cls.baseline_model.eval()
@@ -74,10 +71,8 @@ class TestModuleFullOtherConfigs(parameterized.TestCase):
     @parameterized.named_parameters(("tiny", "tiny"), ("small", "small"), ("base", "base"))
     def test_full(self, model_size):
         model_name = f"facebook/convnext-{model_size}-224"
-        model_ckpt_path = snapshot_download(model_name)
 
-        bonsai_config = getattr(model_lib.ModelConfig, f"convnext_{model_size}_224")()
-        bonsai_model = params.create_convnext_from_pretrained(model_ckpt_path, bonsai_config)
+        bonsai_model = model_lib.ConvNeXt.from_pretrained(model_name)
         baseline_model = ConvNextForImageClassification.from_pretrained(model_name)
         baseline_model.eval()
 

bonsai/models/densenet121/modeling.py

@@ -34,6 +34,22 @@ def densenet_121(cls):
             growth_rate=32,
         )
 
+    @classmethod
+    def densenet_169(cls):
+        return cls(
+            num_classes=1000,
+            dense_block_layers=[6, 12, 32, 32],
+            growth_rate=32,
+        )
+
+    @classmethod
+    def densenet_201(cls):
+        return cls(
+            num_classes=1000,
+            dense_block_layers=[6, 12, 48, 32],
+            growth_rate=32,
+        )
+
 
 class DenseBlock(nnx.Module):
     def __init__(self, num_layers: int, in_channels: int, growth_rate: int, *, rngs: nnx.Rngs):
@@ -119,6 +135,27 @@ def __call__(self, x):
 
         return x
 
+    @classmethod
+    def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
+        """model_name the *model id* of a pretrained model hosted inside
+        a model repo on huggingface.co. For example, "keras/densenet_121_imagenet"
+        """
+        from huggingface_hub import snapshot_download
+        from bonsai.models.densenet121 import params
+
+        if config is None:
+            config_map = {
+                "keras/densenet_121_imagenet": ModelConfig.densenet_121,
+                "keras/densenet_169_imagenet": ModelConfig.densenet_169,
+                "keras/densenet_201_imagenet": ModelConfig.densenet_201,
+            }
+            if model_name not in config_map:
+                raise ValueError(f"Model name '{model_name}' is unknown, please provide config argument")
+            config = config_map[model_name]()
+
+        model_ckpt_path = snapshot_download(repo_id=model_name, allow_patterns="*.h5")
+        return params.create_model_from_h5(model_ckpt_path, config)
+
 
 @jax.jit
 def forward(graphdef: nnx.GraphDef, state: nnx.State, x: jax.Array) -> jax.Array:

bonsai/models/densenet121/params.py

@@ -85,7 +85,7 @@ def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig):
 
     # final bn
     for i in range(len(bn_params)):
-        key = rf"^layers/dense_net_backbone/layers/batch_normalization_120/vars/{i}$"
+        key = rf"^layers/dense_net_backbone/layers/batch_normalization_{end_index + 1}/vars/{i}$"
         mapping[key] = (f"final_bn.{bn_params[i]}", None)
 
     # linear

bonsai/models/densenet121/tests/run_model.py

@@ -17,18 +17,16 @@
 import jax
 import jax.numpy as jnp
 from flax import nnx
-from huggingface_hub import snapshot_download
 
-from bonsai.models.densenet121 import modeling, params
+from bonsai.models.densenet121 import modeling
 
 
 def run_model():
-    # 1. Download h5 file
-    model_ckpt_path = snapshot_download("keras/densenet_121_imagenet")
+    # 1. Define pretrained model name:
+    model_name = "keras/densenet_121_imagenet"
 
     # 2. Load pretrained model
-    config = modeling.ModelConfig.densenet_121()
-    model = params.create_model_from_h5(model_ckpt_path, config)
+    model = modeling.DenseNet.from_pretrained(model_name)
     graphdef, state = nnx.split(model)
     state = jax.tree.leaves(state)
 

bonsai/models/densenet121/tests/test_outputs_densenet121.py

@@ -19,9 +19,7 @@
 import tensorflow as tf
 from absl.testing import absltest
 from flax import nnx
-from huggingface_hub import snapshot_download
-
-from bonsai.models.densenet121 import modeling, params
+from bonsai.models.densenet121 import modeling
 
 
 class TestModuleForwardPasses(absltest.TestCase):
@@ -30,10 +28,7 @@ def setUp(self):
         jax.config.update("jax_default_matmul_precision", "float32")
         try:
             self.ref_model = keras_hub.models.ImageClassifier.from_preset("densenet_121_imagenet")
-            model_ckpt_path = snapshot_download("keras/densenet_121_imagenet")
-            graph_def, state = nnx.split(
-                params.create_model_from_h5(model_ckpt_path, modeling.ModelConfig.densenet_121())
-            )
+            graph_def, state = nnx.split(modeling.DenseNet.from_pretrained("keras/densenet_121_imagenet"))
             state = jax.tree.map(lambda x: x.astype(jnp.float32) if isinstance(x, jax.Array) else x, state)
             self.nnx_model = nnx.merge(graph_def, state)
 

bonsai/models/dinov3/modeling.py

@@ -325,6 +325,31 @@ def __call__(self, pixel_values: Array):
 
         return {"last_hidden_state": sequence_output, "pooler_output": pooled_output}
 
+    @classmethod
+    def from_pretrained(cls, model_name: str, config: ModelConfig | None = None):
+        """model_name the *model id* of a pretrained model hosted inside
+        a model repo on huggingface.co. For example, "facebook/dinov3-vits16-pretrain-lvd1689m".
+        Note that access to the model is restricted and you need to be authorized to access it.
+        """
+        from huggingface_hub import snapshot_download
+        from bonsai.models.dinov3 import params
+
+        if config is None:
+            config_map = {
+                "facebook/dinov3-vits16-pretrain-lvd1689m": ModelConfig.dinov3_vits16,
+                "facebook/dinov3-vits16plus-pretrain-lvd1689m": ModelConfig.dinov3_vits16plus,
+                "facebook/dinov3-vitb16-pretrain-lvd1689m": ModelConfig.dinov3_vitb16,
+                "facebook/dinov3-vitl16-pretrain-lvd1689m": ModelConfig.dinov3_vitl16,
+                "facebook/dinov3-vith16plus-pretrain-lvd1689m": ModelConfig.dinov3_vith16plus,
+                "facebook/dinov3-vit7b16-pretrain-lvd1689m": ModelConfig.dinov3_vit7b16,
+            }
+            if model_name not in config_map:
+                raise ValueError(f"Model name '{model_name}' is unknown, please provide config argument")
+            config = config_map[model_name]()
+
+        model_ckpt_path = snapshot_download(repo_id=model_name, allow_patterns="*.safetensors")
+        return params.create_model_from_safe_tensors(model_ckpt_path, config)
+
 
 @jax.jit()
 def forward(model: Dinov3ViTModel, inputs: Array):

bonsai/models/dinov3/tests/run_model.py

@@ -0,0 +1,45 @@
+import time
+
+import jax
+import jax.numpy as jnp
+
+from bonsai.models.dinov3 import modeling
+
+
+def run_model():
+    # 1. Create model
+    model = modeling.Dinov3ViTModel.from_pretrained("facebook/dinov3-vits16-pretrain-lvd1689m")
+    config = model.config
+
+    # 2. Prepare dummy input
+    batch_size = 4
+    image_size = 224
+    dummy_input = jnp.ones((batch_size, 3, image_size, image_size), dtype=jnp.float32)
+
+    # 3. Warmup (triggers JIT compilation)
+    modeling.forward(model, dummy_input)["pooler_output"].block_until_ready()
+
+    # Profile a few steps
+    jax.profiler.start_trace("/tmp/profile-dinov3")
+    for _ in range(5):
+        logits = modeling.forward(model, dummy_input)["pooler_output"]
+        jax.block_until_ready(logits)
+    jax.profiler.stop_trace()
+
+    # 4. Timed execution for inference
+    num_runs = 10
+    t0 = time.perf_counter()
+    for _ in range(num_runs):
+        logits = modeling.forward(model, dummy_input)["pooler_output"]
+        jax.block_until_ready(logits)
+    t1 = time.perf_counter()
+    print(f"{num_runs} inference runs took {t1 - t0:.4f} s")
+    print(f"Average inference time: {(t1 - t0) / num_runs * 1000:.2f} ms")
+
+    # 5. Show output shape
+    print(f"\nInput shape: {dummy_input.shape}")
+    print(f"Output logits shape: {logits.shape}")
+
+
+if __name__ == "__main__":
+    run_model()

bonsai/models/dinov3/tests/test_outputs_dinov3.py

@@ -44,7 +44,7 @@ def test_input_embeddings(self):
         nnx_emb = self.bonsai_model.embeddings
 
         jx = jax.random.normal(jax.random.key(0), self.image_shape, dtype=jnp.float32)
-        tx = torch.tensor(jx, dtype=torch.float32)
+        tx = torch.tensor(np.asarray(jx), dtype=torch.float32)
 
         with torch.inference_mode():
             ty = torch_emb(tx)
@@ -61,7 +61,7 @@ def test_first_layer(self):
         nnx_layer = self.bonsai_model.layer[0]
 
         jx = jax.random.normal(jax.random.key(0), self.image_shape, dtype=jnp.float32)
-        tx = torch.tensor(jx, dtype=torch.float32)
+        tx = torch.tensor(np.asarray(jx), dtype=torch.float32)
 
         jhs = nnx_emb(jx)
         jpe = nnx_pe(jx)
@@ -77,7 +77,7 @@ def test_first_layer(self):
 
     def test_last_hidden_state(self):
         jx = jax.random.normal(jax.random.key(0), self.image_shape, dtype=jnp.float32)
-        tx = torch.tensor(jx, dtype=torch.float32)
+        tx = torch.tensor(np.asarray(jx), dtype=torch.float32)
 
         with torch.inference_mode():
             ty = self.baseline_model(tx).last_hidden_state
@@ -87,7 +87,7 @@ def test_last_hidden_state(self):
 
     def test_pooled_output_embeddings(self):
         jx = jax.random.normal(jax.random.key(0), self.image_shape, dtype=jnp.float32)
-        tx = torch.tensor(jx, dtype=torch.float32)
+        tx = torch.tensor(np.asarray(jx), dtype=torch.float32)
 
         with torch.inference_mode():
             ty = self.baseline_model(tx).pooler_output