extensions/Hypernetwork-MonkeyPatch-Extension/patches/external_pr/ui.py

import gc
import html
import json
import os
import random

from modules import shared, sd_hijack, devices
from modules.call_queue import wrap_gradio_call
from modules.paths import script_path
from modules.ui import create_refresh_button, gr_show
from webui import wrap_gradio_gpu_call
from .textual_inversion import train_embedding as train_embedding_external
from .hypernetwork import train_hypernetwork as train_hypernetwork_external, train_hypernetwork_tuning
import gradio as gr


def train_hypernetwork_ui(*args):
    initial_hypernetwork = None
    if hasattr(shared, 'loaded_hypernetwork'):
        initial_hypernetwork = shared.loaded_hypernetwork
    else:
        shared.loaded_hypernetworks = []
    assert not shared.cmd_opts.lowvram, 'Training models with lowvram is not possible'

    try:
        sd_hijack.undo_optimizations()

        hypernetwork, filename = train_hypernetwork_external(*args)

        res = f"""
Training {'interrupted' if shared.state.interrupted else 'finished'} at {hypernetwork.step} steps.
Hypernetwork saved to {html.escape(filename)}
"""
        return res, ""
    except Exception:
        raise
    finally:
        if hasattr(shared, 'loaded_hypernetwork'):
            shared.loaded_hypernetwork = initial_hypernetwork
        else:
            shared.loaded_hypernetworks = []
        # check hypernetwork is bounded then delete it
        if locals().get('hypernetwork', None) is not None:
            del hypernetwork
        gc.collect()
        shared.sd_model.cond_stage_model.to(devices.device)
        shared.sd_model.first_stage_model.to(devices.device)
        sd_hijack.apply_optimizations()


def train_hypernetwork_ui_tuning(*args):
    initial_hypernetwork = None
    if hasattr(shared, 'loaded_hypernetwork'):
        initial_hypernetwork = shared.loaded_hypernetwork
    else:
        shared.loaded_hypernetworks = []

    assert not shared.cmd_opts.lowvram, 'Training models with lowvram is not possible'

    try:
        sd_hijack.undo_optimizations()

        train_hypernetwork_tuning(*args)

        res = f"""
Training {'interrupted' if shared.state.interrupted else 'finished'}.
"""
        return res, ""
    except Exception:
        raise
    finally:
        if hasattr(shared, 'loaded_hypernetwork'):
            shared.loaded_hypernetwork = initial_hypernetwork
        else:
            shared.loaded_hypernetworks = []
        shared.sd_model.cond_stage_model.to(devices.device)
        shared.sd_model.first_stage_model.to(devices.device)
        sd_hijack.apply_optimizations()


def save_training_setting(*args):
    save_file_name, learn_rate, batch_size, gradient_step, training_width, \
    training_height, steps, shuffle_tags, tag_drop_out, latent_sampling_method, \
    template_file, use_beta_scheduler, beta_repeat_epoch, epoch_mult, warmup, min_lr, \
    gamma_rate, use_beta_adamW_checkbox, save_when_converge, create_when_converge, \
    adamw_weight_decay, adamw_beta_1, adamw_beta_2, adamw_eps, show_gradient_clip_checkbox, \
    gradient_clip_opt, optional_gradient_clip_value, optional_gradient_norm_type, latent_sampling_std,\
    noise_training_scheduler_enabled, noise_training_scheduler_repeat, noise_training_scheduler_cycle, loss_opt, use_dadaptation, dadapt_growth_factor, use_weight = args
    dumped_locals = locals()
    dumped_locals.pop('args')
    filename = (str(random.randint(0, 1024)) if save_file_name == '' else save_file_name) + '_train_' + '.json'
    filename = os.path.join(shared.cmd_opts.hypernetwork_dir, filename)
    with open(filename, 'w') as file:
        print(dumped_locals)
        json.dump(dumped_locals, file)
        print(f"File saved as {filename}")
    return filename, ""


