import traceback
from pathlib import Path
import gradio as gr
from PIL import Image

from src import backbone, video_mode
from src.core import core_generation_funnel, unload_models, run_makevideo
from src.depthmap_generation import ModelHolder
from src.gradio_args_transport import GradioComponentBundle
from src.misc import *
from src.common_constants import GenerationOptions as go

# Ugly workaround to fix gradio tempfile issue
def ensure_gradio_temp_directory():
    try:
        import tempfile
        path = os.path.join(tempfile.gettempdir(), 'gradio')
        if not (os.path.exists(path)):
            os.mkdir(path)
    except Exception as e:
        traceback.print_exc()


ensure_gradio_temp_directory()


def main_ui_panel(is_depth_tab):
    inp = GradioComponentBundle()
    # TODO: Greater visual separation
    with gr.Blocks():
        with gr.Row() as cur_option_root:
            inp -= 'depthmap_gen_row_0', cur_option_root
            inp += go.COMPUTE_DEVICE, gr.Radio(label="Compute on", choices=['GPU', 'CPU'], value='GPU')
            # TODO: Should return value instead of index. Maybe Enum should be used?
            inp += go.MODEL_TYPE, gr.Dropdown(label="Model",
                                             choices=['res101', 'dpt_beit_large_512 (midas 3.1)',
                                                      'dpt_beit_large_384 (midas 3.1)', 'dpt_large_384 (midas 3.0)',
                                                      'dpt_hybrid_384 (midas 3.0)',
                                                      'midas_v21', 'midas_v21_small',
                                                      'zoedepth_n (indoor)', 'zoedepth_k (outdoor)', 'zoedepth_nk',
                                                      'Marigold v1', 'Depth Anything', 'Depth Anything v2 Small',
                                                      'Depth Anything v2 Base', 'Depth Anything v2 Large'],
                                              value='Depth Anything v2 Base', type="index")
        with gr.Box() as cur_option_root:
            inp -= 'depthmap_gen_row_1', cur_option_root
            with gr.Row():
                inp += go.BOOST, gr.Checkbox(label="BOOST",
                                             info="Generate depth map parts in a mosaic fashion - very slow",
                                             value=False)
                inp += go.NET_SIZE_MATCH, gr.Checkbox(label="Match net size to input size",
                                                      info="Net size affects quality, performance and VRAM usage")
            with gr.Row() as options_depend_on_match_size:
                inp += go.NET_WIDTH, gr.Slider(minimum=64, maximum=2048, step=64, label='Net width')
                inp += go.NET_HEIGHT, gr.Slider(minimum=64, maximum=2048, step=64, label='Net height')
            with gr.Row():
                inp += go.TILING_MODE, gr.Checkbox(
                    label='Tiling mode', info='Reduces seams that appear if the depthmap is tiled into a grid'
                )

        with gr.Box() as cur_option_root:
            inp -= 'depthmap_gen_row_2', cur_option_root
            with gr.Row():
                with gr.Group():  # 50% of width
                    inp += "save_outputs", gr.Checkbox(label="Save Outputs", value=True)
                with gr.Group():  # 50% of width
                    inp += go.DO_OUTPUT_DEPTH, gr.Checkbox(label="Output DepthMap")
                    inp += go.OUTPUT_DEPTH_INVERT, gr.Checkbox(label="Invert (black=near, white=far)")
            with gr.Row() as options_depend_on_output_depth_1:
                inp += go.OUTPUT_DEPTH_COMBINE, gr.Checkbox(
                    label="Combine input and depthmap into one image")
                inp += go.OUTPUT_DEPTH_COMBINE_AXIS, gr.Radio(
                    label="Combine axis", choices=['Vertical', 'Horizontal'], type="value", visible=False)

        with gr.Box() as cur_option_root:
            inp -= 'depthmap_gen_row_3', cur_option_root
            with gr.Row():
                inp += go.CLIPDEPTH, gr.Checkbox(label="Clip and renormalize DepthMap")
                inp += go.CLIPDEPTH_MODE,\
                    gr.Dropdown(label="Mode", choices=['Range', 'Outliers'], type="value", visible=False)
            with gr.Row(visible=False) as clip_options_row_1:
                inp += go.CLIPDEPTH_FAR, gr.Slider(minimum=0, maximum=1, step=0.001, label='Far clip')
                inp += go.CLIPDEPTH_NEAR, gr.Slider(minimum=0, maximum=1, step=0.001, label='Near clip')

        with gr.Box():
            with gr.Row():
                inp += go.GEN_STEREO, gr.Checkbox(label="Generate stereoscopic (3D) image(s)")
            with gr.Column(visible=False) as stereo_options:
                with gr.Row():
                    inp += go.STEREO_MODES, gr.CheckboxGroup(
                        ["left-right", "right-left", "top-bottom", "bottom-top", "red-cyan-anaglyph",
                         "left-only", "only-right", "cyan-red-reverseanaglyph"
                         ][0:8 if backbone.get_opt('depthmap_script_extra_stereomodes', False) else 5], label="Output")
                with gr.Row():
                    inp += go.STEREO_DIVERGENCE, gr.Slider(minimum=0.05, maximum=15.005, step=0.01,
                                                          label='Divergence (3D effect)')
                    inp += go.STEREO_SEPARATION, gr.Slider(minimum=-5.0, maximum=5.0, step=0.01,
                                                          label='Separation (moves images apart)')
                with gr.Row():
                    inp += go.STEREO_FILL_ALGO, gr.Dropdown(label="Gap fill technique",
                                                      choices=['none', 'naive', 'naive_interpolating', 'polylines_soft',
                                                               'polylines_sharp'],
                                                      type="value")
                    inp += go.STEREO_OFFSET_EXPONENT, gr.Slider(label="Magic exponent", minimum=1, maximum=2, step=1)
                    inp += go.STEREO_BALANCE, gr.Slider(minimum=-1.0, maximum=1.0, step=0.05,
                                                       label='Balance between eyes')

        with gr.Box():
            with gr.Row():
                inp += go.GEN_NORMALMAP, gr.Checkbox(label="Generate NormalMap")
            with gr.Column(visible=False) as normalmap_options:
                with gr.Row():
                    inp += go.NORMALMAP_PRE_BLUR, gr.Checkbox(label="Smooth before calculating normals")
                    inp += go.NORMALMAP_PRE_BLUR_KERNEL, gr.Slider(minimum=1, maximum=31, step=2, label='Pre-smooth kernel size', visible=False)
                    inp.add_rule(go.NORMALMAP_PRE_BLUR_KERNEL, 'visible-if', go.NORMALMAP_PRE_BLUR)
                with gr.Row():
                    inp += go.NORMALMAP_SOBEL, gr.Checkbox(label="Sobel gradient")
                    inp += go.NORMALMAP_SOBEL_KERNEL, gr.Slider(minimum=1, maximum=31, step=2, label='Sobel kernel size')
                    inp.add_rule(go.NORMALMAP_SOBEL_KERNEL, 'visible-if', go.NORMALMAP_SOBEL)
                with gr.Row():
                    inp += go.NORMALMAP_POST_BLUR, gr.Checkbox(label="Smooth after calculating normals")
                    inp += go.NORMALMAP_POST_BLUR_KERNEL, gr.Slider(minimum=1, maximum=31, step=2, label='Post-smooth kernel size', visible=False)
                    inp.add_rule(go.NORMALMAP_POST_BLUR_KERNEL, 'visible-if', go.NORMALMAP_POST_BLUR)
                with gr.Row():
                    inp += go.NORMALMAP_INVERT, gr.Checkbox(label="Invert")

        if backbone.get_opt('depthmap_script_gen_heatmap_from_ui', False):
            with gr.Box():
                with gr.Row():
                    inp += go.GEN_HEATMAP, gr.Checkbox(label="Generate HeatMap")

        with gr.Box():
            with gr.Column():
                inp += go.GEN_SIMPLE_MESH, gr.Checkbox(label="Generate simple 3D mesh")
            with gr.Column(visible=False) as mesh_options:
                with gr.Row():
                    gr.HTML(value="Generates fast, accurate only with ZoeDepth models and no boost, no custom maps.")
                with gr.Row():
                    inp += go.SIMPLE_MESH_OCCLUDE, gr.Checkbox(label="Remove occluded edges")
                    inp += go.SIMPLE_MESH_SPHERICAL, gr.Checkbox(label="Equirectangular projection")

        if is_depth_tab:
            with gr.Box():
                with gr.Column():
                    inp += go.GEN_INPAINTED_MESH, gr.Checkbox(
                        label="Generate 3D inpainted mesh")
                with gr.Column(visible=False) as inpaint_options_row_0:
                    gr.HTML("Generation is sloooow. Required for generating videos from mesh.")
                    inp += go.GEN_INPAINTED_MESH_DEMOS, gr.Checkbox(
                        label="Generate 4 demo videos with 3D inpainted mesh.")
                    gr.HTML("More options for generating video can be found in the Generate video tab.")

        with gr.Box():
            # 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 += go.GEN_REMBG, gr.Checkbox(label="Remove background")
            with gr.Column(visible=False) as bgrem_options:
                with gr.Row():
                    inp += go.SAVE_BACKGROUND_REMOVAL_MASKS, gr.Checkbox(label="Save the foreground masks")
                    inp += go.PRE_DEPTH_BACKGROUND_REMOVAL, gr.Checkbox(label="Pre-depth background removal")
                with gr.Row():
                    inp += go.REMBG_MODEL, gr.Dropdown(
                        label="Rembg Model", type="value",
                        choices=['u2net', 'u2netp', 'u2net_human_seg', 'silueta', "isnet-general-use", "isnet-anime"])

        with gr.Box():
            gr.HTML(f"{SCRIPT_FULL_NAME}<br/>")
            gr.HTML("Information, comment and share @ <a "
                    "href='https://github.com/thygate/stable-diffusion-webui-depthmap-script'>"
                    "https://github.com/thygate/stable-diffusion-webui-depthmap-script</a>")

        def update_default_net_size(model_type):
            w, h = ModelHolder.get_default_net_size(model_type)
            return inp[go.NET_WIDTH].update(value=w), inp[go.NET_HEIGHT].update(value=h)

        inp[go.MODEL_TYPE].change(
            fn=update_default_net_size,
            inputs=inp[go.MODEL_TYPE],
            outputs=[inp[go.NET_WIDTH], inp[go.NET_HEIGHT]]
        )

        inp[go.BOOST].change(  # Go boost! Wroom!..
            fn=lambda a, b: (inp[go.NET_SIZE_MATCH].update(visible=not a),
                             options_depend_on_match_size.update(visible=not a and not b)),
            inputs=[inp[go.BOOST], inp[go.NET_SIZE_MATCH]],
            outputs=[inp[go.NET_SIZE_MATCH], options_depend_on_match_size]
        )
        inp.add_rule(options_depend_on_match_size, 'visible-if-not', go.NET_SIZE_MATCH)
        inp[go.TILING_MODE].change(  # Go boost! Wroom!..
            fn=lambda a: (
                inp[go.BOOST].update(value=False), inp[go.NET_SIZE_MATCH].update(value=True)
            ) if a else (inp[go.BOOST].update(), inp[go.NET_SIZE_MATCH].update()),
            inputs=[inp[go.TILING_MODE]],
            outputs=[inp[go.BOOST], inp[go.NET_SIZE_MATCH]]
        )

        inp.add_rule(options_depend_on_output_depth_1, 'visible-if', go.DO_OUTPUT_DEPTH)
        inp.add_rule(go.OUTPUT_DEPTH_INVERT, 'visible-if', go.DO_OUTPUT_DEPTH)
        inp.add_rule(go.OUTPUT_DEPTH_COMBINE_AXIS, 'visible-if', go.OUTPUT_DEPTH_COMBINE)
        inp.add_rule(go.CLIPDEPTH_MODE, 'visible-if', go.CLIPDEPTH)
        inp.add_rule(clip_options_row_1, 'visible-if', go.CLIPDEPTH)

        inp[go.CLIPDEPTH_FAR].change(
            fn=lambda a, b: a if b < a else b,
            inputs=[inp[go.CLIPDEPTH_FAR], inp[go.CLIPDEPTH_NEAR]],
            outputs=[inp[go.CLIPDEPTH_NEAR]],
            show_progress=False
        )
        inp[go.CLIPDEPTH_NEAR].change(
            fn=lambda a, b: a if b > a else b,
            inputs=[inp[go.CLIPDEPTH_NEAR], inp[go.CLIPDEPTH_FAR]],
            outputs=[inp[go.CLIPDEPTH_FAR]],
            show_progress=False
        )

        inp.add_rule(stereo_options, 'visible-if', go.GEN_STEREO)
        inp.add_rule(normalmap_options, 'visible-if', go.GEN_NORMALMAP)
        inp.add_rule(mesh_options, 'visible-if', go.GEN_SIMPLE_MESH)
        if is_depth_tab:
            inp.add_rule(inpaint_options_row_0, 'visible-if', go.GEN_INPAINTED_MESH)
        inp.add_rule(bgrem_options, 'visible-if', go.GEN_REMBG)

    return inp

