hololens/stable-diffusion-webui-depthmap-script
1
1import traceback
2from pathlib import Path
3import gradio as gr
4from PIL import Image
5
6from src import backbone, video_mode
7from src.core import core_generation_funnel, unload_models, run_makevideo
8from src.depthmap_generation import ModelHolder
9from src.gradio_args_transport import GradioComponentBundle
10from src.misc import *
11from src.common_constants import GenerationOptions as go
12
13# Ugly workaround to fix gradio tempfile issue
14def ensure_gradio_temp_directory():
15 try:
16 import tempfile
17 path = os.path.join(tempfile.gettempdir(), 'gradio')
18 if not (os.path.exists(path)):
19 os.mkdir(path)
20 except Exception as e:
21 traceback.print_exc()
22
23
24ensure_gradio_temp_directory()
25
26
27def main_ui_panel(is_depth_tab):
28 inp = GradioComponentBundle()
29 # TODO: Greater visual separation
30 with gr.Blocks():
31 with gr.Row() as cur_option_root:
32 inp -= 'depthmap_gen_row_0', cur_option_root
33 inp += go.COMPUTE_DEVICE, gr.Radio(label="Compute on", choices=['GPU', 'CPU'], value='GPU')
34 # TODO: Should return value instead of index. Maybe Enum should be used?
35 inp += go.MODEL_TYPE, gr.Dropdown(label="Model",
36 choices=['res101', 'dpt_beit_large_512 (midas 3.1)',
37 'dpt_beit_large_384 (midas 3.1)', 'dpt_large_384 (midas 3.0)',
38 'dpt_hybrid_384 (midas 3.0)',
39 'midas_v21', 'midas_v21_small',
40 'zoedepth_n (indoor)', 'zoedepth_k (outdoor)', 'zoedepth_nk',
41 'Marigold v1', 'Depth Anything', 'Depth Anything v2 Small',
42 'Depth Anything v2 Base', 'Depth Anything v2 Large'],
43 value='Depth Anything v2 Base', type="index")
44 with gr.Box() as cur_option_root:
45 inp -= 'depthmap_gen_row_1', cur_option_root
46 with gr.Row():
47 inp += go.BOOST, gr.Checkbox(label="BOOST",
48 info="Generate depth map parts in a mosaic fashion - very slow",
49 value=False)
50 inp += go.NET_SIZE_MATCH, gr.Checkbox(label="Match net size to input size",
51 info="Net size affects quality, performance and VRAM usage")
52 with gr.Row() as options_depend_on_match_size:
53 inp += go.NET_WIDTH, gr.Slider(minimum=64, maximum=2048, step=64, label='Net width')
54 inp += go.NET_HEIGHT, gr.Slider(minimum=64, maximum=2048, step=64, label='Net height')
55 with gr.Row():
56 inp += go.TILING_MODE, gr.Checkbox(
57 label='Tiling mode', info='Reduces seams that appear if the depthmap is tiled into a grid'
58 )
59
60 with gr.Box() as cur_option_root:
61 inp -= 'depthmap_gen_row_2', cur_option_root
62 with gr.Row():
63 with gr.Group(): # 50% of width
64 inp += "save_outputs", gr.Checkbox(label="Save Outputs", value=True)
65 with gr.Group(): # 50% of width
66 inp += go.DO_OUTPUT_DEPTH, gr.Checkbox(label="Output DepthMap")
67 inp += go.OUTPUT_DEPTH_INVERT, gr.Checkbox(label="Invert (black=near, white=far)")
68 with gr.Row() as options_depend_on_output_depth_1:
69 inp += go.OUTPUT_DEPTH_COMBINE, gr.Checkbox(
70 label="Combine input and depthmap into one image")
71 inp += go.OUTPUT_DEPTH_COMBINE_AXIS, gr.Radio(
72 label="Combine axis", choices=['Vertical', 'Horizontal'], type="value", visible=False)
73
74 with gr.Box() as cur_option_root:
75 inp -= 'depthmap_gen_row_3', cur_option_root
76 with gr.Row():
77 inp += go.CLIPDEPTH, gr.Checkbox(label="Clip and renormalize DepthMap")
78 inp += go.CLIPDEPTH_MODE,\
79 gr.Dropdown(label="Mode", choices=['Range', 'Outliers'], type="value", visible=False)
80 with gr.Row(visible=False) as clip_options_row_1:
81 inp += go.CLIPDEPTH_FAR, gr.Slider(minimum=0, maximum=1, step=0.001, label='Far clip')
82 inp += go.CLIPDEPTH_NEAR, gr.Slider(minimum=0, maximum=1, step=0.001, label='Near clip')
83
84 with gr.Box():
85 with gr.Row():
86 inp += go.GEN_STEREO, gr.Checkbox(label="Generate stereoscopic (3D) image(s)")
87 with gr.Column(visible=False) as stereo_options:
88 with gr.Row():
89 inp += go.STEREO_MODES, gr.CheckboxGroup(
90 ["left-right", "right-left", "top-bottom", "bottom-top", "red-cyan-anaglyph",
91 "left-only", "only-right", "cyan-red-reverseanaglyph"
92 ][0:8 if backbone.get_opt('depthmap_script_extra_stereomodes', False) else 5], label="Output")
93 with gr.Row():
94 inp += go.STEREO_DIVERGENCE, gr.Slider(minimum=0.05, maximum=15.005, step=0.01,
95 label='Divergence (3D effect)')
96 inp += go.STEREO_SEPARATION, gr.Slider(minimum=-5.0, maximum=5.0, step=0.01,
97 label='Separation (moves images apart)')
98 with gr.Row():
99 inp += go.STEREO_FILL_ALGO, gr.Dropdown(label="Gap fill technique",
100 choices=['none', 'naive', 'naive_interpolating', 'polylines_soft',
101 'polylines_sharp'],
102 type="value")
103 inp += go.STEREO_OFFSET_EXPONENT, gr.Slider(label="Magic exponent", minimum=1, maximum=2, step=1)
104 inp += go.STEREO_BALANCE, gr.Slider(minimum=-1.0, maximum=1.0, step=0.05,
105 label='Balance between eyes')
106
107 with gr.Box():
108 with gr.Row():
109 inp += go.GEN_NORMALMAP, gr.Checkbox(label="Generate NormalMap")
110 with gr.Column(visible=False) as normalmap_options:
111 with gr.Row():
112 inp += go.NORMALMAP_PRE_BLUR, gr.Checkbox(label="Smooth before calculating normals")
113 inp += go.NORMALMAP_PRE_BLUR_KERNEL, gr.Slider(minimum=1, maximum=31, step=2, label='Pre-smooth kernel size', visible=False)
114 inp.add_rule(go.NORMALMAP_PRE_BLUR_KERNEL, 'visible-if', go.NORMALMAP_PRE_BLUR)
115 with gr.Row():
116 inp += go.NORMALMAP_SOBEL, gr.Checkbox(label="Sobel gradient")
117 inp += go.NORMALMAP_SOBEL_KERNEL, gr.Slider(minimum=1, maximum=31, step=2, label='Sobel kernel size')
118 inp.add_rule(go.NORMALMAP_SOBEL_KERNEL, 'visible-if', go.NORMALMAP_SOBEL)
119 with gr.Row():
120 inp += go.NORMALMAP_POST_BLUR, gr.Checkbox(label="Smooth after calculating normals")
121 inp += go.NORMALMAP_POST_BLUR_KERNEL, gr.Slider(minimum=1, maximum=31, step=2, label='Post-smooth kernel size', visible=False)
122 inp.add_rule(go.NORMALMAP_POST_BLUR_KERNEL, 'visible-if', go.NORMALMAP_POST_BLUR)
123 with gr.Row():
124 inp += go.NORMALMAP_INVERT, gr.Checkbox(label="Invert")
125
126 if backbone.get_opt('depthmap_script_gen_heatmap_from_ui', False):
127 with gr.Box():
128 with gr.Row():
129 inp += go.GEN_HEATMAP, gr.Checkbox(label="Generate HeatMap")
130
131 with gr.Box():
132 with gr.Column():
133 inp += go.GEN_SIMPLE_MESH, gr.Checkbox(label="Generate simple 3D mesh")
134 with gr.Column(visible=False) as mesh_options:
135 with gr.Row():
136 gr.HTML(value="Generates fast, accurate only with ZoeDepth models and no boost, no custom maps.")
137 with gr.Row():
138 inp += go.SIMPLE_MESH_OCCLUDE, gr.Checkbox(label="Remove occluded edges")
139 inp += go.SIMPLE_MESH_SPHERICAL, gr.Checkbox(label="Equirectangular projection")
140
141 if is_depth_tab:
142 with gr.Box():
143 with gr.Column():
144 inp += go.GEN_INPAINTED_MESH, gr.Checkbox(
145 label="Generate 3D inpainted mesh")
146 with gr.Column(visible=False) as inpaint_options_row_0:
147 gr.HTML("Generation is sloooow. Required for generating videos from mesh.")
148 inp += go.GEN_INPAINTED_MESH_DEMOS, gr.Checkbox(
149 label="Generate 4 demo videos with 3D inpainted mesh.")
150 gr.HTML("More options for generating video can be found in the Generate video tab.")
151
152 with gr.Box():
153 # TODO: it should be clear from the UI that there is an option of the background removal
154 # that does not use the model selected above
155 with gr.Row():
156 inp += go.GEN_REMBG, gr.Checkbox(label="Remove background")
157 with gr.Column(visible=False) as bgrem_options:
158 with gr.Row():
159 inp += go.SAVE_BACKGROUND_REMOVAL_MASKS, gr.Checkbox(label="Save the foreground masks")
160 inp += go.PRE_DEPTH_BACKGROUND_REMOVAL, gr.Checkbox(label="Pre-depth background removal")
161 with gr.Row():
162 inp += go.REMBG_MODEL, gr.Dropdown(
163 label="Rembg Model", type="value",
164 choices=['u2net', 'u2netp', 'u2net_human_seg', 'silueta', "isnet-general-use", "isnet-anime"])
165
166 with gr.Box():
167 gr.HTML(f"{SCRIPT_FULL_NAME}<br/>")
168 gr.HTML("Information, comment and share @ <a "
169 "href='https://github.com/thygate/stable-diffusion-webui-depthmap-script'>"
170 "https://github.com/thygate/stable-diffusion-webui-depthmap-script</a>")
171
172 def update_default_net_size(model_type):
173 w, h = ModelHolder.get_default_net_size(model_type)
174 return inp[go.NET_WIDTH].update(value=w), inp[go.NET_HEIGHT].update(value=h)
175
176 inp[go.MODEL_TYPE].change(
177 fn=update_default_net_size,
178 inputs=inp[go.MODEL_TYPE],
179 outputs=[inp[go.NET_WIDTH], inp[go.NET_HEIGHT]]
180 )
181
182 inp[go.BOOST].change( # Go boost! Wroom!..
183 fn=lambda a, b: (inp[go.NET_SIZE_MATCH].update(visible=not a),
184 options_depend_on_match_size.update(visible=not a and not b)),
185 inputs=[inp[go.BOOST], inp[go.NET_SIZE_MATCH]],
186 outputs=[inp[go.NET_SIZE_MATCH], options_depend_on_match_size]
187 )
188 inp.add_rule(options_depend_on_match_size, 'visible-if-not', go.NET_SIZE_MATCH)
189 inp[go.TILING_MODE].change( # Go boost! Wroom!..
190 fn=lambda a: (
191 inp[go.BOOST].update(value=False), inp[go.NET_SIZE_MATCH].update(value=True)
192 ) if a else (inp[go.BOOST].update(), inp[go.NET_SIZE_MATCH].update()),
193 inputs=[inp[go.TILING_MODE]],
194 outputs=[inp[go.BOOST], inp[go.NET_SIZE_MATCH]]
195 )
196
197 inp.add_rule(options_depend_on_output_depth_1, 'visible-if', go.DO_OUTPUT_DEPTH)
198 inp.add_rule(go.OUTPUT_DEPTH_INVERT, 'visible-if', go.DO_OUTPUT_DEPTH)
199 inp.add_rule(go.OUTPUT_DEPTH_COMBINE_AXIS, 'visible-if', go.OUTPUT_DEPTH_COMBINE)
200 inp.add_rule(go.CLIPDEPTH_MODE, 'visible-if', go.CLIPDEPTH)
201 inp.add_rule(clip_options_row_1, 'visible-if', go.CLIPDEPTH)
202
203 inp[go.CLIPDEPTH_FAR].change(
204 fn=lambda a, b: a if b < a else b,
205 inputs=[inp[go.CLIPDEPTH_FAR], inp[go.CLIPDEPTH_NEAR]],
206 outputs=[inp[go.CLIPDEPTH_NEAR]],
207 show_progress=False
208 )
209 inp[go.CLIPDEPTH_NEAR].change(
210 fn=lambda a, b: a if b > a else b,
211 inputs=[inp[go.CLIPDEPTH_NEAR], inp[go.CLIPDEPTH_FAR]],
212 outputs=[inp[go.CLIPDEPTH_FAR]],
213 show_progress=False
214 )
215
216 inp.add_rule(stereo_options, 'visible-if', go.GEN_STEREO)
217 inp.add_rule(normalmap_options, 'visible-if', go.GEN_NORMALMAP)
218 inp.add_rule(mesh_options, 'visible-if', go.GEN_SIMPLE_MESH)
219 if is_depth_tab:
220 inp.add_rule(inpaint_options_row_0, 'visible-if', go.GEN_INPAINTED_MESH)
221 inp.add_rule(bgrem_options, 'visible-if', go.GEN_REMBG)
222
223 return inp
224
225def open_folder_action():
226 # Adapted from stable-diffusion-webui
227 f = backbone.get_outpath()
228 if backbone.get_cmd_opt('hide_ui_dir_config', False):
229 return
230 if not os.path.exists(f) or not os.path.isdir(f):
231 raise Exception("Couldn't open output folder") # .isdir is security-related, do not remove!
232 import platform
233 import subprocess as sp
234 path = os.path.normpath(f)
235 if platform.system() == "Windows":
236 os.startfile(path)
237 elif platform.system() == "Darwin":
238 sp.Popen(["open", path])
239 elif "microsoft-standard-WSL2" in platform.uname().release:
240 sp.Popen(["wsl-open", path])
241 else:
242 sp.Popen(["xdg-open", path])
243
244
245def depthmap_mode_video(inp):
246 gr.HTML(value="Single video mode allows generating videos from videos. Please "
247 "keep in mind that all the frames of the video need to be processed - therefore it is important to "
248 "pick settings so that the generation is not too slow. For the best results, "
249 "use a zoedepth model, since they provide the highest level of coherency between frames.")
250 inp += gr.File(elem_id='depthmap_vm_input', label="Video or animated file",
251 file_count="single", interactive=True, type="file")
252 inp += gr.Checkbox(elem_id="depthmap_vm_custom_checkbox",
253 label="Use custom/pregenerated DepthMap video", value=False)
254 inp += gr.Dropdown(elem_id="depthmap_vm_smoothening_mode", label="Smoothening",
255 type="value", choices=['none', 'experimental'], value='experimental')
256 inp += gr.File(elem_id='depthmap_vm_custom', file_count="single",
257 interactive=True, type="file", visible=False)
258 with gr.Row():
259 inp += gr.Checkbox(elem_id='depthmap_vm_compress_checkbox', label="Compress colorvideos?", value=False)
260 inp += gr.Slider(elem_id='depthmap_vm_compress_bitrate', label="Bitrate (kbit)", visible=False,
261 minimum=1000, value=15000, maximum=50000, step=250)
262
263 inp.add_rule('depthmap_vm_custom', 'visible-if', 'depthmap_vm_custom_checkbox')
264 inp.add_rule('depthmap_vm_smoothening_mode', 'visible-if-not', 'depthmap_vm_custom_checkbox')
265 inp.add_rule('depthmap_vm_compress_bitrate', 'visible-if', 'depthmap_vm_compress_checkbox')
266
267 return inp
268
269
270custom_css = """
271#depthmap_vm_input {height: 75px}
272#depthmap_vm_custom {height: 75px}
273"""
274
275
276def on_ui_tabs():
277 inp = GradioComponentBundle()
278 with gr.Blocks(analytics_enabled=False, title="DepthMap", css=custom_css) as depthmap_interface:
279 with gr.Row(equal_height=False):
280 with gr.Column(variant='panel'):
281 inp += 'depthmap_mode', gr.HTML(visible=False, value='0')
282 with gr.Tabs():
283 with gr.TabItem('Single Image') as depthmap_mode_0:
284 with gr.Group():
285 with gr.Row():
286 inp += gr.Image(label="Source", source="upload", interactive=True, type="pil",
287 elem_id="depthmap_input_image")
288 # TODO: depthmap generation settings should disappear when using this
289 inp += gr.File(label="Custom DepthMap", file_count="single", interactive=True,
290 type="file", elem_id='custom_depthmap_img', visible=False)
291 inp += gr.Checkbox(elem_id="custom_depthmap", label="Use custom DepthMap", value=False)
292 with gr.TabItem('Batch Process') as depthmap_mode_1:
293 inp += gr.File(elem_id='image_batch', label="Batch Process", file_count="multiple",
294 interactive=True, type="file")
295 with gr.TabItem('Batch from Directory') as depthmap_mode_2:
296 inp += gr.Textbox(elem_id="depthmap_batch_input_dir", label="Input directory",
297 **backbone.get_hide_dirs(),
298 placeholder="A directory on the same machine where the server is running.")
299 inp += gr.Textbox(elem_id="depthmap_batch_output_dir", label="Output directory",
300 **backbone.get_hide_dirs(),
301 placeholder="Leave blank to save images to the default path.")
302 gr.HTML("Files in the output directory may be overwritten.")
303 inp += gr.Checkbox(elem_id="depthmap_batch_reuse",
304 label="Skip generation and use (edited/custom) depthmaps "
305 "in output directory when a file already exists.",
306 value=True)
307 with gr.TabItem('Single Video') as depthmap_mode_3:
308 inp = depthmap_mode_video(inp)
309 submit = gr.Button('Generate', elem_id="depthmap_generate", variant='primary')
310 inp |= main_ui_panel(True) # Main panel is inserted here
311 unloadmodels = gr.Button('Unload models', elem_id="depthmap_unloadmodels")
312
313 with gr.Column(variant='panel'):
314 with gr.Tabs(elem_id="mode_depthmap_output"):
315 with gr.TabItem('Depth Output'):
316 with gr.Group():
317 result_images = gr.Gallery(label='Output', show_label=False,
318 elem_id=f"depthmap_gallery", columns=4)
319 with gr.Column():
320 html_info = gr.HTML()
321 folder_symbol = '\U0001f4c2' # 📂
322 gr.Button(folder_symbol, visible=not backbone.get_cmd_opt('hide_ui_dir_config', False)).click(
323 fn=lambda: open_folder_action(), inputs=[], outputs=[],
324 )
325
326 with gr.TabItem('3D Mesh'):
327 with gr.Group():
328 result_depthmesh = gr.Model3D(label="3d Mesh", clear_color=[1.0, 1.0, 1.0, 1.0])
329 with gr.Row():
330 # loadmesh = gr.Button('Load')
331 clearmesh = gr.Button('Clear')
332
333 with gr.TabItem('Generate video'):
334 # generate video
335 with gr.Group():
336 with gr.Row():
337 gr.Markdown("Generate video from inpainted(!) mesh.")
338 with gr.Row():
339 depth_vid = gr.Video(interactive=False)
340 with gr.Column():
341 vid_html_info_x = gr.HTML()
342 vid_html_info = gr.HTML()
343 fn_mesh = gr.Textbox(label="Input Mesh (.ply | .obj)", **backbone.get_hide_dirs(),
344 placeholder="A file on the same machine where "
345 "the server is running.")
346 with gr.Row():
347 vid_numframes = gr.Textbox(label="Number of frames", value="300")
348 vid_fps = gr.Textbox(label="Framerate", value="40")
349 vid_format = gr.Dropdown(label="Format", choices=['mp4', 'webm'], value='mp4',
350 type="value", elem_id="video_format")
351 vid_ssaa = gr.Dropdown(label="SSAA", choices=['1', '2', '3', '4'], value='3',
352 type="value", elem_id="video_ssaa")
353 with gr.Row():
354 vid_traj = gr.Dropdown(label="Trajectory",
355 choices=['straight-line', 'double-straight-line', 'circle'],
356 value='double-straight-line', type="index",
357 elem_id="video_trajectory")
358 vid_shift = gr.Textbox(label="Translate: x, y, z", value="-0.015, 0.0, -0.05")
359 vid_border = gr.Textbox(label="Crop: top, left, bottom, right",
360 value="0.03, 0.03, 0.05, 0.03")
361 vid_dolly = gr.Checkbox(label="Dolly", value=False, elem_classes="smalltxt")
362 with gr.Row():
363 submit_vid = gr.Button('Generate Video', elem_id="depthmap_generatevideo",
364 variant='primary')
365
366 inp += inp.enkey_tail()
367
368 depthmap_mode_0.select(lambda: '0', None, inp['depthmap_mode'])
369 depthmap_mode_1.select(lambda: '1', None, inp['depthmap_mode'])
370 depthmap_mode_2.select(lambda: '2', None, inp['depthmap_mode'])
371 depthmap_mode_3.select(lambda: '3', None, inp['depthmap_mode'])
372
373 def custom_depthmap_change_fn(mode, zero_on, three_on):
374 hide = mode == '0' and zero_on or mode == '3' and three_on
375 return inp['custom_depthmap_img'].update(visible=hide), \
376 inp['depthmap_gen_row_0'].update(visible=not hide), \
377 inp['depthmap_gen_row_1'].update(visible=not hide), \
378 inp['depthmap_gen_row_3'].update(visible=not hide), not hide
379 custom_depthmap_change_els = ['depthmap_mode', 'custom_depthmap', 'depthmap_vm_custom_checkbox']
380 for el in custom_depthmap_change_els:
381 inp[el].change(
382 fn=custom_depthmap_change_fn,
383 inputs=[inp[el] for el in custom_depthmap_change_els],
384 outputs=[inp[st] for st in [
385 'custom_depthmap_img', 'depthmap_gen_row_0', 'depthmap_gen_row_1', 'depthmap_gen_row_3',
386 go.DO_OUTPUT_DEPTH]])
387
388 unloadmodels.click(
389 fn=unload_models,
390 inputs=[],
391 outputs=[]
392 )
393
394 clearmesh.click(
395 fn=lambda: None,
396 inputs=[],
397 outputs=[result_depthmesh]
398 )
399
400 submit.click(
401 fn=backbone.wrap_gradio_gpu_call(run_generate),
402 inputs=inp.enkey_body(),
403 outputs=[
404 result_images,
405 fn_mesh,
406 result_depthmesh,
407 html_info
408 ]
409 )
410
411 submit_vid.click(
412 fn=backbone.wrap_gradio_gpu_call(run_makevideo),
413 inputs=[
414 fn_mesh,
415 vid_numframes,
416 vid_fps,
417 vid_traj,
418 vid_shift,
419 vid_border,
420 vid_dolly,
421 vid_format,
422 vid_ssaa
423 ],
424 outputs=[
425 depth_vid,
426 vid_html_info_x,
427 vid_html_info
428 ]
429 )
430
431 return depthmap_interface
432
433
434def format_exception(e: Exception):
435 traceback.print_exc()
436 msg = '<h3>' + 'ERROR: ' + str(e) + '</h3>' + '\n'
437 if 'out of GPU memory' in msg:
438 pass
439 elif "torch.hub.load('facebookresearch/dinov2'," in traceback.format_exc():
440 msg += ('<h4>To use Depth Anything integration in WebUI mode, please add "--disable-safe-unpickle" to the command line flags. '
441 'Alternatively, use Standalone mode. This is a known issue.')
442 elif "Error(s) in loading state_dict " in traceback.format_exc():
443 msg += ('<h4>There was issue during loading the model.'
444 'Please add "--disable-safe-unpickle" to the command line flags. This is a known issue.')
445 elif 'out of GPU memory' not in msg:
446 msg += \
447 'Please report this issue ' \
448 f'<a href="https://github.com/thygate/{REPOSITORY_NAME}/issues">here</a>. ' \
449 'Make sure to provide the full stacktrace: \n'
450 msg += '<code style="white-space: pre;">' + traceback.format_exc() + '</code>'
451 return msg
452
453
454def run_generate(*inputs):
455 inputs = GradioComponentBundle.enkey_to_dict(inputs)
456 depthmap_mode = inputs['depthmap_mode']
457 depthmap_batch_input_dir = inputs['depthmap_batch_input_dir']
458 image_batch = inputs['image_batch']
459 depthmap_input_image = inputs['depthmap_input_image']
460 depthmap_batch_output_dir = inputs['depthmap_batch_output_dir']
461 depthmap_batch_reuse = inputs['depthmap_batch_reuse']
462 custom_depthmap = inputs['custom_depthmap']
463 custom_depthmap_img = inputs['custom_depthmap_img']
464
465 inputimages = []
466 inputdepthmaps = [] # Allow supplying custom depthmaps
467 inputnames = [] # Also keep track of original file names
468
469 if depthmap_mode == '3':
470 try:
471 custom_depthmap = inputs['depthmap_vm_custom'] \
472 if inputs['depthmap_vm_custom_checkbox'] else None
473 colorvids_bitrate = inputs['depthmap_vm_compress_bitrate'] \
474 if inputs['depthmap_vm_compress_checkbox'] else None
475 ret = video_mode.gen_video(
476 inputs['depthmap_vm_input'], backbone.get_outpath(), inputs, custom_depthmap, colorvids_bitrate,
477 inputs['depthmap_vm_smoothening_mode'])
478 return [], None, None, ret
479 except Exception as e:
480 ret = format_exception(e)
481 return [], None, None, ret
482
483 if depthmap_mode == '2' and depthmap_batch_output_dir != '':
484 outpath = depthmap_batch_output_dir
485 else:
486 outpath = backbone.get_outpath()
487
488 if depthmap_mode == '0': # Single image
489 if depthmap_input_image is None:
490 return [], None, None, "Please select an input image"
491 inputimages.append(depthmap_input_image)
492 inputnames.append(None)
493 if custom_depthmap:
494 if custom_depthmap_img is None:
495 return [], None, None, \
496 "Custom depthmap is not specified. Please either supply it or disable this option."
497 inputdepthmaps.append(Image.open(os.path.abspath(custom_depthmap_img.name)))
498 else:
499 inputdepthmaps.append(None)
500 if depthmap_mode == '1': # Batch Process
501 if image_batch is None:
502 return [], None, None, "Please select input images", ""
503 for img in image_batch:
504 image = Image.open(os.path.abspath(img.name))
505 inputimages.append(image)
506 inputnames.append(os.path.splitext(img.orig_name)[0])
507 print(f'{len(inputimages)} images will be processed')
508 elif depthmap_mode == '2': # Batch from Directory
509 # TODO: There is a RAM leak when we process batches, I can smell it! Or maybe it is gone.
510 assert not backbone.get_cmd_opt('hide_ui_dir_config', False), '--hide-ui-dir-config option must be disabled'
511 if depthmap_batch_input_dir == '':
512 return [], None, None, "Please select an input directory."
513 if depthmap_batch_input_dir == depthmap_batch_output_dir:
514 return [], None, None, "Please pick different directories for batch processing."
515 image_list = backbone.listfiles(depthmap_batch_input_dir)
516 for path in image_list:
517 try:
518 inputimages.append(Image.open(path))
519 inputnames.append(path)
520
521 custom_depthmap = None
522 if depthmap_batch_reuse:
523 basename = Path(path).stem
524 # Custom names are not used in samples directory
525 if outpath != backbone.get_opt('outdir_extras_samples', None):
526 # Possible filenames that the custom depthmaps may have
527 name_candidates = [f'{basename}-0000.{backbone.get_opt("samples_format", "png")}', # current format
528 f'{basename}.png', # human-intuitive format
529 f'{Path(path).name}'] # human-intuitive format (worse)
530 for fn_cand in name_candidates:
531 path_cand = os.path.join(outpath, fn_cand)
532 if os.path.isfile(path_cand):
533 custom_depthmap = Image.open(os.path.abspath(path_cand))
534 break
535 inputdepthmaps.append(custom_depthmap)
536 except Exception as e:
537 print(f'Failed to load {path}, ignoring. Exception: {str(e)}')
538 inputdepthmaps_n = len([1 for x in inputdepthmaps if x is not None])
539 print(f'{len(inputimages)} images will be processed, {inputdepthmaps_n} existing depthmaps will be reused')
540
541 gen_obj = core_generation_funnel(outpath, inputimages, inputdepthmaps, inputnames, inputs, backbone.gather_ops())
542
543 # Saving images
544 img_results = []
545 results_total = 0
546 inpainted_mesh_fi = mesh_simple_fi = None
547 msg = "" # Empty string is never returned
548 while True:
549 try:
550 input_i, type, result = next(gen_obj)
551 results_total += 1
552 except StopIteration:
553 # TODO: return more info
554 msg = '<h3>Successfully generated</h3>' if results_total > 0 else \
555 '<h3>Successfully generated nothing - please check the settings and try again</h3>'
556 break
557 except Exception as e:
558 msg = format_exception(e)
559 break
560 if type == 'simple_mesh':
561 mesh_simple_fi = result
562 continue
563 if type == 'inpainted_mesh':
564 inpainted_mesh_fi = result
565 continue
566 if not isinstance(result, Image.Image):
567 print(f'This is not supposed to happen! Somehow output type {type} is not supported! Input_i: {input_i}.')
568 continue
569 img_results += [(input_i, type, result)]
570
571 if inputs["save_outputs"]:
572 try:
573 basename = 'depthmap'
574 if depthmap_mode == '2' and inputnames[input_i] is not None:
575 if outpath != backbone.get_opt('outdir_extras_samples', None):
576 basename = Path(inputnames[input_i]).stem
577 suffix = "" if type == "depth" else f"{type}"
578 backbone.save_image(result, path=outpath, basename=basename, seed=None,
579 prompt=None, extension=backbone.get_opt('samples_format', 'png'), short_filename=True,
580 no_prompt=True, grid=False, pnginfo_section_name="extras",
581 suffix=suffix)
582 except Exception as e:
583 if not ('image has wrong mode' in str(e) or 'I;16' in str(e)):
584 raise e
585 print('Catched exception: image has wrong mode!')
586 traceback.print_exc()
587
588 # Deciding what mesh to display (and if)
589 display_mesh_fi = None
590 if backbone.get_opt('depthmap_script_show_3d', True):
591 display_mesh_fi = mesh_simple_fi
592 if backbone.get_opt('depthmap_script_show_3d_inpaint', True):
593 if inpainted_mesh_fi is not None and len(inpainted_mesh_fi) > 0:
594 display_mesh_fi = inpainted_mesh_fi
595 return map(lambda x: x[2], img_results), inpainted_mesh_fi, display_mesh_fi, msg.replace('\n', '<br>')
596