def save_hypernetwork_setting(*args):
    save_file_name, enable_sizes, overwrite_old, layer_structure, activation_func, weight_init, add_layer_norm, use_dropout, dropout_structure, optional_info, weight_init_seed, normal_std, skip_connection = args
    dumped_locals = locals()
    dumped_locals.pop('args')
    filename = (str(random.randint(0, 1024)) if save_file_name == '' else save_file_name) + '_hypernetwork_' + '.json'
    filename = os.path.join(shared.cmd_opts.hypernetwork_dir, filename)
    with open(filename, 'w') as file:
        print(dumped_locals)
        json.dump(dumped_locals, file)
        print(f"File saved as {filename}")
    return filename, ""


def on_train_gamma_tab(params=None):
    dummy_component = gr.Label(visible=False)
    with gr.Tab(label="Train_Gamma") as train_gamma:
        gr.HTML(
            value="<p style='margin-bottom: 0.7em'>Train an embedding or Hypernetwork; you must specify a directory <a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\" style=\"font-weight:bold;\">[wiki]</a></p>")
        with gr.Row():
            train_embedding_name = gr.Dropdown(label='Embedding', elem_id="train_embedding", choices=sorted(
                sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))
            create_refresh_button(train_embedding_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.Row():
            train_hypernetwork_name = gr.Dropdown(label='Hypernetwork', elem_id="train_hypernetwork",
                                                  choices=[x for x in shared.hypernetworks.keys()])
            create_refresh_button(train_hypernetwork_name, shared.reload_hypernetworks,
                                  lambda: {"choices": sorted([x for x in shared.hypernetworks.keys()])},
                                  "refresh_train_hypernetwork_name")
        with gr.Row():
            embedding_learn_rate = gr.Textbox(label='Embedding Learning rate',
                                              placeholder="Embedding Learning rate", value="0.005")
            hypernetwork_learn_rate = gr.Textbox(label='Hypernetwork Learning rate',
                                                 placeholder="Hypernetwork Learning rate", value="0.00004")
            use_beta_scheduler_checkbox = gr.Checkbox(
                label='Show advanced learn rate scheduler options')
            use_beta_adamW_checkbox = gr.Checkbox(
                label='Show advanced adamW parameter options)')
            show_gradient_clip_checkbox = gr.Checkbox(
                label='Show Gradient Clipping Options(for both)')
            show_noise_options = gr.Checkbox(
                label='Show Noise Scheduler Options(for both)')
        with gr.Row(visible=False) as adamW_options:
            use_dadaptation = gr.Checkbox(label="Uses D-Adaptation(LR Free) AdamW. Recommended LR is 1.0 at base")
            adamw_weight_decay = gr.Textbox(label="AdamW weight decay parameter", placeholder="default = 0.01",
                                            value="0.01")
            adamw_beta_1 = gr.Textbox(label="AdamW beta1 parameter", placeholder="default = 0.9", value="0.9")
            adamw_beta_2 = gr.Textbox(label="AdamW beta2 parameter", placeholder="default = 0.99", value="0.99")
            adamw_eps = gr.Textbox(label="AdamW epsilon parameter", placeholder="default = 1e-8", value="1e-8")
            with gr.Row(visible=False) as dadapt_growth_options:
                dadapt_growth_factor = gr.Number(value=-1, label='Growth factor limiting, use value like 1.02 or leave it as -1')
        with gr.Row(visible=False) as beta_scheduler_options:
            use_beta_scheduler = gr.Checkbox(label='Use CosineAnnealingWarmupRestarts Scheduler')
            beta_repeat_epoch = gr.Textbox(label='Steps for cycle', placeholder="Cycles every nth Step", value="64")
            epoch_mult = gr.Textbox(label='Step multiplier per cycle', placeholder="Step length multiplier every cycle",
                                    value="1")
            warmup = gr.Textbox(label='Warmup step per cycle', placeholder="CosineAnnealing lr increase step",
                                value="5")
            min_lr = gr.Textbox(label='Minimum learning rate',
                                placeholder="restricts decay value, but does not restrict gamma rate decay",
                                value="6e-7")
            gamma_rate = gr.Textbox(label='Decays learning rate every cycle',
                                    placeholder="Value should be in (0-1]", value="1")
        with gr.Row(visible=False) as beta_scheduler_options2:
            save_converge_opt = gr.Checkbox(label="Saves when every cycle finishes")
            generate_converge_opt = gr.Checkbox(label="Generates image when every cycle finishes")
        with gr.Row(visible=False) as gradient_clip_options:
            gradient_clip_opt = gr.Radio(label="Gradient Clipping Options", choices=["None", "limit", "norm"])
            optional_gradient_clip_value = gr.Textbox(label="Limiting value", value="1e-1")
            optional_gradient_norm_type = gr.Textbox(label="Norm type", value="2")
        with gr.Row(visible=False) as noise_scheduler_options:
            noise_training_scheduler_enabled = gr.Checkbox(label="Use Noise training scheduler(test)")
            noise_training_scheduler_repeat = gr.Checkbox(label="Restarts noise scheduler, or linear")
            noise_training_scheduler_cycle = gr.Number(label="Restarts noise scheduler every nth epoch")
        use_weight = gr.Checkbox(label="Uses image alpha(transparency) channel for adjusting loss")
        # change by feedback
        use_dadaptation.change(
            fn=lambda show: gr_show(show),
            inputs=[use_dadaptation],
            outputs=[dadapt_growth_options]
        )
        show_noise_options.change(
            fn = lambda show:gr_show(show),
            inputs = [show_noise_options],
            outputs = [noise_scheduler_options]
        )
        use_beta_adamW_checkbox.change(
            fn=lambda show: gr_show(show),
            inputs=[use_beta_adamW_checkbox],
            outputs=[adamW_options],
        )
        use_beta_scheduler_checkbox.change(
            fn=lambda show: gr_show(show),
            inputs=[use_beta_scheduler_checkbox],
            outputs=[beta_scheduler_options],
        )
        use_beta_scheduler_checkbox.change(
            fn=lambda show: gr_show(show),
            inputs=[use_beta_scheduler_checkbox],
            outputs=[beta_scheduler_options2],
        )
        show_gradient_clip_checkbox.change(
            fn=lambda show: gr_show(show),
            inputs=[show_gradient_clip_checkbox],
            outputs=[gradient_clip_options],
        )
        move_optim_when_generate = gr.Checkbox(label="Unload Optimizer when generating preview(hypernetwork)",
                                               value=True)
        batch_size = gr.Number(label='Batch size', value=1, precision=0)
        gradient_step = gr.Number(label='Gradient accumulation steps', value=1, precision=0)
        dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
        log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs",
                                   value="textual_inversion")
        template_file = gr.Textbox(label='Prompt template file',
                                   value=os.path.join(script_path, "textual_inversion_templates",
                                                      "style_filewords.txt"))
        training_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)
        training_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)
        steps = gr.Number(label='Max steps', value=100000, precision=0)
        create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable',
                                       value=500, precision=0)
        save_embedding_every = gr.Number(
            label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0)
        save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True)
        preview_from_txt2img = gr.Checkbox(
            label='Read parameters (prompt, etc...) from txt2img tab when making previews', value=False)
        with gr.Row():
            shuffle_tags = gr.Checkbox(label="Shuffle tags by ',' when creating prompts.", value=False)
            tag_drop_out = gr.Slider(minimum=0, maximum=1, step=0.1, label="Drop out tags when creating prompts.",
                                     value=0)
        with gr.Row():
            latent_sampling_method = gr.Radio(label='Choose latent sampling method', value="once",
                                              choices=['once', 'deterministic', 'random'])
            latent_sampling_std_value = gr.Number(label="Standard deviation for sampling", value=-1)
        with gr.Row():
            loss_opt = gr.Radio(label="loss type", value="loss",
                                choices=['loss', 'loss_simple', 'loss_vlb'])
        with gr.Row():
            save_training_option = gr.Button(value="Save training setting")
            save_file_name = gr.Textbox(label="File name to save setting as", value="")
            load_training_option = gr.Textbox(
                label="Load training option from saved json file. This will override settings above", value="")
        with gr.Row():
            interrupt_training = gr.Button(value="Interrupt")
            train_hypernetwork = gr.Button(value="Train Hypernetwork", variant='primary')
            train_embedding = gr.Button(value="Train Embedding", variant='primary')
        ti_output = gr.Text(elem_id="ti_output3", value="", show_label=False)
        ti_outcome = gr.HTML(elem_id="ti_error3", value="")

    # Full path to .json or simple names are recommended.
    save_training_option.click(
        fn=wrap_gradio_call(save_training_setting),
        inputs=[
            save_file_name,
            hypernetwork_learn_rate,
            batch_size,
            gradient_step,
            training_width,
            training_height,
            steps,
            shuffle_tags,
            tag_drop_out,
            latent_sampling_method,
            template_file,
            use_beta_scheduler,
            beta_repeat_epoch,
            epoch_mult,
            warmup,
            min_lr,
            gamma_rate,
            use_beta_adamW_checkbox,
            save_converge_opt,
            generate_converge_opt,
            adamw_weight_decay,
            adamw_beta_1,
            adamw_beta_2,
            adamw_eps,
            show_gradient_clip_checkbox,
            gradient_clip_opt,
            optional_gradient_clip_value,
            optional_gradient_norm_type,
            latent_sampling_std_value,
        noise_training_scheduler_enabled,
        noise_training_scheduler_repeat,
        noise_training_scheduler_cycle,
        loss_opt,
        use_dadaptation,
        dadapt_growth_factor,
        use_weight
        ],
        outputs=[
            ti_output,
            ti_outcome,
        ]
    )
    train_embedding.click(
        fn=wrap_gradio_gpu_call(train_embedding_external, extra_outputs=[gr.update()]),
        _js="start_training_textual_inversion",
        inputs=[
            dummy_component,
            train_embedding_name,
            embedding_learn_rate,
            batch_size,
            gradient_step,
            dataset_directory,
            log_directory,
            training_width,
            training_height,
            steps,
            shuffle_tags,
            tag_drop_out,
            latent_sampling_method,
            create_image_every,
            save_embedding_every,
            template_file,
            save_image_with_stored_embedding,
            preview_from_txt2img,
            *params.txt2img_preview_params,
            use_beta_scheduler,
            beta_repeat_epoch,
            epoch_mult,
            warmup,
            min_lr,
            gamma_rate,
            save_converge_opt,
            generate_converge_opt,
            move_optim_when_generate,
            use_beta_adamW_checkbox,
            adamw_weight_decay,
            adamw_beta_1,
            adamw_beta_2,
            adamw_eps,
            show_gradient_clip_checkbox,
            gradient_clip_opt,
            optional_gradient_clip_value,
            optional_gradient_norm_type,
            latent_sampling_std_value,
            use_weight
        ],
        outputs=[
            ti_output,
            ti_outcome,
        ]
    )

    train_hypernetwork.click(
        fn=wrap_gradio_gpu_call(train_hypernetwork_ui, extra_outputs=[gr.update()]),
        _js="start_training_textual_inversion",
        inputs=[
            dummy_component,
            train_hypernetwork_name,
            hypernetwork_learn_rate,
            batch_size,
            gradient_step,
            dataset_directory,
            log_directory,
            training_width,
            training_height,
            steps,
            shuffle_tags,
            tag_drop_out,
            latent_sampling_method,
            create_image_every,
            save_embedding_every,
            template_file,
            preview_from_txt2img,
            *params.txt2img_preview_params,
            use_beta_scheduler,
            beta_repeat_epoch,
            epoch_mult,
            warmup,
            min_lr,
            gamma_rate,
            save_converge_opt,
            generate_converge_opt,
            move_optim_when_generate,
            use_beta_adamW_checkbox,
            adamw_weight_decay,
            adamw_beta_1,
            adamw_beta_2,
            adamw_eps,
            show_gradient_clip_checkbox,
            gradient_clip_opt,
            optional_gradient_clip_value,
            optional_gradient_norm_type,
            latent_sampling_std_value,
        noise_training_scheduler_enabled,
        noise_training_scheduler_repeat,
        noise_training_scheduler_cycle,
            load_training_option,
            loss_opt,
            use_dadaptation,
            dadapt_growth_factor,
            use_weight
        ],
        outputs=[
            ti_output,
            ti_outcome,
        ]
    )

    interrupt_training.click(
        fn=lambda: shared.state.interrupt(),
        inputs=[],
        outputs=[],
    )
    return [(train_gamma, "Train Gamma", "train_gamma")]


