diff options
| -rw-r--r-- | webui.py | 30 |
1 files changed, 24 insertions, 6 deletions
@@ -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
|