remove training

modules/hypernetworks/hypernetwork.py

@@ -1,23 +1,13 @@
-import datetime
-import html
 import os
-from collections import deque
 import inspect
 from statistics import stdev, mean
 from rich import progress
-import tqdm
 import torch
 from torch import einsum
 from torch.nn.init import normal_, xavier_normal_, xavier_uniform_, kaiming_normal_, kaiming_uniform_, zeros_
 from einops import rearrange, repeat
 from ldm.util import default
-from modules import devices, processing, sd_models, shared, hashes, errors, files_cache
-import modules.textual_inversion.dataset
-from modules.textual_inversion import textual_inversion, ti_logging
-from modules.textual_inversion.learn_schedule import LearnRateScheduler
-
-
-optimizer_dict = {optim_name : cls_obj for optim_name, cls_obj in inspect.getmembers(torch.optim, inspect.isclass) if optim_name != "Optimizer"}
+from modules import devices, shared, hashes, errors, files_cache
 
 
 class HypernetworkModule(torch.nn.Module):
@@ -410,341 +400,3 @@ def report_statistics(loss_info:dict):
             print(recent)
         except Exception as e:
             print(e)
-
-
-def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None):
-    # Remove illegal characters from name.
-    name = "".join( x for x in name if (x.isalnum() or x in "._- "))
-    assert name, "Name cannot be empty!"
-    fn = os.path.join(shared.opts.hypernetwork_dir, f"{name}.pt")
-    if not overwrite_old:
-        assert not os.path.exists(fn), f"file {fn} already exists"
-    if type(layer_structure) == str:
-        layer_structure = [float(x.strip()) for x in layer_structure.split(",")]
-    if use_dropout and dropout_structure and type(dropout_structure) == str:
-        dropout_structure = [float(x.strip()) for x in dropout_structure.split(",")]
-    else:
-        dropout_structure = [0] * len(layer_structure)
-    hypernet = modules.hypernetworks.hypernetwork.Hypernetwork(
-        name=name,
-        enable_sizes=[int(x) for x in enable_sizes],
-        layer_structure=layer_structure,
-        activation_func=activation_func,
-        weight_init=weight_init,
-        add_layer_norm=add_layer_norm,
-        use_dropout=use_dropout,
-        dropout_structure=dropout_structure
-    )
-    hypernet.save(fn)
-    shared.reload_hypernetworks()
-    return name
-
-
-def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, use_weight, create_image_every, save_hypernetwork_every, template_filename, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height): # pylint: disable=unused-argument
-    # images allows training previews to have infotext. Importing it at the top causes a circular import problem.
-    from modules import images, sd_hijack_checkpoint
-
-    save_hypernetwork_every = save_hypernetwork_every or 0
-    create_image_every = create_image_every or 0
-    template_file = textual_inversion.textual_inversion_templates.get(template_filename, None)
-    textual_inversion.validate_train_inputs(hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_hypernetwork_every, create_image_every, name="hypernetwork")
-    template_file = template_file.path
-
-    path = shared.hypernetworks.get(hypernetwork_name, None)
-    hypernetwork = Hypernetwork()
-    hypernetwork.load(path)
-    shared.loaded_hypernetworks = [hypernetwork]
-
-    shared.state.job = "train"
-    shared.state.textinfo = "Initializing hypernetwork training..."
-    shared.state.job_count = steps
-
-    hypernetwork_name = hypernetwork_name.rsplit('(', 1)[0]
-    filename = os.path.join(shared.opts.hypernetwork_dir, f'{hypernetwork_name}.pt')
-
-    log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), hypernetwork_name)
-    unload = shared.opts.unload_models_when_training
-
-    if save_hypernetwork_every > 0:
-        hypernetwork_dir = os.path.join(log_directory, "hypernetworks")
-        os.makedirs(hypernetwork_dir, exist_ok=True)
-    else:
-        hypernetwork_dir = None
-
-    if create_image_every > 0:
-        images_dir = os.path.join(log_directory, "images")
-        os.makedirs(images_dir, exist_ok=True)
-    else:
-        images_dir = None
-
-    checkpoint = sd_models.select_checkpoint()
-
-    initial_step = hypernetwork.step or 0
-    if initial_step >= steps:
-        shared.state.textinfo = "Model has already been trained beyond specified max steps"
-        return hypernetwork, filename
-
-    scheduler = LearnRateScheduler(learn_rate, steps, initial_step)
-
-    clip_grad = torch.nn.utils.clip_grad_value_ if clip_grad_mode == "value" else torch.nn.utils.clip_grad_norm_ if clip_grad_mode == "norm" else None
-    if clip_grad:
-        clip_grad_sched = LearnRateScheduler(clip_grad_value, steps, initial_step, verbose=False)
-
-    if shared.opts.training_enable_tensorboard:
-        tensorboard_writer = textual_inversion.tensorboard_setup(log_directory)
-
-    # dataset loading may take a while, so input validations and early returns should be done before this
-    shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
-
-    pin_memory = shared.opts.pin_memory
-
-    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
-
-    if shared.opts.save_training_settings_to_txt:
-        saved_params = dict(
-            model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds),
-            **{field: getattr(hypernetwork, field) for field in ['layer_structure', 'activation_func', 'weight_init', 'add_layer_norm', 'use_dropout', ]}
-        )
-        ti_logging.save_settings_to_file(log_directory, {**saved_params, **locals()})
-
-    latent_sampling_method = ds.latent_sampling_method
-
-    dl = modules.textual_inversion.dataset.PersonalizedDataLoader(ds, latent_sampling_method=latent_sampling_method, batch_size=ds.batch_size, pin_memory=pin_memory)
-
-    old_parallel_processing_allowed = shared.parallel_processing_allowed
-
-    if unload:
-        shared.parallel_processing_allowed = False
-        shared.sd_model.cond_stage_model.to(devices.cpu)
-        shared.sd_model.first_stage_model.to(devices.cpu)
-
-    weights = hypernetwork.weights()
-    hypernetwork.train()
-
-    # Here we use optimizer from saved HN, or we can specify as UI option.
-    if hypernetwork.optimizer_name in optimizer_dict:
-        optimizer = optimizer_dict[hypernetwork.optimizer_name](params=weights, lr=scheduler.learn_rate)
-        optimizer_name = hypernetwork.optimizer_name
-    else:
-        print(f"Optimizer type {hypernetwork.optimizer_name} is not defined!")
-        optimizer = torch.optim.AdamW(params=weights, lr=scheduler.learn_rate)
-        optimizer_name = 'AdamW'
-
-    if hypernetwork.optimizer_state_dict:  # This line must be changed if Optimizer type can be different from saved optimizer.
-        try:
-            optimizer.load_state_dict(hypernetwork.optimizer_state_dict)
-        except RuntimeError as e:
-            print("Cannot resume from saved optimizer!")
-            print(e)
-
-    scaler = torch.cuda.amp.GradScaler()
-
-    batch_size = ds.batch_size
-    gradient_step = ds.gradient_step
-    # n steps = batch_size * gradient_step * n image processed
-    steps_per_epoch = len(ds) // batch_size // gradient_step
-    max_steps_per_epoch = len(ds) // batch_size - (len(ds) // batch_size) % gradient_step
-    loss_step = 0
-    _loss_step = 0 #internal
-    # size = len(ds.indexes)
-    # loss_dict = defaultdict(lambda : deque(maxlen = 1024))
-    loss_logging = deque(maxlen=len(ds) * 3)  # this should be configurable parameter, this is 3 * epoch(dataset size)
-    # losses = torch.zeros((size,))
-    # previous_mean_losses = [0]
-    # previous_mean_loss = 0
-    # print("Mean loss of {} elements".format(size))
-
-    _steps_without_grad = 0
-
-    last_saved_file = "<none>"
-    last_saved_image = "<none>"
-    forced_filename = "<none>"
-
-    pbar = tqdm.tqdm(total=steps - initial_step)
-    try:
-        sd_hijack_checkpoint.add()
-
-        for _i in range((steps-initial_step) * gradient_step):
-            if scheduler.finished:
-                break
-            if shared.state.interrupted:
-                break
-            for j, batch in enumerate(dl):
-                # works as a drop_last=True for gradient accumulation
-                if j == max_steps_per_epoch:
-                    break
-                scheduler.apply(optimizer, hypernetwork.step)
-                if scheduler.finished:
-                    break
-                if shared.state.interrupted:
-                    break
-
-                if clip_grad:
-                    clip_grad_sched.step(hypernetwork.step)
-
-                with devices.autocast():
-                    x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)
-                    if use_weight:
-                        w = batch.weight.to(devices.device, non_blocking=pin_memory)
-                    if tag_drop_out != 0 or shuffle_tags:
-                        shared.sd_model.cond_stage_model.to(devices.device)
-                        c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
-                        shared.sd_model.cond_stage_model.to(devices.cpu)
-                    else:
-                        c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory)
-                    if use_weight:
-                        loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step
-                        del w
-                    else:
-                        loss = shared.sd_model.forward(x, c)[0] / gradient_step
-                    del x
-                    del c
-                    _loss_step += loss.item()
-
-                scaler.scale(loss).backward()
-                # go back until we reach gradient accumulation steps
-                if (j + 1) % gradient_step != 0:
-                    continue
-                loss_logging.append(_loss_step)
-                if clip_grad:
-                    clip_grad(weights, clip_grad_sched.learn_rate)
-
-                scaler.step(optimizer)
-                scaler.update()
-                hypernetwork.step += 1
-                pbar.update()
-                optimizer.zero_grad(set_to_none=True)
-                loss_step = _loss_step
-                _loss_step = 0
-                steps_done = hypernetwork.step + 1
-                epoch_num = hypernetwork.step // steps_per_epoch
-                epoch_step = hypernetwork.step % steps_per_epoch
-
-                description = f"Training hypernetwork [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}"
-                pbar.set_description(description)
-                if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0:
-                    # Before saving, change name to match current checkpoint.
-                    hypernetwork_name_every = f'{hypernetwork_name}-{steps_done}'
-                    last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork_name_every}.pt')
-                    hypernetwork.optimizer_name = optimizer_name
-                    if shared.opts.save_optimizer_state:
-                        hypernetwork.optimizer_state_dict = optimizer.state_dict()
-                    save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, last_saved_file)
-                    hypernetwork.optimizer_state_dict = None  # dereference it after saving, to save memory.
-
-
-
-                if shared.opts.training_enable_tensorboard:
-                    epoch_num = hypernetwork.step // len(ds)
-                    epoch_step = hypernetwork.step - (epoch_num * len(ds)) + 1
-                    mean_loss = sum(loss_logging) / len(loss_logging)
-                    textual_inversion.tensorboard_add(tensorboard_writer, loss=mean_loss, global_step=hypernetwork.step, step=epoch_step, learn_rate=scheduler.learn_rate, epoch_num=epoch_num)
-
-                textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, steps_per_epoch, {
-                    "loss": f"{loss_step:.7f}",
-                    "learn_rate": scheduler.learn_rate
-                })
-
-                if images_dir is not None and steps_done % create_image_every == 0:
-                    forced_filename = f'{hypernetwork_name}-{steps_done}'
-                    last_saved_image = os.path.join(images_dir, forced_filename)
-                    hypernetwork.eval()
-                    rng_state = torch.get_rng_state()
-                    cuda_rng_state = None
-                    cuda_rng_state = torch.cuda.get_rng_state_all()
-                    shared.sd_model.cond_stage_model.to(devices.device)
-                    shared.sd_model.first_stage_model.to(devices.device)
-
-                    p = processing.StableDiffusionProcessingTxt2Img(
-                        sd_model=shared.sd_model,
-                        do_not_save_grid=True,
-                        do_not_save_samples=True,
-                    )
-
-                    p.disable_extra_networks = True
-
-                    if preview_from_txt2img:
-                        p.prompt = preview_prompt
-                        p.negative_prompt = preview_negative_prompt
-                        p.steps = preview_steps
-                        p.sampler_name = processing.get_sampler_name(preview_sampler_index)
-                        p.cfg_scale = preview_cfg_scale
-                        p.seed = preview_seed
-                        p.width = preview_width
-                        p.height = preview_height
-                    else:
-                        p.prompt = batch.cond_text[0]
-                        p.steps = 20
-                        p.width = training_width
-                        p.height = training_height
-
-                    preview_text = p.prompt
-
-                    processed = processing.process_images(p)
-                    image = processed.images[0] if len(processed.images) > 0 else None
-
-                    if unload:
-                        shared.sd_model.cond_stage_model.to(devices.cpu)
-                        shared.sd_model.first_stage_model.to(devices.cpu)
-                    torch.set_rng_state(rng_state)
-                    torch.cuda.set_rng_state_all(cuda_rng_state)
-                    hypernetwork.train()
-                    if image is not None:
-                        shared.state.assign_current_image(image)
-                        if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images:
-                            textual_inversion.tensorboard_add_image(tensorboard_writer,
-                                                                    f"Validation at epoch {epoch_num}", image,
-                                                                    hypernetwork.step)
-                        last_saved_image, _last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)
-                        last_saved_image += f", prompt: {preview_text}"
-
-                shared.state.job_no = hypernetwork.step
-
-                shared.state.textinfo = f"""
-<p>
-Loss: {loss_step:.7f}<br/>
-Step: {steps_done}<br/>
-Last prompt: {html.escape(batch.cond_text[0])}<br/>
-Last saved hypernetwork: {html.escape(last_saved_file)}<br/>
-Last saved image: {html.escape(last_saved_image)}<br/>
-</p>
-"""
-    except Exception as e:
-        errors.display(e, 'hypernetwork train')
-    finally:
-        pbar.leave = False
-        pbar.close()
-        hypernetwork.eval()
-        #report_statistics(loss_dict)
-        sd_hijack_checkpoint.remove()
-
-
-
-    filename = os.path.join(shared.opts.hypernetwork_dir, f'{hypernetwork_name}.pt')
-    hypernetwork.optimizer_name = optimizer_name
-    if shared.opts.save_optimizer_state:
-        hypernetwork.optimizer_state_dict = optimizer.state_dict()
-    save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename)
-
-    del optimizer
-    hypernetwork.optimizer_state_dict = None  # dereference it after saving, to save memory.
-    shared.sd_model.cond_stage_model.to(devices.device)
-    shared.sd_model.first_stage_model.to(devices.device)
-    shared.parallel_processing_allowed = old_parallel_processing_allowed
-
-    return hypernetwork, filename
-
-def save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename):
-    old_hypernetwork_name = hypernetwork.name
-    old_sd_checkpoint = hypernetwork.sd_checkpoint if hasattr(hypernetwork, "sd_checkpoint") else None
-    old_sd_checkpoint_name = hypernetwork.sd_checkpoint_name if hasattr(hypernetwork, "sd_checkpoint_name") else None
-    try:
-        hypernetwork.sd_checkpoint = checkpoint.shorthash
-        hypernetwork.sd_checkpoint_name = checkpoint.model_name
-        hypernetwork.name = hypernetwork_name
-        hypernetwork.save(filename)
-    except Exception:
-        hypernetwork.sd_checkpoint = old_sd_checkpoint
-        hypernetwork.sd_checkpoint_name = old_sd_checkpoint_name
-        hypernetwork.name = old_hypernetwork_name
-        raise