Skip to content

forge and gradio4 compatibility #467

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions scripts/depthmap_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ async def process(

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")
print(f"Processing {str(len(depth_input_images))} images through the API")

pil_images = []
for input_image in depth_input_images:
Expand All @@ -81,7 +81,7 @@ async def process_video(
):
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")
print(f"Processing {str(len(depth_input_images))} images through the API")

# You can use either these strings, or integers
available_models = {
Expand Down
38 changes: 28 additions & 10 deletions src/backbone.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,21 +60,39 @@ def get_outpath():

def unload_sd_model():
from modules import shared, devices
if shared.sd_model is not None:
if shared.sd_model.cond_stage_model is not None:
shared.sd_model.cond_stage_model.to(devices.cpu)
if shared.sd_model.first_stage_model is not None:
shared.sd_model.first_stage_model.to(devices.cpu)
try:
if shared.sd_model is not None:
if shared.sd_model.cond_stage_model is not None:
shared.sd_model.cond_stage_model.to(devices.cpu)
if shared.sd_model.first_stage_model is not None:
shared.sd_model.first_stage_model.to(devices.cpu)
except Exception as e:
if type(e)== AttributeError:
print('trying to catch forge (might be a attribute error)')
from backend import memory_management
memory_management.unload_all_models()
memory_management.soft_empty_cache()
else:
raise
# Maybe something else???


def reload_sd_model():
from modules import shared, devices
if shared.sd_model is not None:
if shared.sd_model.cond_stage_model is not None:
shared.sd_model.cond_stage_model.to(devices.device)
if shared.sd_model.first_stage_model:
shared.sd_model.first_stage_model.to(devices.device)
try:
if shared.sd_model is not None:
if shared.sd_model.cond_stage_model is not None:
shared.sd_model.cond_stage_model.to(devices.device)
if shared.sd_model.first_stage_model:
shared.sd_model.first_stage_model.to(devices.device)
except Exception as e:
if type(e)== AttributeError:
print('trying to catch forge (might be a attribute error)')
from backend import memory_management
memory_management.unload_all_models()
memory_management.soft_empty_cache()
else:
raise
# Maybe something else???

def get_hide_dirs():
Expand Down
91 changes: 63 additions & 28 deletions src/common_ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,15 @@ def ensure_gradio_temp_directory():

ensure_gradio_temp_directory()

is_gradio4 = int(gr.__version__[0])>3

def main_ui_panel(is_depth_tab):

if is_gradio4:
Box = gr.Group
else:
Box = gr.Box

inp = GradioComponentBundle()
# TODO: Greater visual separation
with gr.Blocks():
Expand All @@ -41,7 +48,7 @@ def main_ui_panel(is_depth_tab):
'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:
with Box() as cur_option_root:
inp -= 'depthmap_gen_row_1', cur_option_root
with gr.Row():
inp += go.BOOST, gr.Checkbox(label="BOOST",
Expand All @@ -57,7 +64,7 @@ def main_ui_panel(is_depth_tab):
label='Tiling mode', info='Reduces seams that appear if the depthmap is tiled into a grid'
)

with gr.Box() as cur_option_root:
with Box() as cur_option_root:
inp -= 'depthmap_gen_row_2', cur_option_root
with gr.Row():
with gr.Group(): # 50% of width
Expand All @@ -71,7 +78,7 @@ def main_ui_panel(is_depth_tab):
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:
with 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")
Expand All @@ -81,7 +88,7 @@ def main_ui_panel(is_depth_tab):
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 Box():
with gr.Row():
inp += go.GEN_STEREO, gr.Checkbox(label="Generate stereoscopic (3D) image(s)")
with gr.Column(visible=False) as stereo_options:
Expand All @@ -104,7 +111,7 @@ def main_ui_panel(is_depth_tab):
inp += go.STEREO_BALANCE, gr.Slider(minimum=-1.0, maximum=1.0, step=0.05,
label='Balance between eyes')

with gr.Box():
with Box():
with gr.Row():
inp += go.GEN_NORMALMAP, gr.Checkbox(label="Generate NormalMap")
with gr.Column(visible=False) as normalmap_options:
Expand All @@ -124,11 +131,11 @@ def main_ui_panel(is_depth_tab):
inp += go.NORMALMAP_INVERT, gr.Checkbox(label="Invert")

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

with gr.Box():
with Box():
with gr.Column():
inp += go.GEN_SIMPLE_MESH, gr.Checkbox(label="Generate simple 3D mesh")
with gr.Column(visible=False) as mesh_options:
Expand All @@ -139,7 +146,7 @@ def main_ui_panel(is_depth_tab):
inp += go.SIMPLE_MESH_SPHERICAL, gr.Checkbox(label="Equirectangular projection")

if is_depth_tab:
with gr.Box():
with Box():
with gr.Column():
inp += go.GEN_INPAINTED_MESH, gr.Checkbox(
label="Generate 3D inpainted mesh")
Expand All @@ -149,7 +156,7 @@ def main_ui_panel(is_depth_tab):
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():
with 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():
Expand All @@ -163,33 +170,49 @@ def main_ui_panel(is_depth_tab):
label="Rembg Model", type="value",
choices=['u2net', 'u2netp', 'u2net_human_seg', 'silueta', "isnet-general-use", "isnet-anime"])

with gr.Box():
with 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)

