Standalone interface (barely works)

Author: semjon00

main.py

@@ -0,0 +1,10 @@
+# This launches DepthMap without the AUTOMATIC1111/stable-diffusion-webui
+import argparse
+import src.common_ui
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--listen", help="Create public link")
+    args = parser.parse_args()
+
+    src.common_ui.on_ui_tabs().launch(share=args.listen)

scripts/depthmap.py

@@ -1,2 +1,96 @@
 # This file contains stable-duiffusion-webui stuff that the plugin relies on.
 # Eventually, when we have a standalone interface, this will load either standalone backbone or webui backbone.
+try:
+    # stable-duiffusion-webui backbone
+    from modules.images import save_image  # Should fail if not on stable-duiffusion-webui
+    from modules.devices import torch_gc  # TODO: is this really sufficient?
+    from modules.images import get_next_sequence_number
+    from modules.call_queue import wrap_gradio_gpu_call
+    from modules.shared import listfiles
+
+    def get_opt(name, default):
+        from modules.shared import opts
+
+        if hasattr(opts, name):
+            return opts.__getattr__(name)
+        return default
+
+
+    def gather_ops():
+        from modules.shared import cmd_opts
+        ops = {}
+        if get_opt('depthmap_script_boost_rmax', None) is not None:
+            ops['boost_whole_size_threshold'] = get_opt('depthmap_script_boost_rmax', None)
+        ops['precision'] = cmd_opts.precision
+        ops['no_half'] = cmd_opts.no_half
+        return ops
+
+
+    def get_outpath():
+        path = get_opt('outdir_samples', None)
+        if path is None or len(path) == 0:
+            path = get_opt('outdir_extras_samples', None)
+        assert path is not None and len(path) > 0
+        return path
+
+
+    def unload_sd_model():
+        from modules import shared, devices
+        if shared.sd_model is not None:
+            shared.sd_model.cond_stage_model.to(devices.cpu)
+            shared.sd_model.first_stage_model.to(devices.cpu)
+
+
+    def reload_sd_model():
+        from modules import shared, devices
+        if shared.sd_model is not None:
+            shared.sd_model.cond_stage_model.to(devices.device)
+            shared.sd_model.first_stage_model.to(devices.device)
+
+    def get_hide_dirs():
+        import modules.shared
+        return modules.shared.hide_dirs
+except:
+    # Standalone backbone
+    print("DepthMap did not detect stable-duiffusion-webui; launching with the standalone backbone.\n"
+          "The standalone backbone is not on par with the stable-duiffusion-webui backbone.\n"
+          "Some features may be missing or work differently.\n")
+
+    def save_image(image, path, basename, **kwargs):
+        import os
+        os.makedirs(path, exist_ok=True)
+        fullfn = os.path.join(path, f"{get_next_sequence_number()}-{basename}.{kwargs['extension']}")
+        image.save(fullfn, format=get_opt('samples_format', 'png'))
+
+    def torch_gc():
+        # TODO: is this really sufficient?
+        import torch
+        if torch.cuda.is_available():
+            with torch.cuda.device('cuda'):
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+
+    def get_next_sequence_number():
+        # Don't really care what the number will be... As long as it is unique.
+        from datetime import datetime, timezone
+        import random
+        return f"{int(datetime.now(timezone.utc).timestamp())}-{random.randint(1000,9999)}"
+
+    def wrap_gradio_gpu_call(f): return f  # Displaying various stats is not supported
+
+    def listfiles(dirname):
+        import os
+        filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname)) if not x.startswith(".")]
+        return [file for file in filenames if os.path.isfile(file)]
+
+    def get_opt(name, default): return default  # Configuring is not supported
+
+    def gather_ops(): return {}  # Configuring is not supported
+
+    def get_outpath(): return '.'
+
+    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

src/backbone.py

@@ -105,7 +105,8 @@ def main_ui_panel(is_depth_tab):
                     gr.HTML("More options for generating video can be found in the Generate video tab")
 
         with gr.Group():
-            # TODO: it should be clear from the UI that the background removal does not use the model selected above
+            # TODO: it should be clear from the UI that there is an option of the background removal
+            #  that does not use the model selected above
             with gr.Row():
                 inp += "background_removal", gr.Checkbox(label="Remove background", value=False)
             with gr.Row(visible=False) as bgrem_options_row_1:
