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:
2026-08-17 14:53:25 +08:00
parent c5b77c5c85
commit 1c98b21235
44 changed files with 637 additions and 2709 deletions

View File

@@ -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",