support for running custom code (primarily to generate various labeled grids)

export for 4chan option

README.md

@@ -205,3 +205,46 @@ image will be upscaled to twice the original width and height, while width and h
 will specify the size of individual tiles. At the moment this method does not support batch size.
 
 ![](images/sd-upscale.jpg)
+
+### User scripts
+If the program is launched with `--allow-code` option, an extra text input field for script code
+is available in txt2img interface. It allows you to input python code that will do the work with
+image. If this field is not empty, the processing that would happen normally is skipped.
+
+In code, access parameters from web UI using the `p` variable, and provide outputs for web UI
+using the `display(images, seed, info)` function. All globals from script are also accessible.
+
+As an example, here is a script that draws a chart seen below (and also saves it as `test/gnomeplot/gnome5.png`):
+
+```python
+steps = [4, 8,12,16, 20]
+cfg_scales = [5.0,10.0,15.0]
+
+def cell(x, y, p=p):
+	p.steps = x
+	p.cfg_scale = y
+	return process_images(p).images[0]
+
+images = [draw_xy_grid(
+	xs = steps,
+	ys = cfg_scales,
+	x_label = lambda x: f'Steps = {x}',
+	y_label = lambda y: f'CFG = {y}',
+	cell = cell
+)]
+
+save_image(images[0], 'test/gnomeplot', 'gnome5')
+display(images)
+```
+
+![](images/scripting.jpg)
+
+A more simple script that would just process the image and output it normally:
+
+```python
+processed = process_images(p)
+
+print("Seed was: " + str(processed.seed))
+
+display(processed.images, processed.seed, processed.info)
+```

webui.py

@@ -50,6 +50,7 @@ parser.add_argument("--no-half", action='store_true', help="do not switch the mo
 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("--embeddings-dir", type=str, default='embeddings', help="embeddings dirtectory for textual inversion (default: embeddings)")
+parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
 
 cmd_opts = parser.parse_args()
 
@@ -132,6 +133,7 @@ class Options:
         "grid_extended_filename": OptionInfo(False, "Add extended info (seed, prompt) to filename when saving grid"),
         "n_rows": OptionInfo(-1, "Grid row count; use -1 for autodetect and 0 for it to be same as batch size", gr.Slider, {"minimum": -1, "maximum": 16, "step": 1}),
         "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
+        "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),
         "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),
         "prompt_matrix_add_to_start": OptionInfo(True, "In prompt matrix, add the variable combination of text to the start of the prompt, rather than the end"),
         "sd_upscale_upscaler_index": OptionInfo("RealESRGAN", "Upscaler to use for SD upscale", gr.Radio, {"choices": list(sd_upscalers.keys())}),
@@ -206,13 +208,12 @@ def torch_gc():
     torch.cuda.ipc_collect()
 
 
-def save_image(image, path, basename, seed, prompt, extension, info=None, short_filename=False):
-    prompt = sanitize_filename_part(prompt)
+def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False):
 
-    if short_filename:
-        filename = f"{basename}.{extension}"
+    if short_filename or prompt is None or seed is None:
+        filename = f"{basename}"
     else:
-        filename = f"{basename}-{seed}-{prompt[:128]}.{extension}"
+        filename = f"{basename}-{seed}-{sanitize_filename_part(prompt)[:128]}"
 
     if extension == 'png' and opts.enable_pnginfo and info is not None:
         pnginfo = PngImagePlugin.PngInfo()
