Rework core_generation_funnel to be generator

Closes thygate#302

Author: semjon00

scripts/depthmap.py

@@ -43,30 +43,34 @@ def run(self, p, *inputs):
                 continue
             inputimages.append(processed.images[count])
 
-        outputs, mesh_fi, meshsimple_fi = core_generation_funnel(p.outpath_samples, inputimages, None, None, inputs, backbone.gather_ops())
+        gen_obj = core_generation_funnel(p.outpath_samples, inputimages, None, None, inputs, backbone.gather_ops())
+
+        for input_i, type, result in gen_obj:
+            if type in ['simple_mesh', 'inpainted_mesh']:
+                continue  # We are in script mode: do nothing with the filenames
 
-        for input_i, imgs in enumerate(outputs):
             # get generation parameters
+            # TODO: could reuse
             if hasattr(processed, 'all_prompts') and shared.opts.enable_pnginfo:
-                info = create_infotext(processed, processed.all_prompts, processed.all_seeds, processed.all_subseeds,
-                                       "", 0, input_i)
+                info = create_infotext(
+                    processed, processed.all_prompts, processed.all_seeds, processed.all_subseeds, "", 0, input_i)
             else:
                 info = None
-            for image_type, image in list(imgs.items()):
-                processed.images.append(image)
-                if inputs["save_outputs"]:
-                    try:
-                        suffix = "" if image_type == "depth" else f"{image_type}"
-                        backbone.save_image(image, path=p.outpath_samples, basename="", seed=processed.all_seeds[input_i],
-                                   prompt=processed.all_prompts[input_i], extension=shared.opts.samples_format,
-                                   info=info,
-                                   p=processed,
-                                   suffix=suffix)
-                    except Exception as e:
-                        if not ('image has wrong mode' in str(e) or 'I;16' in str(e)):
-                            raise e
-                        print('Catched exception: image has wrong mode!')
-                        traceback.print_exc()
+
+            processed.images.append(result)
+            if inputs["save_outputs"]:
+                try:
+                    suffix = "" if type == "depth" else f"{type}"
+                    backbone.save_image(result, path=p.outpath_samples, basename="", seed=processed.all_seeds[input_i],
+                               prompt=processed.all_prompts[input_i], extension=shared.opts.samples_format,
+                               info=info,
+                               p=processed,
+                               suffix=suffix)
+                except Exception as e:
+                    if not ('image has wrong mode' in str(e) or 'I;16' in str(e)):
+                        raise e
+                    print('Catched exception: image has wrong mode!')
+                    traceback.print_exc()
         return processed
 
 

scripts/depthmap_api.py

@@ -53,23 +53,24 @@ async def process(
         depth_input_images: List[str] = Body([], title='Input Images'),
         options: Dict[str, object] = Body("options", title='Generation options'),
     ):
-        if len(depth_input_images) == 0:
-            raise HTTPException(status_code=422, detail="No image supplied")
+        # TODO: restrict mesh options
 
-        print(f"Processing {str(len(depth_input_images))} images trough the API.")
+        if len(depth_input_images) == 0:
+            raise HTTPException(status_code=422, detail="No images supplied")
+        print(f"Processing {str(len(depth_input_images))} images trough the API")
 
-        PIL_images = []
+        pil_images = []
         for input_image in depth_input_images:
-            PIL_images.append(to_base64_PIL(input_image))
-
+            pil_images.append(to_base64_PIL(input_image))
         outpath = backbone.get_outpath()
-        results, _, _ = core_generation_funnel(outpath, PIL_images, None, None, options)
+        gen_obj = core_generation_funnel(outpath, pil_images, None, None, options)
 
-        # TODO: Fix: this just keeps depth image throws everything else away
-        results = [img['depth'] for img in results]
-        results64 = list(map(encode_to_base64, results))
+        results_based = []
+        for count, type, result in gen_obj:
+            if type not in ['simple_mesh', 'inpainted_mesh']:
+                results_based += [encode_to_base64(result)]
+        return {"images": results_based, "info": "Success"}
 
-        return {"images": results64, "info": "Success"}
 
 try:
     import modules.script_callbacks as script_callbacks

src/backbone.py

@@ -3,6 +3,13 @@
 # must be resided in this file (or in the scripts folder).
 import pathlib
 from datetime import datetime
+import enum
+
+
+class BackboneType(enum.Enum):
+    WEBUI = 1
+    STANDALONE = 2
+
 
 try:
     # stable-diffusion-webui backbone
@@ -61,6 +68,8 @@ def reload_sd_model():
     def get_hide_dirs():
         import modules.shared
         return modules.shared.hide_dirs
