422 lines
17 KiB
Python
422 lines
17 KiB
Python
|
import json
|
||
|
|
import urllib3
|
||
|
|
import urllib.request
|
||
|
|
import os
|
||
|
|
import cgi
|
||
|
|
from pathlib import Path
|
||
|
|
import argparse
|
||
|
|
import string
|
||
|
|
#import getpass
|
||
|
|
#print("Env thinks the user is [%s]" % (os.getlogin()))
|
||
|
|
#print("Effective user is [%s]" % (getpass.getuser()))
|
||
|
|
|
||
|
|
import erisscan
|
||
|
|
|
||
|
|
|
||
|
|
def filter_str(s: str) -> str:
|
||
|
|
return "".join(filter(lambda x: x in string.printable, s))
|
||
|
|
|
||
|
|
|
||
|
|
def downloadcivitai(modelid: str, type: str, image: bool) -> bool:
|
||
|
|
|
||
|
|
match type:
|
||
|
|
case "lora":
|
||
|
|
dlpath = f"./models/Lora/"
|
||
|
|
case "sd":
|
||
|
|
dlpath = f"./models/Stable-diffusion/"
|
||
|
|
case "emb":
|
||
|
|
dlpath = f"./embeddings/"
|
||
|
|
|
||
|
|
http = urllib3.PoolManager()
|
||
|
|
urlid = f"https://civitai.com/api/v1/models/{modelid}"
|
||
|
|
resp = http.request('GET', urlid)
|
||
|
|
modelarray = json.loads(resp.data.decode('utf-8'))
|
||
|
|
modelsha256 = modelarray.get('modelVersions')[0].get('files')[0].get('hashes').get("SHA256")
|
||
|
|
urlhash = f"https://civitai.com/api/v1/model-versions/by-hash/{modelsha256}"
|
||
|
|
modelimageurl = modelarray.get('modelVersions')[0].get('images')[0].get('url')
|
||
|
|
modelfilename = modelarray.get('modelVersions')[0].get('files')[0].get('name')
|
||
|
|
#modelfilesize = int(modelarray.get('modelVersions')[0].get('files')[0].get('sizeKB'))
|
||
|
|
modelname = Path(modelfilename).stem
|
||
|
|
modelfileurl = modelarray.get('modelVersions')[0].get('files')[0].get('downloadUrl')
|
||
|
|
#modeldisplayname = modelarray.get('name')
|
||
|
|
#trainedwords = modelarray.get('modelVersions')[0].get('trainedWords')
|
||
|
|
|
||
|
|
modelimagename = [f"{modelname}.preview.jpg", f"{modelname}.preview.jpeg", f"{modelname}.preview.png"]
|
||
|
|
#modelimagename = ""
|
||
|
|
modelinfofile = f"{modelname}.civitai.info"
|
||
|
|
|
||
|
|
|
||
|
|
def downloadmodel(url: str, dlpath: str) -> bool:
|
||
|
|
print(f"Download {modelfilename} to {dlpath}")
|
||
|
|
dlrequest = http.request('GET', url, preload_content=False)
|
||
|
|
filesize = dlrequest.getheader('content-length')
|
||
|
|
contentdis = dlrequest.getheader('Content-Disposition')
|
||
|
|
try:
|
||
|
|
type, filename = cgi.parse_header(contentdis)
|
||
|
|
#print("Found file", dict(filename).get('filename'))
|
||
|
|
filename = dict(filename).get('filename')
|
||
|
|
#filename = filename.encode('utf-8')
|
||
|
|
#filename = filter_str(filename)
|
||
|
|
except Exception as e:
|
||
|
|
print("Cannot get filename", e)
|
||
|
|
if filesize:
|
||
|
|
filesize = int(filesize)
|
||
|
|
blocksize = max(4096, filesize//100)
|
||
|
|
print(f"Filesize: {filesize//1024} kB | {filesize//1024//1024} MB")
|
||
|
|
else:
|
||
|
|
blocksize = 10000
|
||
|
|
#with open(f"{dlpath}{filename}", 'wb') as out:
|
||
|
|
with open(Path(dlpath) / filename, 'wb') as out:
|
||
|
|
size = 0
|
||
|
|
while True:
|
||
|
|
data = dlrequest.read(blocksize)
|
||
|
|
if not data:
|
||
|
|
break
|
||
|
|
size += len(data)
|
||
|
|
out.write(data)
|
||
|
|
if filesize:
|
||
|
|
Percent = int((size / filesize)*100)
|
||
|
|
if Percent == 100:
|
||
|
|
print(f"Download: {Percent}%")
|
||
|
|
else:
|
||
|
|
print(f"Download: {Percent}%", end='\r')
|
||
|
|
dlrequest.release_conn()
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def downloadimage(url: str, dlpath: str) -> bool:
|
||
|
|
print(f"Download Image for {modelfilename}")
|
||
|
|
imagerequest = http.request('GET', url, preload_content=False)
|
||
|
|
#path = f"{os.getcwd()}/models/Loras/"
|
||
|
|
filesize = imagerequest.getheader('content-length')
|
||
|
|
#contentdis = imagerequest.getheader('Content-Disposition')
|
||
|
|
#type, filename = cgi.parse_header(contentdis)
|
||
|
|
filename = f"{Path(modelfilename).stem}.preview{Path(filename).suffix}"
|
||
|
|
global modelimagename
|
||
|
|
modelimagename = filename
|
||
|
|
#print(filename['filename'])
|
||
|
|
if filesize:
|
||
|
|
filesize = int(filesize)
|
||
|
|
blocksize = max(4096, filesize//100)
|
||
|
|
print(f"Filesize: {filesize//1024} kB | {filesize//1024//1024} MB")
|
||
|
|
else:
|
||
|
|
blocksize = 1000
|
||
|
|
with open(f"{dlpath}{filename}", 'wb') as out:
|
||
|
|
size = 0
|
||
|
|
while True:
|
||
|
|
data = imagerequest.read(blocksize)
|
||
|
|
if not data:
|
||
|
|
break
|
||
|
|
size += len(data)
|
||
|
|
out.write(data)
|
||
|
|
if filesize:
|
||
|
|
Percent = int((size / filesize)*100)
|
||
|
|
if Percent == 100:
|
||
|
|
print(f"Download: {Percent}%")
|
||
|
|
else:
|
||
|
|
print(f"Download: {Percent}%", end='\r')
|
||
|
|
imagerequest.release_conn()
|
||
|
|
|
||
|
|
def downloadinfo(url: str, dlfullpath: str) -> bool:
|
||
|
|
print(f"Download Infofile for {modelfilename}")
|
||
|
|
resp = http.request('GET', url)
|
||
|
|
infofile = open(f"{dlfullpath}", "w")
|
||
|
|
infofile.write(resp.data.decode('utf-8'))
|
||
|
|
infofile.close()
|
||
|
|
|
||
|
|
# download Model file
|
||
|
|
try:
|
||
|
|
if not os.path.exists(f"{dlpath}{modelfilename}"):
|
||
|
|
downloadmodel(modelfileurl, dlpath)
|
||
|
|
else:
|
||
|
|
print(f"{modelfilename} already exist, skip download")
|
||
|
|
except Exception as e:
|
||
|
|
print(e)
|
||
|
|
finally:
|
||
|
|
if not os.path.exists(f"{dlpath}{modelfilename}"):
|
||
|
|
print(f"{modelfilename} could not be saved, abort.")
|
||
|
|
return False
|
||
|
|
# download Preview image, if not disabled
|
||
|
|
if image:
|
||
|
|
try:
|
||
|
|
if os.path.exists(f"{dlpath}{modelimagename[0]}") == False and os.path.exists(f"{dlpath}{modelimagename[1]}") == False and os.path.exists(f"{dlpath}{modelimagename[2]}") == False:
|
||
|
|
downloadimage(modelimageurl, dlpath)
|
||
|
|
else:
|
||
|
|
print(f"Preview Image for {modelfilename} already exist, skip download")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(e)
|
||
|
|
finally:
|
||
|
|
if isinstance(modelfilename, str):
|
||
|
|
if not os.path.exists(f"{dlpath}{modelimagename}"):
|
||
|
|
print(f"{modelfilename} could not be saved, continue.")
|
||
|
|
# save infofile
|
||
|
|
try:
|
||
|
|
if not os.path.exists(f"{dlpath}{modelinfofile}"):
|
||
|
|
downloadinfo(urlhash, f"{dlpath}{modelinfofile}")
|
||
|
|
else:
|
||
|
|
print(f"{modelinfofile} already exist, skip creation")
|
||
|
|
except Exception as e:
|
||
|
|
print(e)
|
||
|
|
finally:
|
||
|
|
if not os.path.exists(f"{dlpath}{modelinfofile}"):
|
||
|
|
print(f"{modelinfofile} could not be created, abort.")
|
||
|
|
return False
|
||
|
|
return True
|
||
|
|
|
||
|
|
# ______________________________________________________________________
|
||
|
|
|
||
|
|
def downloadsideload(name: str, modeltype: str) -> bool:
|
||
|
|
sideloadserver = f"https://eris.hitmare.me/eris/{modeltype}/"
|
||
|
|
|
||
|
|
match modeltype:
|
||
|
|
case "lora":
|
||
|
|
dlpath = "./models/lora/"
|
||
|
|
case "sd":
|
||
|
|
dlpath = "./models/stable-diffusion/"
|
||
|
|
case "emb":
|
||
|
|
dlpath = "./embeddings/"
|
||
|
|
|
||
|
|
http = urllib3.PoolManager()
|
||
|
|
url = f"{sideloadserver}{name}"
|
||
|
|
try:
|
||
|
|
print(f"Download {name}")
|
||
|
|
dlrequest = http.request('GET', url, preload_content=False)
|
||
|
|
|
||
|
|
if not int(dlrequest.status) == 200:
|
||
|
|
print(f"File {name} not found")
|
||
|
|
return False
|
||
|
|
|
||
|
|
filesize = dlrequest.getheader('content-length')
|
||
|
|
#print(filesize)
|
||
|
|
#contentdis = dlrequest.getheader('Content-Disposition')
|
||
|
|
#type, filename = cgi.parse_header(contentdis)
|
||
|
|
filename = name
|
||
|
|
if filesize:
|
||
|
|
filesize = int(filesize)
|
||
|
|
blocksize = max(4096, filesize//100)
|
||
|
|
print(f"Filesize: {filesize//1024} kB | {filesize//1024//1024} MB")
|
||
|
|
else:
|
||
|
|
blocksize = 10000
|
||
|
|
with open(f"{dlpath}{filename}", 'wb') as out:
|
||
|
|
size = 0
|
||
|
|
while True:
|
||
|
|
data = dlrequest.read(blocksize)
|
||
|
|
if not data:
|
||
|
|
break
|
||
|
|
size += len(data)
|
||
|
|
out.write(data)
|
||
|
|
if filesize:
|
||
|
|
Percent = int((size / filesize)*100)
|
||
|
|
if Percent == 100:
|
||
|
|
print(f"Download: {Percent}%")
|
||
|
|
else:
|
||
|
|
print(f"Download: {Percent}%", end='\r')
|
||
|
|
print(f"Downloaded {dlpath}{filename}")
|
||
|
|
dlrequest.release_conn()
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(name, "not found on Sideloadserver", " | Python Error:" ,e)
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
# ______________________________________________________________________
|
||
|
|
# check current status of online and local lora
|
||
|
|
def checkcivitai(locallist: list, erislist: list) -> tuple:
|
||
|
|
#print("Check Civitai")
|
||
|
|
#url = "https://eris.hitmare.me/lora.txt"
|
||
|
|
erismodels = erislist
|
||
|
|
localmodels = locallist
|
||
|
|
missingmodels = set()
|
||
|
|
additionalmodels = set()
|
||
|
|
savedmodels = set()
|
||
|
|
# check online list against local list
|
||
|
|
for loraid in erismodels:
|
||
|
|
if not loraid in localmodels:
|
||
|
|
missingmodels.add(loraid)
|
||
|
|
else:
|
||
|
|
savedmodels.add(loraid)
|
||
|
|
|
||
|
|
# check local list against online list
|
||
|
|
for loraid in localmodels:
|
||
|
|
if not loraid in erismodels:
|
||
|
|
additionalmodels.add(loraid)
|
||
|
|
|
||
|
|
return tuple((missingmodels, additionalmodels, savedmodels))
|
||
|
|
|
||
|
|
|
||
|
|
def checksideload(locallist: list, erislist: list) -> tuple:
|
||
|
|
#print("Check Sideload")
|
||
|
|
#url = "https://eris.hitmare.me/lora.txt"
|
||
|
|
erismodels = erislist
|
||
|
|
localmodels = locallist
|
||
|
|
missingmodels = set()
|
||
|
|
additionalmodels = set()
|
||
|
|
savedmodels = set()
|
||
|
|
# check online list against local list
|
||
|
|
for loraid in erismodels:
|
||
|
|
if not loraid in localmodels:
|
||
|
|
missingmodels.add(loraid)
|
||
|
|
else:
|
||
|
|
savedmodels.add(loraid)
|
||
|
|
|
||
|
|
# check local list against online list
|
||
|
|
for loraid in localmodels:
|
||
|
|
if not loraid in erismodels:
|
||
|
|
additionalmodels.add(loraid)
|
||
|
|
|
||
|
|
return tuple((missingmodels, additionalmodels, savedmodels))
|
||
|
|
|
||
|
|
def deleteFiles(nameid: str, modeltype: str, localfile: bool = False) -> bool:
|
||
|
|
match modeltype:
|
||
|
|
case "lora":
|
||
|
|
dlpath = "./models/lora/"
|
||
|
|
case "sd":
|
||
|
|
dlpath = "./models/stable-diffusion/"
|
||
|
|
case "emb":
|
||
|
|
dlpath = "./embeddings/"
|
||
|
|
|
||
|
|
try:
|
||
|
|
if localfile:
|
||
|
|
modelfilename = nameid
|
||
|
|
else:
|
||
|
|
http = urllib3.PoolManager()
|
||
|
|
urlid = f"https://civitai.com/api/v1/models/{nameid}"
|
||
|
|
resp = http.request('GET', urlid)
|
||
|
|
modelarray = json.loads(resp.data.decode('utf-8'))
|
||
|
|
modelfilename = modelarray.get('modelVersions')[0].get('files')[0].get('name')
|
||
|
|
|
||
|
|
modelfiles = os.listdir(dlpath)
|
||
|
|
for modelfile in modelfiles:
|
||
|
|
if Path(modelfile).stem == Path(modelfilename).stem:
|
||
|
|
os.remove(f"{dlpath}{modelfile}")
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Could not delete {nameid}. Error: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
prog='Eris File Updater',
|
||
|
|
description='Updates Lora, Embeddings and Stable-diffusion Models',
|
||
|
|
epilog='Thanks for being a part of the Eris Bot ^-^')
|
||
|
|
|
||
|
|
parser.add_argument('-si', '--skipimage', help=r"Skips the download of the preview images, Applies to all Model Types", action="store_true")
|
||
|
|
parser.add_argument('-sl', '--skiplora', help=r"Skips check of lora models", action="store_true")
|
||
|
|
parser.add_argument('-ss', '--skipsdmodel', help=r"Skips check of stable-diffusion models", action="store_true")
|
||
|
|
parser.add_argument('-se', '--skipemb', help=r"Skips check of embedding models", action="store_true")
|
||
|
|
parser.add_argument('-c', '--cleanup', help=r"Deletes any models that are locally but not on the eris model list. Affected by Skip Triggers", action="store_true")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
if not os.path.exists("./eris.json"):
|
||
|
|
# start scan python or exit
|
||
|
|
#print("Please run the erisscan bat/sh file first")
|
||
|
|
#exit()
|
||
|
|
print("No eris.json file found. Starting scan ...")
|
||
|
|
erisscan.scanall()
|
||
|
|
|
||
|
|
|
||
|
|
# get the eris list from the web and local
|
||
|
|
http = urllib3.PoolManager()
|
||
|
|
erisurl = "https://eris.hitmare.me/eris.json"
|
||
|
|
resp = http.request('GET', erisurl)
|
||
|
|
erisremote = json.loads(resp.data.decode('utf-8'))
|
||
|
|
erislocal = json.load(open("./eris.json", "r"))
|
||
|
|
|
||
|
|
if args.skipimage:
|
||
|
|
dlimage = False
|
||
|
|
else:
|
||
|
|
dlimage = True
|
||
|
|
|
||
|
|
if args.skiplora:
|
||
|
|
print("Skip Lora Model check")
|
||
|
|
else:
|
||
|
|
lorastate = checkcivitai(erislocal.get('lora')['civitai'], erisremote.get('lora')['civitai'])
|
||
|
|
if len(lorastate[0]) > 0 :
|
||
|
|
print(f"Missing lora ids: {lorastate[0]}")
|
||
|
|
for lora in lorastate[0]:
|
||
|
|
if downloadcivitai(lora, 'lora', dlimage):
|
||
|
|
erislocal.get('lora')['civitai'].append(lora)
|
||
|
|
lorasideloadstate = checksideload(erislocal.get('lora')['sideload'], erisremote.get('lora')['sideload'])
|
||
|
|
if len(lorasideloadstate[0]) > 0 :
|
||
|
|
print(f"Missing lora files: {lorasideloadstate[0]}")
|
||
|
|
for lora in lorasideloadstate[0]:
|
||
|
|
if downloadsideload(lora, 'lora'):
|
||
|
|
erislocal.get('lora')['sideload'].append(lora)
|
||
|
|
if args.cleanup:
|
||
|
|
if len(lorastate[1]) > 0:
|
||
|
|
print("Additional CivitAI lora ids:", lorastate[1])
|
||
|
|
for dellora in lorastate[1]:
|
||
|
|
if deleteFiles(dellora, 'lora'):
|
||
|
|
erislocal.get('lora')['civitai'].remove(dellora)
|
||
|
|
if len(lorasideloadstate[1]) > 0:
|
||
|
|
print("Additional CivitAI lora files:", lorastate[1])
|
||
|
|
for delslora in lorasideloadstate[1]:
|
||
|
|
if deleteFiles(delslora, 'lora', True):
|
||
|
|
erislocal.get('lora')['sideload'].remove(delslora)
|
||
|
|
#print(len(lorastate[1]))
|
||
|
|
if args.skipsdmodel:
|
||
|
|
print("Skip Stable-Diffusion Model check")
|
||
|
|
else:
|
||
|
|
sdstate = checkcivitai(erislocal.get('sd')['civitai'], erisremote.get('sd')['civitai'])
|
||
|
|
if len(sdstate[0]) > 0 :
|
||
|
|
print(f"Missing CivitAI stable-diffusion ids: {sdstate[0]}")
|
||
|
|
for sd in sdstate[0]:
|
||
|
|
if downloadcivitai(sd, 'sd', dlimage):
|
||
|
|
erislocal.get('sd')['civitai'].append(sd)
|
||
|
|
sdsideloadstate = checksideload(erislocal.get('sd')['sideload'], erisremote.get('sd')['sideload'])
|
||
|
|
if len(sdsideloadstate[0]) > 0 :
|
||
|
|
print(f"Missing Sideload stable-diffusion files: {sdsideloadstate[0]}")
|
||
|
|
for sd in sdsideloadstate[0]:
|
||
|
|
if downloadsideload(sd, 'sd'):
|
||
|
|
erislocal.get('sd')['sideload'].append(sd)
|
||
|
|
if args.cleanup:
|
||
|
|
if len(sdstate[1]) > 0:
|
||
|
|
print("Additional CivitAI stable-diffusion ids:", sdstate[1])
|
||
|
|
for delsd in sdstate[1]:
|
||
|
|
if deleteFiles(delsd, 'lora'):
|
||
|
|
erislocal.get('sd')['civitai'].remove(delsd)
|
||
|
|
if len(sdsideloadstate[1]) > 0:
|
||
|
|
print("Additional Sideload stable-diffusion files:", sdstate[1])
|
||
|
|
for delssd in sdsideloadstate[1]:
|
||
|
|
if deleteFiles(delssd, 'lora', True):
|
||
|
|
erislocal.get('sd')['sideload'].remove(delssd)
|
||
|
|
|
||
|
|
if args.skipemb:
|
||
|
|
print("Skip Embedding Model check")
|
||
|
|
else:
|
||
|
|
embstate = checkcivitai(erislocal.get('emb')['civitai'], erisremote.get('emb')['civitai'])
|
||
|
|
if len(embstate[0]) > 0 :
|
||
|
|
print(f"Missing embedding ids: {embstate[0]}")
|
||
|
|
for emb in embstate[0]:
|
||
|
|
if downloadcivitai(emb, 'emb', dlimage):
|
||
|
|
erislocal.get('emb')['civitai'].append(emb)
|
||
|
|
embsideloadstate = checksideload(erislocal.get('emb')['sideload'], erisremote.get('emb')['sideload'])
|
||
|
|
if len(embsideloadstate[0]) > 0 :
|
||
|
|
print(f"Missing embedding files: {embsideloadstate[0]}")
|
||
|
|
for emb in embsideloadstate[0]:
|
||
|
|
if downloadsideload(emb, 'emb'):
|
||
|
|
erislocal.get('emb')['sideload'].append(emb)
|
||
|
|
if args.cleanup:
|
||
|
|
if len(embstate[1]) > 0:
|
||
|
|
print("Additional CivitAI embedding ids:", embstate[1])
|
||
|
|
for delemb in embstate[1]:
|
||
|
|
if deleteFiles(delemb, 'lora'):
|
||
|
|
erislocal.get('emb')['civitai'].remove(delemb)
|
||
|
|
if len(embsideloadstate[1]) > 0:
|
||
|
|
print("Additional Sideload embedding files:", embsideloadstate[1])
|
||
|
|
for delsemb in embsideloadstate[1]:
|
||
|
|
if deleteFiles(delsemb, 'lora', True):
|
||
|
|
erislocal.get('emb')['sideload'].remove(delsemb)
|
||
|
|
|
||
|
|
erisjsonfile = open("./eris.json", "w")
|
||
|
|
erisjsonfile.write(json.dumps(erislocal, indent=2))
|
||
|
|
erisjsonfile.close()
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|