Enhance device management and model handling in SnapForm

- Updated frame extraction logic to allow for single frame extraction.
- Integrated model readiness checks for BEN2 and DA3 models in relevant operators.
- Added directory name attribute to model specifications for better organization.
- Improved device selection logic for production and debug tasks.
- Refactored preferences to expose device settings and ensure they reflect available hardware.
- Enhanced runtime device detection and management, including support for CUDA and MPS.
- Updated tests to cover new device management features and model directory handling.
- Removed deprecated download model tests to streamline codebase.
This commit is contained in:
2026-08-08 20:16:11 +08:00
parent aee460fe7f
commit 9c896213a5
18 changed files with 699 additions and 196 deletions

View File

@@ -385,79 +385,6 @@ class WorkerTest(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
self.worker.prepare_input(input_path, 16, directory)
def test_download_passes_optional_file_patterns(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: calls.append(
options
)
args = SimpleNamespace(
repository="PramaLLC/BEN2",
model_dir="model",
status="status.json",
allow_pattern=["config.json", "model.safetensors"],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
"create_download_progress_class",
return_value=object,
):
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.assertEqual(
calls[0]["allow_patterns"],
["config.json", "model.safetensors"],
)
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
def test_download_model_falls_back_to_modelscope_when_direct_fails(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
def fake_snapshot_download(**options):
calls.append(options)
raise OSError("Connection reset by peer")
fake_hugging_face.snapshot_download = fake_snapshot_download
args = SimpleNamespace(
repository="PramaLLC/BEN2",
model_dir="model",
status="status.json",
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.worker,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"download_from_modelscope",
) as mocked_modelscope:
def fake_download_from_modelscope(*args, **kwargs):
model_dir = Path(args[1])
model_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "config.json").write_text("{}")
(model_dir / "model.safetensors").write_bytes(b"x")
mocked_modelscope.side_effect = fake_download_from_modelscope
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.assertEqual(len(calls), 1)
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
self.assertEqual(mocked_modelscope.call_count, 1)
def test_download_model_skips_unreachable_endpoint(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")