attention.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
  2. from dataclasses import dataclass
  3. from typing import Optional
  4. import torch
  5. import torch.nn.functional as F
  6. from torch import nn
  7. from diffusers.configuration_utils import ConfigMixin, register_to_config
  8. from diffusers.modeling_utils import ModelMixin
  9. from diffusers.utils import BaseOutput
  10. from diffusers.utils.import_utils import is_xformers_available
  11. from diffusers.models.attention import CrossAttention, FeedForward, AdaLayerNorm
  12. from einops import rearrange, repeat
  13. import pdb
  14. @dataclass
  15. class Transformer3DModelOutput(BaseOutput):
  16. sample: torch.FloatTensor
  17. if is_xformers_available():
  18. import xformers
  19. import xformers.ops
  20. else:
  21. xformers = None
  22. class Transformer3DModel(ModelMixin, ConfigMixin):
  23. @register_to_config
  24. def __init__(
  25. self,
  26. num_attention_heads: int = 16,
  27. attention_head_dim: int = 88,
  28. in_channels: Optional[int] = None,
  29. num_layers: int = 1,
  30. dropout: float = 0.0,
  31. norm_num_groups: int = 32,
  32. cross_attention_dim: Optional[int] = None,
  33. attention_bias: bool = False,
  34. activation_fn: str = "geglu",
  35. num_embeds_ada_norm: Optional[int] = None,
  36. use_linear_projection: bool = False,
  37. only_cross_attention: bool = False,
  38. upcast_attention: bool = False,
  39. unet_use_cross_frame_attention=None,
  40. unet_use_temporal_attention=None,
  41. ):
  42. super().__init__()
  43. self.use_linear_projection = use_linear_projection
  44. self.num_attention_heads = num_attention_heads
  45. self.attention_head_dim = attention_head_dim
  46. inner_dim = num_attention_heads * attention_head_dim
  47. # Define input layers
  48. self.in_channels = in_channels
  49. self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
  50. if use_linear_projection:
  51. self.proj_in = nn.Linear(in_channels, inner_dim)
  52. else:
  53. self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
  54. # Define transformers blocks
  55. self.transformer_blocks = nn.ModuleList(
  56. [
  57. BasicTransformerBlock(
  58. inner_dim,
  59. num_attention_heads,
  60. attention_head_dim,
  61. dropout=dropout,
  62. cross_attention_dim=cross_attention_dim,
  63. activation_fn=activation_fn,
  64. num_embeds_ada_norm=num_embeds_ada_norm,
  65. attention_bias=attention_bias,
  66. only_cross_attention=only_cross_attention,
  67. upcast_attention=upcast_attention,
  68. unet_use_cross_frame_attention=unet_use_cross_frame_attention,
  69. unet_use_temporal_attention=unet_use_temporal_attention,
  70. )
  71. for d in range(num_layers)
  72. ]
  73. )
  74. # 4. Define output layers
  75. if use_linear_projection:
  76. self.proj_out = nn.Linear(in_channels, inner_dim)
  77. else:
  78. self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
  79. def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True):
  80. # Input
  81. assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
  82. video_length = hidden_states.shape[2]
  83. hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
  84. encoder_hidden_states = repeat(encoder_hidden_states, 'b n c -> (b f) n c', f=video_length)
  85. batch, channel, height, weight = hidden_states.shape
  86. residual = hidden_states
  87. hidden_states = self.norm(hidden_states)
  88. if not self.use_linear_projection:
  89. hidden_states = self.proj_in(hidden_states)
  90. inner_dim = hidden_states.shape[1]
  91. hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
  92. else:
  93. inner_dim = hidden_states.shape[1]
  94. hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim)
  95. hidden_states = self.proj_in(hidden_states)
  96. # Blocks
  97. for block in self.transformer_blocks:
  98. hidden_states = block(
  99. hidden_states,
  100. encoder_hidden_states=encoder_hidden_states,
  101. timestep=timestep,
  102. video_length=video_length
  103. )
  104. # Output
  105. if not self.use_linear_projection:
  106. hidden_states = (
  107. hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
  108. )
  109. hidden_states = self.proj_out(hidden_states)
  110. else:
  111. hidden_states = self.proj_out(hidden_states)
  112. hidden_states = (
  113. hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous()
  114. )
  115. output = hidden_states + residual
  116. output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
  117. if not return_dict:
  118. return (output,)
  119. return Transformer3DModelOutput(sample=output)
  120. class BasicTransformerBlock(nn.Module):
  121. def __init__(
  122. self,
  123. dim: int,
  124. num_attention_heads: int,
  125. attention_head_dim: int,
  126. dropout=0.0,
  127. cross_attention_dim: Optional[int] = None,
  128. activation_fn: str = "geglu",
  129. num_embeds_ada_norm: Optional[int] = None,
  130. attention_bias: bool = False,
  131. only_cross_attention: bool = False,
  132. upcast_attention: bool = False,
  133. unet_use_cross_frame_attention = None,
  134. unet_use_temporal_attention = None,
  135. ):
  136. super().__init__()
  137. self.only_cross_attention = only_cross_attention
  138. self.use_ada_layer_norm = num_embeds_ada_norm is not None
  139. self.unet_use_cross_frame_attention = unet_use_cross_frame_attention
  140. self.unet_use_temporal_attention = unet_use_temporal_attention
  141. # SC-Attn
  142. assert unet_use_cross_frame_attention is not None
  143. if unet_use_cross_frame_attention:
  144. self.attn1 = SparseCausalAttention2D(
  145. query_dim=dim,
  146. heads=num_attention_heads,
  147. dim_head=attention_head_dim,
  148. dropout=dropout,
  149. bias=attention_bias,
  150. cross_attention_dim=cross_attention_dim if only_cross_attention else None,
  151. upcast_attention=upcast_attention,
  152. )
  153. else:
  154. self.attn1 = CrossAttention(
  155. query_dim=dim,
  156. heads=num_attention_heads,
  157. dim_head=attention_head_dim,
  158. dropout=dropout,
  159. bias=attention_bias,
  160. upcast_attention=upcast_attention,
  161. )
  162. self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
  163. # Cross-Attn
  164. if cross_attention_dim is not None:
  165. self.attn2 = CrossAttention(
  166. query_dim=dim,
  167. cross_attention_dim=cross_attention_dim,
  168. heads=num_attention_heads,
  169. dim_head=attention_head_dim,
  170. dropout=dropout,
  171. bias=attention_bias,
  172. upcast_attention=upcast_attention,
  173. )
  174. else:
  175. self.attn2 = None
  176. if cross_attention_dim is not None:
  177. self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
  178. else:
  179. self.norm2 = None
  180. # Feed-forward
  181. self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
  182. self.norm3 = nn.LayerNorm(dim)
  183. # Temp-Attn
  184. assert unet_use_temporal_attention is not None
  185. if unet_use_temporal_attention:
  186. self.attn_temp = CrossAttention(
  187. query_dim=dim,
  188. heads=num_attention_heads,
  189. dim_head=attention_head_dim,
  190. dropout=dropout,
  191. bias=attention_bias,
  192. upcast_attention=upcast_attention,
  193. )
  194. nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
  195. self.norm_temp = AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
  196. def set_use_memory_efficient_attention_xformers(self, use_memory_efficient_attention_xformers: bool):
  197. if not is_xformers_available():
  198. print("Here is how to install it")
  199. raise ModuleNotFoundError(
  200. "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
  201. " xformers",
  202. name="xformers",
  203. )
  204. elif not torch.cuda.is_available():
  205. raise ValueError(
  206. "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is only"
  207. " available for GPU "
  208. )
  209. else:
  210. try:
  211. # Make sure we can run the memory efficient attention
  212. _ = xformers.ops.memory_efficient_attention(
  213. torch.randn((1, 2, 40), device="cuda"),
  214. torch.randn((1, 2, 40), device="cuda"),
  215. torch.randn((1, 2, 40), device="cuda"),
  216. )
  217. except Exception as e:
  218. raise e
  219. self.attn1._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
  220. if self.attn2 is not None:
  221. self.attn2._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
  222. # self.attn_temp._use_memory_efficient_attention_xformers = use_memory_efficient_attention_xformers
  223. def forward(self, hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None):
  224. # SparseCausal-Attention
  225. norm_hidden_states = (
  226. self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)
  227. )
  228. # if self.only_cross_attention:
  229. # hidden_states = (
  230. # self.attn1(norm_hidden_states, encoder_hidden_states, attention_mask=attention_mask) + hidden_states
  231. # )
  232. # else:
  233. # hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states
  234. # pdb.set_trace()
  235. if self.unet_use_cross_frame_attention:
  236. hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask, video_length=video_length) + hidden_states
  237. else:
  238. hidden_states = self.attn1(norm_hidden_states, attention_mask=attention_mask) + hidden_states
  239. if self.attn2 is not None:
  240. # Cross-Attention
  241. norm_hidden_states = (
  242. self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
  243. )
  244. hidden_states = (
  245. self.attn2(
  246. norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
  247. )
  248. + hidden_states
  249. )
  250. # Feed-forward
  251. hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
  252. # Temporal-Attention
  253. if self.unet_use_temporal_attention:
  254. d = hidden_states.shape[1]
  255. hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length)
  256. norm_hidden_states = (
  257. self.norm_temp(hidden_states, timestep) if self.use_ada_layer_norm else self.norm_temp(hidden_states)
  258. )
  259. hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
  260. hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
  261. return hidden_states