Skip to content

Commit

Permalink
Add resolution field in BiRefNet extras tab section
Browse files Browse the repository at this point in the history
  • Loading branch information
dimitribarbot committed Aug 26, 2024
1 parent 59243f1 commit 70b087b
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions scripts/postprocessing_rembg.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,8 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
def ui(self):
with ui_components.InputAccordion(False, label="Remove background with BiRefNet") as enable:
with FormRow():
model = gr.Dropdown(label="Remove background", choices=models, value="None")
model = gr.Dropdown(label="Remove background model", choices=models, value="None", info="Choose a BiRefNet model. Each model gives a different result.")
resolution = gr.Textbox(label="Resolution", value="", placeholder="1024x1024", info="If left empty, it will take image size rounded to the nearest multiple of 32")
return_foreground = gr.Checkbox(label="Return foreground", value=False)
return_edge_mask = gr.Checkbox(label="Return edge mask", value=False)

Expand All @@ -53,12 +54,13 @@ def ui(self):
return {
"enable": enable,
"model": model,
"resolution": resolution,
"return_foreground": return_foreground,
"return_edge_mask": return_edge_mask,
"edge_mask_width": edge_mask_width,
}

def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, model, return_foreground, return_edge_mask, edge_mask_width):
def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, model, resolution, return_foreground, return_edge_mask, edge_mask_width):
if not enable:
return

Expand All @@ -69,7 +71,7 @@ def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, model,

mask, output_image, edge_mask = birefnet.process(
pp.image.convert("RGB"),
resolution='',
resolution=resolution,
return_foreground=return_foreground,
return_edge_mask=return_edge_mask,
edge_mask_width=edge_mask_width
Expand Down

0 comments on commit 70b087b

Please sign in to comment.