1580158
/
stable-diffusion-webui-aesthetic-image-scorer
mirror of https://github.com/tsngo/stable-diffusion-webui-aesthetic-image-scorer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add more logging and update tools to skip empty tags / categories

main
Trung Ngo 2022-10-27 21:46:38 -05:00
parent 5ccd3004f9
commit 36486df1b2
2 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
scripts/image_scorer.py
View File

@ -13,12 +13,13 @@ import clip
import platform
from launch import is_installed, run_pip
extension_name = "Aesthetic Image Scorer"
if platform.system() == "Windows" and not is_installed("pywin32"):
run_pip(f"install pywin32", "pywin32")
try:
from tools.add_tags import tag_files
except:
print("Aesthetic Image Scorer: Unable to load Windows tagging script")
print(f"{extension_name}: Unable to load Windows tagging script from tools directory")
tag_files = None
state_name = "sac+logos+ava1-l14-linearMSE.pth"
@ -55,7 +56,7 @@ except:
force_cpu = False
if force_cpu:
print("Aesthtic Image Scorer: Forcing prediction model to run on CPU")
print(f"{extension_name}: Forcing prediction model to run on CPU")
device = "cuda" if not force_cpu and torch.cuda.is_available() else "cpu"
# load the model you trained previously or the model available in this repo
pt_state = torch.load(state_name, map_location=torch.device(device=device))
@ -87,7 +88,7 @@ def get_score(image):
def on_ui_settings():
options = {}
options.update(shared.options_section(('ais', "Aesthetic Image Scorer"), {
options.update(shared.options_section(('ais', extension_name), {
"ais_add_exif": OptionInfo(False, "Save score as EXIF or PNG Info Chunk"),
"ais_windows_tag": OptionInfo(False, "Save score as tag (Windows Only)"),
"ais_force_cpu": OptionInfo(False, "Force CPU (Requires Custom Script Reload)"),
@ -115,14 +116,14 @@ def on_image_saved(params: ImageSaveParams):
if score is not None and opts.ais_windows_tag:
if tag_files is not None:
tags = [f"aesthetic_score_{score}"]
tag_files(filename=filename, tags=tags)
tag_files(filename=filename, tags=tags, log_prefix=f"{extension_name}: ")
else:
print("Aesthetic Image Scorer: Unable to load Windows tagging script")
print(f"{extension_name}: Unable to load tagging script")
class AestheticImageScorer(scripts.Script):
def title(self):
return "Aesthetic Image Scorer"
return extension_name
def show(self, is_img2img):
return scripts.AlwaysVisible

2
tools

@ -1 +1 @@
Subproject commit c2bf6907100616d55f38f66327dd89e8e7c2c96e
Subproject commit cb92f1a566ba9c6e0dcbb2393d69f7bc4bd0a9e2