Spaces:
Running
Running
| """不下载模型权重，验证 v1.0 推理适配层与官方 pipeline 的调用契约。""" | |
| import base64 | |
| import importlib.util | |
| import io | |
| import os | |
| import sys | |
| import unittest | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| from unittest.mock import MagicMock, Mock, patch | |
| from PIL import Image | |
| class EndpointHandlerTest(unittest.TestCase): | |
| """覆盖标签映射、真实分数、预处理委托和本地加载边界。""" | |
| def setUpClass(cls): | |
| """隔离导入推理依赖，让适配层测试可在没有 GPU 和权重时运行。""" | |
| spec = importlib.util.spec_from_file_location( | |
| "pixai_handler_test", Path(__file__).resolve().parents[1] / "handler.py" | |
| ) | |
| cls.module = importlib.util.module_from_spec(spec) | |
| cls.torch = SimpleNamespace( | |
| cuda=SimpleNamespace(is_available=Mock(return_value=False)), | |
| zeros=Mock(), | |
| ) | |
| with patch.dict( | |
| sys.modules, | |
| {"torch": cls.torch, "transformers": SimpleNamespace(pipeline=Mock())}, | |
| ): | |
| spec.loader.exec_module(cls.module) | |
| def setUp(self): | |
| """构造各类别独立的分数，避免将版权或风格误当成通用标签。""" | |
| self.results = { | |
| "general": {"blue_hair": 0.7, "1girl": 0.99, "hat": 0.2}, | |
| "character": {"character_b": 0.3, "character_a": 0.9}, | |
| "copyright": {"original": 0.64}, | |
| "style": {"style_a": 0.6}, | |
| "meta": {"highres": 0.8}, | |
| "rating": {"rating:g": 0.88}, | |
| } | |
| self.pipe = Mock(return_value={"results": self.results}) | |
| self.factory = Mock(return_value=self.pipe) | |
| self.torch.cuda.is_available.return_value = False | |
| self.torch.zeros.reset_mock() | |
| self.torch.zeros.return_value.to.side_effect = None | |
| self.model_path = Path(__file__).resolve().parents[1] | |
| with patch.object(Path, "is_file", return_value=True), patch.object( | |
| self.module, "pipeline", self.factory | |
| ): | |
| self.handler = self.module.EndpointHandler(str(self.model_path)) | |
| def test_local_snapshot_and_official_processor(self): | |
| """本地四文件快照与官方自定义任务共同决定模型和预处理版本。""" | |
| self.factory.assert_called_once_with( | |
| task="tagger", | |
| model=str(self.model_path), | |
| image_processor=str(self.model_path), | |
| trust_remote_code=True, | |
| token=False, | |
| model_kwargs={"local_files_only": True}, | |
| device="cpu", | |
| ) | |
| image = Image.new("RGBA", (80, 40), (255, 0, 0, 0)) | |
| data = {"inputs": image, "parameters": {"include_scores": True}} | |
| result = self.handler(data) | |
| self.pipe.assert_called_once_with( | |
| image, threshold={"general": 0.17, "character": 0.27} | |
| ) | |
| self.assertEqual(image.mode, "RGBA") | |
| self.assertEqual(image.size, (80, 40)) | |
| self.assertIs(data["inputs"], image) | |
| self.assertEqual(data["parameters"], {"include_scores": True}) | |
| self.assertEqual(result["feature"], ["1girl", "blue_hair", "hat"]) | |
| self.assertEqual(result["feature_scores"], self.results["general"]) | |
| for category in ("character", "style", "meta", "rating"): | |
| self.assertEqual(result[f"{category}_scores"], self.results[category]) | |
| self.assertEqual(result["ip"], ["original"]) | |
| self.assertEqual(result["ip_scores"], {"original": 0.64}) | |
| self.assertEqual(result["_params"]["general_threshold"], 0.17) | |
| self.assertEqual(set(result["_timings"]), {"fetch_s", "inference_s", "post_process_s", "total_s"}) | |
| def test_custom_thresholds_and_optional_scores(self): | |
| """阈值交由官方流程筛选，关闭分数时不泄漏分数字段。""" | |
| image = Image.new("RGB", (32, 32)) | |
| result = self.handler( | |
| {"inputs": image, "parameters": {"general_threshold": 0.4, "character_threshold": 0.6}} | |
| ) | |
| self.pipe.assert_called_once_with(image, threshold={"general": 0.4, "character": 0.6}) | |
| self.assertFalse(any(name.endswith("_scores") for name in result)) | |
| def test_topk_ranks_categories_independently(self): | |
| """Top-K 不先经过默认阈值，且保留版权等其他分类的原始分数。""" | |
| image = Image.new("RGB", (32, 32)) | |
| result = self.handler( | |
| { | |
| "inputs": image, | |
| "parameters": {"mode": "topk", "topk_general": 2, "topk_character": 1, "include_scores": True}, | |
| } | |
| ) | |
| self.pipe.assert_called_once_with(image, threshold={"general": -1.0, "character": -1.0}) | |
| self.assertEqual(result["feature"], ["1girl", "blue_hair"]) | |
| self.assertEqual(result["character"], ["character_a"]) | |
| self.assertEqual(result["character_scores"], {"character_a": 0.9}) | |
| self.assertEqual(result["ip_scores"], {"original": 0.64}) | |
| zero_result = self.handler( | |
| {"inputs": image, "parameters": {"mode": "topk", "topk_general": 0, "topk_character": 0}} | |
| ) | |
| self.assertEqual(zero_result["feature"], []) | |
| self.assertEqual(zero_result["character"], []) | |
| def test_base64_and_remote_images_are_decoded_before_inference(self): | |
| """URL 图片在响应关闭前解码，base64 和透明图片也交由官方处理器。""" | |
| image_bytes = io.BytesIO() | |
| Image.new("RGBA", (80, 40), (0, 255, 0, 0)).save(image_bytes, format="PNG") | |
| image_content = image_bytes.getvalue() | |
| self.handler({"inputs": {"image": base64.b64encode(image_content).decode()}}) | |
| decoded_image = self.pipe.call_args.args[0] | |
| self.assertEqual((decoded_image.mode, decoded_image.size), ("RGBA", (80, 40))) | |
| response = MagicMock() | |
| response.raw = io.BytesIO(image_content) | |
| response.__enter__.return_value = response | |
| response.__exit__.side_effect = lambda *_: response.raw.close() | |
| with patch.object(self.module.requests, "get", return_value=response) as get: | |
| self.handler({"inputs": {"url": "https://example.org/image.png"}}) | |
| get.assert_called_once_with("https://example.org/image.png", stream=True, timeout=5.0) | |
| response.raise_for_status.assert_called_once_with() | |
| self.assertTrue(response.raw.closed) | |
| self.assertEqual(self.pipe.call_args.args[0].getpixel((0, 0)), (0, 255, 0, 0)) | |
| def test_incomplete_snapshot_fails_before_model_loading(self): | |
| """缺失新版本预处理器时直接报错，避免加载旧版或远程补齐混用。""" | |
| with patch.object( | |
| Path, "is_file", autospec=True, side_effect=lambda path: path.name != "preprocessor_config.json" | |
| ), patch.object(self.module, "pipeline") as factory: | |
| with self.assertRaisesRegex(FileNotFoundError, "preprocessor_config.json"): | |
| self.module.EndpointHandler(str(self.model_path)) | |
| factory.assert_not_called() | |
| def test_cuda_probe_failure_uses_cpu(self): | |
| """存在 CUDA 接口但驱动不能工作时，仍能使用 CPU 加载。""" | |
| self.torch.cuda.is_available.return_value = True | |
| self.torch.zeros.return_value.to.side_effect = RuntimeError("driver unavailable") | |
| with patch.dict(os.environ, {"FORCE_CPU": "0"}), patch.object( | |
| Path, "is_file", return_value=True | |
| ), patch.object(self.module, "pipeline", self.factory): | |
| handler = self.module.EndpointHandler(str(self.model_path)) | |
| self.assertEqual(handler.device, "cpu") | |
| self.assertEqual(self.factory.call_args.kwargs["device"], "cpu") | |
| if __name__ == "__main__": | |
| unittest.main() | |