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