Spaces:
markmagic
/
Runtime error

Omost / lib_omost /pipeline.py
layerdiffusion
i
9ab270d
raw
history blame
17.1 kB
import numpy as np
import copy
from tqdm.auto import trange
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import *
from diffusers.models.transformers import Transformer2DModel
original_Transformer2DModel_forward = Transformer2DModel.forward
def hacked_Transformer2DModel_forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
timestep: Optional[torch.LongTensor] = None,
added_cond_kwargs: Dict[str, torch.Tensor] = None,
class_labels: Optional[torch.LongTensor] = None,
cross_attention_kwargs: Dict[str, Any] = None,
attention_mask: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
return_dict: bool = True,
):
cross_attention_kwargs = cross_attention_kwargs or {}
cross_attention_kwargs['hidden_states_original_shape'] = hidden_states.shape
return original_Transformer2DModel_forward(
self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, class_labels, cross_attention_kwargs,
attention_mask, encoder_attention_mask, return_dict)
Transformer2DModel.forward = hacked_Transformer2DModel_forward
@torch.no_grad()
def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
"""DPM-Solver++(2M)."""
extra_args = {} if extra_args is None else extra_args
s_in = x.new_ones([x.shape[0]])
sigma_fn = lambda t: t.neg().exp()
t_fn = lambda sigma: sigma.log().neg()
old_denoised = None
for i in trange(len(sigmas) - 1, disable=disable):
denoised = model(x, sigmas[i] * s_in, **extra_args)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
h = t_next - t
if old_denoised is None or sigmas[i + 1] == 0:
x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
else:
h_last = t - t_fn(sigmas[i - 1])
r = h_last / h
denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
old_denoised = denoised
return x
class KModel:
def __init__(self, unet, timesteps=1000, linear_start=0.00085, linear_end=0.012):
betas = torch.linspace(linear_start ** 0.5, linear_end ** 0.5, timesteps, dtype=torch.float64) ** 2
alphas = 1. - betas
alphas_cumprod = torch.tensor(np.cumprod(alphas, axis=0), dtype=torch.float32)
self.sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
self.log_sigmas = self.sigmas.log()
self.sigma_data = 1.0
self.unet = unet
return
@property
def sigma_min(self):
return self.sigmas[0]
@property
def sigma_max(self):
return self.sigmas[-1]
def timestep(self, sigma):
log_sigma = sigma.log()
dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
return dists.abs().argmin(dim=0).view(sigma.shape).to(sigma.device)
def get_sigmas_karras(self, n, rho=7.):
ramp = torch.linspace(0, 1, n)
min_inv_rho = self.sigma_min ** (1 / rho)
max_inv_rho = self.sigma_max ** (1 / rho)
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
return torch.cat([sigmas, sigmas.new_zeros([1])])
def __call__(self, x, sigma, **extra_args):
x_ddim_space = x / (sigma[:, None, None, None] ** 2 + self.sigma_data ** 2) ** 0.5
t = self.timestep(sigma)
cfg_scale = extra_args['cfg_scale']
eps_positive = self.unet(x_ddim_space, t, return_dict=False, **extra_args['positive'])[0]
eps_negative = self.unet(x_ddim_space, t, return_dict=False, **extra_args['negative'])[0]
noise_pred = eps_negative + cfg_scale * (eps_positive - eps_negative)
return x - noise_pred * sigma[:, None, None, None]
class OmostSelfAttnProcessor:
def __call__(self, attn, hidden_states, encoder_hidden_states, hidden_states_original_shape, *args, **kwargs):
batch_size, sequence_length, _ = hidden_states.shape
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
hidden_states = torch.nn.functional.scaled_dot_product_attention(
query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class OmostCrossAttnProcessor:
def __call__(self, attn, hidden_states, encoder_hidden_states, hidden_states_original_shape, *args, **kwargs):
B, C, H, W = hidden_states_original_shape
conds = []
masks = []
for m, c in encoder_hidden_states:
m = torch.nn.functional.interpolate(m[None, None, :, :], (H, W), mode='nearest-exact').flatten().unsqueeze(1).repeat(1, c.size(1))
conds.append(c)
masks.append(m)
conds = torch.cat(conds, dim=1)
masks = torch.cat(masks, dim=1)
mask_bool = masks > 0.5
mask_scale = (H * W) / torch.sum(masks, dim=0, keepdim=True)
batch_size, sequence_length, _ = conds.shape
query = attn.to_q(hidden_states)
key = attn.to_k(conds)
value = attn.to_v(conds)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
mask_bool = mask_bool[None, None, :, :].repeat(query.size(0), query.size(1), 1, 1)
mask_scale = mask_scale[None, None, :, :].repeat(query.size(0), query.size(1), 1, 1)
sim = query @ key.transpose(-2, -1) * attn.scale
sim = sim * mask_scale.to(sim)
sim.masked_fill_(mask_bool.logical_not(), float("-inf"))
sim = sim.softmax(dim=-1)
h = sim @ value
h = h.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
h = attn.to_out[0](h)
h = attn.to_out[1](h)
return h
class StableDiffusionXLOmostPipeline(StableDiffusionXLImg2ImgPipeline):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.k_model = KModel(unet=self.unet)
attn_procs = {}
for name in self.unet.attn_processors.keys():
if name.endswith("attn2.processor"):
attn_procs[name] = OmostCrossAttnProcessor()
else:
attn_procs[name] = OmostSelfAttnProcessor()
self.unet.set_attn_processor(attn_procs)
return
@torch.inference_mode()
def encode_bag_of_subprompts_greedy(self, prefixes: list[str], suffixes: list[str]):
device = self.text_encoder.device
@torch.inference_mode()
def greedy_partition(items, max_sum):
bags = []
current_bag = []
current_sum = 0
for item in items:
num = item['length']
if current_sum + num > max_sum:
if current_bag:
bags.append(current_bag)
current_bag = [item]
current_sum = num
else:
current_bag.append(item)
current_sum += num
if current_bag:
bags.append(current_bag)
return bags
@torch.inference_mode()
def get_77_tokens_in_torch(subprompt_inds, tokenizer):
# Note that all subprompt are theoretically less than 75 tokens (without bos/eos)
result = [tokenizer.bos_token_id] + subprompt_inds[:75] + [tokenizer.eos_token_id] + [tokenizer.pad_token_id] * 75
result = result[:77]
result = torch.tensor([result]).to(device=device, dtype=torch.int64)
return result
@torch.inference_mode()
def merge_with_prefix(bag):
merged_ids_t1 = copy.deepcopy(prefix_ids_t1)
merged_ids_t2 = copy.deepcopy(prefix_ids_t2)
for item in bag:
merged_ids_t1.extend(item['ids_t1'])
merged_ids_t2.extend(item['ids_t2'])
return dict(
ids_t1=get_77_tokens_in_torch(merged_ids_t1, self.tokenizer),
ids_t2=get_77_tokens_in_torch(merged_ids_t2, self.tokenizer_2)
)
@torch.inference_mode()
def double_encode(pair_of_inds):
inds = [pair_of_inds['ids_t1'], pair_of_inds['ids_t2']]
text_encoders = [self.text_encoder, self.text_encoder_2]
pooled_prompt_embeds = None
prompt_embeds_list = []
for text_input_ids, text_encoder in zip(inds, text_encoders):
prompt_embeds = text_encoder(text_input_ids, output_hidden_states=True)
# Only last pooler_output is needed
pooled_prompt_embeds = prompt_embeds.pooler_output
# "2" because SDXL always indexes from the penultimate layer.
prompt_embeds = prompt_embeds.hidden_states[-2]
prompt_embeds_list.append(prompt_embeds)
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
return prompt_embeds, pooled_prompt_embeds
# Begin with tokenizing prefixes
prefix_length = 0
prefix_ids_t1 = []
prefix_ids_t2 = []
for prefix in prefixes:
ids_t1 = self.tokenizer(prefix, truncation=False, add_special_tokens=False).input_ids
ids_t2 = self.tokenizer_2(prefix, truncation=False, add_special_tokens=False).input_ids
assert len(ids_t1) == len(ids_t2)
prefix_length += len(ids_t1)
prefix_ids_t1 += ids_t1
prefix_ids_t2 += ids_t2
# Then tokenizing suffixes
allowed_suffix_length = 75 - prefix_length
suffix_targets = []
for subprompt in suffixes:
# Note that all subprompt are theoretically less than 75 tokens (without bos/eos)
# So we can safely just crop it to 75
ids_t1 = self.tokenizer(subprompt, truncation=False, add_special_tokens=False).input_ids[:75]
ids_t2 = self.tokenizer_2(subprompt, truncation=False, add_special_tokens=False).input_ids[:75]
assert len(ids_t1) == len(ids_t2)
suffix_targets.append(dict(
length=len(ids_t1),
ids_t1=ids_t1,
ids_t2=ids_t2
))
# Then merge prefix and suffix tokens
suffix_targets = greedy_partition(suffix_targets, max_sum=allowed_suffix_length)
targets = [merge_with_prefix(b) for b in suffix_targets]
# Encode!
conds, poolers = [], []
for target in targets:
cond, pooler = double_encode(target)
conds.append(cond)
poolers.append(pooler)
conds_merged = torch.concat(conds, dim=1)
poolers_merged = poolers[0]
return dict(cond=conds_merged, pooler=poolers_merged)
@torch.inference_mode()
def all_conds_from_canvas(self, canvas_outputs, negative_prompt):
mask_all = torch.ones(size=(90, 90), dtype=torch.float32)
negative_cond, negative_pooler = self.encode_cropped_prompt_77tokens(negative_prompt)
negative_result = [(mask_all, negative_cond)]
positive_result = []
positive_pooler = None
for item in canvas_outputs['bag_of_conditions']:
current_mask = torch.from_numpy(item['mask']).to(torch.float32)
current_prefixes = item['prefixes']
current_suffixes = item['suffixes']
current_cond = self.encode_bag_of_subprompts_greedy(prefixes=current_prefixes, suffixes=current_suffixes)
if positive_pooler is None:
positive_pooler = current_cond['pooler']
positive_result.append((current_mask, current_cond['cond']))
return positive_result, positive_pooler, negative_result, negative_pooler
@torch.inference_mode()
def encode_cropped_prompt_77tokens(self, prompt: str):
device = self.text_encoder.device
tokenizers = [self.tokenizer, self.tokenizer_2]
text_encoders = [self.text_encoder, self.text_encoder_2]
pooled_prompt_embeds = None
prompt_embeds_list = []
for tokenizer, text_encoder in zip(tokenizers, text_encoders):
text_input_ids = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
).input_ids
prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
# Only last pooler_output is needed
pooled_prompt_embeds = prompt_embeds.pooler_output
# "2" because SDXL always indexes from the penultimate layer.
prompt_embeds = prompt_embeds.hidden_states[-2]
prompt_embeds_list.append(prompt_embeds)
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device)
return prompt_embeds, pooled_prompt_embeds
@torch.inference_mode()
def __call__(
self,
initial_latent: torch.FloatTensor = None,
strength: float = 1.0,
num_inference_steps: int = 25,
guidance_scale: float = 5.0,
batch_size: Optional[int] = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
cross_attention_kwargs: Optional[dict] = None,
):
device = self.unet.device
cross_attention_kwargs = cross_attention_kwargs or {}
# Sigmas
sigmas = self.k_model.get_sigmas_karras(int(num_inference_steps / strength))
sigmas = sigmas[-(num_inference_steps + 1):].to(device)
# Initial latents
_, C, H, W = initial_latent.shape
noise = randn_tensor((batch_size, C, H, W), generator=generator, device=device, dtype=self.unet.dtype)
latents = initial_latent.to(noise) + noise * sigmas[0].to(noise)
# Shape
height, width = latents.shape[-2:]
height = height * self.vae_scale_factor
width = width * self.vae_scale_factor
add_time_ids = list((height, width) + (0, 0) + (height, width))
add_time_ids = torch.tensor([add_time_ids], dtype=self.unet.dtype)
add_neg_time_ids = add_time_ids.clone()
# Batch
latents = latents.to(device)
add_time_ids = add_time_ids.repeat(batch_size, 1).to(device)
add_neg_time_ids = add_neg_time_ids.repeat(batch_size, 1).to(device)
prompt_embeds = [(k.to(device), v.repeat(batch_size, 1, 1).to(noise)) for k, v in prompt_embeds]
negative_prompt_embeds = [(k.to(device), v.repeat(batch_size, 1, 1).to(noise)) for k, v in negative_prompt_embeds]
pooled_prompt_embeds = pooled_prompt_embeds.repeat(batch_size, 1).to(noise)
negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1).to(noise)
# Feeds
sampler_kwargs = dict(
cfg_scale=guidance_scale,
positive=dict(
encoder_hidden_states=prompt_embeds,
added_cond_kwargs={"text_embeds": pooled_prompt_embeds, "time_ids": add_time_ids},
cross_attention_kwargs=cross_attention_kwargs
),
negative=dict(
encoder_hidden_states=negative_prompt_embeds,
added_cond_kwargs={"text_embeds": negative_pooled_prompt_embeds, "time_ids": add_neg_time_ids},
cross_attention_kwargs=cross_attention_kwargs
)
)
# Sample
results = sample_dpmpp_2m(self.k_model, latents, sigmas, extra_args=sampler_kwargs, disable=False)
return StableDiffusionXLPipelineOutput(images=results)