mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
9 lines
288 B
Python
9 lines
288 B
Python
from toolkit.stable_diffusion_model import StableDiffusion
|
|
from toolkit.config_modules import ModelConfig
|
|
|
|
def get_model_class(config: ModelConfig):
|
|
if config.arch == "wan21":
|
|
from toolkit.models.wan21 import Wan21
|
|
return Wan21
|
|
else:
|
|
return StableDiffusion |