From 9dcbd25e253dbf5e1bc336b73aaa3e37abc6a1f3 Mon Sep 17 00:00:00 2001 From: Haoming Date: Mon, 28 Jul 2025 11:30:08 +0800 Subject: [PATCH] yeet --- modules/extras.py | 31 +- modules/hypernetworks/hypernetwork.py | 348 +--- modules/import_hook.py | 16 - modules/initialize.py | 2 +- modules/mac_specific.py | 98 -- modules/models/diffusion/ddpm_edit.py | 1460 ----------------- modules/npu_specific.py | 31 - modules/sd_disable_initialization.py | 232 --- modules/sd_hijack.py | 230 --- modules/sd_hijack_checkpoint.py | 46 - modules/sd_hijack_clip.py | 384 ----- modules/sd_hijack_clip_old.py | 82 - modules/sd_hijack_ip2p.py | 10 - modules/sd_hijack_open_clip.py | 71 - modules/sd_hijack_optimizations.py | 677 -------- modules/sd_hijack_unet.py | 154 -- modules/sd_hijack_xlmr.py | 32 - modules/sd_models.py | 11 +- modules/sd_models_config.py | 137 -- modules/sd_models_types.py | 39 - modules/sd_models_xl.py | 115 -- modules/sd_samplers_compvis.py | 0 modules/shared.py | 2 +- .../textual_inversion/textual_inversion.py | 2 +- modules/xlmr.py | 140 -- modules/xlmr_m18.py | 166 -- modules/xpu_specific.py | 138 -- 27 files changed, 6 insertions(+), 4648 deletions(-) delete mode 100644 modules/import_hook.py delete mode 100644 modules/mac_specific.py delete mode 100644 modules/models/diffusion/ddpm_edit.py delete mode 100644 modules/npu_specific.py delete mode 100644 modules/sd_disable_initialization.py delete mode 100644 modules/sd_hijack_checkpoint.py delete mode 100644 modules/sd_hijack_clip.py delete mode 100644 modules/sd_hijack_clip_old.py delete mode 100644 modules/sd_hijack_ip2p.py delete mode 100644 modules/sd_hijack_open_clip.py delete mode 100644 modules/sd_hijack_optimizations.py delete mode 100644 modules/sd_hijack_unet.py delete mode 100644 modules/sd_hijack_xlmr.py delete mode 100644 modules/sd_models_config.py delete mode 100644 modules/sd_models_types.py delete mode 100644 modules/sd_models_xl.py delete mode 100644 modules/sd_samplers_compvis.py delete mode 100644 modules/xlmr.py delete mode 100644 modules/xlmr_m18.py delete mode 100644 modules/xpu_specific.py diff --git a/modules/extras.py b/modules/extras.py index 306c5609..b13dab43 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -7,7 +7,7 @@ import json import torch import tqdm -from modules import shared, images, sd_models, sd_vae, sd_models_config, errors +from modules import shared, images, sd_models, sd_vae, errors from modules.ui_common import plaintext_to_html import gradio as gr import safetensors.torch @@ -36,32 +36,6 @@ def run_pnginfo(image): return '', geninfo, info -def create_config(ckpt_result, config_source, a, b, c): - def config(x): - res = sd_models_config.find_checkpoint_config_near_filename(x) if x else None - return res if res != shared.sd_default_config else None - - if config_source == 0: - cfg = config(a) or config(b) or config(c) - elif config_source == 1: - cfg = config(b) - elif config_source == 2: - cfg = config(c) - else: - cfg = None - - if cfg is None: - return - - filename, _ = os.path.splitext(ckpt_result) - checkpoint_filename = filename + ".yaml" - - print("Copying config:") - print(" from:", cfg) - print(" to:", checkpoint_filename) - shutil.copyfile(cfg, checkpoint_filename) - - checkpoint_dict_skip_on_merge = ["cond_stage_model.transformer.text_model.embeddings.position_ids"] @@ -321,9 +295,6 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ if created_model: created_model.calculate_shorthash() - # TODO inside create_config() sd_models_config.find_checkpoint_config_near_filename() is called which has been commented out - #create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info) - print(f"Checkpoint saved to {output_modelname}.") shared.state.textinfo = "Checkpoint saved" shared.state.end() diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py index 5a13d14c..5674d906 100644 --- a/modules/hypernetworks/hypernetwork.py +++ b/modules/hypernetworks/hypernetwork.py @@ -9,7 +9,7 @@ import torch import tqdm from einops import rearrange, repeat from backend.nn.unet import default -from modules import devices, sd_models, shared, sd_samplers, hashes, sd_hijack_checkpoint, errors +from modules import devices, sd_models, shared, sd_samplers, hashes, errors from modules.textual_inversion import textual_inversion from torch import einsum from torch.nn.init import normal_, xavier_normal_, xavier_uniform_, kaiming_normal_, kaiming_uniform_, zeros_ @@ -433,349 +433,3 @@ def statistics(data): std = stdev(recent_data) recent_information = f"recent 32 loss:{mean(recent_data):.3f}" + u"\u00B1" + f"({std / (len(recent_data) ** 0.5):.3f})" return total_information, recent_information - -# -# 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.cmd_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() -# -# -# def train_hypernetwork(id_task, hypernetwork_name: str, learn_rate: float, batch_size: int, gradient_step: int, data_root: str, log_directory: str, training_width: int, training_height: int, varsize: bool, steps: int, clip_grad_mode: str, clip_grad_value: float, shuffle_tags: bool, tag_drop_out: bool, latent_sampling_method: str, use_weight: bool, create_image_every: int, save_hypernetwork_every: int, template_filename: str, preview_from_txt2img: bool, preview_prompt: str, preview_negative_prompt: str, preview_steps: int, preview_sampler_name: str, preview_cfg_scale: float, preview_seed: int, preview_width: int, preview_height: int): -# from modules import images, processing -# -# 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, log_directory, 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-hypernetwork" -# shared.state.textinfo = "Initializing hypernetwork training..." -# shared.state.job_count = steps -# -# hypernetwork_name = hypernetwork_name.rsplit('(', 1)[0] -# filename = os.path.join(shared.cmd_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', ]} -# ) -# saving_settings.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 = "" -# last_saved_image = "" -# forced_filename = "" -# -# pbar = tqdm.tqdm(total=steps - initial_step) -# try: -# sd_hijack_checkpoint.add() -# -# for _ 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 -# if torch.cuda.is_available(): -# 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 = sd_samplers.samplers_map[preview_sampler_name.lower()] -# 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 -# -# with closing(p): -# 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) -# if torch.cuda.is_available(): -# 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""" -#

-# Loss: {loss_step:.7f}
-# Step: {steps_done}
-# Last prompt: {html.escape(batch.cond_text[0])}
-# Last saved hypernetwork: {html.escape(last_saved_file)}
-# Last saved image: {html.escape(last_saved_image)}
-#