def open_folder_action():
    # Adapted from stable-diffusion-webui
    f = backbone.get_outpath()
    if backbone.get_cmd_opt('hide_ui_dir_config', False):
        return
    if not os.path.exists(f) or not os.path.isdir(f):
        raise Exception("Couldn't open output folder")  # .isdir is security-related, do not remove!
    import platform
    import subprocess as sp
    path = os.path.normpath(f)
    if platform.system() == "Windows":
        os.startfile(path)
    elif platform.system() == "Darwin":
        sp.Popen(["open", path])
    elif "microsoft-standard-WSL2" in platform.uname().release:
        sp.Popen(["wsl-open", path])
    else:
        sp.Popen(["xdg-open", path])


def depthmap_mode_video(inp):
    gr.HTML(value="Single video mode allows generating videos from videos. Please "
                  "keep in mind that all the frames of the video need to be processed - therefore it is important to "
                  "pick settings so that the generation is not too slow. For the best results, "
                  "use a zoedepth model, since they provide the highest level of coherency between frames.")
    inp += gr.File(elem_id='depthmap_vm_input', label="Video or animated file",
                   file_count="single", interactive=True, type="file")
    inp += gr.Checkbox(elem_id="depthmap_vm_custom_checkbox",
                       label="Use custom/pregenerated DepthMap video", value=False)
    inp += gr.Dropdown(elem_id="depthmap_vm_smoothening_mode", label="Smoothening",
                       type="value", choices=['none', 'experimental'], value='experimental')
    inp += gr.File(elem_id='depthmap_vm_custom', file_count="single",
                   interactive=True, type="file", visible=False)
    with gr.Row():
        inp += gr.Checkbox(elem_id='depthmap_vm_compress_checkbox', label="Compress colorvideos?", value=False)
        inp += gr.Slider(elem_id='depthmap_vm_compress_bitrate', label="Bitrate (kbit)", visible=False,
                         minimum=1000, value=15000, maximum=50000, step=250)

    inp.add_rule('depthmap_vm_custom', 'visible-if', 'depthmap_vm_custom_checkbox')
    inp.add_rule('depthmap_vm_smoothening_mode', 'visible-if-not', 'depthmap_vm_custom_checkbox')
    inp.add_rule('depthmap_vm_compress_bitrate', 'visible-if', 'depthmap_vm_compress_checkbox')

    return inp


