304 lines
13 KiB
Python
304 lines
13 KiB
Python
import os
|
|
import gradio as gr
|
|
import datetime
|
|
import requests
|
|
import shutil
|
|
import json
|
|
|
|
from . import util
|
|
from . import model
|
|
from . import setting
|
|
from . import civitai
|
|
|
|
from . import ishortcut
|
|
from . import ishortcut_action
|
|
|
|
def on_scan_ui():
|
|
with gr.Column():
|
|
with gr.Row():
|
|
with gr.Accordion("Scan models for Civitai", open=True):
|
|
with gr.Row():
|
|
with gr.Column():
|
|
fix_information_filename = gr.Checkbox(label="Fix version information filename", value=False , visible=False)
|
|
scan_models_btn = gr.Button(value="Scan Models",variant="primary")
|
|
gr.Markdown(value="This feature targets models that do not have information files available in the saved models. It calculates the hash value and searches for the model in Civitai, registering it as a shortcut. Calculating the hash value can take a significant amount of time.", visible=True)
|
|
with gr.Box(elem_classes="cs_box", visible=False) as scanned_result:
|
|
scan_models_result = gr.CheckboxGroup(visible=True, label="Scanned Model List").style(item_container=True,container=True)
|
|
with gr.Row(visible=False) as update_information:
|
|
with gr.Column():
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
scan_register_shortcut = gr.Checkbox(label="Register a shortcut when creating the model information file.", value=True)
|
|
with gr.Column(scale=1):
|
|
with gr.Row():
|
|
scan_save_modelfolder = gr.Checkbox(label="Create a model folder corresponding to the model type.", value=False)
|
|
scan_save_vsfolder = gr.Checkbox(label="Create individual model version folder.", value=False, interactive=False)
|
|
with gr.Row():
|
|
with gr.Column():
|
|
create_models_info_btn = gr.Button(value="Create Model Information",variant="primary")
|
|
with gr.Row():
|
|
with gr.Accordion("Update Shortcuts", open=True):
|
|
with gr.Row():
|
|
with gr.Column():
|
|
scan_to_shortcut_btn = gr.Button(value="Scan downloaded models for shortcut registration",variant="primary")
|
|
scan_progress = gr.Markdown(value="This feature scans for models that have information files available and registers a shortcut for them, downloading any necessary images in the process. If there is no information available for a particular model, please use the 'Scan Models' feature.", visible=True)
|
|
|
|
with gr.Row():
|
|
with gr.Column():
|
|
update_all_shortcuts_btn = gr.Button(value="Update the model information for the shortcut",variant="primary")
|
|
update_progress = gr.Markdown(value="This feature updates registered shortcuts with the latest information and downloads any new images if available.", visible=True)
|
|
|
|
scan_save_modelfolder.change(
|
|
fn=on_scan_save_modelfolder_change,
|
|
inputs=[
|
|
scan_save_modelfolder
|
|
],
|
|
outputs=[
|
|
scan_save_vsfolder
|
|
]
|
|
)
|
|
|
|
create_models_info_btn.click(
|
|
fn=on_create_models_info_btn_click,
|
|
inputs=[
|
|
scan_models_result,
|
|
scan_save_modelfolder,
|
|
scan_save_vsfolder,
|
|
scan_register_shortcut
|
|
],
|
|
outputs=[
|
|
scan_models_result,
|
|
scanned_result,
|
|
update_information
|
|
]
|
|
)
|
|
|
|
scan_models_btn.click(
|
|
fn=on_scan_models_btn_click,
|
|
inputs=[fix_information_filename],
|
|
outputs=[
|
|
scan_models_result,
|
|
scanned_result,
|
|
update_information,
|
|
scan_save_modelfolder,
|
|
scan_save_vsfolder
|
|
]
|
|
)
|
|
|
|
update_all_shortcuts_btn.click(
|
|
fn=on_update_all_shortcuts_btn,
|
|
inputs=None,
|
|
outputs=[
|
|
update_progress,
|
|
]
|
|
)
|
|
|
|
update_progress.change(
|
|
fn=on_update_progress_change,
|
|
inputs=None,
|
|
outputs=[update_progress]
|
|
)
|
|
|
|
scan_to_shortcut_btn.click(
|
|
fn=on_scan_to_shortcut_click,
|
|
inputs=None,
|
|
outputs=[
|
|
scan_progress,
|
|
]
|
|
)
|
|
|
|
scan_progress.change(
|
|
fn=on_scan_progress_change,
|
|
inputs=None,
|
|
outputs=[scan_progress]
|
|
)
|
|
|
|
def create_models_information(files, mfolder, vs_folder, register_shortcut, progress=gr.Progress()):
|
|
|
|
non_list = list()
|
|
if not files:
|
|
return None
|
|
|
|
for file_path in progress.tqdm(files, desc=f"Create Models Information"):
|
|
if os.path.isfile(file_path):
|
|
util.printD(f"Generate SHA256: {file_path}")
|
|
hash = util.calculate_sha256(file_path)
|
|
version_info = civitai.get_version_info_by_hash(hash)
|
|
|
|
if not version_info:
|
|
# These models are not registered with Civitai.
|
|
non_list.append(file_path)
|
|
continue
|
|
|
|
vfolder , vfile = os.path.split(file_path)
|
|
basename , ext = os.path.splitext(vfile)
|
|
|
|
# 저장할 폴더 생성
|
|
if mfolder:
|
|
model_folder = util.make_download_model_folder(version_info, True, vs_folder)
|
|
# 다정하면 임의의 분류뒤에 모델폴더를 생성하고 그뒤에 버전까지 생성가능
|
|
# model_folder = make_download_model_folder(version_info, ms_folder=True, vs_folder=True, vs_foldername=None, cs_foldername=None):
|
|
# model_folder = util.make_version_folder(version_info, vs_folder)
|
|
else:
|
|
model_folder = vfolder
|
|
|
|
# version info file name 으로 교체시
|
|
# savefile_base = downloader.get_save_base_name(version_info)
|
|
# basename = savefile_base
|
|
# destination = os.path.join(model_folder, f"{basename}{ext}")
|
|
|
|
# save info
|
|
info_path = os.path.join(model_folder, f"{basename}{setting.info_suffix}{setting.info_ext}")
|
|
result = civitai.write_version_info(info_path, version_info)
|
|
if result:
|
|
util.printD(f"Wrote version info : {info_path}")
|
|
|
|
# save preview
|
|
if "images" in version_info.keys():
|
|
description_img = os.path.join(model_folder, f"{basename}{setting.preview_image_suffix}{setting.preview_image_ext}")
|
|
try:
|
|
img_dict = version_info["images"][0]
|
|
if "url" in img_dict:
|
|
img_url = img_dict["url"]
|
|
if "width" in img_dict:
|
|
if img_dict["width"]:
|
|
img_url = util.change_width_from_image_url(img_url, img_dict["width"])
|
|
# get image
|
|
with requests.get(img_url, stream=True) as img_r:
|
|
if not img_r.ok:
|
|
util.printD("Get error code: " + str(img_r.status_code))
|
|
return
|
|
|
|
with open(description_img, 'wb') as f:
|
|
img_r.raw.decode_content = True
|
|
shutil.copyfileobj(img_r.raw, f)
|
|
util.printD(f"Downloaded preview image : {description_img}")
|
|
except Exception as e:
|
|
pass
|
|
|
|
# 파일 이동
|
|
if mfolder:
|
|
destination = os.path.join(model_folder, vfile)
|
|
if file_path != destination:
|
|
if not os.path.isfile(destination):
|
|
os.rename(file_path, destination)
|
|
else:
|
|
util.printD(f"The target file already exists : target {destination}")
|
|
|
|
# 숏컷 추가
|
|
if register_shortcut:
|
|
if version_info['modelId']:
|
|
ishortcut.update_shortcut(version_info['modelId'], progress)
|
|
model.update_downloaded_model()
|
|
|
|
return non_list
|
|
|
|
def is_filename_in_version_info_in_directory(directory, filename):
|
|
|
|
file_list = []
|
|
for file in os.listdir(directory):
|
|
if file.endswith(f"{setting.info_suffix}{setting.info_ext}"):
|
|
file_list.append(os.path.join(directory,file))
|
|
|
|
if not file_list:
|
|
return False
|
|
|
|
for file in file_list:
|
|
try:
|
|
with open(file, 'r') as f:
|
|
json_data = json.load(f)
|
|
if "files" in json_data.keys():
|
|
files = json_data['files']
|
|
for file in files:
|
|
if file['name'] == filename:
|
|
return True
|
|
except:
|
|
pass
|
|
|
|
return False
|
|
|
|
def scan_models(fix_information_filename, progress=gr.Progress()):
|
|
root_dirs = list(set(setting.model_folders.values()))
|
|
file_list = util.search_file(root_dirs,None,setting.model_exts)
|
|
|
|
result = list()
|
|
|
|
if fix_information_filename:
|
|
# fix_version_information_filename()
|
|
pass
|
|
|
|
for file_path in progress.tqdm(file_list, desc=f"Scan Models for Civitai"):
|
|
|
|
vfolder , vfile = os.path.split(file_path)
|
|
basename , ext = os.path.splitext(vfile)
|
|
info = os.path.join(vfolder, f"{basename}{setting.info_suffix}{setting.info_ext}")
|
|
|
|
if not os.path.isfile(info):
|
|
# result.append(file_path)
|
|
if not is_filename_in_version_info_in_directory(vfolder, vfile):
|
|
# util.printD(f"{file_path} : {vfile}: no info")
|
|
result.append(file_path)
|
|
|
|
return result
|
|
|
|
# def fix_version_information_filename():
|
|
# root_dirs = list(set(setting.model_folders.values()))
|
|
# file_list = util.search_file(root_dirs,None,[setting.info_ext])
|
|
|
|
# version_info = None
|
|
# if not file_list:
|
|
# return
|
|
|
|
# for file_path in file_list:
|
|
|
|
# try:
|
|
# with open(file_path, 'r') as f:
|
|
# json_data = json.load(f)
|
|
|
|
# if 'id' in json_data.keys():
|
|
# version_info = json_data
|
|
|
|
# file_path = file_path.strip()
|
|
# vfolder , vfile = os.path.split(file_path)
|
|
# savefile_base = downloader.get_save_base_name(version_info)
|
|
# info_file = os.path.join(vfolder, f"{util.replace_filename(savefile_base)}{setting.info_suffix}{setting.info_ext}")
|
|
|
|
# if file_path != info_file:
|
|
# if not os.path.isfile(info_file):
|
|
# os.rename(file_path, info_file)
|
|
|
|
# except:
|
|
# pass
|
|
|
|
def on_create_models_info_btn_click(files, mfolder, vsfolder, register_shortcut, progress=gr.Progress()):
|
|
remain_files = create_models_information(files,mfolder,vsfolder,register_shortcut, progress)
|
|
if remain_files and len(remain_files) > 0:
|
|
return gr.update(choices=remain_files, value=remain_files, interactive=True, label="These models are not registered with Civitai."),gr.update(visible=True),gr.update(visible=True)
|
|
return gr.update(choices=[], value=[], interactive=True),gr.update(visible=False),gr.update(visible=False)
|
|
|
|
def on_update_progress_change():
|
|
current_time = datetime.datetime.now()
|
|
return gr.update(value=current_time)
|
|
|
|
def on_scan_progress_change():
|
|
current_time = datetime.datetime.now()
|
|
return gr.update(value=current_time)
|
|
|
|
def on_scan_models_btn_click(fix_information_filename, progress=gr.Progress()):
|
|
files = scan_models(fix_information_filename, progress)
|
|
return gr.update(choices=files,value=files,interactive=True,label="Scanned Model List"),gr.update(visible=True),gr.update(visible=True),gr.update(value=False, interactive=True),gr.update(value=False, interactive=False)
|
|
|
|
def on_scan_to_shortcut_click(progress=gr.Progress()):
|
|
model.update_downloaded_model()
|
|
ishortcut_action.scan_downloadedmodel_to_shortcut(progress)
|
|
return gr.update(value="This feature scans for models that have information files available and registers a shortcut for them, downloading any necessary images in the process. If there is no information available for a particular model, please use the 'Scan Models' feature.")
|
|
|
|
def on_update_all_shortcuts_btn(progress=gr.Progress()):
|
|
ishortcut.update_all_shortcut_informations(progress)
|
|
return gr.update(value="This feature updates registered shortcuts with the latest information and downloads any new images if available.")
|
|
|
|
def on_scan_save_modelfolder_change(scan_save_modelfolder):
|
|
if scan_save_modelfolder:
|
|
return gr.update(interactive=True)
|
|
return gr.update(value=False, interactive=False)
|
|
|