Spaces:
Running
Running
File size: 1,762 Bytes
fcc02a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import os
from typing import List
from toolkit.models.base_model import BaseModel
from toolkit.stable_diffusion_model import StableDiffusion
from toolkit.config_modules import ModelConfig
from toolkit.paths import TOOLKIT_ROOT
import importlib
import pkgutil
from toolkit.models.wan21 import Wan21, Wan21I2V
from toolkit.models.cogview4 import CogView4
BUILT_IN_MODELS = [
Wan21,
Wan21I2V,
CogView4,
]
def get_all_models() -> List[BaseModel]:
extension_folders = ['extensions', 'extensions_built_in']
# This will hold the classes from all extension modules
all_model_classes: List[BaseModel] = BUILT_IN_MODELS
# Iterate over all directories (i.e., packages) in the "extensions" directory
for sub_dir in extension_folders:
extensions_dir = os.path.join(TOOLKIT_ROOT, sub_dir)
for (_, name, _) in pkgutil.iter_modules([extensions_dir]):
try:
# Import the module
module = importlib.import_module(f"{sub_dir}.{name}")
# Get the value of the AI_TOOLKIT_MODELS variable
models = getattr(module, "AI_TOOLKIT_MODELS", None)
# Check if the value is a list
if isinstance(models, list):
# Iterate over the list and add the classes to the main list
all_model_classes.extend(models)
except ImportError as e:
print(f"Failed to import the {name} module. Error: {str(e)}")
return all_model_classes
def get_model_class(config: ModelConfig):
all_models = get_all_models()
for ModelClass in all_models:
if ModelClass.arch == config.arch:
return ModelClass
# default to the legacy model
return StableDiffusion
|