def on_train_tuning(params=None):
    dummy_component = gr.Label(visible=False)
    with gr.Tab(label="Train_Tuning") as train_tuning:
        gr.HTML(
            value="<p style='margin-bottom: 0.7em'>Train Hypernetwork; you must specify a directory <a href=\"https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Textual-Inversion\" style=\"font-weight:bold;\">[wiki]</a></p>")
        with gr.Row():
            train_hypernetwork_name = gr.Dropdown(label='Hypernetwork', elem_id="train_hypernetwork",
                                                  choices=[x for x in shared.hypernetworks.keys()])
            create_refresh_button(train_hypernetwork_name, shared.reload_hypernetworks,
                                  lambda: {"choices": sorted([x for x in shared.hypernetworks.keys()])},
                                  "refresh_train_hypernetwork_name")
            optional_new_hypernetwork_name = gr.Textbox(
                label="Hypernetwork name to create, leave it empty to use selected", value="")
        with gr.Row():
            load_hypernetworks_option = gr.Textbox(
                label="Load Hypernetwork creation option from saved json file",
                placeholder=". filename cannot have ',' inside, and files should be splitted by ','.", value="")
        with gr.Row():
            load_training_options = gr.Textbox(
                label="Load training option(s) from saved json file",
                placeholder=". filename cannot have ',' inside, and files should be splitted by ','.", value="")
        move_optim_when_generate = gr.Checkbox(label="Unload Optimizer when generating preview(hypernetwork)",
                                               value=True)
        dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")
        log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs",
                                   value="textual_inversion")
        create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable',
                                       value=500, precision=0)
        save_model_every = gr.Number(
            label='Save a copy of model to log directory every N steps, 0 to disable', value=500, precision=0)
        preview_from_txt2img = gr.Checkbox(
            label='Read parameters (prompt, etc...) from txt2img tab when making previews', value=False)
        manual_dataset_seed = gr.Number(
            label="Manual dataset seed", value=-1, precision=0
        )
        with gr.Row():
            interrupt_training = gr.Button(value="Interrupt")
            train_hypernetwork = gr.Button(value="Train Hypernetwork", variant='primary')
        ti_output = gr.Text(elem_id="ti_output4", value="", show_label=False)
        ti_outcome = gr.HTML(elem_id="ti_error4", value="")
    train_hypernetwork.click(
        fn=wrap_gradio_gpu_call(train_hypernetwork_ui_tuning, extra_outputs=[gr.update()]),
        _js="start_training_textual_inversion",
        inputs=[
            dummy_component,
            train_hypernetwork_name,
            dataset_directory,
            log_directory,
            create_image_every,
            save_model_every,
            preview_from_txt2img,
            *params.txt2img_preview_params,
            move_optim_when_generate,
            optional_new_hypernetwork_name,
            load_hypernetworks_option,
            load_training_options,
            manual_dataset_seed
        ],
        outputs=[
            ti_output,
            ti_outcome,
        ]
    )

    interrupt_training.click(
        fn=lambda: shared.state.interrupt(),
        inputs=[],
        outputs=[],
    )
    return [(train_tuning, "Train Tuning", "train_tuning")]