mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
62 lines
1.7 KiB
Python
Executable File
62 lines
1.7 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
"""
|
|
Scan the models directory and print out a new models.yaml
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
|
|
from pathlib import Path
|
|
from omegaconf import OmegaConf
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Model directory scanner")
|
|
parser.add_argument("models_directory")
|
|
parser.add_argument(
|
|
"--all-models",
|
|
default=False,
|
|
action="store_true",
|
|
help="If true, then generates stanzas for all models; otherwise just diffusers",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
directory = args.models_directory
|
|
|
|
conf = OmegaConf.create()
|
|
conf["_version"] = "3.0.0"
|
|
|
|
for root, dirs, files in os.walk(directory):
|
|
parents = root.split("/")
|
|
subpaths = parents[parents.index("models") + 1 :]
|
|
if len(subpaths) < 2:
|
|
continue
|
|
base, model_type, *_ = subpaths
|
|
|
|
if args.all_models or model_type == "diffusers":
|
|
for d in dirs:
|
|
conf[f"{base}/{model_type}/{d}"] = dict(
|
|
path=os.path.join(root, d),
|
|
description=f"{model_type} model {d}",
|
|
format="folder",
|
|
base=base,
|
|
)
|
|
|
|
for f in files:
|
|
basename = Path(f).stem
|
|
format = Path(f).suffix[1:]
|
|
conf[f"{base}/{model_type}/{basename}"] = dict(
|
|
path=os.path.join(root, f),
|
|
description=f"{model_type} model {basename}",
|
|
format=format,
|
|
base=base,
|
|
)
|
|
|
|
OmegaConf.save(config=dict(sorted(conf.items())), f=sys.stdout)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|