-# """ -# except Exception: -# errors.report("Exception in training hypernetwork", exc_info=True) -# finally: -# pbar.leave = False -# pbar.close() -# hypernetwork.eval() -# sd_hijack_checkpoint.remove() -# -# -# -# filename = os.path.join(shared.cmd_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: -# hypernetwork.sd_checkpoint = old_sd_checkpoint -# hypernetwork.sd_checkpoint_name = old_sd_checkpoint_name -# hypernetwork.name = old_hypernetwork_name -# raise diff --git a/modules/import_hook.py b/modules/import_hook.py deleted file mode 100644 index f0197ffa..00000000 --- a/modules/import_hook.py +++ /dev/null @@ -1,16 +0,0 @@ -# import sys -# -# # this will break any attempt to import xformers which will prevent stability diffusion repo from trying to use it -# if "--xformers" not in "".join(sys.argv): -# sys.modules["xformers"] = None -# -# # Hack to fix a changed import in torchvision 0.17+, which otherwise breaks -# # basicsr; see https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13985 -# try: -# import torchvision.transforms.functional_tensor # noqa: F401 -# except ImportError: -# try: -# import torchvision.transforms.functional as functional -# sys.modules["torchvision.transforms.functional_tensor"] = functional -# except ImportError: -# pass # shrug... diff --git a/modules/initialize.py b/modules/initialize.py index c9a3d9f0..536b4959 100644 --- a/modules/initialize.py +++ b/modules/initialize.py @@ -22,7 +22,7 @@ def imports(): import gradio # noqa: F401 startup_timer.record("import gradio") - from modules import paths, timer, import_hook, errors # noqa: F401 + from modules import paths, timer, errors # noqa: F401 startup_timer.record("setup paths") from modules import shared_init diff --git a/modules/mac_specific.py b/modules/mac_specific.py deleted file mode 100644 index a42741f6..00000000 --- a/modules/mac_specific.py +++ /dev/null @@ -1,98 +0,0 @@ -# import logging -# -# import torch -# from torch import Tensor -# import platform -# from modules.sd_hijack_utils import CondFunc -# from packaging import version -# from modules import shared -# -# log = logging.getLogger(__name__) -# -# -# # before torch version 1.13, has_mps is only available in nightly pytorch and macOS 12.3+, -# # use check `getattr` and try it for compatibility. -# # in torch version 1.13, backends.mps.is_available() and backends.mps.is_built() are introduced in to check mps availability, -# # since torch 2.0.1+ nightly build, getattr(torch, 'has_mps', False) was deprecated, see https://github.com/pytorch/pytorch/pull/103279 -# def check_for_mps() -> bool: -# if version.parse(torch.__version__) <= version.parse("2.0.1"): -# if not getattr(torch, 'has_mps', False): -# return False -# try: -# torch.zeros(1).to(torch.device("mps")) -# return True -# except Exception: -# return False -# else: -# return torch.backends.mps.is_available() and torch.backends.mps.is_built() -# -# -# has_mps = check_for_mps() -# -# -# def torch_mps_gc() -> None: -# try: -# if shared.state.current_latent is not None: -# log.debug("`current_latent` is set, skipping MPS garbage collection") -# return -# from torch.mps import empty_cache -# empty_cache() -# except Exception: -# log.warning("MPS garbage collection failed", exc_info=True) -# -# -# # MPS workaround for https://github.com/pytorch/pytorch/issues/89784 -# def cumsum_fix(input, cumsum_func, *args, **kwargs): -# if input.device.type == 'mps': -# output_dtype = kwargs.get('dtype', input.dtype) -# if output_dtype == torch.int64: -# return cumsum_func(input.cpu(), *args, **kwargs).to(input.device) -# elif output_dtype == torch.bool or cumsum_needs_int_fix and (output_dtype == torch.int8 or output_dtype == torch.int16): -# return cumsum_func(input.to(torch.int32), *args, **kwargs).to(torch.int64) -# return cumsum_func(input, *args, **kwargs) -# -# -# # MPS workaround for https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14046 -# def interpolate_with_fp32_fallback(orig_func, *args, **kwargs) -> Tensor: -# try: -# return orig_func(*args, **kwargs) -# except RuntimeError as e: -# if "not implemented for" in str(e) and "Half" in str(e): -# input_tensor = args[0] -# return orig_func(input_tensor.to(torch.float32), *args[1:], **kwargs).to(input_tensor.dtype) -# else: -# print(f"An unexpected RuntimeError occurred: {str(e)}") -# -# if has_mps: -# if platform.mac_ver()[0].startswith("13.2."): -# # MPS workaround for https://github.com/pytorch/pytorch/issues/95188, thanks to danieldk (https://github.com/explosion/curated-transformers/pull/124) -# CondFunc('torch.nn.functional.linear', lambda _, input, weight, bias: (torch.matmul(input, weight.t()) + bias) if bias is not None else torch.matmul(input, weight.t()), lambda _, input, weight, bias: input.numel() > 10485760) -# -# if version.parse(torch.__version__) < version.parse("1.13"): -# # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working -# -# # MPS workaround for https://github.com/pytorch/pytorch/issues/79383 -# CondFunc('torch.Tensor.to', lambda orig_func, self, *args, **kwargs: orig_func(self.contiguous(), *args, **kwargs), -# lambda _, self, *args, **kwargs: self.device.type != 'mps' and (args and isinstance(args[0], torch.device) and args[0].type == 'mps' or isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps')) -# # MPS workaround for https://github.com/pytorch/pytorch/issues/80800 -# CondFunc('torch.nn.functional.layer_norm', lambda orig_func, *args, **kwargs: orig_func(*([args[0].contiguous()] + list(args[1:])), **kwargs), -# lambda _, *args, **kwargs: args and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps') -# # MPS workaround for https://github.com/pytorch/pytorch/issues/90532 -# CondFunc('torch.Tensor.numpy', lambda orig_func, self, *args, **kwargs: orig_func(self.detach(), *args, **kwargs), lambda _, self, *args, **kwargs: self.requires_grad) -# elif version.parse(torch.__version__) > version.parse("1.13.1"): -# cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0)) -# cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs) -# CondFunc('torch.cumsum', cumsum_fix_func, None) -# CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None) -# CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None) -# -# # MPS workaround for https://github.com/pytorch/pytorch/issues/96113 -# CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps') -# -# # MPS workaround for https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14046 -# CondFunc('torch.nn.functional.interpolate', interpolate_with_fp32_fallback, None) -# -# # MPS workaround for https://github.com/pytorch/pytorch/issues/92311 -# if platform.processor() == 'i386': -# for funcName in ['torch.argmax', 'torch.Tensor.argmax']: -# CondFunc(funcName, lambda _, input, *args, **kwargs: torch.max(input.float() if input.dtype == torch.int64 else input, *args, **kwargs)[1], lambda _, input, *args, **kwargs: input.device.type == 'mps') diff --git a/modules/models/diffusion/ddpm_edit.py b/modules/models/diffusion/ddpm_edit.py deleted file mode 100644 index aa894eac..00000000 --- a/modules/models/diffusion/ddpm_edit.py +++ /dev/null @@ -1,1460 +0,0 @@ -# """ -# wild mixture of -# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py -# https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py -# https://github.com/CompVis/taming-transformers -# -- merci -# """ -# -# # File modified by authors of InstructPix2Pix from original (https://github.com/CompVis/stable-diffusion). -# # See more details in LICENSE. -# -# import torch -# import torch.nn as nn -# import numpy as np -# import pytorch_lightning as pl -# from torch.optim.lr_scheduler import LambdaLR -# from einops import rearrange, repeat -# from contextlib import contextmanager -# from functools import partial -# from tqdm import tqdm -# from torchvision.utils import make_grid -# from pytorch_lightning.utilities.distributed import rank_zero_only -# -# from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config -# from ldm.modules.ema import LitEma -# from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution -# from ldm.models.autoencoder import IdentityFirstStage, AutoencoderKL -# from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like -# from ldm.models.diffusion.ddim import DDIMSampler -# -# try: -# from ldm.models.autoencoder import VQModelInterface -# except Exception: -# class VQModelInterface: -# pass -# -# __conditioning_keys__ = {'concat': 'c_concat', -# 'crossattn': 'c_crossattn', -# 'adm': 'y'} -# -# -# def disabled_train(self, mode=True): -# """Overwrite model.train with this function to make sure train/eval mode -# does not change anymore.""" -# return self -# -# -# def uniform_on_device(r1, r2, shape, device): -# return (r1 - r2) * torch.rand(*shape, device=device) + r2 -# -# -# class DDPM(pl.LightningModule): -# # classic DDPM with Gaussian diffusion, in image space -# def __init__(self, -# unet_config, -# timesteps=1000, -# beta_schedule="linear", -# loss_type="l2", -# ckpt_path=None, -# ignore_keys=None, -# load_only_unet=False, -# monitor="val/loss", -# use_ema=True, -# first_stage_key="image", -# image_size=256, -# channels=3, -# log_every_t=100, -# clip_denoised=True, -# linear_start=1e-4, -# linear_end=2e-2, -# cosine_s=8e-3, -# given_betas=None, -# original_elbo_weight=0., -# v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta -# l_simple_weight=1., -# conditioning_key=None, -# parameterization="eps", # all assuming fixed variance schedules -# scheduler_config=None, -# use_positional_encodings=False, -# learn_logvar=False, -# logvar_init=0., -# load_ema=True, -# ): -# super().__init__() -# assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"' -# self.parameterization = parameterization -# print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode") -# self.cond_stage_model = None -# self.clip_denoised = clip_denoised -# self.log_every_t = log_every_t -# self.first_stage_key = first_stage_key -# self.image_size = image_size # try conv? -# self.channels = channels -# self.use_positional_encodings = use_positional_encodings -# self.model = DiffusionWrapper(unet_config, conditioning_key) -# count_params(self.model, verbose=True) -# self.use_ema = use_ema -# -# self.use_scheduler = scheduler_config is not None -# if self.use_scheduler: -# self.scheduler_config = scheduler_config -# -# self.v_posterior = v_posterior -# self.original_elbo_weight = original_elbo_weight -# self.l_simple_weight = l_simple_weight -# -# if monitor is not None: -# self.monitor = monitor -# -# if self.use_ema and load_ema: -# self.model_ema = LitEma(self.model) -# print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") -# -# if ckpt_path is not None: -# self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or [], only_model=load_only_unet) -# -# # If initialing from EMA-only checkpoint, create EMA model after loading. -# if self.use_ema and not load_ema: -# self.model_ema = LitEma(self.model) -# print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") -# -# self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps, -# linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s) -# -# self.loss_type = loss_type -# -# self.learn_logvar = learn_logvar -# self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,)) -# if self.learn_logvar: -# self.logvar = nn.Parameter(self.logvar, requires_grad=True) -# -# -# def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000, -# linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): -# if exists(given_betas): -# betas = given_betas -# else: -# betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, -# cosine_s=cosine_s) -# alphas = 1. - betas -# alphas_cumprod = np.cumprod(alphas, axis=0) -# alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) -# -# timesteps, = betas.shape -# self.num_timesteps = int(timesteps) -# self.linear_start = linear_start -# self.linear_end = linear_end -# assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' -# -# to_torch = partial(torch.tensor, dtype=torch.float32) -# -# self.register_buffer('betas', to_torch(betas)) -# self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) -# self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) -# -# # calculations for diffusion q(x_t | x_{t-1}) and others -# self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) -# self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) -# self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) -# self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) -# self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) -# -# # calculations for posterior q(x_{t-1} | x_t, x_0) -# posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / ( -# 1. - alphas_cumprod) + self.v_posterior * betas -# # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) -# self.register_buffer('posterior_variance', to_torch(posterior_variance)) -# # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain -# self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) -# self.register_buffer('posterior_mean_coef1', to_torch( -# betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) -# self.register_buffer('posterior_mean_coef2', to_torch( -# (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) -# -# if self.parameterization == "eps": -# lvlb_weights = self.betas ** 2 / ( -# 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) -# elif self.parameterization == "x0": -# lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) -# else: -# raise NotImplementedError("mu not supported") -# # TODO how to choose this term -# lvlb_weights[0] = lvlb_weights[1] -# self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) -# assert not torch.isnan(self.lvlb_weights).all() -# -# @contextmanager -# def ema_scope(self, context=None): -# if self.use_ema: -# self.model_ema.store(self.model.parameters()) -# self.model_ema.copy_to(self.model) -# if context is not None: -# print(f"{context}: Switched to EMA weights") -# try: -# yield None -# finally: -# if self.use_ema: -# self.model_ema.restore(self.model.parameters()) -# if context is not None: -# print(f"{context}: Restored training weights") -# -# def init_from_ckpt(self, path, ignore_keys=None, only_model=False): -# ignore_keys = ignore_keys or [] -# -# sd = torch.load(path, map_location="cpu") -# if "state_dict" in list(sd.keys()): -# sd = sd["state_dict"] -# keys = list(sd.keys()) -# -# # Our model adds additional channels to the first layer to condition on an input image. -# # For the first layer, copy existing channel weights and initialize new channel weights to zero. -# input_keys = [ -# "model.diffusion_model.input_blocks.0.0.weight", -# "model_ema.diffusion_modelinput_blocks00weight", -# ] -# -# self_sd = self.state_dict() -# for input_key in input_keys: -# if input_key not in sd or input_key not in self_sd: -# continue -# -# input_weight = self_sd[input_key] -# -# if input_weight.size() != sd[input_key].size(): -# print(f"Manual init: {input_key}") -# input_weight.zero_() -# input_weight[:, :4, :, :].copy_(sd[input_key]) -# ignore_keys.append(input_key) -# -# for k in keys: -# for ik in ignore_keys: -# if k.startswith(ik): -# print(f"Deleting key {k} from state_dict.") -# del sd[k] -# missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( -# sd, strict=False) -# print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") -# if missing: -# print(f"Missing Keys: {missing}") -# if unexpected: -# print(f"Unexpected Keys: {unexpected}") -# -# def q_mean_variance(self, x_start, t): -# """ -# Get the distribution q(x_t | x_0). -# :param x_start: the [N x C x ...] tensor of noiseless inputs. -# :param t: the number of diffusion steps (minus 1). Here, 0 means one step. -# :return: A tuple (mean, variance, log_variance), all of x_start's shape. -# """ -# mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start) -# variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) -# log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) -# return mean, variance, log_variance -# -# def predict_start_from_noise(self, x_t, t, noise): -# return ( -# extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - -# extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise -# ) -# -# def q_posterior(self, x_start, x_t, t): -# posterior_mean = ( -# extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + -# extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t -# ) -# posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape) -# posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape) -# return posterior_mean, posterior_variance, posterior_log_variance_clipped -# -# def p_mean_variance(self, x, t, clip_denoised: bool): -# model_out = self.model(x, t) -# if self.parameterization == "eps": -# x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) -# elif self.parameterization == "x0": -# x_recon = model_out -# if clip_denoised: -# x_recon.clamp_(-1., 1.) -# -# model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) -# return model_mean, posterior_variance, posterior_log_variance -# -# @torch.no_grad() -# def p_sample(self, x, t, clip_denoised=True, repeat_noise=False): -# b, *_, device = *x.shape, x.device -# model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised) -# noise = noise_like(x.shape, device, repeat_noise) -# # no noise when t == 0 -# nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) -# return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise -# -# @torch.no_grad() -# def p_sample_loop(self, shape, return_intermediates=False): -# device = self.betas.device -# b = shape[0] -# img = torch.randn(shape, device=device) -# intermediates = [img] -# for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps): -# img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long), -# clip_denoised=self.clip_denoised) -# if i % self.log_every_t == 0 or i == self.num_timesteps - 1: -# intermediates.append(img) -# if return_intermediates: -# return img, intermediates -# return img -# -# @torch.no_grad() -# def sample(self, batch_size=16, return_intermediates=False): -# image_size = self.image_size -# channels = self.channels -# return self.p_sample_loop((batch_size, channels, image_size, image_size), -# return_intermediates=return_intermediates) -# -# def q_sample(self, x_start, t, noise=None): -# noise = default(noise, lambda: torch.randn_like(x_start)) -# return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + -# extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) -# -# def get_loss(self, pred, target, mean=True): -# if self.loss_type == 'l1': -# loss = (target - pred).abs() -# if mean: -# loss = loss.mean() -# elif self.loss_type == 'l2': -# if mean: -# loss = torch.nn.functional.mse_loss(target, pred) -# else: -# loss = torch.nn.functional.mse_loss(target, pred, reduction='none') -# else: -# raise NotImplementedError("unknown loss type '{loss_type}'") -# -# return loss -# -# def p_losses(self, x_start, t, noise=None): -# noise = default(noise, lambda: torch.randn_like(x_start)) -# x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) -# model_out = self.model(x_noisy, t) -# -# loss_dict = {} -# if self.parameterization == "eps": -# target = noise -# elif self.parameterization == "x0": -# target = x_start -# else: -# raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported") -# -# loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3]) -# -# log_prefix = 'train' if self.training else 'val' -# -# loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()}) -# loss_simple = loss.mean() * self.l_simple_weight -# -# loss_vlb = (self.lvlb_weights[t] * loss).mean() -# loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb}) -# -# loss = loss_simple + self.original_elbo_weight * loss_vlb -# -# loss_dict.update({f'{log_prefix}/loss': loss}) -# -# return loss, loss_dict -# -# def forward(self, x, *args, **kwargs): -# # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size -# # assert h == img_size and w == img_size, f'height and width of image must be {img_size}' -# t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() -# return self.p_losses(x, t, *args, **kwargs) -# -# def get_input(self, batch, k): -# return batch[k] -# -# def shared_step(self, batch): -# x = self.get_input(batch, self.first_stage_key) -# loss, loss_dict = self(x) -# return loss, loss_dict -# -# def training_step(self, batch, batch_idx): -# loss, loss_dict = self.shared_step(batch) -# -# self.log_dict(loss_dict, prog_bar=True, -# logger=True, on_step=True, on_epoch=True) -# -# self.log("global_step", self.global_step, -# prog_bar=True, logger=True, on_step=True, on_epoch=False) -# -# if self.use_scheduler: -# lr = self.optimizers().param_groups[0]['lr'] -# self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False) -# -# return loss -# -# @torch.no_grad() -# def validation_step(self, batch, batch_idx): -# _, loss_dict_no_ema = self.shared_step(batch) -# with self.ema_scope(): -# _, loss_dict_ema = self.shared_step(batch) -# loss_dict_ema = {f"{key}_ema": loss_dict_ema[key] for key in loss_dict_ema} -# self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) -# self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) -# -# def on_train_batch_end(self, *args, **kwargs): -# if self.use_ema: -# self.model_ema(self.model) -# -# def _get_rows_from_list(self, samples): -# n_imgs_per_row = len(samples) -# denoise_grid = rearrange(samples, 'n b c h w -> b n c h w') -# denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') -# denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) -# return denoise_grid -# -# @torch.no_grad() -# def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs): -# log = {} -# x = self.get_input(batch, self.first_stage_key) -# N = min(x.shape[0], N) -# n_row = min(x.shape[0], n_row) -# x = x.to(self.device)[:N] -# log["inputs"] = x -# -# # get diffusion row -# diffusion_row = [] -# x_start = x[:n_row] -# -# for t in range(self.num_timesteps): -# if t % self.log_every_t == 0 or t == self.num_timesteps - 1: -# t = repeat(torch.tensor([t]), '1 -> b', b=n_row) -# t = t.to(self.device).long() -# noise = torch.randn_like(x_start) -# x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) -# diffusion_row.append(x_noisy) -# -# log["diffusion_row"] = self._get_rows_from_list(diffusion_row) -# -# if sample: -# # get denoise row -# with self.ema_scope("Plotting"): -# samples, denoise_row = self.sample(batch_size=N, return_intermediates=True) -# -# log["samples"] = samples -# log["denoise_row"] = self._get_rows_from_list(denoise_row) -# -# if return_keys: -# if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: -# return log -# else: -# return {key: log[key] for key in return_keys} -# return log -# -# def configure_optimizers(self): -# lr = self.learning_rate -# params = list(self.model.parameters()) -# if self.learn_logvar: -# params = params + [self.logvar] -# opt = torch.optim.AdamW(params, lr=lr) -# return opt -# -# -# class LatentDiffusion(DDPM): -# """main class""" -# def __init__(self, -# first_stage_config, -# cond_stage_config, -# num_timesteps_cond=None, -# cond_stage_key="image", -# cond_stage_trainable=False, -# concat_mode=True, -# cond_stage_forward=None, -# conditioning_key=None, -# scale_factor=1.0, -# scale_by_std=False, -# load_ema=True, -# *args, **kwargs): -# self.num_timesteps_cond = default(num_timesteps_cond, 1) -# self.scale_by_std = scale_by_std -# assert self.num_timesteps_cond <= kwargs['timesteps'] -# # for backwards compatibility after implementation of DiffusionWrapper -# if conditioning_key is None: -# conditioning_key = 'concat' if concat_mode else 'crossattn' -# if cond_stage_config == '__is_unconditional__': -# conditioning_key = None -# ckpt_path = kwargs.pop("ckpt_path", None) -# ignore_keys = kwargs.pop("ignore_keys", []) -# super().__init__(*args, conditioning_key=conditioning_key, load_ema=load_ema, **kwargs) -# self.concat_mode = concat_mode -# self.cond_stage_trainable = cond_stage_trainable -# self.cond_stage_key = cond_stage_key -# try: -# self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 -# except Exception: -# self.num_downs = 0 -# if not scale_by_std: -# self.scale_factor = scale_factor -# else: -# self.register_buffer('scale_factor', torch.tensor(scale_factor)) -# self.instantiate_first_stage(first_stage_config) -# self.instantiate_cond_stage(cond_stage_config) -# self.cond_stage_forward = cond_stage_forward -# self.clip_denoised = False -# self.bbox_tokenizer = None -# -# self.restarted_from_ckpt = False -# if ckpt_path is not None: -# self.init_from_ckpt(ckpt_path, ignore_keys) -# self.restarted_from_ckpt = True -# -# if self.use_ema and not load_ema: -# self.model_ema = LitEma(self.model) -# print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") -# -# def make_cond_schedule(self, ): -# self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long) -# ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long() -# self.cond_ids[:self.num_timesteps_cond] = ids -# -# @rank_zero_only -# @torch.no_grad() -# def on_train_batch_start(self, batch, batch_idx, dataloader_idx): -# # only for very first batch -# if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt: -# assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously' -# # set rescale weight to 1./std of encodings -# print("### USING STD-RESCALING ###") -# x = super().get_input(batch, self.first_stage_key) -# x = x.to(self.device) -# encoder_posterior = self.encode_first_stage(x) -# z = self.get_first_stage_encoding(encoder_posterior).detach() -# del self.scale_factor -# self.register_buffer('scale_factor', 1. / z.flatten().std()) -# print(f"setting self.scale_factor to {self.scale_factor}") -# print("### USING STD-RESCALING ###") -# -# def register_schedule(self, -# given_betas=None, beta_schedule="linear", timesteps=1000, -# linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): -# super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s) -# -# self.shorten_cond_schedule = self.num_timesteps_cond > 1 -# if self.shorten_cond_schedule: -# self.make_cond_schedule() -# -# def instantiate_first_stage(self, config): -# model = instantiate_from_config(config) -# self.first_stage_model = model.eval() -# self.first_stage_model.train = disabled_train -# for param in self.first_stage_model.parameters(): -# param.requires_grad = False -# -# def instantiate_cond_stage(self, config): -# if not self.cond_stage_trainable: -# if config == "__is_first_stage__": -# print("Using first stage also as cond stage.") -# self.cond_stage_model = self.first_stage_model -# elif config == "__is_unconditional__": -# print(f"Training {self.__class__.__name__} as an unconditional model.") -# self.cond_stage_model = None -# # self.be_unconditional = True -# else: -# model = instantiate_from_config(config) -# self.cond_stage_model = model.eval() -# self.cond_stage_model.train = disabled_train -# for param in self.cond_stage_model.parameters(): -# param.requires_grad = False -# else: -# assert config != '__is_first_stage__' -# assert config != '__is_unconditional__' -# model = instantiate_from_config(config) -# self.cond_stage_model = model -# -# def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False): -# denoise_row = [] -# for zd in tqdm(samples, desc=desc): -# denoise_row.append(self.decode_first_stage(zd.to(self.device), -# force_not_quantize=force_no_decoder_quantization)) -# n_imgs_per_row = len(denoise_row) -# denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W -# denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w') -# denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') -# denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) -# return denoise_grid -# -# def get_first_stage_encoding(self, encoder_posterior): -# if isinstance(encoder_posterior, DiagonalGaussianDistribution): -# z = encoder_posterior.sample() -# elif isinstance(encoder_posterior, torch.Tensor): -# z = encoder_posterior -# else: -# raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented") -# return self.scale_factor * z -# -# def get_learned_conditioning(self, c): -# if self.cond_stage_forward is None: -# if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode): -# c = self.cond_stage_model.encode(c) -# if isinstance(c, DiagonalGaussianDistribution): -# c = c.mode() -# else: -# c = self.cond_stage_model(c) -# else: -# assert hasattr(self.cond_stage_model, self.cond_stage_forward) -# c = getattr(self.cond_stage_model, self.cond_stage_forward)(c) -# return c -# -# def meshgrid(self, h, w): -# y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1) -# x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1) -# -# arr = torch.cat([y, x], dim=-1) -# return arr -# -# def delta_border(self, h, w): -# """ -# :param h: height -# :param w: width -# :return: normalized distance to image border, -# wtith min distance = 0 at border and max dist = 0.5 at image center -# """ -# lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2) -# arr = self.meshgrid(h, w) / lower_right_corner -# dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0] -# dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0] -# edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0] -# return edge_dist -# -# def get_weighting(self, h, w, Ly, Lx, device): -# weighting = self.delta_border(h, w) -# weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"], -# self.split_input_params["clip_max_weight"], ) -# weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device) -# -# if self.split_input_params["tie_braker"]: -# L_weighting = self.delta_border(Ly, Lx) -# L_weighting = torch.clip(L_weighting, -# self.split_input_params["clip_min_tie_weight"], -# self.split_input_params["clip_max_tie_weight"]) -# -# L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device) -# weighting = weighting * L_weighting -# return weighting -# -# def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code -# """ -# :param x: img of size (bs, c, h, w) -# :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1]) -# """ -# bs, nc, h, w = x.shape -# -# # number of crops in image -# Ly = (h - kernel_size[0]) // stride[0] + 1 -# Lx = (w - kernel_size[1]) // stride[1] + 1 -# -# if uf == 1 and df == 1: -# fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) -# unfold = torch.nn.Unfold(**fold_params) -# -# fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params) -# -# weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype) -# normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap -# weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx)) -# -# elif uf > 1 and df == 1: -# fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) -# unfold = torch.nn.Unfold(**fold_params) -# -# fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf), -# dilation=1, padding=0, -# stride=(stride[0] * uf, stride[1] * uf)) -# fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2) -# -# weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype) -# normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap -# weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx)) -# -# elif df > 1 and uf == 1: -# fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) -# unfold = torch.nn.Unfold(**fold_params) -# -# fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df), -# dilation=1, padding=0, -# stride=(stride[0] // df, stride[1] // df)) -# fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2) -# -# weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype) -# normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap -# weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx)) -# -# else: -# raise NotImplementedError -# -# return fold, unfold, normalization, weighting -# -# @torch.no_grad() -# def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False, -# cond_key=None, return_original_cond=False, bs=None, uncond=0.05): -# x = super().get_input(batch, k) -# if bs is not None: -# x = x[:bs] -# x = x.to(self.device) -# encoder_posterior = self.encode_first_stage(x) -# z = self.get_first_stage_encoding(encoder_posterior).detach() -# cond_key = cond_key or self.cond_stage_key -# xc = super().get_input(batch, cond_key) -# if bs is not None: -# xc["c_crossattn"] = xc["c_crossattn"][:bs] -# xc["c_concat"] = xc["c_concat"][:bs] -# cond = {} -# -# # To support classifier-free guidance, randomly drop out only text conditioning 5%, only image conditioning 5%, and both 5%. -# random = torch.rand(x.size(0), device=x.device) -# prompt_mask = rearrange(random < 2 * uncond, "n -> n 1 1") -# input_mask = 1 - rearrange((random >= uncond).float() * (random < 3 * uncond).float(), "n -> n 1 1 1") -# -# null_prompt = self.get_learned_conditioning([""]) -# cond["c_crossattn"] = [torch.where(prompt_mask, null_prompt, self.get_learned_conditioning(xc["c_crossattn"]).detach())] -# cond["c_concat"] = [input_mask * self.encode_first_stage((xc["c_concat"].to(self.device))).mode().detach()] -# -# out = [z, cond] -# if return_first_stage_outputs: -# xrec = self.decode_first_stage(z) -# out.extend([x, xrec]) -# if return_original_cond: -# out.append(xc) -# return out -# -# @torch.no_grad() -# def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): -# if predict_cids: -# if z.dim() == 4: -# z = torch.argmax(z.exp(), dim=1).long() -# z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) -# z = rearrange(z, 'b h w c -> b c h w').contiguous() -# -# z = 1. / self.scale_factor * z -# -# if hasattr(self, "split_input_params"): -# if self.split_input_params["patch_distributed_vq"]: -# ks = self.split_input_params["ks"] # eg. (128, 128) -# stride = self.split_input_params["stride"] # eg. (64, 64) -# uf = self.split_input_params["vqf"] -# bs, nc, h, w = z.shape -# if ks[0] > h or ks[1] > w: -# ks = (min(ks[0], h), min(ks[1], w)) -# print("reducing Kernel") -# -# if stride[0] > h or stride[1] > w: -# stride = (min(stride[0], h), min(stride[1], w)) -# print("reducing stride") -# -# fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) -# -# z = unfold(z) # (bn, nc * prod(**ks), L) -# # 1. Reshape to img shape -# z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) -# -# # 2. apply model loop over last dim -# if isinstance(self.first_stage_model, VQModelInterface): -# output_list = [self.first_stage_model.decode(z[:, :, :, :, i], -# force_not_quantize=predict_cids or force_not_quantize) -# for i in range(z.shape[-1])] -# else: -# -# output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) -# for i in range(z.shape[-1])] -# -# o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) -# o = o * weighting -# # Reverse 1. reshape to img shape -# o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) -# # stitch crops together -# decoded = fold(o) -# decoded = decoded / normalization # norm is shape (1, 1, h, w) -# return decoded -# else: -# if isinstance(self.first_stage_model, VQModelInterface): -# return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) -# else: -# return self.first_stage_model.decode(z) -# -# else: -# if isinstance(self.first_stage_model, VQModelInterface): -# return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) -# else: -# return self.first_stage_model.decode(z) -# -# # same as above but without decorator -# def differentiable_decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): -# if predict_cids: -# if z.dim() == 4: -# z = torch.argmax(z.exp(), dim=1).long() -# z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) -# z = rearrange(z, 'b h w c -> b c h w').contiguous() -# -# z = 1. / self.scale_factor * z -# -# if hasattr(self, "split_input_params"): -# if self.split_input_params["patch_distributed_vq"]: -# ks = self.split_input_params["ks"] # eg. (128, 128) -# stride = self.split_input_params["stride"] # eg. (64, 64) -# uf = self.split_input_params["vqf"] -# bs, nc, h, w = z.shape -# if ks[0] > h or ks[1] > w: -# ks = (min(ks[0], h), min(ks[1], w)) -# print("reducing Kernel") -# -# if stride[0] > h or stride[1] > w: -# stride = (min(stride[0], h), min(stride[1], w)) -# print("reducing stride") -# -# fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) -# -# z = unfold(z) # (bn, nc * prod(**ks), L) -# # 1. Reshape to img shape -# z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) -# -# # 2. apply model loop over last dim -# if isinstance(self.first_stage_model, VQModelInterface): -# output_list = [self.first_stage_model.decode(z[:, :, :, :, i], -# force_not_quantize=predict_cids or force_not_quantize) -# for i in range(z.shape[-1])] -# else: -# -# output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) -# for i in range(z.shape[-1])] -# -# o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) -# o = o * weighting -# # Reverse 1. reshape to img shape -# o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) -# # stitch crops together -# decoded = fold(o) -# decoded = decoded / normalization # norm is shape (1, 1, h, w) -# return decoded -# else: -# if isinstance(self.first_stage_model, VQModelInterface): -# return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) -# else: -# return self.first_stage_model.decode(z) -# -# else: -# if isinstance(self.first_stage_model, VQModelInterface): -# return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) -# else: -# return self.first_stage_model.decode(z) -# -# @torch.no_grad() -# def encode_first_stage(self, x): -# if hasattr(self, "split_input_params"): -# if self.split_input_params["patch_distributed_vq"]: -# ks = self.split_input_params["ks"] # eg. (128, 128) -# stride = self.split_input_params["stride"] # eg. (64, 64) -# df = self.split_input_params["vqf"] -# self.split_input_params['original_image_size'] = x.shape[-2:] -# bs, nc, h, w = x.shape -# if ks[0] > h or ks[1] > w: -# ks = (min(ks[0], h), min(ks[1], w)) -# print("reducing Kernel") -# -# if stride[0] > h or stride[1] > w: -# stride = (min(stride[0], h), min(stride[1], w)) -# print("reducing stride") -# -# fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df) -# z = unfold(x) # (bn, nc * prod(**ks), L) -# # Reshape to img shape -# z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) -# -# output_list = [self.first_stage_model.encode(z[:, :, :, :, i]) -# for i in range(z.shape[-1])] -# -# o = torch.stack(output_list, axis=-1) -# o = o * weighting -# -# # Reverse reshape to img shape -# o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) -# # stitch crops together -# decoded = fold(o) -# decoded = decoded / normalization -# return decoded -# -# else: -# return self.first_stage_model.encode(x) -# else: -# return self.first_stage_model.encode(x) -# -# def shared_step(self, batch, **kwargs): -# x, c = self.get_input(batch, self.first_stage_key) -# loss = self(x, c) -# return loss -# -# def forward(self, x, c, *args, **kwargs): -# t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() -# if self.model.conditioning_key is not None: -# assert c is not None -# if self.cond_stage_trainable: -# c = self.get_learned_conditioning(c) -# if self.shorten_cond_schedule: # TODO: drop this option -# tc = self.cond_ids[t].to(self.device) -# c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float())) -# return self.p_losses(x, c, t, *args, **kwargs) -# -# def apply_model(self, x_noisy, t, cond, return_ids=False): -# -# if isinstance(cond, dict): -# # hybrid case, cond is expected to be a dict -# pass -# else: -# if not isinstance(cond, list): -# cond = [cond] -# key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn' -# cond = {key: cond} -# -# if hasattr(self, "split_input_params"): -# assert len(cond) == 1 # todo can only deal with one conditioning atm -# assert not return_ids -# ks = self.split_input_params["ks"] # eg. (128, 128) -# stride = self.split_input_params["stride"] # eg. (64, 64) -# -# h, w = x_noisy.shape[-2:] -# -# fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride) -# -# z = unfold(x_noisy) # (bn, nc * prod(**ks), L) -# # Reshape to img shape -# z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) -# z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])] -# -# if self.cond_stage_key in ["image", "LR_image", "segmentation", -# 'bbox_img'] and self.model.conditioning_key: # todo check for completeness -# c_key = next(iter(cond.keys())) # get key -# c = next(iter(cond.values())) # get value -# assert (len(c) == 1) # todo extend to list with more than one elem -# c = c[0] # get element -# -# c = unfold(c) -# c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1])) # (bn, nc, ks[0], ks[1], L ) -# -# cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])] -# -# elif self.cond_stage_key == 'coordinates_bbox': -# assert 'original_image_size' in self.split_input_params, 'BoundingBoxRescaling is missing original_image_size' -# -# # assuming padding of unfold is always 0 and its dilation is always 1 -# n_patches_per_row = int((w - ks[0]) / stride[0] + 1) -# full_img_h, full_img_w = self.split_input_params['original_image_size'] -# # as we are operating on latents, we need the factor from the original image size to the -# # spatial latent size to properly rescale the crops for regenerating the bbox annotations -# num_downs = self.first_stage_model.encoder.num_resolutions - 1 -# rescale_latent = 2 ** (num_downs) -# -# # get top left positions of patches as conforming for the bbbox tokenizer, therefore we -# # need to rescale the tl patch coordinates to be in between (0,1) -# tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w, -# rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h) -# for patch_nr in range(z.shape[-1])] -# -# # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w) -# patch_limits = [(x_tl, y_tl, -# rescale_latent * ks[0] / full_img_w, -# rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates] -# # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates] -# -# # tokenize crop coordinates for the bounding boxes of the respective patches -# patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer._crop_encoder(bbox))[None].to(self.device) -# for bbox in patch_limits] # list of length l with tensors of shape (1, 2) -# print(patch_limits_tknzd[0].shape) -# # cut tknzd crop position from conditioning -# assert isinstance(cond, dict), 'cond must be dict to be fed into model' -# cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device) -# print(cut_cond.shape) -# -# adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd]) -# adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n') -# print(adapted_cond.shape) -# adapted_cond = self.get_learned_conditioning(adapted_cond) -# print(adapted_cond.shape) -# adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1]) -# print(adapted_cond.shape) -# -# cond_list = [{'c_crossattn': [e]} for e in adapted_cond] -# -# else: -# cond_list = [cond for i in range(z.shape[-1])] # Todo make this more efficient -# -# # apply model by loop over crops -# output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])] -# assert not isinstance(output_list[0], -# tuple) # todo cant deal with multiple model outputs check this never happens -# -# o = torch.stack(output_list, axis=-1) -# o = o * weighting -# # Reverse reshape to img shape -# o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) -# # stitch crops together -# x_recon = fold(o) / normalization -# -# else: -# x_recon = self.model(x_noisy, t, **cond) -# -# if isinstance(x_recon, tuple) and not return_ids: -# return x_recon[0] -# else: -# return x_recon -# -# def _predict_eps_from_xstart(self, x_t, t, pred_xstart): -# return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \ -# extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) -# -# def _prior_bpd(self, x_start): -# """ -# Get the prior KL term for the variational lower-bound, measured in -# bits-per-dim. -# This term can't be optimized, as it only depends on the encoder. -# :param x_start: the [N x C x ...] tensor of inputs. -# :return: a batch of [N] KL values (in bits), one per batch element. -# """ -# batch_size = x_start.shape[0] -# t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) -# qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) -# kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0) -# return mean_flat(kl_prior) / np.log(2.0) -# -# def p_losses(self, x_start, cond, t, noise=None): -# noise = default(noise, lambda: torch.randn_like(x_start)) -# x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) -# model_output = self.apply_model(x_noisy, t, cond) -# -# loss_dict = {} -# prefix = 'train' if self.training else 'val' -# -# if self.parameterization == "x0": -# target = x_start -# elif self.parameterization == "eps": -# target = noise -# else: -# raise NotImplementedError() -# -# loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) -# loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()}) -# -# logvar_t = self.logvar[t].to(self.device) -# loss = loss_simple / torch.exp(logvar_t) + logvar_t -# # loss = loss_simple / torch.exp(self.logvar) + self.logvar -# if self.learn_logvar: -# loss_dict.update({f'{prefix}/loss_gamma': loss.mean()}) -# loss_dict.update({'logvar': self.logvar.data.mean()}) -# -# loss = self.l_simple_weight * loss.mean() -# -# loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) -# loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() -# loss_dict.update({f'{prefix}/loss_vlb': loss_vlb}) -# loss += (self.original_elbo_weight * loss_vlb) -# loss_dict.update({f'{prefix}/loss': loss}) -# -# return loss, loss_dict -# -# def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, -# return_x0=False, score_corrector=None, corrector_kwargs=None): -# t_in = t -# model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids) -# -# if score_corrector is not None: -# assert self.parameterization == "eps" -# model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs) -# -# if return_codebook_ids: -# model_out, logits = model_out -# -# if self.parameterization == "eps": -# x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) -# elif self.parameterization == "x0": -# x_recon = model_out -# else: -# raise NotImplementedError() -# -# if clip_denoised: -# x_recon.clamp_(-1., 1.) -# if quantize_denoised: -# x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon) -# model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) -# if return_codebook_ids: -# return model_mean, posterior_variance, posterior_log_variance, logits -# elif return_x0: -# return model_mean, posterior_variance, posterior_log_variance, x_recon -# else: -# return model_mean, posterior_variance, posterior_log_variance -# -# @torch.no_grad() -# def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False, -# return_codebook_ids=False, quantize_denoised=False, return_x0=False, -# temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None): -# b, *_, device = *x.shape, x.device -# outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised, -# return_codebook_ids=return_codebook_ids, -# quantize_denoised=quantize_denoised, -# return_x0=return_x0, -# score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) -# if return_codebook_ids: -# raise DeprecationWarning("Support dropped.") -# model_mean, _, model_log_variance, logits = outputs -# elif return_x0: -# model_mean, _, model_log_variance, x0 = outputs -# else: -# model_mean, _, model_log_variance = outputs -# -# noise = noise_like(x.shape, device, repeat_noise) * temperature -# if noise_dropout > 0.: -# noise = torch.nn.functional.dropout(noise, p=noise_dropout) -# # no noise when t == 0 -# nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) -# -# if return_codebook_ids: -# return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1) -# if return_x0: -# return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0 -# else: -# return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise -# -# @torch.no_grad() -# def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False, -# img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0., -# score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, -# log_every_t=None): -# if not log_every_t: -# log_every_t = self.log_every_t -# timesteps = self.num_timesteps -# if batch_size is not None: -# b = batch_size if batch_size is not None else shape[0] -# shape = [batch_size] + list(shape) -# else: -# b = batch_size = shape[0] -# if x_T is None: -# img = torch.randn(shape, device=self.device) -# else: -# img = x_T -# intermediates = [] -# if cond is not None: -# if isinstance(cond, dict): -# cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else -# [x[:batch_size] for x in cond[key]] for key in cond} -# else: -# cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] -# -# if start_T is not None: -# timesteps = min(timesteps, start_T) -# iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation', -# total=timesteps) if verbose else reversed( -# range(0, timesteps)) -# if type(temperature) == float: -# temperature = [temperature] * timesteps -# -# for i in iterator: -# ts = torch.full((b,), i, device=self.device, dtype=torch.long) -# if self.shorten_cond_schedule: -# assert self.model.conditioning_key != 'hybrid' -# tc = self.cond_ids[ts].to(cond.device) -# cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) -# -# img, x0_partial = self.p_sample(img, cond, ts, -# clip_denoised=self.clip_denoised, -# quantize_denoised=quantize_denoised, return_x0=True, -# temperature=temperature[i], noise_dropout=noise_dropout, -# score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) -# if mask is not None: -# assert x0 is not None -# img_orig = self.q_sample(x0, ts) -# img = img_orig * mask + (1. - mask) * img -# -# if i % log_every_t == 0 or i == timesteps - 1: -# intermediates.append(x0_partial) -# if callback: -# callback(i) -# if img_callback: -# img_callback(img, i) -# return img, intermediates -# -# @torch.no_grad() -# def p_sample_loop(self, cond, shape, return_intermediates=False, -# x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, -# mask=None, x0=None, img_callback=None, start_T=None, -# log_every_t=None): -# -# if not log_every_t: -# log_every_t = self.log_every_t -# device = self.betas.device -# b = shape[0] -# if x_T is None: -# img = torch.randn(shape, device=device) -# else: -# img = x_T -# -# intermediates = [img] -# if timesteps is None: -# timesteps = self.num_timesteps -# -# if start_T is not None: -# timesteps = min(timesteps, start_T) -# iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( -# range(0, timesteps)) -# -# if mask is not None: -# assert x0 is not None -# assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match -# -# for i in iterator: -# ts = torch.full((b,), i, device=device, dtype=torch.long) -# if self.shorten_cond_schedule: -# assert self.model.conditioning_key != 'hybrid' -# tc = self.cond_ids[ts].to(cond.device) -# cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) -# -# img = self.p_sample(img, cond, ts, -# clip_denoised=self.clip_denoised, -# quantize_denoised=quantize_denoised) -# if mask is not None: -# img_orig = self.q_sample(x0, ts) -# img = img_orig * mask + (1. - mask) * img -# -# if i % log_every_t == 0 or i == timesteps - 1: -# intermediates.append(img) -# if callback: -# callback(i) -# if img_callback: -# img_callback(img, i) -# -# if return_intermediates: -# return img, intermediates -# return img -# -# @torch.no_grad() -# def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, -# verbose=True, timesteps=None, quantize_denoised=False, -# mask=None, x0=None, shape=None,**kwargs): -# if shape is None: -# shape = (batch_size, self.channels, self.image_size, self.image_size) -# if cond is not None: -# if isinstance(cond, dict): -# cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else -# [x[:batch_size] for x in cond[key]] for key in cond} -# else: -# cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] -# return self.p_sample_loop(cond, -# shape, -# return_intermediates=return_intermediates, x_T=x_T, -# verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised, -# mask=mask, x0=x0) -# -# @torch.no_grad() -# def sample_log(self,cond,batch_size,ddim, ddim_steps,**kwargs): -# -# if ddim: -# ddim_sampler = DDIMSampler(self) -# shape = (self.channels, self.image_size, self.image_size) -# samples, intermediates =ddim_sampler.sample(ddim_steps,batch_size, -# shape,cond,verbose=False,**kwargs) -# -# else: -# samples, intermediates = self.sample(cond=cond, batch_size=batch_size, -# return_intermediates=True,**kwargs) -# -# return samples, intermediates -# -# -# @torch.no_grad() -# def log_images(self, batch, N=4, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, -# quantize_denoised=True, inpaint=False, plot_denoise_rows=False, plot_progressive_rows=False, -# plot_diffusion_rows=False, **kwargs): -# -# use_ddim = False -# -# log = {} -# z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, -# return_first_stage_outputs=True, -# force_c_encode=True, -# return_original_cond=True, -# bs=N, uncond=0) -# N = min(x.shape[0], N) -# n_row = min(x.shape[0], n_row) -# log["inputs"] = x -# log["reals"] = xc["c_concat"] -# log["reconstruction"] = xrec -# if self.model.conditioning_key is not None: -# if hasattr(self.cond_stage_model, "decode"): -# xc = self.cond_stage_model.decode(c) -# log["conditioning"] = xc -# elif self.cond_stage_key in ["caption"]: -# xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"]) -# log["conditioning"] = xc -# elif self.cond_stage_key == 'class_label': -# xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"]) -# log['conditioning'] = xc -# elif isimage(xc): -# log["conditioning"] = xc -# if ismap(xc): -# log["original_conditioning"] = self.to_rgb(xc) -# -# if plot_diffusion_rows: -# # get diffusion row -# diffusion_row = [] -# z_start = z[:n_row] -# for t in range(self.num_timesteps): -# if t % self.log_every_t == 0 or t == self.num_timesteps - 1: -# t = repeat(torch.tensor([t]), '1 -> b', b=n_row) -# t = t.to(self.device).long() -# noise = torch.randn_like(z_start) -# z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) -# diffusion_row.append(self.decode_first_stage(z_noisy)) -# -# diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W -# diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') -# diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') -# diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) -# log["diffusion_row"] = diffusion_grid -# -# if sample: -# # get denoise row -# with self.ema_scope("Plotting"): -# samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, -# ddim_steps=ddim_steps,eta=ddim_eta) -# # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True) -# x_samples = self.decode_first_stage(samples) -# log["samples"] = x_samples -# if plot_denoise_rows: -# denoise_grid = self._get_denoise_row_from_list(z_denoise_row) -# log["denoise_row"] = denoise_grid -# -# if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance( -# self.first_stage_model, IdentityFirstStage): -# # also display when quantizing x0 while sampling -# with self.ema_scope("Plotting Quantized Denoised"): -# samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, -# ddim_steps=ddim_steps,eta=ddim_eta, -# quantize_denoised=True) -# # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True, -# # quantize_denoised=True) -# x_samples = self.decode_first_stage(samples.to(self.device)) -# log["samples_x0_quantized"] = x_samples -# -# if inpaint: -# # make a simple center square -# h, w = z.shape[2], z.shape[3] -# mask = torch.ones(N, h, w).to(self.device) -# # zeros will be filled in -# mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0. -# mask = mask[:, None, ...] -# with self.ema_scope("Plotting Inpaint"): -# -# samples, _ = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, eta=ddim_eta, -# ddim_steps=ddim_steps, x0=z[:N], mask=mask) -# x_samples = self.decode_first_stage(samples.to(self.device)) -# log["samples_inpainting"] = x_samples -# log["mask"] = mask -# -# # outpaint -# with self.ema_scope("Plotting Outpaint"): -# samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,eta=ddim_eta, -# ddim_steps=ddim_steps, x0=z[:N], mask=mask) -# x_samples = self.decode_first_stage(samples.to(self.device)) -# log["samples_outpainting"] = x_samples -# -# if plot_progressive_rows: -# with self.ema_scope("Plotting Progressives"): -# img, progressives = self.progressive_denoising(c, -# shape=(self.channels, self.image_size, self.image_size), -# batch_size=N) -# prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation") -# log["progressive_row"] = prog_row -# -# if return_keys: -# if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: -# return log -# else: -# return {key: log[key] for key in return_keys} -# return log -# -# def configure_optimizers(self): -# lr = self.learning_rate -# params = list(self.model.parameters()) -# if self.cond_stage_trainable: -# print(f"{self.__class__.__name__}: Also optimizing conditioner params!") -# params = params + list(self.cond_stage_model.parameters()) -# if self.learn_logvar: -# print('Diffusion model optimizing logvar') -# params.append(self.logvar) -# opt = torch.optim.AdamW(params, lr=lr) -# if self.use_scheduler: -# assert 'target' in self.scheduler_config -# scheduler = instantiate_from_config(self.scheduler_config) -# -# print("Setting up LambdaLR scheduler...") -# scheduler = [ -# { -# 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule), -# 'interval': 'step', -# 'frequency': 1 -# }] -# return [opt], scheduler -# return opt -# -# @torch.no_grad() -# def to_rgb(self, x): -# x = x.float() -# if not hasattr(self, "colorize"): -# self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x) -# x = nn.functional.conv2d(x, weight=self.colorize) -# x = 2. * (x - x.min()) / (x.max() - x.min()) - 1. -# return x -# -# -# class DiffusionWrapper(pl.LightningModule): -# def __init__(self, diff_model_config, conditioning_key): -# super().__init__() -# self.diffusion_model = instantiate_from_config(diff_model_config) -# self.conditioning_key = conditioning_key -# assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm'] -# -# def forward(self, x, t, c_concat: list = None, c_crossattn: list = None): -# if self.conditioning_key is None: -# out = self.diffusion_model(x, t) -# elif self.conditioning_key == 'concat': -# xc = torch.cat([x] + c_concat, dim=1) -# out = self.diffusion_model(xc, t) -# elif self.conditioning_key == 'crossattn': -# cc = torch.cat(c_crossattn, 1) -# out = self.diffusion_model(x, t, context=cc) -# elif self.conditioning_key == 'hybrid': -# xc = torch.cat([x] + c_concat, dim=1) -# cc = torch.cat(c_crossattn, 1) -# out = self.diffusion_model(xc, t, context=cc) -# elif self.conditioning_key == 'adm': -# cc = c_crossattn[0] -# out = self.diffusion_model(x, t, y=cc) -# else: -# raise NotImplementedError() -# -# return out -# -# -# class Layout2ImgDiffusion(LatentDiffusion): -# # TODO: move all layout-specific hacks to this class -# def __init__(self, cond_stage_key, *args, **kwargs): -# assert cond_stage_key == 'coordinates_bbox', 'Layout2ImgDiffusion only for cond_stage_key="coordinates_bbox"' -# super().__init__(*args, cond_stage_key=cond_stage_key, **kwargs) -# -# def log_images(self, batch, N=8, *args, **kwargs): -# logs = super().log_images(*args, batch=batch, N=N, **kwargs) -# -# key = 'train' if self.training else 'validation' -# dset = self.trainer.datamodule.datasets[key] -# mapper = dset.conditional_builders[self.cond_stage_key] -# -# bbox_imgs = [] -# map_fn = lambda catno: dset.get_textual_label(dset.get_category_id(catno)) -# for tknzd_bbox in batch[self.cond_stage_key][:N]: -# bboximg = mapper.plot(tknzd_bbox.detach().cpu(), map_fn, (256, 256)) -# bbox_imgs.append(bboximg) -# -# cond_img = torch.stack(bbox_imgs, dim=0) -# logs['bbox_image'] = cond_img -# return logs diff --git a/modules/npu_specific.py b/modules/npu_specific.py deleted file mode 100644 index 66ba3102..00000000 --- a/modules/npu_specific.py +++ /dev/null @@ -1,31 +0,0 @@ -# import importlib -# import torch -# -# from modules import shared -# -# -# def check_for_npu(): -# if importlib.util.find_spec("torch_npu") is None: -# return False -# import torch_npu -# -# try: -# # Will raise a RuntimeError if no NPU is found -# _ = torch_npu.npu.device_count() -# return torch.npu.is_available() -# except RuntimeError: -# return False -# -# -# def get_npu_device_string(): -# if shared.cmd_opts.device_id is not None: -# return f"npu:{shared.cmd_opts.device_id}" -# return "npu:0" -# -# -# def torch_npu_gc(): -# with torch.npu.device(get_npu_device_string()): -# torch.npu.empty_cache() -# -# -# has_npu = check_for_npu() diff --git a/modules/sd_disable_initialization.py b/modules/sd_disable_initialization.py deleted file mode 100644 index 3bff9255..00000000 --- a/modules/sd_disable_initialization.py +++ /dev/null @@ -1,232 +0,0 @@ -# import ldm.modules.encoders.modules -# import open_clip -# import torch -# import transformers.utils.hub -# -# from modules import shared -# -# -# class ReplaceHelper: -# def __init__(self): -# self.replaced = [] -# -# def replace(self, obj, field, func): -# original = getattr(obj, field, None) -# if original is None: -# return None -# -# self.replaced.append((obj, field, original)) -# setattr(obj, field, func) -# -# return original -# -# def restore(self): -# for obj, field, original in self.replaced: -# setattr(obj, field, original) -# -# self.replaced.clear() -# -# -# class DisableInitialization(ReplaceHelper): -# """ -# When an object of this class enters a `with` block, it starts: -# - preventing torch's layer initialization functions from working -# - changes CLIP and OpenCLIP to not download model weights -# - changes CLIP to not make requests to check if there is a new version of a file you already have -# -# When it leaves the block, it reverts everything to how it was before. -# -# Use it like this: -# ``` -# with DisableInitialization(): -# do_things() -# ``` -# """ -# -# def __init__(self, disable_clip=True): -# super().__init__() -# self.disable_clip = disable_clip -# -# def replace(self, obj, field, func): -# original = getattr(obj, field, None) -# if original is None: -# return None -# -# self.replaced.append((obj, field, original)) -# setattr(obj, field, func) -# -# return original -# -# def __enter__(self): -# def do_nothing(*args, **kwargs): -# pass -# -# def create_model_and_transforms_without_pretrained(*args, pretrained=None, **kwargs): -# return self.create_model_and_transforms(*args, pretrained=None, **kwargs) -# -# def CLIPTextModel_from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs): -# res = self.CLIPTextModel_from_pretrained(None, *model_args, config=pretrained_model_name_or_path, state_dict={}, **kwargs) -# res.name_or_path = pretrained_model_name_or_path -# return res -# -# def transformers_modeling_utils_load_pretrained_model(*args, **kwargs): -# args = args[0:3] + ('/', ) + args[4:] # resolved_archive_file; must set it to something to prevent what seems to be a bug -# return self.transformers_modeling_utils_load_pretrained_model(*args, **kwargs) -# -# def transformers_utils_hub_get_file_from_cache(original, url, *args, **kwargs): -# -# # this file is always 404, prevent making request -# if url == 'https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/added_tokens.json' or url == 'openai/clip-vit-large-patch14' and args[0] == 'added_tokens.json': -# return None -# -# try: -# res = original(url, *args, local_files_only=True, **kwargs) -# if res is None: -# res = original(url, *args, local_files_only=False, **kwargs) -# return res -# except Exception: -# return original(url, *args, local_files_only=False, **kwargs) -# -# def transformers_utils_hub_get_from_cache(url, *args, local_files_only=False, **kwargs): -# return transformers_utils_hub_get_file_from_cache(self.transformers_utils_hub_get_from_cache, url, *args, **kwargs) -# -# def transformers_tokenization_utils_base_cached_file(url, *args, local_files_only=False, **kwargs): -# return transformers_utils_hub_get_file_from_cache(self.transformers_tokenization_utils_base_cached_file, url, *args, **kwargs) -# -# def transformers_configuration_utils_cached_file(url, *args, local_files_only=False, **kwargs): -# return transformers_utils_hub_get_file_from_cache(self.transformers_configuration_utils_cached_file, url, *args, **kwargs) -# -# self.replace(torch.nn.init, 'kaiming_uniform_', do_nothing) -# self.replace(torch.nn.init, '_no_grad_normal_', do_nothing) -# self.replace(torch.nn.init, '_no_grad_uniform_', do_nothing) -# -# if self.disable_clip: -# self.create_model_and_transforms = self.replace(open_clip, 'create_model_and_transforms', create_model_and_transforms_without_pretrained) -# self.CLIPTextModel_from_pretrained = self.replace(ldm.modules.encoders.modules.CLIPTextModel, 'from_pretrained', CLIPTextModel_from_pretrained) -# self.transformers_modeling_utils_load_pretrained_model = self.replace(transformers.modeling_utils.PreTrainedModel, '_load_pretrained_model', transformers_modeling_utils_load_pretrained_model) -# self.transformers_tokenization_utils_base_cached_file = self.replace(transformers.tokenization_utils_base, 'cached_file', transformers_tokenization_utils_base_cached_file) -# self.transformers_configuration_utils_cached_file = self.replace(transformers.configuration_utils, 'cached_file', transformers_configuration_utils_cached_file) -# self.transformers_utils_hub_get_from_cache = self.replace(transformers.utils.hub, 'get_from_cache', transformers_utils_hub_get_from_cache) -# -# def __exit__(self, exc_type, exc_val, exc_tb): -# self.restore() -# -# -# class InitializeOnMeta(ReplaceHelper): -# """ -# Context manager that causes all parameters for linear/conv2d/mha layers to be allocated on meta device, -# which results in those parameters having no values and taking no memory. model.to() will be broken and -# will need to be repaired by using LoadStateDictOnMeta below when loading params from state dict. -# -# Usage: -# ``` -# with sd_disable_initialization.InitializeOnMeta(): -# sd_model = instantiate_from_config(sd_config.model) -# ``` -# """ -# -# def __enter__(self): -# if shared.cmd_opts.disable_model_loading_ram_optimization: -# return -# -# def set_device(x): -# x["device"] = "meta" -# return x -# -# linear_init = self.replace(torch.nn.Linear, '__init__', lambda *args, **kwargs: linear_init(*args, **set_device(kwargs))) -# conv2d_init = self.replace(torch.nn.Conv2d, '__init__', lambda *args, **kwargs: conv2d_init(*args, **set_device(kwargs))) -# mha_init = self.replace(torch.nn.MultiheadAttention, '__init__', lambda *args, **kwargs: mha_init(*args, **set_device(kwargs))) -# self.replace(torch.nn.Module, 'to', lambda *args, **kwargs: None) -# -# def __exit__(self, exc_type, exc_val, exc_tb): -# self.restore() -# -# -# class LoadStateDictOnMeta(ReplaceHelper): -# """ -# Context manager that allows to read parameters from state_dict into a model that has some of its parameters in the meta device. -# As those parameters are read from state_dict, they will be deleted from it, so by the end state_dict will be mostly empty, to save memory. -# Meant to be used together with InitializeOnMeta above. -# -# Usage: -# ``` -# with sd_disable_initialization.LoadStateDictOnMeta(state_dict): -# model.load_state_dict(state_dict, strict=False) -# ``` -# """ -# -# def __init__(self, state_dict, device, weight_dtype_conversion=None): -# super().__init__() -# self.state_dict = state_dict -# self.device = device -# self.weight_dtype_conversion = weight_dtype_conversion or {} -# self.default_dtype = self.weight_dtype_conversion.get('') -# -# def get_weight_dtype(self, key): -# key_first_term, _ = key.split('.', 1) -# return self.weight_dtype_conversion.get(key_first_term, self.default_dtype) -# -# def __enter__(self): -# if shared.cmd_opts.disable_model_loading_ram_optimization: -# return -# -# sd = self.state_dict -# device = self.device -# -# def load_from_state_dict(original, module, state_dict, prefix, *args, **kwargs): -# used_param_keys = [] -# -# for name, param in module._parameters.items(): -# if param is None: -# continue -# -# key = prefix + name -# sd_param = sd.pop(key, None) -# if sd_param is not None: -# state_dict[key] = sd_param.to(dtype=self.get_weight_dtype(key)) -# used_param_keys.append(key) -# -# if param.is_meta: -# dtype = sd_param.dtype if sd_param is not None else param.dtype -# module._parameters[name] = torch.nn.parameter.Parameter(torch.zeros_like(param, device=device, dtype=dtype), requires_grad=param.requires_grad) -# -# for name in module._buffers: -# key = prefix + name -# -# sd_param = sd.pop(key, None) -# if sd_param is not None: -# state_dict[key] = sd_param -# used_param_keys.append(key) -# -# original(module, state_dict, prefix, *args, **kwargs) -# -# for key in used_param_keys: -# state_dict.pop(key, None) -# -# def load_state_dict(original, module, state_dict, strict=True): -# """torch makes a lot of copies of the dictionary with weights, so just deleting entries from state_dict does not help -# because the same values are stored in multiple copies of the dict. The trick used here is to give torch a dict with -# all weights on meta device, i.e. deleted, and then it doesn't matter how many copies torch makes. -# -# In _load_from_state_dict, the correct weight will be obtained from a single dict with the right weights (sd). -# -# The dangerous thing about this is if _load_from_state_dict is not called, (if some exotic module overloads -# the function and does not call the original) the state dict will just fail to load because weights -# would be on the meta device. -# """ -# -# if state_dict is sd: -# state_dict = {k: v.to(device="meta", dtype=v.dtype) for k, v in state_dict.items()} -# -# original(module, state_dict, strict=strict) -# -# module_load_state_dict = self.replace(torch.nn.Module, 'load_state_dict', lambda *args, **kwargs: load_state_dict(module_load_state_dict, *args, **kwargs)) -# module_load_from_state_dict = self.replace(torch.nn.Module, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(module_load_from_state_dict, *args, **kwargs)) -# linear_load_from_state_dict = self.replace(torch.nn.Linear, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(linear_load_from_state_dict, *args, **kwargs)) -# conv2d_load_from_state_dict = self.replace(torch.nn.Conv2d, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(conv2d_load_from_state_dict, *args, **kwargs)) -# mha_load_from_state_dict = self.replace(torch.nn.MultiheadAttention, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(mha_load_from_state_dict, *args, **kwargs)) -# layer_norm_load_from_state_dict = self.replace(torch.nn.LayerNorm, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(layer_norm_load_from_state_dict, *args, **kwargs)) -# group_norm_load_from_state_dict = self.replace(torch.nn.GroupNorm, '_load_from_state_dict', lambda *args, **kwargs: load_from_state_dict(group_norm_load_from_state_dict, *args, **kwargs)) -# -# def __exit__(self, exc_type, exc_val, exc_tb): -# self.restore() diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 654f576b..7702a7f5 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -27,233 +27,3 @@ class StableDiffusionModelHijack: model_hijack = StableDiffusionModelHijack() - -# import torch -# from torch.nn.functional import silu -# from types import MethodType -# -# from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet, patches -# from modules.hypernetworks import hypernetwork -# from modules.shared import cmd_opts -# from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr, xlmr_m18 -# -# import ldm.modules.attention -# import ldm.modules.diffusionmodules.model -# import ldm.modules.diffusionmodules.openaimodel -# import ldm.models.diffusion.ddpm -# import ldm.models.diffusion.ddim -# import ldm.models.diffusion.plms -# import ldm.modules.encoders.modules -# -# import sgm.modules.attention -# import sgm.modules.diffusionmodules.model -# import sgm.modules.diffusionmodules.openaimodel -# import sgm.modules.encoders.modules -# -# attention_CrossAttention_forward = ldm.modules.attention.CrossAttention.forward -# diffusionmodules_model_nonlinearity = ldm.modules.diffusionmodules.model.nonlinearity -# diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.AttnBlock.forward -# -# # new memory efficient cross attention blocks do not support hypernets and we already -# # have memory efficient cross attention anyway, so this disables SD2.0's memory efficient cross attention -# ldm.modules.attention.MemoryEfficientCrossAttention = ldm.modules.attention.CrossAttention -# ldm.modules.attention.BasicTransformerBlock.ATTENTION_MODES["softmax-xformers"] = ldm.modules.attention.CrossAttention -# -# # silence new console spam from SD2 -# ldm.modules.attention.print = shared.ldm_print -# ldm.modules.diffusionmodules.model.print = shared.ldm_print -# ldm.util.print = shared.ldm_print -# ldm.models.diffusion.ddpm.print = shared.ldm_print -# -# optimizers = [] -# current_optimizer: sd_hijack_optimizations.SdOptimization = None -# -# ldm_patched_forward = sd_unet.create_unet_forward(ldm.modules.diffusionmodules.openaimodel.UNetModel.forward) -# ldm_original_forward = patches.patch(__file__, ldm.modules.diffusionmodules.openaimodel.UNetModel, "forward", ldm_patched_forward) -# -# sgm_patched_forward = sd_unet.create_unet_forward(sgm.modules.diffusionmodules.openaimodel.UNetModel.forward) -# sgm_original_forward = patches.patch(__file__, sgm.modules.diffusionmodules.openaimodel.UNetModel, "forward", sgm_patched_forward) -# -# -# def list_optimizers(): -# new_optimizers = script_callbacks.list_optimizers_callback() -# -# new_optimizers = [x for x in new_optimizers if x.is_available()] -# -# new_optimizers = sorted(new_optimizers, key=lambda x: x.priority, reverse=True) -# -# optimizers.clear() -# optimizers.extend(new_optimizers) -# -# -# def apply_optimizations(option=None): -# return -# -# -# def undo_optimizations(): -# return -# -# -# def fix_checkpoint(): -# """checkpoints are now added and removed in embedding/hypernet code, since torch doesn't want -# checkpoints to be added when not training (there's a warning)""" -# -# pass -# -# -# def weighted_loss(sd_model, pred, target, mean=True): -# #Calculate the weight normally, but ignore the mean -# loss = sd_model._old_get_loss(pred, target, mean=False) -# -# #Check if we have weights available -# weight = getattr(sd_model, '_custom_loss_weight', None) -# if weight is not None: -# loss *= weight -# -# #Return the loss, as mean if specified -# return loss.mean() if mean else loss -# -# def weighted_forward(sd_model, x, c, w, *args, **kwargs): -# try: -# #Temporarily append weights to a place accessible during loss calc -# sd_model._custom_loss_weight = w -# -# #Replace 'get_loss' with a weight-aware one. Otherwise we need to reimplement 'forward' completely -# #Keep 'get_loss', but don't overwrite the previous old_get_loss if it's already set -# if not hasattr(sd_model, '_old_get_loss'): -# sd_model._old_get_loss = sd_model.get_loss -# sd_model.get_loss = MethodType(weighted_loss, sd_model) -# -# #Run the standard forward function, but with the patched 'get_loss' -# return sd_model.forward(x, c, *args, **kwargs) -# finally: -# try: -# #Delete temporary weights if appended -# del sd_model._custom_loss_weight -# except AttributeError: -# pass -# -# #If we have an old loss function, reset the loss function to the original one -# if hasattr(sd_model, '_old_get_loss'): -# sd_model.get_loss = sd_model._old_get_loss -# del sd_model._old_get_loss -# -# def apply_weighted_forward(sd_model): -# #Add new function 'weighted_forward' that can be called to calc weighted loss -# sd_model.weighted_forward = MethodType(weighted_forward, sd_model) -# -# def undo_weighted_forward(sd_model): -# try: -# del sd_model.weighted_forward -# except AttributeError: -# pass -# -# -# class StableDiffusionModelHijack: -# fixes = None -# layers = None -# circular_enabled = False -# clip = None -# optimization_method = None -# -# def __init__(self): -# self.extra_generation_params = {} -# self.comments = [] -# -# def apply_optimizations(self, option=None): -# pass -# -# def convert_sdxl_to_ssd(self, m): -# pass -# -# def hijack(self, m): -# pass -# -# def undo_hijack(self, m): -# pass -# -# def apply_circular(self, enable): -# pass -# -# def clear_comments(self): -# self.comments = [] -# self.extra_generation_params = {} -# -# def get_prompt_lengths(self, text, cond_stage_model): -# pass -# -# def redo_hijack(self, m): -# pass -# -# -# class EmbeddingsWithFixes(torch.nn.Module): -# def __init__(self, wrapped, embeddings, textual_inversion_key='clip_l'): -# super().__init__() -# self.wrapped = wrapped -# self.embeddings = embeddings -# self.textual_inversion_key = textual_inversion_key -# self.weight = self.wrapped.weight -# -# def forward(self, input_ids): -# batch_fixes = self.embeddings.fixes -# self.embeddings.fixes = None -# -# inputs_embeds = self.wrapped(input_ids) -# -# if batch_fixes is None or len(batch_fixes) == 0 or max([len(x) for x in batch_fixes]) == 0: -# return inputs_embeds -# -# vecs = [] -# for fixes, tensor in zip(batch_fixes, inputs_embeds): -# for offset, embedding in fixes: -# vec = embedding.vec[self.textual_inversion_key] if isinstance(embedding.vec, dict) else embedding.vec -# emb = devices.cond_cast_unet(vec) -# emb_len = min(tensor.shape[0] - offset - 1, emb.shape[0]) -# tensor = torch.cat([tensor[0:offset + 1], emb[0:emb_len], tensor[offset + 1 + emb_len:]]).to(dtype=inputs_embeds.dtype) -# -# vecs.append(tensor) -# -# return torch.stack(vecs) -# -# -# class TextualInversionEmbeddings(torch.nn.Embedding): -# def __init__(self, num_embeddings: int, embedding_dim: int, textual_inversion_key='clip_l', **kwargs): -# super().__init__(num_embeddings, embedding_dim, **kwargs) -# -# self.embeddings = model_hijack -# self.textual_inversion_key = textual_inversion_key -# -# @property -# def wrapped(self): -# return super().forward -# -# def forward(self, input_ids): -# return EmbeddingsWithFixes.forward(self, input_ids) -# -# -# def add_circular_option_to_conv_2d(): -# conv2d_constructor = torch.nn.Conv2d.__init__ -# -# def conv2d_constructor_circular(self, *args, **kwargs): -# return conv2d_constructor(self, *args, padding_mode='circular', **kwargs) -# -# torch.nn.Conv2d.__init__ = conv2d_constructor_circular -# -# -# -# -# -# def register_buffer(self, name, attr): -# """ -# Fix register buffer bug for Mac OS. -# """ -# -# if type(attr) == torch.Tensor: -# if attr.device != devices.device: -# attr = attr.to(device=devices.device, dtype=(torch.float32 if devices.device.type == 'mps' else None)) -# -# setattr(self, name, attr) -# -# -# ldm.models.diffusion.ddim.DDIMSampler.register_buffer = register_buffer -# ldm.models.diffusion.plms.PLMSSampler.register_buffer = register_buffer diff --git a/modules/sd_hijack_checkpoint.py b/modules/sd_hijack_checkpoint.py deleted file mode 100644 index 8b11b443..00000000 --- a/modules/sd_hijack_checkpoint.py +++ /dev/null @@ -1,46 +0,0 @@ -# from torch.utils.checkpoint import checkpoint -# -# import ldm.modules.attention -# import ldm.modules.diffusionmodules.openaimodel -# -# -# def BasicTransformerBlock_forward(self, x, context=None): -# return checkpoint(self._forward, x, context) -# -# -# def AttentionBlock_forward(self, x): -# return checkpoint(self._forward, x) -# -# -# def ResBlock_forward(self, x, emb): -# return checkpoint(self._forward, x, emb) -# -# -# stored = [] -# -# -# def add(): -# if len(stored) != 0: -# return -# -# stored.extend([ -# ldm.modules.attention.BasicTransformerBlock.forward, -# ldm.modules.diffusionmodules.openaimodel.ResBlock.forward, -# ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward -# ]) -# -# ldm.modules.attention.BasicTransformerBlock.forward = BasicTransformerBlock_forward -# ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = ResBlock_forward -# ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = AttentionBlock_forward -# -# -# def remove(): -# if len(stored) == 0: -# return -# -# ldm.modules.attention.BasicTransformerBlock.forward = stored[0] -# ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = stored[1] -# ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = stored[2] -# -# stored.clear() -# diff --git a/modules/sd_hijack_clip.py b/modules/sd_hijack_clip.py deleted file mode 100644 index bafc6782..00000000 --- a/modules/sd_hijack_clip.py +++ /dev/null @@ -1,384 +0,0 @@ -# import math -# from collections import namedtuple -# -# import torch -# -# from modules import prompt_parser, devices, sd_hijack, sd_emphasis -# from modules.shared import opts -# -# -# class PromptChunk: -# """ -# This object contains token ids, weight (multipliers:1.4) and textual inversion embedding info for a chunk of prompt. -# If a prompt is short, it is represented by one PromptChunk, otherwise, multiple are necessary. -# Each PromptChunk contains an exact amount of tokens - 77, which includes one for start and end token, -# so just 75 tokens from prompt. -# """ -# -# def __init__(self): -# self.tokens = [] -# self.multipliers = [] -# self.fixes = [] -# -# -# PromptChunkFix = namedtuple('PromptChunkFix', ['offset', 'embedding']) -# """An object of this type is a marker showing that textual inversion embedding's vectors have to placed at offset in the prompt -# chunk. Those objects are found in PromptChunk.fixes and, are placed into FrozenCLIPEmbedderWithCustomWordsBase.hijack.fixes, and finally -# are applied by sd_hijack.EmbeddingsWithFixes's forward function.""" -# -# -# class TextConditionalModel(torch.nn.Module): -# def __init__(self): -# super().__init__() -# -# self.hijack = sd_hijack.model_hijack -# self.chunk_length = 75 -# -# self.is_trainable = False -# self.input_key = 'txt' -# self.return_pooled = False -# -# self.comma_token = None -# self.id_start = None -# self.id_end = None -# self.id_pad = None -# -# def empty_chunk(self): -# """creates an empty PromptChunk and returns it""" -# -# chunk = PromptChunk() -# chunk.tokens = [self.id_start] + [self.id_end] * (self.chunk_length + 1) -# chunk.multipliers = [1.0] * (self.chunk_length + 2) -# return chunk -# -# def get_target_prompt_token_count(self, token_count): -# """returns the maximum number of tokens a prompt of a known length can have before it requires one more PromptChunk to be represented""" -# -# return math.ceil(max(token_count, 1) / self.chunk_length) * self.chunk_length -# -# def tokenize(self, texts): -# """Converts a batch of texts into a batch of token ids""" -# -# raise NotImplementedError -# -# def encode_with_transformers(self, tokens): -# """ -# converts a batch of token ids (in python lists) into a single tensor with numeric representation of those tokens; -# All python lists with tokens are assumed to have same length, usually 77. -# if input is a list with B elements and each element has T tokens, expected output shape is (B, T, C), where C depends on -# model - can be 768 and 1024. -# Among other things, this call will read self.hijack.fixes, apply it to its inputs, and clear it (setting it to None). -# """ -# -# raise NotImplementedError -# -# def encode_embedding_init_text(self, init_text, nvpt): -# """Converts text into a tensor with this text's tokens' embeddings. Note that those are embeddings before they are passed through -# transformers. nvpt is used as a maximum length in tokens. If text produces less teokens than nvpt, only this many is returned.""" -# -# raise NotImplementedError -# -# def tokenize_line(self, line): -# """ -# this transforms a single prompt into a list of PromptChunk objects - as many as needed to -# represent the prompt. -# Returns the list and the total number of tokens in the prompt. -# """ -# -# if opts.emphasis != "None": -# parsed = prompt_parser.parse_prompt_attention(line) -# else: -# parsed = [[line, 1.0]] -# -# tokenized = self.tokenize([text for text, _ in parsed]) -# -# chunks = [] -# chunk = PromptChunk() -# token_count = 0 -# last_comma = -1 -# -# def next_chunk(is_last=False): -# """puts current chunk into the list of results and produces the next one - empty; -# if is_last is true, tokens tokens at the end won't add to token_count""" -# nonlocal token_count -# nonlocal last_comma -# nonlocal chunk -# -# if is_last: -# token_count += len(chunk.tokens) -# else: -# token_count += self.chunk_length -# -# to_add = self.chunk_length - len(chunk.tokens) -# if to_add > 0: -# chunk.tokens += [self.id_end] * to_add -# chunk.multipliers += [1.0] * to_add -# -# chunk.tokens = [self.id_start] + chunk.tokens + [self.id_end] -# chunk.multipliers = [1.0] + chunk.multipliers + [1.0] -# -# last_comma = -1 -# chunks.append(chunk) -# chunk = PromptChunk() -# -# for tokens, (text, weight) in zip(tokenized, parsed): -# if text == 'BREAK' and weight == -1: -# next_chunk() -# continue -# -# position = 0 -# while position < len(tokens): -# token = tokens[position] -# -# if token == self.comma_token: -# last_comma = len(chunk.tokens) -# -# # this is when we are at the end of allotted 75 tokens for the current chunk, and the current token is not a comma. opts.comma_padding_backtrack -# # is a setting that specifies that if there is a comma nearby, the text after the comma should be moved out of this chunk and into the next. -# elif opts.comma_padding_backtrack != 0 and len(chunk.tokens) == self.chunk_length and last_comma != -1 and len(chunk.tokens) - last_comma <= opts.comma_padding_backtrack: -# break_location = last_comma + 1 -# -# reloc_tokens = chunk.tokens[break_location:] -# reloc_mults = chunk.multipliers[break_location:] -# -# chunk.tokens = chunk.tokens[:break_location] -# chunk.multipliers = chunk.multipliers[:break_location] -# -# next_chunk() -# chunk.tokens = reloc_tokens -# chunk.multipliers = reloc_mults -# -# if len(chunk.tokens) == self.chunk_length: -# next_chunk() -# -# embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, position) -# if embedding is None: -# chunk.tokens.append(token) -# chunk.multipliers.append(weight) -# position += 1 -# continue -# -# emb_len = int(embedding.vectors) -# if len(chunk.tokens) + emb_len > self.chunk_length: -# next_chunk() -# -# chunk.fixes.append(PromptChunkFix(len(chunk.tokens), embedding)) -# -# chunk.tokens += [0] * emb_len -# chunk.multipliers += [weight] * emb_len -# position += embedding_length_in_tokens -# -# if chunk.tokens or not chunks: -# next_chunk(is_last=True) -# -# return chunks, token_count -# -# def process_texts(self, texts): -# """ -# Accepts a list of texts and calls tokenize_line() on each, with cache. Returns the list of results and maximum -# length, in tokens, of all texts. -# """ -# -# token_count = 0 -# -# cache = {} -# batch_chunks = [] -# for line in texts: -# if line in cache: -# chunks = cache[line] -# else: -# chunks, current_token_count = self.tokenize_line(line) -# token_count = max(current_token_count, token_count) -# -# cache[line] = chunks -# -# batch_chunks.append(chunks) -# -# return batch_chunks, token_count -# -# def forward(self, texts): -# """ -# Accepts an array of texts; Passes texts through transformers network to create a tensor with numerical representation of those texts. -# Returns a tensor with shape of (B, T, C), where B is length of the array; T is length, in tokens, of texts (including padding) - T will -# be a multiple of 77; and C is dimensionality of each token - for SD1 it's 768, for SD2 it's 1024, and for SDXL it's 1280. -# An example shape returned by this function can be: (2, 77, 768). -# For SDXL, instead of returning one tensor avobe, it returns a tuple with two: the other one with shape (B, 1280) with pooled values. -# Webui usually sends just one text at a time through this function - the only time when texts is an array with more than one element -# is when you do prompt editing: "a picture of a [cat:dog:0.4] eating ice cream" -# """ -# -# batch_chunks, token_count = self.process_texts(texts) -# -# used_embeddings = {} -# chunk_count = max([len(x) for x in batch_chunks]) -# -# zs = [] -# for i in range(chunk_count): -# batch_chunk = [chunks[i] if i < len(chunks) else self.empty_chunk() for chunks in batch_chunks] -# -# tokens = [x.tokens for x in batch_chunk] -# multipliers = [x.multipliers for x in batch_chunk] -# self.hijack.fixes = [x.fixes for x in batch_chunk] -# -# for fixes in self.hijack.fixes: -# for _position, embedding in fixes: -# used_embeddings[embedding.name] = embedding -# devices.torch_npu_set_device() -# z = self.process_tokens(tokens, multipliers) -# zs.append(z) -# -# if opts.textual_inversion_add_hashes_to_infotext and used_embeddings: -# hashes = [] -# for name, embedding in used_embeddings.items(): -# shorthash = embedding.shorthash -# if not shorthash: -# continue -# -# name = name.replace(":", "").replace(",", "") -# hashes.append(f"{name}: {shorthash}") -# -# if hashes: -# if self.hijack.extra_generation_params.get("TI hashes"): -# hashes.append(self.hijack.extra_generation_params.get("TI hashes")) -# self.hijack.extra_generation_params["TI hashes"] = ", ".join(hashes) -# -# if any(x for x in texts if "(" in x or "[" in x) and opts.emphasis != "Original": -# self.hijack.extra_generation_params["Emphasis"] = opts.emphasis -# -# if self.return_pooled: -# return torch.hstack(zs), zs[0].pooled -# else: -# return torch.hstack(zs) -# -# def process_tokens(self, remade_batch_tokens, batch_multipliers): -# """ -# sends one single prompt chunk to be encoded by transformers neural network. -# remade_batch_tokens is a batch of tokens - a list, where every element is a list of tokens; usually -# there are exactly 77 tokens in the list. batch_multipliers is the same but for multipliers instead of tokens. -# Multipliers are used to give more or less weight to the outputs of transformers network. Each multiplier -# corresponds to one token. -# """ -# tokens = torch.asarray(remade_batch_tokens).to(devices.device) -# -# # this is for SD2: SD1 uses the same token for padding and end of text, while SD2 uses different ones. -# if self.id_end != self.id_pad: -# for batch_pos in range(len(remade_batch_tokens)): -# index = remade_batch_tokens[batch_pos].index(self.id_end) -# tokens[batch_pos, index+1:tokens.shape[1]] = self.id_pad -# -# z = self.encode_with_transformers(tokens) -# -# pooled = getattr(z, 'pooled', None) -# -# emphasis = sd_emphasis.get_current_option(opts.emphasis)() -# emphasis.tokens = remade_batch_tokens -# emphasis.multipliers = torch.asarray(batch_multipliers).to(devices.device) -# emphasis.z = z -# -# emphasis.after_transformers() -# -# z = emphasis.z -# -# if pooled is not None: -# z.pooled = pooled -# -# return z -# -# -# class FrozenCLIPEmbedderWithCustomWordsBase(TextConditionalModel): -# """A pytorch module that is a wrapper for FrozenCLIPEmbedder module. it enhances FrozenCLIPEmbedder, making it possible to -# have unlimited prompt length and assign weights to tokens in prompt. -# """ -# -# def __init__(self, wrapped, hijack): -# super().__init__() -# -# self.hijack = hijack -# -# self.wrapped = wrapped -# """Original FrozenCLIPEmbedder module; can also be FrozenOpenCLIPEmbedder or xlmr.BertSeriesModelWithTransformation, -# depending on model.""" -# -# self.is_trainable = getattr(wrapped, 'is_trainable', False) -# self.input_key = getattr(wrapped, 'input_key', 'txt') -# self.return_pooled = getattr(self.wrapped, 'return_pooled', False) -# -# self.legacy_ucg_val = None # for sgm codebase -# -# def forward(self, texts): -# if opts.use_old_emphasis_implementation: -# import modules.sd_hijack_clip_old -# return modules.sd_hijack_clip_old.forward_old(self, texts) -# -# return super().forward(texts) -# -# -# class FrozenCLIPEmbedderWithCustomWords(FrozenCLIPEmbedderWithCustomWordsBase): -# def __init__(self, wrapped, hijack): -# super().__init__(wrapped, hijack) -# self.tokenizer = wrapped.tokenizer -# -# vocab = self.tokenizer.get_vocab() -# -# self.comma_token = vocab.get(',', None) -# -# self.token_mults = {} -# tokens_with_parens = [(k, v) for k, v in vocab.items() if '(' in k or ')' in k or '[' in k or ']' in k] -# for text, ident in tokens_with_parens: -# mult = 1.0 -# for c in text: -# if c == '[': -# mult /= 1.1 -# if c == ']': -# mult *= 1.1 -# if c == '(': -# mult *= 1.1 -# if c == ')': -# mult /= 1.1 -# -# if mult != 1.0: -# self.token_mults[ident] = mult -# -# self.id_start = self.wrapped.tokenizer.bos_token_id -# self.id_end = self.wrapped.tokenizer.eos_token_id -# self.id_pad = self.id_end -# -# def tokenize(self, texts): -# tokenized = self.wrapped.tokenizer(texts, truncation=False, add_special_tokens=False)["input_ids"] -# -# return tokenized -# -# def encode_with_transformers(self, tokens): -# outputs = self.wrapped.transformer(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers) -# -# if opts.CLIP_stop_at_last_layers > 1: -# z = outputs.hidden_states[-opts.CLIP_stop_at_last_layers] -# z = self.wrapped.transformer.text_model.final_layer_norm(z) -# else: -# z = outputs.last_hidden_state -# -# return z -# -# def encode_embedding_init_text(self, init_text, nvpt): -# embedding_layer = self.wrapped.transformer.text_model.embeddings -# ids = self.wrapped.tokenizer(init_text, max_length=nvpt, return_tensors="pt", add_special_tokens=False)["input_ids"] -# embedded = embedding_layer.token_embedding.wrapped(ids.to(embedding_layer.token_embedding.wrapped.weight.device)).squeeze(0) -# -# return embedded -# -# -# class FrozenCLIPEmbedderForSDXLWithCustomWords(FrozenCLIPEmbedderWithCustomWords): -# def __init__(self, wrapped, hijack): -# super().__init__(wrapped, hijack) -# -# def encode_with_transformers(self, tokens): -# outputs = self.wrapped.transformer(input_ids=tokens, output_hidden_states=self.wrapped.layer == "hidden") -# -# if opts.sdxl_clip_l_skip is True: -# z = outputs.hidden_states[-opts.CLIP_stop_at_last_layers] -# elif self.wrapped.layer == "last": -# z = outputs.last_hidden_state -# else: -# z = outputs.hidden_states[self.wrapped.layer_idx] -# -# return z diff --git a/modules/sd_hijack_clip_old.py b/modules/sd_hijack_clip_old.py deleted file mode 100644 index 803a723d..00000000 --- a/modules/sd_hijack_clip_old.py +++ /dev/null @@ -1,82 +0,0 @@ -# from modules import sd_hijack_clip -# from modules import shared -# -# -# def process_text_old(self: sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase, texts): -# id_start = self.id_start -# id_end = self.id_end -# maxlen = self.wrapped.max_length # you get to stay at 77 -# used_custom_terms = [] -# remade_batch_tokens = [] -# hijack_comments = [] -# hijack_fixes = [] -# token_count = 0 -# -# cache = {} -# batch_tokens = self.tokenize(texts) -# batch_multipliers = [] -# for tokens in batch_tokens: -# tuple_tokens = tuple(tokens) -# -# if tuple_tokens in cache: -# remade_tokens, fixes, multipliers = cache[tuple_tokens] -# else: -# fixes = [] -# remade_tokens = [] -# multipliers = [] -# mult = 1.0 -# -# i = 0 -# while i < len(tokens): -# token = tokens[i] -# -# embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i) -# -# mult_change = self.token_mults.get(token) if shared.opts.emphasis != "None" else None -# if mult_change is not None: -# mult *= mult_change -# i += 1 -# elif embedding is None: -# remade_tokens.append(token) -# multipliers.append(mult) -# i += 1 -# else: -# emb_len = int(embedding.vec.shape[0]) -# fixes.append((len(remade_tokens), embedding)) -# remade_tokens += [0] * emb_len -# multipliers += [mult] * emb_len -# used_custom_terms.append((embedding.name, embedding.checksum())) -# i += embedding_length_in_tokens -# -# if len(remade_tokens) > maxlen - 2: -# vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()} -# ovf = remade_tokens[maxlen - 2:] -# overflowing_words = [vocab.get(int(x), "") for x in ovf] -# overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words)) -# hijack_comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n") -# -# token_count = len(remade_tokens) -# remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens)) -# remade_tokens = [id_start] + remade_tokens[0:maxlen - 2] + [id_end] -# cache[tuple_tokens] = (remade_tokens, fixes, multipliers) -# -# multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers)) -# multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0] -# -# remade_batch_tokens.append(remade_tokens) -# hijack_fixes.append(fixes) -# batch_multipliers.append(multipliers) -# return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count -# -# -# def forward_old(self: sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase, texts): -# batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = process_text_old(self, texts) -# -# self.hijack.comments += hijack_comments -# -# if used_custom_terms: -# embedding_names = ", ".join(f"{word} [{checksum}]" for word, checksum in used_custom_terms) -# self.hijack.comments.append(f"Used embeddings: {embedding_names}") -# -# self.hijack.fixes = hijack_fixes -# return self.process_tokens(remade_batch_tokens, batch_multipliers) diff --git a/modules/sd_hijack_ip2p.py b/modules/sd_hijack_ip2p.py deleted file mode 100644 index 9a4d838b..00000000 --- a/modules/sd_hijack_ip2p.py +++ /dev/null @@ -1,10 +0,0 @@ -# import os.path -# -# -# def should_hijack_ip2p(checkpoint_info): -# from modules import sd_models_config -# -# ckpt_basename = os.path.basename(checkpoint_info.filename).lower() -# cfg_basename = os.path.basename(sd_models_config.find_checkpoint_config_near_filename(checkpoint_info)).lower() -# -# return "pix2pix" in ckpt_basename and "pix2pix" not in cfg_basename diff --git a/modules/sd_hijack_open_clip.py b/modules/sd_hijack_open_clip.py deleted file mode 100644 index 96f33720..00000000 --- a/modules/sd_hijack_open_clip.py +++ /dev/null @@ -1,71 +0,0 @@ -# import open_clip.tokenizer -# import torch -# -# from modules import sd_hijack_clip, devices -# from modules.shared import opts -# -# tokenizer = open_clip.tokenizer._tokenizer -# -# -# class FrozenOpenCLIPEmbedderWithCustomWords(sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase): -# def __init__(self, wrapped, hijack): -# super().__init__(wrapped, hijack) -# -# self.comma_token = [v for k, v in tokenizer.encoder.items() if k == ','][0] -# self.id_start = tokenizer.encoder[""] -# self.id_end = tokenizer.encoder[""] -# self.id_pad = 0 -# -# def tokenize(self, texts): -# assert not opts.use_old_emphasis_implementation, 'Old emphasis implementation not supported for Open Clip' -# -# tokenized = [tokenizer.encode(text) for text in texts] -# -# return tokenized -# -# def encode_with_transformers(self, tokens): -# # set self.wrapped.layer_idx here according to opts.CLIP_stop_at_last_layers -# z = self.wrapped.encode_with_transformer(tokens) -# -# return z -# -# def encode_embedding_init_text(self, init_text, nvpt): -# ids = tokenizer.encode(init_text) -# ids = torch.asarray([ids], device=devices.device, dtype=torch.int) -# embedded = self.wrapped.model.token_embedding.wrapped(ids).squeeze(0) -# -# return embedded -# -# -# class FrozenOpenCLIPEmbedder2WithCustomWords(sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase): -# def __init__(self, wrapped, hijack): -# super().__init__(wrapped, hijack) -# -# self.comma_token = [v for k, v in tokenizer.encoder.items() if k == ','][0] -# self.id_start = tokenizer.encoder[""] -# self.id_end = tokenizer.encoder[""] -# self.id_pad = 0 -# -# def tokenize(self, texts): -# assert not opts.use_old_emphasis_implementation, 'Old emphasis implementation not supported for Open Clip' -# -# tokenized = [tokenizer.encode(text) for text in texts] -# -# return tokenized -# -# def encode_with_transformers(self, tokens): -# d = self.wrapped.encode_with_transformer(tokens) -# z = d[self.wrapped.layer] -# -# pooled = d.get("pooled") -# if pooled is not None: -# z.pooled = pooled -# -# return z -# -# def encode_embedding_init_text(self, init_text, nvpt): -# ids = tokenizer.encode(init_text) -# ids = torch.asarray([ids], device=devices.device, dtype=torch.int) -# embedded = self.wrapped.model.token_embedding.wrapped(ids.to(self.wrapped.model.token_embedding.wrapped.weight.device)).squeeze(0) -# -# return embedded diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py deleted file mode 100644 index 696835ad..00000000 --- a/modules/sd_hijack_optimizations.py +++ /dev/null @@ -1,677 +0,0 @@ -# from __future__ import annotations -# import math -# import psutil -# import platform -# -# import torch -# from torch import einsum -# -# from ldm.util import default -# from einops import rearrange -# -# from modules import shared, errors, devices, sub_quadratic_attention -# from modules.hypernetworks import hypernetwork -# -# import ldm.modules.attention -# import ldm.modules.diffusionmodules.model -# -# import sgm.modules.attention -# import sgm.modules.diffusionmodules.model -# -# diffusionmodules_model_AttnBlock_forward = ldm.modules.diffusionmodules.model.AttnBlock.forward -# sgm_diffusionmodules_model_AttnBlock_forward = sgm.modules.diffusionmodules.model.AttnBlock.forward -# -# -# class SdOptimization: -# name: str = None -# label: str | None = None -# cmd_opt: str | None = None -# priority: int = 0 -# -# def title(self): -# if self.label is None: -# return self.name -# -# return f"{self.name} - {self.label}" -# -# def is_available(self): -# return True -# -# def apply(self): -# pass -# -# def undo(self): -# ldm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward -# ldm.modules.diffusionmodules.model.AttnBlock.forward = diffusionmodules_model_AttnBlock_forward -# -# sgm.modules.attention.CrossAttention.forward = hypernetwork.attention_CrossAttention_forward -# sgm.modules.diffusionmodules.model.AttnBlock.forward = sgm_diffusionmodules_model_AttnBlock_forward -# -# -# class SdOptimizationXformers(SdOptimization): -# name = "xformers" -# cmd_opt = "xformers" -# priority = 100 -# -# def is_available(self): -# return shared.cmd_opts.force_enable_xformers or (shared.xformers_available and torch.cuda.is_available() and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (9, 0)) -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = xformers_attention_forward -# ldm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward -# sgm.modules.attention.CrossAttention.forward = xformers_attention_forward -# sgm.modules.diffusionmodules.model.AttnBlock.forward = xformers_attnblock_forward -# -# -# class SdOptimizationSdpNoMem(SdOptimization): -# name = "sdp-no-mem" -# label = "scaled dot product without memory efficient attention" -# cmd_opt = "opt_sdp_no_mem_attention" -# priority = 80 -# -# def is_available(self): -# return hasattr(torch.nn.functional, "scaled_dot_product_attention") and callable(torch.nn.functional.scaled_dot_product_attention) -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = scaled_dot_product_no_mem_attention_forward -# ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_no_mem_attnblock_forward -# sgm.modules.attention.CrossAttention.forward = scaled_dot_product_no_mem_attention_forward -# sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_no_mem_attnblock_forward -# -# -# class SdOptimizationSdp(SdOptimizationSdpNoMem): -# name = "sdp" -# label = "scaled dot product" -# cmd_opt = "opt_sdp_attention" -# priority = 70 -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward -# ldm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward -# sgm.modules.attention.CrossAttention.forward = scaled_dot_product_attention_forward -# sgm.modules.diffusionmodules.model.AttnBlock.forward = sdp_attnblock_forward -# -# -# class SdOptimizationSubQuad(SdOptimization): -# name = "sub-quadratic" -# cmd_opt = "opt_sub_quad_attention" -# -# @property -# def priority(self): -# return 1000 if shared.device.type == 'mps' else 10 -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = sub_quad_attention_forward -# ldm.modules.diffusionmodules.model.AttnBlock.forward = sub_quad_attnblock_forward -# sgm.modules.attention.CrossAttention.forward = sub_quad_attention_forward -# sgm.modules.diffusionmodules.model.AttnBlock.forward = sub_quad_attnblock_forward -# -# -# class SdOptimizationV1(SdOptimization): -# name = "V1" -# label = "original v1" -# cmd_opt = "opt_split_attention_v1" -# priority = 10 -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1 -# sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward_v1 -# -# -# class SdOptimizationInvokeAI(SdOptimization): -# name = "InvokeAI" -# cmd_opt = "opt_split_attention_invokeai" -# -# @property -# def priority(self): -# return 1000 if shared.device.type != 'mps' and not torch.cuda.is_available() else 10 -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeAI -# sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward_invokeAI -# -# -# class SdOptimizationDoggettx(SdOptimization): -# name = "Doggettx" -# cmd_opt = "opt_split_attention" -# priority = 90 -# -# def apply(self): -# ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward -# ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward -# sgm.modules.attention.CrossAttention.forward = split_cross_attention_forward -# sgm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward -# -# -# def list_optimizers(res): -# res.extend([ -# SdOptimizationXformers(), -# SdOptimizationSdpNoMem(), -# SdOptimizationSdp(), -# SdOptimizationSubQuad(), -# SdOptimizationV1(), -# SdOptimizationInvokeAI(), -# SdOptimizationDoggettx(), -# ]) -# -# -# if shared.cmd_opts.xformers or shared.cmd_opts.force_enable_xformers: -# try: -# import xformers.ops -# shared.xformers_available = True -# except Exception: -# errors.report("Cannot import xformers", exc_info=True) -# -# -# def get_available_vram(): -# if shared.device.type == 'cuda': -# stats = torch.cuda.memory_stats(shared.device) -# mem_active = stats['active_bytes.all.current'] -# mem_reserved = stats['reserved_bytes.all.current'] -# mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device()) -# mem_free_torch = mem_reserved - mem_active -# mem_free_total = mem_free_cuda + mem_free_torch -# return mem_free_total -# else: -# return psutil.virtual_memory().available -# -# -# # see https://github.com/basujindal/stable-diffusion/pull/117 for discussion -# def split_cross_attention_forward_v1(self, x, context=None, mask=None, **kwargs): -# h = self.heads -# -# q_in = self.to_q(x) -# context = default(context, x) -# -# context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) -# k_in = self.to_k(context_k) -# v_in = self.to_v(context_v) -# del context, context_k, context_v, x -# -# q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q_in, k_in, v_in)) -# del q_in, k_in, v_in -# -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k, v = q.float(), k.float(), v.float() -# -# with devices.without_autocast(disable=not shared.opts.upcast_attn): -# r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) -# for i in range(0, q.shape[0], 2): -# end = i + 2 -# s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end]) -# s1 *= self.scale -# -# s2 = s1.softmax(dim=-1) -# del s1 -# -# r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end]) -# del s2 -# del q, k, v -# -# r1 = r1.to(dtype) -# -# r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h) -# del r1 -# -# return self.to_out(r2) -# -# -# # taken from https://github.com/Doggettx/stable-diffusion and modified -# def split_cross_attention_forward(self, x, context=None, mask=None, **kwargs): -# h = self.heads -# -# q_in = self.to_q(x) -# context = default(context, x) -# -# context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) -# k_in = self.to_k(context_k) -# v_in = self.to_v(context_v) -# -# dtype = q_in.dtype -# if shared.opts.upcast_attn: -# q_in, k_in, v_in = q_in.float(), k_in.float(), v_in if v_in.device.type == 'mps' else v_in.float() -# -# with devices.without_autocast(disable=not shared.opts.upcast_attn): -# k_in = k_in * self.scale -# -# del context, x -# -# q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q_in, k_in, v_in)) -# del q_in, k_in, v_in -# -# r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) -# -# mem_free_total = get_available_vram() -# -# gb = 1024 ** 3 -# tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() -# modifier = 3 if q.element_size() == 2 else 2.5 -# mem_required = tensor_size * modifier -# steps = 1 -# -# if mem_required > mem_free_total: -# steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2))) -# # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB " -# # f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}") -# -# if steps > 64: -# max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64 -# raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). ' -# f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free') -# -# slice_size = q.shape[1] // steps -# for i in range(0, q.shape[1], slice_size): -# end = min(i + slice_size, q.shape[1]) -# s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) -# -# s2 = s1.softmax(dim=-1, dtype=q.dtype) -# del s1 -# -# r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v) -# del s2 -# -# del q, k, v -# -# r1 = r1.to(dtype) -# -# r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h) -# del r1 -# -# return self.to_out(r2) -# -# -# # -- Taken from https://github.com/invoke-ai/InvokeAI and modified -- -# mem_total_gb = psutil.virtual_memory().total // (1 << 30) -# -# -# def einsum_op_compvis(q, k, v): -# s = einsum('b i d, b j d -> b i j', q, k) -# s = s.softmax(dim=-1, dtype=s.dtype) -# return einsum('b i j, b j d -> b i d', s, v) -# -# -# def einsum_op_slice_0(q, k, v, slice_size): -# r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) -# for i in range(0, q.shape[0], slice_size): -# end = i + slice_size -# r[i:end] = einsum_op_compvis(q[i:end], k[i:end], v[i:end]) -# return r -# -# -# def einsum_op_slice_1(q, k, v, slice_size): -# r = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) -# for i in range(0, q.shape[1], slice_size): -# end = i + slice_size -# r[:, i:end] = einsum_op_compvis(q[:, i:end], k, v) -# return r -# -# -# def einsum_op_mps_v1(q, k, v): -# if q.shape[0] * q.shape[1] <= 2**16: # (512x512) max q.shape[1]: 4096 -# return einsum_op_compvis(q, k, v) -# else: -# slice_size = math.floor(2**30 / (q.shape[0] * q.shape[1])) -# if slice_size % 4096 == 0: -# slice_size -= 1 -# return einsum_op_slice_1(q, k, v, slice_size) -# -# -# def einsum_op_mps_v2(q, k, v): -# if mem_total_gb > 8 and q.shape[0] * q.shape[1] <= 2**16: -# return einsum_op_compvis(q, k, v) -# else: -# return einsum_op_slice_0(q, k, v, 1) -# -# -# def einsum_op_tensor_mem(q, k, v, max_tensor_mb): -# size_mb = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size() // (1 << 20) -# if size_mb <= max_tensor_mb: -# return einsum_op_compvis(q, k, v) -# div = 1 << int((size_mb - 1) / max_tensor_mb).bit_length() -# if div <= q.shape[0]: -# return einsum_op_slice_0(q, k, v, q.shape[0] // div) -# return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1)) -# -# -# def einsum_op_cuda(q, k, v): -# stats = torch.cuda.memory_stats(q.device) -# mem_active = stats['active_bytes.all.current'] -# mem_reserved = stats['reserved_bytes.all.current'] -# mem_free_cuda, _ = torch.cuda.mem_get_info(q.device) -# mem_free_torch = mem_reserved - mem_active -# mem_free_total = mem_free_cuda + mem_free_torch -# # Divide factor of safety as there's copying and fragmentation -# return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) -# -# -# def einsum_op(q, k, v): -# if q.device.type == 'cuda': -# return einsum_op_cuda(q, k, v) -# -# if q.device.type == 'mps': -# if mem_total_gb >= 32 and q.shape[0] % 32 != 0 and q.shape[0] * q.shape[1] < 2**18: -# return einsum_op_mps_v1(q, k, v) -# return einsum_op_mps_v2(q, k, v) -# -# # Smaller slices are faster due to L2/L3/SLC caches. -# # Tested on i7 with 8MB L3 cache. -# return einsum_op_tensor_mem(q, k, v, 32) -# -# -# def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None, **kwargs): -# h = self.heads -# -# q = self.to_q(x) -# context = default(context, x) -# -# context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) -# k = self.to_k(context_k) -# v = self.to_v(context_v) -# del context, context_k, context_v, x -# -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k, v = q.float(), k.float(), v if v.device.type == 'mps' else v.float() -# -# with devices.without_autocast(disable=not shared.opts.upcast_attn): -# k = k * self.scale -# -# q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q, k, v)) -# r = einsum_op(q, k, v) -# r = r.to(dtype) -# return self.to_out(rearrange(r, '(b h) n d -> b n (h d)', h=h)) -# -# # -- End of code from https://github.com/invoke-ai/InvokeAI -- -# -# -# # Based on Birch-san's modified implementation of sub-quadratic attention from https://github.com/Birch-san/diffusers/pull/1 -# # The sub_quad_attention_forward function is under the MIT License listed under Memory Efficient Attention in the Licenses section of the web UI interface -# def sub_quad_attention_forward(self, x, context=None, mask=None, **kwargs): -# assert mask is None, "attention-mask not currently implemented for SubQuadraticCrossAttnProcessor." -# -# h = self.heads -# -# q = self.to_q(x) -# context = default(context, x) -# -# context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) -# k = self.to_k(context_k) -# v = self.to_v(context_v) -# del context, context_k, context_v, x -# -# q = q.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1) -# k = k.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1) -# v = v.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1) -# -# if q.device.type == 'mps': -# q, k, v = q.contiguous(), k.contiguous(), v.contiguous() -# -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k = q.float(), k.float() -# -# x = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training) -# -# x = x.to(dtype) -# -# x = x.unflatten(0, (-1, h)).transpose(1,2).flatten(start_dim=2) -# -# out_proj, dropout = self.to_out -# x = out_proj(x) -# x = dropout(x) -# -# return x -# -# -# def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_size_min=None, chunk_threshold=None, use_checkpoint=True): -# bytes_per_token = torch.finfo(q.dtype).bits//8 -# batch_x_heads, q_tokens, _ = q.shape -# _, k_tokens, _ = k.shape -# qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens -# -# if chunk_threshold is None: -# if q.device.type == 'mps': -# chunk_threshold_bytes = 268435456 * (2 if platform.processor() == 'i386' else bytes_per_token) -# else: -# chunk_threshold_bytes = int(get_available_vram() * 0.7) -# elif chunk_threshold == 0: -# chunk_threshold_bytes = None -# else: -# chunk_threshold_bytes = int(0.01 * chunk_threshold * get_available_vram()) -# -# if kv_chunk_size_min is None and chunk_threshold_bytes is not None: -# kv_chunk_size_min = chunk_threshold_bytes // (batch_x_heads * bytes_per_token * (k.shape[2] + v.shape[2])) -# elif kv_chunk_size_min == 0: -# kv_chunk_size_min = None -# -# if chunk_threshold_bytes is not None and qk_matmul_size_bytes <= chunk_threshold_bytes: -# # the big matmul fits into our memory limit; do everything in 1 chunk, -# # i.e. send it down the unchunked fast-path -# kv_chunk_size = k_tokens -# -# with devices.without_autocast(disable=q.dtype == v.dtype): -# return sub_quadratic_attention.efficient_dot_product_attention( -# q, -# k, -# v, -# query_chunk_size=q_chunk_size, -# kv_chunk_size=kv_chunk_size, -# kv_chunk_size_min = kv_chunk_size_min, -# use_checkpoint=use_checkpoint, -# ) -# -# -# def get_xformers_flash_attention_op(q, k, v): -# if not shared.cmd_opts.xformers_flash_attention: -# return None -# -# try: -# flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp -# fw, bw = flash_attention_op -# if fw.supports(xformers.ops.fmha.Inputs(query=q, key=k, value=v, attn_bias=None)): -# return flash_attention_op -# except Exception as e: -# errors.display_once(e, "enabling flash attention") -# -# return None -# -# -# def xformers_attention_forward(self, x, context=None, mask=None, **kwargs): -# h = self.heads -# q_in = self.to_q(x) -# context = default(context, x) -# -# context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) -# k_in = self.to_k(context_k) -# v_in = self.to_v(context_v) -# -# q, k, v = (t.reshape(t.shape[0], t.shape[1], h, -1) for t in (q_in, k_in, v_in)) -# -# del q_in, k_in, v_in -# -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k, v = q.float(), k.float(), v.float() -# -# out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=get_xformers_flash_attention_op(q, k, v)) -# -# out = out.to(dtype) -# -# b, n, h, d = out.shape -# out = out.reshape(b, n, h * d) -# return self.to_out(out) -# -# -# # Based on Diffusers usage of scaled dot product attention from https://github.com/huggingface/diffusers/blob/c7da8fd23359a22d0df2741688b5b4f33c26df21/src/diffusers/models/cross_attention.py -# # The scaled_dot_product_attention_forward function contains parts of code under Apache-2.0 license listed under Scaled Dot Product Attention in the Licenses section of the web UI interface -# def scaled_dot_product_attention_forward(self, x, context=None, mask=None, **kwargs): -# batch_size, sequence_length, inner_dim = x.shape -# -# if mask is not None: -# mask = self.prepare_attention_mask(mask, sequence_length, batch_size) -# mask = mask.view(batch_size, self.heads, -1, mask.shape[-1]) -# -# h = self.heads -# q_in = self.to_q(x) -# context = default(context, x) -# -# context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) -# k_in = self.to_k(context_k) -# v_in = self.to_v(context_v) -# -# head_dim = inner_dim // h -# q = q_in.view(batch_size, -1, h, head_dim).transpose(1, 2) -# k = k_in.view(batch_size, -1, h, head_dim).transpose(1, 2) -# v = v_in.view(batch_size, -1, h, head_dim).transpose(1, 2) -# -# del q_in, k_in, v_in -# -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k, v = q.float(), k.float(), v.float() -# -# # the output of sdp = (batch, num_heads, seq_len, head_dim) -# hidden_states = torch.nn.functional.scaled_dot_product_attention( -# q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False -# ) -# -# hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, h * head_dim) -# hidden_states = hidden_states.to(dtype) -# -# # linear proj -# hidden_states = self.to_out[0](hidden_states) -# # dropout -# hidden_states = self.to_out[1](hidden_states) -# return hidden_states -# -# -# def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None, **kwargs): -# with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): -# return scaled_dot_product_attention_forward(self, x, context, mask) -# -# -# def cross_attention_attnblock_forward(self, x): -# h_ = x -# h_ = self.norm(h_) -# q1 = self.q(h_) -# k1 = self.k(h_) -# v = self.v(h_) -# -# # compute attention -# b, c, h, w = q1.shape -# -# q2 = q1.reshape(b, c, h*w) -# del q1 -# -# q = q2.permute(0, 2, 1) # b,hw,c -# del q2 -# -# k = k1.reshape(b, c, h*w) # b,c,hw -# del k1 -# -# h_ = torch.zeros_like(k, device=q.device) -# -# mem_free_total = get_available_vram() -# -# tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size() -# mem_required = tensor_size * 2.5 -# steps = 1 -# -# if mem_required > mem_free_total: -# steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2))) -# -# slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1] -# for i in range(0, q.shape[1], slice_size): -# end = i + slice_size -# -# w1 = torch.bmm(q[:, i:end], k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] -# w2 = w1 * (int(c)**(-0.5)) -# del w1 -# w3 = torch.nn.functional.softmax(w2, dim=2, dtype=q.dtype) -# del w2 -# -# # attend to values -# v1 = v.reshape(b, c, h*w) -# w4 = w3.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) -# del w3 -# -# h_[:, :, i:end] = torch.bmm(v1, w4) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] -# del v1, w4 -# -# h2 = h_.reshape(b, c, h, w) -# del h_ -# -# h3 = self.proj_out(h2) -# del h2 -# -# h3 += x -# -# return h3 -# -# -# def xformers_attnblock_forward(self, x): -# try: -# h_ = x -# h_ = self.norm(h_) -# q = self.q(h_) -# k = self.k(h_) -# v = self.v(h_) -# b, c, h, w = q.shape -# q, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v)) -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k = q.float(), k.float() -# q = q.contiguous() -# k = k.contiguous() -# v = v.contiguous() -# out = xformers.ops.memory_efficient_attention(q, k, v, op=get_xformers_flash_attention_op(q, k, v)) -# out = out.to(dtype) -# out = rearrange(out, 'b (h w) c -> b c h w', h=h) -# out = self.proj_out(out) -# return x + out -# except NotImplementedError: -# return cross_attention_attnblock_forward(self, x) -# -# -# def sdp_attnblock_forward(self, x): -# h_ = x -# h_ = self.norm(h_) -# q = self.q(h_) -# k = self.k(h_) -# v = self.v(h_) -# b, c, h, w = q.shape -# q, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v)) -# dtype = q.dtype -# if shared.opts.upcast_attn: -# q, k, v = q.float(), k.float(), v.float() -# q = q.contiguous() -# k = k.contiguous() -# v = v.contiguous() -# out = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) -# out = out.to(dtype) -# out = rearrange(out, 'b (h w) c -> b c h w', h=h) -# out = self.proj_out(out) -# return x + out -# -# -# def sdp_no_mem_attnblock_forward(self, x): -# with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): -# return sdp_attnblock_forward(self, x) -# -# -# def sub_quad_attnblock_forward(self, x): -# h_ = x -# h_ = self.norm(h_) -# q = self.q(h_) -# k = self.k(h_) -# v = self.v(h_) -# b, c, h, w = q.shape -# q, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v)) -# q = q.contiguous() -# k = k.contiguous() -# v = v.contiguous() -# out = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training) -# out = rearrange(out, 'b (h w) c -> b c h w', h=h) -# out = self.proj_out(out) -# return x + out diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py deleted file mode 100644 index eb4a0af4..00000000 --- a/modules/sd_hijack_unet.py +++ /dev/null @@ -1,154 +0,0 @@ -# import torch -# from packaging import version -# from einops import repeat -# import math -# -# from modules import devices -# from modules.sd_hijack_utils import CondFunc -# -# -# class TorchHijackForUnet: -# """ -# This is torch, but with cat that resizes tensors to appropriate dimensions if they do not match; -# this makes it possible to create pictures with dimensions that are multiples of 8 rather than 64 -# """ -# -# def __getattr__(self, item): -# if item == 'cat': -# return self.cat -# -# if hasattr(torch, item): -# return getattr(torch, item) -# -# raise AttributeError(f"'{type(self).__name__}' object has no attribute '{item}'") -# -# def cat(self, tensors, *args, **kwargs): -# if len(tensors) == 2: -# a, b = tensors -# if a.shape[-2:] != b.shape[-2:]: -# a = torch.nn.functional.interpolate(a, b.shape[-2:], mode="nearest") -# -# tensors = (a, b) -# -# return torch.cat(tensors, *args, **kwargs) -# -# -# th = TorchHijackForUnet() -# -# -# # Below are monkey patches to enable upcasting a float16 UNet for float32 sampling -# def apply_model(orig_func, self, x_noisy, t, cond, **kwargs): -# """Always make sure inputs to unet are in correct dtype.""" -# if isinstance(cond, dict): -# for y in cond.keys(): -# if isinstance(cond[y], list): -# cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]] -# else: -# cond[y] = cond[y].to(devices.dtype_unet) if isinstance(cond[y], torch.Tensor) else cond[y] -# -# with devices.autocast(): -# result = orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs) -# if devices.unet_needs_upcast: -# return result.float() -# else: -# return result -# -# -# # Monkey patch to create timestep embed tensor on device, avoiding a block. -# def timestep_embedding(_, timesteps, dim, max_period=10000, repeat_only=False): -# """ -# Create sinusoidal timestep embeddings. -# :param timesteps: a 1-D Tensor of N indices, one per batch element. -# These may be fractional. -# :param dim: the dimension of the output. -# :param max_period: controls the minimum frequency of the embeddings. -# :return: an [N x dim] Tensor of positional embeddings. -# """ -# if not repeat_only: -# half = dim // 2 -# freqs = torch.exp( -# -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half -# ) -# args = timesteps[:, None].float() * freqs[None] -# embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) -# if dim % 2: -# embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) -# else: -# embedding = repeat(timesteps, 'b -> b d', d=dim) -# return embedding -# -# -# # Monkey patch to SpatialTransformer removing unnecessary contiguous calls. -# # Prevents a lot of unnecessary aten::copy_ calls -# def spatial_transformer_forward(_, self, x: torch.Tensor, context=None): -# # note: if no context is given, cross-attention defaults to self-attention -# if not isinstance(context, list): -# context = [context] -# b, c, h, w = x.shape -# x_in = x -# x = self.norm(x) -# if not self.use_linear: -# x = self.proj_in(x) -# x = x.permute(0, 2, 3, 1).reshape(b, h * w, c) -# if self.use_linear: -# x = self.proj_in(x) -# for i, block in enumerate(self.transformer_blocks): -# x = block(x, context=context[i]) -# if self.use_linear: -# x = self.proj_out(x) -# x = x.view(b, h, w, c).permute(0, 3, 1, 2) -# if not self.use_linear: -# x = self.proj_out(x) -# return x + x_in -# -# -# class GELUHijack(torch.nn.GELU, torch.nn.Module): -# def __init__(self, *args, **kwargs): -# torch.nn.GELU.__init__(self, *args, **kwargs) -# def forward(self, x): -# if devices.unet_needs_upcast: -# return torch.nn.GELU.forward(self.float(), x.float()).to(devices.dtype_unet) -# else: -# return torch.nn.GELU.forward(self, x) -# -# -# ddpm_edit_hijack = None -# def hijack_ddpm_edit(): -# global ddpm_edit_hijack -# if not ddpm_edit_hijack: -# CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond) -# CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond) -# ddpm_edit_hijack = CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.apply_model', apply_model) -# -# -# unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast -# CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast) -# CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding) -# CondFunc('ldm.modules.attention.SpatialTransformer.forward', spatial_transformer_forward) -# CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast) -# -# if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available(): -# CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast) -# CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast) -# CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU) -# -# first_stage_cond = lambda _, self, *args, **kwargs: devices.unet_needs_upcast and self.model.diffusion_model.dtype == torch.float16 -# first_stage_sub = lambda orig_func, self, x, **kwargs: orig_func(self, x.to(devices.dtype_vae), **kwargs) -# CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond) -# CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond) -# CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.get_first_stage_encoding', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).float(), first_stage_cond) -# -# CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model) -# CondFunc('sgm.modules.diffusionmodules.wrappers.OpenAIWrapper.forward', apply_model) -# -# -# def timestep_embedding_cast_result(orig_func, timesteps, *args, **kwargs): -# if devices.unet_needs_upcast and timesteps.dtype == torch.int64: -# dtype = torch.float32 -# else: -# dtype = devices.dtype_unet -# return orig_func(timesteps, *args, **kwargs).to(dtype=dtype) -# -# -# CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding_cast_result) -# CondFunc('sgm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding_cast_result) diff --git a/modules/sd_hijack_xlmr.py b/modules/sd_hijack_xlmr.py deleted file mode 100644 index 091de420..00000000 --- a/modules/sd_hijack_xlmr.py +++ /dev/null @@ -1,32 +0,0 @@ -# import torch -# -# from modules import sd_hijack_clip, devices -# -# -# class FrozenXLMREmbedderWithCustomWords(sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords): -# def __init__(self, wrapped, hijack): -# super().__init__(wrapped, hijack) -# -# self.id_start = wrapped.config.bos_token_id -# self.id_end = wrapped.config.eos_token_id -# self.id_pad = wrapped.config.pad_token_id -# -# self.comma_token = self.tokenizer.get_vocab().get(',', None) # alt diffusion doesn't have bits for comma -# -# def encode_with_transformers(self, tokens): -# # there's no CLIP Skip here because all hidden layers have size of 1024 and the last one uses a -# # trained layer to transform those 1024 into 768 for unet; so you can't choose which transformer -# # layer to work with - you have to use the last -# -# attention_mask = (tokens != self.id_pad).to(device=tokens.device, dtype=torch.int64) -# features = self.wrapped(input_ids=tokens, attention_mask=attention_mask) -# z = features['projection_state'] -# -# return z -# -# def encode_embedding_init_text(self, init_text, nvpt): -# embedding_layer = self.wrapped.roberta.embeddings -# ids = self.wrapped.tokenizer(init_text, max_length=nvpt, return_tensors="pt", add_special_tokens=False)["input_ids"] -# embedded = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0) -# -# return embedded diff --git a/modules/sd_models.py b/modules/sd_models.py index 76940ecc..012ae6af 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -14,7 +14,7 @@ from urllib import request import gc import contextlib -from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config, sd_unet, sd_models_xl, cache, extra_networks, processing, lowvram, sd_hijack, patches +from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, errors, hashes, sd_unet, cache, extra_networks, processing, lowvram, sd_hijack, patches from modules.shared import opts, cmd_opts from modules.timer import Timer import numpy as np @@ -140,15 +140,6 @@ class CheckpointInfo: return str(dict(filename=self.filename, hash=self.hash)) -# try: -# # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start. -# from transformers import logging, CLIPModel # noqa: F401 -# -# logging.set_verbosity_error() -# except Exception: -# pass - - def setup_model(): """called once at startup to do various one-time tasks related to SD models""" diff --git a/modules/sd_models_config.py b/modules/sd_models_config.py deleted file mode 100644 index c8e1f9f5..00000000 --- a/modules/sd_models_config.py +++ /dev/null @@ -1,137 +0,0 @@ -# import os -# -# import torch -# -# from modules import shared, paths, sd_disable_initialization, devices -# -# sd_configs_path = shared.sd_configs_path -# # sd_repo_configs_path = os.path.join(paths.paths['Stable Diffusion'], "configs", "stable-diffusion") -# # sd_xl_repo_configs_path = os.path.join(paths.paths['Stable Diffusion XL'], "configs", "inference") -# -# -# config_default = shared.sd_default_config -# # config_sd2 = os.path.join(sd_repo_configs_path, "v2-inference.yaml") -# config_sd2v = os.path.join(sd_repo_configs_path, "v2-inference-v.yaml") -# config_sd2_inpainting = os.path.join(sd_repo_configs_path, "v2-inpainting-inference.yaml") -# config_sdxl = os.path.join(sd_xl_repo_configs_path, "sd_xl_base.yaml") -# config_sdxl_refiner = os.path.join(sd_xl_repo_configs_path, "sd_xl_refiner.yaml") -# config_sdxl_inpainting = os.path.join(sd_configs_path, "sd_xl_inpaint.yaml") -# config_depth_model = os.path.join(sd_repo_configs_path, "v2-midas-inference.yaml") -# config_unclip = os.path.join(sd_repo_configs_path, "v2-1-stable-unclip-l-inference.yaml") -# config_unopenclip = os.path.join(sd_repo_configs_path, "v2-1-stable-unclip-h-inference.yaml") -# config_inpainting = os.path.join(sd_configs_path, "v1-inpainting-inference.yaml") -# config_instruct_pix2pix = os.path.join(sd_configs_path, "instruct-pix2pix.yaml") -# config_alt_diffusion = os.path.join(sd_configs_path, "alt-diffusion-inference.yaml") -# config_alt_diffusion_m18 = os.path.join(sd_configs_path, "alt-diffusion-m18-inference.yaml") -# config_sd3 = os.path.join(sd_configs_path, "sd3-inference.yaml") -# -# -# def is_using_v_parameterization_for_sd2(state_dict): -# """ -# Detects whether unet in state_dict is using v-parameterization. Returns True if it is. You're welcome. -# """ -# -# import ldm.modules.diffusionmodules.openaimodel -# -# device = devices.device -# -# with sd_disable_initialization.DisableInitialization(): -# unet = ldm.modules.diffusionmodules.openaimodel.UNetModel( -# use_checkpoint=False, -# use_fp16=False, -# image_size=32, -# in_channels=4, -# out_channels=4, -# model_channels=320, -# attention_resolutions=[4, 2, 1], -# num_res_blocks=2, -# channel_mult=[1, 2, 4, 4], -# num_head_channels=64, -# use_spatial_transformer=True, -# use_linear_in_transformer=True, -# transformer_depth=1, -# context_dim=1024, -# legacy=False -# ) -# unet.eval() -# -# with torch.no_grad(): -# unet_sd = {k.replace("model.diffusion_model.", ""): v for k, v in state_dict.items() if "model.diffusion_model." in k} -# unet.load_state_dict(unet_sd, strict=True) -# unet.to(device=device, dtype=devices.dtype_unet) -# -# test_cond = torch.ones((1, 2, 1024), device=device) * 0.5 -# x_test = torch.ones((1, 4, 8, 8), device=device) * 0.5 -# -# with devices.autocast(): -# out = (unet(x_test, torch.asarray([999], device=device), context=test_cond) - x_test).mean().cpu().item() -# -# return out < -1 -# -# -# def guess_model_config_from_state_dict(sd, filename): -# sd2_cond_proj_weight = sd.get('cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight', None) -# diffusion_model_input = sd.get('model.diffusion_model.input_blocks.0.0.weight', None) -# sd2_variations_weight = sd.get('embedder.model.ln_final.weight', None) -# -# if "model.diffusion_model.x_embedder.proj.weight" in sd: -# return config_sd3 -# -# if sd.get('conditioner.embedders.1.model.ln_final.weight', None) is not None: -# if diffusion_model_input.shape[1] == 9: -# return config_sdxl_inpainting -# else: -# return config_sdxl -# -# if sd.get('conditioner.embedders.0.model.ln_final.weight', None) is not None: -# return config_sdxl_refiner -# elif sd.get('depth_model.model.pretrained.act_postprocess3.0.project.0.bias', None) is not None: -# return config_depth_model -# elif sd2_variations_weight is not None and sd2_variations_weight.shape[0] == 768: -# return config_unclip -# elif sd2_variations_weight is not None and sd2_variations_weight.shape[0] == 1024: -# return config_unopenclip -# -# if sd2_cond_proj_weight is not None and sd2_cond_proj_weight.shape[1] == 1024: -# if diffusion_model_input.shape[1] == 9: -# return config_sd2_inpainting -# # elif is_using_v_parameterization_for_sd2(sd): -# # return config_sd2v -# else: -# return config_sd2v -# -# if diffusion_model_input is not None: -# if diffusion_model_input.shape[1] == 9: -# return config_inpainting -# if diffusion_model_input.shape[1] == 8: -# return config_instruct_pix2pix -# -# if sd.get('cond_stage_model.roberta.embeddings.word_embeddings.weight', None) is not None: -# if sd.get('cond_stage_model.transformation.weight').size()[0] == 1024: -# return config_alt_diffusion_m18 -# return config_alt_diffusion -# -# return config_default -# -# -# def find_checkpoint_config(state_dict, info): -# if info is None: -# return guess_model_config_from_state_dict(state_dict, "") -# -# config = find_checkpoint_config_near_filename(info) -# if config is not None: -# return config -# -# return guess_model_config_from_state_dict(state_dict, info.filename) -# -# -# def find_checkpoint_config_near_filename(info): -# if info is None: -# return None -# -# config = f"{os.path.splitext(info.filename)[0]}.yaml" -# if os.path.exists(config): -# return config -# -# return None -# diff --git a/modules/sd_models_types.py b/modules/sd_models_types.py deleted file mode 100644 index ba271aa1..00000000 --- a/modules/sd_models_types.py +++ /dev/null @@ -1,39 +0,0 @@ -# from typing import TYPE_CHECKING -# -# -# if TYPE_CHECKING: -# from modules.sd_models import CheckpointInfo -# -# -# class WebuiSdModel: -# """This class is not actually instantinated, but its fields are created and fieeld by webui""" -# -# lowvram: bool -# """True if lowvram/medvram optimizations are enabled -- see modules.lowvram for more info""" -# -# sd_model_hash: str -# """short hash, 10 first characters of SHA1 hash of the model file; may be None if --no-hashing flag is used""" -# -# sd_model_checkpoint: str -# """path to the file on disk that model weights were obtained from""" -# -# sd_checkpoint_info: 'CheckpointInfo' -# """structure with additional information about the file with model's weights""" -# -# is_sdxl: bool -# """True if the model's architecture is SDXL or SSD""" -# -# is_ssd: bool -# """True if the model is SSD""" -# -# is_sd2: bool -# """True if the model's architecture is SD 2.x""" -# -# is_sd1: bool -# """True if the model's architecture is SD 1.x""" -# -# is_sd3: bool -# """True if the model's architecture is SD 3""" -# -# latent_channels: int -# """number of layer in latent image representation; will be 16 in SD3 and 4 in other version""" diff --git a/modules/sd_models_xl.py b/modules/sd_models_xl.py deleted file mode 100644 index 0b84f2fc..00000000 --- a/modules/sd_models_xl.py +++ /dev/null @@ -1,115 +0,0 @@ -# from __future__ import annotations -# -# import torch -# -# import sgm.models.diffusion -# import sgm.modules.diffusionmodules.denoiser_scaling -# import sgm.modules.diffusionmodules.discretizer -# from modules import devices, shared, prompt_parser -# from modules import torch_utils -# -# from backend import memory_management -# -# -# def get_learned_conditioning(self: sgm.models.diffusion.DiffusionEngine, batch: prompt_parser.SdConditioning | list[str]): -# -# for embedder in self.conditioner.embedders: -# embedder.ucg_rate = 0.0 -# -# width = getattr(batch, 'width', 1024) or 1024 -# height = getattr(batch, 'height', 1024) or 1024 -# is_negative_prompt = getattr(batch, 'is_negative_prompt', False) -# aesthetic_score = shared.opts.sdxl_refiner_low_aesthetic_score if is_negative_prompt else shared.opts.sdxl_refiner_high_aesthetic_score -# -# devices_args = dict(device=self.forge_objects.clip.patcher.current_device, dtype=memory_management.text_encoder_dtype()) -# -# sdxl_conds = { -# "txt": batch, -# "original_size_as_tuple": torch.tensor([height, width], **devices_args).repeat(len(batch), 1), -# "crop_coords_top_left": torch.tensor([shared.opts.sdxl_crop_top, shared.opts.sdxl_crop_left], **devices_args).repeat(len(batch), 1), -# "target_size_as_tuple": torch.tensor([height, width], **devices_args).repeat(len(batch), 1), -# "aesthetic_score": torch.tensor([aesthetic_score], **devices_args).repeat(len(batch), 1), -# } -# -# force_zero_negative_prompt = is_negative_prompt and all(x == '' for x in batch) -# c = self.conditioner(sdxl_conds, force_zero_embeddings=['txt'] if force_zero_negative_prompt else []) -# -# return c -# -# -# def apply_model(self: sgm.models.diffusion.DiffusionEngine, x, t, cond, *args, **kwargs): -# if self.model.diffusion_model.in_channels == 9: -# x = torch.cat([x] + cond['c_concat'], dim=1) -# -# return self.model(x, t, cond, *args, **kwargs) -# -# -# def get_first_stage_encoding(self, x): # SDXL's encode_first_stage does everything so get_first_stage_encoding is just there for compatibility -# return x -# -# -# sgm.models.diffusion.DiffusionEngine.get_learned_conditioning = get_learned_conditioning -# sgm.models.diffusion.DiffusionEngine.apply_model = apply_model -# sgm.models.diffusion.DiffusionEngine.get_first_stage_encoding = get_first_stage_encoding -# -# -# def encode_embedding_init_text(self: sgm.modules.GeneralConditioner, init_text, nvpt): -# res = [] -# -# for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'encode_embedding_init_text')]: -# encoded = embedder.encode_embedding_init_text(init_text, nvpt) -# res.append(encoded) -# -# return torch.cat(res, dim=1) -# -# -# def tokenize(self: sgm.modules.GeneralConditioner, texts): -# for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'tokenize')]: -# return embedder.tokenize(texts) -# -# raise AssertionError('no tokenizer available') -# -# -# -# def process_texts(self, texts): -# for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'process_texts')]: -# return embedder.process_texts(texts) -# -# -# def get_target_prompt_token_count(self, token_count): -# for embedder in [embedder for embedder in self.embedders if hasattr(embedder, 'get_target_prompt_token_count')]: -# return embedder.get_target_prompt_token_count(token_count) -# -# -# # those additions to GeneralConditioner make it possible to use it as model.cond_stage_model from SD1.5 in exist -# sgm.modules.GeneralConditioner.encode_embedding_init_text = encode_embedding_init_text -# sgm.modules.GeneralConditioner.tokenize = tokenize -# sgm.modules.GeneralConditioner.process_texts = process_texts -# sgm.modules.GeneralConditioner.get_target_prompt_token_count = get_target_prompt_token_count -# -# -# def extend_sdxl(model): -# """this adds a bunch of parameters to make SDXL model look a bit more like SD1.5 to the rest of the codebase.""" -# -# dtype = torch_utils.get_param(model.model.diffusion_model).dtype -# model.model.diffusion_model.dtype = dtype -# model.model.conditioning_key = 'crossattn' -# model.cond_stage_key = 'txt' -# # model.cond_stage_model will be set in sd_hijack -# -# model.parameterization = "v" if isinstance(model.denoiser.scaling, sgm.modules.diffusionmodules.denoiser_scaling.VScaling) else "eps" -# -# discretization = sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization() -# model.alphas_cumprod = torch.asarray(discretization.alphas_cumprod, device=devices.device, dtype=torch.float32) -# -# model.conditioner.wrapped = torch.nn.Module() -# -# -# sgm.modules.attention.print = shared.ldm_print -# sgm.modules.diffusionmodules.model.print = shared.ldm_print -# sgm.modules.diffusionmodules.openaimodel.print = shared.ldm_print -# sgm.modules.encoders.modules.print = shared.ldm_print -# -# # this gets the code to load the vanilla attention that we override -# sgm.modules.attention.SDP_IS_AVAILABLE = True -# sgm.modules.attention.XFORMERS_IS_AVAILABLE = False diff --git a/modules/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py deleted file mode 100644 index e69de29b..00000000 diff --git a/modules/shared.py b/modules/shared.py index ab6e412f..01cdd82e 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -3,7 +3,7 @@ import sys import gradio as gr -from modules import shared_cmd_options, shared_gradio_themes, options, shared_items, sd_models_types +from modules import shared_cmd_options, shared_gradio_themes, options, shared_items from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # noqa: F401 from modules import util from typing import TYPE_CHECKING diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index c6f6ea5d..6cd38be3 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -12,7 +12,7 @@ import safetensors.torch import numpy as np from PIL import Image, PngImagePlugin -from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors, hashes +from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, errors, hashes from modules.textual_inversion.image_embedding import embedding_to_b64, embedding_from_b64, insert_image_data_embed, extract_image_data_embed, caption_image_overlay diff --git a/modules/xlmr.py b/modules/xlmr.py deleted file mode 100644 index 029821be..00000000 --- a/modules/xlmr.py +++ /dev/null @@ -1,140 +0,0 @@ -# from transformers import BertPreTrainedModel, BertConfig -# import torch.nn as nn -# import torch -# from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig -# from transformers import XLMRobertaModel,XLMRobertaTokenizer -# from typing import Optional -# -# from modules import torch_utils -# -# -# class BertSeriesConfig(BertConfig): -# def __init__(self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=0, position_embedding_type="absolute", use_cache=True, classifier_dropout=None,project_dim=512, pooler_fn="average",learn_encoder=False,model_type='bert',**kwargs): -# -# super().__init__(vocab_size, hidden_size, num_hidden_layers, num_attention_heads, intermediate_size, hidden_act, hidden_dropout_prob, attention_probs_dropout_prob, max_position_embeddings, type_vocab_size, initializer_range, layer_norm_eps, pad_token_id, position_embedding_type, use_cache, classifier_dropout, **kwargs) -# self.project_dim = project_dim -# self.pooler_fn = pooler_fn -# self.learn_encoder = learn_encoder -# -# class RobertaSeriesConfig(XLMRobertaConfig): -# def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2,project_dim=512,pooler_fn='cls',learn_encoder=False, **kwargs): -# super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) -# self.project_dim = project_dim -# self.pooler_fn = pooler_fn -# self.learn_encoder = learn_encoder -# -# -# class BertSeriesModelWithTransformation(BertPreTrainedModel): -# -# _keys_to_ignore_on_load_unexpected = [r"pooler"] -# _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"] -# config_class = BertSeriesConfig -# -# def __init__(self, config=None, **kargs): -# # modify initialization for autoloading -# if config is None: -# config = XLMRobertaConfig() -# config.attention_probs_dropout_prob= 0.1 -# config.bos_token_id=0 -# config.eos_token_id=2 -# config.hidden_act='gelu' -# config.hidden_dropout_prob=0.1 -# config.hidden_size=1024 -# config.initializer_range=0.02 -# config.intermediate_size=4096 -# config.layer_norm_eps=1e-05 -# config.max_position_embeddings=514 -# -# config.num_attention_heads=16 -# config.num_hidden_layers=24 -# config.output_past=True -# config.pad_token_id=1 -# config.position_embedding_type= "absolute" -# -# config.type_vocab_size= 1 -# config.use_cache=True -# config.vocab_size= 250002 -# config.project_dim = 768 -# config.learn_encoder = False -# super().__init__(config) -# self.roberta = XLMRobertaModel(config) -# self.transformation = nn.Linear(config.hidden_size,config.project_dim) -# self.pre_LN=nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) -# self.tokenizer = XLMRobertaTokenizer.from_pretrained('xlm-roberta-large') -# self.pooler = lambda x: x[:,0] -# self.post_init() -# -# def encode(self,c): -# device = torch_utils.get_param(self).device -# text = self.tokenizer(c, -# truncation=True, -# max_length=77, -# return_length=False, -# return_overflowing_tokens=False, -# padding="max_length", -# return_tensors="pt") -# text["input_ids"] = torch.tensor(text["input_ids"]).to(device) -# text["attention_mask"] = torch.tensor( -# text['attention_mask']).to(device) -# features = self(**text) -# return features['projection_state'] -# -# def forward( -# self, -# input_ids: Optional[torch.Tensor] = None, -# attention_mask: Optional[torch.Tensor] = None, -# token_type_ids: Optional[torch.Tensor] = None, -# position_ids: Optional[torch.Tensor] = None, -# head_mask: Optional[torch.Tensor] = None, -# inputs_embeds: Optional[torch.Tensor] = None, -# encoder_hidden_states: Optional[torch.Tensor] = None, -# encoder_attention_mask: Optional[torch.Tensor] = None, -# output_attentions: Optional[bool] = None, -# return_dict: Optional[bool] = None, -# output_hidden_states: Optional[bool] = None, -# ) : -# r""" -# """ -# -# return_dict = return_dict if return_dict is not None else self.config.use_return_dict -# -# -# outputs = self.roberta( -# input_ids=input_ids, -# attention_mask=attention_mask, -# token_type_ids=token_type_ids, -# position_ids=position_ids, -# head_mask=head_mask, -# inputs_embeds=inputs_embeds, -# encoder_hidden_states=encoder_hidden_states, -# encoder_attention_mask=encoder_attention_mask, -# output_attentions=output_attentions, -# output_hidden_states=True, -# return_dict=return_dict, -# ) -# -# # last module outputs -# sequence_output = outputs[0] -# -# -# # project every module -# sequence_output_ln = self.pre_LN(sequence_output) -# -# # pooler -# pooler_output = self.pooler(sequence_output_ln) -# pooler_output = self.transformation(pooler_output) -# projection_state = self.transformation(outputs.last_hidden_state) -# -# return { -# 'pooler_output':pooler_output, -# 'last_hidden_state':outputs.last_hidden_state, -# 'hidden_states':outputs.hidden_states, -# 'attentions':outputs.attentions, -# 'projection_state':projection_state, -# 'sequence_out': sequence_output -# } -# -# -# class RobertaSeriesModelWithTransformation(BertSeriesModelWithTransformation): -# base_model_prefix = 'roberta' -# config_class= RobertaSeriesConfig diff --git a/modules/xlmr_m18.py b/modules/xlmr_m18.py deleted file mode 100644 index f92afd16..00000000 --- a/modules/xlmr_m18.py +++ /dev/null @@ -1,166 +0,0 @@ -# from transformers import BertPreTrainedModel,BertConfig -# import torch.nn as nn -# import torch -# from transformers.models.xlm_roberta.configuration_xlm_roberta import XLMRobertaConfig -# from transformers import XLMRobertaModel,XLMRobertaTokenizer -# from typing import Optional -# from modules import torch_utils -# -# -# class BertSeriesConfig(BertConfig): -# def __init__(self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act="gelu", hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1, max_position_embeddings=512, type_vocab_size=2, initializer_range=0.02, layer_norm_eps=1e-12, pad_token_id=0, position_embedding_type="absolute", use_cache=True, classifier_dropout=None,project_dim=512, pooler_fn="average",learn_encoder=False,model_type='bert',**kwargs): -# -# super().__init__(vocab_size, hidden_size, num_hidden_layers, num_attention_heads, intermediate_size, hidden_act, hidden_dropout_prob, attention_probs_dropout_prob, max_position_embeddings, type_vocab_size, initializer_range, layer_norm_eps, pad_token_id, position_embedding_type, use_cache, classifier_dropout, **kwargs) -# self.project_dim = project_dim -# self.pooler_fn = pooler_fn -# self.learn_encoder = learn_encoder -# -# class RobertaSeriesConfig(XLMRobertaConfig): -# def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2,project_dim=512,pooler_fn='cls',learn_encoder=False, **kwargs): -# super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) -# self.project_dim = project_dim -# self.pooler_fn = pooler_fn -# self.learn_encoder = learn_encoder -# -# -# class BertSeriesModelWithTransformation(BertPreTrainedModel): -# -# _keys_to_ignore_on_load_unexpected = [r"pooler"] -# _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"] -# config_class = BertSeriesConfig -# -# def __init__(self, config=None, **kargs): -# # modify initialization for autoloading -# if config is None: -# config = XLMRobertaConfig() -# config.attention_probs_dropout_prob= 0.1 -# config.bos_token_id=0 -# config.eos_token_id=2 -# config.hidden_act='gelu' -# config.hidden_dropout_prob=0.1 -# config.hidden_size=1024 -# config.initializer_range=0.02 -# config.intermediate_size=4096 -# config.layer_norm_eps=1e-05 -# config.max_position_embeddings=514 -# -# config.num_attention_heads=16 -# config.num_hidden_layers=24 -# config.output_past=True -# config.pad_token_id=1 -# config.position_embedding_type= "absolute" -# -# config.type_vocab_size= 1 -# config.use_cache=True -# config.vocab_size= 250002 -# config.project_dim = 1024 -# config.learn_encoder = False -# super().__init__(config) -# self.roberta = XLMRobertaModel(config) -# self.transformation = nn.Linear(config.hidden_size,config.project_dim) -# # self.pre_LN=nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) -# self.tokenizer = XLMRobertaTokenizer.from_pretrained('xlm-roberta-large') -# # self.pooler = lambda x: x[:,0] -# # self.post_init() -# -# self.has_pre_transformation = True -# if self.has_pre_transformation: -# self.transformation_pre = nn.Linear(config.hidden_size, config.project_dim) -# self.pre_LN = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) -# self.post_init() -# -# def encode(self,c): -# device = torch_utils.get_param(self).device -# text = self.tokenizer(c, -# truncation=True, -# max_length=77, -# return_length=False, -# return_overflowing_tokens=False, -# padding="max_length", -# return_tensors="pt") -# text["input_ids"] = torch.tensor(text["input_ids"]).to(device) -# text["attention_mask"] = torch.tensor( -# text['attention_mask']).to(device) -# features = self(**text) -# return features['projection_state'] -# -# def forward( -# self, -# input_ids: Optional[torch.Tensor] = None, -# attention_mask: Optional[torch.Tensor] = None, -# token_type_ids: Optional[torch.Tensor] = None, -# position_ids: Optional[torch.Tensor] = None, -# head_mask: Optional[torch.Tensor] = None, -# inputs_embeds: Optional[torch.Tensor] = None, -# encoder_hidden_states: Optional[torch.Tensor] = None, -# encoder_attention_mask: Optional[torch.Tensor] = None, -# output_attentions: Optional[bool] = None, -# return_dict: Optional[bool] = None, -# output_hidden_states: Optional[bool] = None, -# ) : -# r""" -# """ -# -# return_dict = return_dict if return_dict is not None else self.config.use_return_dict -# -# -# outputs = self.roberta( -# input_ids=input_ids, -# attention_mask=attention_mask, -# token_type_ids=token_type_ids, -# position_ids=position_ids, -# head_mask=head_mask, -# inputs_embeds=inputs_embeds, -# encoder_hidden_states=encoder_hidden_states, -# encoder_attention_mask=encoder_attention_mask, -# output_attentions=output_attentions, -# output_hidden_states=True, -# return_dict=return_dict, -# ) -# -# # # last module outputs -# # sequence_output = outputs[0] -# -# -# # # project every module -# # sequence_output_ln = self.pre_LN(sequence_output) -# -# # # pooler -# # pooler_output = self.pooler(sequence_output_ln) -# # pooler_output = self.transformation(pooler_output) -# # projection_state = self.transformation(outputs.last_hidden_state) -# -# if self.has_pre_transformation: -# sequence_output2 = outputs["hidden_states"][-2] -# sequence_output2 = self.pre_LN(sequence_output2) -# projection_state2 = self.transformation_pre(sequence_output2) -# -# return { -# "projection_state": projection_state2, -# "last_hidden_state": outputs.last_hidden_state, -# "hidden_states": outputs.hidden_states, -# "attentions": outputs.attentions, -# } -# else: -# projection_state = self.transformation(outputs.last_hidden_state) -# return { -# "projection_state": projection_state, -# "last_hidden_state": outputs.last_hidden_state, -# "hidden_states": outputs.hidden_states, -# "attentions": outputs.attentions, -# } -# -# -# # return { -# # 'pooler_output':pooler_output, -# # 'last_hidden_state':outputs.last_hidden_state, -# # 'hidden_states':outputs.hidden_states, -# # 'attentions':outputs.attentions, -# # 'projection_state':projection_state, -# # 'sequence_out': sequence_output -# # } -# -# -# class RobertaSeriesModelWithTransformation(BertSeriesModelWithTransformation): -# base_model_prefix = 'roberta' -# config_class= RobertaSeriesConfig diff --git a/modules/xpu_specific.py b/modules/xpu_specific.py deleted file mode 100644 index c5b445f8..00000000 --- a/modules/xpu_specific.py +++ /dev/null @@ -1,138 +0,0 @@ -# from modules import shared -# from modules.sd_hijack_utils import CondFunc -# -# has_ipex = False -# try: -# import torch -# import intel_extension_for_pytorch as ipex # noqa: F401 -# has_ipex = True -# except Exception: -# pass -# -# -# def check_for_xpu(): -# return has_ipex and hasattr(torch, 'xpu') and torch.xpu.is_available() -# -# -# def get_xpu_device_string(): -# if shared.cmd_opts.device_id is not None: -# return f"xpu:{shared.cmd_opts.device_id}" -# return "xpu" -# -# -# def torch_xpu_gc(): -# with torch.xpu.device(get_xpu_device_string()): -# torch.xpu.empty_cache() -# -# -# has_xpu = check_for_xpu() -# -# -# # Arc GPU cannot allocate a single block larger than 4GB: https://github.com/intel/compute-runtime/issues/627 -# # Here we implement a slicing algorithm to split large batch size into smaller chunks, -# # so that SDPA of each chunk wouldn't require any allocation larger than ARC_SINGLE_ALLOCATION_LIMIT. -# # The heuristic limit (TOTAL_VRAM // 8) is tuned for Intel Arc A770 16G and Arc A750 8G, -# # which is the best trade-off between VRAM usage and performance. -# ARC_SINGLE_ALLOCATION_LIMIT = {} -# orig_sdp_attn_func = torch.nn.functional.scaled_dot_product_attention -# def torch_xpu_scaled_dot_product_attention( -# query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *args, **kwargs -# ): -# # cast to same dtype first -# key = key.to(query.dtype) -# value = value.to(query.dtype) -# if attn_mask is not None and attn_mask.dtype != torch.bool: -# attn_mask = attn_mask.to(query.dtype) -# -# N = query.shape[:-2] # Batch size -# L = query.size(-2) # Target sequence length -# E = query.size(-1) # Embedding dimension of the query and key -# S = key.size(-2) # Source sequence length -# Ev = value.size(-1) # Embedding dimension of the value -# -# total_batch_size = torch.numel(torch.empty(N)) -# device_id = query.device.index -# if device_id not in ARC_SINGLE_ALLOCATION_LIMIT: -# ARC_SINGLE_ALLOCATION_LIMIT[device_id] = min(torch.xpu.get_device_properties(device_id).total_memory // 8, 4 * 1024 * 1024 * 1024) -# batch_size_limit = max(1, ARC_SINGLE_ALLOCATION_LIMIT[device_id] // (L * S * query.element_size())) -# -# if total_batch_size <= batch_size_limit: -# return orig_sdp_attn_func( -# query, -# key, -# value, -# attn_mask, -# dropout_p, -# is_causal, -# *args, **kwargs -# ) -# -# query = torch.reshape(query, (-1, L, E)) -# key = torch.reshape(key, (-1, S, E)) -# value = torch.reshape(value, (-1, S, Ev)) -# if attn_mask is not None: -# attn_mask = attn_mask.view(-1, L, S) -# chunk_count = (total_batch_size + batch_size_limit - 1) // batch_size_limit -# outputs = [] -# for i in range(chunk_count): -# attn_mask_chunk = ( -# None -# if attn_mask is None -# else attn_mask[i * batch_size_limit : (i + 1) * batch_size_limit, :, :] -# ) -# chunk_output = orig_sdp_attn_func( -# query[i * batch_size_limit : (i + 1) * batch_size_limit, :, :], -# key[i * batch_size_limit : (i + 1) * batch_size_limit, :, :], -# value[i * batch_size_limit : (i + 1) * batch_size_limit, :, :], -# attn_mask_chunk, -# dropout_p, -# is_causal, -# *args, **kwargs -# ) -# outputs.append(chunk_output) -# result = torch.cat(outputs, dim=0) -# return torch.reshape(result, (*N, L, Ev)) -# -# -# def is_xpu_device(device: str | torch.device = None): -# if device is None: -# return False -# if isinstance(device, str): -# return device.startswith("xpu") -# return device.type == "xpu" -# -# -# if has_xpu: -# try: -# # torch.Generator supports "xpu" device since 2.1 -# torch.Generator("xpu") -# except RuntimeError: -# # W/A for https://github.com/intel/intel-extension-for-pytorch/issues/452: torch.Generator API doesn't support XPU device (for torch < 2.1) -# CondFunc('torch.Generator', -# lambda orig_func, device=None: torch.xpu.Generator(device), -# lambda orig_func, device=None: is_xpu_device(device)) -# -# # W/A for some OPs that could not handle different input dtypes -# CondFunc('torch.nn.functional.layer_norm', -# lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs: -# orig_func(input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs), -# lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs: -# weight is not None and input.dtype != weight.data.dtype) -# CondFunc('torch.nn.modules.GroupNorm.forward', -# lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)), -# lambda orig_func, self, input: input.dtype != self.weight.data.dtype) -# CondFunc('torch.nn.modules.linear.Linear.forward', -# lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)), -# lambda orig_func, self, input: input.dtype != self.weight.data.dtype) -# CondFunc('torch.nn.modules.conv.Conv2d.forward', -# lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)), -# lambda orig_func, self, input: input.dtype != self.weight.data.dtype) -# CondFunc('torch.bmm', -# lambda orig_func, input, mat2, out=None: orig_func(input.to(mat2.dtype), mat2, out=out), -# lambda orig_func, input, mat2, out=None: input.dtype != mat2.dtype) -# CondFunc('torch.cat', -# lambda orig_func, tensors, dim=0, out=None: orig_func([t.to(tensors[0].dtype) for t in tensors], dim=dim, out=out), -# lambda orig_func, tensors, dim=0, out=None: not all(t.dtype == tensors[0].dtype for t in tensors)) -# CondFunc('torch.nn.functional.scaled_dot_product_attention', -# lambda orig_func, *args, **kwargs: torch_xpu_scaled_dot_product_attention(*args, **kwargs), -# lambda orig_func, query, *args, **kwargs: query.is_xpu)