segment prototype

This commit is contained in:
Vladimir Mandic
2024-01-13 21:20:31 -05:00
parent 3bb7d04599
commit 4cef19892e
8 changed files with 277 additions and 238 deletions
+53 -20
View File
@@ -21,6 +21,7 @@ MODELS = {
'Facebook SAM ViT Huge': 'facebook/sam-vit-huge',
'SlimSAM Uniform': 'Zigeng/SlimSAM-uniform-50',
}
COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen']
cache_dir = 'models/control/segment'
loaded_model = None
model: SamModel = None
@@ -49,43 +50,75 @@ def init(selected_model: str, input_image: gr.Image):
# run as auto-mask with all possible masks
def run_segment(selected_model: str, input_image: gr.Image):
def run_segment(selected_model: str, input_image: gr.Image, points_per_batch=64, pred_iou_thresh=0.75, stability_score_thresh=0.85, crops_nms_thresh=0.5, crop_overlap_ratio=0.3, topK=25, colormap='jet', erode=0, dilate=0):
if not init(selected_model, input_image):
return input_image
return gr.update(), None
input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 255)
input_mask = input_mask.convert('L')
input_image = input_image.get('image', None)
generator: MaskGenerationPipeline = MaskGenerationPipeline(model=model, image_processor=processor, device=devices.device)
with devices.inference_context():
outputs = generator(
input_image,
points_per_batch=64,
pred_iou_thresh=0.75,
stability_score_thresh=0.95,
crops_nms_thresh=0.7,
crop_overlap_ratio=0.3,
points_per_batch=points_per_batch,
pred_iou_thresh=pred_iou_thresh,
stability_score_thresh=stability_score_thresh,
crops_nms_thresh=crops_nms_thresh,
crop_overlap_ratio=crop_overlap_ratio,
)
combined_mask = np.zeros(input_mask.size, dtype='uint8')
for i, mask in enumerate(outputs['masks']):
mask = mask.astype('uint8') * i * 10
input_mask = np.array(input_mask) // 255
input_mask_size = np.count_nonzero(input_mask)
print('HERE', input_mask.shape, input_mask_size)
i = 1
for mask in outputs['masks']:
mask = mask.astype('uint8')
mask_size = np.count_nonzero(mask)
if mask_size == 0:
continue
overlap = 0
if input_mask_size > 0:
overlap = cv2.bitwise_and(mask, input_mask)
overlap = np.count_nonzero(overlap)
if overlap == 0:
continue
# TODO erode,dilate
if erode > 0:
mask = cv2.erode(mask, np.ones((erode, erode), np.uint8), iterations=2) # remove noise
if dilate > 0:
mask = cv2.dilate(mask, np.ones((dilate, dilate), np.uint8), iterations=2) # expand area
mask = (topK + 1 - i) * mask * (255 // topK) # set grayscale intensity so we can recolor
combined_mask = combined_mask + mask
total_size = np.prod(mask.shape)
area_size = np.count_nonzero(mask)
shared.log.debug(f'Segment mask: i={i} area={area_size/total_size:.2f} score={outputs["scores"][i].item():.2f}')
if i > 25:
i += 1
if i > topK:
break
combined_mask = cv2.applyColorMap(combined_mask, cv2.COLORMAP_JET)
combined_image = cv2.addWeighted(np.array(input_image), 0.6, combined_mask, 0.4, 0)
mask_size = np.count_nonzero(combined_mask)
total_size = np.prod(combined_mask.shape)
area_size = np.count_nonzero(combined_mask)
shared.log.debug(f'Segment mask: size={input_image.width}x{input_image.height} input={input_mask_size}px masked={mask_size}px area={area_size/total_size:.2f}')
colored_mask = cv2.applyColorMap(combined_mask, COLORMAP.index(colormap)) # recolor mask
combined_image = cv2.addWeighted(np.array(input_image), 0.6, colored_mask, 0.4, 0)
_thres, binary_mask = cv2.threshold(combined_mask, 1, 255, cv2.THRESH_BINARY_INV) # create mask
binary_mask = np.invert(binary_mask)
binary_mask = Image.fromarray(binary_mask)
combined_mask = Image.fromarray(combined_mask)
combined_image = Image.fromarray(combined_image)
colored_mask = Image.fromarray(colored_mask)
overlay_image = Image.fromarray(combined_image)
# TODO return type
binary_mask.save('/tmp/mask-binary.png')
combined_mask.save('/tmp/mask-combined.png')
combined_image.save('/tmp/mask-combined-image.png')
combined_mask.save('/tmp/mask-colored.png')
overlay_image.save('/tmp/mask-overlay.png')
return input_image, overlay_image
# run with sam model directly needing set of points
def run_segment_points(selected_model: str, input_image: gr.Image):
if not init(selected_model, input_image):
return input_image
input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 0)
# input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 0)
input_image = input_image.get('image', None)
with devices.inference_context():
inputs = processor(
@@ -118,7 +151,7 @@ def run_segment_points(selected_model: str, input_image: gr.Image):
output_masks.append(mask)
def create_segment_ui(input_image: gr.Image):
def create_segment_ui(input_image: gr.Image, preview_image: gr.Image):
selected = gr.Dropdown(label="Segment", choices=MODELS.keys(), value='None')
selected.change(fn=run_segment, inputs=[selected, input_image], outputs=[])
selected.change(fn=run_segment, inputs=[selected, input_image], outputs=[input_image, preview_image])
return selected
+2 -2
View File
@@ -22,11 +22,11 @@ class Timer:
self.total += e + extra_time
def summary(self, min_time=0.05):
res = f"{self.total:.2f}"
res = f"{self.total:.2f} "
additions = [x for x in self.records.items() if x[1] >= min_time]
if not additions:
return res
res += " { " + " ".join([f"{category}={time_taken:.2f}" for category, time_taken in additions]) + " }"
res += " ".join([f"{category}={time_taken:.2f}" for category, time_taken in additions])
return res
def reset(self):
+3 -4
View File
@@ -141,10 +141,9 @@ def create_ui(startup_timer = None):
timer.startup.record("ui-extras")
with gr.Blocks(analytics_enabled=False) as train_interface:
if shared.backend == shared.Backend.ORIGINAL:
from modules import ui_train
ui_train.create_ui()
timer.startup.record("ui-train")
from modules import ui_train
ui_train.create_ui()
timer.startup.record("ui-train")
with gr.Blocks(analytics_enabled=False) as models_interface:
from modules import ui_models
+5 -3
View File
@@ -353,7 +353,7 @@ def create_ui(_blocks: gr.Blocks=None):
with gr.Row(variant='compact', elem_id="control_extra_networks", visible=False) as extra_networks_ui:
from modules import timer, ui_extra_networks
extra_networks_ui = ui_extra_networks.create_ui(extra_networks_ui, btn_extra, 'control', skip_indexing=shared.opts.extra_network_skip_indexing)
timer.startup.record('ui-extra-networks')
timer.startup.record('ui-en')
with gr.Row(elem_id='control_status'):
result_txt = gr.HTML(elem_classes=['control-result'], elem_id='control-result')
@@ -368,8 +368,6 @@ def create_ui(_blocks: gr.Blocks=None):
input_resize = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="select", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_resize')
input_inpaint = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=64, mask_opacity=0.6)
interrogate_clip, interrogate_booru = ui_sections.create_interrogate_buttons('control')
# with gr.Row():
# segment_ui = segment.create_segment_ui(input_inpaint)
with gr.Row():
input_buttons = [gr.Button('Select', visible=True, interactive=False), gr.Button('Inpaint', visible=True, interactive=True), gr.Button('Outpaint', visible=True, interactive=True)]
with gr.Tab('Video', id='in-video') as tab_video:
@@ -404,6 +402,10 @@ def create_ui(_blocks: gr.Blocks=None):
with gr.Tab('Preview', id='preview-image') as tab_image:
preview_process = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=gr_height, visible=True)
# TODO segment as accordian
# with gr.Row():
# segment_ui = segment.create_segment_ui(input_inpaint, preview_process)
with gr.Tabs(elem_id='control-tabs') as _tabs_control_type:
with gr.Tab('ControlNet') as _tab_controlnet:
+2 -1
View File
@@ -3,7 +3,7 @@ from PIL import Image
import gradio as gr
import numpy as np
from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call
from modules import shared, ui_common, ui_sections, generation_parameters_copypaste
from modules import timer, shared, ui_common, ui_sections, generation_parameters_copypaste
from modules.ui_components import FormRow, FormGroup
@@ -45,6 +45,7 @@ def create_ui():
with FormRow(variant='compact', elem_id="img2img_extra_networks", visible=False) as extra_networks_ui:
from modules import ui_extra_networks
extra_networks_ui_img2img = ui_extra_networks.create_ui(extra_networks_ui, img2img_extra_networks_button, 'img2img', skip_indexing=shared.opts.extra_network_skip_indexing)
timer.startup.record('ui-en')
with FormRow(elem_id="img2img_interface", equal_height=False):
with gr.Column(variant='compact', elem_id="img2img_settings"):
+208 -204
View File
@@ -1,6 +1,6 @@
import os
import gradio as gr
from modules import sd_hijack, script_callbacks, shared
from modules import script_callbacks, shared
from modules.ui_components import FormRow
from modules.ui_common import create_refresh_button
from modules.ui_sections import create_sampler_inputs
@@ -150,227 +150,231 @@ def create_ui():
)
### train embedding tab
with gr.Tab(label="Train embedding", id="train_embedding_tab") as tab_ti:
tab_ti.select(fn=lambda: train_tab_change('ti'), inputs=[], outputs=[action_pp, action_ti, action_hn])
def get_textual_inversion_template_names():
return sorted(textual_inversion.textual_inversion_templates)
if shared.backend == shared.Backend.ORIGINAL:
from modules import sd_hijack
with gr.Tab(label="Train embedding", id="train_embedding_tab") as tab_ti:
tab_ti.select(fn=lambda: train_tab_change('ti'), inputs=[], outputs=[action_pp, action_ti, action_hn])
def get_textual_inversion_template_names():
return sorted(textual_inversion.textual_inversion_templates)
gr.HTML('<h2>Select existing embedding to continue training or create a new one</h2>')
with FormRow():
with gr.Column():
with gr.Row():
ti_name = gr.Dropdown(label='Select embedding', choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))
create_refresh_button(ti_name, sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings, lambda: {"choices": sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys())}, "refresh_train_embedding_name")
with gr.Column():
ti_new_name = gr.Textbox(label="Create emebedding")
ti_init_text = gr.Textbox(label="Initialization text", value="*")
ti_vectors = gr.Slider(label="Number of vectors per token", minimum=1, maximum=75, step=1, value=1)
ti_overwrite = gr.Checkbox(value=False, label="Overwrite Old Embedding")
with gr.Row():
ti_create = gr.Button(value="Create embedding", variant='secondary')
with gr.Box():
gr.HTML('<h2>Training parameters</h2>')
ti_learn_rate = gr.Textbox(label='Embedding Learning rate', placeholder="Embedding Learning rate", value="0.005")
gr.HTML('<h2>Select existing embedding to continue training or create a new one</h2>')
with FormRow():
ti_clip_grad_mode = gr.Dropdown(value="disabled", label="Gradient Clipping", choices=["disabled", "value", "norm"])
ti_clip_grad_value = gr.Number(label="Gradient clip value", value=0.1)
ti_batch_size = gr.Number(label='Batch size', value=1, precision=0)
ti_gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0)
ti_steps = gr.Number(label='Max steps', value=1000, precision=0)
with gr.Column():
with gr.Row():
ti_name = gr.Dropdown(label='Select embedding', choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))
create_refresh_button(ti_name, sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings, lambda: {"choices": sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys())}, "refresh_train_embedding_name")
with gr.Column():
ti_new_name = gr.Textbox(label="Create emebedding")
ti_init_text = gr.Textbox(label="Initialization text", value="*")
ti_vectors = gr.Slider(label="Number of vectors per token", minimum=1, maximum=75, step=1, value=1)
ti_overwrite = gr.Checkbox(value=False, label="Overwrite Old Embedding")
with gr.Row():
ti_create = gr.Button(value="Create embedding", variant='secondary')
with gr.Box():
gr.HTML('<h2>Training images</h2>')
ti_dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
with FormRow():
ti_varsize = gr.Checkbox(label="Do not resize images", value=False)
ti_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)
ti_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)
ti_use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False)
with gr.Box():
gr.HTML('<h2>Training parameters</h2>')
ti_learn_rate = gr.Textbox(label='Embedding Learning rate', placeholder="Embedding Learning rate", value="0.005")
with FormRow():
ti_clip_grad_mode = gr.Dropdown(value="disabled", label="Gradient Clipping", choices=["disabled", "value", "norm"])
ti_clip_grad_value = gr.Number(label="Gradient clip value", value=0.1)
ti_batch_size = gr.Number(label='Batch size', value=1, precision=0)
ti_gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0)
ti_steps = gr.Number(label='Max steps', value=1000, precision=0)
with gr.Box():
gr.HTML('<h2>Dataset processing</h2>')
with FormRow():
ti_template = gr.Dropdown(label='Prompt template', value="style_filewords.txt", choices=get_textual_inversion_template_names())
create_refresh_button(ti_template, textual_inversion.list_textual_inversion_templates, lambda: {"choices": get_textual_inversion_template_names()}, "refrsh_train_template_file")
ti_shuffle = gr.Checkbox(label="Shuffle tags", value=False)
ti_tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts", value=0)
ti_latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'])
with gr.Box():
gr.HTML('<h2>Training images</h2>')
ti_dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
with FormRow():
ti_varsize = gr.Checkbox(label="Do not resize images", value=False)
ti_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)
ti_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)
ti_use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False)
with gr.Box():
gr.HTML('<h2>Training outputs</h2>')
with FormRow():
ti_create_every = gr.Number(label='Create interim images', value=500, precision=0)
ti_save_every = gr.Number(label='Create interim embeddings', value=500, precision=0)
ti_save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True)
ti_preview_from_txt2img = gr.Checkbox(label='Use current settings for previews', value=False)
ti_log_directory = gr.Textbox(label='Log directory', placeholder="Defaults to train/log/embedding", value="")
with gr.Box():
gr.HTML('<h2>Dataset processing</h2>')
with FormRow():
ti_template = gr.Dropdown(label='Prompt template', value="style_filewords.txt", choices=get_textual_inversion_template_names())
create_refresh_button(ti_template, textual_inversion.list_textual_inversion_templates, lambda: {"choices": get_textual_inversion_template_names()}, "refrsh_train_template_file")
ti_shuffle = gr.Checkbox(label="Shuffle tags", value=False)
ti_tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts", value=0)
ti_latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'])
ti_stop.click(fn=lambda: shared.state.interrupt(), inputs=[], outputs=[])
with gr.Box():
gr.HTML('<h2>Training outputs</h2>')
with FormRow():
ti_create_every = gr.Number(label='Create interim images', value=500, precision=0)
ti_save_every = gr.Number(label='Create interim embeddings', value=500, precision=0)
ti_save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True)
ti_preview_from_txt2img = gr.Checkbox(label='Use current settings for previews', value=False)
ti_log_directory = gr.Textbox(label='Log directory', placeholder="Defaults to train/log/embedding", value="")
ti_create.click(
fn=modules.textual_inversion.ui.create_embedding,
inputs=[
ti_new_name,
ti_init_text,
ti_vectors,
ti_overwrite,
],
outputs=[
ti_name,
train_output,
train_outcome,
]
)
ti_stop.click(fn=lambda: shared.state.interrupt(), inputs=[], outputs=[])
ti_train.click(
fn=wrap_gradio_gpu_call(modules.textual_inversion.ui.train_embedding, extra_outputs=[gr.update()]),
_js="startTrainMonitor",
inputs=[
dummy_component,
ti_name,
ti_learn_rate,
ti_batch_size,
ti_gradient_step,
ti_dataset_directory,
ti_log_directory,
ti_width,
ti_height,
ti_varsize,
ti_steps,
ti_clip_grad_mode,
ti_clip_grad_value,
ti_shuffle,
ti_tag_drop_out,
ti_latent_sampling_method,
ti_use_weight,
ti_create_every,
ti_save_every,
ti_template,
ti_save_image_with_stored_embedding,
ti_preview_from_txt2img,
*txt2img_preview_params,
],
outputs=[
train_output,
train_outcome,
]
)
ti_create.click(
fn=modules.textual_inversion.ui.create_embedding,
inputs=[
ti_new_name,
ti_init_text,
ti_vectors,
ti_overwrite,
],
outputs=[
ti_name,
train_output,
train_outcome,
]
)
ti_train.click(
fn=wrap_gradio_gpu_call(modules.textual_inversion.ui.train_embedding, extra_outputs=[gr.update()]),
_js="startTrainMonitor",
inputs=[
dummy_component,
ti_name,
ti_learn_rate,
ti_batch_size,
ti_gradient_step,
ti_dataset_directory,
ti_log_directory,
ti_width,
ti_height,
ti_varsize,
ti_steps,
ti_clip_grad_mode,
ti_clip_grad_value,
ti_shuffle,
ti_tag_drop_out,
ti_latent_sampling_method,
ti_use_weight,
ti_create_every,
ti_save_every,
ti_template,
ti_save_image_with_stored_embedding,
ti_preview_from_txt2img,
*txt2img_preview_params,
],
outputs=[
train_output,
train_outcome,
]
)
### train hypernetwork tab
with gr.Tab(label="Train hypernetwork", id="train_hypernetwork_tab") as tab_hn:
tab_hn.select(fn=lambda: train_tab_change('hn'), inputs=[], outputs=[action_pp, action_ti, action_hn])
gr.HTML('<h2>Select existing hypernetwork to continue training or create a new one</h2>')
with FormRow():
with gr.Column():
if shared.backend == shared.Backend.ORIGINAL:
from modules import sd_hijack
with gr.Tab(label="Train hypernetwork", id="train_hypernetwork_tab") as tab_hn:
tab_hn.select(fn=lambda: train_tab_change('hn'), inputs=[], outputs=[action_pp, action_ti, action_hn])
gr.HTML('<h2>Select existing hypernetwork to continue training or create a new one</h2>')
with FormRow():
with gr.Column():
with FormRow():
hn_name = gr.Dropdown(label='Hypernetwork', choices=sorted(shared.hypernetworks))
create_refresh_button(hn_name, shared.reload_hypernetworks, lambda: {"choices": sorted(shared.hypernetworks)}, "refresh_train_hypernetwork_name")
with gr.Column():
hn_new_name = gr.Textbox(label="Name")
hn_new_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "1024", "320", "640", "1280"])
hn_new_layer_structure = gr.Textbox("1, 2, 1", label="Enter hypernetwork layer structure", placeholder="1st and last digit must be 1. ex:'1, 2, 1'")
with gr.Row():
hn_new_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork", choices=modules.hypernetworks.ui.keys)
hn_new_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization", choices=["Normal", "KaimingUniform", "KaimingNormal", "XavierUniform", "XavierNormal"])
hn_new_add_layer_norm = gr.Checkbox(label="Add layer normalization")
hn_new_use_dropout = gr.Checkbox(label="Use dropout")
hn_new_dropout_structure = gr.Textbox("0, 0, 0", label="Enter hypernetwork Dropout structure", placeholder="1st and last digit must be 0 and values should be between 0 and 1. ex:'0, 0.01, 0'")
hn_overwrite = gr.Checkbox(value=False, label="Overwrite Old Hypernetwork")
with gr.Row():
hn_create = gr.Button(value="Create hypernetwork", variant='secondary')
with gr.Box():
gr.HTML('<h2>Training parameters</h2>')
hn_learn_rate = gr.Textbox(label='Hypernetwork Learning rate', placeholder="Hypernetwork Learning rate", value="0.00001")
with FormRow():
hn_name = gr.Dropdown(label='Hypernetwork', choices=sorted(shared.hypernetworks))
create_refresh_button(hn_name, shared.reload_hypernetworks, lambda: {"choices": sorted(shared.hypernetworks)}, "refresh_train_hypernetwork_name")
with gr.Column():
hn_new_name = gr.Textbox(label="Name")
hn_new_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "1024", "320", "640", "1280"])
hn_new_layer_structure = gr.Textbox("1, 2, 1", label="Enter hypernetwork layer structure", placeholder="1st and last digit must be 1. ex:'1, 2, 1'")
with gr.Row():
hn_new_activation_func = gr.Dropdown(value="linear", label="Select activation function of hypernetwork", choices=modules.hypernetworks.ui.keys)
hn_new_initialization_option = gr.Dropdown(value = "Normal", label="Select Layer weights initialization", choices=["Normal", "KaimingUniform", "KaimingNormal", "XavierUniform", "XavierNormal"])
hn_new_add_layer_norm = gr.Checkbox(label="Add layer normalization")
hn_new_use_dropout = gr.Checkbox(label="Use dropout")
hn_new_dropout_structure = gr.Textbox("0, 0, 0", label="Enter hypernetwork Dropout structure", placeholder="1st and last digit must be 0 and values should be between 0 and 1. ex:'0, 0.01, 0'")
hn_overwrite = gr.Checkbox(value=False, label="Overwrite Old Hypernetwork")
with gr.Row():
hn_create = gr.Button(value="Create hypernetwork", variant='secondary')
hn_clip_grad_mode = gr.Dropdown(value="disabled", label="Gradient Clipping", choices=["disabled", "value", "norm"])
hn_clip_grad_value = gr.Number(label="Gradient clip value", value=0.1)
hn_batch_size = gr.Number(label='Batch size', value=1, precision=0)
hn_gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0)
hn_steps = gr.Number(label='Max steps', value=1000, precision=0)
with gr.Box():
gr.HTML('<h2>Training parameters</h2>')
hn_learn_rate = gr.Textbox(label='Hypernetwork Learning rate', placeholder="Hypernetwork Learning rate", value="0.00001")
with FormRow():
hn_clip_grad_mode = gr.Dropdown(value="disabled", label="Gradient Clipping", choices=["disabled", "value", "norm"])
hn_clip_grad_value = gr.Number(label="Gradient clip value", value=0.1)
hn_batch_size = gr.Number(label='Batch size', value=1, precision=0)
hn_gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0)
hn_steps = gr.Number(label='Max steps', value=1000, precision=0)
with gr.Box():
gr.HTML('<h2>Training images</h2>')
hn_dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
with FormRow():
hn_varsize = gr.Checkbox(label="Do not resize images", value=False)
hn_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)
hn_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)
hn_use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False)
with gr.Box():
gr.HTML('<h2>Training images</h2>')
hn_dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
with FormRow():
hn_varsize = gr.Checkbox(label="Do not resize images", value=False)
hn_width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512)
hn_height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512)
hn_use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False)
with gr.Box():
gr.HTML('<h2>Dataset processing</h2>')
with FormRow():
hn_template = gr.Dropdown(label='Prompt template', value="style_filewords.txt", choices=get_textual_inversion_template_names())
create_refresh_button(hn_template, textual_inversion.list_textual_inversion_templates, lambda: {"choices": get_textual_inversion_template_names()}, "refrsh_train_template_file")
hn_shuffle_tags = gr.Checkbox(label="Shuffle tags by ',' when creating prompts.", value=False)
hn_tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts", value=0)
hn_latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'])
with gr.Box():
gr.HTML('<h2>Dataset processing</h2>')
with FormRow():
hn_template = gr.Dropdown(label='Prompt template', value="style_filewords.txt", choices=get_textual_inversion_template_names())
create_refresh_button(hn_template, textual_inversion.list_textual_inversion_templates, lambda: {"choices": get_textual_inversion_template_names()}, "refrsh_train_template_file")
hn_shuffle_tags = gr.Checkbox(label="Shuffle tags by ',' when creating prompts.", value=False)
hn_tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts", value=0)
hn_latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once", choices=['once', 'deterministic', 'random'])
with gr.Box():
gr.HTML('<h2>Training outputs</h2>')
with FormRow():
hn_create_every = gr.Number(label='Create interim images', value=500, precision=0)
hn_save_every = gr.Number(label='Create interim hypernetworks', value=500, precision=0)
hn_preview_from_txt2img = gr.Checkbox(label='Use current settings for previews', value=False)
hn_log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value=f"{os.path.join('cmd_opts.data_dir', 'train/log/embeddings')}")
with gr.Box():
gr.HTML('<h2>Training outputs</h2>')
with FormRow():
hn_create_every = gr.Number(label='Create interim images', value=500, precision=0)
hn_save_every = gr.Number(label='Create interim hypernetworks', value=500, precision=0)
hn_preview_from_txt2img = gr.Checkbox(label='Use current settings for previews', value=False)
hn_log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value=f"{os.path.join('cmd_opts.data_dir', 'train/log/embeddings')}")
hn_stop.click(fn=lambda: shared.state.interrupt(), inputs=[], outputs=[])
hn_stop.click(fn=lambda: shared.state.interrupt(), inputs=[], outputs=[])
hn_create.click(
fn=modules.hypernetworks.ui.create_hypernetwork,
inputs=[
hn_new_name,
hn_new_sizes,
hn_overwrite,
hn_new_layer_structure,
hn_new_activation_func,
hn_new_initialization_option,
hn_new_add_layer_norm,
hn_new_use_dropout,
hn_new_dropout_structure
],
outputs=[
hn_name,
train_output,
train_outcome,
]
)
hn_create.click(
fn=modules.hypernetworks.ui.create_hypernetwork,
inputs=[
hn_new_name,
hn_new_sizes,
hn_overwrite,
hn_new_layer_structure,
hn_new_activation_func,
hn_new_initialization_option,
hn_new_add_layer_norm,
hn_new_use_dropout,
hn_new_dropout_structure
],
outputs=[
hn_name,
train_output,
train_outcome,
]
)
hn_train.click(
fn=wrap_gradio_gpu_call(modules.hypernetworks.ui.train_hypernetwork, extra_outputs=[gr.update()]),
_js="startTrainMonitor",
inputs=[
dummy_component,
hn_name,
hn_learn_rate,
hn_batch_size,
hn_gradient_step,
hn_dataset_directory,
hn_log_directory,
hn_width,
hn_height,
hn_varsize,
hn_steps,
hn_clip_grad_mode,
hn_clip_grad_value,
hn_shuffle_tags,
hn_tag_drop_out,
hn_latent_sampling_method,
hn_use_weight,
hn_create_every,
hn_save_every,
hn_template,
hn_preview_from_txt2img,
*txt2img_preview_params,
],
outputs=[
train_output,
train_outcome,
]
)
hn_train.click(
fn=wrap_gradio_gpu_call(modules.hypernetworks.ui.train_hypernetwork, extra_outputs=[gr.update()]),
_js="startTrainMonitor",
inputs=[
dummy_component,
hn_name,
hn_learn_rate,
hn_batch_size,
hn_gradient_step,
hn_dataset_directory,
hn_log_directory,
hn_width,
hn_height,
hn_varsize,
hn_steps,
hn_clip_grad_mode,
hn_clip_grad_value,
hn_shuffle_tags,
hn_tag_drop_out,
hn_latent_sampling_method,
hn_use_weight,
hn_create_every,
hn_save_every,
hn_template,
hn_preview_from_txt2img,
*txt2img_preview_params,
],
outputs=[
train_output,
train_outcome,
]
)
params = script_callbacks.UiTrainTabParams(txt2img_preview_params)
script_callbacks.ui_train_tabs_callback(params)
+1 -1
View File
@@ -29,7 +29,7 @@ def create_ui():
with FormRow(variant='compact', elem_id="txt2img_extra_networks", visible=False) as extra_networks_ui:
from modules import ui_extra_networks
extra_networks_ui = ui_extra_networks.create_ui(extra_networks_ui, txt2img_extra_networks_button, 'txt2img', skip_indexing=shared.opts.extra_network_skip_indexing)
timer.startup.record('ui-extra-networks')
timer.startup.record('ui-en')
with gr.Row(elem_id="txt2img_interface", equal_height=False):
with gr.Column(variant='compact', elem_id="txt2img_settings"):
+3 -3
View File
@@ -122,7 +122,7 @@ def initialize():
ui_extra_networks.register_pages()
extra_networks.initialize()
extra_networks.register_default_extra_networks()
timer.startup.record("extra-networks")
timer.startup.record("networks")
if cmd_opts.tls_keyfile is not None and cmd_opts.tls_certfile is not None:
try:
@@ -311,7 +311,6 @@ def webui(restart=False):
if cmd_opts.profile:
for k, v in modules.script_callbacks.callback_map.items():
shared.log.debug(f'Registered callbacks: {k}={len(v)} {[c.script for c in v]}')
log.info(f"Startup time: {timer.startup.summary()}")
debug = log.trace if os.environ.get('SD_SCRIPT_DEBUG', None) is not None else lambda *args, **kwargs: None
debug('Trace: SCRIPTS')
for m in modules.scripts.scripts_data:
@@ -319,8 +318,9 @@ def webui(restart=False):
debug('Loaded postprocessing scripts:')
for m in modules.scripts.postprocessing_scripts_data:
debug(f' {m}')
timer.startup.reset()
modules.script_callbacks.print_timers()
log.info(f"Startup time: {timer.startup.summary()}")
timer.startup.reset()
if not restart:
# override all loggers to use the same handlers as the main logger