VisualStylePrompting / pipelines /inverted_ve_pipeline.py
taki0112's picture
add code
4f4656c
raw
history blame
75.2 kB
from __future__ import annotations
from diffusers import StableDiffusionPipeline
import torch
from dataclasses import dataclass
from typing import Callable, List, Optional, Union, Any, Dict
import numpy as np
from diffusers.utils import deprecate, logging, BaseOutput
from einops import rearrange, repeat
from torch.nn.functional import grid_sample
from torch.nn import functional as nnf
import torchvision.transforms as T
from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
from diffusers.models import AutoencoderKL, UNet2DConditionModel, attention_processor
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
import PIL
from PIL import Image
from kornia.morphology import dilation
from collections import OrderedDict
from packaging import version
import inspect
from diffusers.utils import (
deprecate,
is_accelerate_available,
is_accelerate_version,
logging,
replace_example_docstring,
)
from diffusers.utils.torch_utils import randn_tensor
import torch.nn as nn
T = torch.Tensor
@dataclass(frozen=True)
class StyleAlignedArgs:
share_group_norm: bool = True
share_layer_norm: bool = True,
share_attention: bool = True
adain_queries: bool = True
adain_keys: bool = True
adain_values: bool = False
full_attention_share: bool = False
keys_scale: float = 1.
only_self_level: float = 0.
def expand_first(feat: T, scale=1., ) -> T:
b = feat.shape[0]
feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1)
if scale == 1:
feat_style = feat_style.expand(2, b // 2, *feat.shape[1:])
else:
feat_style = feat_style.repeat(1, b // 2, 1, 1, 1)
feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1)
return feat_style.reshape(*feat.shape)
def concat_first(feat: T, dim=2, scale=1.) -> T:
feat_style = expand_first(feat, scale=scale)
return torch.cat((feat, feat_style), dim=dim)
def calc_mean_std(feat, eps: float = 1e-5) -> tuple[T, T]:
feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt()
feat_mean = feat.mean(dim=-2, keepdims=True)
return feat_mean, feat_std
def adain(feat: T) -> T:
feat_mean, feat_std = calc_mean_std(feat)
feat_style_mean = expand_first(feat_mean)
feat_style_std = expand_first(feat_std)
feat = (feat - feat_mean) / feat_std
feat = feat * feat_style_std + feat_style_mean
return feat
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
>>> from diffusers import StableDiffusionPipeline
>>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
>>> pipe = pipe.to("cuda")
>>> prompt = "a photo of an astronaut riding a horse on mars"
>>> image = pipe(prompt).images[0]
```
"""
# ACTIVATE_STEP_CANDIDATE = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 1]
def create_image_grid(image_list, rows, cols, padding=10):
# Ensure the number of rows and columns doesn't exceed the number of images
rows = min(rows, len(image_list))
cols = min(cols, len(image_list))
# Get the dimensions of a single image
image_width, image_height = image_list[0].size
# Calculate the size of the output image
grid_width = cols * (image_width + padding) - padding
grid_height = rows * (image_height + padding) - padding
# Create an empty grid image
grid_image = Image.new('RGB', (grid_width, grid_height), (255, 255, 255))
# Paste images into the grid
for i, img in enumerate(image_list[:rows * cols]):
row = i // cols
col = i % cols
x = col * (image_width + padding)
y = row * (image_height + padding)
grid_image.paste(img, (x, y))
return grid_image
class CrossFrameAttnProcessor_backup:
def __init__(self, unet_chunk_size=2):
self.unet_chunk_size = unet_chunk_size
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None):
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
query = attn.to_q(hidden_states)
is_cross_attention = encoder_hidden_states is not None
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
# elif attn.cross_attention_norm:
# encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
# Sparse Attention
if not is_cross_attention:
video_length = key.size()[0] // self.unet_chunk_size
# former_frame_index = torch.arange(video_length) - 1
# former_frame_index[0] = 0
# import pdb; pdb.set_trace()
# if video_length > 3:
# import pdb; pdb.set_trace()
former_frame_index = [0] * video_length
key = rearrange(key, "(b f) d c -> b f d c", f=video_length)
key = key[:, former_frame_index]
key = rearrange(key, "b f d c -> (b f) d c")
value = rearrange(value, "(b f) d c -> b f d c", f=video_length)
value = value[:, former_frame_index]
value = rearrange(value, "b f d c -> (b f) d c")
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class SharedAttentionProcessor:
def __init__(self,
adain_keys=True,
adain_queries=True,
adain_values=False,
keys_scale=1.,
attn_map_save_steps=[]):
super().__init__()
self.adain_queries = adain_queries
self.adain_keys = adain_keys
self.adain_values = adain_values
# self.full_attention_share = style_aligned_args.full_attention_share
self.keys_scale = keys_scale
self.attn_map_save_steps = attn_map_save_steps
def __call__(
self,
attn: attention_processor.Attention,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
**kwargs
):
if not hasattr(attn, "attn_map"):
setattr(attn, "attn_map", {})
setattr(attn, "inference_step", 0)
else:
attn.inference_step += 1
residual = hidden_states
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
is_cross_attention = encoder_hidden_states is not None
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
# elif attn.cross_attention_norm:
# encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_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)
# if self.step >= self.start_inject:
if not is_cross_attention:# and self.share_attention:
if self.adain_queries:
query = adain(query)
if self.adain_keys:
key = adain(key)
if self.adain_values:
value = adain(value)
key = concat_first(key, -2, scale=self.keys_scale)
value = concat_first(value, -2)
hidden_states = nnf.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
else:
hidden_states = nnf.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
# hidden_states = adain(hidden_states)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
class SharedAttentionProcessor_v2:
def __init__(self,
adain_keys=True,
adain_queries=True,
adain_values=False,
keys_scale=1.,
attn_map_save_steps=[]):
super().__init__()
self.adain_queries = adain_queries
self.adain_keys = adain_keys
self.adain_values = adain_values
# self.full_attention_share = style_aligned_args.full_attention_share
self.keys_scale = keys_scale
self.attn_map_save_steps = attn_map_save_steps
def __call__(
self,
attn: attention_processor.Attention,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
**kwargs
):
if not hasattr(attn, "attn_map"):
setattr(attn, "attn_map", {})
setattr(attn, "inference_step", 0)
else:
attn.inference_step += 1
residual = hidden_states
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
is_cross_attention = encoder_hidden_states is not None
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
# elif attn.cross_attention_norm:
# encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
tmp_query_shape = query.shape
tmp_key_shape = key.shape
tmp_value_shape = value.shape
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)
# if self.step >= self.start_inject:
if not is_cross_attention:# and self.share_attention:
if self.adain_queries:
query = adain(query)
if self.adain_keys:
key = adain(key)
if self.adain_values:
value = adain(value)
key = concat_first(key, -2, scale=self.keys_scale)
value = concat_first(value, -2)
# hidden_states = nnf.scaled_dot_product_attention(
# query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
# )
if attn.inference_step in self.attn_map_save_steps:
query = query.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
key = key.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
value = value.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
if attn.inference_step in self.attn_map_save_steps:
attn.attn_map[attn.inference_step] = attention_probs.clone().cpu().detach()
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
else:
hidden_states = nnf.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
# hidden_states = adain(hidden_states)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
else:
hidden_states = nnf.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
# hidden_states = adain(hidden_states)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
if attn.inference_step == 49:
#initialize inference step
attn.inference_step = -1
return hidden_states
def swapping_attention(key, value, chunk_size=2):
chunk_length = key.size()[0] // chunk_size # [text-condition, null-condition]
reference_image_index = [0] * chunk_length # [0 0 0 0 0]
key = rearrange(key, "(b f) d c -> b f d c", f=chunk_length)
key = key[:, reference_image_index] # ref to all
key = rearrange(key, "b f d c -> (b f) d c")
value = rearrange(value, "(b f) d c -> b f d c", f=chunk_length)
value = value[:, reference_image_index] # ref to all
value = rearrange(value, "b f d c -> (b f) d c")
return key, value
class CrossFrameAttnProcessor:
def __init__(self, unet_chunk_size=2, attn_map_save_steps=[],activate_step_indices=None):
self.unet_chunk_size = unet_chunk_size
self.attn_map_save_steps = attn_map_save_steps
self.activate_step_indices = activate_step_indices
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None):
if not hasattr(attn, "attn_map"):
setattr(attn, "attn_map", {})
setattr(attn, "inference_step", 0)
else:
attn.inference_step += 1
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
query = attn.to_q(hidden_states)
is_cross_attention = encoder_hidden_states is not None
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
is_in_inference_step = False
if self.activate_step_indices is not None:
for activate_step_index in self.activate_step_indices:
if attn.inference_step >= activate_step_index[0] and attn.inference_step <= activate_step_index[1]:
is_in_inference_step = True
break
# Swapping Attention
if not is_cross_attention and is_in_inference_step:
key, value = swapping_attention(key, value, self.unet_chunk_size)
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
if attn.inference_step in self.attn_map_save_steps:
attn.attn_map[attn.inference_step] = attention_probs.clone().cpu().detach()
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if attn.inference_step == 49:
attn.inference_step = -1
return hidden_states
class CrossFrameAttnProcessor4Inversion:
def __init__(self, unet_chunk_size=2, attn_map_save_steps=[],activate_step_indices=None):
self.unet_chunk_size = unet_chunk_size
self.attn_map_save_steps = attn_map_save_steps
self.activate_step_indices = activate_step_indices
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None):
if not hasattr(attn, "attn_map"):
setattr(attn, "attn_map", {})
setattr(attn, "inference_step", 0)
else:
attn.inference_step += 1
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
query = attn.to_q(hidden_states)
is_cross_attention = encoder_hidden_states is not None
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
# elif attn.cross_attention_norm:
# encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
is_in_inference_step = False
if self.activate_step_indices is not None:
for activate_step_index in self.activate_step_indices:
if attn.inference_step >= activate_step_index[0] and attn.inference_step <= activate_step_index[1]:
is_in_inference_step = True
break
# Swapping Attention
if not is_cross_attention and is_in_inference_step:
key, value = swapping_attention(key, value, self.unet_chunk_size)
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
# if attn.inference_step > 45 and attn.inference_step < 50:
# if attn.inference_step == 42 or attn.inference_step==49:
if attn.inference_step in self.attn_map_save_steps:
attn.attn_map[attn.inference_step] = attention_probs.clone().cpu().detach()
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if attn.inference_step == 49:
#initialize inference step
attn.inference_step = -1
return hidden_states
class CrossFrameAttnProcessor_store:
def __init__(self, unet_chunk_size=2, attn_map_save_steps=[]):
self.unet_chunk_size = unet_chunk_size
self.attn_map_save_steps = attn_map_save_steps
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None):
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
query = attn.to_q(hidden_states)
is_cross_attention = encoder_hidden_states is not None
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
# elif attn.cross_attention_norm:
# encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
# Swapping Attention
if not is_cross_attention:
key, value = swapping_attention(key, value, self.unet_chunk_size)
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
if not hasattr(attn, "attn_map"):
setattr(attn, "attn_map", {})
setattr(attn, "inference_step", 0)
else:
attn.inference_step += 1
# if attn.inference_step > 45 and attn.inference_step < 50:
# if attn.inference_step == 42 or attn.inference_step==49:
if attn.inference_step in self.attn_map_save_steps:
attn.attn_map[attn.inference_step] = attention_probs.clone().cpu().detach()
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class InvertedVEAttnProcessor:
def __init__(self, unet_chunk_size=2, scale=1.0):
self.unet_chunk_size = unet_chunk_size
self.scale = scale
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None):
batch_size, sequence_length, _ = hidden_states.shape
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
query = attn.to_q(hidden_states)
is_cross_attention = encoder_hidden_states is not None
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.cross_attention_norm:
encoder_hidden_states = attn.norm_cross(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
#Dual Attention
if not is_cross_attention:
ve_key = key.clone()
ve_value = value.clone()
video_length = ve_key.size()[0] // self.unet_chunk_size
former_frame_index = [0] * video_length
ve_key = rearrange(ve_key, "(b f) d c -> b f d c", f=video_length)
ve_key = ve_key[:, former_frame_index]
ve_key = rearrange(ve_key, "b f d c -> (b f) d c")
ve_value = rearrange(ve_value, "(b f) d c -> b f d c", f=video_length)
ve_value = ve_value[:, former_frame_index]
ve_value = rearrange(ve_value, "b f d c -> (b f) d c")
ve_key = attn.head_to_batch_dim(ve_key)
ve_value = attn.head_to_batch_dim(ve_value)
ve_query = attn.head_to_batch_dim(query)
ve_attention_probs = attn.get_attention_scores(ve_query, ve_key, attention_mask)
ve_hidden_states = torch.bmm(ve_attention_probs, ve_value)
ve_hidden_states = attn.batch_to_head_dim(ve_hidden_states)
ve_hidden_states[0,...] = 0
ve_hidden_states[video_length,...] = 0
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
hidden_states = hidden_states + ve_hidden_states * self.scale
else:
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
class AttnProcessor(nn.Module):
r"""
Default processor for performing attention-related computations.
"""
def __init__(
self,
hidden_size=None,
cross_attention_dim=None,
):
super().__init__()
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
temb=None,
):
residual = hidden_states
# import pdb; pdb.set_trace()
# if attn.spatial_norm is not None:
# hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# if attn.group_norm is not None:
# hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
attention_probs = attn.get_attention_scores(query, key, attention_mask)
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
@dataclass
class StableDiffusionPipelineOutput(BaseOutput):
"""
Output class for Stable Diffusion pipelines.
Args:
images (`List[PIL.Image.Image]` or `np.ndarray`)
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
nsfw_content_detected (`List[bool]`)
List of flags denoting whether the corresponding generated image likely represents "not-safe-for-work"
(nsfw) content, or `None` if safety checking could not be performed.
"""
images: Union[List[PIL.Image.Image], np.ndarray]
nsfw_content_detected: Optional[List[bool]]
class FrozenDict(OrderedDict):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
for key, value in self.items():
setattr(self, key, value)
self.__frozen = True
def __delitem__(self, *args, **kwargs):
raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
def setdefault(self, *args, **kwargs):
raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
def pop(self, *args, **kwargs):
raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
def update(self, *args, **kwargs):
raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
def __setattr__(self, name, value):
if hasattr(self, "__frozen") and self.__frozen:
raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
super().__setattr__(name, value)
def __setitem__(self, name, value):
if hasattr(self, "__frozen") and self.__frozen:
raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
super().__setitem__(name, value)
class InvertedVEPipeline(StableDiffusionPipeline):
r"""
Pipeline for text-to-image generation using Stable Diffusion.
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
Args:
vae ([`AutoencoderKL`]):
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
text_encoder ([`CLIPTextModel`]):
Frozen text-encoder. Stable Diffusion uses the text portion of
[CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
tokenizer (`CLIPTokenizer`):
Tokenizer of class
[CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
scheduler ([`SchedulerMixin`]):
A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
[`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
safety_checker ([`StableDiffusionSafetyChecker`]):
Classification module that estimates whether generated images could be considered offensive or harmful.
Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.
feature_extractor ([`CLIPFeatureExtractor`]):
Model that extracts features from generated images to be used as inputs for the `safety_checker`.
"""
_optional_components = ["safety_checker", "feature_extractor"]
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers,
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPFeatureExtractor,
requires_safety_checker: bool = True,
):
# super().__init__()
super().__init__(vae, text_encoder, tokenizer, unet, scheduler,
safety_checker, feature_extractor, requires_safety_checker)
if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
deprecation_message = (
f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
"to update the config accordingly as leaving `steps_offset` might led to incorrect results"
" in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
" it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
" file"
)
deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
new_config = dict(scheduler.config)
new_config["steps_offset"] = 1
scheduler._internal_dict = FrozenDict(new_config)
if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
deprecation_message = (
f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
" `clip_sample` should be set to False in the configuration file. Please make sure to update the"
" config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
" future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
" nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
)
deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
new_config = dict(scheduler.config)
new_config["clip_sample"] = False
scheduler._internal_dict = FrozenDict(new_config)
if safety_checker is None and requires_safety_checker:
logger.warning(
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
)
if safety_checker is not None and feature_extractor is None:
raise ValueError(
"Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
" checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
)
is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
version.parse(unet.config._diffusers_version).base_version
) < version.parse("0.9.0.dev0")
is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
deprecation_message = (
"The configuration file of the unet has set the default `sample_size` to smaller than"
" 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
" following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
" CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
" \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
" configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
" in the config might lead to incorrect results in future versions. If you have downloaded this"
" checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
" the `unet/config.json` file"
)
deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
new_config = dict(unet.config)
new_config["sample_size"] = 64
unet._internal_dict = FrozenDict(new_config)
self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.register_to_config(requires_safety_checker=requires_safety_checker)
def enable_vae_slicing(self):
r"""
Enable sliced VAE decoding.
When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
steps. This is useful to save some memory and allow larger batch sizes.
"""
self.vae.enable_slicing()
def disable_vae_slicing(self):
r"""
Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
computing decoding in one step.
"""
self.vae.disable_slicing()
def enable_vae_tiling(self):
r"""
Enable tiled VAE decoding.
When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
"""
self.vae.enable_tiling()
def disable_vae_tiling(self):
r"""
Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
computing decoding in one step.
"""
self.vae.disable_tiling()
def enable_sequential_cpu_offload(self, gpu_id=0):
r"""
Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
`torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
Note that offloading happens on a submodule basis. Memory savings are higher than with
`enable_model_cpu_offload`, but performance is lower.
"""
if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
from accelerate import cpu_offload
else:
raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
device = torch.device(f"cuda:{gpu_id}")
if self.device.type != "cpu":
self.to("cpu", silence_dtype_warnings=True)
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
cpu_offload(cpu_offloaded_model, device)
if self.safety_checker is not None:
cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)
def enable_model_cpu_offload(self, gpu_id=0):
r"""
Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
`enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
"""
if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
from accelerate import cpu_offload_with_hook
else:
raise ImportError("`enable_model_offload` requires `accelerate v0.17.0` or higher.")
device = torch.device(f"cuda:{gpu_id}")
if self.device.type != "cpu":
self.to("cpu", silence_dtype_warnings=True)
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
hook = None
for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
_, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
if self.safety_checker is not None:
_, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)
# We'll offload the last model manually.
self.final_offload_hook = hook
@property
def _execution_device(self):
r"""
Returns the device on which the pipeline's models will be executed. After calling
`pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
hooks.
"""
if not hasattr(self.unet, "_hf_hook"):
return self.device
for module in self.unet.modules():
if (
hasattr(module, "_hf_hook")
and hasattr(module._hf_hook, "execution_device")
and module._hf_hook.execution_device is not None
):
return torch.device(module._hf_hook.execution_device)
return self.device
def run_safety_checker(self, image, device, dtype):
if self.safety_checker is not None:
safety_checker_input = self.feature_extractor(self.numpy_to_pil(image), return_tensors="pt").to(device)
image, has_nsfw_concept = self.safety_checker(
images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
)
else:
has_nsfw_concept = None
return image, has_nsfw_concept
def decode_latents(self, latents):
latents = 1 / self.vae.config.scaling_factor * latents
image = self.vae.decode(latents).sample
image = (image / 2 + 0.5).clamp(0, 1)
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
image = image.cpu().permute(0, 2, 3, 1).float().numpy()
return image
def prepare_extra_step_kwargs(self, generator, eta):
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
# and should be between [0, 1]
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
# check if the scheduler accepts generator
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
if accepts_generator:
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
def check_inputs(
self,
prompt,
height,
width,
callback_steps,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
if (callback_steps is None) or (
callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
):
raise ValueError(
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
f" {type(callback_steps)}."
)
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents = latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
return latents
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
negative_prompt: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: int = 1,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
target_prompt: Optional[str] = None,
# device: Optional[Union[str, torch.device]] = "cpu",
):
r"""
Function invoked when calling the pipeline for generation.
Args:
prompt (`str` or `List[str]`, *optional*):
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
instead.
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
The width in pixels of the generated image.
num_inference_steps (`int`, *optional*, defaults to 50):
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference.
guidance_scale (`float`, *optional*, defaults to 7.5):
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
`guidance_scale` is defined as `w` of equation 2. of [Imagen
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
usually at the expense of lower image quality.
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
num_images_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
eta (`float`, *optional*, defaults to 0.0):
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
[`schedulers.DDIMScheduler`], will be ignored for others.
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
to make generation deterministic.
latents (`torch.FloatTensor`, *optional*):
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor will ge generated by sampling using the supplied random `generator`.
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generate image. Choose between
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
plain tuple.
callback (`Callable`, *optional*):
A function that will be called every `callback_steps` steps during inference. The function will be
called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
callback_steps (`int`, *optional*, defaults to 1):
The frequency at which the `callback` function will be called. If not specified, the callback will be
called at every step.
cross_attention_kwargs (`dict`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under
`self.processor` in
[diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
Examples:
Returns:
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
When returning a tuple, the first element is a list with the generated images, and the second element is a
list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
(nsfw) content, according to the `safety_checker`.
"""
# 0. Default height and width to unet
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
)
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input prompt
# import pdb; pdb.set_trace()
prompt_embeds = self._encode_prompt(
prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
)
# import pdb; pdb.set_trace()
if target_prompt is not None:
target_prompt_embeds = self._encode_prompt(
target_prompt,
device,
num_images_per_prompt,
do_classifier_free_guidance,
negative_prompt,
prompt_embeds=None,
negative_prompt_embeds=negative_prompt_embeds,
)
prompt_embeds[num_images_per_prompt+1: ] = target_prompt_embeds[num_images_per_prompt+1:]
import pdb; pdb.set_trace()
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
# 5. Prepare latent variables
num_channels_latents = self.unet.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
# 7. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
cross_attention_kwargs=cross_attention_kwargs,
).sample
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
# call the callback, if provided
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
callback(i, t, latents)
if output_type == "latent":
image = latents
has_nsfw_concept = None
elif output_type == "pil":
# 8. Post-processing
image = self.decode_latents(latents)
# 9. Run safety checker
image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
# 10. Convert to PIL
image = self.numpy_to_pil(image)
else:
# 8. Post-processing
image = self.decode_latents(latents)
# 9. Run safety checker
image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
ACTIVATE_LAYER_CANDIDATE= [
'down_blocks.1.attentions.0.transformer_blocks.0.attn1.processor',
'down_blocks.1.attentions.0.transformer_blocks.0.attn2.processor',
'down_blocks.1.attentions.0.transformer_blocks.1.attn1.processor',
'down_blocks.1.attentions.0.transformer_blocks.1.attn2.processor',
'down_blocks.1.attentions.1.transformer_blocks.0.attn1.processor',
'down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor',
'down_blocks.1.attentions.1.transformer_blocks.1.attn1.processor',
'down_blocks.1.attentions.1.transformer_blocks.1.attn2.processor', #8
'down_blocks.2.attentions.0.transformer_blocks.0.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.0.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.1.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.1.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.2.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.2.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.3.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.3.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.4.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.4.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.5.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.5.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.6.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.6.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.7.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.7.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.8.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.8.attn2.processor',
'down_blocks.2.attentions.0.transformer_blocks.9.attn1.processor',
'down_blocks.2.attentions.0.transformer_blocks.9.attn2.processor', #20
'down_blocks.2.attentions.1.transformer_blocks.0.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.0.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.1.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.1.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.2.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.2.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.3.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.3.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.4.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.4.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.5.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.5.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.6.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.6.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.7.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.7.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.8.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.8.attn2.processor',
'down_blocks.2.attentions.1.transformer_blocks.9.attn1.processor',
'down_blocks.2.attentions.1.transformer_blocks.9.attn2.processor',#20
'mid_block.attentions.0.transformer_blocks.0.attn1.processor',
'mid_block.attentions.0.transformer_blocks.0.attn2.processor',
'mid_block.attentions.0.transformer_blocks.1.attn1.processor',
'mid_block.attentions.0.transformer_blocks.1.attn2.processor',
'mid_block.attentions.0.transformer_blocks.2.attn1.processor',
'mid_block.attentions.0.transformer_blocks.2.attn2.processor',
'mid_block.attentions.0.transformer_blocks.3.attn1.processor',
'mid_block.attentions.0.transformer_blocks.3.attn2.processor',
'mid_block.attentions.0.transformer_blocks.4.attn1.processor',
'mid_block.attentions.0.transformer_blocks.4.attn2.processor',
'mid_block.attentions.0.transformer_blocks.5.attn1.processor',
'mid_block.attentions.0.transformer_blocks.5.attn2.processor',
'mid_block.attentions.0.transformer_blocks.6.attn1.processor',
'mid_block.attentions.0.transformer_blocks.6.attn2.processor',
'mid_block.attentions.0.transformer_blocks.7.attn1.processor',
'mid_block.attentions.0.transformer_blocks.7.attn2.processor',
'mid_block.attentions.0.transformer_blocks.8.attn1.processor',
'mid_block.attentions.0.transformer_blocks.8.attn2.processor',
'mid_block.attentions.0.transformer_blocks.9.attn1.processor',
'mid_block.attentions.0.transformer_blocks.9.attn2.processor', #20
'up_blocks.0.attentions.0.transformer_blocks.0.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.0.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.1.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.1.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.2.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.2.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.3.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.3.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.4.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.4.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.5.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.5.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.6.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.6.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.7.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.7.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.8.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.8.attn2.processor',
'up_blocks.0.attentions.0.transformer_blocks.9.attn1.processor',
'up_blocks.0.attentions.0.transformer_blocks.9.attn2.processor',#20
'up_blocks.0.attentions.1.transformer_blocks.0.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.0.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.1.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.1.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.2.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.2.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.3.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.3.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.4.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.4.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.5.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.5.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.6.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.6.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.7.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.7.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.8.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.8.attn2.processor',
'up_blocks.0.attentions.1.transformer_blocks.9.attn1.processor',
'up_blocks.0.attentions.1.transformer_blocks.9.attn2.processor',#20
'up_blocks.0.attentions.2.transformer_blocks.0.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.0.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.1.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.1.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.2.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.2.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.3.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.3.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.4.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.4.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.5.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.5.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.6.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.6.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.7.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.7.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.8.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.8.attn2.processor',
'up_blocks.0.attentions.2.transformer_blocks.9.attn1.processor',
'up_blocks.0.attentions.2.transformer_blocks.9.attn2.processor', #20
'up_blocks.1.attentions.0.transformer_blocks.0.attn1.processor',
'up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor',
'up_blocks.1.attentions.0.transformer_blocks.1.attn1.processor',
'up_blocks.1.attentions.0.transformer_blocks.1.attn2.processor',
'up_blocks.1.attentions.1.transformer_blocks.0.attn1.processor',
'up_blocks.1.attentions.1.transformer_blocks.0.attn2.processor',
'up_blocks.1.attentions.1.transformer_blocks.1.attn1.processor',
'up_blocks.1.attentions.1.transformer_blocks.1.attn2.processor',
'up_blocks.1.attentions.2.transformer_blocks.0.attn1.processor',
'up_blocks.1.attentions.2.transformer_blocks.0.attn2.processor',
'up_blocks.1.attentions.2.transformer_blocks.1.attn1.processor',
'up_blocks.1.attentions.2.transformer_blocks.1.attn2.processor',#12
]
STYLE_DESCRIPTION_DICT = {
"chinese-ink-paint":("{object} in colorful chinese ink paintings style",""),
"cloud":("Photography of {object}, realistic",""),
"digital-art":("{object} in digital glitch arts style",""),
"fire":("{object} photography, realistic, black background'",""),
"klimt":("{object} in style of Gustav Klimt",""),
"line-art":("line art drawing of {object} . professional, sleek, modern, minimalist, graphic, line art, vector graphics",""),
"low-poly":("low-poly style of {object} . low-poly game art, polygon mesh, jagged, blocky, wireframe edges, centered composition",
"noisy, sloppy, messy, grainy, highly detailed, ultra textured, photo"),
"munch":("{object} in Edvard Munch style",""),
"van-gogh":("{object}, Van Gogh",""),
"totoro":("{object}, art by studio ghibli, cinematic, masterpiece,key visual, studio anime, highly detailed",
"photo, deformed, black and white, realism, disfigured, low contrast"),
"realistic": ("A portrait of {object}, photorealistic, 35mm film, realistic",
"gray, ugly, deformed, noisy, blurry"),
"line_art": ("line art drawing of {object} . professional, sleek, modern, minimalist, graphic, line art, vector graphics",
"anime, photorealistic, 35mm film, deformed, glitch, blurry, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, mutated, realism, realistic, impressionism, expressionism, oil, acrylic"
) ,
"anime": ("anime artwork of {object} . anime style, key visual, vibrant, studio anime, highly detailed",
"photo, deformed, black and white, realism, disfigured, low contrast"
),
"Artstyle_Pop_Art" : ("pop Art style of {object} . bright colors, bold outlines, popular culture themes, ironic or kitsch",
"ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, minimalist"
),
"Artstyle_Pointillism": ("pointillism style of {object} . composed entirely of small, distinct dots of color, vibrant, highly detailed",
"line drawing, smooth shading, large color fields, simplistic"
),
"origami": ("origami style of {object} . paper art, pleated paper, folded, origami art, pleats, cut and fold, centered composition",
"noisy, sloppy, messy, grainy, highly detailed, ultra textured, photo"
),
"craft_clay": ("play-doh style of {object} . sculpture, clay art, centered composition, Claymation",
"sloppy, messy, grainy, highly detailed, ultra textured, photo"
),
"low_poly" : ("low-poly style of {object} . low-poly game art, polygon mesh, jagged, blocky, wireframe edges, centered composition",
"noisy, sloppy, messy, grainy, highly detailed, ultra textured, photo"
),
"Artstyle_watercolor": ("watercolor painting of {object} . vibrant, beautiful, painterly, detailed, textural, artistic",
"anime, photorealistic, 35mm film, deformed, glitch, low contrast, noisy"
),
"Papercraft_Collage" : ("collage style of {object} . mixed media, layered, textural, detailed, artistic",
"ugly, deformed, noisy, blurry, low contrast, realism, photorealistic"
),
"Artstyle_Impressionist" : ("impressionist painting of {object} . loose brushwork, vibrant color, light and shadow play, captures feeling over form",
"anime, photorealistic, 35mm film, deformed, glitch, low contrast, noisy"
),
"realistic_bg_black":("{object} photography, realistic, black background",
""),
"photography_realistic":("Photography of {object}, realistic",
""),
"digital_art":("{object} in digital glitch arts style.",
""
),
"chinese_painting":("{object} in traditional a chinese ink painting style.",
""
),
"no_style":("{object}",
""),
"kid_drawing":("{object} in kid crayon drawings style.",""),
"onepiece":("{object}, wanostyle, angry looking, straw hat, looking at viewer, solo, upper body, masterpiece, best quality, (extremely detailed), watercolor, illustration, depth of field, sketch, dark intense shadows, sharp focus, soft lighting, hdr, colorful, good composition, fire all around, spectacular, closed shirt",
" watermark, text, error, blurry, jpeg artifacts, many objects, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature")
}