@@ -220,7 +221,23 @@ def save_image(image, path, basename, seed, prompt, extension, info=None, short_
     else:
         pnginfo = None
 
-    image.save(os.path.join(path, filename), quality=opts.jpeg_quality, pnginfo=pnginfo)
+    os.makedirs(path, exist_ok=True)
+    fullfn = os.path.join(path, f"{filename}.{extension}")
+    image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo)
+
+    target_side_length = 4000
+    oversize = image.width > target_side_length or image.height > target_side_length
+    if opts.export_for_4chan and (oversize or os.stat(fullfn).st_size > 4 * 1024 * 1024):
+        ratio = image.width / image.height
+
+        if oversize and ratio > 1:
+            image = image.resize((target_side_length, image.height * target_side_length // image.width), LANCZOS)
+        elif oversize:
+            image = image.resize((image.width * target_side_length // image.height, target_side_length), LANCZOS)
+
+        image.save(os.path.join(path, f"{filename}.jpg"), quality=opts.jpeg_quality, pnginfo=pnginfo)
+
+
 
 
 def sanitize_filename_part(text):
@@ -244,10 +261,9 @@ def load_gfpgan():
     return GFPGANer(model_path=model_path, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None)
 
 
-def image_grid(imgs, batch_size, force_n_rows=None):
-    if force_n_rows is not None:
-        rows = force_n_rows
-    elif opts.n_rows > 0:
+def image_grid(imgs, batch_size=1, rows=None):
+    if rows is None:
+        if opts.n_rows > 0:
             rows = opts.n_rows
         elif opts.n_rows == 0:
             rows = batch_size
@@ -427,6 +443,22 @@ def draw_prompt_matrix(im, width, height, all_prompts):
     return draw_grid_annotations(im, width, height, hor_texts, ver_texts)
 
 
+def draw_xy_grid(xs, ys, x_label, y_label, cell):
+    res = []
+
+    ver_texts = [[GridAnnotation(y_label(y))] for y in ys]
+    hor_texts = [[GridAnnotation(x_label(x))] for x in xs]
+
+    for y in ys:
+        for x in xs:
+            res.append(cell(x, y))
+
+
+    grid = image_grid(res, rows=len(ys))
+    grid = draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
+
+    return grid
+
 def resize_image(resize_mode, im, width, height):
     if resize_mode == 0:
         res = im.resize((width, height), resample=LANCZOS)
@@ -742,7 +774,10 @@ class KDiffusionSampler:
         return samples_ddim
 
 
-def process_images(p: StableDiffusionProcessing):
+Processed = namedtuple('Processed', ['images','seed', 'info'])
+
+
+def process_images(p: StableDiffusionProcessing) -> Processed:
     """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""
 
     prompt = p.prompt
@@ -753,10 +788,7 @@ def process_images(p: StableDiffusionProcessing):
 
     seed = int(random.randrange(4294967294) if p.seed == -1 else p.seed)
 
-    os.makedirs(p.outpath, exist_ok=True)
-
     sample_path = os.path.join(p.outpath, "samples")
-    os.makedirs(sample_path, exist_ok=True)
     base_count = len(os.listdir(sample_path))
     grid_count = len(os.listdir(p.outpath)) - 1
 
@@ -846,7 +878,7 @@ def process_images(p: StableDiffusionProcessing):
 
         if (p.prompt_matrix or opts.grid_save) and not p.do_not_save_grid:
             if p.prompt_matrix:
-                grid = image_grid(output_images, p.batch_size, force_n_rows=1 << ((len(prompt_matrix_parts)-1)//2))
+                grid = image_grid(output_images, p.batch_size, rows=1 << ((len(prompt_matrix_parts)-1)//2))
 
                 try:
                     grid = draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts)
@@ -863,7 +895,7 @@ def process_images(p: StableDiffusionProcessing):
             grid_count += 1
 
     torch_gc()
-    return output_images, seed, infotext()
+    return Processed(output_images, seed, infotext())
 
 
 class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
@@ -876,8 +908,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         samples_ddim = self.sampler.sample(self, x, conditioning, unconditional_conditioning)
         return samples_ddim
 
-
-def txt2img(prompt: str, ddim_steps: int, sampler_index: int, use_GFPGAN: bool, prompt_matrix: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, height: int, width: int):
+def txt2img(prompt: str, steps: int, sampler_index: int, use_GFPGAN: bool, prompt_matrix: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, height: int, width: int, code: str):
     outpath = opts.outdir or "outputs/txt2img-samples"
 
     p = StableDiffusionProcessingTxt2Img(
@@ -887,7 +918,7 @@ def txt2img(prompt: str, ddim_steps: int, sampler_index: int, use_GFPGAN: bool,
         sampler_index=sampler_index,
         batch_size=batch_size,
         n_iter=n_iter,
-        steps=ddim_steps,
+        steps=steps,
         cfg_scale=cfg_scale,
         width=width,
         height=height,
@@ -895,9 +926,29 @@ def txt2img(prompt: str, ddim_steps: int, sampler_index: int, use_GFPGAN: bool,
         use_GFPGAN=use_GFPGAN
     )
 
-    output_images, seed, info = process_images(p)
+    if code != '' and cmd_opts.allow_code:
+        p.do_not_save_grid = True
+        p.do_not_save_samples = True
+
+        display_result_data = [[], -1, ""]
+        def display(imgs, s=display_result_data[1], i=display_result_data[2]):
+            display_result_data[0] = imgs
+            display_result_data[1] = s
+            display_result_data[2] = i
+
+        from types import ModuleType
+        compiled = compile(code, '', 'exec')
+        module = ModuleType("testmodule")
+        module.__dict__.update(globals())
+        module.p = p
+        module.display = display
+        exec(compiled, module.__dict__)
+
+        processed = Processed(*display_result_data)
+    else:
+        processed = process_images(p)
 
-    return output_images, seed, plaintext_to_html(info)
+    return processed.images, processed.seed, plaintext_to_html(processed.info)
 
 
 class Flagging(gr.FlaggingCallback):
@@ -911,7 +962,7 @@ class Flagging(gr.FlaggingCallback):
         os.makedirs("log/images", exist_ok=True)
 
         # those must match the "txt2img" function
-        prompt, ddim_steps, sampler_name, use_gfpgan, prompt_matrix, ddim_eta, n_iter, n_samples, cfg_scale, request_seed, height, width, images, seed, comment = flag_data
+        prompt, ddim_steps, sampler_name, use_gfpgan, prompt_matrix, ddim_eta, n_iter, n_samples, cfg_scale, request_seed, height, width, code, images, seed, comment = flag_data
 
         filenames = []
 
@@ -955,6 +1006,7 @@ txt2img_interface = gr.Interface(
         gr.Number(label='Seed', value=-1),
         gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512),
         gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512),
+        gr.Textbox(label="Python script", visible=cmd_opts.allow_code, lines=1)
     ],
     outputs=[
         gr.Gallery(label="Images"),
@@ -1042,29 +1094,30 @@ def img2img(prompt: str, init_img, ddim_steps: int, sampler_index: int, use_GFPG
         output_images, info = None, None
         history = []
         initial_seed = None
+        initial_info = None
 
         for i in range(n_iter):
             p.n_iter = 1
             p.batch_size = 1
             p.do_not_save_grid = True
 
-            output_images, seed, info = process_images(p)
+            processed = process_images(p)
 
             if initial_seed is None:
-                initial_seed = seed
+                initial_seed = processed.seed
+                initial_info = processed.info
 
-            p.init_img = output_images[0]
-            p.seed = seed + 1
+            p.init_img = processed.images[0]
+            p.seed = processed.seed + 1
             p.denoising_strength = max(p.denoising_strength * 0.95, 0.1)
-            history.append(output_images[0])
+            history.append(processed.images[0])
 
         grid_count = len(os.listdir(outpath)) - 1
-        grid = image_grid(history, batch_size, force_n_rows=1)
+        grid = image_grid(history, batch_size, rows=1)
 
         save_image(grid, outpath, f"grid-{grid_count:04}", initial_seed, prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename)
 
-        output_images = history
-        seed = initial_seed
+        processed = Processed(history, initial_seed, initial_info)
 
     elif sd_upscale:
         initial_seed = None
@@ -1094,14 +1147,14 @@ def img2img(prompt: str, init_img, ddim_steps: int, sampler_index: int, use_GFPG
         for i in range(batch_count):
             p.init_images = work[i*p.batch_size:(i+1)*p.batch_size]
 
-            output_images, seed, info = process_images(p)
+            processed = process_images(p)
 
             if initial_seed is None:
-                initial_seed = seed
-                initial_info = info
+                initial_seed = processed.seed
+                initial_info = processed.info
 
-            p.seed = seed + 1
-            work_results += output_images
+            p.seed = processed.seed + 1
+            work_results += processed.images
 
         image_index = 0
         for y, h, row in grid.tiles:
@@ -1114,14 +1167,12 @@ def img2img(prompt: str, init_img, ddim_steps: int, sampler_index: int, use_GFPG
         grid_count = len(os.listdir(outpath)) - 1
         save_image(combined_image, outpath, f"grid-{grid_count:04}", initial_seed, prompt, opts.grid_format, info=initial_info, short_filename=not opts.grid_extended_filename)
 
-        output_images = [combined_image]
-        seed = initial_seed
-        info = initial_info
+        processed = Processed([combined_image], initial_seed, initial_info)
 
     else:
-        output_images, seed, info = process_images(p)
+        processed = process_images(p)
 
-    return output_images, seed, plaintext_to_html(info)
+    return processed.images, processed.seed, plaintext_to_html(processed.info)
 
 
 sample_img2img = "assets/stable-samples/img2img/sketch-mountains-input.jpg"
@@ -1192,9 +1243,7 @@ def run_extras(image, GFPGAN_strength, RealESRGAN_upscaling, RealESRGAN_model_in
     if have_realesrgan and RealESRGAN_upscaling != 1.0:
         image = upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index)
 
-    os.makedirs(outpath, exist_ok=True)
     base_count = len(os.listdir(outpath))
-
     save_image(image, outpath, f"{base_count:05}", None, '', opts.samples_format, short_filename=True)
 
     return image, 0, ''