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:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user