29 lines
864 B
Python
29 lines
864 B
Python
|
import json
|
||
|
from pathlib import Path
|
||
|
from typing import Dict, List
|
||
|
|
||
|
|
||
|
def folder_name_to_show_name(name: str) -> str:
|
||
|
return name.replace("models--", "").replace("--", "/")
|
||
|
|
||
|
|
||
|
def _scan_models(cache_dir, class_name: str) -> List[str]:
|
||
|
cache_dir = Path(cache_dir)
|
||
|
res = []
|
||
|
for it in cache_dir.glob("**/*/model_index.json"):
|
||
|
with open(it, "r", encoding="utf-8") as f:
|
||
|
data = json.load(f)
|
||
|
if data["_class_name"] == class_name:
|
||
|
name = folder_name_to_show_name(it.parent.parent.parent.name)
|
||
|
if name not in res:
|
||
|
res.append(name)
|
||
|
return res
|
||
|
|
||
|
|
||
|
def scan_models(cache_dir) -> List[str]:
|
||
|
return _scan_models(cache_dir, "StableDiffusionPipeline")
|
||
|
|
||
|
|
||
|
def scan_inpainting_models(cache_dir) -> List[str]:
|
||
|
return _scan_models(cache_dir, "StableDiffusionInpaintPipeline")
|