mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
1546 lines
63 KiB
Python
1546 lines
63 KiB
Python
import copy
|
|
import gc
|
|
import inspect
|
|
import json
|
|
import random
|
|
import shutil
|
|
import typing
|
|
from typing import Optional, Union, List, Literal
|
|
import os
|
|
from collections import OrderedDict
|
|
import copy
|
|
import yaml
|
|
from PIL import Image
|
|
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
|
|
from torch.nn import Parameter
|
|
from tqdm import tqdm
|
|
from torchvision.transforms import Resize, transforms
|
|
|
|
from toolkit.clip_vision_adapter import ClipVisionAdapter
|
|
from toolkit.custom_adapter import CustomAdapter
|
|
from toolkit.ip_adapter import IPAdapter
|
|
from toolkit.config_modules import ModelConfig, GenerateImageConfig, ModelArch
|
|
from toolkit.models.decorator import Decorator
|
|
from toolkit.paths import KEYMAPS_ROOT
|
|
from toolkit.prompt_utils import inject_trigger_into_prompt, PromptEmbeds, concat_prompt_embeds
|
|
from toolkit.reference_adapter import ReferenceAdapter
|
|
from toolkit.sd_device_states_presets import empty_preset
|
|
from toolkit.train_tools import get_torch_dtype, apply_noise_offset
|
|
import torch
|
|
from toolkit.pipelines import CustomStableDiffusionXLPipeline
|
|
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, DDPMScheduler, \
|
|
LCMScheduler, Transformer2DModel, AutoencoderTiny, ControlNetModel
|
|
import diffusers
|
|
from diffusers import \
|
|
AutoencoderKL, \
|
|
UNet2DConditionModel
|
|
from diffusers import PixArtAlphaPipeline
|
|
from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextModelWithProjection
|
|
from torchvision.transforms import functional as TF
|
|
|
|
from toolkit.accelerator import get_accelerator, unwrap_model
|
|
from typing import TYPE_CHECKING
|
|
from toolkit.print import print_acc
|
|
from toolkit.memory_management import MemoryManager
|
|
|
|
if TYPE_CHECKING:
|
|
from toolkit.lora_special import LoRASpecialNetwork
|
|
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
|
|
|
|
# tell it to shut up
|
|
diffusers.logging.set_verbosity(diffusers.logging.ERROR)
|
|
|
|
SD_PREFIX_VAE = "vae"
|
|
SD_PREFIX_UNET = "unet"
|
|
SD_PREFIX_REFINER_UNET = "refiner_unet"
|
|
SD_PREFIX_TEXT_ENCODER = "te"
|
|
|
|
SD_PREFIX_TEXT_ENCODER1 = "te0"
|
|
SD_PREFIX_TEXT_ENCODER2 = "te1"
|
|
|
|
# prefixed diffusers keys
|
|
DO_NOT_TRAIN_WEIGHTS = [
|
|
"unet_time_embedding.linear_1.bias",
|
|
"unet_time_embedding.linear_1.weight",
|
|
"unet_time_embedding.linear_2.bias",
|
|
"unet_time_embedding.linear_2.weight",
|
|
"refiner_unet_time_embedding.linear_1.bias",
|
|
"refiner_unet_time_embedding.linear_1.weight",
|
|
"refiner_unet_time_embedding.linear_2.bias",
|
|
"refiner_unet_time_embedding.linear_2.weight",
|
|
]
|
|
|
|
DeviceStatePreset = Literal['cache_latents', 'generate']
|
|
|
|
|
|
class BlankNetwork:
|
|
|
|
def __init__(self):
|
|
self.multiplier = 1.0
|
|
self.is_active = True
|
|
self.is_merged_in = False
|
|
self.can_merge_in = False
|
|
|
|
def __enter__(self):
|
|
self.is_active = True
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.is_active = False
|
|
|
|
def train(self):
|
|
pass
|
|
|
|
|
|
def flush():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
|
|
UNET_IN_CHANNELS = 4 # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
|
|
# VAE_SCALE_FACTOR = 8 # 2 ** (len(vae.config.block_out_channels) - 1) = 8
|
|
|
|
|
|
class BaseModel:
|
|
# override these in child classes
|
|
arch = None
|
|
|
|
def __init__(
|
|
self,
|
|
device,
|
|
model_config: ModelConfig,
|
|
dtype='fp16',
|
|
custom_pipeline=None,
|
|
noise_scheduler=None,
|
|
**kwargs
|
|
):
|
|
self.accelerator = get_accelerator()
|
|
self.custom_pipeline = custom_pipeline
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.torch_dtype = get_torch_dtype(dtype)
|
|
self.device_torch = torch.device(device)
|
|
|
|
self.vae_device_torch = torch.device(device)
|
|
self.vae_torch_dtype = get_torch_dtype(model_config.vae_dtype)
|
|
|
|
self.te_device_torch = torch.device(device)
|
|
self.te_torch_dtype = get_torch_dtype(model_config.te_dtype)
|
|
|
|
self.model_config = model_config
|
|
self.prediction_type = "v_prediction" if self.model_config.is_v_pred else "epsilon"
|
|
|
|
self.device_state = None
|
|
|
|
self.pipeline: Union[None, 'StableDiffusionPipeline',
|
|
'CustomStableDiffusionXLPipeline', 'PixArtAlphaPipeline']
|
|
self.vae: Union[None, 'AutoencoderKL']
|
|
self.model: Union[None, 'Transformer2DModel', 'UNet2DConditionModel']
|
|
self.text_encoder: Union[None, 'CLIPTextModel',
|
|
List[Union['CLIPTextModel', 'CLIPTextModelWithProjection']]]
|
|
self.tokenizer: Union[None, 'CLIPTokenizer', List['CLIPTokenizer']]
|
|
self.noise_scheduler: Union[None, 'DDPMScheduler'] = noise_scheduler
|
|
|
|
self.refiner_unet: Union[None, 'UNet2DConditionModel'] = None
|
|
self.assistant_lora: Union[None, 'LoRASpecialNetwork'] = None
|
|
|
|
# sdxl stuff
|
|
self.logit_scale = None
|
|
self.ckppt_info = None
|
|
self.is_loaded = False
|
|
|
|
# to hold network if there is one
|
|
self.network = None
|
|
self.adapter: Union['ControlNetModel', 'T2IAdapter',
|
|
'IPAdapter', 'ReferenceAdapter', None] = None
|
|
self.decorator: Union[Decorator, None] = None
|
|
self.arch: ModelArch = model_config.arch
|
|
|
|
self.use_text_encoder_1 = model_config.use_text_encoder_1
|
|
self.use_text_encoder_2 = model_config.use_text_encoder_2
|
|
|
|
self.config_file = None
|
|
|
|
self.is_flow_matching = False
|
|
|
|
self.quantize_device = self.device_torch
|
|
self.low_vram = self.model_config.low_vram
|
|
|
|
# merge in and preview active with -1 weight
|
|
self.invert_assistant_lora = False
|
|
self._after_sample_img_hooks = []
|
|
self._status_update_hooks = []
|
|
self.is_transformer = False
|
|
|
|
self.sample_prompts_cache = None
|
|
|
|
self.accuracy_recovery_adapter: Union[None, 'LoRASpecialNetwork'] = None
|
|
self.is_multistage = False
|
|
# a list of multistage boundaries starting with train step 1000 to first idx
|
|
self.multistage_boundaries: List[float] = [0.0]
|
|
# a list of trainable multistage boundaries
|
|
self.trainable_multistage_boundaries: List[int] = [0]
|
|
|
|
# set true for models that encode control image into text embeddings
|
|
self.encode_control_in_text_embeddings = False
|
|
# control images will come in as a list for encoding some things if true
|
|
self.has_multiple_control_images = False
|
|
# do not resize control images
|
|
self.use_raw_control_images = False
|
|
|
|
self.memory_manager = MemoryManager(self)
|
|
|
|
# properties for old arch for backwards compatibility
|
|
@property
|
|
def unet(self):
|
|
return self.model
|
|
|
|
# set unet to model
|
|
@unet.setter
|
|
def unet(self, value):
|
|
self.model = value
|
|
|
|
@property
|
|
def transformer(self):
|
|
return self.model
|
|
|
|
@transformer.setter
|
|
def transformer(self, value):
|
|
self.model = value
|
|
|
|
@property
|
|
def unet_unwrapped(self):
|
|
return unwrap_model(self.model)
|
|
|
|
@property
|
|
def model_unwrapped(self):
|
|
return unwrap_model(self.model)
|
|
|
|
@property
|
|
def is_xl(self):
|
|
return self.arch == 'sdxl'
|
|
|
|
@property
|
|
def is_v2(self):
|
|
return self.arch == 'sd2'
|
|
|
|
@property
|
|
def is_ssd(self):
|
|
return self.arch == 'ssd'
|
|
|
|
@property
|
|
def is_v3(self):
|
|
return self.arch == 'sd3'
|
|
|
|
@property
|
|
def is_vega(self):
|
|
return self.arch == 'vega'
|
|
|
|
@property
|
|
def is_pixart(self):
|
|
return self.arch == 'pixart'
|
|
|
|
@property
|
|
def is_auraflow(self):
|
|
return self.arch == 'auraflow'
|
|
|
|
@property
|
|
def is_flux(self):
|
|
return self.arch == 'flux'
|
|
|
|
@property
|
|
def is_lumina2(self):
|
|
return self.arch == 'lumina2'
|
|
|
|
def get_bucket_divisibility(self):
|
|
if self.vae is None:
|
|
return 8
|
|
try:
|
|
divisibility = 2 ** (len(self.vae.config['block_out_channels']) - 1)
|
|
except:
|
|
# if we have a custom vae, it might not have this
|
|
divisibility = 8
|
|
|
|
# flux packs this again,
|
|
if self.is_flux:
|
|
divisibility = divisibility * 2
|
|
return divisibility
|
|
|
|
# these must be implemented in child classes
|
|
def load_model(self):
|
|
# override this in child classes
|
|
raise NotImplementedError(
|
|
"load_model must be implemented in child classes")
|
|
|
|
def get_generation_pipeline(self):
|
|
# override this in child classes
|
|
raise NotImplementedError(
|
|
"get_generation_pipeline must be implemented in child classes")
|
|
|
|
def generate_single_image(
|
|
self,
|
|
pipeline,
|
|
gen_config: GenerateImageConfig,
|
|
conditional_embeds: PromptEmbeds,
|
|
unconditional_embeds: PromptEmbeds,
|
|
generator: torch.Generator,
|
|
extra: dict,
|
|
):
|
|
# override this in child classes
|
|
raise NotImplementedError(
|
|
"generate_single_image must be implemented in child classes")
|
|
|
|
def get_noise_prediction(
|
|
latent_model_input: torch.Tensor,
|
|
timestep: torch.Tensor, # 0 to 1000 scale
|
|
text_embeddings: PromptEmbeds,
|
|
**kwargs
|
|
):
|
|
raise NotImplementedError(
|
|
"get_noise_prediction must be implemented in child classes")
|
|
|
|
def get_prompt_embeds(self, prompt: str, control_images=None) -> PromptEmbeds:
|
|
raise NotImplementedError(
|
|
"get_prompt_embeds must be implemented in child classes")
|
|
|
|
def get_model_has_grad(self):
|
|
raise NotImplementedError(
|
|
"get_model_has_grad must be implemented in child classes")
|
|
|
|
def get_te_has_grad(self):
|
|
raise NotImplementedError(
|
|
"get_te_has_grad must be implemented in child classes")
|
|
|
|
def save_model(self, output_path, meta, save_dtype):
|
|
# todo handle dtype without overloading anything (vram, cpu, etc)
|
|
unwrap_model(self.pipeline).save_pretrained(
|
|
save_directory=output_path,
|
|
safe_serialization=True,
|
|
)
|
|
# save out meta config
|
|
meta_path = os.path.join(output_path, 'aitk_meta.yaml')
|
|
with open(meta_path, 'w') as f:
|
|
yaml.dump(meta, f)
|
|
# end must be implemented in child classes
|
|
|
|
def te_train(self):
|
|
if isinstance(self.text_encoder, list):
|
|
for te in self.text_encoder:
|
|
te.train()
|
|
elif self.text_encoder is not None:
|
|
self.text_encoder.train()
|
|
|
|
def te_eval(self):
|
|
if isinstance(self.text_encoder, list):
|
|
for te in self.text_encoder:
|
|
te.eval()
|
|
elif self.text_encoder is not None:
|
|
self.text_encoder.eval()
|
|
|
|
def _after_sample_image(self, img_num, total_imgs):
|
|
# process all hooks
|
|
for hook in self._after_sample_img_hooks:
|
|
hook(img_num, total_imgs)
|
|
|
|
def add_after_sample_image_hook(self, func):
|
|
self._after_sample_img_hooks.append(func)
|
|
|
|
def _status_update(self, status: str):
|
|
for hook in self._status_update_hooks:
|
|
hook(status)
|
|
|
|
def print_and_status_update(self, status: str):
|
|
print_acc(status)
|
|
self._status_update(status)
|
|
|
|
def add_status_update_hook(self, func):
|
|
self._status_update_hooks.append(func)
|
|
|
|
@torch.no_grad()
|
|
def generate_images(
|
|
self,
|
|
image_configs: List[GenerateImageConfig],
|
|
sampler=None,
|
|
pipeline: Union[None, StableDiffusionPipeline,
|
|
StableDiffusionXLPipeline] = None,
|
|
):
|
|
network = self.network
|
|
merge_multiplier = 1.0
|
|
flush()
|
|
# if using assistant, unfuse it
|
|
if self.model_config.assistant_lora_path is not None:
|
|
print_acc("Unloading assistant lora")
|
|
if self.invert_assistant_lora:
|
|
self.assistant_lora.is_active = True
|
|
# move weights on to the device
|
|
self.assistant_lora.force_to(
|
|
self.device_torch, self.torch_dtype)
|
|
else:
|
|
self.assistant_lora.is_active = False
|
|
|
|
if self.model_config.inference_lora_path is not None:
|
|
print_acc("Loading inference lora")
|
|
self.assistant_lora.is_active = True
|
|
# move weights on to the device
|
|
self.assistant_lora.force_to(self.device_torch, self.torch_dtype)
|
|
|
|
if network is not None:
|
|
network = unwrap_model(self.network)
|
|
network.eval()
|
|
# check if we have the same network weight for all samples. If we do, we can merge in th
|
|
# the network to drastically speed up inference
|
|
unique_network_weights = set(
|
|
[x.network_multiplier for x in image_configs])
|
|
if len(unique_network_weights) == 1 and network.can_merge_in:
|
|
can_merge_in = True
|
|
merge_multiplier = unique_network_weights.pop()
|
|
network.merge_in(merge_weight=merge_multiplier)
|
|
else:
|
|
network = BlankNetwork()
|
|
|
|
self.save_device_state()
|
|
self.set_device_state_preset('generate')
|
|
|
|
# save current seed state for training
|
|
rng_state = torch.get_rng_state()
|
|
cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None
|
|
|
|
if pipeline is None:
|
|
pipeline = self.get_generation_pipeline()
|
|
try:
|
|
pipeline.set_progress_bar_config(disable=True)
|
|
except:
|
|
pass
|
|
|
|
start_multiplier = 1.0
|
|
if network is not None:
|
|
start_multiplier = network.multiplier
|
|
|
|
# pipeline.to(self.device_torch)
|
|
|
|
with network:
|
|
with torch.no_grad():
|
|
if network is not None:
|
|
assert network.is_active
|
|
|
|
for i in tqdm(range(len(image_configs)), desc=f"Generating Images", leave=False):
|
|
gen_config = image_configs[i]
|
|
|
|
extra = {}
|
|
validation_image = None
|
|
if self.adapter is not None and gen_config.adapter_image_path is not None:
|
|
validation_image = Image.open(gen_config.adapter_image_path)
|
|
if ".inpaint." not in gen_config.adapter_image_path:
|
|
validation_image = validation_image.convert("RGB")
|
|
else:
|
|
# make sure it has an alpha
|
|
if validation_image.mode != "RGBA":
|
|
raise ValueError("Inpainting images must have an alpha channel")
|
|
if isinstance(self.adapter, T2IAdapter):
|
|
# not sure why this is double??
|
|
validation_image = validation_image.resize(
|
|
(gen_config.width * 2, gen_config.height * 2))
|
|
extra['image'] = validation_image
|
|
extra['adapter_conditioning_scale'] = gen_config.adapter_conditioning_scale
|
|
if isinstance(self.adapter, ControlNetModel):
|
|
validation_image = validation_image.resize(
|
|
(gen_config.width, gen_config.height))
|
|
extra['image'] = validation_image
|
|
extra['controlnet_conditioning_scale'] = gen_config.adapter_conditioning_scale
|
|
if isinstance(self.adapter, CustomAdapter) and self.adapter.control_lora is not None:
|
|
validation_image = validation_image.resize((gen_config.width, gen_config.height))
|
|
extra['control_image'] = validation_image
|
|
extra['control_image_idx'] = gen_config.ctrl_idx
|
|
if isinstance(self.adapter, IPAdapter) or isinstance(self.adapter, ClipVisionAdapter):
|
|
transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
])
|
|
validation_image = transform(validation_image)
|
|
if isinstance(self.adapter, CustomAdapter):
|
|
# todo allow loading multiple
|
|
transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
])
|
|
validation_image = transform(validation_image)
|
|
self.adapter.num_images = 1
|
|
if isinstance(self.adapter, ReferenceAdapter):
|
|
# need -1 to 1
|
|
validation_image = transforms.ToTensor()(validation_image)
|
|
validation_image = validation_image * 2.0 - 1.0
|
|
validation_image = validation_image.unsqueeze(0)
|
|
self.adapter.set_reference_images(validation_image)
|
|
|
|
if network is not None:
|
|
network.multiplier = gen_config.network_multiplier
|
|
torch.manual_seed(gen_config.seed)
|
|
torch.cuda.manual_seed(gen_config.seed)
|
|
|
|
generator = torch.manual_seed(gen_config.seed)
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, ClipVisionAdapter) \
|
|
and gen_config.adapter_image_path is not None:
|
|
# run through the adapter to saturate the embeds
|
|
conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
|
|
validation_image)
|
|
self.adapter(conditional_clip_embeds)
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, CustomAdapter):
|
|
# handle condition the prompts
|
|
gen_config.prompt = self.adapter.condition_prompt(
|
|
gen_config.prompt,
|
|
is_unconditional=False,
|
|
)
|
|
gen_config.prompt_2 = gen_config.prompt
|
|
gen_config.negative_prompt = self.adapter.condition_prompt(
|
|
gen_config.negative_prompt,
|
|
is_unconditional=True,
|
|
)
|
|
gen_config.negative_prompt_2 = gen_config.negative_prompt
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, CustomAdapter) and validation_image is not None:
|
|
self.adapter.trigger_pre_te(
|
|
tensors_0_1=validation_image,
|
|
is_training=False,
|
|
has_been_preprocessed=False,
|
|
quad_count=4
|
|
)
|
|
|
|
if self.sample_prompts_cache is not None:
|
|
conditional_embeds = self.sample_prompts_cache[i]['conditional'].to(self.device_torch, dtype=self.torch_dtype)
|
|
unconditional_embeds = self.sample_prompts_cache[i]['unconditional'].to(self.device_torch, dtype=self.torch_dtype)
|
|
else:
|
|
ctrl_img = None
|
|
# load the control image if out model uses it in text encoding
|
|
if gen_config.ctrl_img is not None and self.encode_control_in_text_embeddings:
|
|
ctrl_img = Image.open(gen_config.ctrl_img).convert("RGB")
|
|
# convert to 0 to 1 tensor
|
|
ctrl_img = (
|
|
TF.to_tensor(ctrl_img)
|
|
.unsqueeze(0)
|
|
.to(self.device_torch, dtype=self.torch_dtype)
|
|
)
|
|
# encode the prompt ourselves so we can do fun stuff with embeddings
|
|
if isinstance(self.adapter, CustomAdapter):
|
|
self.adapter.is_unconditional_run = False
|
|
conditional_embeds = self.encode_prompt(
|
|
gen_config.prompt,
|
|
gen_config.prompt_2,
|
|
force_all=True,
|
|
control_images=ctrl_img
|
|
)
|
|
|
|
if isinstance(self.adapter, CustomAdapter):
|
|
self.adapter.is_unconditional_run = True
|
|
unconditional_embeds = self.encode_prompt(
|
|
gen_config.negative_prompt,
|
|
gen_config.negative_prompt_2,
|
|
force_all=True,
|
|
control_images=ctrl_img
|
|
)
|
|
if isinstance(self.adapter, CustomAdapter):
|
|
self.adapter.is_unconditional_run = False
|
|
|
|
# allow any manipulations to take place to embeddings
|
|
gen_config.post_process_embeddings(
|
|
conditional_embeds,
|
|
unconditional_embeds,
|
|
)
|
|
|
|
if self.decorator is not None:
|
|
# apply the decorator to the embeddings
|
|
conditional_embeds.text_embeds = self.decorator(
|
|
conditional_embeds.text_embeds)
|
|
unconditional_embeds.text_embeds = self.decorator(
|
|
unconditional_embeds.text_embeds, is_unconditional=True)
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, IPAdapter) \
|
|
and gen_config.adapter_image_path is not None:
|
|
# apply the image projection
|
|
conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(
|
|
validation_image)
|
|
unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(validation_image,
|
|
True)
|
|
conditional_embeds = self.adapter(
|
|
conditional_embeds, conditional_clip_embeds, is_unconditional=False)
|
|
unconditional_embeds = self.adapter(
|
|
unconditional_embeds, unconditional_clip_embeds, is_unconditional=True)
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, CustomAdapter):
|
|
conditional_embeds = self.adapter.condition_encoded_embeds(
|
|
tensors_0_1=validation_image,
|
|
prompt_embeds=conditional_embeds,
|
|
is_training=False,
|
|
has_been_preprocessed=False,
|
|
is_generating_samples=True,
|
|
)
|
|
unconditional_embeds = self.adapter.condition_encoded_embeds(
|
|
tensors_0_1=validation_image,
|
|
prompt_embeds=unconditional_embeds,
|
|
is_training=False,
|
|
has_been_preprocessed=False,
|
|
is_unconditional=True,
|
|
is_generating_samples=True,
|
|
)
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, CustomAdapter) and len(
|
|
gen_config.extra_values) > 0:
|
|
extra_values = torch.tensor([gen_config.extra_values], device=self.device_torch,
|
|
dtype=self.torch_dtype)
|
|
# apply extra values to the embeddings
|
|
self.adapter.add_extra_values(
|
|
extra_values, is_unconditional=False)
|
|
self.adapter.add_extra_values(torch.zeros_like(
|
|
extra_values), is_unconditional=True)
|
|
pass # todo remove, for debugging
|
|
|
|
if self.refiner_unet is not None and gen_config.refiner_start_at < 1.0:
|
|
# if we have a refiner loaded, set the denoising end at the refiner start
|
|
extra['denoising_end'] = gen_config.refiner_start_at
|
|
extra['output_type'] = 'latent'
|
|
if not self.is_xl:
|
|
raise ValueError(
|
|
"Refiner is only supported for XL models")
|
|
|
|
conditional_embeds = conditional_embeds.to(
|
|
self.device_torch, dtype=self.unet.dtype)
|
|
unconditional_embeds = unconditional_embeds.to(
|
|
self.device_torch, dtype=self.unet.dtype)
|
|
|
|
img = self.generate_single_image(
|
|
pipeline,
|
|
gen_config,
|
|
conditional_embeds,
|
|
unconditional_embeds,
|
|
generator,
|
|
extra,
|
|
)
|
|
|
|
gen_config.save_image(img, i)
|
|
gen_config.log_image(img, i)
|
|
self._after_sample_image(i, len(image_configs))
|
|
flush()
|
|
|
|
if self.adapter is not None and isinstance(self.adapter, ReferenceAdapter):
|
|
self.adapter.clear_memory()
|
|
|
|
# clear pipeline and cache to reduce vram usage
|
|
del pipeline
|
|
torch.cuda.empty_cache()
|
|
|
|
# restore training state
|
|
torch.set_rng_state(rng_state)
|
|
if cuda_rng_state is not None:
|
|
torch.cuda.set_rng_state(cuda_rng_state)
|
|
|
|
self.restore_device_state()
|
|
if network is not None:
|
|
network.train()
|
|
network.multiplier = start_multiplier
|
|
|
|
self.unet.to(self.device_torch, dtype=self.torch_dtype)
|
|
if network.is_merged_in:
|
|
network.merge_out(merge_multiplier)
|
|
# self.tokenizer.to(original_device_dict['tokenizer'])
|
|
|
|
# refuse loras
|
|
if self.model_config.assistant_lora_path is not None:
|
|
print_acc("Loading assistant lora")
|
|
if self.invert_assistant_lora:
|
|
self.assistant_lora.is_active = False
|
|
# move weights off the device
|
|
self.assistant_lora.force_to('cpu', self.torch_dtype)
|
|
else:
|
|
self.assistant_lora.is_active = True
|
|
|
|
if self.model_config.inference_lora_path is not None:
|
|
print_acc("Unloading inference lora")
|
|
self.assistant_lora.is_active = False
|
|
# move weights off the device
|
|
self.assistant_lora.force_to('cpu', self.torch_dtype)
|
|
flush()
|
|
|
|
def get_latent_noise(
|
|
self,
|
|
height=None,
|
|
width=None,
|
|
pixel_height=None,
|
|
pixel_width=None,
|
|
batch_size=1,
|
|
noise_offset=0.0,
|
|
):
|
|
VAE_SCALE_FACTOR = 2 ** (
|
|
len(self.vae.config['block_out_channels']) - 1)
|
|
if height is None and pixel_height is None:
|
|
raise ValueError("height or pixel_height must be specified")
|
|
if width is None and pixel_width is None:
|
|
raise ValueError("width or pixel_width must be specified")
|
|
if height is None:
|
|
height = pixel_height // VAE_SCALE_FACTOR
|
|
if width is None:
|
|
width = pixel_width // VAE_SCALE_FACTOR
|
|
|
|
num_channels = self.unet_unwrapped.config['in_channels']
|
|
if self.is_flux:
|
|
# has 64 channels in for some reason
|
|
num_channels = 16
|
|
noise = torch.randn(
|
|
(
|
|
batch_size,
|
|
num_channels,
|
|
height,
|
|
width,
|
|
),
|
|
device=self.unet.device,
|
|
)
|
|
noise = apply_noise_offset(noise, noise_offset)
|
|
return noise
|
|
|
|
def get_latent_noise_from_latents(
|
|
self,
|
|
latents: torch.Tensor,
|
|
noise_offset=0.0
|
|
):
|
|
noise = torch.randn_like(latents)
|
|
noise = apply_noise_offset(noise, noise_offset)
|
|
return noise
|
|
|
|
def add_noise(
|
|
self,
|
|
original_samples: torch.FloatTensor,
|
|
noise: torch.FloatTensor,
|
|
timesteps: torch.IntTensor,
|
|
**kwargs,
|
|
) -> torch.FloatTensor:
|
|
original_samples_chunks = torch.chunk(
|
|
original_samples, original_samples.shape[0], dim=0)
|
|
noise_chunks = torch.chunk(noise, noise.shape[0], dim=0)
|
|
timesteps_chunks = torch.chunk(timesteps, timesteps.shape[0], dim=0)
|
|
|
|
if len(timesteps_chunks) == 1 and len(timesteps_chunks) != len(original_samples_chunks):
|
|
timesteps_chunks = [timesteps_chunks[0]] * \
|
|
len(original_samples_chunks)
|
|
|
|
noisy_latents_chunks = []
|
|
|
|
for idx in range(original_samples.shape[0]):
|
|
noisy_latents = self.noise_scheduler.add_noise(original_samples_chunks[idx], noise_chunks[idx],
|
|
timesteps_chunks[idx])
|
|
noisy_latents_chunks.append(noisy_latents)
|
|
|
|
noisy_latents = torch.cat(noisy_latents_chunks, dim=0)
|
|
return noisy_latents
|
|
|
|
def predict_noise(
|
|
self,
|
|
latents: torch.Tensor,
|
|
text_embeddings: Union[PromptEmbeds, None] = None,
|
|
timestep: Union[int, torch.Tensor] = 1,
|
|
guidance_scale=7.5,
|
|
guidance_rescale=0,
|
|
add_time_ids=None,
|
|
conditional_embeddings: Union[PromptEmbeds, None] = None,
|
|
unconditional_embeddings: Union[PromptEmbeds, None] = None,
|
|
is_input_scaled=False,
|
|
detach_unconditional=False,
|
|
rescale_cfg=None,
|
|
return_conditional_pred=False,
|
|
guidance_embedding_scale=1.0,
|
|
bypass_guidance_embedding=False,
|
|
batch: Union[None, 'DataLoaderBatchDTO'] = None,
|
|
**kwargs,
|
|
):
|
|
conditional_pred = None
|
|
# get the embeddings
|
|
if text_embeddings is None and conditional_embeddings is None:
|
|
raise ValueError(
|
|
"Either text_embeddings or conditional_embeddings must be specified")
|
|
if text_embeddings is None and unconditional_embeddings is not None:
|
|
text_embeddings = concat_prompt_embeds([
|
|
unconditional_embeddings, # negative embedding
|
|
conditional_embeddings, # positive embedding
|
|
])
|
|
elif text_embeddings is None and conditional_embeddings is not None:
|
|
# not doing cfg
|
|
text_embeddings = conditional_embeddings
|
|
|
|
# CFG is comparing neg and positive, if we have concatenated embeddings
|
|
# then we are doing it, otherwise we are not and takes half the time.
|
|
do_classifier_free_guidance = True
|
|
|
|
# check if batch size of embeddings matches batch size of latents
|
|
if isinstance(text_embeddings.text_embeds, list):
|
|
te_batch_size = text_embeddings.text_embeds[0].shape[0]
|
|
else:
|
|
te_batch_size = text_embeddings.text_embeds.shape[0]
|
|
if latents.shape[0] == te_batch_size:
|
|
do_classifier_free_guidance = False
|
|
elif latents.shape[0] * 2 != te_batch_size:
|
|
raise ValueError(
|
|
"Batch size of latents must be the same or half the batch size of text embeddings")
|
|
latents = latents.to(self.device_torch)
|
|
text_embeddings = text_embeddings.to(self.device_torch)
|
|
timestep = timestep.to(self.device_torch)
|
|
|
|
# if timestep is zero dim, unsqueeze it
|
|
if len(timestep.shape) == 0:
|
|
timestep = timestep.unsqueeze(0)
|
|
|
|
# if we only have 1 timestep, we can just use the same timestep for all
|
|
if timestep.shape[0] == 1 and latents.shape[0] > 1:
|
|
# check if it is rank 1 or 2
|
|
if len(timestep.shape) == 1:
|
|
timestep = timestep.repeat(latents.shape[0])
|
|
else:
|
|
timestep = timestep.repeat(latents.shape[0], 0)
|
|
|
|
# handle t2i adapters
|
|
if 'down_intrablock_additional_residuals' in kwargs:
|
|
# go through each item and concat if doing cfg and it doesnt have the same shape
|
|
for idx, item in enumerate(kwargs['down_intrablock_additional_residuals']):
|
|
if do_classifier_free_guidance and item.shape[0] != text_embeddings.text_embeds.shape[0]:
|
|
kwargs['down_intrablock_additional_residuals'][idx] = torch.cat([
|
|
item] * 2, dim=0)
|
|
|
|
# handle controlnet
|
|
if 'down_block_additional_residuals' in kwargs and 'mid_block_additional_residual' in kwargs:
|
|
# go through each item and concat if doing cfg and it doesnt have the same shape
|
|
for idx, item in enumerate(kwargs['down_block_additional_residuals']):
|
|
if do_classifier_free_guidance and item.shape[0] != text_embeddings.text_embeds.shape[0]:
|
|
kwargs['down_block_additional_residuals'][idx] = torch.cat([
|
|
item] * 2, dim=0)
|
|
for idx, item in enumerate(kwargs['mid_block_additional_residual']):
|
|
if do_classifier_free_guidance and item.shape[0] != text_embeddings.text_embeds.shape[0]:
|
|
kwargs['mid_block_additional_residual'][idx] = torch.cat(
|
|
[item] * 2, dim=0)
|
|
|
|
def scale_model_input(model_input, timestep_tensor):
|
|
if is_input_scaled:
|
|
return model_input
|
|
mi_chunks = torch.chunk(model_input, model_input.shape[0], dim=0)
|
|
timestep_chunks = torch.chunk(
|
|
timestep_tensor, timestep_tensor.shape[0], dim=0)
|
|
out_chunks = []
|
|
# unsqueeze if timestep is zero dim
|
|
for idx in range(model_input.shape[0]):
|
|
# if scheduler has step_index
|
|
if hasattr(self.noise_scheduler, '_step_index'):
|
|
self.noise_scheduler._step_index = None
|
|
out_chunks.append(
|
|
self.noise_scheduler.scale_model_input(
|
|
mi_chunks[idx], timestep_chunks[idx])
|
|
)
|
|
return torch.cat(out_chunks, dim=0)
|
|
|
|
with torch.no_grad():
|
|
if do_classifier_free_guidance:
|
|
# if we are doing classifier free guidance, need to double up
|
|
latent_model_input = torch.cat([latents] * 2, dim=0)
|
|
timestep = torch.cat([timestep] * 2)
|
|
else:
|
|
latent_model_input = latents
|
|
|
|
latent_model_input = scale_model_input(
|
|
latent_model_input, timestep)
|
|
|
|
# check if we need to concat timesteps
|
|
if isinstance(timestep, torch.Tensor) and len(timestep.shape) > 1:
|
|
ts_bs = timestep.shape[0]
|
|
if ts_bs != latent_model_input.shape[0]:
|
|
if ts_bs == 1:
|
|
timestep = torch.cat(
|
|
[timestep] * latent_model_input.shape[0])
|
|
elif ts_bs * 2 == latent_model_input.shape[0]:
|
|
timestep = torch.cat([timestep] * 2, dim=0)
|
|
else:
|
|
raise ValueError(
|
|
f"Batch size of latents {latent_model_input.shape[0]} must be the same or half the batch size of timesteps {timestep.shape[0]}")
|
|
|
|
# predict the noise residual
|
|
if self.unet.device != self.device_torch:
|
|
try:
|
|
self.unet.to(self.device_torch)
|
|
except Exception as e:
|
|
pass
|
|
if self.unet.dtype != self.torch_dtype:
|
|
self.unet = self.unet.to(dtype=self.torch_dtype)
|
|
|
|
# check if get_noise prediction has guidance_embedding_scale
|
|
# if it does not, we dont pass it
|
|
signatures = inspect.signature(self.get_noise_prediction).parameters
|
|
|
|
if 'guidance_embedding_scale' in signatures:
|
|
kwargs['guidance_embedding_scale'] = guidance_embedding_scale
|
|
if 'bypass_guidance_embedding' in signatures:
|
|
kwargs['bypass_guidance_embedding'] = bypass_guidance_embedding
|
|
if 'batch' in signatures:
|
|
kwargs['batch'] = batch
|
|
|
|
noise_pred = self.get_noise_prediction(
|
|
latent_model_input=latent_model_input,
|
|
timestep=timestep,
|
|
text_embeddings=text_embeddings,
|
|
**kwargs
|
|
)
|
|
|
|
conditional_pred = noise_pred
|
|
|
|
if do_classifier_free_guidance:
|
|
# perform guidance
|
|
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0)
|
|
conditional_pred = noise_pred_text
|
|
if detach_unconditional:
|
|
noise_pred_uncond = noise_pred_uncond.detach()
|
|
noise_pred = noise_pred_uncond + guidance_scale * (
|
|
noise_pred_text - noise_pred_uncond
|
|
)
|
|
if rescale_cfg is not None and rescale_cfg != guidance_scale:
|
|
with torch.no_grad():
|
|
# do cfg at the target rescale so we can match it
|
|
target_pred_mean_std = noise_pred_uncond + rescale_cfg * (
|
|
noise_pred_text - noise_pred_uncond
|
|
)
|
|
target_mean = target_pred_mean_std.mean(
|
|
[1, 2, 3], keepdim=True).detach()
|
|
target_std = target_pred_mean_std.std(
|
|
[1, 2, 3], keepdim=True).detach()
|
|
|
|
pred_mean = noise_pred.mean(
|
|
[1, 2, 3], keepdim=True).detach()
|
|
pred_std = noise_pred.std([1, 2, 3], keepdim=True).detach()
|
|
|
|
# match the mean and std
|
|
noise_pred = (noise_pred - pred_mean) / pred_std
|
|
noise_pred = (noise_pred * target_std) + target_mean
|
|
|
|
# https://github.com/huggingface/diffusers/blob/7a91ea6c2b53f94da930a61ed571364022b21044/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L775
|
|
if guidance_rescale > 0.0:
|
|
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
|
|
noise_pred = rescale_noise_cfg(
|
|
noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
|
|
|
|
if return_conditional_pred:
|
|
return noise_pred, conditional_pred
|
|
return noise_pred
|
|
|
|
def step_scheduler(self, model_input, latent_input, timestep_tensor, noise_scheduler=None):
|
|
if noise_scheduler is None:
|
|
noise_scheduler = self.noise_scheduler
|
|
# // sometimes they are on the wrong device, no idea why
|
|
if isinstance(noise_scheduler, DDPMScheduler) or isinstance(noise_scheduler, LCMScheduler):
|
|
try:
|
|
noise_scheduler.betas = noise_scheduler.betas.to(
|
|
self.device_torch)
|
|
noise_scheduler.alphas = noise_scheduler.alphas.to(
|
|
self.device_torch)
|
|
noise_scheduler.alphas_cumprod = noise_scheduler.alphas_cumprod.to(
|
|
self.device_torch)
|
|
except Exception as e:
|
|
pass
|
|
|
|
mi_chunks = torch.chunk(model_input, model_input.shape[0], dim=0)
|
|
latent_chunks = torch.chunk(latent_input, latent_input.shape[0], dim=0)
|
|
timestep_chunks = torch.chunk(
|
|
timestep_tensor, timestep_tensor.shape[0], dim=0)
|
|
out_chunks = []
|
|
if len(timestep_chunks) == 1 and len(mi_chunks) > 1:
|
|
# expand timestep to match
|
|
timestep_chunks = timestep_chunks * len(mi_chunks)
|
|
|
|
for idx in range(model_input.shape[0]):
|
|
# Reset it so it is unique for the
|
|
if hasattr(noise_scheduler, '_step_index'):
|
|
noise_scheduler._step_index = None
|
|
if hasattr(noise_scheduler, 'is_scale_input_called'):
|
|
noise_scheduler.is_scale_input_called = True
|
|
out_chunks.append(
|
|
noise_scheduler.step(mi_chunks[idx], timestep_chunks[idx], latent_chunks[idx], return_dict=False)[
|
|
0]
|
|
)
|
|
return torch.cat(out_chunks, dim=0)
|
|
|
|
# ref: https://github.com/huggingface/diffusers/blob/0bab447670f47c28df60fbd2f6a0f833f75a16f5/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L746
|
|
def diffuse_some_steps(
|
|
self,
|
|
latents: torch.FloatTensor,
|
|
text_embeddings: PromptEmbeds,
|
|
total_timesteps: int = 1000,
|
|
start_timesteps=0,
|
|
guidance_scale=1,
|
|
add_time_ids=None,
|
|
bleed_ratio: float = 0.5,
|
|
bleed_latents: torch.FloatTensor = None,
|
|
is_input_scaled=False,
|
|
return_first_prediction=False,
|
|
**kwargs,
|
|
):
|
|
timesteps_to_run = self.noise_scheduler.timesteps[start_timesteps:total_timesteps]
|
|
|
|
first_prediction = None
|
|
|
|
for timestep in tqdm(timesteps_to_run, leave=False):
|
|
timestep = timestep.unsqueeze_(0)
|
|
noise_pred, conditional_pred = self.predict_noise(
|
|
latents,
|
|
text_embeddings,
|
|
timestep,
|
|
guidance_scale=guidance_scale,
|
|
add_time_ids=add_time_ids,
|
|
is_input_scaled=is_input_scaled,
|
|
return_conditional_pred=True,
|
|
**kwargs,
|
|
)
|
|
# some schedulers need to run separately, so do that. (euler for example)
|
|
|
|
if return_first_prediction and first_prediction is None:
|
|
first_prediction = conditional_pred
|
|
|
|
latents = self.step_scheduler(noise_pred, latents, timestep)
|
|
|
|
# if not last step, and bleeding, bleed in some latents
|
|
if bleed_latents is not None and timestep != self.noise_scheduler.timesteps[-1]:
|
|
latents = (latents * (1 - bleed_ratio)) + \
|
|
(bleed_latents * bleed_ratio)
|
|
|
|
# only skip first scaling
|
|
is_input_scaled = False
|
|
|
|
# return latents_steps
|
|
if return_first_prediction:
|
|
return latents, first_prediction
|
|
return latents
|
|
|
|
def encode_prompt(
|
|
self,
|
|
prompt,
|
|
prompt2=None,
|
|
num_images_per_prompt=1,
|
|
force_all=False,
|
|
long_prompts=False,
|
|
max_length=None,
|
|
dropout_prob=0.0,
|
|
control_images=None,
|
|
) -> PromptEmbeds:
|
|
# sd1.5 embeddings are (bs, 77, 768)
|
|
prompt = prompt
|
|
# if it is not a list, make it one
|
|
if not isinstance(prompt, list):
|
|
prompt = [prompt]
|
|
|
|
if prompt2 is not None and not isinstance(prompt2, list):
|
|
prompt2 = [prompt2]
|
|
# if control_images in the signature, pass it. This keep from breaking plugins
|
|
if self.encode_control_in_text_embeddings:
|
|
return self.get_prompt_embeds(prompt, control_images=control_images)
|
|
|
|
return self.get_prompt_embeds(prompt)
|
|
|
|
@torch.no_grad()
|
|
def encode_images(
|
|
self,
|
|
image_list: List[torch.Tensor],
|
|
device=None,
|
|
dtype=None
|
|
):
|
|
if device is None:
|
|
device = self.vae_device_torch
|
|
if dtype is None:
|
|
dtype = self.vae_torch_dtype
|
|
|
|
latent_list = []
|
|
# Move to vae to device if on cpu
|
|
if self.vae.device == 'cpu':
|
|
self.vae.to(device)
|
|
self.vae.eval()
|
|
self.vae.requires_grad_(False)
|
|
# move to device and dtype
|
|
image_list = [image.to(device, dtype=dtype) for image in image_list]
|
|
|
|
VAE_SCALE_FACTOR = 2 ** (
|
|
len(self.vae.config['block_out_channels']) - 1)
|
|
|
|
# resize images if not divisible by 8
|
|
for i in range(len(image_list)):
|
|
image = image_list[i]
|
|
if image.shape[1] % VAE_SCALE_FACTOR != 0 or image.shape[2] % VAE_SCALE_FACTOR != 0:
|
|
image_list[i] = Resize((image.shape[1] // VAE_SCALE_FACTOR * VAE_SCALE_FACTOR,
|
|
image.shape[2] // VAE_SCALE_FACTOR * VAE_SCALE_FACTOR))(image)
|
|
|
|
images = torch.stack(image_list).to(device, dtype=dtype)
|
|
if isinstance(self.vae, AutoencoderTiny):
|
|
latents = self.vae.encode(images, return_dict=False)[0]
|
|
else:
|
|
latents = self.vae.encode(images).latent_dist.sample()
|
|
shift = self.vae.config['shift_factor'] if self.vae.config['shift_factor'] is not None else 0
|
|
|
|
# flux ref https://github.com/black-forest-labs/flux/blob/c23ae247225daba30fbd56058d247cc1b1fc20a3/src/flux/modules/autoencoder.py#L303
|
|
# z = self.scale_factor * (z - self.shift_factor)
|
|
latents = self.vae.config['scaling_factor'] * (latents - shift)
|
|
latents = latents.to(device, dtype=dtype)
|
|
|
|
return latents
|
|
|
|
def decode_latents(
|
|
self,
|
|
latents: torch.Tensor,
|
|
device=None,
|
|
dtype=None
|
|
):
|
|
if device is None:
|
|
device = self.device
|
|
if dtype is None:
|
|
dtype = self.torch_dtype
|
|
|
|
# Move to vae to device if on cpu
|
|
if self.vae.device == 'cpu':
|
|
self.vae.to(self.device)
|
|
latents = latents.to(device, dtype=dtype)
|
|
latents = (
|
|
latents / self.vae.config['scaling_factor']) + self.vae.config['shift_factor']
|
|
images = self.vae.decode(latents).sample
|
|
images = images.to(device, dtype=dtype)
|
|
|
|
return images
|
|
|
|
def encode_image_prompt_pairs(
|
|
self,
|
|
prompt_list: List[str],
|
|
image_list: List[torch.Tensor],
|
|
device=None,
|
|
dtype=None
|
|
):
|
|
# todo check image types and expand and rescale as needed
|
|
# device and dtype are for outputs
|
|
if device is None:
|
|
device = self.device
|
|
if dtype is None:
|
|
dtype = self.torch_dtype
|
|
|
|
embedding_list = []
|
|
latent_list = []
|
|
# embed the prompts
|
|
for prompt in prompt_list:
|
|
embedding = self.encode_prompt(prompt).to(
|
|
self.device_torch, dtype=dtype)
|
|
embedding_list.append(embedding)
|
|
|
|
return embedding_list, latent_list
|
|
|
|
def get_weight_by_name(self, name):
|
|
# weights begin with te{te_num}_ for text encoder
|
|
# weights begin with unet_ for unet_
|
|
if name.startswith('te'):
|
|
key = name[4:]
|
|
# text encoder
|
|
te_num = int(name[2])
|
|
if isinstance(self.text_encoder, list):
|
|
return self.text_encoder[te_num].state_dict()[key]
|
|
else:
|
|
return self.text_encoder.state_dict()[key]
|
|
elif name.startswith('unet'):
|
|
key = name[5:]
|
|
# unet
|
|
return self.unet.state_dict()[key]
|
|
|
|
raise ValueError(f"Unknown weight name: {name}")
|
|
|
|
def inject_trigger_into_prompt(self, prompt, trigger=None, to_replace_list=None, add_if_not_present=False):
|
|
return inject_trigger_into_prompt(
|
|
prompt,
|
|
trigger=trigger,
|
|
to_replace_list=to_replace_list,
|
|
add_if_not_present=add_if_not_present,
|
|
)
|
|
|
|
def state_dict(self, vae=True, text_encoder=True, unet=True):
|
|
state_dict = OrderedDict()
|
|
if vae:
|
|
for k, v in self.vae.state_dict().items():
|
|
new_key = k if k.startswith(
|
|
f"{SD_PREFIX_VAE}") else f"{SD_PREFIX_VAE}_{k}"
|
|
state_dict[new_key] = v
|
|
if text_encoder:
|
|
if isinstance(self.text_encoder, list):
|
|
for i, encoder in enumerate(self.text_encoder):
|
|
for k, v in encoder.state_dict().items():
|
|
new_key = k if k.startswith(
|
|
f"{SD_PREFIX_TEXT_ENCODER}{i}_") else f"{SD_PREFIX_TEXT_ENCODER}{i}_{k}"
|
|
state_dict[new_key] = v
|
|
else:
|
|
for k, v in self.text_encoder.state_dict().items():
|
|
new_key = k if k.startswith(
|
|
f"{SD_PREFIX_TEXT_ENCODER}_") else f"{SD_PREFIX_TEXT_ENCODER}_{k}"
|
|
state_dict[new_key] = v
|
|
if unet:
|
|
for k, v in self.unet.state_dict().items():
|
|
new_key = k if k.startswith(
|
|
f"{SD_PREFIX_UNET}_") else f"{SD_PREFIX_UNET}_{k}"
|
|
state_dict[new_key] = v
|
|
return state_dict
|
|
|
|
def named_parameters(self, vae=True, text_encoder=True, unet=True, refiner=False, state_dict_keys=False) -> \
|
|
OrderedDict[
|
|
str, Parameter]:
|
|
named_params: OrderedDict[str, Parameter] = OrderedDict()
|
|
if vae:
|
|
for name, param in self.vae.named_parameters(recurse=True, prefix=f"{SD_PREFIX_VAE}"):
|
|
named_params[name] = param
|
|
if text_encoder:
|
|
if isinstance(self.text_encoder, list):
|
|
for i, encoder in enumerate(self.text_encoder):
|
|
if self.is_xl and not self.model_config.use_text_encoder_1 and i == 0:
|
|
# dont add these params
|
|
continue
|
|
if self.is_xl and not self.model_config.use_text_encoder_2 and i == 1:
|
|
# dont add these params
|
|
continue
|
|
|
|
for name, param in encoder.named_parameters(recurse=True, prefix=f"{SD_PREFIX_TEXT_ENCODER}{i}"):
|
|
named_params[name] = param
|
|
else:
|
|
for name, param in self.text_encoder.named_parameters(recurse=True, prefix=f"{SD_PREFIX_TEXT_ENCODER}"):
|
|
named_params[name] = param
|
|
if unet:
|
|
if self.is_flux or self.is_lumina2 or self.is_transformer:
|
|
for name, param in self.unet.named_parameters(recurse=True, prefix="transformer"):
|
|
named_params[name] = param
|
|
else:
|
|
for name, param in self.unet.named_parameters(recurse=True, prefix=f"{SD_PREFIX_UNET}"):
|
|
named_params[name] = param
|
|
|
|
if self.model_config.ignore_if_contains is not None:
|
|
# remove params that contain the ignore_if_contains from named params
|
|
for key in list(named_params.keys()):
|
|
if any([s in f"transformer.{key}" for s in self.model_config.ignore_if_contains]):
|
|
del named_params[key]
|
|
if self.model_config.only_if_contains is not None:
|
|
# remove params that do not contain the only_if_contains from named params
|
|
for key in list(named_params.keys()):
|
|
if not any([s in f"transformer.{key}" for s in self.model_config.only_if_contains]):
|
|
del named_params[key]
|
|
|
|
if refiner:
|
|
for name, param in self.refiner_unet.named_parameters(recurse=True, prefix=f"{SD_PREFIX_REFINER_UNET}"):
|
|
named_params[name] = param
|
|
|
|
# convert to state dict keys, jsut replace . with _ on keys
|
|
if state_dict_keys:
|
|
new_named_params = OrderedDict()
|
|
for k, v in named_params.items():
|
|
# replace only the first . with an _
|
|
new_key = k.replace('.', '_', 1)
|
|
new_named_params[new_key] = v
|
|
named_params = new_named_params
|
|
|
|
return named_params
|
|
|
|
def save(self, output_file: str, meta: OrderedDict, save_dtype=get_torch_dtype('fp16'), logit_scale=None):
|
|
self.save_model(
|
|
output_path=output_file,
|
|
meta=meta,
|
|
save_dtype=save_dtype
|
|
)
|
|
|
|
def prepare_optimizer_params(
|
|
self,
|
|
unet=False,
|
|
text_encoder=False,
|
|
text_encoder_lr=None,
|
|
unet_lr=None,
|
|
refiner_lr=None,
|
|
refiner=False,
|
|
default_lr=1e-6,
|
|
):
|
|
# todo maybe only get locon ones?
|
|
# not all items are saved, to make it match, we need to match out save mappings
|
|
# and not train anything not mapped. Also add learning rate
|
|
version = 'sd1'
|
|
if self.is_xl:
|
|
version = 'sdxl'
|
|
if self.is_v2:
|
|
version = 'sd2'
|
|
mapping_filename = f"stable_diffusion_{version}.json"
|
|
mapping_path = os.path.join(KEYMAPS_ROOT, mapping_filename)
|
|
with open(mapping_path, 'r') as f:
|
|
mapping = json.load(f)
|
|
ldm_diffusers_keymap = mapping['ldm_diffusers_keymap']
|
|
|
|
trainable_parameters = []
|
|
|
|
# we use state dict to find params
|
|
|
|
if unet:
|
|
named_params = self.named_parameters(
|
|
vae=False, unet=unet, text_encoder=False, state_dict_keys=True)
|
|
unet_lr = unet_lr if unet_lr is not None else default_lr
|
|
params = []
|
|
for param in named_params.values():
|
|
if param.requires_grad:
|
|
params.append(param)
|
|
|
|
param_data = {"params": params, "lr": unet_lr}
|
|
trainable_parameters.append(param_data)
|
|
print_acc(f"Found {len(params)} trainable parameter in unet")
|
|
|
|
if text_encoder:
|
|
named_params = self.named_parameters(
|
|
vae=False, unet=False, text_encoder=text_encoder, state_dict_keys=True)
|
|
text_encoder_lr = text_encoder_lr if text_encoder_lr is not None else default_lr
|
|
params = []
|
|
for key, diffusers_key in ldm_diffusers_keymap.items():
|
|
if diffusers_key in named_params and diffusers_key not in DO_NOT_TRAIN_WEIGHTS:
|
|
if named_params[diffusers_key].requires_grad:
|
|
params.append(named_params[diffusers_key])
|
|
param_data = {"params": params, "lr": text_encoder_lr}
|
|
trainable_parameters.append(param_data)
|
|
|
|
print_acc(
|
|
f"Found {len(params)} trainable parameter in text encoder")
|
|
|
|
if refiner:
|
|
named_params = self.named_parameters(vae=False, unet=False, text_encoder=False, refiner=True,
|
|
state_dict_keys=True)
|
|
refiner_lr = refiner_lr if refiner_lr is not None else default_lr
|
|
params = []
|
|
for key, diffusers_key in ldm_diffusers_keymap.items():
|
|
diffusers_key = f"refiner_{diffusers_key}"
|
|
if diffusers_key in named_params and diffusers_key not in DO_NOT_TRAIN_WEIGHTS:
|
|
if named_params[diffusers_key].requires_grad:
|
|
params.append(named_params[diffusers_key])
|
|
param_data = {"params": params, "lr": refiner_lr}
|
|
trainable_parameters.append(param_data)
|
|
|
|
print_acc(f"Found {len(params)} trainable parameter in refiner")
|
|
|
|
return trainable_parameters
|
|
|
|
def save_device_state(self):
|
|
# saves the current device state for all modules
|
|
# this is useful for when we want to alter the state and restore it
|
|
unet_has_grad = self.get_model_has_grad()
|
|
|
|
self.device_state = {
|
|
**empty_preset,
|
|
'vae': {
|
|
'training': self.vae.training,
|
|
'device': self.vae.device,
|
|
},
|
|
'unet': {
|
|
'training': self.unet.training,
|
|
'device': self.unet.device,
|
|
'requires_grad': unet_has_grad,
|
|
},
|
|
}
|
|
if isinstance(self.text_encoder, list):
|
|
self.device_state['text_encoder']: List[dict] = []
|
|
for encoder in self.text_encoder:
|
|
te_has_grad = self.get_te_has_grad()
|
|
self.device_state['text_encoder'].append({
|
|
'training': encoder.training,
|
|
'device': encoder.device,
|
|
# todo there has to be a better way to do this
|
|
'requires_grad': te_has_grad
|
|
})
|
|
else:
|
|
te_has_grad = self.get_te_has_grad()
|
|
|
|
self.device_state['text_encoder'] = {
|
|
'training': self.text_encoder.training,
|
|
'device': self.text_encoder.device,
|
|
'requires_grad': te_has_grad
|
|
}
|
|
if self.adapter is not None:
|
|
if isinstance(self.adapter, IPAdapter):
|
|
requires_grad = self.adapter.image_proj_model.training
|
|
adapter_device = self.unet.device
|
|
elif isinstance(self.adapter, T2IAdapter):
|
|
requires_grad = self.adapter.adapter.conv_in.weight.requires_grad
|
|
adapter_device = self.adapter.device
|
|
elif isinstance(self.adapter, ControlNetModel):
|
|
requires_grad = self.adapter.conv_in.training
|
|
adapter_device = self.adapter.device
|
|
elif isinstance(self.adapter, ClipVisionAdapter):
|
|
requires_grad = self.adapter.embedder.training
|
|
adapter_device = self.adapter.device
|
|
elif isinstance(self.adapter, CustomAdapter):
|
|
requires_grad = self.adapter.training
|
|
adapter_device = self.adapter.device
|
|
elif isinstance(self.adapter, ReferenceAdapter):
|
|
# todo update this!!
|
|
requires_grad = True
|
|
adapter_device = self.adapter.device
|
|
else:
|
|
raise ValueError(f"Unknown adapter type: {type(self.adapter)}")
|
|
self.device_state['adapter'] = {
|
|
'training': self.adapter.training,
|
|
'device': adapter_device,
|
|
'requires_grad': requires_grad,
|
|
}
|
|
|
|
if self.refiner_unet is not None:
|
|
self.device_state['refiner_unet'] = {
|
|
'training': self.refiner_unet.training,
|
|
'device': self.refiner_unet.device,
|
|
'requires_grad': self.refiner_unet.conv_in.weight.requires_grad,
|
|
}
|
|
|
|
def restore_device_state(self):
|
|
# restores the device state for all modules
|
|
# this is useful for when we want to alter the state and restore it
|
|
if self.device_state is None:
|
|
return
|
|
self.set_device_state(self.device_state)
|
|
self.device_state = None
|
|
|
|
def set_device_state(self, state):
|
|
if state['vae']['training']:
|
|
self.vae.train()
|
|
else:
|
|
self.vae.eval()
|
|
self.vae.to(state['vae']['device'])
|
|
if state['unet']['training']:
|
|
self.unet.train()
|
|
else:
|
|
self.unet.eval()
|
|
self.unet.to(state['unet']['device'])
|
|
if state['unet']['requires_grad']:
|
|
self.unet.requires_grad_(True)
|
|
else:
|
|
self.unet.requires_grad_(False)
|
|
if isinstance(self.text_encoder, list):
|
|
for i, encoder in enumerate(self.text_encoder):
|
|
if isinstance(state['text_encoder'], list):
|
|
if state['text_encoder'][i]['training']:
|
|
encoder.train()
|
|
else:
|
|
encoder.eval()
|
|
encoder.to(state['text_encoder'][i]['device'])
|
|
encoder.requires_grad_(
|
|
state['text_encoder'][i]['requires_grad'])
|
|
else:
|
|
if state['text_encoder']['training']:
|
|
encoder.train()
|
|
else:
|
|
encoder.eval()
|
|
encoder.to(state['text_encoder']['device'])
|
|
encoder.requires_grad_(
|
|
state['text_encoder']['requires_grad'])
|
|
else:
|
|
if state['text_encoder']['training']:
|
|
self.text_encoder.train()
|
|
else:
|
|
self.text_encoder.eval()
|
|
self.text_encoder.to(state['text_encoder']['device'])
|
|
self.text_encoder.requires_grad_(
|
|
state['text_encoder']['requires_grad'])
|
|
|
|
if self.adapter is not None:
|
|
self.adapter.to(state['adapter']['device'])
|
|
self.adapter.requires_grad_(state['adapter']['requires_grad'])
|
|
if state['adapter']['training']:
|
|
self.adapter.train()
|
|
else:
|
|
self.adapter.eval()
|
|
|
|
if self.refiner_unet is not None:
|
|
self.refiner_unet.to(state['refiner_unet']['device'])
|
|
self.refiner_unet.requires_grad_(
|
|
state['refiner_unet']['requires_grad'])
|
|
if state['refiner_unet']['training']:
|
|
self.refiner_unet.train()
|
|
else:
|
|
self.refiner_unet.eval()
|
|
flush()
|
|
|
|
def set_device_state_preset(self, device_state_preset: DeviceStatePreset):
|
|
# sets a preset for device state
|
|
|
|
# save current state first
|
|
self.save_device_state()
|
|
|
|
active_modules = []
|
|
training_modules = []
|
|
if device_state_preset in ['cache_latents']:
|
|
active_modules = ['vae']
|
|
if device_state_preset in ['cache_clip']:
|
|
active_modules = ['clip']
|
|
if device_state_preset in ['generate']:
|
|
active_modules = ['vae', 'unet',
|
|
'text_encoder', 'adapter', 'refiner_unet']
|
|
|
|
state = copy.deepcopy(empty_preset)
|
|
# vae
|
|
state['vae'] = {
|
|
'training': 'vae' in training_modules,
|
|
'device': self.vae_device_torch if 'vae' in active_modules else 'cpu',
|
|
'requires_grad': 'vae' in training_modules,
|
|
}
|
|
|
|
# unet
|
|
state['unet'] = {
|
|
'training': 'unet' in training_modules,
|
|
'device': self.device_torch if 'unet' in active_modules else 'cpu',
|
|
'requires_grad': 'unet' in training_modules,
|
|
}
|
|
|
|
if self.refiner_unet is not None:
|
|
state['refiner_unet'] = {
|
|
'training': 'refiner_unet' in training_modules,
|
|
'device': self.device_torch if 'refiner_unet' in active_modules else 'cpu',
|
|
'requires_grad': 'refiner_unet' in training_modules,
|
|
}
|
|
|
|
# text encoder
|
|
if isinstance(self.text_encoder, list):
|
|
state['text_encoder'] = []
|
|
for i, encoder in enumerate(self.text_encoder):
|
|
state['text_encoder'].append({
|
|
'training': 'text_encoder' in training_modules,
|
|
'device': self.te_device_torch if 'text_encoder' in active_modules else 'cpu',
|
|
'requires_grad': 'text_encoder' in training_modules,
|
|
})
|
|
else:
|
|
state['text_encoder'] = {
|
|
'training': 'text_encoder' in training_modules,
|
|
'device': self.te_device_torch if 'text_encoder' in active_modules else 'cpu',
|
|
'requires_grad': 'text_encoder' in training_modules,
|
|
}
|
|
|
|
if self.adapter is not None:
|
|
state['adapter'] = {
|
|
'training': 'adapter' in training_modules,
|
|
'device': self.device_torch if 'adapter' in active_modules else 'cpu',
|
|
'requires_grad': 'adapter' in training_modules,
|
|
}
|
|
|
|
self.set_device_state(state)
|
|
|
|
def text_encoder_to(self, *args, **kwargs):
|
|
if isinstance(self.text_encoder, list):
|
|
for encoder in self.text_encoder:
|
|
encoder.to(*args, **kwargs)
|
|
else:
|
|
self.text_encoder.to(*args, **kwargs)
|
|
|
|
def convert_lora_weights_before_save(self, state_dict):
|
|
# can be overridden in child classes to convert weights before saving
|
|
return state_dict
|
|
|
|
def convert_lora_weights_before_load(self, state_dict):
|
|
# can be overridden in child classes to convert weights before loading
|
|
return state_dict
|
|
|
|
def condition_noisy_latents(self, latents: torch.Tensor, batch:'DataLoaderBatchDTO'):
|
|
# can be overridden in child classes to condition latents before noise prediction
|
|
return latents
|
|
|
|
def get_transformer_block_names(self) -> Optional[List[str]]:
|
|
# override in child classes to get transformer block names for lora targeting
|
|
return None
|
|
|
|
def get_base_model_version(self) -> str:
|
|
# override in child classes to get the base model version
|
|
return "unknown"
|
|
|
|
def get_model_to_train(self):
|
|
# called to get model to attach LoRAs to. Can be overridden in child classes
|
|
return self.unet
|