+
+    USED_BACKBONE = BackboneType.WEBUI
 except:
     # Standalone backbone
     print(  # "  DepthMap did not detect stable-diffusion-webui; launching with the standalone backbone.\n"
@@ -116,3 +125,6 @@ def unload_sd_model(): pass  # Not needed
     def reload_sd_model(): pass  # Not needed
 
     def get_hide_dirs(): return {}  # Directories will not be hidden from traversal (except when starts with the dot)
+
+
+    USED_BACKBONE = BackboneType.STANDALONE

src/common_ui.py

@@ -470,36 +470,42 @@ def run_generate(*inputs):
         inputdepthmaps_n = len([1 for x in inputdepthmaps if x is not None])
         print(f'{len(inputimages)} images will be processed, {inputdepthmaps_n} existing depthmaps will be reused')
 
-    outputs, fn_mesh, display_mesh = core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inputs, backbone.gather_ops())
+    gen_obj = core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inputs, backbone.gather_ops())
 
-    # Saving images
     show_images = []
-    for input_i, imgs in enumerate(outputs):
+    inpainted_mesh_fi = mesh_simple_fi = None
+    for input_i, type, result in gen_obj:
+        if type == 'simple_mesh':
+            mesh_simple_fi = result
+            continue
+        if type == 'inpainted_mesh':
+            inpainted_mesh_fi = result
+            continue
+
         basename = 'depthmap'
         if depthmap_mode == '2' and inputnames[input_i] is not None and outpath != backbone.get_opt('outdir_extras_samples', None):
             basename = Path(inputnames[input_i]).stem
 
-        for image_type, image in list(imgs.items()):
-            show_images += [image]
-            if inputs["save_outputs"]:
-                try:
-                    suffix = "" if image_type == "depth" else f"{image_type}"
-                    backbone.save_image(image, path=outpath, basename=basename, seed=None,
-                               prompt=None, extension=backbone.get_opt('samples_format', 'png'), short_filename=True,
-                               no_prompt=True, grid=False, pnginfo_section_name="extras",
-                               suffix=suffix)
-                except Exception as e:
-                    if not ('image has wrong mode' in str(e) or 'I;16' in str(e)):
-                        raise e
-                    print('Catched exception: image has wrong mode!')
-                    traceback.print_exc()
-
-    display_mesh = None
-    # use inpainted 3d mesh to show in 3d model output when enabled in settings
-    if backbone.get_opt('depthmap_script_show_3d_inpaint', True) and fn_mesh is not None and len(fn_mesh) > 0:
-        display_mesh = fn_mesh
-    # however, don't show 3dmodel when disabled in settings
+        show_images += [result]
+        if inputs["save_outputs"]:
+            try:
+                suffix = "" if type == "depth" else f"{type}"
+                backbone.save_image(result, path=outpath, basename=basename, seed=None,
+                           prompt=None, extension=backbone.get_opt('samples_format', 'png'), short_filename=True,
+                           no_prompt=True, grid=False, pnginfo_section_name="extras",
+                           suffix=suffix)
+            except Exception as e:
+                if not ('image has wrong mode' in str(e) or 'I;16' in str(e)):
+                    raise e
+                print('Catched exception: image has wrong mode!')
+                traceback.print_exc()
+
+    # Deciding what mesh to display (and if)
+    display_mesh_fi = None
     if not backbone.get_opt('depthmap_script_show_3d', True):
-        display_mesh = None
+        display_mesh_fi = mesh_simple_fi
+        if backbone.get_opt('depthmap_script_show_3d_inpaint', True):
+            if inpainted_mesh_fi is not None and len(inpainted_mesh_fi) > 0:
+                display_mesh_fi = inpainted_mesh_fi
     # TODO: return more info
-    return show_images, fn_mesh, display_mesh, 'Generated!'
+    return show_images, inpainted_mesh_fi, display_mesh_fi, 'Generated!'

src/core.py

@@ -66,7 +66,7 @@ def __getattr__(self, item):
 
 def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp, ops=None):
     if len(inputimages) == 0 or inputimages[0] is None:
-        return [], '', ''
+        return
     if inputdepthmaps is None or len(inputdepthmaps) == 0:
         inputdepthmaps: list[Image] = [None for _ in range(len(inputimages))]
     inputdepthmaps_complete = all([x is not None for x in inputdepthmaps])
@@ -103,12 +103,7 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
         device = torch.device("cpu")
     print("device: %s" % device)
 
-    generated_images = [{} for _ in range(len(inputimages))]
-    """Images that will be returned.
-    Every array element corresponds to particular input image.
-    Dictionary keys are types of images that were derived from the input image."""
-    # TODO: ???
-    meshsimple_fi = None
+    # TODO: This should not be here
     inpaint_imgs = []
     inpaint_depths = []
 
@@ -192,14 +187,14 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
                         background_removed_array[:, :, 2] == 0) & (background_removed_array[:, :, 3] <= 0.2)
                 img_output[bg_mask] = 0  # far value
 
