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", encoding='utf-8') 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", encoding='utf-8') 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()