from pathlib import Path import os import sys import random import time from typing import List, Dict import modules.scripts as scripts import gradio as gr from modules import processing, shared, sd_samplers, images, ui from modules.processing import Processed from modules.sd_samplers import samplers from modules.shared import opts, cmd_opts, state from PIL import Image, ImageFile, UnidentifiedImageError from tagger import format from tagger.preset import Preset from tagger.interrogator import Interrogator, DeepDanbooruInterrogator, WaifuDiffusionInterrogator from tagger.utils import split_str def refresh_interrogators() -> List[str]: global interrogators interrogators = {} # load waifu diffusion 1.4 tagger models interrogators['wd14-vit'] = WaifuDiffusionInterrogator('wd14-vit') interrogators['wd14-convnext'] = WaifuDiffusionInterrogator( 'wd14-convnext', repo='SmilingWolf/wd-v1-4-convnext-tagger' ) # load deepdanbooru project os.makedirs( shared.cmd_opts.deepdanbooru_projects_path, exist_ok=True ) for path in os.scandir(shared.cmd_opts.deepdanbooru_projects_path): if not path.is_dir(): continue if not Path(path, 'project.json').is_file(): continue interrogators[path.name] = DeepDanbooruInterrogator(path.name, path) print(interrogators) return sorted(interrogators.keys()) class Script(scripts.Script): # if you do not initialize the Image object # Image.registered_extensions() returns only PNG Image.init() # PIL spits errors when loading a truncated image by default # https://pillow.readthedocs.io/en/stable/reference/ImageFile.html#PIL.ImageFile.LOAD_TRUNCATED_IMAGES ImageFile.LOAD_TRUNCATED_IMAGES = True def title(self): return "Waifu Tagger to Prompt" def show(self, is_img2img): return is_img2img def ui(self, is_img2img): loops = gr.Slider(minimum=1, maximum=512, step=1, label='Iters', value=4) denoising_strength_change_factor = gr.Slider(minimum=0.9, maximum=1.1, step=0.005, label='Denoising strength change factor', value=1) thresholdGr = gr.Slider(minimum=0, maximum=1, step=0.01, label='Threshold', value=0.35) #split_str(additional_tags), #split_str(exclude_tags), sort_alphabetical = gr.Checkbox(label='Sort alphabetically', value=False) #sort_by_alphabetical_order = False, confident_weight = gr.Checkbox(label='Set confidence as weight', value=False) #add_confident_as_weight = False, swap_underscore = gr.Checkbox(label='Replace underscore', value=True) #replace_underscore = True, escape_char = gr.Checkbox(label='Escape parenthesis', value=True) #escape_tag = True interrogator_names = refresh_interrogators() with gr.Row(variant='compact'): interrogator = gr.Dropdown(interrogator_names, value=( None if len(interrogator_names) < 1 else interrogator_names[-1] ) ) ui.create_refresh_button( interrogator, lambda: None, lambda: {'choices': refresh_interrogators()}, 'refresh_interrogator' ) startPrompt = gr.Textbox(label="Prompt at start") endPrompt = gr.Textbox(label="Prompt at end") outdir = gr.Textbox(label='Output Directory') #negPrompt = gr.Textbox(label="Add to end of negative propmt") unload_all_models = gr.Button( value='Unload all interrogate models (Placebo)' ) return [loops, denoising_strength_change_factor, thresholdGr, sort_alphabetical, confident_weight, swap_underscore, escape_char, interrogator, startPrompt, endPrompt, outdir]#, negPrompt] def run(self, p, loops, denoising_strength_change_factor, thresholdGr, sort_alphabetical, confident_weight, swap_underscore, escape_char, interrogator, startPrompt, endPrompt, outdir=str(time.time())):#, negPrompt): history = [] #interrogator = WaifuDiffusionInterrogator('wd14-vit') p.do_not_save_grid = True job_count = 0 jobs = [] processing.fix_seed(p) batch_count = p.n_iter p.batch_size = 1 p.n_iter = 1 output_images, info = None, None initial_seed = None initial_info = None grids = [] all_images = [] #original_init_image = p.init_images state.job_count = loops * batch_count initial_color_corrections = [processing.setup_color_correction(p.init_images[0])] #this does not feel right. interrogator = interrogators[interrogator] p.outpath_samples = "outputs/waifutagger/" + outdir #i should probably figure something better out for this LOOOOOOOOL for i in range(loops): p.n_iter = 1 p.batch_size = 1 p.do_not_save_grid = True if opts.img2img_color_correction: p.color_corrections = initial_color_corrections #state.job = f"Iteration {i + 1}/{loops}, batch {n + 1}/{batch_count}" processed = processing.process_images(p) if initial_seed is None: initial_seed = processed.seed initial_info = processed.info init_img = processed.images[0] p.init_images = [init_img] p.seed = processed.seed + 1 p.denoising_strength = min(max(p.denoising_strength * denoising_strength_change_factor, 0.1), 1) history.append(processed.images[0]) ratings, tags = interrogator.interrogate(init_img) processed_tags = Interrogator.postprocess_tags( tags, threshold = thresholdGr, #split_str(additional_tags), #split_str(exclude_tags), sort_by_alphabetical_order = sort_alphabetical, add_confident_as_weight = confident_weight, replace_underscore = swap_underscore, escape_tag = escape_char #*postprocess_opts ) #did you know they made .format() for shit like this? newprompt = startPrompt + ', ' + ', '.join(processed_tags) + ', ' + endPrompt print(newprompt) p.prompt = newprompt grid = images.image_grid(history, rows=1) if opts.grid_save: images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p) grids.append(grid) all_images += history if opts.return_grid: all_images = grids + all_images processed = Processed(p, all_images, initial_seed, initial_info) return processed