diff --git a/webui.py b/webui.py index 14b0dd24..7f8c928c 100644 --- a/webui.py +++ b/webui.py @@ -50,7 +50,10 @@ parser.add_argument("--no-verify-input", action='store_true', help="do not verif parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats") parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware accleration in browser)") parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI") +parser.add_argument("--save-format", type=str, default='png', help="file format for saved indiviual samples; can be png or jpg") parser.add_argument("--grid-format", type=str, default='png', help="file format for saved grids; can be png or jpg") +parser.add_argument("--grid-extended-filename", action='store_true', help="save grid images to filenames with extended info: seed, prompt") +parser.add_argument("--jpeg-quality", type=int, default=80, help="quality for saved jpeg images") opt = parser.parse_args() GFPGAN_dir = opt.gfpgan_dir @@ -128,10 +131,27 @@ def create_random_tensors(shape, seeds): x = torch.stack(xs) return x + def torch_gc(): torch.cuda.empty_cache() torch.cuda.ipc_collect() + +def sanitize_filename_part(text): + return text.replace(' ', '_').translate({ord(x): '' for x in invalid_filename_chars})[:128] + + +def save_image(image, path, basename, seed, prompt, extension, short_filename=False): + prompt = sanitize_filename_part(prompt) + + if short_filename: + filename = f"{basename}.{extension}" + else: + filename = f"{basename}-{seed}-{prompt[:128]}.{extension}" + + image.save(os.path.join(path, filename), quality=opt.jpeg_quality) + + def load_GFPGAN(): model_name = 'GFPGANv1.3' model_path = os.path.join(GFPGAN_dir, 'experiments/pretrained_models', model_name + '.pth') @@ -387,9 +407,7 @@ def process_images(outpath, func_init, func_sample, prompt, seed, sampler_name, x_sample = restored_img image = Image.fromarray(x_sample) - filename = f"{base_count:05}-{seeds[i]}_{prompts[i].replace(' ', '_').translate({ord(x): '' for x in invalid_filename_chars})[:128]}.png" - - image.save(os.path.join(sample_path, filename)) + save_image(image, sample_path, f"{base_count:05}", seeds[i], prompts[i], opt.save_format) output_images.append(image) base_count += 1 @@ -408,8 +426,7 @@ def process_images(outpath, func_init, func_sample, prompt, seed, sampler_name, output_images.insert(0, grid) - - grid.save(os.path.join(outpath, f'grid-{grid_count:04}.{opt.grid_format}')) + save_image(grid, outpath, f"grid-{grid_count:04}", seed, prompt, opt.grid_format, short_filename=not opt.grid_extended_filename) grid_count += 1 info = f""" @@ -601,7 +618,8 @@ def img2img(prompt: str, init_img, ddim_steps: int, use_GFPGAN: bool, prompt_mat grid_count = len(os.listdir(outpath)) - 1 grid = image_grid(history, batch_size, force_n_rows=1) - grid.save(os.path.join(outpath, f'grid-{grid_count:04}.{opt.grid_format}')) + + save_image(grid, outpath, f"grid-{grid_count:04}", initial_seed, prompt, opt.grid_format, short_filename=not opt.grid_extended_filename) output_images = history seed = initial_seed