| | import torch |
| |
|
| | from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter |
| | from modules_forge.shared import add_supported_preprocessor |
| | from ldm_patched.modules.samplers import sampling_function |
| | import ldm_patched.ldm.modules.attention as attention |
| |
|
| |
|
| | def sdp(q, k, v, transformer_options): |
| | if q.shape[0] == 0: |
| | return q |
| |
|
| | return attention.optimized_attention(q, k, v, heads=transformer_options["n_heads"], mask=None) |
| |
|
| |
|
| | def adain(x, target_std, target_mean): |
| | if x.shape[0] == 0: |
| | return x |
| |
|
| | std, mean = torch.std_mean(x, dim=(2, 3), keepdim=True, correction=0) |
| | return (((x - mean) / std) * target_std) + target_mean |
| |
|
| |
|
| | def zero_cat(a, b, dim): |
| | if a.shape[0] == 0: |
| | return b |
| | if b.shape[0] == 0: |
| | return a |
| | return torch.cat([a, b], dim=dim) |
| |
|
| |
|
| | class PreprocessorReference(Preprocessor): |
| | def __init__(self, name, use_attn=True, use_adain=True, priority=0): |
| | super().__init__() |
| | self.name = name |
| | self.use_attn = use_attn |
| | self.use_adain = use_adain |
| | self.sorting_priority = priority |
| | self.tags = ['Reference'] |
| | self.slider_resolution = PreprocessorParameter(visible=False) |
| | self.slider_1 = PreprocessorParameter(label='Style Fidelity', value=0.5, minimum=0.0, maximum=1.0, step=0.01, visible=True) |
| | self.show_control_mode = False |
| | self.corp_image_with_a1111_mask_when_in_img2img_inpaint_tab = False |
| | self.do_not_need_model = True |
| |
|
| | self.is_recording_style = False |
| | self.recorded_attn1 = {} |
| | self.recorded_h = {} |
| |
|
| | def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): |
| | unit = kwargs['unit'] |
| | weight = float(unit.weight) |
| | style_fidelity = float(unit.threshold_a) |
| | start_percent = float(unit.guidance_start) |
| | end_percent = float(unit.guidance_end) |
| |
|
| | if process.sd_model.is_sdxl: |
| | style_fidelity = style_fidelity ** 3.0 |
| |
|
| | vae = process.sd_model.forge_objects.vae |
| | |
| |
|
| | latent_image = vae.encode(cond.movedim(1, -1)) |
| | latent_image = process.sd_model.forge_objects.unet.model.latent_format.process_in(latent_image) |
| |
|
| | gen_seed = process.seeds[0] + 1 |
| | gen_cpu = torch.Generator().manual_seed(gen_seed) |
| |
|
| | unet = process.sd_model.forge_objects.unet.clone() |
| | sigma_max = unet.model.model_sampling.percent_to_sigma(start_percent) |
| | sigma_min = unet.model.model_sampling.percent_to_sigma(end_percent) |
| |
|
| | self.recorded_attn1 = {} |
| | self.recorded_h = {} |
| |
|
| | def conditioning_modifier(model, x, timestep, uncond, cond, cond_scale, model_options, seed): |
| | sigma = timestep[0].item() |
| | if not (sigma_min <= sigma <= sigma_max): |
| | return model, x, timestep, uncond, cond, cond_scale, model_options, seed |
| |
|
| | self.is_recording_style = True |
| |
|
| | xt = latent_image.to(x) + torch.randn(x.size(), dtype=x.dtype, generator=gen_cpu).to(x) * sigma |
| | sampling_function(model, xt, timestep, uncond, cond, 1, model_options, seed) |
| |
|
| | self.is_recording_style = False |
| |
|
| | return model, x, timestep, uncond, cond, cond_scale, model_options, seed |
| |
|
| | def block_proc(h, flag, transformer_options): |
| | if not self.use_adain: |
| | return h |
| |
|
| | if flag != 'after': |
| | return h |
| |
|
| | location = transformer_options['block'] |
| |
|
| | sigma = transformer_options["sigmas"][0].item() |
| | if not (sigma_min <= sigma <= sigma_max): |
| | return h |
| |
|
| | channel = int(h.shape[1]) |
| | minimal_channel = 1500 - 1000 * weight |
| |
|
| | if channel < minimal_channel: |
| | return h |
| |
|
| | if self.is_recording_style: |
| | self.recorded_h[location] = torch.std_mean(h, dim=(2, 3), keepdim=True, correction=0) |
| | return h |
| | else: |
| | cond_indices = transformer_options['cond_indices'] |
| | uncond_indices = transformer_options['uncond_indices'] |
| | cond_or_uncond = transformer_options['cond_or_uncond'] |
| | r_std, r_mean = self.recorded_h[location] |
| |
|
| | h_c = h[cond_indices] |
| | h_uc = h[uncond_indices] |
| |
|
| | o_c = adain(h_c, r_std, r_mean) |
| | o_uc_strong = h_uc |
| | o_uc_weak = adain(h_uc, r_std, r_mean) |
| | o_uc = o_uc_weak + (o_uc_strong - o_uc_weak) * style_fidelity |
| |
|
| | recon = [] |
| | for cx in cond_or_uncond: |
| | if cx == 0: |
| | recon.append(o_c) |
| | else: |
| | recon.append(o_uc) |
| |
|
| | o = torch.cat(recon, dim=0) |
| | return o |
| |
|
| | def attn1_proc(q, k, v, transformer_options): |
| | if not self.use_attn: |
| | return sdp(q, k, v, transformer_options) |
| |
|
| | sigma = transformer_options["sigmas"][0].item() |
| | if not (sigma_min <= sigma <= sigma_max): |
| | return sdp(q, k, v, transformer_options) |
| |
|
| | location = (transformer_options['block'][0], transformer_options['block'][1], |
| | transformer_options['block_index']) |
| |
|
| | channel = int(q.shape[2]) |
| | minimal_channel = 1500 - 1280 * weight |
| |
|
| | if channel < minimal_channel: |
| | return sdp(q, k, v, transformer_options) |
| |
|
| | if self.is_recording_style: |
| | self.recorded_attn1[location] = (k, v) |
| | return sdp(q, k, v, transformer_options) |
| | else: |
| | cond_indices = transformer_options['cond_indices'] |
| | uncond_indices = transformer_options['uncond_indices'] |
| | cond_or_uncond = transformer_options['cond_or_uncond'] |
| |
|
| | q_c = q[cond_indices] |
| | q_uc = q[uncond_indices] |
| |
|
| | k_c = k[cond_indices] |
| | k_uc = k[uncond_indices] |
| |
|
| | v_c = v[cond_indices] |
| | v_uc = v[uncond_indices] |
| |
|
| | k_r, v_r = self.recorded_attn1[location] |
| |
|
| | o_c = sdp(q_c, zero_cat(k_c, k_r, dim=1), zero_cat(v_c, v_r, dim=1), transformer_options) |
| | o_uc_strong = sdp(q_uc, k_uc, v_uc, transformer_options) |
| | o_uc_weak = sdp(q_uc, zero_cat(k_uc, k_r, dim=1), zero_cat(v_uc, v_r, dim=1), transformer_options) |
| | o_uc = o_uc_weak + (o_uc_strong - o_uc_weak) * style_fidelity |
| |
|
| | recon = [] |
| | for cx in cond_or_uncond: |
| | if cx == 0: |
| | recon.append(o_c) |
| | else: |
| | recon.append(o_uc) |
| |
|
| | o = torch.cat(recon, dim=0) |
| | return o |
| |
|
| | unet.add_block_modifier(block_proc) |
| | unet.add_conditioning_modifier(conditioning_modifier) |
| | unet.set_model_replace_all(attn1_proc, 'attn1') |
| |
|
| | process.sd_model.forge_objects.unet = unet |
| |
|
| | return cond, mask |
| |
|
| |
|
| | add_supported_preprocessor(PreprocessorReference( |
| | name='reference_only', |
| | use_attn=True, |
| | use_adain=False, |
| | priority=100 |
| | )) |
| |
|
| | add_supported_preprocessor(PreprocessorReference( |
| | name='reference_adain', |
| | use_attn=False, |
| | use_adain=True |
| | )) |
| |
|
| | add_supported_preprocessor(PreprocessorReference( |
| | name='reference_adain+attn', |
| | use_attn=True, |
| | use_adain=True |
| | )) |
| |
|
