Add model catalog with specifications for DA3 and MoGe-2 models
- Introduced ModelSpec dataclass to define model attributes. - Added multiple models including DA3 Small, Base, Giant 1.1, Nested Giant Large 1.1, and MoGe-2 variants. - Implemented functions to retrieve model details and metadata. - Included BEN2 model specifications for background removal.
This commit is contained in:
@@ -44,6 +44,35 @@ class ServerTest(unittest.TestCase):
|
||||
self.addCleanup(self.jobs_temp.cleanup)
|
||||
self.storage_root = Path(self.jobs_temp.name)
|
||||
|
||||
def test_required_model_download_skips_ready_model(self):
|
||||
manager = self.server_module.ModelManager(Path("."))
|
||||
downloads = []
|
||||
context = SimpleNamespace(progress=lambda *_args: None)
|
||||
with patch.object(
|
||||
manager,
|
||||
"ready",
|
||||
side_effect=lambda key: key == "BEN2_BASE",
|
||||
), patch.object(
|
||||
manager,
|
||||
"_download",
|
||||
side_effect=lambda _context, key, _progress: downloads.append(key),
|
||||
):
|
||||
result = manager.download_missing(
|
||||
context,
|
||||
("MOGE2_VITS_NORMAL", "BEN2_BASE"),
|
||||
)
|
||||
|
||||
self.assertEqual(downloads, ["MOGE2_VITS_NORMAL"])
|
||||
self.assertEqual(result, {"downloaded": ["MOGE2_VITS_NORMAL"]})
|
||||
|
||||
def test_models_are_owned_by_the_server_resource(self):
|
||||
self.assertFalse((PROJECT_ROOT / "src" / "snapform" / "models.py").exists())
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
manager = self.server_module.ModelManager(Path(directory))
|
||||
snapshot = manager.snapshot()
|
||||
self.assertIn("catalog", snapshot)
|
||||
self.assertTrue(snapshot["catalog"])
|
||||
|
||||
def create_cache(self):
|
||||
cache = self.server_module.ModelSessionCache.__new__(
|
||||
self.server_module.ModelSessionCache
|
||||
@@ -178,6 +207,7 @@ class ServerTest(unittest.TestCase):
|
||||
set(server.handlers),
|
||||
{
|
||||
"download-model",
|
||||
"download-required-models",
|
||||
"remove-background",
|
||||
"extract-frames",
|
||||
"generate-da3-depth",
|
||||
|
||||
Reference in New Issue
Block a user