wip: update ldm_patched

currently issues with calculate_sigmas call
This commit is contained in:
Manuel Schmid
2024-06-02 20:52:52 +02:00
parent ab01104d42
commit 3881fa0806
84 changed files with 5380 additions and 1923 deletions
-1
View File
@@ -16,7 +16,6 @@ from ldm_patched.modules.sd import load_checkpoint_guess_config
from ldm_patched.contrib.external import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, \
ControlNetApplyAdvanced
from ldm_patched.contrib.external_freelunch import FreeU_V2
from ldm_patched.modules.sample import prepare_mask
from modules.lora import match_lora
from modules.util import get_file_from_folder_list
from ldm_patched.modules.lora import model_lora_keys_unet, model_lora_keys_clip
+2 -2
View File
@@ -288,14 +288,14 @@ def vae_parse(latent):
@torch.no_grad()
@torch.inference_mode()
def calculate_sigmas_all(sampler, model, scheduler, steps):
from ldm_patched.modules.samplers import calculate_sigmas_scheduler
from ldm_patched.modules.samplers import calculate_sigmas
discard_penultimate_sigma = False
if sampler in ['dpm_2', 'dpm_2_ancestral']:
steps += 1
discard_penultimate_sigma = True
sigmas = calculate_sigmas_scheduler(model, scheduler, steps)
sigmas = calculate_sigmas(model, scheduler, steps)
if discard_penultimate_sigma:
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
+25 -34
View File
@@ -1,5 +1,6 @@
# Consistent with Kohya/A1111 to reduce differences between model training and inference.
import json
import os
import torch
import ldm_patched.controlnet.cldm
@@ -62,48 +63,42 @@ def patched_encode_token_weights(self, token_weight_pairs):
return torch.cat(output, dim=-2).to(ldm_patched.modules.model_management.intermediate_device()), first_pooled
def patched_SDClipModel__init__(self, max_length=77, freeze=True, layer="last", layer_idx=None,
textmodel_json_config=None, dtype=None, special_tokens=None,
layer_norm_hidden_state=True, **kwargs):
def patched_SDClipModel__init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=ldm_patched.modules.clip_model.CLIPTextModel,
special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, return_projected_pooled=True): # clip-vit-base-patch32
torch.nn.Module.__init__(self)
assert layer in self.LAYERS
if special_tokens is None:
special_tokens = {"start": 49406, "end": 49407, "pad": 49407}
if textmodel_json_config is None:
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(ldm_patched.modules.sd1_clip.__file__)),
"sd1_clip_config.json")
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(ldm_patched.modules.sd1_clip.__file__)), "sd1_clip_config.json")
config = CLIPTextConfig.from_json_file(textmodel_json_config)
self.num_layers = config.num_hidden_layers
with use_patched_ops(ops.manual_cast):
with modeling_utils.no_init_weights():
self.transformer = CLIPTextModel(config)
if dtype is not None:
self.transformer.to(dtype)
with open(textmodel_json_config) as f:
config = json.load(f)
self.transformer = model_class(config, dtype, device, ldm_patched.modules.ops.manual_cast)
self.num_layers = self.transformer.num_layers
self.transformer.text_model.embeddings.to(torch.float32)
self.max_length = max_length
if freeze:
self.freeze()
self.max_length = max_length
self.layer = layer
self.layer_idx = None
self.special_tokens = special_tokens
# TODO check if necessary
self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
self.enable_attention_masks = False
self.enable_attention_masks = enable_attention_masks
self.layer_norm_hidden_state = layer_norm_hidden_state
self.return_projected_pooled = return_projected_pooled
if layer == "hidden":
assert layer_idx is not None
assert abs(layer_idx) < self.num_layers
self.clip_layer(layer_idx)
self.layer_default = (self.layer, self.layer_idx)
self.set_clip_options({"layer": layer_idx})
self.options_default = (self.layer, self.layer_idx, self.return_projected_pooled)
def patched_SDClipModel_forward(self, tokens):
@@ -122,8 +117,7 @@ def patched_SDClipModel_forward(self, tokens):
if tokens[x, y] == max_token:
break
outputs = self.transformer(input_ids=tokens, attention_mask=attention_mask,
output_hidden_states=self.layer == "hidden")
outputs = self.transformer(tokens, attention_mask, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
self.transformer.set_input_embeddings(backup_embeds)
if self.layer == "last":
@@ -131,17 +125,14 @@ def patched_SDClipModel_forward(self, tokens):
elif self.layer == "pooled":
z = outputs.pooler_output[:, None, :]
else:
z = outputs.hidden_states[self.layer_idx]
if self.layer_norm_hidden_state:
z = self.transformer.text_model.final_layer_norm(z)
z = outputs[1]
if hasattr(outputs, "pooler_output"):
pooled_output = outputs.pooler_output.float()
else:
pooled_output = None
if self.text_projection is not None and pooled_output is not None:
pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
pooled_output = None
if len(outputs) >= 3:
if not self.return_projected_pooled and len(outputs) >= 4 and outputs[3] is not None:
pooled_output = outputs[3].float()
elif outputs[2] is not None:
pooled_output = outputs[2].float()
return z.float(), pooled_output
-2
View File
@@ -51,8 +51,6 @@ def patched_register_schedule(self, given_betas=None, beta_schedule="linear", ti
self.linear_end = linear_end
sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32)
self.set_sigmas(sigmas)
alphas_cumprod = torch.tensor(alphas_cumprod, dtype=torch.float32)
self.set_alphas_cumprod(alphas_cumprod)
return
+128 -66
View File
@@ -9,9 +9,10 @@ from ldm_patched.k_diffusion import sampling as k_diffusion_sampling
from ldm_patched.modules.samplers import normal_scheduler, simple_scheduler, ddim_scheduler
from ldm_patched.modules.model_base import SDXLRefiner, SDXL
from ldm_patched.modules.conds import CONDRegular
from ldm_patched.modules.sample import get_additional_models, get_models_from_cond, cleanup_additional_models
from ldm_patched.modules.samplers import resolve_areas_and_cond_masks, wrap_model, calculate_start_end_timesteps, \
create_cond_with_same_area_if_none, pre_run_control, apply_empty_x_to_equal_area, encode_model_conds
from ldm_patched.modules.sampler_helpers import get_additional_models, get_models_from_cond, cleanup_additional_models
from ldm_patched.modules.samplers import resolve_areas_and_cond_masks, calculate_start_end_timesteps, \
create_cond_with_same_area_if_none, pre_run_control, apply_empty_x_to_equal_area, encode_model_conds, CFGGuider, \
process_conds
current_refiner = None
@@ -84,79 +85,140 @@ def clip_separate_after_preparation(cond, target_model=None, target_clip=None):
return results
@torch.no_grad()
@torch.inference_mode()
# @torch.no_grad()
# @torch.inference_mode()
# def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
# global current_refiner
#
# positive = positive[:]
# negative = negative[:]
#
# resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
# resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)
#
# model_wrap = wrap_model(model)
#
# calculate_start_end_timesteps(model, negative)
# calculate_start_end_timesteps(model, positive)
#
# if latent_image is not None:
# latent_image = model.process_latent_in(latent_image)
#
# if hasattr(model, 'extra_conds'):
# positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
# negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
#
# #make sure each cond area has an opposite one with the same area
# for c in positive:
# create_cond_with_same_area_if_none(negative, c)
# for c in negative:
# create_cond_with_same_area_if_none(positive, c)
#
# # pre_run_control(model, negative + positive)
# pre_run_control(model, positive) # negative is not necessary in Fooocus, 0.5s faster.
#
# apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
# apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
#
# extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
#
# if current_refiner is not None and hasattr(current_refiner.model, 'extra_conds'):
# positive_refiner = clip_separate_after_preparation(positive, target_model=current_refiner.model)
# negative_refiner = clip_separate_after_preparation(negative, target_model=current_refiner.model)
#
# positive_refiner = encode_model_conds(current_refiner.model.extra_conds, positive_refiner, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
# negative_refiner = encode_model_conds(current_refiner.model.extra_conds, negative_refiner, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
#
# def refiner_switch():
# cleanup_additional_models(set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control")))
#
# extra_args["cond"] = positive_refiner
# extra_args["uncond"] = negative_refiner
#
# # clear ip-adapter for refiner
# extra_args['model_options'] = {k: {} if k == 'transformer_options' else v for k, v in extra_args['model_options'].items()}
#
# models, inference_memory = get_additional_models(positive_refiner, negative_refiner, current_refiner.model_dtype())
# ldm_patched.modules.model_management.load_models_gpu(
# [current_refiner] + models,
# model.memory_required([noise.shape[0] * 2] + list(noise.shape[1:])) + inference_memory)
#
# model_wrap.inner_model = current_refiner.model
# print('Refiner Swapped')
# return
#
# def callback_wrap(step, x0, x, total_steps):
# if step == refiner_switch_step and current_refiner is not None:
# refiner_switch()
# if callback is not None:
# # residual_noise_preview = x - x0
# # residual_noise_preview /= residual_noise_preview.std()
# # residual_noise_preview *= x0.std()
# callback(step, x0, x, total_steps)
#
# samples = sampler.sample(model_wrap, sigmas, extra_args, callback_wrap, noise, latent_image, denoise_mask, disable_pbar)
# return model.process_latent_out(samples.to(torch.float32))
def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
global current_refiner
cfg_guider = CFGGuiderHacked(model)
cfg_guider.set_conds(positive, negative)
# TODO cfg_guider.inner_set_conds({"positive": positive}) # negative is not necessary in Fooocus, 0.5s faster.
cfg_guider.set_cfg(cfg)
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
positive = positive[:]
negative = negative[:]
resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)
class CFGGuiderHacked(CFGGuider):
def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
global current_refiner
model_wrap = wrap_model(model)
if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
latent_image = self.inner_model.process_latent_in(latent_image)
calculate_start_end_timesteps(model, negative)
calculate_start_end_timesteps(model, positive)
self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)
if latent_image is not None:
latent_image = model.process_latent_in(latent_image)
extra_args = {"model_options": self.model_options, "seed":seed}
if hasattr(model, 'extra_conds'):
positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
if current_refiner is not None and hasattr(current_refiner.model, 'extra_conds'):
positive_refiner = clip_separate_after_preparation(self.conds['positive'], target_model=current_refiner.model)
negative_refiner = clip_separate_after_preparation(self.conds['negative'], target_model=current_refiner.model)
#make sure each cond area has an opposite one with the same area
for c in positive:
create_cond_with_same_area_if_none(negative, c)
for c in negative:
create_cond_with_same_area_if_none(positive, c)
positive_refiner = encode_model_conds(current_refiner.model.extra_conds, positive_refiner, noise, device,
"positive", latent_image=latent_image, denoise_mask=denoise_mask)
negative_refiner = encode_model_conds(current_refiner.model.extra_conds, negative_refiner, noise, device,
"negative", latent_image=latent_image, denoise_mask=denoise_mask)
# pre_run_control(model, negative + positive)
pre_run_control(model, positive) # negative is not necessary in Fooocus, 0.5s faster.
def refiner_switch():
cleanup_additional_models(
set(get_models_from_cond(self.conds['positive'], "control") + get_models_from_cond(self.conds['negative'], "control")))
apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
extra_args["cond"] = positive_refiner
extra_args["uncond"] = negative_refiner
extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
# clear ip-adapter for refiner
extra_args['model_options'] = {k: {} if k == 'transformer_options' else v for k, v in
extra_args['model_options'].items()}
if current_refiner is not None and hasattr(current_refiner.model, 'extra_conds'):
positive_refiner = clip_separate_after_preparation(positive, target_model=current_refiner.model)
negative_refiner = clip_separate_after_preparation(negative, target_model=current_refiner.model)
models, inference_memory = get_additional_models(positive_refiner, negative_refiner,
current_refiner.model_dtype())
ldm_patched.modules.model_management.load_models_gpu(
[current_refiner] + models,
self.model_patcher.memory_required([noise.shape[0] * 2] + list(noise.shape[1:])) + inference_memory)
positive_refiner = encode_model_conds(current_refiner.model.extra_conds, positive_refiner, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
negative_refiner = encode_model_conds(current_refiner.model.extra_conds, negative_refiner, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
self.inner_model = current_refiner.model
print('Refiner Swapped')
return
def refiner_switch():
cleanup_additional_models(set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control")))
def callback_wrap(step, x0, x, total_steps):
if step == refiner_switch_step and current_refiner is not None:
refiner_switch()
if callback is not None:
# residual_noise_preview = x - x0
# residual_noise_preview /= residual_noise_preview.std()
# residual_noise_preview *= x0.std()
callback(step, x0, x, total_steps)
extra_args["cond"] = positive_refiner
extra_args["uncond"] = negative_refiner
# clear ip-adapter for refiner
extra_args['model_options'] = {k: {} if k == 'transformer_options' else v for k, v in extra_args['model_options'].items()}
models, inference_memory = get_additional_models(positive_refiner, negative_refiner, current_refiner.model_dtype())
ldm_patched.modules.model_management.load_models_gpu(
[current_refiner] + models,
model.memory_required([noise.shape[0] * 2] + list(noise.shape[1:])) + inference_memory)
model_wrap.inner_model = current_refiner.model
print('Refiner Swapped')
return
def callback_wrap(step, x0, x, total_steps):
if step == refiner_switch_step and current_refiner is not None:
refiner_switch()
if callback is not None:
# residual_noise_preview = x - x0
# residual_noise_preview /= residual_noise_preview.std()
# residual_noise_preview *= x0.std()
callback(step, x0, x, total_steps)
samples = sampler.sample(model_wrap, sigmas, extra_args, callback_wrap, noise, latent_image, denoise_mask, disable_pbar)
return model.process_latent_out(samples.to(torch.float32))
samples = sampler.sample(self, sigmas, extra_args, callback_wrap, noise, latent_image, denoise_mask, disable_pbar)
return self.inner_model.process_latent_out(samples.to(torch.float32))
@torch.no_grad()
@@ -167,13 +229,13 @@ def calculate_sigmas_scheduler_hacked(model, scheduler_name, steps):
elif scheduler_name == "exponential":
sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
elif scheduler_name == "normal":
sigmas = normal_scheduler(model, steps)
sigmas = normal_scheduler(model.model_sampling, steps)
elif scheduler_name == "simple":
sigmas = simple_scheduler(model, steps)
sigmas = simple_scheduler(model.model_sampling, steps)
elif scheduler_name == "ddim_uniform":
sigmas = ddim_scheduler(model, steps)
sigmas = ddim_scheduler(model.model_sampling, steps)
elif scheduler_name == "sgm_uniform":
sigmas = normal_scheduler(model, steps, sgm=True)
sigmas = normal_scheduler(model.model_sampling, steps, sgm=True)
elif scheduler_name == "turbo":
sigmas = SDTurboScheduler().get_sigmas(model=model, steps=steps, denoise=1.0)[0]
elif scheduler_name == "align_your_steps":
@@ -184,5 +246,5 @@ def calculate_sigmas_scheduler_hacked(model, scheduler_name, steps):
return sigmas
ldm_patched.modules.samplers.calculate_sigmas_scheduler = calculate_sigmas_scheduler_hacked
ldm_patched.modules.samplers.calculate_sigmas = calculate_sigmas_scheduler_hacked
ldm_patched.modules.samplers.sample = sample_hacked