This commit is contained in:
tastelikefeet
2026-06-18 17:16:11 +08:00
committed by GitHub
parent c1d88e6ab7
commit cecb699cd0

View File

@@ -131,7 +131,7 @@ def check_model_from_owner_group(model_dir: str,
"""This checking is for the torch.load, this function may eval malicious code into memory
Args:
model_dir: The local model_dir
model_dir: The local model_dir or model_id
owner_group: The owner group to trust
Returns:
@@ -142,6 +142,15 @@ def check_model_from_owner_group(model_dir: str,
if owner_group is None:
owner_group = ['iic', 'damo']
model_dir = model_dir.rstrip('/').rstrip('\\')
model_dir = os.path.dirname(model_dir)
group = os.path.basename(model_dir)
return group in owner_group
parent_dir = os.path.dirname(model_dir)
group = os.path.basename(parent_dir)
if group in owner_group:
return True
# Also check cache path pattern: {cache_root}/{owner}--{model_name}/snapshots/{revision}
# Only check the grandparent directory to avoid path traversal attacks
grandparent = os.path.basename(os.path.dirname(parent_dir))
if '--' in grandparent:
prefix = grandparent.split('--', 1)[0]
if prefix in owner_group:
return True
return False