nan
c5b77c5c85
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing. - Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing. - Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Developed unit tests for server functionalities in `test_server.py`, including model caching and job management. - Enhanced job server to handle frame extraction without requiring models.
129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_download_module():
|
|
addon_directory = PROJECT_ROOT / "src" / "snapform"
|
|
sys.path.insert(0, str(addon_directory))
|
|
try:
|
|
return __import__(
|
|
"server.model_download",
|
|
fromlist=["model_download"],
|
|
)
|
|
finally:
|
|
sys.path.remove(str(addon_directory))
|
|
|
|
|
|
class FakeTqdm:
|
|
def __init__(self, *args, **kwargs):
|
|
self.n = kwargs.get("initial", 0)
|
|
self.total = kwargs.get("total", 0)
|
|
|
|
def update(self, amount=1):
|
|
self.n += amount
|
|
return True
|
|
|
|
def refresh(self, *args, **kwargs):
|
|
return True
|
|
|
|
|
|
class DownloadProgressTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.download = load_download_module()
|
|
|
|
def create_progress(self, description="Downloading bytes", **options):
|
|
reports = []
|
|
progress_class = self.download.create_download_progress_class(
|
|
lambda value, message: reports.append((value, message)),
|
|
"owner/model",
|
|
tqdm_type=options.pop("tqdm_type", FakeTqdm),
|
|
clock=options.pop("clock", lambda: 1.0),
|
|
**options,
|
|
)
|
|
progress = progress_class(
|
|
desc=description,
|
|
total=100,
|
|
initial=0,
|
|
unit="B",
|
|
)
|
|
return progress, reports
|
|
|
|
def test_reports_aggregated_download_bytes(self):
|
|
progress, reports = self.create_progress()
|
|
|
|
progress.update(25)
|
|
|
|
self.assertAlmostEqual(reports[-1][0], 0.275)
|
|
self.assertEqual(reports[-1][1], "Downloading owner/model")
|
|
|
|
def test_reports_hugging_face_aggregated_bar(self):
|
|
progress, reports = self.create_progress(
|
|
"Downloading (incomplete total...)"
|
|
)
|
|
|
|
progress.update(25)
|
|
|
|
self.assertAlmostEqual(reports[-1][0], 0.275)
|
|
|
|
def test_reports_download_speed(self):
|
|
class FakeTqdmWithRate(FakeTqdm):
|
|
@property
|
|
def format_dict(self):
|
|
return {"rate": 14_000_000}
|
|
|
|
progress, reports = self.create_progress(
|
|
tqdm_type=FakeTqdmWithRate,
|
|
label="da3 small",
|
|
)
|
|
|
|
progress.update(50)
|
|
|
|
self.assertEqual(
|
|
reports[-1][1],
|
|
"Downloading da3 small 13.4 MB/s",
|
|
)
|
|
|
|
def test_progress_callback_has_no_eta_or_state_protocol(self):
|
|
times = iter((0.0, 10.0))
|
|
progress, reports = self.create_progress(clock=lambda: next(times))
|
|
|
|
progress.update(50)
|
|
|
|
self.assertEqual(len(reports[-1]), 2)
|
|
|
|
def test_ignores_unrelated_progress_bars(self):
|
|
progress, reports = self.create_progress("Reconstructing")
|
|
|
|
progress.update(25)
|
|
|
|
self.assertEqual(reports, [])
|
|
|
|
def test_download_endpoints_uses_direct_endpoint(self):
|
|
with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
|
|
endpoints = self.download.download_endpoints()
|
|
|
|
self.assertEqual(endpoints, ["https://huggingface.co"])
|
|
|
|
def test_download_endpoints_respects_environment_override(self):
|
|
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}):
|
|
endpoints = self.download.download_endpoints()
|
|
|
|
self.assertEqual(endpoints, ["https://hf.example"])
|
|
|
|
def test_download_endpoints_strips_trailing_slash(self):
|
|
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}):
|
|
endpoints = self.download.download_endpoints()
|
|
|
|
self.assertEqual(endpoints, ["https://hf.example"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|