1580158
/
stable-diffusion-aws-extension
mirror of https://github.com/awslabs/stable-diffusion-aws-extension.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

update move mdel to tmp

feature/de-webui
yuxiao 2024-01-15 13:57:04 +00:00
parent 7be57d3a40
commit c56280d3f9
1 changed files with 50 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
build_scripts/api/modules/api/api.py
View File

@ -525,44 +525,62 @@ def get_file_md5_dict(path):
file_dict[file] = md5(os.path.join(root, file))
return file_dict
def move_model_to_tmp(_, app: FastAPI):
# Move model dir to /tmp
print('!!!!!!!!!!!!!')
logging.info("Copy model dir to tmp")
model_tmp_dir = f"models_{time.time()}"
os.system(f"cp -rL models /tmp/{model_tmp_dir}")
src_file_dict = get_file_md5_dict("models")
tgt_file_dict = get_file_md5_dict(f"/tmp/{model_tmp_dir}")
is_complete = True
for file in src_file_dict:
logging.info(f"Src file {file} md5 {src_file_dict[file]}")
if file not in tgt_file_dict:
is_complete = False
break
if src_file_dict[file] != tgt_file_dict[file]:
is_complete = False
break
if is_complete:
os.system(f"rm -rf models")
logging.info("Link model dir")
os.system(f"ln -s /tmp/{model_tmp_dir} models")
else:
logging.info("Failed to copy model dir, use the original dir")
logging.info("Check disk usage on app started")
os.system("df -h")
# def move_model_to_tmp(_, app: FastAPI):
# # Move model dir to /tmp
# print('!!!!!!!!!!!!!')
# logging.info("Copy model dir to tmp")
# model_tmp_dir = f"models_{time.time()}"
# os.system(f"cp -rL models /tmp/{model_tmp_dir}")
# src_file_dict = get_file_md5_dict("models")
# tgt_file_dict = get_file_md5_dict(f"/tmp/{model_tmp_dir}")
# is_complete = True
# for file in src_file_dict:
# logging.info(f"Src file {file} md5 {src_file_dict[file]}")
# if file not in tgt_file_dict:
# is_complete = False
# break
# if src_file_dict[file] != tgt_file_dict[file]:
# is_complete = False
# break
# if is_complete:
# os.system(f"rm -rf models")
# logging.info("Link model dir")
# os.system(f"ln -s /tmp/{model_tmp_dir} models")
# else:
# logging.info("Failed to copy model dir, use the original dir")
# logging.info("Check disk usage on app started")
# os.system("df -h")
try:
print('!!!!!!!!! on docker decision')
import modules.script_callbacks as script_callbacks
print('!!!!!!!!! on docker condition')
on_docker = os.environ.get('ON_DOCKER', "false")
if on_docker == "true":
print('!!!!!!!!! move model to tmp')
from modules import shared
shared.opts.data.update(control_net_max_models_num=10)
script_callbacks.on_app_started(move_model_to_tmp)
logger.debug("Diffusers API layer loaded")
# Move model dir to /tmp
logging.info("Copy model dir to tmp")
model_tmp_dir = f"models_{time.time()}"
os.system(f"cp -rL models /tmp/{model_tmp_dir}")
src_file_dict = get_file_md5_dict("models")
tgt_file_dict = get_file_md5_dict(f"/tmp/{model_tmp_dir}")
is_complete = True
for file in src_file_dict:
logging.info(f"Src file {file} md5 {src_file_dict[file]}")
if file not in tgt_file_dict:
is_complete = False
break
if src_file_dict[file] != tgt_file_dict[file]:
is_complete = False
break
if is_complete:
print('!!!!!!!!!, link to tmp')
os.system(f"rm -rf models")
logging.info("Link model dir")
os.system(f"ln -s /tmp/{model_tmp_dir} models")
else:
logging.info("Failed to copy model dir, use the original dir")
logging.info("Check disk usage on app started")
os.system("df -h")
except Exception as e:
logger.error(e)
logger.debug("Unable to import script callbacks.")
pass
pass