Skip to content

Commit

Permalink
SD/UI: Use a single model selection box on UI tabs (#1906)
Browse files Browse the repository at this point in the history 
* Allow entry of a huggingface model id or civitai download url to be
done in the main model selection dropdown on SD tabs
* Remove separate textbox for entering huggingface model id or civitai
download url on SD Tabs
* Remove 'None' option from the model selection dropdown (no longer
needed) on SD tabs
* Update png metadata drop zone on txt2img tab to work with a single
argument for model selection
* Update UI generate functions on SD tabs to work with single argument
model selection
* Update API code for changes to the UI generate functions
* Move info about the custom model path to the logging textarea on SD
tabs
  • Loading branch information
@one-lithe-rune
one-lithe-rune authored Oct 21, 2023
1 parent e6cb5ce commit 7cd14fd
Show file tree
Hide file tree
Showing 8 changed files with 107 additions and 221 deletions.
15 changes: 5 additions & 10 deletions apps/stable_diffusion/web/index.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,6 @@ def resource_path(relative_path):
from apps.stable_diffusion.web.ui import (
txt2img_web,
txt2img_custom_model,
txt2img_hf_model_id,
txt2img_gallery,
txt2img_png_info_img,
txt2img_status,
Expand All @@ -122,7 +121,6 @@ def resource_path(relative_path):
# h2ogpt_web,
img2img_web,
img2img_custom_model,
img2img_hf_model_id,
img2img_gallery,
img2img_init_image,
img2img_status,
Expand All @@ -131,7 +129,6 @@ def resource_path(relative_path):
img2img_sendto_upscaler,
inpaint_web,
inpaint_custom_model,
inpaint_hf_model_id,
inpaint_gallery,
inpaint_init_image,
inpaint_status,
Expand All @@ -140,7 +137,6 @@ def resource_path(relative_path):
inpaint_sendto_upscaler,
outpaint_web,
outpaint_custom_model,
outpaint_hf_model_id,
outpaint_gallery,
outpaint_init_image,
outpaint_status,
Expand All @@ -149,7 +145,6 @@ def resource_path(relative_path):
outpaint_sendto_upscaler,
upscaler_web,
upscaler_custom_model,
upscaler_hf_model_id,
upscaler_gallery,
upscaler_init_image,
upscaler_status,
Expand Down Expand Up @@ -399,31 +394,31 @@ def register_outputgallery_button(button, selectedid, inputs, outputs):
modelmanager_sendto_txt2img,
0,
[hf_models],
[txt2img_custom_model, txt2img_hf_model_id, tabs],
[txt2img_custom_model, tabs],
)
register_modelmanager_button(
modelmanager_sendto_img2img,
1,
[hf_models],
[img2img_custom_model, img2img_hf_model_id, tabs],
[img2img_custom_model, tabs],
)
register_modelmanager_button(
modelmanager_sendto_inpaint,
2,
[hf_models],
[inpaint_custom_model, inpaint_hf_model_id, tabs],
[inpaint_custom_model, tabs],
)
register_modelmanager_button(
modelmanager_sendto_outpaint,
3,
[hf_models],
[outpaint_custom_model, outpaint_hf_model_id, tabs],
[outpaint_custom_model, tabs],
)
register_modelmanager_button(
modelmanager_sendto_upscaler,
4,
[hf_models],
[upscaler_custom_model, upscaler_hf_model_id, tabs],
[upscaler_custom_model, tabs],
)

sd_web.queue()
Expand Down
5 changes: 0 additions & 5 deletions apps/stable_diffusion/web/ui/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
txt2img_api,
txt2img_web,
txt2img_custom_model,
txt2img_hf_model_id,
txt2img_gallery,
txt2img_png_info_img,
txt2img_status,
Expand All @@ -17,7 +16,6 @@
img2img_api,
img2img_web,
img2img_custom_model,
img2img_hf_model_id,
img2img_gallery,
img2img_init_image,
img2img_status,
Expand All @@ -30,7 +28,6 @@
inpaint_api,
inpaint_web,
inpaint_custom_model,
inpaint_hf_model_id,
inpaint_gallery,
inpaint_init_image,
inpaint_status,
Expand All @@ -43,7 +40,6 @@
outpaint_api,
outpaint_web,
outpaint_custom_model,
outpaint_hf_model_id,
outpaint_gallery,
outpaint_init_image,
outpaint_status,
Expand All @@ -56,7 +52,6 @@
upscaler_api,
upscaler_web,
upscaler_custom_model,
upscaler_hf_model_id,
upscaler_gallery,
upscaler_init_image,
upscaler_status,
Expand Down
60 changes: 21 additions & 39 deletions apps/stable_diffusion/web/ui/img2img_ui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,7 @@ def img2img_inf(
batch_count: int,
batch_size: int,
scheduler: str,
custom_model: str,
hf_model_id: str,
model_id: str,
custom_vae: str,
precision: str,
device: str,
Expand Down Expand Up @@ -103,21 +102,17 @@ def img2img_inf(
args.ckpt_loc = ""
args.hf_model_id = ""
args.custom_vae = ""
if custom_model == "None":
if not hf_model_id:
return (
None,
"Please provide either custom model or huggingface model ID, "
"both must not be empty.",
)
if "civitai" in hf_model_id:
args.ckpt_loc = hf_model_id
else:
args.hf_model_id = hf_model_id
elif ".ckpt" in custom_model or ".safetensors" in custom_model:
args.ckpt_loc = get_custom_model_pathfile(custom_model)

# .safetensor or .chkpt on the custom model path
if model_id in get_custom_model_files():
args.ckpt_loc = get_custom_model_pathfile(model_id)
# civitai download
elif "civitai" in model_id:
args.ckpt_loc = model_id
# either predefined or huggingface
else:
args.hf_model_id = custom_model
args.hf_model_id = model_id

if custom_vae != "None":
args.custom_vae = get_custom_model_pathfile(custom_vae, model="vae")

Expand Down Expand Up @@ -334,8 +329,7 @@ def img2img_api(
batch_count=1,
batch_size=1,
scheduler="EulerDiscrete",
custom_model="None",
hf_model_id=InputData["hf_model_id"]
model_id=InputData["hf_model_id"]
if "hf_model_id" in InputData.keys()
else "stabilityai/stable-diffusion-2-1-base",
custom_vae="None",
Expand Down Expand Up @@ -382,32 +376,19 @@ def img2img_api(
with gr.Column(scale=1, min_width=600):
with gr.Row():
# janky fix for overflowing text
i2i_model_info = (str(get_custom_model_path())).replace(
"\\", "\n\\"
i2i_model_info = (
f"Custom Model Path: {str(get_custom_model_path())}"
)
i2i_model_info = f"Custom Model Path: {i2i_model_info}"
img2img_custom_model = gr.Dropdown(
label=f"Models",
info=i2i_model_info,
info="Select, or enter HuggingFace Model ID or Civitai model download URL",
elem_id="custom_model",
value=os.path.basename(args.ckpt_loc)
if args.ckpt_loc
else "stabilityai/stable-diffusion-2-1-base",
choices=["None"]
+ get_custom_model_files()
+ predefined_models,
choices=get_custom_model_files() + predefined_models,
allow_custom_value=True,
)
img2img_hf_model_id = gr.Textbox(
elem_id="hf_model_id",
placeholder="Select 'None' in the Models dropdown "
"on the left and enter model ID here "
"e.g: SG161222/Realistic_Vision_V1.3, "
"https://civitai.com/api/download/models/15236",
value="",
label="HuggingFace Model ID or Civitai model "
"download URL",
lines=3,
scale=2,
)
# janky fix for overflowing text
i2i_vae_info = (str(get_custom_model_path("vae"))).replace(
Expand All @@ -423,6 +404,7 @@ def img2img_api(
else "None",
choices=["None"] + get_custom_model_files("vae"),
allow_custom_value=True,
scale=1,
)

with gr.Group(elem_id="prompt_box_outer"):
Expand Down Expand Up @@ -677,9 +659,10 @@ def create_canvas(w, h):
object_fit="contain",
)
std_output = gr.Textbox(
value=f"Images will be saved at "
value=f"{i2i_model_info}\n"
f"Images will be saved at "
f"{get_generated_imgs_path()}",
lines=1,
lines=2,
elem_id="std_output",
show_label=False,
)
Expand Down Expand Up @@ -709,7 +692,6 @@ def create_canvas(w, h):
batch_size,
scheduler,
img2img_custom_model,
img2img_hf_model_id,
custom_vae,
precision,
device,
Expand Down
59 changes: 20 additions & 39 deletions apps/stable_diffusion/web/ui/inpaint_ui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,7 @@ def inpaint_inf(
batch_count: int,
batch_size: int,
scheduler: str,
custom_model: str,
hf_model_id: str,
model_id: str,
custom_vae: str,
precision: str,
device: str,
Expand Down Expand Up @@ -89,21 +88,17 @@ def inpaint_inf(
args.ckpt_loc = ""
args.hf_model_id = ""
args.custom_vae = ""
if custom_model == "None":
if not hf_model_id:
return (
None,
"Please provide either custom model or huggingface model ID, "
"both must not be empty.",
)
if "civitai" in hf_model_id:
args.ckpt_loc = hf_model_id
else:
args.hf_model_id = hf_model_id
elif ".ckpt" in custom_model or ".safetensors" in custom_model:
args.ckpt_loc = get_custom_model_pathfile(custom_model)

# .safetensor or .chkpt on the custom model path
if model_id in get_custom_model_files(custom_checkpoint_type="inpainting"):
args.ckpt_loc = get_custom_model_pathfile(model_id)
# civitai download
elif "civitai" in model_id:
args.ckpt_loc = model_id
# either predefined or huggingface
else:
args.hf_model_id = custom_model
args.hf_model_id = model_id

if custom_vae != "None":
args.custom_vae = get_custom_model_pathfile(custom_vae, model="vae")

Expand Down Expand Up @@ -282,8 +277,7 @@ def inpaint_api(
batch_count=1,
batch_size=1,
scheduler="EulerDiscrete",
custom_model="None",
hf_model_id=InputData["hf_model_id"]
model_id=InputData["hf_model_id"]
if "hf_model_id" in InputData.keys()
else "stabilityai/stable-diffusion-2-inpainting",
custom_vae="None",
Expand Down Expand Up @@ -327,35 +321,21 @@ def inpaint_api(
with gr.Row():
# janky fix for overflowing text
inpaint_model_info = (
str(get_custom_model_path())
).replace("\\", "\n\\")
inpaint_model_info = (
f"Custom Model Path: {inpaint_model_info}"
f"Custom Model Path: {str(get_custom_model_path())}"
)
inpaint_custom_model = gr.Dropdown(
label=f"Models",
info=inpaint_model_info,
info="Select, or enter HuggingFace Model ID or Civitai model download URL",
elem_id="custom_model",
value=os.path.basename(args.ckpt_loc)
if args.ckpt_loc
else "stabilityai/stable-diffusion-2-inpainting",
choices=["None"]
+ get_custom_model_files(
choices=get_custom_model_files(
custom_checkpoint_type="inpainting"
)
+ predefined_paint_models,
allow_custom_value=True,
)
inpaint_hf_model_id = gr.Textbox(
elem_id="hf_model_id",
placeholder="Select 'None' in the Models dropdown "
"on the left and enter model ID here "
"e.g: ghunkins/stable-diffusion-liberty-inpainting, "
"https://civitai.com/api/download/models/3433",
value="",
label="HuggingFace Model ID or Civitai model "
"download URL",
lines=3,
scale=2,
)
# janky fix for overflowing text
inpaint_vae_info = (
Expand All @@ -371,6 +351,7 @@ def inpaint_api(
else "None",
choices=["None"] + get_custom_model_files("vae"),
allow_custom_value=True,
scale=1,
)

with gr.Group(elem_id="prompt_box_outer"):
Expand Down Expand Up @@ -554,9 +535,10 @@ def inpaint_api(
object_fit="contain",
)
std_output = gr.Textbox(
value=f"Images will be saved at "
value=f"{inpaint_model_info}\n"
"Images will be saved at "
f"{get_generated_imgs_path()}",
lines=1,
lines=2,
elem_id="std_output",
show_label=False,
)
Expand Down Expand Up @@ -588,7 +570,6 @@ def inpaint_api(
batch_size,
scheduler,
inpaint_custom_model,
inpaint_hf_model_id,
custom_vae,
precision,
device,
Expand Down
Loading

0 comments on commit 7cd14fd

Please sign in to comment.