custom_css = """
#depthmap_vm_input {height: 75px}
#depthmap_vm_custom {height: 75px}
"""


def on_ui_tabs():
    inp = GradioComponentBundle()
    with gr.Blocks(analytics_enabled=False, title="DepthMap", css=custom_css) as depthmap_interface:
        with gr.Row(equal_height=False):
            with gr.Column(variant='panel'):
                inp += 'depthmap_mode', gr.HTML(visible=False, value='0')
                with gr.Tabs():
                    with gr.TabItem('Single Image') as depthmap_mode_0:
                        with gr.Group():
                            with gr.Row():
                                inp += gr.Image(label="Source", source="upload", interactive=True, type="pil",
                                                elem_id="depthmap_input_image")
                                # TODO: depthmap generation settings should disappear when using this
                                inp += gr.File(label="Custom DepthMap", file_count="single", interactive=True,
                                               type="file", elem_id='custom_depthmap_img', visible=False)
                        inp += gr.Checkbox(elem_id="custom_depthmap", label="Use custom DepthMap", value=False)
                    with gr.TabItem('Batch Process') as depthmap_mode_1:
                        inp += gr.File(elem_id='image_batch', label="Batch Process", file_count="multiple",
                                       interactive=True, type="file")
                    with gr.TabItem('Batch from Directory') as depthmap_mode_2:
                        inp += gr.Textbox(elem_id="depthmap_batch_input_dir", label="Input directory",
                                          **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(),
                                          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",
                                           label="Skip generation and use (edited/custom) depthmaps "
                                                 "in output directory when a file already exists.",
                                           value=True)
                    with gr.TabItem('Single Video') as depthmap_mode_3:
                        inp = depthmap_mode_video(inp)
                submit = gr.Button('Generate', elem_id="depthmap_generate", variant='primary')
                inp |= main_ui_panel(True)  # Main panel is inserted here
                unloadmodels = gr.Button('Unload models', elem_id="depthmap_unloadmodels")

            with gr.Column(variant='panel'):
                with gr.Tabs(elem_id="mode_depthmap_output"):
                    with gr.TabItem('Depth Output'):
                        with gr.Group():
                            result_images = gr.Gallery(label='Output', show_label=False,
                                                       elem_id=f"depthmap_gallery", columns=4)
                        with gr.Column():
                            html_info = gr.HTML()
                        folder_symbol = '\U0001f4c2'  # 📂
                        gr.Button(folder_symbol, visible=not backbone.get_cmd_opt('hide_ui_dir_config', False)).click(
                            fn=lambda: open_folder_action(), inputs=[], outputs=[],
                        )

                    with gr.TabItem('3D Mesh'):
                        with gr.Group():
                            result_depthmesh = gr.Model3D(label="3d Mesh", clear_color=[1.0, 1.0, 1.0, 1.0])
                            with gr.Row():
                                # loadmesh = gr.Button('Load')
                                clearmesh = gr.Button('Clear')

                    with gr.TabItem('Generate video'):
                        # generate video
                        with gr.Group():
                            with gr.Row():
                                gr.Markdown("Generate video from inpainted(!) mesh.")
                            with gr.Row():
                                depth_vid = gr.Video(interactive=False)
                            with gr.Column():
                                vid_html_info_x = gr.HTML()
                                vid_html_info = gr.HTML()
                                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():
                                vid_numframes = gr.Textbox(label="Number of frames", value="300")
                                vid_fps = gr.Textbox(label="Framerate", value="40")
                                vid_format = gr.Dropdown(label="Format", choices=['mp4', 'webm'], value='mp4',
                                                         type="value", elem_id="video_format")
                                vid_ssaa = gr.Dropdown(label="SSAA", choices=['1', '2', '3', '4'], value='3',
                                                       type="value", elem_id="video_ssaa")
                            with gr.Row():
                                vid_traj = gr.Dropdown(label="Trajectory",
                                                       choices=['straight-line', 'double-straight-line', 'circle'],
                                                       value='double-straight-line', type="index",
                                                       elem_id="video_trajectory")
                                vid_shift = gr.Textbox(label="Translate: x, y, z", value="-0.015, 0.0, -0.05")
                                vid_border = gr.Textbox(label="Crop: top, left, bottom, right",
                                                        value="0.03, 0.03, 0.05, 0.03")
                                vid_dolly = gr.Checkbox(label="Dolly", value=False, elem_classes="smalltxt")
                            with gr.Row():
                                submit_vid = gr.Button('Generate Video', elem_id="depthmap_generatevideo",
                                                       variant='primary')

        inp += inp.enkey_tail()

        depthmap_mode_0.select(lambda: '0', None, inp['depthmap_mode'])
        depthmap_mode_1.select(lambda: '1', None, inp['depthmap_mode'])
        depthmap_mode_2.select(lambda: '2', None, inp['depthmap_mode'])
        depthmap_mode_3.select(lambda: '3', None, inp['depthmap_mode'])

        def custom_depthmap_change_fn(mode, zero_on, three_on):
            hide = mode == '0' and zero_on or mode == '3' and three_on
            return inp['custom_depthmap_img'].update(visible=hide), \
                inp['depthmap_gen_row_0'].update(visible=not hide), \
                inp['depthmap_gen_row_1'].update(visible=not hide), \
                inp['depthmap_gen_row_3'].update(visible=not hide), not hide
        custom_depthmap_change_els = ['depthmap_mode', 'custom_depthmap', 'depthmap_vm_custom_checkbox']
        for el in custom_depthmap_change_els:
            inp[el].change(
            fn=custom_depthmap_change_fn,
            inputs=[inp[el] for el in custom_depthmap_change_els],
            outputs=[inp[st] for st in [
                'custom_depthmap_img', 'depthmap_gen_row_0', 'depthmap_gen_row_1', 'depthmap_gen_row_3',
                go.DO_OUTPUT_DEPTH]])

        unloadmodels.click(
            fn=unload_models,
            inputs=[],
            outputs=[]
        )

        clearmesh.click(
            fn=lambda: None,
            inputs=[],
            outputs=[result_depthmesh]
        )

        submit.click(
            fn=backbone.wrap_gradio_gpu_call(run_generate),
            inputs=inp.enkey_body(),
            outputs=[
                result_images,
                fn_mesh,
                result_depthmesh,
                html_info
            ]
        )

        submit_vid.click(
            fn=backbone.wrap_gradio_gpu_call(run_makevideo),
            inputs=[
                fn_mesh,
                vid_numframes,
                vid_fps,
                vid_traj,
                vid_shift,
                vid_border,
                vid_dolly,
                vid_format,
                vid_ssaa
            ],
            outputs=[
                depth_vid,
                vid_html_info_x,
                vid_html_info
            ]
        )

    return depthmap_interface


