|
|
|
from typing import Any, Dict, Optional
|
|
|
|
import torch
|
|
from einops import rearrange
|
|
|
|
from .attention import BasicTransformerBlock
|
|
from .attention import TemporalBasicTransformerBlock
|
|
|
|
|
|
def torch_dfs(model: torch.nn.Module):
|
|
result = [model]
|
|
for child in model.children():
|
|
result += torch_dfs(child)
|
|
return result
|
|
|
|
|
|
class ReferenceAttentionControl:
|
|
def __init__(
|
|
self,
|
|
unet,
|
|
mode="write",
|
|
do_classifier_free_guidance=False,
|
|
attention_auto_machine_weight=float("inf"),
|
|
gn_auto_machine_weight=1.0,
|
|
style_fidelity=1.0,
|
|
reference_attn=True,
|
|
reference_adain=False,
|
|
fusion_blocks="midup",
|
|
batch_size=1,
|
|
reference_attention_weight=1.,
|
|
audio_attention_weight=1.,
|
|
) -> None:
|
|
|
|
self.unet = unet
|
|
assert mode in ["read", "write"]
|
|
assert fusion_blocks in ["midup", "full"]
|
|
self.reference_attn = reference_attn
|
|
self.reference_adain = reference_adain
|
|
self.fusion_blocks = fusion_blocks
|
|
self.reference_attention_weight = reference_attention_weight
|
|
self.audio_attention_weight = audio_attention_weight
|
|
self.register_reference_hooks(
|
|
mode,
|
|
do_classifier_free_guidance,
|
|
attention_auto_machine_weight,
|
|
gn_auto_machine_weight,
|
|
style_fidelity,
|
|
reference_attn,
|
|
reference_adain,
|
|
fusion_blocks,
|
|
batch_size=batch_size,
|
|
)
|
|
|
|
def register_reference_hooks(
|
|
self,
|
|
mode,
|
|
do_classifier_free_guidance,
|
|
attention_auto_machine_weight,
|
|
gn_auto_machine_weight,
|
|
style_fidelity,
|
|
reference_attn,
|
|
reference_adain,
|
|
dtype=torch.float16,
|
|
batch_size=1,
|
|
num_images_per_prompt=1,
|
|
device=torch.device("cpu"),
|
|
fusion_blocks="midup",
|
|
):
|
|
MODE = mode
|
|
do_classifier_free_guidance = do_classifier_free_guidance
|
|
attention_auto_machine_weight = attention_auto_machine_weight
|
|
gn_auto_machine_weight = gn_auto_machine_weight
|
|
style_fidelity = style_fidelity
|
|
reference_attn = reference_attn
|
|
reference_adain = reference_adain
|
|
fusion_blocks = fusion_blocks
|
|
num_images_per_prompt = num_images_per_prompt
|
|
reference_attention_weight = self.reference_attention_weight
|
|
audio_attention_weight = self.audio_attention_weight
|
|
dtype = dtype
|
|
if do_classifier_free_guidance:
|
|
uc_mask = (
|
|
torch.Tensor(
|
|
[1] * batch_size * num_images_per_prompt * 16
|
|
+ [0] * batch_size * num_images_per_prompt * 16
|
|
)
|
|
.to(device)
|
|
.bool()
|
|
)
|
|
else:
|
|
uc_mask = (
|
|
torch.Tensor([0] * batch_size * num_images_per_prompt * 2)
|
|
.to(device)
|
|
.bool()
|
|
)
|
|
|
|
def hacked_basic_transformer_inner_forward(
|
|
self,
|
|
hidden_states: torch.FloatTensor,
|
|
attention_mask: Optional[torch.FloatTensor] = None,
|
|
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
|
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
|
timestep: Optional[torch.LongTensor] = None,
|
|
cross_attention_kwargs: Dict[str, Any] = None,
|
|
class_labels: Optional[torch.LongTensor] = None,
|
|
video_length=None,
|
|
):
|
|
if self.use_ada_layer_norm:
|
|
norm_hidden_states = self.norm1(hidden_states, timestep)
|
|
elif self.use_ada_layer_norm_zero:
|
|
(
|
|
norm_hidden_states,
|
|
gate_msa,
|
|
shift_mlp,
|
|
scale_mlp,
|
|
gate_mlp,
|
|
) = self.norm1(
|
|
hidden_states,
|
|
timestep,
|
|
class_labels,
|
|
hidden_dtype=hidden_states.dtype,
|
|
)
|
|
else:
|
|
norm_hidden_states = self.norm1(hidden_states)
|
|
|
|
|
|
|
|
cross_attention_kwargs = (
|
|
cross_attention_kwargs if cross_attention_kwargs is not None else {}
|
|
)
|
|
if self.only_cross_attention:
|
|
attn_output = self.attn1(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states
|
|
if self.only_cross_attention
|
|
else None,
|
|
attention_mask=attention_mask,
|
|
**cross_attention_kwargs,
|
|
)
|
|
else:
|
|
if MODE == "write":
|
|
attn_output = self.attn1(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states
|
|
if self.only_cross_attention
|
|
else None,
|
|
attention_mask=attention_mask,
|
|
**cross_attention_kwargs,
|
|
)
|
|
|
|
if self.use_ada_layer_norm_zero:
|
|
attn_output = gate_msa.unsqueeze(1) * attn_output
|
|
hidden_states = attn_output + hidden_states
|
|
|
|
if self.attn2 is not None:
|
|
norm_hidden_states = (
|
|
self.norm2(hidden_states, timestep)
|
|
if self.use_ada_layer_norm
|
|
else self.norm2(hidden_states)
|
|
)
|
|
self.bank.append(norm_hidden_states.clone())
|
|
|
|
|
|
attn_output = self.attn2(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
attention_mask=encoder_attention_mask,
|
|
**cross_attention_kwargs,
|
|
)
|
|
hidden_states = attn_output + hidden_states
|
|
|
|
if MODE == "read":
|
|
hidden_states = (
|
|
self.attn1(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=norm_hidden_states,
|
|
attention_mask=attention_mask,
|
|
)
|
|
+ hidden_states
|
|
)
|
|
|
|
if self.use_ada_layer_norm:
|
|
norm_hidden_states = self.norm1_5(hidden_states, timestep)
|
|
elif self.use_ada_layer_norm_zero:
|
|
(
|
|
norm_hidden_states,
|
|
gate_msa,
|
|
shift_mlp,
|
|
scale_mlp,
|
|
gate_mlp,
|
|
) = self.norm1_5(
|
|
hidden_states,
|
|
timestep,
|
|
class_labels,
|
|
hidden_dtype=hidden_states.dtype,
|
|
)
|
|
else:
|
|
norm_hidden_states = self.norm1_5(hidden_states)
|
|
|
|
bank_fea = []
|
|
for d in self.bank:
|
|
if len(d.shape) == 3:
|
|
d = d.unsqueeze(1).repeat(1, video_length, 1, 1)
|
|
bank_fea.append(rearrange(d, "b t l c -> (b t) l c"))
|
|
|
|
attn_hidden_states = self.attn1_5(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=bank_fea[0],
|
|
attention_mask=attention_mask,
|
|
)
|
|
|
|
if reference_attention_weight != 1.:
|
|
attn_hidden_states *= reference_attention_weight
|
|
|
|
hidden_states = (attn_hidden_states + hidden_states)
|
|
|
|
|
|
if self.attn2 is not None:
|
|
|
|
norm_hidden_states = (
|
|
self.norm2(hidden_states, timestep)
|
|
if self.use_ada_layer_norm
|
|
else self.norm2(hidden_states)
|
|
)
|
|
|
|
attn_hidden_states = self.attn2(
|
|
norm_hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
attention_mask=attention_mask,
|
|
)
|
|
|
|
if audio_attention_weight != 1.:
|
|
attn_hidden_states *= audio_attention_weight
|
|
|
|
hidden_states = (attn_hidden_states + hidden_states)
|
|
|
|
|
|
hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
|
|
|
|
|
|
if self.unet_use_temporal_attention:
|
|
d = hidden_states.shape[1]
|
|
hidden_states = rearrange(
|
|
hidden_states, "(b f) d c -> (b d) f c", f=video_length
|
|
)
|
|
norm_hidden_states = (
|
|
self.norm_temp(hidden_states, timestep)
|
|
if self.use_ada_layer_norm
|
|
else self.norm_temp(hidden_states)
|
|
)
|
|
hidden_states = (
|
|
self.attn_temp(norm_hidden_states) + hidden_states
|
|
)
|
|
hidden_states = rearrange(
|
|
hidden_states, "(b d) f c -> (b f) d c", d=d
|
|
)
|
|
|
|
return hidden_states
|
|
|
|
|
|
norm_hidden_states = self.norm3(hidden_states)
|
|
|
|
if self.use_ada_layer_norm_zero:
|
|
norm_hidden_states = (
|
|
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
|
)
|
|
|
|
ff_output = self.ff(norm_hidden_states)
|
|
|
|
if self.use_ada_layer_norm_zero:
|
|
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
|
|
|
hidden_states = ff_output + hidden_states
|
|
|
|
return hidden_states
|
|
|
|
if self.reference_attn:
|
|
if self.fusion_blocks == "midup":
|
|
attn_modules = [
|
|
module
|
|
for module in (
|
|
torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
|
|
)
|
|
if isinstance(module, BasicTransformerBlock)
|
|
or isinstance(module, TemporalBasicTransformerBlock)
|
|
]
|
|
elif self.fusion_blocks == "full":
|
|
attn_modules = [
|
|
module
|
|
for module in torch_dfs(self.unet)
|
|
if isinstance(module, BasicTransformerBlock)
|
|
or isinstance(module, TemporalBasicTransformerBlock)
|
|
]
|
|
attn_modules = sorted(
|
|
attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
|
)
|
|
|
|
for i, module in enumerate(attn_modules):
|
|
module._original_inner_forward = module.forward
|
|
if isinstance(module, BasicTransformerBlock):
|
|
module.forward = hacked_basic_transformer_inner_forward.__get__(
|
|
module, BasicTransformerBlock
|
|
)
|
|
if isinstance(module, TemporalBasicTransformerBlock):
|
|
module.forward = hacked_basic_transformer_inner_forward.__get__(
|
|
module, TemporalBasicTransformerBlock
|
|
)
|
|
|
|
module.bank = []
|
|
module.attn_weight = float(i) / float(len(attn_modules))
|
|
|
|
def update(
|
|
self,
|
|
writer,
|
|
do_classifier_free_guidance=True,
|
|
dtype=torch.float16,
|
|
):
|
|
if self.reference_attn:
|
|
if self.fusion_blocks == "midup":
|
|
reader_attn_modules = [
|
|
module
|
|
for module in (torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks))
|
|
if isinstance(module, TemporalBasicTransformerBlock)
|
|
]
|
|
writer_attn_modules = [
|
|
module
|
|
for module in (torch_dfs(writer.unet.mid_block) + torch_dfs(writer.unet.up_blocks))
|
|
if isinstance(module, BasicTransformerBlock)
|
|
]
|
|
elif self.fusion_blocks == "full":
|
|
reader_attn_modules = [
|
|
module
|
|
for module in torch_dfs(self.unet)
|
|
if isinstance(module, TemporalBasicTransformerBlock)
|
|
]
|
|
writer_attn_modules = [
|
|
module
|
|
for module in torch_dfs(writer.unet)
|
|
if isinstance(module, BasicTransformerBlock)
|
|
]
|
|
reader_attn_modules = sorted(
|
|
reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
|
)
|
|
writer_attn_modules = sorted(
|
|
writer_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
|
)
|
|
for r, w in zip(reader_attn_modules, writer_attn_modules):
|
|
if do_classifier_free_guidance:
|
|
r.bank = [torch.cat([torch.zeros_like(v), v]).to(dtype) for v in w.bank]
|
|
else:
|
|
r.bank = [v.clone().to(dtype) for v in w.bank]
|
|
|
|
def clear(self):
|
|
if self.reference_attn:
|
|
if self.fusion_blocks == "midup":
|
|
reader_attn_modules = [
|
|
module
|
|
for module in (
|
|
torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
|
|
)
|
|
if isinstance(module, BasicTransformerBlock)
|
|
or isinstance(module, TemporalBasicTransformerBlock)
|
|
]
|
|
elif self.fusion_blocks == "full":
|
|
reader_attn_modules = [
|
|
module
|
|
for module in torch_dfs(self.unet)
|
|
if isinstance(module, BasicTransformerBlock)
|
|
or isinstance(module, TemporalBasicTransformerBlock)
|
|
]
|
|
reader_attn_modules = sorted(
|
|
reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
|
)
|
|
for r in reader_attn_modules:
|
|
r.bank.clear()
|
|
|