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

@@ -3,7 +3,7 @@ import time
from pathlib import Path
from types import SimpleNamespace
from . import model_catalog
import model_catalog
from .model_download import download_model
@@ -157,11 +157,28 @@ class ModelManager:
self.sessions.close()
def download(self, context, model_key):
return self._download(context, model_key, context.progress)
def download_missing(self, context, model_keys):
missing = [key for key in model_keys if not self.ready(key)]
if not missing:
context.progress(1.0, "Required models are ready")
return {"downloaded": []}
total = len(missing)
for index, model_key in enumerate(missing):
def report(value, message, *, offset=index):
context.progress((offset + value) / total, message)
self._download(context, model_key, report)
return {"downloaded": missing}
def _download(self, context, model_key, progress):
model = model_catalog.get_downloadable_model(model_key)
args = SimpleNamespace(
repository=model.repository,
model_dir=str(self.directory(model_key)),
progress=context.progress,
progress=progress,
cancel_check=context.check_cancelled,
label=model.label,
allow_pattern=list(model.download_patterns),