File size: 3,835 Bytes
4cf7d37 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | from pathlib import Path
from src.model_scan import analyze_model_metadata, normalize_model_id
def root() -> Path:
return Path(__file__).resolve().parents[1]
def test_normalize_model_id_accepts_url_and_id():
assert normalize_model_id("https://huggingface.co/org/model-name?x=1") == "org/model-name"
assert normalize_model_id("org/model-name") == "org/model-name"
def test_model_prescan_safe_candidate():
result = analyze_model_metadata(
model_id="org/safe-model",
pipeline_tag="text-generation",
library_name="transformers",
tags=["safetensors"],
siblings=[{"rfilename": "config.json"}, {"rfilename": "model.safetensors"}, {"rfilename": "README.md"}],
readme="This model card has useful documentation. " * 30,
)
assert result["verdict"] in {"safe", "caution"}
assert result["metadata"]["has_safetensors"] is True
assert any("Safetensors" in item for item in result["good_signals"])
def test_model_prescan_risky_custom_code_and_gated():
result = analyze_model_metadata(
model_id="org/risky-model",
pipeline_tag=None,
library_name=None,
gated="manual",
siblings=[{"rfilename": "model.bin"}, {"rfilename": "modeling_custom.py"}],
config={"auto_map": {"AutoModel": "modeling_custom.Custom"}},
readme="tiny",
)
assert result["verdict"] in {"risky", "unsupported"}
assert result["metadata"]["has_custom_code_signal"] is True
assert any("gated" in item.lower() for item in result["risk_signals"])
def test_prescan_route_and_ui_are_present():
app = (root() / "app.py").read_text(encoding="utf-8")
html = (root() / "web" / "index.html").read_text(encoding="utf-8")
js = (root() / "web" / "static" / "app.js").read_text(encoding="utf-8")
css = (root() / "web" / "static" / "app.css").read_text(encoding="utf-8")
assert '"/api/models/pre-scan"' in app
assert "modelPreScanCard" in html
assert "scanModel" in js
assert "modelScanMatchesCurrent" in js
assert "Run the model pre-scan before launching" in js
assert ".model-prescan-card" in css
def test_model_prescan_known_good_z_image_turbo_is_safe():
readme = """
# Z-Image-Turbo
This is a text-to-image Diffusers model.
```python
import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"Tongyi-MAI/Z-Image-Turbo",
torch_dtype=torch.bfloat16,
device_map="cuda",
)
image = pipe(prompt="A small robot building a Hugging Face Space", num_inference_steps=8).images[0]
```
""" * 5
result = analyze_model_metadata(
model_id="Tongyi-MAI/Z-Image-Turbo",
pipeline_tag="text-to-image",
library_name="diffusers",
tags=["diffusers", "safetensors", "text-to-image"],
siblings=[
{"rfilename": "README.md"},
{"rfilename": "model_index.json"},
{"rfilename": "transformer/diffusion_pytorch_model.safetensors"},
{"rfilename": "vae/diffusion_pytorch_model.safetensors"},
{"rfilename": "scheduler/scheduler_config.json"},
],
model_index={"_class_name": "ZImagePipeline"},
readme=readme,
)
assert result["verdict"] == "safe"
assert result["score"] >= 82
assert result["metadata"]["has_diffusers_example"] is True
assert result["metadata"]["diffusers_standard"] is True
assert result["metadata"]["pipeline_class"] == "ZImagePipeline"
assert any("Diffusers example" in item for item in result["good_signals"])
def test_default_model_is_known_good_z_image_turbo():
html = (root() / "web" / "index.html").read_text(encoding="utf-8")
assert 'value="Tongyi-MAI/Z-Image-Turbo"' in html
assert 'placeholder="z-image-turbo-demo"' in html
|