Replace numpy with keras.ops across model code

zeromodels/models/deepseek_v4/deepseek_v4_layers.py

@@ -1,5 +1,4 @@
 import keras
-import numpy as np
 from keras import layers, ops
 
 MASK_NEG = -1e9
@@ -450,7 +449,9 @@ def __init__(
         self.compress_inv_freq = (
             None
             if compress_inv_freq is None
-            else np.asarray(compress_inv_freq, dtype="float32")
+            else ops.convert_to_numpy(
+                ops.convert_to_tensor(compress_inv_freq, dtype="float32")
+            )
         )
         self.norm_eps = norm_eps
         self.scaling = head_dim**-0.5

zeromodels/models/efficientdet/efficientdet_image_processor.py

@@ -1,9 +1,7 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, Optional, Tuple
 
 import keras
-import numpy as np
 from keras import ops
-from PIL import Image
 
 from zeromodels.base import BaseImageProcessor
 from zeromodels.utils.image_util import get_data_format, load_image
@@ -95,9 +93,7 @@ def preprocess_one(self, image):
             t = ops.transpose(t, (0, 3, 1, 2))
         return t, scale, (orig_h, orig_w)
 
-    def call(
-        self, image: Union[str, np.ndarray, Image.Image, List]
-    ) -> Dict[str, Union[keras.KerasTensor, np.ndarray]]:
+    def call(self, image) -> Dict:
         items = list(image) if isinstance(image, (list, tuple)) else [image]
         tensors, scales, sizes = [], [], []
         for item in items:

zeromodels/models/gemma3n/gemma3n_layers.py

@@ -1,7 +1,6 @@
 import math
 
 import keras
-import numpy as np
 from keras import layers, ops
 
 
@@ -720,9 +719,13 @@ def __init__(self, hidden_size, num_heads, context_left, context_right, **kwargs
         )
         num_timescales = hidden_size // 2
         log_increment = math.log(1.0e4) / max(num_timescales - 1, 1)
-        inv = np.exp(np.arange(num_timescales) * -log_increment).astype("float32")
+        inv = ops.convert_to_numpy(
+            ops.exp(ops.arange(num_timescales, dtype="float32") * -log_increment)
+        )
         self.inv_timescales = inv[None, None, :]  # [1, 1, num_timescales]
-        pos = np.arange(self.max_backward, -self.max_forward - 1, -1, dtype="float32")
+        pos = ops.convert_to_numpy(
+            ops.arange(self.max_backward, -self.max_forward - 1, -1, dtype="float32")
+        )
         self.pos_indices = pos[None]  # [1, F_span]
         self.max_span_plus_1 = self.pos_indices.shape[1]
 
@@ -815,14 +818,16 @@ def __init__(
 
         # Static local causal validity mask [W, C].
         w, c = self.chunk_size, self.context_size
-        lower = np.tril(np.ones((c, w), dtype=bool), k=0).T
-        upper = np.tril(
-            np.ones((w, c), dtype=bool),
+        lower = ops.transpose(ops.tril(ops.ones((c, w)), k=0))
+        upper = ops.tril(
+            ops.ones((w, c)),
             k=self.max_past_horizon + self.max_future_horizon,
         )
-        self.local_causal_valid_mask = np.ones((w, c), dtype=bool) & lower & upper
+        self.local_causal_valid_mask = ops.convert_to_numpy(
+            ops.logical_and(ops.cast(lower, "bool"), ops.cast(upper, "bool"))
+        )
         # Block gather start indices; sliced per call by num_blocks.
-        self._starts = np.arange(4096)
+        self._starts = ops.convert_to_numpy(ops.arange(4096))
 
     def build(self, input_shape):
         self.per_dim_scale = self.add_weight(
@@ -845,7 +850,7 @@ def extract_block_context(self, x, num_blocks, seq_len):
         x = ops.pad(x, pad_cfg)
         idx = (
             self._starts[:num_blocks, None] * self.chunk_size
-            + np.arange(self.context_size)[None, :]
+            + ops.convert_to_numpy(ops.arange(self.context_size))[None, :]
         )
         idx = ops.convert_to_tensor(idx.astype("int32"))
         return ops.take(x, idx, axis=1)  # [B, U, C, ...]
@@ -879,7 +884,7 @@ def call(self, hidden_states, mask=None):
 
         logits = self.relative_position_embedding(query_blocks, key_blocks)
         logits = ops.tanh(logits / self.logit_cap) * self.logit_cap
-        neg_inf = ops.cast(float(np.finfo(np.float32).min), "float32")
+        neg_inf = ops.cast(-3.4028234663852886e38, "float32")  # most negative float32
         logits = ops.where(final_cond, logits, neg_inf)
         probs = ops.softmax(logits, axis=-1)  # [B,N,U,W,C]
 

zeromodels/models/gemma3n/gemma3n_processor.py

@@ -1,5 +1,4 @@
 import keras
-import numpy as np
 from keras import ops
 
 from zeromodels.base import BaseProcessor
@@ -95,7 +94,11 @@ def load_image(self, item):
 
     def load_audio(self, item):
         if item.get("audio") is not None:
-            return np.asarray(item["audio"], dtype="float32").reshape(-1)
+            return ops.convert_to_numpy(
+                ops.reshape(
+                    ops.convert_to_tensor(item["audio"], dtype="float32"), (-1,)
+                )
+            )
         if item.get("path") is not None:
             import soundfile as sf
 
@@ -185,13 +188,10 @@ def call(
         bos = tok.bos_token_id
         ids = [[bos] + tok.encode(t) for t in texts]
         max_len = max(len(x) for x in ids)
-        input_ids = np.zeros((len(ids), max_len), dtype="int32")
-        attention_mask = np.zeros((len(ids), max_len), dtype="int32")
-        for i, seq_ids in enumerate(ids):
-            input_ids[i, : len(seq_ids)] = seq_ids
-            attention_mask[i, : len(seq_ids)] = 1
-        out["input_ids"] = ops.convert_to_tensor(input_ids)
-        out["attention_mask"] = ops.convert_to_tensor(attention_mask)
+        input_ids = [row + [0] * (max_len - len(row)) for row in ids]
+        attention_mask = [[1] * len(row) + [0] * (max_len - len(row)) for row in ids]
+        out["input_ids"] = ops.convert_to_tensor(input_ids, dtype="int32")
+        out["attention_mask"] = ops.convert_to_tensor(attention_mask, dtype="int32")
         return out
 
     def get_config(self):

zeromodels/models/glm5_moe/glm5_moe_layers.py

@@ -1,5 +1,4 @@
 import keras
-import numpy as np
 from keras import layers, ops
 
 from zeromodels.base.base_attention import fused_attention
@@ -198,7 +197,7 @@ def call(self, hidden_states):
             ops.one_hot(group_idx, self.n_group, dtype="float32"), axis=1
         )
         score_mask = ops.repeat(group_mask, self.num_experts // self.n_group, axis=-1)
-        choice = ops.where(score_mask > 0, biased, -np.inf)
+        choice = ops.where(score_mask > 0, biased, float("-inf"))
         _, top_idx = ops.top_k(choice, self.num_experts_per_tok)
         top_vals = ops.take_along_axis(scores, top_idx, axis=-1)
         if self.norm_topk_prob:

zeromodels/models/granite_speech5/granite_speech5_processor.py

@@ -1,5 +1,4 @@
 import keras
-import numpy as np
 from keras import ops
 
 from zeromodels.base import BaseProcessor
@@ -55,10 +54,8 @@ def call(self, audio=None, text=None, sampling_rate=16000):
             label_ids = [self.tokenizer.tokenize(t) for t in texts]
             max_len = max(len(x) for x in label_ids)
             pad_id = self.tokenizer.pad_token_id
-            labels = np.full((len(label_ids), max_len), pad_id, dtype="int32")
-            for i, seq in enumerate(label_ids):
-                labels[i, : len(seq)] = seq
-            out["labels"] = ops.convert_to_tensor(labels)
+            labels = [list(seq) + [pad_id] * (max_len - len(seq)) for seq in label_ids]
+            out["labels"] = ops.convert_to_tensor(labels, dtype="int32")
         return out
 
     def batch_decode(
@@ -69,7 +66,7 @@ def batch_decode(
                 token_ids, skip_special_tokens=skip_special_tokens
             )
         # Word-level timestamps: one dict per clip, mirroring Whisper's shape.
-        token_ids = np.asarray(ops.convert_to_numpy(token_ids)).tolist()
+        token_ids = ops.convert_to_numpy(token_ids).tolist()
         fs = self.frame_seconds
         return [
             {

zeromodels/models/locateanything/locateanything_processor.py

@@ -1,5 +1,4 @@
 import keras
-import numpy as np
 from keras import ops
 
 from zeromodels.base import BaseProcessor
@@ -198,19 +197,18 @@ def call(
             out["image_grid_hws"] = ops.convert_to_tensor(
                 image_inputs["image_grid_hws"]
             )
-            grid = [tuple(g) for g in np.asarray(image_inputs["image_grid_hws"])]
+            grid = [
+                tuple(g) for g in ops.convert_to_numpy(image_inputs["image_grid_hws"])
+            ]
             per_text = self.deal_per_text(texts, self.image_token, grid)
             texts = [self.expand_image_tokens(t, g) for t, g in zip(texts, per_text)]
 
         ids = [self.tokenizer.encode(t) for t in texts]
         max_len = max(len(x) for x in ids)
-        input_ids = np.zeros((len(ids), max_len), dtype="int32")
-        attention_mask = np.zeros((len(ids), max_len), dtype="int32")
-        for i, seq in enumerate(ids):
-            input_ids[i, : len(seq)] = seq
-            attention_mask[i, : len(seq)] = 1
-        out["input_ids"] = ops.convert_to_tensor(input_ids)
-        out["attention_mask"] = ops.convert_to_tensor(attention_mask)
+        input_ids = [list(seq) + [0] * (max_len - len(seq)) for seq in ids]
+        attention_mask = [[1] * len(seq) + [0] * (max_len - len(seq)) for seq in ids]
+        out["input_ids"] = ops.convert_to_tensor(input_ids, dtype="int32")
+        out["attention_mask"] = ops.convert_to_tensor(attention_mask, dtype="int32")
         return out
 
     def post_process_generation(self, generated, task=None, image_size=None, text=None):
@@ -229,10 +227,17 @@ def post_process_generation(self, generated, task=None, image_size=None, text=No
         [x1, y1, x2, y2]}`` or ``{"label": ..., "point": [x, y]}``.
         """
         try:
-            arr = np.asarray(ops.convert_to_numpy(generated))
+            arr = ops.convert_to_numpy(generated)
         except (TypeError, ValueError):
-            arr = np.asarray(generated)
-        sequences = [arr.tolist()] if arr.ndim == 1 else [row.tolist() for row in arr]
+            arr = None
+        if arr is not None:
+            sequences = (
+                [arr.tolist()] if arr.ndim == 1 else [row.tolist() for row in arr]
+            )
+        elif generated and isinstance(generated[0], (list, tuple)):
+            sequences = [list(row) for row in generated]
+        else:
+            sequences = [list(generated)]
 
         results = []
         for seq in sequences: