diff --git a/README.md b/README.md index 5ec6061..a8ebd72 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,22 @@ # Aesthetic Image Scorer +![](tag_group_by.png) + Calculates aestetic score for generated images using [CLIP+MLP Aesthetic Score Predictor](https://github.com/christophschuhmann/improved-aesthetic-predictor) based on [Chad Scorer](https://github.com/grexzen/SD-Chad/blob/main/chad_scorer.py) See [Discussion](https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/1831) -Saves score to windows tags with other options planned \ No newline at end of file +Saves score to windows tags with other options planned + +## Installation +Clone the repo into the `extensions` directory and restart the web ui: + +```commandline +git clone https://github.com/tsngo/stable-diffusion-webui-aesthetic-image-scorer extensions/aesthetic-image-scorer +``` + +## Features +- Save score as EXIF or PNG Info Chunk (Only PNG Chunk Info currently) +- Save score as tag (Windows Only) + - Added to tags as `aesthetic_score_5.9` + - JPG supports by default. PNG tags requires a 3rd party software like [File Metadata](https://github.com/Dijji/FileMeta/releases) diff --git a/scripts/image_scorer.py b/scripts/image_scorer.py index 303d1bd..eef2a1e 100644 --- a/scripts/image_scorer.py +++ b/scripts/image_scorer.py @@ -15,7 +15,7 @@ if platform.system() == "Windows" and not is_installed("pywin32"): try: from tools.add_tags import tag_files except: - print("Unable to load") + print("Aesthetic Image Scorer: Unable to load Windows tagging script") tag_files = None state_name = "sac+logos+ava1-l14-linearMSE.pth" @@ -26,19 +26,6 @@ if not Path(state_name).exists(): with open(state_name, "wb") as f: f.write(r.content) -class AestheticImageScorer: - def __init__(self): - self.ais_windows_tag = False - - def set_params(self, p, ais_windows_tag=False): - self.ais_windows_tag = ais_windows_tag - - p.extra_generation_params.update({ - "AIS Windows Tag": ais_windows_tag, - }) - -ais = AestheticImageScorer() - class AestheticPredictor(nn.Module): def __init__(self, input_size): super().__init__() @@ -57,7 +44,14 @@ class AestheticPredictor(nn.Module): def forward(self, x): return self.layers(x) -device = "cuda" if torch.cuda.is_available() else "cpu" +try: + force_cpu = opts.ais_force_cpu +except: + force_cpu = False + +if force_cpu: + print("Aesthtic Image Scorer: 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,19 +81,37 @@ def get_score(image): def on_ui_settings(): options = {} options.update(shared.options_section(('ais', "Aesthetic Image Scorer"), { + "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)"), })) + + opts.add_option("ais_add_exif", options["ais_add_exif"]) opts.add_option("ais_windows_tag", options["ais_windows_tag"]) + opts.add_option("ais_force_cpu", options["ais_force_cpu"]) -def on_save_imaged(image, p, fullfn, txt_fullfn): - score = round(get_score(image), 1) - if opts.ais_windows_tag: +def on_before_image_saved(image, p, **kwargs): + if opts.ais_add_exif: + score = round(get_score(image), 1) + if "existing_info" not in kwargs or kwargs["existing_info"] is None: + kwargs["existing_info"] = {} + kwargs["existing_info"].update({ + "aesthetic_score": score, + }) + return image, p, kwargs + +def on_image_saved(image, p, fullfn, txt_fullfn, **kwargs): + if "existing_info" in kwargs and kwargs["existing_info"] is not None and "aesthetic_score" in kwargs["existing_info"]: + score = kwargs["existing_info"]["aesthetic_score"] + else: + score = round(get_score(image), 1) + if score is not None and opts.ais_windows_tag: if tag_files is not None: tags = [f"aesthetic_score_{score}"] tag_files(filename=fullfn, tags=tags) else: - print("Unable to load windows tagging script") + print("Aesthetic Image Scorer: Unable to load Windows tagging script") class AestheticImageScorer(scripts.Script): def title(self): @@ -112,7 +124,8 @@ class AestheticImageScorer(scripts.Script): return [] def process(self, p): - ais.set_params(p, bool(opts.ais_windows_tag)) + pass script_callbacks.on_ui_settings(on_ui_settings) -script_callbacks.on_save_imaged(on_save_imaged) +script_callbacks.on_before_image_saved(on_before_image_saved) +script_callbacks.on_image_saved(on_image_saved) diff --git a/tag_group_by.png b/tag_group_by.png new file mode 100644 index 0000000..c573ea3 Binary files /dev/null and b/tag_group_by.png differ diff --git a/tools b/tools index 21b75ff..2d2e2d7 160000 --- a/tools +++ b/tools @@ -1 +1 @@ -Subproject commit 21b75ff7655c12986fdd5a36a9a3f49b0630b9d2 +Subproject commit 2d2e2d7f673fe8abdebfa29beda07312c12992de