From 00d3d1b4b31b2effa32f6eb96f8e5caf6368f8e3 Mon Sep 17 00:00:00 2001 From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Date: Sat, 18 May 2024 15:50:28 +0200 Subject: [PATCH 01/12] feat: add nsfw image censoring via config and checkbox (#958) * add nsfw image censoring activatable via config, uses CompVis/stable-diffusion-safety-checker * fix progressbar call for nsfw output * use config to set cache dir for safety checker * add checkbox black_out_nsfw makes both enabling via config and checkbox possible, where config overrides the checkbox value * fix: add missing diffusers package * feat: extract safety checker, remove dependency to diffusers * feat: make code compatible again after merge with main * feat: move censor to extras, optimize safety checker file handling * refactor: rename folder safety_checker_models to safety_checker --- extras/censor.py | 56 ++++++ extras/safety_checker/configs/config.json | 171 ++++++++++++++++++ .../configs/preprocessor_config.json | 20 ++ .../safety_checker/models/safety_checker.py | 126 +++++++++++++ language/en.json | 2 + .../put_safety_checker_models_here | 0 modules/async_worker.py | 32 +++- modules/config.py | 14 ++ webui.py | 14 +- 9 files changed, 424 insertions(+), 11 deletions(-) create mode 100644 extras/censor.py create mode 100644 extras/safety_checker/configs/config.json create mode 100644 extras/safety_checker/configs/preprocessor_config.json create mode 100644 extras/safety_checker/models/safety_checker.py create mode 100644 models/safety_checker/put_safety_checker_models_here diff --git a/extras/censor.py b/extras/censor.py new file mode 100644 index 00000000..2047db24 --- /dev/null +++ b/extras/censor.py @@ -0,0 +1,56 @@ +# modified version of https://github.com/AUTOMATIC1111/stable-diffusion-webui-nsfw-censor/blob/master/scripts/censor.py +import numpy as np +import os + +from extras.safety_checker.models.safety_checker import StableDiffusionSafetyChecker +from transformers import CLIPFeatureExtractor, CLIPConfig +from PIL import Image +import modules.config + +safety_checker_repo_root = os.path.join(os.path.dirname(__file__), 'safety_checker') +config_path = os.path.join(safety_checker_repo_root, "configs", "config.json") +preprocessor_config_path = os.path.join(safety_checker_repo_root, "configs", "preprocessor_config.json") + +safety_feature_extractor = None +safety_checker = None + + +def numpy_to_pil(image): + image = (image * 255).round().astype("uint8") + pil_image = Image.fromarray(image) + + return pil_image + + +# check and replace nsfw content +def check_safety(x_image): + global safety_feature_extractor, safety_checker + + if safety_feature_extractor is None or safety_checker is None: + safety_checker_model = modules.config.downloading_safety_checker_model() + safety_feature_extractor = CLIPFeatureExtractor.from_json_file(preprocessor_config_path) + clip_config = CLIPConfig.from_json_file(config_path) + safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_checker_model, config=clip_config) + + safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt") + x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values) + + return x_checked_image, has_nsfw_concept + + +def censor_single(x): + x_checked_image, has_nsfw_concept = check_safety(x) + + # replace image with black pixels, keep dimensions + # workaround due to different numpy / pytorch image matrix format + if has_nsfw_concept[0]: + imageshape = x_checked_image.shape + x_checked_image = np.zeros((imageshape[0], imageshape[1], 3), dtype = np.uint8) + + return x_checked_image + + +def censor_batch(images): + images = [censor_single(image) for image in images] + + return images \ No newline at end of file diff --git a/extras/safety_checker/configs/config.json b/extras/safety_checker/configs/config.json new file mode 100644 index 00000000..aa454d22 --- /dev/null +++ b/extras/safety_checker/configs/config.json @@ -0,0 +1,171 @@ +{ + "_name_or_path": "clip-vit-large-patch14/", + "architectures": [ + "SafetyChecker" + ], + "initializer_factor": 1.0, + "logit_scale_init_value": 2.6592, + "model_type": "clip", + "projection_dim": 768, + "text_config": { + "_name_or_path": "", + "add_cross_attention": false, + "architectures": null, + "attention_dropout": 0.0, + "bad_words_ids": null, + "bos_token_id": 0, + "chunk_size_feed_forward": 0, + "cross_attention_hidden_size": null, + "decoder_start_token_id": null, + "diversity_penalty": 0.0, + "do_sample": false, + "dropout": 0.0, + "early_stopping": false, + "encoder_no_repeat_ngram_size": 0, + "eos_token_id": 2, + "exponential_decay_length_penalty": null, + "finetuning_task": null, + "forced_bos_token_id": null, + "forced_eos_token_id": null, + "hidden_act": "quick_gelu", + "hidden_size": 768, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_factor": 1.0, + "initializer_range": 0.02, + "intermediate_size": 3072, + "is_decoder": false, + "is_encoder_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "layer_norm_eps": 1e-05, + "length_penalty": 1.0, + "max_length": 20, + "max_position_embeddings": 77, + "min_length": 0, + "model_type": "clip_text_model", + "no_repeat_ngram_size": 0, + "num_attention_heads": 12, + "num_beam_groups": 1, + "num_beams": 1, + "num_hidden_layers": 12, + "num_return_sequences": 1, + "output_attentions": false, + "output_hidden_states": false, + "output_scores": false, + "pad_token_id": 1, + "prefix": null, + "problem_type": null, + "pruned_heads": {}, + "remove_invalid_values": false, + "repetition_penalty": 1.0, + "return_dict": true, + "return_dict_in_generate": false, + "sep_token_id": null, + "task_specific_params": null, + "temperature": 1.0, + "tie_encoder_decoder": false, + "tie_word_embeddings": true, + "tokenizer_class": null, + "top_k": 50, + "top_p": 1.0, + "torch_dtype": null, + "torchscript": false, + "transformers_version": "4.21.0.dev0", + "typical_p": 1.0, + "use_bfloat16": false, + "vocab_size": 49408 + }, + "text_config_dict": { + "hidden_size": 768, + "intermediate_size": 3072, + "num_attention_heads": 12, + "num_hidden_layers": 12 + }, + "torch_dtype": "float32", + "transformers_version": null, + "vision_config": { + "_name_or_path": "", + "add_cross_attention": false, + "architectures": null, + "attention_dropout": 0.0, + "bad_words_ids": null, + "bos_token_id": null, + "chunk_size_feed_forward": 0, + "cross_attention_hidden_size": null, + "decoder_start_token_id": null, + "diversity_penalty": 0.0, + "do_sample": false, + "dropout": 0.0, + "early_stopping": false, + "encoder_no_repeat_ngram_size": 0, + "eos_token_id": null, + "exponential_decay_length_penalty": null, + "finetuning_task": null, + "forced_bos_token_id": null, + "forced_eos_token_id": null, + "hidden_act": "quick_gelu", + "hidden_size": 1024, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "image_size": 224, + "initializer_factor": 1.0, + "initializer_range": 0.02, + "intermediate_size": 4096, + "is_decoder": false, + "is_encoder_decoder": false, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "layer_norm_eps": 1e-05, + "length_penalty": 1.0, + "max_length": 20, + "min_length": 0, + "model_type": "clip_vision_model", + "no_repeat_ngram_size": 0, + "num_attention_heads": 16, + "num_beam_groups": 1, + "num_beams": 1, + "num_hidden_layers": 24, + "num_return_sequences": 1, + "output_attentions": false, + "output_hidden_states": false, + "output_scores": false, + "pad_token_id": null, + "patch_size": 14, + "prefix": null, + "problem_type": null, + "pruned_heads": {}, + "remove_invalid_values": false, + "repetition_penalty": 1.0, + "return_dict": true, + "return_dict_in_generate": false, + "sep_token_id": null, + "task_specific_params": null, + "temperature": 1.0, + "tie_encoder_decoder": false, + "tie_word_embeddings": true, + "tokenizer_class": null, + "top_k": 50, + "top_p": 1.0, + "torch_dtype": null, + "torchscript": false, + "transformers_version": "4.21.0.dev0", + "typical_p": 1.0, + "use_bfloat16": false + }, + "vision_config_dict": { + "hidden_size": 1024, + "intermediate_size": 4096, + "num_attention_heads": 16, + "num_hidden_layers": 24, + "patch_size": 14 + } +} diff --git a/extras/safety_checker/configs/preprocessor_config.json b/extras/safety_checker/configs/preprocessor_config.json new file mode 100644 index 00000000..5294955f --- /dev/null +++ b/extras/safety_checker/configs/preprocessor_config.json @@ -0,0 +1,20 @@ +{ + "crop_size": 224, + "do_center_crop": true, + "do_convert_rgb": true, + "do_normalize": true, + "do_resize": true, + "feature_extractor_type": "CLIPFeatureExtractor", + "image_mean": [ + 0.48145466, + 0.4578275, + 0.40821073 + ], + "image_std": [ + 0.26862954, + 0.26130258, + 0.27577711 + ], + "resample": 3, + "size": 224 +} diff --git a/extras/safety_checker/models/safety_checker.py b/extras/safety_checker/models/safety_checker.py new file mode 100644 index 00000000..ea38bf03 --- /dev/null +++ b/extras/safety_checker/models/safety_checker.py @@ -0,0 +1,126 @@ +# from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py + +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import torch +import torch.nn as nn +from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +def cosine_distance(image_embeds, text_embeds): + normalized_image_embeds = nn.functional.normalize(image_embeds) + normalized_text_embeds = nn.functional.normalize(text_embeds) + return torch.mm(normalized_image_embeds, normalized_text_embeds.t()) + + +class StableDiffusionSafetyChecker(PreTrainedModel): + config_class = CLIPConfig + main_input_name = "clip_input" + + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPConfig): + super().__init__(config) + + self.vision_model = CLIPVisionModel(config.vision_config) + self.visual_projection = nn.Linear(config.vision_config.hidden_size, config.projection_dim, bias=False) + + self.concept_embeds = nn.Parameter(torch.ones(17, config.projection_dim), requires_grad=False) + self.special_care_embeds = nn.Parameter(torch.ones(3, config.projection_dim), requires_grad=False) + + self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False) + self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False) + + @torch.no_grad() + def forward(self, clip_input, images): + pooled_output = self.vision_model(clip_input)[1] # pooled_output + image_embeds = self.visual_projection(pooled_output) + + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().float().numpy() + cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy() + + result = [] + batch_size = image_embeds.shape[0] + for i in range(batch_size): + result_img = {"special_scores": {}, "special_care": [], "concept_scores": {}, "bad_concepts": []} + + # increase this value to create a stronger `nfsw` filter + # at the cost of increasing the possibility of filtering benign images + adjustment = 0.0 + + for concept_idx in range(len(special_cos_dist[0])): + concept_cos = special_cos_dist[i][concept_idx] + concept_threshold = self.special_care_embeds_weights[concept_idx].item() + result_img["special_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3) + if result_img["special_scores"][concept_idx] > 0: + result_img["special_care"].append({concept_idx, result_img["special_scores"][concept_idx]}) + adjustment = 0.01 + + for concept_idx in range(len(cos_dist[0])): + concept_cos = cos_dist[i][concept_idx] + concept_threshold = self.concept_embeds_weights[concept_idx].item() + result_img["concept_scores"][concept_idx] = round(concept_cos - concept_threshold + adjustment, 3) + if result_img["concept_scores"][concept_idx] > 0: + result_img["bad_concepts"].append(concept_idx) + + result.append(result_img) + + has_nsfw_concepts = [len(res["bad_concepts"]) > 0 for res in result] + + for idx, has_nsfw_concept in enumerate(has_nsfw_concepts): + if has_nsfw_concept: + if torch.is_tensor(images) or torch.is_tensor(images[0]): + images[idx] = torch.zeros_like(images[idx]) # black image + else: + images[idx] = np.zeros(images[idx].shape) # black image + + if any(has_nsfw_concepts): + logger.warning( + "Potential NSFW content was detected in one or more images. A black image will be returned instead." + " Try again with a different prompt and/or seed." + ) + + return images, has_nsfw_concepts + + @torch.no_grad() + def forward_onnx(self, clip_input: torch.Tensor, images: torch.Tensor): + pooled_output = self.vision_model(clip_input)[1] # pooled_output + image_embeds = self.visual_projection(pooled_output) + + special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds) + cos_dist = cosine_distance(image_embeds, self.concept_embeds) + + # increase this value to create a stronger `nsfw` filter + # at the cost of increasing the possibility of filtering benign images + adjustment = 0.0 + + special_scores = special_cos_dist - self.special_care_embeds_weights + adjustment + # special_scores = special_scores.round(decimals=3) + special_care = torch.any(special_scores > 0, dim=1) + special_adjustment = special_care * 0.01 + special_adjustment = special_adjustment.unsqueeze(1).expand(-1, cos_dist.shape[1]) + + concept_scores = (cos_dist - self.concept_embeds_weights) + special_adjustment + # concept_scores = concept_scores.round(decimals=3) + has_nsfw_concepts = torch.any(concept_scores > 0, dim=1) + + images[has_nsfw_concepts] = 0.0 # black image + + return images, has_nsfw_concepts diff --git a/language/en.json b/language/en.json index e9cd6b73..3eb5d5e2 100644 --- a/language/en.json +++ b/language/en.json @@ -55,6 +55,8 @@ "Disable seed increment": "Disable seed increment", "Disable automatic seed increment when image number is > 1.": "Disable automatic seed increment when image number is > 1.", "Read wildcards in order": "Read wildcards in order", + "Black Out NSFW": "Black Out NSFW", + "Use black image if NSFW is detected.": "Use black image if NSFW is detected.", "\ud83d\udcda History Log": "\uD83D\uDCDA History Log", "Image Style": "Image Style", "Fooocus V2": "Fooocus V2", diff --git a/models/safety_checker/put_safety_checker_models_here b/models/safety_checker/put_safety_checker_models_here new file mode 100644 index 00000000..e69de29b diff --git a/modules/async_worker.py b/modules/async_worker.py index cde99bdc..6f0b30a9 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -43,6 +43,7 @@ def worker(): import fooocus_version import args_manager + from extras.censor import censor_batch, censor_single from modules.sdxl_styles import apply_style, get_random_style, apply_wildcards, fooocus_expansion, apply_arrays, random_style_name from modules.private_logger import log from extras.expansion import safe_str @@ -68,10 +69,14 @@ def worker(): print(f'[Fooocus] {text}') async_task.yields.append(['preview', (number, text, None)]) - def yield_result(async_task, imgs, do_not_show_finished_images=False): + def yield_result(async_task, imgs, black_out_nsfw, censor=True, do_not_show_finished_images=False, progressbar_index=13): if not isinstance(imgs, list): imgs = [imgs] + if censor and (modules.config.default_black_out_nsfw or black_out_nsfw): + progressbar(async_task, progressbar_index, 'Checking for NSFW content ...') + imgs = censor_batch(imgs) + async_task.results = async_task.results + imgs if do_not_show_finished_images: @@ -160,6 +165,7 @@ def worker(): disable_preview = args.pop() disable_intermediate_results = args.pop() disable_seed_increment = args.pop() + black_out_nsfw = args.pop() adm_scaler_positive = args.pop() adm_scaler_negative = args.pop() adm_scaler_end = args.pop() @@ -578,8 +584,11 @@ def worker(): if direct_return: d = [('Upscale (Fast)', 'upscale_fast', '2x')] + if modules.config.default_black_out_nsfw or black_out_nsfw: + progressbar(async_task, 100, 'Checking for NSFW content ...') + uov_input_image = censor_single(uov_input_image) uov_input_image_path = log(uov_input_image, d, output_format=output_format) - yield_result(async_task, uov_input_image_path, do_not_show_finished_images=True) + yield_result(async_task, uov_input_image_path, black_out_nsfw, False, do_not_show_finished_images=True) return tiled = True @@ -643,8 +652,7 @@ def worker(): ) if debugging_inpaint_preprocessor: - yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), - do_not_show_finished_images=True) + yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), black_out_nsfw, do_not_show_finished_images=True) return progressbar(async_task, 13, 'VAE Inpaint encoding ...') @@ -707,7 +715,7 @@ def worker(): cn_img = HWC3(cn_img) task[0] = core.numpy_to_pytorch(cn_img) if debugging_cn_preprocessor: - yield_result(async_task, cn_img, do_not_show_finished_images=True) + yield_result(async_task, cn_img, black_out_nsfw, do_not_show_finished_images=True) return for task in cn_tasks[flags.cn_cpds]: cn_img, cn_stop, cn_weight = task @@ -719,7 +727,7 @@ def worker(): cn_img = HWC3(cn_img) task[0] = core.numpy_to_pytorch(cn_img) if debugging_cn_preprocessor: - yield_result(async_task, cn_img, do_not_show_finished_images=True) + yield_result(async_task, cn_img, black_out_nsfw, do_not_show_finished_images=True) return for task in cn_tasks[flags.cn_ip]: cn_img, cn_stop, cn_weight = task @@ -730,7 +738,7 @@ def worker(): task[0] = ip_adapter.preprocess(cn_img, ip_adapter_path=ip_adapter_path) if debugging_cn_preprocessor: - yield_result(async_task, cn_img, do_not_show_finished_images=True) + yield_result(async_task, cn_img, black_out_nsfw, do_not_show_finished_images=True) return for task in cn_tasks[flags.cn_ip_face]: cn_img, cn_stop, cn_weight = task @@ -744,7 +752,7 @@ def worker(): task[0] = ip_adapter.preprocess(cn_img, ip_adapter_path=ip_adapter_face_path) if debugging_cn_preprocessor: - yield_result(async_task, cn_img, do_not_show_finished_images=True) + yield_result(async_task, cn_img, black_out_nsfw, do_not_show_finished_images=True) return all_ip_tasks = cn_tasks[flags.cn_ip] + cn_tasks[flags.cn_ip_face] @@ -844,6 +852,12 @@ def worker(): imgs = [inpaint_worker.current_task.post_process(x) for x in imgs] img_paths = [] + + if modules.config.default_black_out_nsfw or black_out_nsfw: + progressbar(async_task, int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)), + 'Checking for NSFW content ...') + imgs = censor_batch(imgs) + for x in imgs: d = [('Prompt', 'prompt', task['log_positive_prompt']), ('Negative Prompt', 'negative_prompt', task['log_negative_prompt']), @@ -895,7 +909,7 @@ def worker(): d.append(('Version', 'version', 'Fooocus v' + fooocus_version.version)) img_paths.append(log(x, d, metadata_parser, output_format, task)) - yield_result(async_task, img_paths, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results) + yield_result(async_task, img_paths, black_out_nsfw, False, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results) except ldm_patched.modules.model_management.InterruptProcessingException as e: if async_task.last_stop == 'skip': print('User skipped') diff --git a/modules/config.py b/modules/config.py index f11460c8..ffb74a23 100644 --- a/modules/config.py +++ b/modules/config.py @@ -196,6 +196,7 @@ path_controlnet = get_dir_or_set_default('path_controlnet', '../models/controlne path_clip_vision = get_dir_or_set_default('path_clip_vision', '../models/clip_vision/') path_fooocus_expansion = get_dir_or_set_default('path_fooocus_expansion', '../models/prompt_expansion/fooocus_expansion') path_wildcards = get_dir_or_set_default('path_wildcards', '../wildcards/') +path_safety_checker = get_dir_or_set_default('path_safety_checker', '../models/safety_checker/') path_outputs = get_path_output() @@ -456,6 +457,11 @@ example_inpaint_prompts = get_config_item_or_set_default( ], validator=lambda x: isinstance(x, list) and all(isinstance(v, str) for v in x) ) +default_black_out_nsfw = get_config_item_or_set_default( + key='default_black_out_nsfw', + default_value=False, + validator=lambda x: isinstance(x, bool) +) default_save_metadata_to_images = get_config_item_or_set_default( key='default_save_metadata_to_images', default_value=False, @@ -691,5 +697,13 @@ def downloading_upscale_model(): ) return os.path.join(path_upscale_models, 'fooocus_upscaler_s409985e5.bin') +def downloading_safety_checker_model(): + load_file_from_url( + url='https://huggingface.co/mashb1t/misc/resolve/main/stable-diffusion-safety-checker.bin', + model_dir=path_safety_checker, + file_name='stable-diffusion-safety-checker.bin' + ) + return os.path.join(path_safety_checker, 'stable-diffusion-safety-checker.bin') + update_files() diff --git a/webui.py b/webui.py index f99ab159..55f3102c 100644 --- a/webui.py +++ b/webui.py @@ -436,7 +436,8 @@ with shared.gradio_root: overwrite_upscale_strength = gr.Slider(label='Forced Overwrite of Denoising Strength of "Upscale"', minimum=-1, maximum=1.0, step=0.001, value=-1, info='Set as negative number to disable. For developer debugging.') - disable_preview = gr.Checkbox(label='Disable Preview', value=False, + disable_preview = gr.Checkbox(label='Disable Preview', value=modules.config.default_black_out_nsfw, + interactive=not modules.config.default_black_out_nsfw, info='Disable preview during generation.') disable_intermediate_results = gr.Checkbox(label='Disable Intermediate Results', value=modules.config.default_performance == flags.Performance.EXTREME_SPEED.value, @@ -447,6 +448,15 @@ with shared.gradio_root: value=False) read_wildcards_in_order = gr.Checkbox(label="Read wildcards in order", value=False) + black_out_nsfw = gr.Checkbox(label='Black Out NSFW', + value=modules.config.default_black_out_nsfw, + interactive=not modules.config.default_black_out_nsfw, + info='Use black image if NSFW is detected.') + + black_out_nsfw.change(lambda x: gr.update(value=x, interactive=not x), + inputs=black_out_nsfw, outputs=disable_preview, queue=False, + show_progress=False) + if not args_manager.args.disable_metadata: save_metadata_to_images = gr.Checkbox(label='Save Metadata to Images', value=modules.config.default_save_metadata_to_images, info='Adds parameters to generated images allowing manual regeneration.') @@ -636,7 +646,7 @@ with shared.gradio_root: ctrls += [input_image_checkbox, current_tab] ctrls += [uov_method, uov_input_image] ctrls += [outpaint_selections, inpaint_input_image, inpaint_additional_prompt, inpaint_mask_image] - ctrls += [disable_preview, disable_intermediate_results, disable_seed_increment] + ctrls += [disable_preview, disable_intermediate_results, disable_seed_increment, black_out_nsfw] ctrls += [adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adaptive_cfg] ctrls += [sampler_name, scheduler_name, vae_name] ctrls += [overwrite_step, overwrite_switch, overwrite_width, overwrite_height, overwrite_vary_strength] From 3a55e7e3910b8ae58f82a5a0e4c11d7d4fa3143f Mon Sep 17 00:00:00 2001 From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Date: Sat, 18 May 2024 15:53:34 +0200 Subject: [PATCH 02/12] feat: add AlignYourStepsScheduler (#2905) --- .../contrib/external_align_your_steps.py | 55 +++++++++++++++++++ modules/flags.py | 2 +- modules/sample_hijack.py | 4 ++ 3 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 ldm_patched/contrib/external_align_your_steps.py diff --git a/ldm_patched/contrib/external_align_your_steps.py b/ldm_patched/contrib/external_align_your_steps.py new file mode 100644 index 00000000..624bbce2 --- /dev/null +++ b/ldm_patched/contrib/external_align_your_steps.py @@ -0,0 +1,55 @@ +# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py + +#from: https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html +import numpy as np +import torch + +def loglinear_interp(t_steps, num_steps): + """ + Performs log-linear interpolation of a given array of decreasing numbers. + """ + xs = np.linspace(0, 1, len(t_steps)) + ys = np.log(t_steps[::-1]) + + new_xs = np.linspace(0, 1, num_steps) + new_ys = np.interp(new_xs, xs, ys) + + interped_ys = np.exp(new_ys)[::-1].copy() + return interped_ys + +NOISE_LEVELS = {"SD1": [14.6146412293, 6.4745760956, 3.8636745985, 2.6946151520, 1.8841921177, 1.3943805092, 0.9642583904, 0.6523686016, 0.3977456272, 0.1515232662, 0.0291671582], + "SDXL":[14.6146412293, 6.3184485287, 3.7681790315, 2.1811480769, 1.3405244945, 0.8620721141, 0.5550693289, 0.3798540708, 0.2332364134, 0.1114188177, 0.0291671582], + "SVD": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002]} + +class AlignYourStepsScheduler: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"model_type": (["SD1", "SDXL", "SVD"], ), + "steps": ("INT", {"default": 10, "min": 10, "max": 10000}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "sampling/custom_sampling/schedulers" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, model_type, steps, denoise): + total_steps = steps + if denoise < 1.0: + if denoise <= 0.0: + return (torch.FloatTensor([]),) + total_steps = round(steps * denoise) + + sigmas = NOISE_LEVELS[model_type][:] + if (steps + 1) != len(sigmas): + sigmas = loglinear_interp(sigmas, steps + 1) + + sigmas = sigmas[-(total_steps + 1):] + sigmas[-1] = 0 + return (torch.FloatTensor(sigmas), ) + +NODE_CLASS_MAPPINGS = { + "AlignYourStepsScheduler": AlignYourStepsScheduler, +} \ No newline at end of file diff --git a/modules/flags.py b/modules/flags.py index 9f2aefb3..0c605439 100644 --- a/modules/flags.py +++ b/modules/flags.py @@ -47,7 +47,7 @@ SAMPLERS = KSAMPLER | SAMPLER_EXTRA KSAMPLER_NAMES = list(KSAMPLER.keys()) -SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo"] +SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo", "align_your_steps"] SAMPLER_NAMES = KSAMPLER_NAMES + list(SAMPLER_EXTRA.keys()) sampler_list = SAMPLER_NAMES diff --git a/modules/sample_hijack.py b/modules/sample_hijack.py index 5936a096..4ab3cbbd 100644 --- a/modules/sample_hijack.py +++ b/modules/sample_hijack.py @@ -3,6 +3,7 @@ import ldm_patched.modules.samplers import ldm_patched.modules.model_management from collections import namedtuple +from ldm_patched.contrib.external_align_your_steps import AlignYourStepsScheduler from ldm_patched.contrib.external_custom_sampler import SDTurboScheduler from ldm_patched.k_diffusion import sampling as k_diffusion_sampling from ldm_patched.modules.samplers import normal_scheduler, simple_scheduler, ddim_scheduler @@ -175,6 +176,9 @@ def calculate_sigmas_scheduler_hacked(model, scheduler_name, steps): sigmas = normal_scheduler(model, steps, sgm=True) elif scheduler_name == "turbo": sigmas = SDTurboScheduler().get_sigmas(namedtuple('Patcher', ['model'])(model=model), steps=steps, denoise=1.0)[0] + elif scheduler_name == "align_your_steps": + model_type = 'SDXL' if isinstance(model.latent_format, ldm_patched.modules.latent_formats.SDXL) else 'SD1' + sigmas = AlignYourStepsScheduler().get_sigmas(model_type=model_type, steps=steps, denoise=1.0)[0] else: raise TypeError("error invalid scheduler") return sigmas From 3bae73e23ecc85430e532ce57a25b04437c7cf67 Mon Sep 17 00:00:00 2001 From: cantor-set <32692347+cantor-set@users.noreply.github.com> Date: Sat, 18 May 2024 11:19:46 -0400 Subject: [PATCH 03/12] feat: add support for lora inline prompt references (#2323) * Adding support to inline prompt references * Added unittests * Added an initial documentation for development guidelines * Added a negative number * renamed parameter * removed wrongly committed file * Code fixes * Fixed circular reference * Fixed typo. Added TODO * Fixed merge * Code cleanup * Added missing refernce function * Removed function from util.py... again... * Update modules/async_worker.py Implemented suggested change Co-authored-by: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> * Removed another circular reference * Renamed module * Addressed PR comments * Added return type to function * refactor: move apply_wildcards to module util * refactor: code cleanup, unify usage of tuples in lora list * docs: add instructions for running unittests on embedded python, code cleanup * refactor: code cleanup, move makedirs_with_log back to util --------- Co-authored-by: cantor-set Co-authored-by: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Co-authored-by: Manuel Schmid --- development.md | 11 ++++++ modules/__init__.py | 0 modules/async_worker.py | 46 ++++++++++++++++-------- modules/config.py | 5 +-- modules/extra_utils.py | 20 +++++++++++ modules/sdxl_styles.py | 36 ++----------------- modules/util.py | 79 +++++++++++++++++++++++++++++------------ tests/__init__.py | 4 +++ tests/test_utils.py | 48 +++++++++++++++++++++++++ 9 files changed, 176 insertions(+), 73 deletions(-) create mode 100644 development.md create mode 100644 modules/__init__.py create mode 100644 modules/extra_utils.py create mode 100644 tests/__init__.py create mode 100644 tests/test_utils.py diff --git a/development.md b/development.md new file mode 100644 index 00000000..bbb3def9 --- /dev/null +++ b/development.md @@ -0,0 +1,11 @@ +## Running unit tests + +Native python: +``` +python -m unittest tests/ +``` + +Embedded python (Windows zip file installation method): +``` +..\python_embeded\python.exe -m unittest +``` diff --git a/modules/__init__.py b/modules/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modules/async_worker.py b/modules/async_worker.py index 6f0b30a9..7f0a46e3 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -4,6 +4,7 @@ from modules.patch import PatchSettings, patch_settings, patch_all patch_all() + class AsyncTask: def __init__(self, args): self.args = args @@ -44,11 +45,12 @@ def worker(): import args_manager from extras.censor import censor_batch, censor_single - from modules.sdxl_styles import apply_style, get_random_style, apply_wildcards, fooocus_expansion, apply_arrays, random_style_name + from modules.sdxl_styles import apply_style, get_random_style, fooocus_expansion, apply_arrays, random_style_name from modules.private_logger import log from extras.expansion import safe_str - from modules.util import remove_empty_str, HWC3, resize_image, get_image_shape_ceil, set_image_shape_ceil, \ - get_shape_ceil, resample_image, erode_or_dilate, ordinal_suffix, get_enabled_loras + from modules.util import (remove_empty_str, HWC3, resize_image, get_image_shape_ceil, set_image_shape_ceil, + get_shape_ceil, resample_image, erode_or_dilate, ordinal_suffix, get_enabled_loras, + parse_lora_references_from_prompt, apply_wildcards) from modules.upscaler import perform_upscale from modules.flags import Performance from modules.meta_parser import get_metadata_parser, MetadataScheme @@ -69,7 +71,8 @@ def worker(): print(f'[Fooocus] {text}') async_task.yields.append(['preview', (number, text, None)]) - def yield_result(async_task, imgs, black_out_nsfw, censor=True, do_not_show_finished_images=False, progressbar_index=13): + def yield_result(async_task, imgs, black_out_nsfw, censor=True, do_not_show_finished_images=False, + progressbar_index=13): if not isinstance(imgs, list): imgs = [imgs] @@ -152,7 +155,8 @@ def worker(): base_model_name = args.pop() refiner_model_name = args.pop() refiner_switch = args.pop() - loras = get_enabled_loras([[bool(args.pop()), str(args.pop()), float(args.pop())] for _ in range(modules.config.default_max_lora_number)]) + loras = get_enabled_loras([(bool(args.pop()), str(args.pop()), float(args.pop())) for _ in + range(modules.config.default_max_lora_number)]) input_image_checkbox = args.pop() current_tab = args.pop() uov_method = args.pop() @@ -202,7 +206,8 @@ def worker(): inpaint_erode_or_dilate = args.pop() save_metadata_to_images = args.pop() if not args_manager.args.disable_metadata else False - metadata_scheme = MetadataScheme(args.pop()) if not args_manager.args.disable_metadata else MetadataScheme.FOOOCUS + metadata_scheme = MetadataScheme( + args.pop()) if not args_manager.args.disable_metadata else MetadataScheme.FOOOCUS cn_tasks = {x: [] for x in flags.ip_list} for _ in range(flags.controlnet_image_count): @@ -433,13 +438,16 @@ def worker(): extra_negative_prompts = negative_prompts[1:] if len(negative_prompts) > 1 else [] progressbar(async_task, 3, 'Loading models ...') + + loras = parse_lora_references_from_prompt(prompt, loras, modules.config.default_max_lora_number) + pipeline.refresh_everything(refiner_model_name=refiner_model_name, base_model_name=base_model_name, loras=loras, base_model_additional_loras=base_model_additional_loras, use_synthetic_refiner=use_synthetic_refiner, vae_name=vae_name) progressbar(async_task, 3, 'Processing prompts ...') tasks = [] - + for i in range(image_number): if disable_seed_increment: task_seed = seed % (constants.MAX_SEED + 1) @@ -450,8 +458,10 @@ def worker(): task_prompt = apply_wildcards(prompt, task_rng, i, read_wildcards_in_order) task_prompt = apply_arrays(task_prompt, i) task_negative_prompt = apply_wildcards(negative_prompt, task_rng, i, read_wildcards_in_order) - task_extra_positive_prompts = [apply_wildcards(pmt, task_rng, i, read_wildcards_in_order) for pmt in extra_positive_prompts] - task_extra_negative_prompts = [apply_wildcards(pmt, task_rng, i, read_wildcards_in_order) for pmt in extra_negative_prompts] + task_extra_positive_prompts = [apply_wildcards(pmt, task_rng, i, read_wildcards_in_order) for pmt in + extra_positive_prompts] + task_extra_negative_prompts = [apply_wildcards(pmt, task_rng, i, read_wildcards_in_order) for pmt in + extra_negative_prompts] positive_basic_workloads = [] negative_basic_workloads = [] @@ -652,7 +662,8 @@ def worker(): ) if debugging_inpaint_preprocessor: - yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), black_out_nsfw, do_not_show_finished_images=True) + yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(), black_out_nsfw, + do_not_show_finished_images=True) return progressbar(async_task, 13, 'VAE Inpaint encoding ...') @@ -807,7 +818,8 @@ def worker(): done_steps = current_task_id * steps + step async_task.yields.append(['preview', ( int(15.0 + 85.0 * float(done_steps) / float(all_steps)), - f'Step {step}/{total_steps} in the {current_task_id + 1}{ordinal_suffix(current_task_id + 1)} Sampling', y)]) + f'Step {step}/{total_steps} in the {current_task_id + 1}{ordinal_suffix(current_task_id + 1)} Sampling', + y)]) for current_task_id, task in enumerate(tasks): execution_start_time = time.perf_counter() @@ -862,7 +874,8 @@ def worker(): d = [('Prompt', 'prompt', task['log_positive_prompt']), ('Negative Prompt', 'negative_prompt', task['log_negative_prompt']), ('Fooocus V2 Expansion', 'prompt_expansion', task['expansion']), - ('Styles', 'styles', str(task['styles'] if not use_expansion else [fooocus_expansion] + task['styles'])), + ('Styles', 'styles', + str(task['styles'] if not use_expansion else [fooocus_expansion] + task['styles'])), ('Performance', 'performance', performance_selection.value)] if performance_selection.steps() != steps: @@ -885,7 +898,8 @@ def worker(): if refiner_swap_method != flags.refiner_swap_method: d.append(('Refiner Swap Method', 'refiner_swap_method', refiner_swap_method)) if modules.patch.patch_settings[pid].adaptive_cfg != modules.config.default_cfg_tsnr: - d.append(('CFG Mimicking from TSNR', 'adaptive_cfg', modules.patch.patch_settings[pid].adaptive_cfg)) + d.append( + ('CFG Mimicking from TSNR', 'adaptive_cfg', modules.patch.patch_settings[pid].adaptive_cfg)) d.append(('Sampler', 'sampler', sampler_name)) d.append(('Scheduler', 'scheduler', scheduler_name)) @@ -905,11 +919,13 @@ def worker(): metadata_parser.set_data(task['log_positive_prompt'], task['positive'], task['log_negative_prompt'], task['negative'], steps, base_model_name, refiner_model_name, loras, vae_name) - d.append(('Metadata Scheme', 'metadata_scheme', metadata_scheme.value if save_metadata_to_images else save_metadata_to_images)) + d.append(('Metadata Scheme', 'metadata_scheme', + metadata_scheme.value if save_metadata_to_images else save_metadata_to_images)) d.append(('Version', 'version', 'Fooocus v' + fooocus_version.version)) img_paths.append(log(x, d, metadata_parser, output_format, task)) - yield_result(async_task, img_paths, black_out_nsfw, False, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results) + yield_result(async_task, img_paths, black_out_nsfw, False, + do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results) except ldm_patched.modules.model_management.InterruptProcessingException as e: if async_task.last_stop == 'skip': print('User skipped') diff --git a/modules/config.py b/modules/config.py index ffb74a23..11fe3181 100644 --- a/modules/config.py +++ b/modules/config.py @@ -8,7 +8,8 @@ import modules.flags import modules.sdxl_styles from modules.model_loader import load_file_from_url -from modules.util import get_files_from_folder, makedirs_with_log +from modules.util import makedirs_with_log +from modules.extra_utils import get_files_from_folder from modules.flags import OutputFormat, Performance, MetadataScheme @@ -20,7 +21,7 @@ def get_config_path(key, default_value): else: return os.path.abspath(default_value) - +wildcards_max_bfs_depth = 64 config_path = get_config_path('config_path', "./config.txt") config_example_path = get_config_path('config_example_path', "config_modification_tutorial.txt") config_dict = {} diff --git a/modules/extra_utils.py b/modules/extra_utils.py new file mode 100644 index 00000000..3e95e8b5 --- /dev/null +++ b/modules/extra_utils.py @@ -0,0 +1,20 @@ +import os + + +def get_files_from_folder(folder_path, extensions=None, name_filter=None): + if not os.path.isdir(folder_path): + raise ValueError("Folder path is not a valid directory.") + + filenames = [] + + for root, _, files in os.walk(folder_path, topdown=False): + relative_path = os.path.relpath(root, folder_path) + if relative_path == ".": + relative_path = "" + for filename in sorted(files, key=lambda s: s.casefold()): + _, file_extension = os.path.splitext(filename) + if (extensions is None or file_extension.lower() in extensions) and (name_filter is None or name_filter in _): + path = os.path.join(relative_path, filename) + filenames.append(path) + + return filenames diff --git a/modules/sdxl_styles.py b/modules/sdxl_styles.py index 5b6afb59..12ab6c5c 100644 --- a/modules/sdxl_styles.py +++ b/modules/sdxl_styles.py @@ -2,14 +2,12 @@ import os import re import json import math -import modules.config -from modules.util import get_files_from_folder +from modules.extra_utils import get_files_from_folder from random import Random # cannot use modules.config - validators causing circular imports styles_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../sdxl_styles/')) -wildcards_max_bfs_depth = 64 def normalize_key(k): @@ -25,7 +23,6 @@ def normalize_key(k): styles = {} - styles_files = get_files_from_folder(styles_path, ['.json']) for x in ['sdxl_styles_fooocus.json', @@ -65,34 +62,7 @@ def apply_style(style, positive): return p.replace('{prompt}', positive).splitlines(), n.splitlines() -def apply_wildcards(wildcard_text, rng, i, read_wildcards_in_order): - for _ in range(wildcards_max_bfs_depth): - placeholders = re.findall(r'__([\w-]+)__', wildcard_text) - if len(placeholders) == 0: - return wildcard_text - - print(f'[Wildcards] processing: {wildcard_text}') - for placeholder in placeholders: - try: - matches = [x for x in modules.config.wildcard_filenames if os.path.splitext(os.path.basename(x))[0] == placeholder] - words = open(os.path.join(modules.config.path_wildcards, matches[0]), encoding='utf-8').read().splitlines() - words = [x for x in words if x != ''] - assert len(words) > 0 - if read_wildcards_in_order: - wildcard_text = wildcard_text.replace(f'__{placeholder}__', words[i % len(words)], 1) - else: - wildcard_text = wildcard_text.replace(f'__{placeholder}__', rng.choice(words), 1) - except: - print(f'[Wildcards] Warning: {placeholder}.txt missing or empty. ' - f'Using "{placeholder}" as a normal word.') - wildcard_text = wildcard_text.replace(f'__{placeholder}__', placeholder) - print(f'[Wildcards] {wildcard_text}') - - print(f'[Wildcards] BFS stack overflow. Current text: {wildcard_text}') - return wildcard_text - - -def get_words(arrays, totalMult, index): +def get_words(arrays, total_mult, index): if len(arrays) == 1: return [arrays[0].split(',')[index]] else: @@ -101,7 +71,7 @@ def get_words(arrays, totalMult, index): index -= index % len(words) index /= len(words) index = math.floor(index) - return [word] + get_words(arrays[1:], math.floor(totalMult/len(words)), index) + return [word] + get_words(arrays[1:], math.floor(total_mult / len(words)), index) def apply_arrays(text, index): diff --git a/modules/util.py b/modules/util.py index d2feecb6..73430230 100644 --- a/modules/util.py +++ b/modules/util.py @@ -1,11 +1,12 @@ -import typing - import numpy as np import datetime import random import math import os import cv2 +import re +from typing import List, Tuple, AnyStr, NamedTuple + import json import hashlib @@ -14,8 +15,16 @@ from PIL import Image import modules.sdxl_styles LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) + + +# Regexp compiled once. Matches entries with the following pattern: +# +# +LORAS_PROMPT_PATTERN = re.compile(r".* .*", re.X) + HASH_SHA256_LENGTH = 10 + def erode_or_dilate(x, k): k = int(k) if k > 0: @@ -163,25 +172,6 @@ def generate_temp_filename(folder='./outputs/', extension='png'): return date_string, os.path.abspath(result), filename -def get_files_from_folder(folder_path, extensions=None, name_filter=None): - if not os.path.isdir(folder_path): - raise ValueError("Folder path is not a valid directory.") - - filenames = [] - - for root, dirs, files in os.walk(folder_path, topdown=False): - relative_path = os.path.relpath(root, folder_path) - if relative_path == ".": - relative_path = "" - for filename in sorted(files, key=lambda s: s.casefold()): - _, file_extension = os.path.splitext(filename) - if (extensions is None or file_extension.lower() in extensions) and (name_filter is None or name_filter in _): - path = os.path.join(relative_path, filename) - filenames.append(path) - - return filenames - - def sha256(filename, use_addnet_hash=False, length=HASH_SHA256_LENGTH): print(f"Calculating sha256 for {filename}: ", end='') if use_addnet_hash: @@ -355,7 +345,7 @@ def extract_styles_from_prompt(prompt, negative_prompt): return list(reversed(extracted)), real_prompt, negative_prompt -class PromptStyle(typing.NamedTuple): +class PromptStyle(NamedTuple): name: str prompt: str negative_prompt: str @@ -394,4 +384,47 @@ def makedirs_with_log(path): def get_enabled_loras(loras: list) -> list: - return [[lora[1], lora[2]] for lora in loras if lora[0]] + return [(lora[1], lora[2]) for lora in loras if lora[0]] + + +def parse_lora_references_from_prompt(prompt: str, loras: List[Tuple[AnyStr, float]], loras_limit: int = 5) -> List[Tuple[AnyStr, float]]: + new_loras = [] + updated_loras = [] + for token in prompt.split(","): + m = LORAS_PROMPT_PATTERN.match(token) + + if m: + new_loras.append((f"{m.group(1)}.safetensors", float(m.group(2)))) + + for lora in loras + new_loras: + if lora[0] != "None": + updated_loras.append(lora) + + return updated_loras[:loras_limit] + + +def apply_wildcards(wildcard_text, rng, i, read_wildcards_in_order) -> str: + for _ in range(modules.config.wildcards_max_bfs_depth): + placeholders = re.findall(r'__([\w-]+)__', wildcard_text) + if len(placeholders) == 0: + return wildcard_text + + print(f'[Wildcards] processing: {wildcard_text}') + for placeholder in placeholders: + try: + matches = [x for x in modules.config.wildcard_filenames if os.path.splitext(os.path.basename(x))[0] == placeholder] + words = open(os.path.join(modules.config.path_wildcards, matches[0]), encoding='utf-8').read().splitlines() + words = [x for x in words if x != ''] + assert len(words) > 0 + if read_wildcards_in_order: + wildcard_text = wildcard_text.replace(f'__{placeholder}__', words[i % len(words)], 1) + else: + wildcard_text = wildcard_text.replace(f'__{placeholder}__', rng.choice(words), 1) + except: + print(f'[Wildcards] Warning: {placeholder}.txt missing or empty. ' + f'Using "{placeholder}" as a normal word.') + wildcard_text = wildcard_text.replace(f'__{placeholder}__', placeholder) + print(f'[Wildcards] {wildcard_text}') + + print(f'[Wildcards] BFS stack overflow. Current text: {wildcard_text}') + return wildcard_text diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..c424468f --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,4 @@ +import sys +import pathlib + +sys.path.append(pathlib.Path(f'{__file__}/../modules').parent.resolve()) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..0698dcc8 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,48 @@ +import unittest + +from modules import util + + +class TestUtils(unittest.TestCase): + def test_can_parse_tokens_with_lora(self): + test_cases = [ + { + "input": ("some prompt, very cool, , cool ", [], 5), + "output": [("hey-lora.safetensors", 0.4), ("you-lora.safetensors", 0.2)], + }, + # Test can not exceed limit + { + "input": ("some prompt, very cool, , cool ", [], 1), + "output": [("hey-lora.safetensors", 0.4)], + }, + # test Loras from UI take precedence over prompt + { + "input": ( + "some prompt, very cool, , , , , , ", + [("hey-lora.safetensors", 0.4)], + 5, + ), + "output": [ + ("hey-lora.safetensors", 0.4), + ("l1.safetensors", 0.4), + ("l2.safetensors", -0.2), + ("l3.safetensors", 0.3), + ("l4.safetensors", 0.5), + ], + }, + # Test lora specification not separated by comma are ignored, only latest specified is used + { + "input": ("some prompt, very cool, ", [], 3), + "output": [("you-lora.safetensors", 0.2)], + }, + { + "input": (", , and ", [], 6), + "output": [] + } + ] + + for test in test_cases: + prompt, loras, loras_limit = test["input"] + expected = test["output"] + actual = util.parse_lora_references_from_prompt(prompt, loras, loras_limit) + self.assertEqual(expected, actual) From 2e2e8f851a501e1ae7870112b5e3144e241c015a Mon Sep 17 00:00:00 2001 From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Date: Sun, 19 May 2024 13:08:33 +0200 Subject: [PATCH 04/12] feat: add tcd sampler and discrete distilled tcd scheduler based on sgm_uniform (same as lcm) (#2907) --- .../contrib/external_custom_sampler.py | 20 +++++++++++++ .../contrib/external_model_advanced.py | 5 +++- ldm_patched/k_diffusion/sampling.py | 28 ++++++++++++++++++- ldm_patched/modules/model_sampling.py | 8 +++--- ldm_patched/modules/samplers.py | 2 +- modules/async_worker.py | 8 +++--- modules/flags.py | 5 ++-- modules/patch_precision.py | 2 ++ 8 files changed, 65 insertions(+), 13 deletions(-) diff --git a/ldm_patched/contrib/external_custom_sampler.py b/ldm_patched/contrib/external_custom_sampler.py index 8f92e841..985b03a0 100644 --- a/ldm_patched/contrib/external_custom_sampler.py +++ b/ldm_patched/contrib/external_custom_sampler.py @@ -230,6 +230,25 @@ class SamplerDPMPP_SDE: sampler = ldm_patched.modules.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise, "r": r}) return (sampler, ) + +class SamplerTCD: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "eta": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + RETURN_TYPES = ("SAMPLER",) + CATEGORY = "sampling/custom_sampling/samplers" + + FUNCTION = "get_sampler" + + def get_sampler(self, eta=0.3): + sampler = ldm_patched.modules.samplers.ksampler("tcd", {"eta": eta}) + return (sampler, ) + + class SamplerCustom: @classmethod def INPUT_TYPES(s): @@ -292,6 +311,7 @@ NODE_CLASS_MAPPINGS = { "KSamplerSelect": KSamplerSelect, "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE, "SamplerDPMPP_SDE": SamplerDPMPP_SDE, + "SamplerTCD": SamplerTCD, "SplitSigmas": SplitSigmas, "FlipSigmas": FlipSigmas, } diff --git a/ldm_patched/contrib/external_model_advanced.py b/ldm_patched/contrib/external_model_advanced.py index 03a2f045..9b52c36b 100644 --- a/ldm_patched/contrib/external_model_advanced.py +++ b/ldm_patched/contrib/external_model_advanced.py @@ -70,7 +70,7 @@ class ModelSamplingDiscrete: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), - "sampling": (["eps", "v_prediction", "lcm"],), + "sampling": (["eps", "v_prediction", "lcm", "tcd"]), "zsnr": ("BOOLEAN", {"default": False}), }} @@ -90,6 +90,9 @@ class ModelSamplingDiscrete: elif sampling == "lcm": sampling_type = LCM sampling_base = ModelSamplingDiscreteDistilled + elif sampling == "tcd": + sampling_type = ldm_patched.modules.model_sampling.EPS + sampling_base = ModelSamplingDiscreteDistilled class ModelSamplingAdvanced(sampling_base, sampling_type): pass diff --git a/ldm_patched/k_diffusion/sampling.py b/ldm_patched/k_diffusion/sampling.py index 761c2e0e..d1bc1e4b 100644 --- a/ldm_patched/k_diffusion/sampling.py +++ b/ldm_patched/k_diffusion/sampling.py @@ -752,7 +752,6 @@ def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, n return x - @torch.no_grad() def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.): # From MIT licensed: https://github.com/Carzit/sd-webui-samplers-scheduler/ @@ -808,3 +807,30 @@ def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=Non d_prime = w1 * d + w2 * d_2 + w3 * d_3 x = x + d_prime * dt return x + + +@torch.no_grad() +def sample_tcd(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, eta=0.3): + extra_args = {} if extra_args is None else extra_args + noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler + s_in = x.new_ones([x.shape[0]]) + + model_sampling = model.inner_model.inner_model.model_sampling + timesteps_s = torch.floor((1 - eta) * model_sampling.timestep(sigmas)).to(dtype=torch.long).detach().cpu() + timesteps_s[-1] = 0 + alpha_prod_s = model_sampling.alphas_cumprod[timesteps_s] + beta_prod_s = 1 - alpha_prod_s + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) # predicted_original_sample + eps = (x - denoised) / sigmas[i] + denoised = alpha_prod_s[i + 1].sqrt() * denoised + beta_prod_s[i + 1].sqrt() * eps + + if callback is not None: + callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised}) + + x = denoised + if eta > 0 and sigmas[i + 1] > 0: + noise = noise_sampler(sigmas[i], sigmas[i + 1]) + x = x / alpha_prod_s[i+1].sqrt() + noise * (sigmas[i+1]**2 + 1 - 1/alpha_prod_s[i+1]).sqrt() + + return x \ No newline at end of file diff --git a/ldm_patched/modules/model_sampling.py b/ldm_patched/modules/model_sampling.py index f39e275d..57f51a00 100644 --- a/ldm_patched/modules/model_sampling.py +++ b/ldm_patched/modules/model_sampling.py @@ -50,17 +50,17 @@ class ModelSamplingDiscrete(torch.nn.Module): self.linear_start = linear_start self.linear_end = linear_end - # self.register_buffer('betas', torch.tensor(betas, dtype=torch.float32)) - # self.register_buffer('alphas_cumprod', torch.tensor(alphas_cumprod, dtype=torch.float32)) - # self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32)) - sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5 self.set_sigmas(sigmas) + self.set_alphas_cumprod(alphas_cumprod.float()) def set_sigmas(self, sigmas): self.register_buffer('sigmas', sigmas) self.register_buffer('log_sigmas', sigmas.log()) + def set_alphas_cumprod(self, alphas_cumprod): + self.register_buffer("alphas_cumprod", alphas_cumprod.float()) + @property def sigma_min(self): return self.sigmas[0] diff --git a/ldm_patched/modules/samplers.py b/ldm_patched/modules/samplers.py index 1f69d2b1..35cb3d73 100644 --- a/ldm_patched/modules/samplers.py +++ b/ldm_patched/modules/samplers.py @@ -523,7 +523,7 @@ class UNIPCBH2(Sampler): KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", - "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"] + "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "tcd"] class KSAMPLER(Sampler): def __init__(self, sampler_function, extra_options={}, inpaint_options={}): diff --git a/modules/async_worker.py b/modules/async_worker.py index 7f0a46e3..1dabf89c 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -798,19 +798,19 @@ def worker(): final_sampler_name = sampler_name final_scheduler_name = scheduler_name - if scheduler_name == 'lcm': + if scheduler_name in ['lcm', 'tcd']: final_scheduler_name = 'sgm_uniform' if pipeline.final_unet is not None: pipeline.final_unet = core.opModelSamplingDiscrete.patch( pipeline.final_unet, - sampling='lcm', + sampling=scheduler_name, zsnr=False)[0] if pipeline.final_refiner_unet is not None: pipeline.final_refiner_unet = core.opModelSamplingDiscrete.patch( pipeline.final_refiner_unet, - sampling='lcm', + sampling=scheduler_name, zsnr=False)[0] - print('Using lcm scheduler.') + print(f'Using {scheduler_name} scheduler.') async_task.yields.append(['preview', (13, 'Moving model to GPU ...', None)]) diff --git a/modules/flags.py b/modules/flags.py index 0c605439..cb4c3ec9 100644 --- a/modules/flags.py +++ b/modules/flags.py @@ -34,7 +34,8 @@ KSAMPLER = { "dpmpp_3m_sde": "", "dpmpp_3m_sde_gpu": "", "ddpm": "", - "lcm": "LCM" + "lcm": "LCM", + "tcd": "TCD" } SAMPLER_EXTRA = { @@ -47,7 +48,7 @@ SAMPLERS = KSAMPLER | SAMPLER_EXTRA KSAMPLER_NAMES = list(KSAMPLER.keys()) -SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo", "align_your_steps"] +SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo", "align_your_steps", "tcd"] SAMPLER_NAMES = KSAMPLER_NAMES + list(SAMPLER_EXTRA.keys()) sampler_list = SAMPLER_NAMES diff --git a/modules/patch_precision.py b/modules/patch_precision.py index 83569bdd..22ffda0a 100644 --- a/modules/patch_precision.py +++ b/modules/patch_precision.py @@ -51,6 +51,8 @@ def patched_register_schedule(self, given_betas=None, beta_schedule="linear", ti self.linear_end = linear_end sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32) self.set_sigmas(sigmas) + alphas_cumprod = torch.tensor(alphas_cumprod, dtype=torch.float32) + self.set_alphas_cumprod(alphas_cumprod) return From 13599edb9b5066649c3ac31bb5a7b15403fd6297 Mon Sep 17 00:00:00 2001 From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Date: Sun, 19 May 2024 13:23:08 +0200 Subject: [PATCH 05/12] feat: add performance hyper-sd based on 4step LoRA (#2812) * feat: add performance hyper-sd based on 4step LoRA * feat: use LoRA weight 0.8, sampler dpmpp_sde_gpu and scheduler_name karras suggested in https://github.com/lllyasviel/Fooocus/discussions/2813#discussioncomment-9245251 results see https://github.com/lllyasviel/Fooocus/discussions/2813#discussioncomment-9275251 * feat: change ByteDance huggingface profile with mashb1t * wip: add hyper-sd 8 step cfg lora with negative prompt support * feat: remove hyper-sd8 performance still waiting for the release of hyper-sd 4step CFG LoRA, not yet satisfied with any of the CFG LoRAs compared to non-cfg ones. see https://huggingface.co/ByteDance/Hyper-SD --- modules/async_worker.py | 27 +++++++++++++++++++++++++++ modules/config.py | 14 ++++++++++++-- modules/flags.py | 5 ++++- 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/modules/async_worker.py b/modules/async_worker.py index 1dabf89c..cf1eda30 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -275,6 +275,33 @@ def worker(): adm_scaler_negative = 1.0 adm_scaler_end = 0.0 + elif performance_selection == Performance.HYPER_SD: + print('Enter Hyper-SD mode.') + progressbar(async_task, 1, 'Downloading Hyper-SD components ...') + loras += [(modules.config.downloading_sdxl_hyper_sd_lora(), 0.8)] + + if refiner_model_name != 'None': + print(f'Refiner disabled in Hyper-SD mode.') + + refiner_model_name = 'None' + sampler_name = 'dpmpp_sde_gpu' + scheduler_name = 'karras' + sharpness = 0.0 + guidance_scale = 1.0 + adaptive_cfg = 1.0 + refiner_switch = 1.0 + adm_scaler_positive = 1.0 + adm_scaler_negative = 1.0 + adm_scaler_end = 0.0 + + elif performance_selection == Performance.HYPER_SD8: + print('Enter Hyper-SD8 mode.') + progressbar(async_task, 1, 'Downloading Hyper-SD components ...') + loras += [(modules.config.downloading_sdxl_hyper_sd_cfg_lora(), 0.3)] + + sampler_name = 'dpmpp_sde_gpu' + scheduler_name = 'normal' + print(f'[Parameters] Adaptive CFG = {adaptive_cfg}') print(f'[Parameters] Sharpness = {sharpness}') print(f'[Parameters] ControlNet Softness = {controlnet_softness}') diff --git a/modules/config.py b/modules/config.py index 11fe3181..db7036c5 100644 --- a/modules/config.py +++ b/modules/config.py @@ -553,7 +553,8 @@ wildcard_filenames = [] sdxl_lcm_lora = 'sdxl_lcm_lora.safetensors' sdxl_lightning_lora = 'sdxl_lightning_4step_lora.safetensors' -loras_metadata_remove = [sdxl_lcm_lora, sdxl_lightning_lora] +sdxl_hyper_sd_lora = 'sdxl_hyper_sd_4step_lora.safetensors' +loras_metadata_remove = [sdxl_lcm_lora, sdxl_lightning_lora, sdxl_hyper_sd_lora] def get_model_filenames(folder_paths, extensions=None, name_filter=None): @@ -627,13 +628,22 @@ def downloading_sdxl_lcm_lora(): def downloading_sdxl_lightning_lora(): load_file_from_url( - url='https://huggingface.co/ByteDance/SDXL-Lightning/resolve/main/sdxl_lightning_4step_lora.safetensors', + url='https://huggingface.co/mashb1t/misc/resolve/main/sdxl_lightning_4step_lora.safetensors', model_dir=paths_loras[0], file_name=sdxl_lightning_lora ) return sdxl_lightning_lora +def downloading_sdxl_hyper_sd_lora(): + load_file_from_url( + url='https://huggingface.co/mashb1t/misc/resolve/main/sdxl_hyper_sd_4step_lora.safetensors', + model_dir=paths_loras[0], + file_name=sdxl_hyper_sd_lora + ) + return sdxl_hyper_sd_lora + + def downloading_controlnet_canny(): load_file_from_url( url='https://huggingface.co/lllyasviel/misc/resolve/main/control-lora-canny-rank128.safetensors', diff --git a/modules/flags.py b/modules/flags.py index cb4c3ec9..77ad012a 100644 --- a/modules/flags.py +++ b/modules/flags.py @@ -110,6 +110,7 @@ class Steps(IntEnum): SPEED = 30 EXTREME_SPEED = 8 LIGHTNING = 4 + HYPER_SD = 4 class StepsUOV(IntEnum): @@ -117,6 +118,7 @@ class StepsUOV(IntEnum): SPEED = 18 EXTREME_SPEED = 8 LIGHTNING = 4 + HYPER_SD = 4 class Performance(Enum): @@ -124,6 +126,7 @@ class Performance(Enum): SPEED = 'Speed' EXTREME_SPEED = 'Extreme Speed' LIGHTNING = 'Lightning' + HYPER_SD = 'Hyper-SD' @classmethod def list(cls) -> list: @@ -133,7 +136,7 @@ class Performance(Enum): def has_restricted_features(cls, x) -> bool: if isinstance(x, Performance): x = x.value - return x in [cls.EXTREME_SPEED.value, cls.LIGHTNING.value] + return x in [cls.EXTREME_SPEED.value, cls.LIGHTNING.value, cls.HYPER_SD.value] def steps(self) -> int | None: return Steps[self.name].value if Steps[self.name] else None From 0466ff944c96632e100a4cc988f7d02288ce7d3b Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sun, 19 May 2024 14:29:10 +0200 Subject: [PATCH 06/12] release: bump version number to 2.4.0-rc1 --- fooocus_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fooocus_version.py b/fooocus_version.py index b2050196..def34a20 100644 --- a/fooocus_version.py +++ b/fooocus_version.py @@ -1 +1 @@ -version = '2.3.1' +version = '2.4.0-rc1' From dad228907e5ae290f441510aafb3b07bd75a98ea Mon Sep 17 00:00:00 2001 From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Date: Sun, 19 May 2024 17:42:46 +0200 Subject: [PATCH 07/12] fix: remove leftover code from hyper-sd8 testing (#2959) --- modules/async_worker.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/modules/async_worker.py b/modules/async_worker.py index cf1eda30..892f99a7 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -294,14 +294,6 @@ def worker(): adm_scaler_negative = 1.0 adm_scaler_end = 0.0 - elif performance_selection == Performance.HYPER_SD8: - print('Enter Hyper-SD8 mode.') - progressbar(async_task, 1, 'Downloading Hyper-SD components ...') - loras += [(modules.config.downloading_sdxl_hyper_sd_cfg_lora(), 0.3)] - - sampler_name = 'dpmpp_sde_gpu' - scheduler_name = 'normal' - print(f'[Parameters] Adaptive CFG = {adaptive_cfg}') print(f'[Parameters] Sharpness = {sharpness}') print(f'[Parameters] ControlNet Softness = {controlnet_softness}') From 35b74dfa64b2efe8dd7652577eec443f45c56939 Mon Sep 17 00:00:00 2001 From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com> Date: Sun, 19 May 2024 18:02:24 +0200 Subject: [PATCH 08/12] feat: optimize model management of image censoring (#2960) now follows general Fooocus model management principles + includes code optimisations for reusability --- extras/censor.py | 96 +++++++++++++++++++++-------------------- modules/async_worker.py | 15 ++++--- 2 files changed, 58 insertions(+), 53 deletions(-) diff --git a/extras/censor.py b/extras/censor.py index 2047db24..45617fd8 100644 --- a/extras/censor.py +++ b/extras/censor.py @@ -1,56 +1,60 @@ -# modified version of https://github.com/AUTOMATIC1111/stable-diffusion-webui-nsfw-censor/blob/master/scripts/censor.py -import numpy as np import os -from extras.safety_checker.models.safety_checker import StableDiffusionSafetyChecker -from transformers import CLIPFeatureExtractor, CLIPConfig -from PIL import Image +import numpy as np +import torch +from transformers import CLIPConfig, CLIPImageProcessor + +import ldm_patched.modules.model_management as model_management import modules.config +from extras.safety_checker.models.safety_checker import StableDiffusionSafetyChecker +from ldm_patched.modules.model_patcher import ModelPatcher safety_checker_repo_root = os.path.join(os.path.dirname(__file__), 'safety_checker') config_path = os.path.join(safety_checker_repo_root, "configs", "config.json") preprocessor_config_path = os.path.join(safety_checker_repo_root, "configs", "preprocessor_config.json") -safety_feature_extractor = None -safety_checker = None + +class Censor: + def __init__(self): + self.safety_checker_model: ModelPatcher | None = None + self.clip_image_processor: CLIPImageProcessor | None = None + self.load_device = torch.device('cpu') + self.offload_device = torch.device('cpu') + + def init(self): + if self.safety_checker_model is None and self.clip_image_processor is None: + safety_checker_model = modules.config.downloading_safety_checker_model() + self.clip_image_processor = CLIPImageProcessor.from_json_file(preprocessor_config_path) + clip_config = CLIPConfig.from_json_file(config_path) + model = StableDiffusionSafetyChecker.from_pretrained(safety_checker_model, config=clip_config) + model.eval() + + self.load_device = model_management.text_encoder_device() + self.offload_device = model_management.text_encoder_offload_device() + + model.to(self.offload_device) + + self.safety_checker_model = ModelPatcher(model, load_device=self.load_device, offload_device=self.offload_device) + + def censor(self, images: list | np.ndarray) -> list | np.ndarray: + self.init() + model_management.load_model_gpu(self.safety_checker_model) + + single = False + if not isinstance(images, list) or isinstance(images, np.ndarray): + images = [images] + single = True + + safety_checker_input = self.clip_image_processor(images, return_tensors="pt") + safety_checker_input.to(device=self.load_device) + checked_images, has_nsfw_concept = self.safety_checker_model.model(images=images, + clip_input=safety_checker_input.pixel_values) + checked_images = [image.astype(np.uint8) for image in checked_images] + + if single: + checked_images = checked_images[0] + + return checked_images -def numpy_to_pil(image): - image = (image * 255).round().astype("uint8") - pil_image = Image.fromarray(image) - - return pil_image - - -# check and replace nsfw content -def check_safety(x_image): - global safety_feature_extractor, safety_checker - - if safety_feature_extractor is None or safety_checker is None: - safety_checker_model = modules.config.downloading_safety_checker_model() - safety_feature_extractor = CLIPFeatureExtractor.from_json_file(preprocessor_config_path) - clip_config = CLIPConfig.from_json_file(config_path) - safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_checker_model, config=clip_config) - - safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt") - x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values) - - return x_checked_image, has_nsfw_concept - - -def censor_single(x): - x_checked_image, has_nsfw_concept = check_safety(x) - - # replace image with black pixels, keep dimensions - # workaround due to different numpy / pytorch image matrix format - if has_nsfw_concept[0]: - imageshape = x_checked_image.shape - x_checked_image = np.zeros((imageshape[0], imageshape[1], 3), dtype = np.uint8) - - return x_checked_image - - -def censor_batch(images): - images = [censor_single(image) for image in images] - - return images \ No newline at end of file +default_censor = Censor().censor diff --git a/modules/async_worker.py b/modules/async_worker.py index 892f99a7..302db84c 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -44,7 +44,7 @@ def worker(): import fooocus_version import args_manager - from extras.censor import censor_batch, censor_single + from extras.censor import default_censor from modules.sdxl_styles import apply_style, get_random_style, fooocus_expansion, apply_arrays, random_style_name from modules.private_logger import log from extras.expansion import safe_str @@ -78,7 +78,7 @@ def worker(): if censor and (modules.config.default_black_out_nsfw or black_out_nsfw): progressbar(async_task, progressbar_index, 'Checking for NSFW content ...') - imgs = censor_batch(imgs) + imgs = default_censor(imgs) async_task.results = async_task.results + imgs @@ -615,7 +615,8 @@ def worker(): d = [('Upscale (Fast)', 'upscale_fast', '2x')] if modules.config.default_black_out_nsfw or black_out_nsfw: progressbar(async_task, 100, 'Checking for NSFW content ...') - uov_input_image = censor_single(uov_input_image) + uov_input_image = default_censor(uov_input_image) + progressbar(async_task, 100, 'Saving image to system ...') uov_input_image_path = log(uov_input_image, d, output_format=output_format) yield_result(async_task, uov_input_image_path, black_out_nsfw, False, do_not_show_finished_images=True) return @@ -883,12 +884,12 @@ def worker(): imgs = [inpaint_worker.current_task.post_process(x) for x in imgs] img_paths = [] - + current_progress = int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)) if modules.config.default_black_out_nsfw or black_out_nsfw: - progressbar(async_task, int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)), - 'Checking for NSFW content ...') - imgs = censor_batch(imgs) + progressbar(async_task, current_progress, 'Checking for NSFW content ...') + imgs = default_censor(imgs) + progressbar(async_task, current_progress, 'Saving image to system ...') for x in imgs: d = [('Prompt', 'prompt', task['log_positive_prompt']), ('Negative Prompt', 'negative_prompt', task['log_negative_prompt']), From e94b97604f245c137bced3c5a941b45667f1483b Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sun, 19 May 2024 18:37:18 +0200 Subject: [PATCH 09/12] release: bump version number to 2.4.0-rc2 --- fooocus_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fooocus_version.py b/fooocus_version.py index def34a20..41556f90 100644 --- a/fooocus_version.py +++ b/fooocus_version.py @@ -1 +1 @@ -version = '2.4.0-rc1' +version = '2.4.0-rc2' From d2255a940c9466887e209e971c9093fb8a5d235a Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sun, 19 May 2024 20:12:28 +0200 Subject: [PATCH 10/12] feat: align progress bar vertically --- css/style.css | 1 + 1 file changed, 1 insertion(+) diff --git a/css/style.css b/css/style.css index b9e6e2ce..f1c7c606 100644 --- a/css/style.css +++ b/css/style.css @@ -27,6 +27,7 @@ progress { border-radius: 5px; /* Round the corners of the progress bar */ background-color: #f3f3f3; /* Light grey background */ width: 100%; + vertical-align: middle !important; } /* Style the progress bar container */ From 888de562908ccb7820689b99b88d49771fd9db63 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sun, 19 May 2024 20:13:29 +0200 Subject: [PATCH 11/12] feat: use fixed width for status text, remove ordinals --- css/style.css | 5 +++++ modules/async_worker.py | 9 +++++---- modules/util.py | 4 ---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/css/style.css b/css/style.css index f1c7c606..b5f7a448 100644 --- a/css/style.css +++ b/css/style.css @@ -70,6 +70,11 @@ progress::after { height: 30px !important; } +.progress-bar span { + text-align: right; + width: 200px; +} + .type_row{ height: 80px !important; } diff --git a/modules/async_worker.py b/modules/async_worker.py index 302db84c..9523ec48 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -49,7 +49,7 @@ def worker(): from modules.private_logger import log from extras.expansion import safe_str from modules.util import (remove_empty_str, HWC3, resize_image, get_image_shape_ceil, set_image_shape_ceil, - get_shape_ceil, resample_image, erode_or_dilate, ordinal_suffix, get_enabled_loras, + get_shape_ceil, resample_image, erode_or_dilate, get_enabled_loras, parse_lora_references_from_prompt, apply_wildcards) from modules.upscaler import perform_upscale from modules.flags import Performance @@ -838,10 +838,11 @@ def worker(): done_steps = current_task_id * steps + step async_task.yields.append(['preview', ( int(15.0 + 85.0 * float(done_steps) / float(all_steps)), - f'Step {step}/{total_steps} in the {current_task_id + 1}{ordinal_suffix(current_task_id + 1)} Sampling', - y)]) + f'Sampling step {step + 1}/{total_steps}, image {current_task_id + 1}/{image_number} ...', y)]) for current_task_id, task in enumerate(tasks): + current_progress = int(15.0 + 85.0 * float(current_task_id * steps) / float(all_steps)) + progressbar(async_task, current_progress, f'Preparing task {current_task_id + 1}/{image_number} ...') execution_start_time = time.perf_counter() try: @@ -889,7 +890,7 @@ def worker(): progressbar(async_task, current_progress, 'Checking for NSFW content ...') imgs = default_censor(imgs) - progressbar(async_task, current_progress, 'Saving image to system ...') + progressbar(async_task, current_progress, f'Saving image {current_task_id + 1}/{image_number} to system ...') for x in imgs: d = [('Prompt', 'prompt', task['log_positive_prompt']), ('Negative Prompt', 'negative_prompt', task['log_negative_prompt']), diff --git a/modules/util.py b/modules/util.py index 73430230..8e85ffbe 100644 --- a/modules/util.py +++ b/modules/util.py @@ -372,10 +372,6 @@ def get_file_from_folder_list(name, folders): return os.path.abspath(os.path.realpath(os.path.join(folders[0], name))) -def ordinal_suffix(number: int) -> str: - return 'th' if 10 <= number % 100 <= 20 else {1: 'st', 2: 'nd', 3: 'rd'}.get(number % 10, 'th') - - def makedirs_with_log(path): try: os.makedirs(path, exist_ok=True) From 5d8eb5235aeb44bcc7b3a133e7ad8bf9551c924a Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sun, 19 May 2024 20:33:21 +0200 Subject: [PATCH 12/12] refactor: align progress to actions --- modules/async_worker.py | 30 +++++++++++++++--------------- modules/flags.py | 1 + 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/modules/async_worker.py b/modules/async_worker.py index 9523ec48..8e4d6d95 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -72,7 +72,7 @@ def worker(): async_task.yields.append(['preview', (number, text, None)]) def yield_result(async_task, imgs, black_out_nsfw, censor=True, do_not_show_finished_images=False, - progressbar_index=13): + progressbar_index=flags.preparation_step_count): if not isinstance(imgs, list): imgs = [imgs] @@ -456,7 +456,7 @@ def worker(): extra_positive_prompts = prompts[1:] if len(prompts) > 1 else [] extra_negative_prompts = negative_prompts[1:] if len(negative_prompts) > 1 else [] - progressbar(async_task, 3, 'Loading models ...') + progressbar(async_task, 2, 'Loading models ...') loras = parse_lora_references_from_prompt(prompt, loras, modules.config.default_max_lora_number) @@ -523,25 +523,25 @@ def worker(): if use_expansion: for i, t in enumerate(tasks): - progressbar(async_task, 5, f'Preparing Fooocus text #{i + 1} ...') + progressbar(async_task, 4, f'Preparing Fooocus text #{i + 1} ...') expansion = pipeline.final_expansion(t['task_prompt'], t['task_seed']) print(f'[Prompt Expansion] {expansion}') t['expansion'] = expansion t['positive'] = copy.deepcopy(t['positive']) + [expansion] # Deep copy. for i, t in enumerate(tasks): - progressbar(async_task, 7, f'Encoding positive #{i + 1} ...') + progressbar(async_task, 5, f'Encoding positive #{i + 1} ...') t['c'] = pipeline.clip_encode(texts=t['positive'], pool_top_k=t['positive_top_k']) for i, t in enumerate(tasks): if abs(float(cfg_scale) - 1.0) < 1e-4: t['uc'] = pipeline.clone_cond(t['c']) else: - progressbar(async_task, 10, f'Encoding negative #{i + 1} ...') + progressbar(async_task, 6, f'Encoding negative #{i + 1} ...') t['uc'] = pipeline.clip_encode(texts=t['negative'], pool_top_k=t['negative_top_k']) if len(goals) > 0: - progressbar(async_task, 13, 'Image processing ...') + progressbar(async_task, 7, 'Image processing ...') if 'vary' in goals: if 'subtle' in uov_method: @@ -562,7 +562,7 @@ def worker(): uov_input_image = set_image_shape_ceil(uov_input_image, shape_ceil) initial_pixels = core.numpy_to_pytorch(uov_input_image) - progressbar(async_task, 13, 'VAE encoding ...') + progressbar(async_task, 8, 'VAE encoding ...') candidate_vae, _ = pipeline.get_candidate_vae( steps=steps, @@ -579,7 +579,7 @@ def worker(): if 'upscale' in goals: H, W, C = uov_input_image.shape - progressbar(async_task, 13, f'Upscaling image from {str((H, W))} ...') + progressbar(async_task, 9, f'Upscaling image from {str((H, W))} ...') uov_input_image = perform_upscale(uov_input_image) print(f'Image upscaled.') @@ -628,7 +628,7 @@ def worker(): denoising_strength = overwrite_upscale_strength initial_pixels = core.numpy_to_pytorch(uov_input_image) - progressbar(async_task, 13, 'VAE encoding ...') + progressbar(async_task, 10, 'VAE encoding ...') candidate_vae, _ = pipeline.get_candidate_vae( steps=steps, @@ -686,7 +686,7 @@ def worker(): do_not_show_finished_images=True) return - progressbar(async_task, 13, 'VAE Inpaint encoding ...') + progressbar(async_task, 11, 'VAE Inpaint encoding ...') inpaint_pixel_fill = core.numpy_to_pytorch(inpaint_worker.current_task.interested_fill) inpaint_pixel_image = core.numpy_to_pytorch(inpaint_worker.current_task.interested_image) @@ -706,7 +706,7 @@ def worker(): latent_swap = None if candidate_vae_swap is not None: - progressbar(async_task, 13, 'VAE SD15 encoding ...') + progressbar(async_task, 12, 'VAE SD15 encoding ...') latent_swap = core.encode_vae( vae=candidate_vae_swap, pixels=inpaint_pixel_fill)['samples'] @@ -832,16 +832,16 @@ def worker(): zsnr=False)[0] print(f'Using {scheduler_name} scheduler.') - async_task.yields.append(['preview', (13, 'Moving model to GPU ...', None)]) + async_task.yields.append(['preview', (flags.preparation_step_count, 'Moving model to GPU ...', None)]) def callback(step, x0, x, total_steps, y): done_steps = current_task_id * steps + step async_task.yields.append(['preview', ( - int(15.0 + 85.0 * float(done_steps) / float(all_steps)), + int(flags.preparation_step_count + (100 - flags.preparation_step_count) * float(done_steps) / float(all_steps)), f'Sampling step {step + 1}/{total_steps}, image {current_task_id + 1}/{image_number} ...', y)]) for current_task_id, task in enumerate(tasks): - current_progress = int(15.0 + 85.0 * float(current_task_id * steps) / float(all_steps)) + current_progress = int(flags.preparation_step_count + (100 - flags.preparation_step_count) * float(current_task_id * steps) / float(all_steps)) progressbar(async_task, current_progress, f'Preparing task {current_task_id + 1}/{image_number} ...') execution_start_time = time.perf_counter() @@ -885,7 +885,7 @@ def worker(): imgs = [inpaint_worker.current_task.post_process(x) for x in imgs] img_paths = [] - current_progress = int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)) + current_progress = int(flags.preparation_step_count + (100 - flags.preparation_step_count) * float((current_task_id + 1) * steps) / float(all_steps)) if modules.config.default_black_out_nsfw or black_out_nsfw: progressbar(async_task, current_progress, 'Checking for NSFW content ...') imgs = default_censor(imgs) diff --git a/modules/flags.py b/modules/flags.py index 77ad012a..7b3ac393 100644 --- a/modules/flags.py +++ b/modules/flags.py @@ -93,6 +93,7 @@ metadata_scheme = [ ] controlnet_image_count = 4 +preparation_step_count = 13 class OutputFormat(Enum):