if is_gradio4:
return gr.Slider(value=w), gr.Slider(value=h)
else:
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]]
)
def update_boost(a, b):
if is_gradio4:
return (gr.Checkbox(visible= not a), gr.Row(visible = not a and not b ))
else:
return (inp[go.NET_SIZE_MATCH].update(visible=not a),
options_depend_on_match_size.update(visible=not a and not b))

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)),
fn=update_boost,
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)

def update_tiling(a):
if is_gradio4:
if a:
return (gr.Checkbox(value=False), gr.Checkbox(value= True))
return (inp[go.BOOST], inp[go.NET_SIZE_MATCH])
else:
(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())

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()),
fn= update_tiling,
inputs=[inp[go.TILING_MODE]],
outputs=[inp[go.BOOST], inp[go.NET_SIZE_MATCH]]
)
Expand Down Expand Up @@ -248,13 +271,13 @@ def depthmap_mode_video(inp):
"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")
file_count="single", interactive=True, type="binary")
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)
interactive=True, type="binary", 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,
Expand Down Expand Up @@ -287,11 +310,11 @@ def on_ui_tabs():
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)
type="binary", 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")
interactive=True, type= "filepath" if is_gradio4 else "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(),
Expand Down Expand Up @@ -372,10 +395,17 @@ def on_ui_tabs():

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
if is_gradio4:
return gr.Row(visible=hide), \
gr.Group(visible = not hide), \
gr.Group(visible = not hide), \
gr.Group(visible = not hide), not hide
else:
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(
Expand Down Expand Up @@ -501,9 +531,14 @@ def run_generate(*inputs):
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])
if not is_gradio4:
image = Image.open(os.path.abspath(img.name))
inputimages.append(image)
inputnames.append(os.path.splitext(img.orig_name)[0])
else:
image = Image.open(os.path.abspath(img))
inputimages.append(image)
inputnames.append(os.path.splitext(img)[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.
Expand Down
15 changes: 13 additions & 2 deletions src/gradio_args_transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,12 +69,23 @@ def enkey_body(self):
return [self.internal[x] for x in sorted(list(self.internal.keys()))]

def add_rule(self, first, rule, second):
is_gradio4 = int(gr.__version__[0])>3
first = self[first] if first in self else first
second = self[second] if second in self else second
if rule == 'visible-if-not':
second.change(fn=lambda v: first.update(visible=not v), inputs=[second], outputs=[first])
def update_visible_if_not(v):
if is_gradio4:
return gr.Column(visible=not v)
else:
return first.update(visible=not v)
second.change(update_visible_if_not, [second], [first])
elif rule == 'visible-if':
second.change(fn=lambda v: first.update(visible=v), inputs=[second], outputs=[first])
def update_visible_if(v):
if is_gradio4:
return gr.Column(visible=v)
else:
return first.update(visible=v)
second.change(update_visible_if, [second], [first])
else:
raise Exception(f'Unknown rule type {rule}')

Expand Down