def format_exception(e: Exception):
    traceback.print_exc()
    msg = '<h3>' + 'ERROR: ' + str(e) + '</h3>' + '\n'
    if 'out of GPU memory' in msg:
        pass
    elif "torch.hub.load('facebookresearch/dinov2'," in traceback.format_exc():
        msg += ('<h4>To use Depth Anything integration in WebUI mode, please add "--disable-safe-unpickle" to the command line flags. '
                'Alternatively, use Standalone mode. This is a known issue.')
    elif "Error(s) in loading state_dict " in traceback.format_exc():
        msg += ('<h4>There was issue during loading the model.'
                'Please add "--disable-safe-unpickle" to the command line flags. This is a known issue.')
    elif 'out of GPU memory' not in msg:
        msg += \
            'Please report this issue ' \
            f'<a href="https://github.com/thygate/{REPOSITORY_NAME}/issues">here</a>. ' \
            'Make sure to provide the full stacktrace: \n'
        msg += '<code style="white-space: pre;">' + traceback.format_exc() + '</code>'
    return msg


def run_generate(*inputs):
    inputs = GradioComponentBundle.enkey_to_dict(inputs)
    depthmap_mode = inputs['depthmap_mode']
    depthmap_batch_input_dir = inputs['depthmap_batch_input_dir']
    image_batch = inputs['image_batch']
    depthmap_input_image = inputs['depthmap_input_image']
    depthmap_batch_output_dir = inputs['depthmap_batch_output_dir']
    depthmap_batch_reuse = inputs['depthmap_batch_reuse']
    custom_depthmap = inputs['custom_depthmap']
    custom_depthmap_img = inputs['custom_depthmap_img']

    inputimages = []
    inputdepthmaps = []  # Allow supplying custom depthmaps
    inputnames = []  # Also keep track of original file names

    if depthmap_mode == '3':
        try:
            custom_depthmap = inputs['depthmap_vm_custom'] \
                if inputs['depthmap_vm_custom_checkbox'] else None
            colorvids_bitrate = inputs['depthmap_vm_compress_bitrate'] \
                if inputs['depthmap_vm_compress_checkbox'] else None
            ret = video_mode.gen_video(
                inputs['depthmap_vm_input'], backbone.get_outpath(), inputs, custom_depthmap, colorvids_bitrate,
                inputs['depthmap_vm_smoothening_mode'])
            return [], None, None, ret
        except Exception as e:
            ret = format_exception(e)
        return [], None, None, ret

    if depthmap_mode == '2' and depthmap_batch_output_dir != '':
        outpath = depthmap_batch_output_dir
    else:
        outpath = backbone.get_outpath()

    if depthmap_mode == '0':  # Single image
        if depthmap_input_image is None:
            return [], None, None, "Please select an input image"
        inputimages.append(depthmap_input_image)
        inputnames.append(None)
        if custom_depthmap:
            if custom_depthmap_img is None:
                return [], None, None, \
                    "Custom depthmap is not specified. Please either supply it or disable this option."
            inputdepthmaps.append(Image.open(os.path.abspath(custom_depthmap_img.name)))
        else:
            inputdepthmaps.append(None)
    if depthmap_mode == '1':  # Batch Process
        if image_batch is None:
            return [], None, None, "Please select input images", ""
        for img in image_batch:
            image = Image.open(os.path.abspath(img.name))
            inputimages.append(image)
            inputnames.append(os.path.splitext(img.orig_name)[0])
        print(f'{len(inputimages)} images will be processed')
    elif depthmap_mode == '2':  # Batch from Directory
        # TODO: There is a RAM leak when we process batches, I can smell it! Or maybe it is gone.
        assert not backbone.get_cmd_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:
            return [], None, None, "Please pick different directories for batch processing."
        image_list = backbone.listfiles(depthmap_batch_input_dir)
        for path in image_list:
            try:
                inputimages.append(Image.open(path))
                inputnames.append(path)

                custom_depthmap = None
                if depthmap_batch_reuse:
                    basename = Path(path).stem
                    # Custom names are not used in samples directory
                    if outpath != backbone.get_opt('outdir_extras_samples', None):
                        # Possible filenames that the custom depthmaps may have
                        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:
                            path_cand = os.path.join(outpath, fn_cand)
                            if os.path.isfile(path_cand):
                                custom_depthmap = Image.open(os.path.abspath(path_cand))
                                break
                inputdepthmaps.append(custom_depthmap)
            except Exception as e:
                print(f'Failed to load {path}, ignoring. Exception: {str(e)}')
        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')

    gen_obj = core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inputs, backbone.gather_ops())

    # Saving images
    img_results = []
    results_total = 0
    inpainted_mesh_fi = mesh_simple_fi = None
    msg = ""  # Empty string is never returned
    while True:
        try:
            input_i, type, result = next(gen_obj)
            results_total += 1
        except StopIteration:
            # TODO: return more info
            msg = '<h3>Successfully generated</h3>' if results_total > 0 else \
                '<h3>Successfully generated nothing - please check the settings and try again</h3>'
            break
        except Exception as e:
            msg = format_exception(e)
            break
        if type == 'simple_mesh':
            mesh_simple_fi = result
            continue
        if type == 'inpainted_mesh':
            inpainted_mesh_fi = result
            continue
        if not isinstance(result, Image.Image):
            print(f'This is not supposed to happen! Somehow output type {type} is not supported! Input_i: {input_i}.')
            continue
        img_results += [(input_i, type, result)]

        if inputs["save_outputs"]:
            try:
                basename = 'depthmap'
                if depthmap_mode == '2' and inputnames[input_i] is not None:
                    if outpath != backbone.get_opt('outdir_extras_samples', None):
                        basename = Path(inputnames[input_i]).stem
                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 backbone.get_opt('depthmap_script_show_3d', True):
        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
    return map(lambda x: x[2], img_results), inpainted_mesh_fi, display_mesh_fi, msg.replace('\n', '<br>')