104 lines
4.9 KiB
Python
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()
|