diff --git a/zeromodels/base/base_mixin.py b/zeromodels/base/base_mixin.py index 247dfa14..093fdfc3 100644 --- a/zeromodels/base/base_mixin.py +++ b/zeromodels/base/base_mixin.py @@ -12,7 +12,10 @@ from huggingface_hub import hf_hub_download from zeromodels.base import base_attention -from zeromodels.conversion import download_weights +from zeromodels.conversion.file_downloader import ( + download_weights, + is_huggingface_url, +) from zeromodels.conversion.hf_download_utils import ( download_hf_state_dict, ) @@ -65,7 +68,7 @@ def _url_exists(url): """True if a range GET on ``url`` succeeds (uses HF_TOKEN for hf.co if set).""" headers = {"User-Agent": "zeromodels", "Range": "bytes=0-0"} token = os.environ.get("HF_TOKEN") - if token and "huggingface.co" in url: + if token and is_huggingface_url(url): headers["Authorization"] = f"Bearer {token}" try: with urllib.request.urlopen( diff --git a/zeromodels/conversion/file_downloader.py b/zeromodels/conversion/file_downloader.py index d165ac12..ecf59067 100644 --- a/zeromodels/conversion/file_downloader.py +++ b/zeromodels/conversion/file_downloader.py @@ -5,6 +5,24 @@ from typing import Optional from urllib.parse import urlparse +_HUGGING_FACE_HOSTS = frozenset({"huggingface.co", "www.huggingface.co"}) + + +def is_huggingface_url(url: str) -> bool: + """Return whether ``url`` is a trusted HTTPS Hugging Face URL. + + Authentication tokens must only be attached after parsing the URL. A substring + check would also match attacker-controlled hosts, paths, and query strings that + merely contain ``huggingface.co``. + """ + try: + parsed = urlparse(url) + return ( + parsed.scheme.lower() == "https" and parsed.hostname in _HUGGING_FACE_HOSTS + ) + except (TypeError, ValueError): + return False + def validate_url(url: str) -> bool: """Validate if the provided URL is well-formed. @@ -30,7 +48,7 @@ def _parse_hf_resolve(url: str): fall back to a plain streamed download. """ parsed = urlparse(url) - if parsed.netloc not in ("huggingface.co", "www.huggingface.co"): + if not is_huggingface_url(url): return None parts = parsed.path.strip("/").split("/") if len(parts) < 5 or parts[2] != "resolve": @@ -94,7 +112,7 @@ def download_file( headers = {"User-Agent": "zeromodels"} token = os.environ.get("HF_TOKEN") - if token and "huggingface.co" in file_url: + if token and is_huggingface_url(file_url): headers["Authorization"] = f"Bearer {token}" # Stream to a sibling ``.incomplete`` file, then atomically move it into diff --git a/zeromodels/models/glm5_moe/glm5_moe_layers.py b/zeromodels/models/glm5_moe/glm5_moe_layers.py index abecb668..498f9182 100644 --- a/zeromodels/models/glm5_moe/glm5_moe_layers.py +++ b/zeromodels/models/glm5_moe/glm5_moe_layers.py @@ -285,7 +285,7 @@ def call(self, hidden_states, q_resid, cos, sin, attention_mask): k_pe = apply_rope(k_pe[:, :, None, :], cos, sin, unsqueeze_axis=2)[:, :, 0, :] k = ops.concatenate([k_pe, k_nope], axis=-1) - weights = self.weights_proj(ops.cast(hidden_states, "float32")) * ( + weights = ops.cast(self.weights_proj(hidden_states), "float32") * ( self.n_heads**-0.5 ) # (B, S, H) scores = ( diff --git a/zeromodels/models/grounding_dino/__init__.py b/zeromodels/models/grounding_dino/__init__.py index bd132b4d..de0e5c18 100644 --- a/zeromodels/models/grounding_dino/__init__.py +++ b/zeromodels/models/grounding_dino/__init__.py @@ -6,7 +6,6 @@ ) from zeromodels.models.grounding_dino.grounding_dino_model import ( GroundingDinoDetect, - GroundingDinoForObjectDetection, GroundingDinoModel, ) from zeromodels.models.grounding_dino.grounding_dino_processor import ( @@ -23,7 +22,6 @@ "GroundingDinoConfig", "GroundingDinoModel", "GroundingDinoDetect", - "GroundingDinoForObjectDetection", "GroundingDinoTextModel", "GroundingDinoTokenizer", "GroundingDinoImageProcessor", diff --git a/zeromodels/models/grounding_dino/grounding_dino_model.py b/zeromodels/models/grounding_dino/grounding_dino_model.py index c8e246b3..9e9ea446 100644 --- a/zeromodels/models/grounding_dino/grounding_dino_model.py +++ b/zeromodels/models/grounding_dino/grounding_dino_model.py @@ -740,7 +740,3 @@ class GroundingDinoDetect(_GroundingDinoFunctional): """ core_class = GroundingDinoDetectCore - - -# Backward-compatible alias for the previous class name. -GroundingDinoForObjectDetection = GroundingDinoDetect