feat: Bump version to 0.11.0 and update dependencies
- Updated ONNX Runtime to 1.24.4 and replaced OpenCV with PyAV 17.0.0 in the installation environment. - Enhanced model downloading logic to prioritize Hugging Face and fallback to OmooLab R2, with improved error handling and status reporting. - Introduced a new environment revision file for shared metadata. - Refactored video frame extraction to utilize the PyAV library instead of OpenCV. - Improved installation process to handle retries across specified PyPI indexes. - Added tests for download progress reporting and model downloading behavior.
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
@@ -25,6 +26,49 @@ def load_worker():
|
||||
return module
|
||||
|
||||
|
||||
def _make_fake_av(frame_count):
|
||||
"""Build a fake av module yielding frame_count integer-valued frames."""
|
||||
class FakeImage:
|
||||
def __init__(self, value):
|
||||
self.value = value
|
||||
|
||||
def save(self, path, format=None):
|
||||
Path(path).write_text(str(self.value), encoding="utf-8")
|
||||
|
||||
class FakeFrame:
|
||||
def __init__(self, value):
|
||||
self.value = value
|
||||
|
||||
def to_image(self):
|
||||
return FakeImage(self.value)
|
||||
|
||||
class FakeContainer:
|
||||
def __init__(self):
|
||||
self.position = 0
|
||||
self.closed = False
|
||||
self.streams = SimpleNamespace(video=[object()])
|
||||
|
||||
def decode(self, _stream):
|
||||
while self.position < frame_count:
|
||||
frame = FakeFrame(self.position)
|
||||
self.position += 1
|
||||
yield frame
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
class FakeAv:
|
||||
def __init__(self):
|
||||
self.containers = []
|
||||
|
||||
def open(self, _path):
|
||||
container = FakeContainer()
|
||||
self.containers.append(container)
|
||||
return container
|
||||
|
||||
return FakeAv()
|
||||
|
||||
|
||||
class WorkerTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
@@ -177,92 +221,32 @@ class WorkerTest(unittest.TestCase):
|
||||
self.assertIsNone(temporary_directory)
|
||||
|
||||
def test_extract_video_frames_reads_frames_sequentially(self):
|
||||
class FakeCapture:
|
||||
def __init__(self):
|
||||
self.position = 0
|
||||
|
||||
def isOpened(self):
|
||||
return True
|
||||
|
||||
def read(self):
|
||||
if self.position >= 100:
|
||||
return False, None
|
||||
frame = self.position
|
||||
self.position += 1
|
||||
return True, frame
|
||||
|
||||
def release(self):
|
||||
return None
|
||||
|
||||
class FakeCv2:
|
||||
def __init__(self):
|
||||
self.captures = []
|
||||
|
||||
def VideoCapture(self, _path):
|
||||
capture = FakeCapture()
|
||||
self.captures.append(capture)
|
||||
return capture
|
||||
|
||||
def imwrite(self, path, frame):
|
||||
Path(path).write_text(str(frame), encoding="utf-8")
|
||||
return True
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
fake_cv2 = FakeCv2()
|
||||
fake_av = _make_fake_av(100)
|
||||
image_paths = self.worker.extract_video_frames(
|
||||
root / "clip.mkv",
|
||||
root / "out",
|
||||
10,
|
||||
fake_cv2,
|
||||
fake_av,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
[path.name for path in image_paths],
|
||||
[f"frame.{index:04d}.jpg" for index in range(1, 11)],
|
||||
)
|
||||
self.assertEqual(len(fake_cv2.captures), 1)
|
||||
self.assertEqual(fake_cv2.captures[0].position, 10)
|
||||
self.assertEqual(len(fake_av.containers), 1)
|
||||
self.assertEqual(fake_av.containers[0].position, 10)
|
||||
|
||||
def test_extract_video_frames_starts_after_offset(self):
|
||||
class FakeCapture:
|
||||
def __init__(self):
|
||||
self.position = 0
|
||||
|
||||
def isOpened(self):
|
||||
return True
|
||||
|
||||
def read(self):
|
||||
if self.position >= 100:
|
||||
return False, None
|
||||
frame = self.position
|
||||
self.position += 1
|
||||
return True, frame
|
||||
|
||||
def release(self):
|
||||
return None
|
||||
|
||||
class FakeCv2:
|
||||
def __init__(self):
|
||||
self.captures = []
|
||||
|
||||
def VideoCapture(self, _path):
|
||||
capture = FakeCapture()
|
||||
self.captures.append(capture)
|
||||
return capture
|
||||
|
||||
def imwrite(self, path, frame):
|
||||
Path(path).write_text(str(frame), encoding="utf-8")
|
||||
return True
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
fake_cv2 = FakeCv2()
|
||||
fake_av = _make_fake_av(100)
|
||||
image_paths = self.worker.extract_video_frames(
|
||||
root / "clip.mkv",
|
||||
root / "out",
|
||||
10,
|
||||
fake_cv2,
|
||||
fake_av,
|
||||
start_frame=71,
|
||||
)
|
||||
extracted = [
|
||||
@@ -271,41 +255,16 @@ class WorkerTest(unittest.TestCase):
|
||||
]
|
||||
|
||||
self.assertEqual(extracted, [str(frame) for frame in range(70, 80)])
|
||||
self.assertEqual(fake_cv2.captures[0].position, 80)
|
||||
self.assertEqual(fake_av.containers[0].position, 80)
|
||||
|
||||
def test_extract_video_frames_tolerates_short_video_tail(self):
|
||||
class ShortCapture:
|
||||
def __init__(self):
|
||||
self.position = 0
|
||||
|
||||
def isOpened(self):
|
||||
return True
|
||||
|
||||
def read(self):
|
||||
if self.position >= 99:
|
||||
return False, None
|
||||
frame = self.position
|
||||
self.position += 1
|
||||
return True, frame
|
||||
|
||||
def release(self):
|
||||
return None
|
||||
|
||||
class FakeCv2:
|
||||
def VideoCapture(self, _path):
|
||||
return ShortCapture()
|
||||
|
||||
def imwrite(self, path, frame):
|
||||
Path(path).write_text(str(frame), encoding="utf-8")
|
||||
return True
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
image_paths = self.worker.extract_video_frames(
|
||||
root / "clip.mkv",
|
||||
root / "out",
|
||||
10,
|
||||
FakeCv2(),
|
||||
_make_fake_av(99),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
@@ -443,38 +402,120 @@ class WorkerTest(unittest.TestCase):
|
||||
self.worker.download_model(args)
|
||||
|
||||
def test_download_model_reports_all_failed_sources(self):
|
||||
fake_hugging_face = ModuleType("huggingface_hub")
|
||||
|
||||
def fake_snapshot_download(**options):
|
||||
raise OSError("direct down")
|
||||
|
||||
fake_hugging_face.snapshot_download = fake_snapshot_download
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
allow_pattern=[],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
||||
sys.modules,
|
||||
{"huggingface_hub": fake_hugging_face},
|
||||
), patch.object(
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
self.worker,
|
||||
"create_download_progress_class",
|
||||
return_value=object,
|
||||
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||
"download_from_hf",
|
||||
side_effect=OSError("direct down"),
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"download_from_r2",
|
||||
side_effect=OSError("mirror down"),
|
||||
):
|
||||
args.r2_directory = "DA3Small"
|
||||
args.r2_file = ["config.json|100|abc"]
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
|
||||
self.worker.download_model(args)
|
||||
|
||||
def test_download_model_prefers_direct_hf_then_r2(self):
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
allow_pattern=["config.json"],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
)
|
||||
order = []
|
||||
|
||||
def fake_download_from_hf(*_args, **_kwargs):
|
||||
order.append("hf")
|
||||
raise OSError("hf down")
|
||||
|
||||
def fake_download_from_r2(*_args, **_kwargs):
|
||||
order.append("r2")
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
self.worker,
|
||||
"download_from_hf",
|
||||
side_effect=fake_download_from_hf,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"download_from_r2",
|
||||
side_effect=fake_download_from_r2,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"_verify_model_complete",
|
||||
return_value=None,
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
self.worker.download_model(args)
|
||||
|
||||
self.assertEqual(order, ["hf", "r2"])
|
||||
|
||||
def test_download_declared_skips_hf_when_r2_preferred(self):
|
||||
root = PROJECT_ROOT / ".tmp-models-from-r2"
|
||||
model_dir = root / "models" / "DA3"
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir=str(model_dir),
|
||||
status="status.json",
|
||||
allow_pattern=[],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
)
|
||||
order = []
|
||||
try:
|
||||
self.worker._mark_models_from_r2(model_dir)
|
||||
with patch.object(
|
||||
self.worker,
|
||||
"download_from_hf",
|
||||
side_effect=lambda *a, **k: order.append("hf"),
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"download_from_r2",
|
||||
side_effect=lambda *a, **k: order.append("r2"),
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"_verify_model_complete",
|
||||
return_value=None,
|
||||
), patch.object(self.worker, "write_status"):
|
||||
self.worker._download_declared_or_fail(
|
||||
args, model_dir, "DA3 Small", 0.0, 1.0, {}
|
||||
)
|
||||
finally:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
|
||||
self.assertEqual(order, ["r2"])
|
||||
|
||||
def test_models_from_r2_roundtrip_preserves_fields(self):
|
||||
root = PROJECT_ROOT / ".tmp-models-from-r2"
|
||||
model_dir = root / "models" / "DA3"
|
||||
discovery = root / "service.json"
|
||||
try:
|
||||
discovery.parent.mkdir(parents=True, exist_ok=True)
|
||||
discovery.write_text(
|
||||
json.dumps({"port": 1, "token": "t"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
self.assertFalse(self.worker._models_from_r2(model_dir))
|
||||
self.worker._mark_models_from_r2(model_dir)
|
||||
self.assertTrue(self.worker._models_from_r2(model_dir))
|
||||
payload = json.loads(discovery.read_text(encoding="utf-8"))
|
||||
self.assertEqual(payload["port"], 1)
|
||||
self.assertTrue(payload["models_from_r2"])
|
||||
finally:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
|
||||
def test_download_from_r2_downloads_declared_files(self):
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
@@ -488,8 +529,8 @@ class WorkerTest(unittest.TestCase):
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
self.worker,
|
||||
"_download_r2_file",
|
||||
side_effect=lambda *args, **kwargs: downloaded.append(args[1]),
|
||||
"_download_file_from_url",
|
||||
side_effect=lambda *args, **kwargs: downloaded.append(args[1].name),
|
||||
), patch.object(self.worker, "write_status"):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
|
||||
Reference in New Issue
Block a user