@@ -236,7 +237,7 @@ def on_ui_tabs():
                                           **backbone.get_hide_dirs(),
                                           placeholder="A directory on the same machine where the server is running.")
                         inp += gr.Textbox(elem_id="depthmap_batch_output_dir", label="Output directory",
-                                          **backbone.get_hide_dirs,
+                                          **backbone.get_hide_dirs(),
                                           placeholder="Leave blank to save images to the default path.")
                         gr.HTML("Files in the output directory may be overwritten.")
                         inp += gr.Checkbox(elem_id="depthmap_batch_reuse",
@@ -273,7 +274,7 @@ def on_ui_tabs():
                             with gr.Column():
                                 vid_html_info_x = gr.HTML()
                                 vid_html_info = gr.HTML()
-                                fn_mesh = gr.Textbox(label="Input Mesh (.ply | .obj)", **shared.hide_dirs,
+                                fn_mesh = gr.Textbox(label="Input Mesh (.ply | .obj)", **backbone.get_hide_dirs(),
                                                      placeholder="A file on the same machine where "
                                                                  "the server is running.")
                             with gr.Row():
@@ -377,7 +378,7 @@ def run_generate(*inputs):
     if depthmap_mode == '2' and depthmap_batch_output_dir != '':
         outpath = depthmap_batch_output_dir
     else:
-        outpath = backbone.opts.outdir_samples or backbone.opts.outdir_extras_samples
+        outpath = backbone.get_outpath()
 
     if depthmap_mode == '0':  # Single image
         if depthmap_input_image is None:
@@ -399,7 +400,7 @@ def run_generate(*inputs):
             inputimages.append(image)
             inputnames.append(os.path.splitext(img.orig_name)[0])
     elif depthmap_mode == '2':  # Batch from Directory
-        assert not backbone.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'
+        assert not backbone.get_opt('hide_ui_dir_config', False), '--hide-ui-dir-config option must be disabled'
         if depthmap_batch_input_dir == '':
             return [], None, None, "Please select an input directory."
         if depthmap_batch_input_dir == depthmap_batch_output_dir:
@@ -414,9 +415,9 @@ def run_generate(*inputs):
                 if depthmap_batch_reuse:
                     basename = Path(path).stem
                     # Custom names are not used in samples directory
-                    if outpath != backbone.opts.outdir_extras_samples:
+                    if outpath != backbone.get_opt('outdir_extras_samples', None):
                         # Possible filenames that the custom depthmaps may have
-                        name_candidates = [f'{basename}-0000.{backbone.opts.samples_format}',  # current format
+                        name_candidates = [f'{basename}-0000.{backbone.get_opt("samples_format", "png")}',  # current format
                                            f'{basename}.png',  # human-intuitive format
                                            f'{Path(path).name}']  # human-intuitive format (worse)
                         for fn_cand in name_candidates:
@@ -430,22 +431,22 @@ 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, mesh_fi, meshsimple_fi = core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inputs, backbone.gather_ops())
+    outputs, fn_mesh, display_mesh = core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inputs, backbone.gather_ops())
 
     # Saving images
     show_images = []
     for input_i, imgs in enumerate(outputs):
         basename = 'depthmap'
-        if depthmap_mode == '2' and inputnames[input_i] is not None and outpath != backbone.opts.outdir_extras_samples:
+        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}"
+                    suffix = "" if image_type == "depth" else f"{image_type}"
                     backbone.save_image(image, path=outpath, basename=basename, seed=None,
-                               prompt=None, extension=backbone.opts.samples_format, short_filename=True,
+                               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:
@@ -454,12 +455,12 @@ def run_generate(*inputs):
                     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 hasattr(backbone.opts, 'depthmap_script_show_3d_inpaint') and backbone.opts.depthmap_script_show_3d_inpaint \
-            and mesh_fi is not None and len(mesh_fi) > 0:
-        meshsimple_fi = mesh_fi
+    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
-    if hasattr(backbone.opts, 'depthmap_script_show_3d') and not backbone.opts.depthmap_script_show_3d:
-        meshsimple_fi = None
+    if not backbone.get_opt('depthmap_script_show_3d', True):
+        display_mesh = None
     # TODO: return more info
-    return show_images, mesh_fi, meshsimple_fi, 'Generated!'
+    return show_images, fn_mesh, display_mesh, 'Generated!'

src/common_ui.py

@@ -1,11 +1,6 @@
 from pathlib import Path
-
 from PIL import Image
 
-from modules import shared, devices
-from modules.images import get_next_sequence_number
-from modules.shared import opts, cmd_opts
-
 try:
     from tqdm import trange
 except:
@@ -24,6 +19,7 @@
 from src.main import *
 from src.stereoimage_generation import create_stereoimages
 from src.depthmap_generation import ModelHolder
+from src import backbone
 
 # 3d-photo-inpainting imports
 from inpaint.mesh import write_mesh, read_mesh, output_3d_photo
@@ -47,18 +43,6 @@ def convert_i16_to_rgb(image, like):
     return output
 
 
-def unload_sd_model():
-    if shared.sd_model is not None:
-        shared.sd_model.cond_stage_model.to(devices.cpu)
-        shared.sd_model.first_stage_model.to(devices.cpu)
-
-
-def reload_sd_model():
-    if shared.sd_model is not None:
-        shared.sd_model.cond_stage_model.to(devices.device)
-        shared.sd_model.first_stage_model.to(devices.device)
-
-
 def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp, ops=None):
     if len(inputimages) == 0 or inputimages[0] is None:
         return [], '', ''
