Eris-Update-script/erisscan.py

104 lines
4.9 KiB
Python

import os
import json
import hashlib
import urllib3
from pathlib import Path
def scanfolder(path: str) -> list:
civitailist = []
sideloaddict = []
modelfolder = path
modelfiles = os.listdir(modelfolder)
for modelfile in modelfiles:
modelname = Path(modelfile).stem
sideload = False
if modelfile.endswith('.safetensors') or modelfile.endswith('.pt'):
#print(f"Found '{lorafile}' - Searching civitai.info")
if os.path.exists(f"{modelfolder}{modelname}.civitai.info"):
civitaiinfo = open(f"{modelfolder}{modelname}.civitai.info", "r", encoding='utf-8')
civitaiinfocontent = civitaiinfo.read()
civitaiinfo.close()
if civitaiinfocontent[:2] == "{}":
os.remove(f"{modelfolder}{modelname}.civitai.info")
sideload = True
else:
with open(f"{modelfolder}{modelname}.civitai.info", "r") as f:
#fstring = "dd"
print("Reading", modelfolder, modelfile)
modelinfoarray = json.load(f)
#print(loraarray)
try:
modelsha256 = modelinfoarray.get('files')[0].get('hashes').get('SHA256')
modelid = modelinfoarray.get('modelId')
civitailist.append(modelid)
#print(lorasha256, " " , loramodelid)
except Exception as e:
print("[Model Step local civitai.info parse] ", path, modelfile , " | ",e)
#print("Content:", civitaiinfocontent[:2])
else:
sideload = True
if sideload:
print(f"{modelfolder}{modelname}.civitai.info not found. Getting SHA256 Hash and fetch data from civitai API")
with open(f"{modelfolder}{modelfile}", "rb") as f:
bytes = f.read() # read entire file as bytes
modelsha256 = hashlib.sha256(bytes).hexdigest();
#print(readable_hash)
http = urllib3.PoolManager()
url = f"https://civitai.com/api/v1/model-versions/by-hash/{modelsha256}"
resp = http.request('GET', url)
modelinfoarray = json.loads(resp.data.decode('utf-8'))
try:
if modelinfoarray.get('error') == "Model not found":
print(f"Model {modelfile} is not a CivitAi Model")
sideloaddict.append(modelfile)
if os.path.exists(f"{modelfolder}{modelname}.preview.jpg"):
sideloaddict.append(f"{modelname}.preview.jpg")
if os.path.exists(f"{modelfolder}{modelname}.preview.jpeg"):
sideloaddict.append(f"{modelname}.preview.jpeg")
if os.path.exists(f"{modelfolder}{modelname}.preview.png"):
sideloaddict.append(f"{modelname}.preview.png")
else:
print(f"Model {modelfile} found on CivitAi")
modelid = modelinfoarray.get('modelId')
civitailist.append(modelid)
print(f"Creating {modelfolder}{modelname}.civitai.info for future use")
modelinfofile = open(f"{modelfolder}{modelname}.civitai.info", "w")
modelinfofile.write(resp.data.decode('utf-8'))
modelinfofile.close()
except Exception as e:
print("[Model Step SHA256 civitai.info parse] ", path, modelfile , " | ",e)
#if os.path.exists(f"{modelfolder}{modelname}.civitai.info"):
#print("Civitai Info file found.")
# do json stuff
#else:
return list([civitailist, sideloaddict])
def scanall():
erislist = {}
erislist['lora'] = {}
result = scanfolder("./models/lora/")
erislist['lora']['civitai'] = result[0]
erislist['lora']['sideload'] = result[1]
erislist['sd'] = {}
result = scanfolder("./models/stable-diffusion/")
erislist['sd']['civitai'] = result[0]
erislist['sd']['sideload'] = result[1]
erislist['emb'] = {}
result = scanfolder("./embeddings/")
erislist['emb']['civitai'] = result[0]
erislist['emb']['sideload'] = result[1]
erisjsonfile = open("./eris.json", "w")
erisjsonfile.write(json.dumps(erislist, indent=2))
erisjsonfile.close()
if __name__ == "__main__":
scanall()