diff --git a/CHANGELOG.md b/CHANGELOG.md index 5590b2ee6..c31c4daa6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,14 +1,14 @@ # Change Log for SD.Next -## Future +## TODO Future - ipadapter multi image - control second pass - diffusers public callbacks - image2video: pia and vgen pipelines +- video2video - wuerstchen v3 [pr](https://github.com/huggingface/diffusers/pull/6487) - more pipelines: -- segmoe: - control api - masking api - preprocess api @@ -21,12 +21,12 @@ - update docs - diffusers 0.26.2 -## Update for 2023-02-04 +## TODO Release notes Another big release, highlights being: - A lot more functionality in the **Control** module: - Inpaint and outpaint support, flexible resizing options, optional hires - - Built-in support for many new processors and models which are auto-downloaded on first use + - Built-in support for many new processors and models, all auto-downloaded on first use - Full support for scripts and extensions - Complete **Face** module implements all variations of **FaceID**, **FaceSwap** and latest **PhotoMaker** and **InstantID** @@ -34,14 +34,21 @@ Another big release, highlights being: - Brand new **Intelligent masking**, manual or automatic Using ML models (*LAMA* object removal, *REMBG* background removal, *SAM* segmentation, etc.) and with live previews With granular blur, erode and dilate controls +- New models and pipelines: + **Segmind SegMoE**, **Mixture Tiling**, **InstaFlow**, **SAG**, **BlipDiffusion** - Massive work integrating latest advances with [OpenVINO](https://github.com/vladmandic/automatic/wiki/OpenVINO), [IPEX](https://github.com/vladmandic/automatic/wiki/Intel-ARC) and [ONNX Olive](https://github.com/vladmandic/automatic/wiki/ONNX-Runtime-&-Olive) -- **New models** and pipelines: *Mixture Tiling*, *SAG*, *InstaFlow*, *BlipDiffusion* - Full control over brightness, sharpness and color during generate process directly in latent space Plus welcome additions to **UI performance, usability and accessibility** and flexibility of deployment And it also includes fixes for all reported issues so far -As of this release, default backend is set to **diffusers** as its more feature rich than **original** and supports many additional models +As of this release, default backend is set to **diffusers** as its more feature rich than **original** and supports many additional models (original backend does remain as fully supported) + +- For basic instructions, see [README](https://github.com/vladmandic/automatic/blob/master/README.md) +- For more details on all new features see full [CHANGELOG](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) +- For documentation, see [WIKI](https://github.com/vladmandic/automatic/wiki) + +## Update for 2023-02-04 - **Control**: - add **inpaint** support @@ -129,6 +136,13 @@ As of this release, default backend is set to **diffusers** as its more feature **SD15**: Base, Base ViT-G, Light, Plus, Plus Face, Full Face **SDXL**: Base SXDL, Base ViT-H SXDL, Plus ViT-H SXDL, Plus Face ViT-H SXDL - enable use via api, thanks @trojaner +- [Segmind SegMoE](https://github.com/segmind/segmoe) + - initial support for reference models + download&load via network -> models -> reference -> **SegMoE SD 4x2** (3.7GB), **SegMoE XL 2x1** (10GB), **SegMoE XL 4x2** + - note: since segmoe is basically sequential mix of unets from multiple models, it can get large + SD 4x2 is ~4GB, XL 2x1 is ~10GB and XL 4x2 is 18GB + - support for create and load custom mixes will be added in the future + - support for lora and other advanced features will be added in the future - [Mixture Tiling](https://arxiv.org/abs/2302.02412) - uses multiple prompts to guide different parts of the grid during diffusion process - can be used ot create complex scenes with multiple subjects diff --git a/extensions-builtin/Lora/network_lora.py b/extensions-builtin/Lora/network_lora.py index 8f31388d4..1405a9d13 100644 --- a/extensions-builtin/Lora/network_lora.py +++ b/extensions-builtin/Lora/network_lora.py @@ -42,7 +42,7 @@ class NetworkModuleLora(network.NetworkModule): elif is_conv and key == "lora_up.weight" or key == "dyn_down": module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False) else: - raise AssertionError(f'Lora layer {self.network_key} matched a layer with unsupported type: {type(self.sd_module).__name__}') + raise AssertionError(f'Lora unsupported: layer={self.network_key} type={type(self.sd_module).__name__}') with torch.no_grad(): if weight.shape != module.weight.shape: weight = weight.reshape(module.weight.shape) diff --git a/html/reference.json b/html/reference.json index eadcec5be..0d48447fd 100644 --- a/html/reference.json +++ b/html/reference.json @@ -88,6 +88,21 @@ "desc": "Segmind's Tiny-SD offers a compact, efficient, and distilled version of Realistic Vision 4.0 and is up to 80% faster than SD1.5", "preview": "segmind--tiny-sd.jpg" }, + "Segmind SegMoE SD 4x2": { + "path": "segmind/SegMoE-SD-4x2-v0", + "desc": "SegMoE-SD-4x2-v0 is an untrained Segmind Mixture of Diffusion Experts Model generated using segmoe from 4 Expert SD1.5 models. SegMoE is a powerful framework for dynamically combining Stable Diffusion Models into a Mixture of Experts within minutes without training", + "preview": "segmind--SegMoE-SD-4x2-v0.jpg" + }, + "Segmind SegMoE XL 2x1": { + "path": "segmind/SegMoE-2x1-v0", + "desc": "SegMoE-2x1-v0 is an untrained Segmind Mixture of Diffusion Experts Model generated using segmoe from 2 Expert SDXL models. SegMoE is a powerful framework for dynamically combining Stable Diffusion Models into a Mixture of Experts within minutes without training", + "preview": "segmind--SegMoE-2x1-v0.jpg" + }, + "Segmind SegMoE XL 4x2": { + "path": "segmind/SegMoE-4x2-v0", + "desc": "SegMoE-4x2-v0 is an untrained Segmind Mixture of Diffusion Experts Model generated using segmoe from 4 Expert SDXL models. SegMoE is a powerful framework for dynamically combining Stable Diffusion Models into a Mixture of Experts within minutes without training", + "preview": "segmind--SegMoE-4x2-v0.jpg" + }, "LCM SD-1.5 Dreamshaper 7": { "path": "SimianLuo/LCM_Dreamshaper_v7", "desc": "Latent Consistencey Models enable swift inference with minimal steps on any pre-trained LDMs, including Stable Diffusion. By distilling classifier-free guidance into the model's input, LCM can generate high-quality images in very short inference time. LCM can generate quality images in as few as 3-4 steps, making it blazingly fast.", diff --git a/models/Reference/segmind--SegMoE-2x1-v0.jpg b/models/Reference/segmind--SegMoE-2x1-v0.jpg new file mode 100644 index 000000000..7d7989ec2 Binary files /dev/null and b/models/Reference/segmind--SegMoE-2x1-v0.jpg differ diff --git a/models/Reference/segmind--SegMoE-4x2-v0.jpg b/models/Reference/segmind--SegMoE-4x2-v0.jpg new file mode 100644 index 000000000..7d7989ec2 Binary files /dev/null and b/models/Reference/segmind--SegMoE-4x2-v0.jpg differ diff --git a/models/Reference/segmind--SegMoE-SD-4x2-v0.jpg b/models/Reference/segmind--SegMoE-SD-4x2-v0.jpg new file mode 100644 index 000000000..7d7989ec2 Binary files /dev/null and b/models/Reference/segmind--SegMoE-SD-4x2-v0.jpg differ diff --git a/modules/control/proc/marigold/util/ensemble.py b/modules/control/proc/marigold/util/ensemble.py index ae254329d..d7a9023ca 100644 --- a/modules/control/proc/marigold/util/ensemble.py +++ b/modules/control/proc/marigold/util/ensemble.py @@ -74,7 +74,7 @@ def ensemble_depths( # objective function def closure(x): - l = len(x) # noqa + l = len(x) s = x[: int(l / 2)] t = x[int(l / 2) :] s = torch.from_numpy(s).to(dtype=dtype).to(device) @@ -102,7 +102,7 @@ def ensemble_depths( closure, x, method="BFGS", tol=tol, options={"maxiter": max_iter, "disp": False} ) x = res.x - l = len(x) # noqa + l = len(x) s = x[: int(l / 2)] t = x[int(l / 2) :] diff --git a/modules/face/instantid_model.py b/modules/face/instantid_model.py index 5ae70b1fa..7f70aff98 100644 --- a/modules/face/instantid_model.py +++ b/modules/face/instantid_model.py @@ -94,7 +94,7 @@ class PerceiverAttention(nn.Module): x = self.norm1(x) latents = self.norm2(latents) - b, l, _ = latents.shape # noqa:E741 + b, l, _ = latents.shape q = self.to_q(latents) kv_input = torch.cat((x, latents), dim=-2) diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index b7ab23a2a..59fb9be9a 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -158,7 +158,9 @@ def process_diffusers(p: processing.StableDiffusionProcessing): if hasattr(model, "set_progress_bar_config"): model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba') args = {} - signature = inspect.signature(type(model).__call__) + if hasattr(model, 'pipe'): # recurse + model = model.pipe + signature = inspect.signature(type(model).__call__, follow_wrapped=True) possible = signature.parameters.keys() debug(f'Diffusers pipeline possible: {possible}') if shared.opts.diffusers_generator_device == "Unset": @@ -201,14 +203,15 @@ def process_diffusers(p: processing.StableDiffusionProcessing): if 'generator' in possible and generator is not None: args['generator'] = generator if 'output_type' in possible: - args['output_type'] = 'np' + if hasattr(model, 'vae'): + args['output_type'] = 'np' # only set latent if model has vae if 'callback_steps' in possible: args['callback_steps'] = 1 if 'callback' in possible: args['callback'] = diffusers_callback_legacy elif 'callback_on_step_end_tensor_inputs' in possible: args['callback_on_step_end'] = diffusers_callback - if 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible: + if 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible and hasattr(model, '_callback_tensor_inputs'): args['callback_on_step_end_tensor_inputs'] = model._callback_tensor_inputs # pylint: disable=protected-access else: args['callback_on_step_end_tensor_inputs'] = ['latents'] @@ -405,7 +408,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): desc='Base', ) update_sampler(shared.sd_model) - shared.state.sampling_steps = base_args['num_inference_steps'] + shared.state.sampling_steps = base_args.get('num_inference_steps', p.steps) p.extra_generation_params['Pipeline'] = shared.sd_model.__class__.__name__ if shared.opts.scheduler_eta is not None and shared.opts.scheduler_eta > 0 and shared.opts.scheduler_eta < 1: p.extra_generation_params["Sampler Eta"] = shared.opts.scheduler_eta diff --git a/modules/sd_models.py b/modules/sd_models.py index fc1abca30..64679973e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -586,6 +586,10 @@ def detect_pipeline(f: str, op: str = 'model', warning=True): if shared.backend == shared.Backend.ORIGINAL: warn(f'Model detected as InstaFlow model, but attempting to load using backend=original: {op}={f} size={size} MB') guess = 'InstaFlow' + if 'SegMoE' in f: + if shared.backend == shared.Backend.ORIGINAL: + warn(f'Model detected as SegMoE model, but attempting to load using backend=original: {op}={f} size={size} MB') + guess = 'SegMoE' if 'PixArt' in f: if shared.backend == shared.Backend.ORIGINAL: warn(f'Model detected as PixArt Alpha model, but attempting to load using backend=original: {op}={f} size={size} MB') @@ -794,6 +798,14 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No except Exception as e: shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}') return + if model_type in ['SegMoE']: # forced pipeline + try: + from modules.segmoe.segmoe_model import SegMoEPipeline + sd_model = SegMoEPipeline(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) + sd_model = sd_model.pipe # segmoe pipe does its stuff in __init__ and __call__ is the original pipeline + except Exception as e: + shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}') + return elif 'ONNX' in model_type: # forced pipeline sd_model = pipeline.from_pretrained(checkpoint_info.path) else: diff --git a/modules/segmoe/segmoe_model.py b/modules/segmoe/segmoe_model.py new file mode 100644 index 000000000..3dd5a8cb1 --- /dev/null +++ b/modules/segmoe/segmoe_model.py @@ -0,0 +1,1315 @@ +import gc +from collections import OrderedDict +from typing import Any, Dict, Callable +import os +from copy import deepcopy +from math import ceil +import json +import safetensors +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers import ( + DiffusionPipeline, + StableDiffusionPipeline, + StableDiffusionXLPipeline, + DDPMScheduler, + UNet2DConditionModel, +) +import tqdm +import yaml + + +def remove_all_forward_hooks(model: torch.nn.Module) -> None: + for _name, child in model._modules.items(): # pylint: disable=protected-access + if child is not None: + if hasattr(child, "_forward_hooks"): + child._forward_hooks: Dict[int, Callable] = OrderedDict() + remove_all_forward_hooks(child) + + +# Inspired from transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock +class SparseMoeBlock(nn.Module): + def __init__(self, config, experts): + super().__init__() + self.hidden_dim = config["hidden_size"] + self.num_experts = config["num_local_experts"] + self.top_k = config["num_experts_per_tok"] + self.out_dim = config.get("out_dim", self.hidden_dim) + + # gating + self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False) + self.experts = nn.ModuleList([deepcopy(exp) for exp in experts]) + + def forward(self, hidden_states: torch.Tensor, scale=None) -> torch.Tensor: # pylint: disable=unused-argument + batch_size, sequence_length, f_map_sz = hidden_states.shape + hidden_states = hidden_states.view(-1, f_map_sz) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + _, selected_experts = torch.topk( + router_logits.sum(dim=0, keepdim=True), self.top_k, dim=1 + ) + routing_weights = F.softmax( + router_logits[:, selected_experts[0]], dim=1, dtype=torch.float + ) + + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, self.out_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Loop over all available experts in the model and perform the computation on each expert + for i, expert_idx in enumerate(selected_experts[0].tolist()): + expert_layer = self.experts[expert_idx] + + current_hidden_states = routing_weights[:, i].view( + batch_size * sequence_length, -1 + ) * expert_layer(hidden_states) + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states = final_hidden_states + current_hidden_states + final_hidden_states = final_hidden_states.reshape( + batch_size, sequence_length, self.out_dim + ) + return final_hidden_states + + +def getActivation(activation, name): + def hook(model, inp, output): # pylint: disable=unused-argument + activation[name] = inp + + return hook + + +class SegMoEPipeline: + def __init__(self, config_or_path, **kwargs) -> Any: + """ + Instantiates the SegMoEPipeline. SegMoEPipeline implements the Segmind Mixture of Diffusion Experts, efficiently combining Stable Diffusion and Stable Diffusion Xl models. + + Usage: + + from segmoe import SegMoEPipeline + pipeline = SegMoEPipeline(config_or_path, **kwargs) + + config_or_path: Path to Config or Directory containing SegMoE checkpoint or HF Card of SegMoE Checkpoint. + + Other Keyword Arguments: + torch_dtype: Data Type to load the pipeline in. (Default: torch.float16) + variant: Variant of the Model. (Default: fp16) + device: Device to load the model on. (Default: cuda) + Other args supported by diffusers.DiffusionPipeline are also supported. + + For more details visit https://github.com/segmind/segmoe. + """ + self.torch_dtype = kwargs.pop("torch_dtype", torch.float16) + self.use_safetensors = kwargs.pop("use_safetensors", True) + self.variant = kwargs.pop("variant", "fp16") + self.device = kwargs.pop("device", "cuda") + if os.path.isfile(config_or_path): + self.load_from_scratch(config_or_path, **kwargs) + else: + if not os.path.isdir(config_or_path): + cached_folder = DiffusionPipeline.download(config_or_path) + else: + cached_folder = config_or_path + unet = self.create_empty(cached_folder) + unet.load_state_dict( + safetensors.torch.load_file( + f"{cached_folder}/unet/diffusion_pytorch_model.safetensors" + ) + ) + self.pipe = DiffusionPipeline.from_pretrained( + cached_folder, + unet=unet, + torch_dtype=self.torch_dtype, + use_safetensors=self.use_safetensors, + ) + self.pipe.to(self.device) + self.pipe.unet.to( + device=self.device, + dtype=self.torch_dtype, + memory_format=torch.channels_last, + ) + + def to(self, *args, **kwargs): # TODO added no-op to avoid error + self.pipe.to(*args, **kwargs) + + def load_from_scratch(self, config: str, **kwargs) -> None: + # Load Config + with open(config, "r", encoding='utf8') as f: + config = yaml.load(f, Loader=yaml.SafeLoader) + self.config = config + if self.config.get("num_experts", None): + self.num_experts = self.config["num_experts"] + else: + if self.config.get("experts", None): + self.num_experts = len(self.config["experts"]) + else: + if self.config.get("loras", None): + self.num_experts = len(self.config["loras"]) + else: + self.num_experts = 1 + num_experts_per_tok = self.config.get("num_experts_per_tok", 1) + self.config["num_experts_per_tok"] = num_experts_per_tok + moe_layers = self.config.get("moe_layers", "attn") + self.config["moe_layers"] = moe_layers + # Load Base Model + if self.config["base_model"].startswith( + "https://civitai.com/api/download/models/" + ): + os.makedirs("base", exist_ok=True) + if not os.path.isfile("base/model.safetensors"): + os.system( + "wget -O " + + "base/model.safetensors" + + self.config["base_model"] + + " --content-disposition" + ) + self.config["base_model"] = "base/model.safetensors" + self.pipe = DiffusionPipeline.from_single_file( + self.config["base_model"], torch_dtype=self.torch_dtype + ) + else: + try: + self.pipe = DiffusionPipeline.from_pretrained( + self.config["base_model"], + torch_dtype=self.torch_dtype, + use_safetensors=self.use_safetensors, + variant=self.variant, + **kwargs, + ) + except Exception: + self.pipe = DiffusionPipeline.from_pretrained( + self.config["base_model"], torch_dtype=self.torch_dtype, **kwargs + ) + if self.pipe.__class__ == StableDiffusionPipeline: + self.up_idx_start = 1 + self.up_idx_end = len(self.pipe.unet.up_blocks) + self.down_idx_start = 0 + self.down_idx_end = len(self.pipe.unet.down_blocks) - 1 + elif self.pipe.__class__ == StableDiffusionXLPipeline: + self.up_idx_start = 0 + self.up_idx_end = len(self.pipe.unet.up_blocks) - 1 + self.down_idx_start = 1 + self.down_idx_end = len(self.pipe.unet.down_blocks) + self.config["up_idx_start"] = self.up_idx_start + self.config["up_idx_end"] = self.up_idx_end + self.config["down_idx_start"] = self.down_idx_start + self.config["down_idx_end"] = self.down_idx_end + + # TODO: Add Support for Scheduler Selection + self.pipe.scheduler = DDPMScheduler.from_config(self.pipe.scheduler.config) + + # Load Experts + experts = [] + positive = [] + negative = [] + if self.config.get("experts", None): + for i, exp in enumerate(self.config["experts"]): + positive.append(exp["positive_prompt"]) + negative.append(exp["negative_prompt"]) + if exp["source_model"].startswith( + "https://civitai.com/api/download/models/" + ): + try: + if not os.path.isfile(f"expert_{i}/model.safetensors"): + os.makedirs(f"expert_{i}", exist_ok=True) + if not os.path.isfile(f"expert_{i}/model.safetensors"): + os.system( + f"wget {exp['source_model']} -O " + + f"expert_{i}/model.safetensors" + + " --content-disposition" + ) + exp["source_model"] = f"expert_{i}/model.safetensors" + expert = DiffusionPipeline.from_single_file( + exp["source_model"], + ).to(self.device, self.torch_dtype) + except Exception as e: + print(f"Expert {i} {exp['source_model']} failed to load") + print("Error:", e) + else: + try: + expert = DiffusionPipeline.from_pretrained( + exp["source_model"], + torch_dtype=self.torch_dtype, + use_safetensors=self.use_safetensors, + variant=self.variant, + **kwargs, + ) + + # TODO: Add Support for Scheduler Selection + expert.scheduler = DDPMScheduler.from_config( + expert.scheduler.config + ) + except Exception: + expert = DiffusionPipeline.from_pretrained( + exp["source_model"], torch_dtype=self.torch_dtype, **kwargs + ) + expert.scheduler = DDPMScheduler.from_config( + expert.scheduler.config + ) + if exp.get("loras", None): + for j, lora in enumerate(exp["loras"]): + if lora.get("positive_prompt", None): + positive[-1] += " " + lora["positive_prompt"] + if lora.get("negative_prompt", None): + negative[-1] += " " + lora["negative_prompt"] + if lora["source_model"].startswith( + "https://civitai.com/api/download/models/" + ): + try: + os.makedirs(f"expert_{i}/lora_{i}", exist_ok=True) + if not os.path.isfile( + f"expert_{i}/lora_{i}/pytorch_lora_weights.safetensors" + ): + os.system( + f"wget {lora['source_model']} -O " + + f"expert_{i}/lora_{j}/pytorch_lora_weights.safetensors" + + " --content-disposition" + ) + lora["source_model"] = f"expert_{j}/lora_{j}" + expert.load_lora_weights(lora["source_model"]) + if len(exp["loras"]) == 1: + expert.fuse_lora() + except Exception as e: + print( + f"Expert{i} LoRA {j} {lora['source_model']} failed to load" + ) + print("Error:", e) + else: + expert.load_lora_weights(lora["source_model"]) + if len(exp["loras"]) == 1: + expert.fuse_lora() + experts.append(expert) + else: + experts = [deepcopy(self.pipe) for _ in range(self.num_experts)] + if self.config.get("experts", None): + if self.config.get("loras", None): + for i, lora in enumerate(self.config["loras"]): + if lora["source_model"].startswith( + "https://civitai.com/api/download/models/" + ): + try: + os.makedirs(f"lora_{i}", exist_ok=True) + if not os.path.isfile( + f"lora_{i}/pytorch_lora_weights.safetensors" + ): + os.system( + f"wget {lora['source_model']} -O " + + f"lora_{i}/pytorch_lora_weights.safetensors" + + " --content-disposition" + ) + lora["source_model"] = f"lora_{i}" + self.pipe.load_lora_weights(lora["source_model"]) + if len(self.config["loras"]) == 1: + self.pipe.fuse_lora() + except Exception as e: + print(f"LoRA {i} {lora['source_model']} failed to load") + print("Error:", e) + else: + self.pipe.load_lora_weights(lora["source_model"]) + if len(self.config["loras"]) == 1: + self.pipe.fuse_lora() + else: + if self.config.get("loras", None): + j = [] + n_loras = len(self.config["loras"]) + i = 0 + positive = [""] * len(experts) + negative = [""] * len(experts) + while n_loras: + n = ceil(n_loras / len(experts)) + j += [i] * n + n_loras -= n + i += 1 + for i, lora in enumerate(self.config["loras"]): + positive[j[i]] += lora["positive_prompt"] + " " + negative[j[i]] += lora["negative_prompt"] + " " + if lora["source_model"].startswith( + "https://civitai.com/api/download/models/" + ): + try: + os.makedirs(f"lora_{i}", exist_ok=True) + if not os.path.isfile( + f"lora_{i}/pytorch_lora_weights.safetensors" + ): + os.system( + f"wget {lora['source_model']} -O " + + f"lora_{i}/pytorch_lora_weights.safetensors" + + " --content-disposition" + ) + lora["source_model"] = f"lora_{i}" + experts[j[i]].load_lora_weights(lora["source_model"]) + experts[j[i]].fuse_lora() + except Exception: + print(f"LoRA {i} {lora['source_model']} failed to load") + else: + experts[j[i]].load_lora_weights(lora["source_model"]) + experts[j[i]].fuse_lora() + + # Replace FF and Attention Layers with Sparse MoE Layers + for i in range(self.down_idx_start, self.down_idx_end): + for j in range(len(self.pipe.unet.down_blocks[i].attentions)): + for k in range( + len(self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks) + ): + if not moe_layers == "attn": + config = { + "hidden_size": next( + self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .ff.parameters() + ).size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + # FF Layers + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .ff + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].ff = SparseMoeBlock(config, layers) + if not moe_layers == "ff": + ## Attns + config = { + "hidden_size": self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": self.num_experts, + } + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_q = SparseMoeBlock(config, layers) + + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_k + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_k = SparseMoeBlock(config, layers) + + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_v + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_v = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_q = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + "out_dim": self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[0], + } + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_k = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[-1], + "out_dim": self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[0], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v + ) + ) + self.pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_v = SparseMoeBlock(config, layers) + + for i in range(self.up_idx_start, self.up_idx_end): + for j in range(len(self.pipe.unet.up_blocks[i].attentions)): + for k in range( + len(self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks) + ): + if not moe_layers == "attn": + config = { + "hidden_size": next( + self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .ff.parameters() + ).size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + # FF Layers + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .ff + ) + ) + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].ff = SparseMoeBlock(config, layers) + + if not moe_layers == "ff": + # Attns + config = { + "hidden_size": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + + layers = [] + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q + ) + ) + + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_q = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_k.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + layers = [] + + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_k + ) + ) + + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_k = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_v.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + layers = [] + + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_v + ) + ) + + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_v = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + layers = [] + + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q + ) + ) + + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_q = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[-1], + "out_dim": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[0], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + + layers = [] + + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k + ) + ) + + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_k = SparseMoeBlock(config, layers) + + config = { + "hidden_size": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[-1], + "out_dim": self.pipe.unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[0], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": len(experts), + } + layers = [] + + for l in range(len(experts)): + layers.append( + deepcopy( + experts[l] + .unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v + ) + ) + + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_v = SparseMoeBlock(config, layers) + + # Routing Weight Initialization + if self.config.get("init", "hidden") == "hidden": + gate_params = self.get_gate_params(experts, positive, negative) + for i in range(self.down_idx_start, self.down_idx_end): + for j in range(len(self.pipe.unet.down_blocks[i].attentions)): + for k in range( + len( + self.pipe.unet.down_blocks[i] + .attentions[j] + .transformer_blocks + ) + ): + # FF Layers + if not moe_layers == "attn": + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[k].ff.gate.weight = nn.Parameter( + gate_params[f"d{i}a{j}t{k}"] + ) + + # Attns + if not moe_layers == "ff": + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn1.to_q.gate.weight = nn.Parameter( + gate_params[f"sattnqd{i}a{j}t{k}"] + ) + + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn1.to_k.gate.weight = nn.Parameter( + gate_params[f"sattnkd{i}a{j}t{k}"] + ) + + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn1.to_v.gate.weight = nn.Parameter( + gate_params[f"sattnvd{i}a{j}t{k}"] + ) + + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn2.to_q.gate.weight = nn.Parameter( + gate_params[f"cattnqd{i}a{j}t{k}"] + ) + + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn2.to_k.gate.weight = nn.Parameter( + gate_params[f"cattnkd{i}a{j}t{k}"] + ) + + self.pipe.unet.down_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn2.to_v.gate.weight = nn.Parameter( + gate_params[f"cattnvd{i}a{j}t{k}"] + ) + + for i in range(self.up_idx_start, self.up_idx_end): + for j in range(len(self.pipe.unet.up_blocks[i].attentions)): + for k in range( + len( + self.pipe.unet.up_blocks[i].attentions[j].transformer_blocks + ) + ): + # FF Layers + if not moe_layers == "attn": + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[k].ff.gate.weight = nn.Parameter( + gate_params[f"u{i}a{j}t{k}"] + ) + if not moe_layers == "ff": + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn1.to_q.gate.weight = nn.Parameter( + gate_params[f"sattnqu{i}a{j}t{k}"] + ) + + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn1.to_k.gate.weight = nn.Parameter( + gate_params[f"sattnku{i}a{j}t{k}"] + ) + + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn1.to_v.gate.weight = nn.Parameter( + gate_params[f"sattnvu{i}a{j}t{k}"] + ) + + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn2.to_q.gate.weight = nn.Parameter( + gate_params[f"cattnqu{i}a{j}t{k}"] + ) + + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn2.to_k.gate.weight = nn.Parameter( + gate_params[f"cattnku{i}a{j}t{k}"] + ) + + self.pipe.unet.up_blocks[i].attentions[ + j + ].transformer_blocks[ + k + ].attn2.to_v.gate.weight = nn.Parameter( + gate_params[f"cattnvu{i}a{j}t{k}"] + ) + self.config["num_experts"] = len(experts) + remove_all_forward_hooks(self.pipe.unet) + try: + del experts + del expert + except Exception: + pass + # Move Model to Device + self.pipe.to(self.device) + self.pipe.unet.to( + device=self.device, + dtype=self.torch_dtype, + memory_format=torch.channels_last, + ) + gc.collect() + torch.cuda.empty_cache() + + def __call__(self, *args: Any, **kwds: Any) -> Any: + """ + Inference the SegMoEPipeline. + + Calls diffusers.DiffusionPipeline forward with the keyword arguments. See https://github.com/segmind/segmoe#usage for detailed usage. + """ + return self.pipe(*args, **kwds) + + def create_empty(self, path): + with open(f"{path}/unet/config.json", encoding='utf8') as f: + config = json.load(f) + self.config = config["segmoe_config"] + unet = UNet2DConditionModel.from_config(config) + num_experts_per_tok = self.config["num_experts_per_tok"] + num_experts = self.config["num_experts"] + moe_layers = self.config["moe_layers"] + self.up_idx_start = self.config["up_idx_start"] + self.up_idx_end = self.config["up_idx_end"] + self.down_idx_start = self.config["down_idx_start"] + self.down_idx_end = self.config["down_idx_end"] + for i in range(self.down_idx_start, self.down_idx_end): + for j in range(len(unet.down_blocks[i].attentions)): + for k in range( + len(unet.down_blocks[i].attentions[j].transformer_blocks) + ): + if not moe_layers == "attn": + config = { + "hidden_size": next( + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .ff.parameters() + ).size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + # FF Layers + layers = [ + unet.down_blocks[i].attentions[j].transformer_blocks[k].ff + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].ff = SparseMoeBlock(config, layers) + if not moe_layers == "ff": + ## Attns + config = { + "hidden_size": unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + layers = [ + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_q = SparseMoeBlock(config, layers) + + layers = [ + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_k + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_k = SparseMoeBlock(config, layers) + + layers = [ + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_v + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_v = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + + layers = [ + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_q = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + "out_dim": unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[0], + } + layers = [ + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_k = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[-1], + "out_dim": unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[0], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + layers = [ + unet.down_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v + ] * num_experts + unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_v = SparseMoeBlock(config, layers) + for i in range(self.up_idx_start, self.up_idx_end): + for j in range(len(unet.up_blocks[i].attentions)): + for k in range(len(unet.up_blocks[i].attentions[j].transformer_blocks)): + if not moe_layers == "attn": + config = { + "hidden_size": next( + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .ff.parameters() + ).size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + # FF Layers + layers = [ + unet.up_blocks[i].attentions[j].transformer_blocks[k].ff + ] * num_experts + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].ff = SparseMoeBlock(config, layers) + + if not moe_layers == "ff": + # Attns + config = { + "hidden_size": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + + layers = [ + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_q + ] * num_experts + + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_q = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_k.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + layers = [ + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_k + ] * num_experts + + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_k = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_v.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + layers = [ + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn1.to_v + ] * num_experts + + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_v = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q.weight.size()[-1], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + layers = [ + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_q + ] * num_experts + + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_q = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[-1], + "out_dim": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k.weight.size()[0], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + + layers = [ + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_k + ] * num_experts + + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_k = SparseMoeBlock(config, layers) + + config = { + "hidden_size": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[-1], + "out_dim": unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v.weight.size()[0], + "num_experts_per_tok": num_experts_per_tok, + "num_local_experts": num_experts, + } + layers = [ + unet.up_blocks[i] + .attentions[j] + .transformer_blocks[k] + .attn2.to_v + ] * num_experts + + unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_v = SparseMoeBlock(config, layers) + return unet + + def save_pretrained(self, path): + """ + Save SegMoEPipeline to Disk. + + Usage: + pipeline.save_pretrained(path) + + Parameters: + path: Path to Directory to save the model in. + """ + for param in self.pipe.unet.parameters(): + param.data = param.data.contiguous() + self.pipe.unet.config["segmoe_config"] = self.config + self.pipe.save_pretrained(path) + safetensors.torch.save_file( + self.pipe.unet.state_dict(), + f"{path}/unet/diffusion_pytorch_model.safetensors", + ) + + def cast_hook(self, pipe, dicts): + for i in range(self.down_idx_start, self.down_idx_end): + for j in range(len(pipe.unet.down_blocks[i].attentions)): + for k in range( + len(pipe.unet.down_blocks[i].attentions[j].transformer_blocks) + ): + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].ff.register_forward_hook(getActivation(dicts, f"d{i}a{j}t{k}")) + + ## Down Self Attns + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_q.register_forward_hook( + getActivation(dicts, f"sattnqd{i}a{j}t{k}") + ) + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_k.register_forward_hook( + getActivation(dicts, f"sattnkd{i}a{j}t{k}") + ) + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_v.register_forward_hook( + getActivation(dicts, f"sattnvd{i}a{j}t{k}") + ) + + ## Down Cross Attns + + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_q.register_forward_hook( + getActivation(dicts, f"cattnqd{i}a{j}t{k}") + ) + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_k.register_forward_hook( + getActivation(dicts, f"cattnkd{i}a{j}t{k}") + ) + pipe.unet.down_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_v.register_forward_hook( + getActivation(dicts, f"cattnvd{i}a{j}t{k}") + ) + + for i in range(self.up_idx_start, self.up_idx_end): + for j in range(len(pipe.unet.up_blocks[i].attentions)): + for k in range( + len(pipe.unet.up_blocks[i].attentions[j].transformer_blocks) + ): + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].ff.register_forward_hook(getActivation(dicts, f"u{i}a{j}t{k}")) + ## Up Self Attns + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_q.register_forward_hook( + getActivation(dicts, f"sattnqu{i}a{j}t{k}") + ) + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_k.register_forward_hook( + getActivation(dicts, f"sattnku{i}a{j}t{k}") + ) + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn1.to_v.register_forward_hook( + getActivation(dicts, f"sattnvu{i}a{j}t{k}") + ) + + ## Up Cross Attns + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_q.register_forward_hook( + getActivation(dicts, f"cattnqu{i}a{j}t{k}") + ) + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_k.register_forward_hook( + getActivation(dicts, f"cattnku{i}a{j}t{k}") + ) + pipe.unet.up_blocks[i].attentions[j].transformer_blocks[ + k + ].attn2.to_v.register_forward_hook( + getActivation(dicts, f"cattnvu{i}a{j}t{k}") + ) + + @torch.no_grad + def get_hidden_states(self, model, positive, negative, average: bool = True): + intermediate = {} + self.cast_hook(model, intermediate) + with torch.no_grad(): + _ = model(positive, negative_prompt=negative, num_inference_steps=25) + hidden = {} + for key in intermediate: + hidden_states = intermediate[key][0][-1] + if average: + # use average over sequence + hidden_states = hidden_states.sum(dim=0) / hidden_states.shape[0] + else: + # take last value + hidden_states = hidden_states[:-1] + hidden[key] = hidden_states.to(self.device) + del intermediate + gc.collect() + torch.cuda.empty_cache() + return hidden + + @torch.no_grad + def get_gate_params( + self, + experts, + positive, + negative, + ): + gate_vects = {} + for i, expert in enumerate(tqdm.tqdm(experts, desc="Expert Prompts")): + expert.to(self.device) + expert.unet.to( + device=self.device, + dtype=self.torch_dtype, + memory_format=torch.channels_last, + ) + hidden_states = self.get_hidden_states(expert, positive[i], negative[i]) + del expert + gc.collect() + torch.cuda.empty_cache() + for h in hidden_states: + if i == 0: + gate_vects[h] = [] + hidden_states[h] /= ( + hidden_states[h].norm(p=2, dim=-1, keepdim=True).clamp(min=1e-8) + ) + gate_vects[h].append(hidden_states[h]) + for h in hidden_states: + gate_vects[h] = torch.stack( + gate_vects[h], dim=0 + ) # (num_expert, num_layer, hidden_size) + gate_vects[h].permute(1, 0) + + return gate_vects diff --git a/modules/shared_items.py b/modules/shared_items.py index a785aeab2..74c4de8ff 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -54,7 +54,8 @@ def get_pipelines(): 'ONNX Stable Diffusion XL': getattr(diffusers, 'OnnxStableDiffusionXLPipeline', None), 'ONNX Stable Diffusion XL Img2Img': getattr(diffusers, 'OnnxStableDiffusionXLImg2ImgPipeline', None), 'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None), - 'InstaFlow': getattr(diffusers, 'StableDiffusionPipeline', None) # dynamically redefined and loaded in sd_models.load_diffuser + 'InstaFlow': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser + 'SegMoE': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser # Segmind SSD-1B, Segmind Tiny } diff --git a/pyproject.toml b/pyproject.toml index a1491bd98..dc5c9c4b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ ignore = [ "C408", # Rewrite as a literal "E402", # Module level import not at top of file "E721", # Do not compare types, use `isinstance()` + "E741", # Do not use variables named `l`, `O`, or `I` "EXE001", # Shebang present "F401", # Imported but unused "ISC003", # Implicit string concatenation