@@ -97,12 +81,14 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
     stereo_modes = inp["stereo_modes"]
     stereo_separation = inp["stereo_separation"]
 
+    if ops is None:
+        ops = {}
     model_holder.update_settings(**ops)
 
     # TODO: ideally, run_depthmap should not save meshes - that makes the function not pure
     print(SCRIPT_FULL_NAME)
 
-    unload_sd_model()
+    backbone.unload_sd_model()
 
     # TODO: this still should not be here
     background_removed_images = []
@@ -308,7 +294,7 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
         else:
             raise e
     finally:
-        if hasattr(opts, 'depthmap_script_keepmodels') and opts.depthmap_script_keepmodels:
+        if backbone.get_opt('depthmap_script_keepmodels', False):
             model_holder.offload()  # Swap to CPU memory
         else:
             if 'model' in locals():
@@ -318,7 +304,7 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
             model_holder.unload_models()
 
         gc.collect()
-        devices.torch_gc()
+        backbone.torch_gc()
 
     # TODO: This should not be here
     mesh_fi = None
@@ -328,14 +314,14 @@ def core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inp
         except Exception as e:
             print(f'{str(e)}, some issue with generating inpainted mesh')
 
-    reload_sd_model()
+    backbone.reload_sd_model()
     print("All done.\n")
     return generated_images, mesh_fi, meshsimple_fi
 
 
 def get_uniquefn(outpath, basename, ext):
     # Inefficient and may fail, maybe use unbounded binary search?
-    basecount = get_next_sequence_number(outpath, basename)
+    basecount = backbone.get_next_sequence_number(outpath, basename)
     if basecount > 0: basecount = basecount - 1
     fullfn = None
     for i in range(500):
@@ -403,10 +389,7 @@ def run_3dphoto(device, img_rgb, img_depth, inputnames, outpath, inpaint_vids, v
         config['repeat_inpaint_edge'] = True
         config['ply_fmt'] = "bin"
 
-        config['save_ply'] = False
-        if hasattr(opts, 'depthmap_script_save_ply') and opts.depthmap_script_save_ply:
-            config['save_ply'] = True
-
+        config['save_ply'] = backbone.get_opt('depthmap_script_save_ply', False)
         config['save_obj'] = True
 
         if device == torch.device("cpu"):
@@ -473,7 +456,7 @@ def run_3dphoto(device, img_rgb, img_depth, inputnames, outpath, inpaint_vids, v
                                    [-0.05, -0.05, -0.05, -0.05],
                                    ['dolly-zoom-in', 'zoom-in', 'circle', 'swing'], False, vid_format, vid_ssaa)
 
-            devices.torch_gc()
+            backbone.torch_gc()
 
     finally:
         del rgb_model
@@ -482,7 +465,7 @@ def run_3dphoto(device, img_rgb, img_depth, inputnames, outpath, inpaint_vids, v
         depth_edge_model = None
         del depth_feat_model
         depth_feat_model = None
-        devices.torch_gc()
+        backbone.torch_gc()
 
     return mesh_fi
 
@@ -604,9 +587,9 @@ def run_makevideo(fn_mesh, vid_numframes, vid_fps, vid_traj, vid_shift, vid_bord
 
     # output path and filename mess ..
     basename = Path(fn_mesh).stem
-    outpath = opts.outdir_samples or opts.outdir_extras_samples
+    outpath = backbone.get_outpath()
     # unique filename
-    basecount = get_next_sequence_number(outpath, basename)
+    basecount = backbone.get_next_sequence_number(outpath, basename)
     if basecount > 0: basecount = basecount - 1
     fullfn = None
     for i in range(500):
@@ -699,9 +682,7 @@ def depth_edges_mask(depth):
 def create_mesh(image, depth, keep_edges=False, spherical=False):
     import trimesh
     from dzoedepth.utils.geometry import depth_to_points, create_triangles
-    maxsize = 1024
-    if hasattr(opts, 'depthmap_script_mesh_maxsize'):
-        maxsize = opts.depthmap_script_mesh_maxsize
+    maxsize = backbone.get_opt('depthmap_script_mesh_maxsize', 2048)
 
     # limit the size of the input image
     image.thumbnail((maxsize, maxsize))