-                generated_images[count]['background_removed'] = background_removed_image
+                yield count, 'background_removed', background_removed_image
 
                 if inp[go.SAVE_BACKGROUND_REMOVAL_MASKS]:
                     bg_array = (1 - bg_mask.astype('int8')) * 255
                     mask_array = np.stack((bg_array, bg_array, bg_array, bg_array), axis=2)
                     mask_image = Image.fromarray(mask_array.astype(np.uint8))
 
-                    generated_images[count]['foreground_mask'] = mask_image
+                    yield count, 'foreground_mask', mask_image
 
             # A weird quirk: if user tries to save depthmap, whereas input depthmap is used,
             # depthmap will be outputed, even if output_depth_combine is used.
@@ -211,9 +206,9 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
                         img_concat = Image.fromarray(np.concatenate(
                             (inputimages[count], convert_i16_to_rgb(img_depth, inputimages[count])),
                             axis=axis))
-                        generated_images[count]['concat_depth'] = img_concat
+                        yield count, 'concat_depth', img_concat
                     else:
-                        generated_images[count]['depth'] = Image.fromarray(img_depth)
+                        yield count, 'depth', Image.fromarray(img_depth)
 
             if inp[go.GEN_STEREO]:
                 print("Generating stereoscopic images..")
@@ -222,21 +217,22 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
                     inp[go.STEREO_DIVERGENCE], inp[go.STEREO_SEPARATION],
                     inp[go.STEREO_MODES], inp[go.STEREO_BALANCE], inp[go.STEREO_FILL_ALGO])
                 for c in range(0, len(stereoimages)):
-                    generated_images[count][inp[go.STEREO_MODES][c]] = stereoimages[c]
+                    yield count, inp[go.STEREO_MODES][c], stereoimages[c]
 
             if inp[go.GEN_NORMALMAP]:
-                generated_images[count]['normalmap'] = create_normalmap(
+                normalmap = create_normalmap(
                     img_output,
                     inp[go.NORMALMAP_PRE_BLUR_KERNEL] if inp[go.NORMALMAP_PRE_BLUR] else None,
                     inp[go.NORMALMAP_SOBEL_KERNEL] if inp[go.NORMALMAP_SOBEL] else None,
                     inp[go.NORMALMAP_POST_BLUR_KERNEL] if inp[go.NORMALMAP_POST_BLUR] else None,
                     inp[go.NORMALMAP_INVERT]
                 )
+                yield count, 'normalmap', normalmap
 
             if inp[go.GEN_HEATMAP]:
                 from dzoedepth.utils.misc import colorize
                 heatmap = Image.fromarray(colorize(img_output, cmap='inferno'))
-                generated_images[count]['heatmap'] = heatmap
+                yield count, 'heatmap', heatmap
 
             # gen mesh
             if inp[go.GEN_SIMPLE_MESH]:
@@ -268,17 +264,25 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
                 mesh = create_mesh(inputimages[count], depthi, keep_edges=not inp[go.SIMPLE_MESH_OCCLUDE],
                                    spherical=(inp[go.SIMPLE_MESH_SPHERICAL]))
                 mesh.export(meshsimple_fi)
+                yield count, 'simple_mesh', meshsimple_fi
 
         print("Computing output(s) done.")
     except RuntimeError as e:
         # TODO: display in UI
         if 'out of memory' in str(e):
-            suggestion = "ERROR: out of memory, could not generate depthmap!\nPlease try a different model"
-            if device != torch.device("cpu"):
-                suggestion += ", or try using the CPU"
+            suggestion = "ERROR: out of GPU memory, could not generate depthmap! " \
+                         "Here are some suggestions to work around this issue:\n"
             if inp[go.BOOST]:
-                suggestion += ", or disable BOOST"
-            print(f"{suggestion}.")
+                suggestion += " * Disable BOOST (generation will be faster, but the depthmap will be less detailed)\n"
+            if backbone.USED_BACKBONE == backbone.BackboneType.WEBUI:
+                suggestion += " * Run DepthMap in the standalone mode - without launching the SD WebUI\n"
+            if device != torch.device("cpu"):
+                suggestion += " * Select CPU as the processing device (this will be slower)\n"
+            if inp[go.MODEL_TYPE] != 6:
+                suggestion += " * Use a different model (this could reduce quality)\n"
+            if not inp[go.BOOST]:
+                suggestion += " * Reduce net size (this could reduce quality)\n"
+            print(f"{suggestion}")
         else:
             raise e
     finally:
@@ -301,12 +305,12 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
             mesh_fi = run_3dphoto(device, inpaint_imgs, inpaint_depths, inputnames, outpath,
                                   inp[go.GEN_INPAINTED_MESH_DEMOS],
                                   1, "mp4")
+            yield 0, 'inpainted_mesh', mesh_fi
         except Exception as e:
             print(f'{str(e)}, some issue with generating inpainted mesh')
 
     backbone.reload_sd_model()
     print("All done.\n")
-    return generated_images, mesh_fi, meshsimple_fi
 
 
 def get_uniquefn(outpath, basename, ext):