Merge pull request #2214 from Symbiomatrix/vlad3

Optional repurpose batch mode for frames

modules/img2img.py

@@ -6,6 +6,7 @@ from modules import sd_samplers, shared, processing, images
 from modules.generation_parameters_copypaste import create_override_settings_dict
 from modules.ui import plaintext_to_html
 from modules.memstats import memory_stats
+import itertools # SBM Batch frames
 
 
 def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args):
@@ -24,45 +25,69 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args)
         is_inpaint_batch = len(inpaint_masks) > 0
     if is_inpaint_batch:
         shared.log.info(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")
+    # SBM Batch frame should actually print btcrept. Or customise the messages.
     shared.log.info(f"Will process {len(image_files)} images, creating {p.n_iter * p.batch_size} new images for each.")
     save_normally = output_dir == ''
     p.do_not_save_grid = True
     p.do_not_save_samples = not save_normally
     shared.state.job_count = len(image_files) * p.n_iter
-    for i, image_file in enumerate(image_files):
-        shared.state.job = f"{i+1} out of {len(image_files)}"
+    # SBM Batch frame mode, take 2.
+    if shared.opts.batch_frame_mode:
+        window_size = p.batch_size
+        btcrept = 1
+        p.seed = [p.seed] * window_size # SBM MONKEYPATCH: Need to change processing to support a fixed seed value.
+        p.subseed = [p.subseed] * window_size # SBM MONKEYPATCH
+    else: # SBM Frame mode is off, standard operation of repeating same images with sequential seed.
+        window_size = 1
+        btcrept = p.batch_size
+    for i in range(0, len(image_files), window_size):
+        shared.state.job = f"{i+1} to {min(i+window_size, len(image_files))} out of {len(image_files)}"
         if shared.state.skipped:
             shared.state.skipped = False
         if shared.state.interrupted:
             break
-        try:
-            img = Image.open(image_file)
-        except UnidentifiedImageError as e:
-            shared.log.error(f"Image error: {e}")
-            continue
-        img = ImageOps.exif_transpose(img)
-        if p.scale_by != 1:
-            p.width = int(img.width * p.scale_by)
-            p.height = int(img.height * p.scale_by)
-        p.init_images = [img] * p.batch_size
+        batch_image_files = image_files[i:i+window_size]
+        batch_images = []
+        for image_file in batch_image_files:
+            try:
+                img = Image.open(image_file)
+                if p.scale_by != 1:
+                    p.width = int(img.width * p.scale_by)
+                    p.height = int(img.height * p.scale_by)
+            except UnidentifiedImageError as e:
+                shared.log.error(f"Image error: {e}")
+                continue
+            img = ImageOps.exif_transpose(img)
+            batch_images.append(img)
+        batch_images = batch_images * btcrept # Standard mode sends the same image per batchsize.
+        p.init_images = batch_images
 
         if is_inpaint_batch:
             # try to find corresponding mask for an image using simple filename matching
-            mask_image_path = os.path.join(inpaint_mask_dir, os.path.basename(image_file))
-            # if not found use first one ("same mask for all images" use-case)
-            if mask_image_path not in inpaint_masks:
-                mask_image_path = inpaint_masks[0]
-            mask_image = Image.open(mask_image_path)
-            p.image_mask = mask_image
+            batch_mask_images = []
+            for image_file in batch_image_files:
+                mask_image_path = os.path.join(inpaint_mask_dir, os.path.basename(image_file))
+                # if not found use first one ("same mask for all images" use-case)
+                if mask_image_path not in inpaint_masks:
+                    mask_image_path = inpaint_masks[0]
+                mask_image = Image.open(mask_image_path)
+                batch_mask_images.append(mask_image)
+            batch_mask_images = batch_mask_images * btcrept
+            p.image_mask = batch_mask_images
+
+        batch_image_files = batch_image_files * btcrept # List used for naming later.
 
         proc = modules.scripts.scripts_img2img.run(p, *args)
         if proc is None:
             proc = processing.process_images(p)
-        for n, image in enumerate(proc.images):
+        for n, (image, image_file) in enumerate(itertools.zip_longest(proc.images,batch_image_files)):
             basename, ext = os.path.splitext(os.path.basename(image_file))
             ext = ext[1:]
             if len(proc.images) > 1:
-                basename = f'{basename}-{n}'
+                if shared.opts.batch_frame_mode: # SBM Frames are numbered globally.
+                    basename = f'{basename}-{n + i}'
+                else: # Images are numbered per rept.
+                    basename = f'{basename}-{n}'
             if not shared.opts.use_original_name_batch:
                 basename = ''
                 ext = shared.opts.samples_format
@@ -74,7 +99,7 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args)
             for k, v in items.items():
                 image.info[k] = v
             images.save_image(image, path=output_dir, basename=basename, seed=None, prompt=None, extension=ext, info=geninfo, short_filename=True, no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=image.info, forced_filename=None)
-        shared.log.debug(f'Processed: images={len(image_files)} memory={memory_stats()} op=batch')
+        shared.log.debug(f'Processed: {len(batch_image_files)} Memory: {memory_stats()} batch')
 
 
 def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, latent_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, full_quality: bool, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, image_cfg_scale: float, diffusers_guidance_rescale: float, refiner_steps: int, refiner_start: float, clip_skip: int, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, selected_scale_tab: int, height: int, width: int, scale_by: float, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_files: list, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, override_settings_texts, *args): # pylint: disable=unused-argument

modules/shared.py

@@ -463,7 +463,8 @@ options_templates.update(options_section(('advanced', "Inference Settings"), {
     "hypertile_unet_enabled": OptionInfo(False, "HyperTile for UNet enabled"),
     "hypertile_unet_tile": OptionInfo(256, "HyperTile for UNet tile size", gr.Slider, {"minimum": 256, "maximum": 1024, "step": 8}),
 
-    "inference_mode_sep": OptionInfo("<h2>Inference mode</h2>", "", gr.HTML),
+    "inference_other_sep": OptionInfo("<h2>Other</h2>", "", gr.HTML),
+    "batch_frame_mode": OptionInfo(False, "Use batchsize to process multiple images in batch mode"),
     "inference_mode": OptionInfo("no-grad", "Torch inference mode", gr.Radio, lambda: {"choices": ["no-grad", "inference-mode", "none"]}),
     "sd_vae_sliced_encode": OptionInfo(False, "VAE Slicing (original)"),
 }))