From 50514a5ab2262f2c791eef9abc905e142caeb2d6 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 15:19:24 -0700 Subject: [PATCH 1/6] Add support for transformer options and ControlNet --- lvdm/modules/networks/openaimodel3d.py | 124 +++++++++++++++++++------ 1 file changed, 98 insertions(+), 26 deletions(-) diff --git a/lvdm/modules/networks/openaimodel3d.py b/lvdm/modules/networks/openaimodel3d.py index 874e21e..cd4e1c6 100644 --- a/lvdm/modules/networks/openaimodel3d.py +++ b/lvdm/modules/networks/openaimodel3d.py @@ -15,6 +15,8 @@ from ...basics import ( ) from ...modules.attention import SpatialTransformer, TemporalTransformer import comfy.ops +import logging + ops = comfy.ops.disable_weight_init class TimestepBlock(nn.Module): @@ -28,20 +30,20 @@ class TimestepBlock(nn.Module): """ #This is needed because accelerate makes a copy of transformer_options which breaks "transformer_index" -def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None, time_context=None, num_video_frames=None, image_only_indicator=None): +def forward_timestep_embed(ts, x, emb, context=None, batch_size=None, transformer_options={}): for layer in ts: if isinstance(layer, TimestepBlock): - x = layer(x, emb) - elif isinstance(layer, SpatialVideoTransformer): - x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options) + x = layer(x, emb, batch_size=batch_size) + elif isinstance(layer, SpatialTransformer): + x = layer(x, context) if "transformer_index" in transformer_options: transformer_options["transformer_index"] += 1 elif isinstance(layer, TemporalTransformer): - x = layer(x, context, transformer_options) + x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size) + x = layer(x, context) if "transformer_index" in transformer_options: transformer_options["transformer_index"] += 1 - elif isinstance(layer, Upsample): - x = layer(x, output_shape=output_shape) + x = rearrange(x, 'b c f h w -> (b f) c h w') else: x = layer(x) return x @@ -52,19 +54,8 @@ class TimestepEmbedSequential(nn.Sequential, TimestepBlock): support it as an extra input. """ - def forward(self, x, emb, context=None, batch_size=None): - for layer in self: - if isinstance(layer, TimestepBlock): - x = layer(x, emb, batch_size=batch_size) - elif isinstance(layer, SpatialTransformer): - x = layer(x, context) - elif isinstance(layer, TemporalTransformer): - x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size) - x = layer(x, context) - x = rearrange(x, 'b c f h w -> (b f) c h w') - else: - x = layer(x) - return x + def forward(self, *args, **kwargs): + return forward_timestep_embed(self, *args, **kwargs) class Downsample(nn.Module): """ @@ -340,6 +331,32 @@ def context_processor(context, t, img_emb=None, temporal_size=16, concat_only=Fa return context +def apply_control(h, control, name, cond_idx=None): + if control is not None and name in control and len(control[name]) > 0: + frames = h.shape[0] + ctrl = control[name].pop() + if ctrl is not None: + try: + if cond_idx is not None and ctrl.shape[0] > frames: + ctrl_frames_list = list(range(ctrl.shape[0])) + ctrl_frames = len(ctrl_frames_list) + + idxs = ( + ctrl_frames_list[ctrl_frames // 2:] if cond_idx == 0 else \ + ctrl_frames_list[:ctrl_frames // 2] + ) + + ctrl = ctrl[idxs] + + h += ctrl + except Exception as e: + if h.shape != ctrl.shape: + logging.warning( + "warning control could not be applied {} {}".format(h.shape, ctrl.shape) + ) + logging.warning(e) + return h + class UNetModel(nn.Module): """ The full UNet model with attention and timestep embedding. @@ -677,13 +694,16 @@ class UNetModel(nn.Module): features_adapter=None, fs=None, img_emb=None, + control=None, transformer_options={}, + cond_idx=None, **kwargs ): if any([fs is None, img_emb is None, cc_concat is None]): raise ValueError("One or more of the required inputs for UNet Forward is None.") + cond_idx = transformer_options.get("cond_idx", None) transformer_options['original_shape'] = list(x.shape) transformer_options['transformer_index'] = 0 transformer_patches = transformer_options.get("patches", {}) @@ -726,21 +746,73 @@ class UNetModel(nn.Module): hs = [] for id, module in enumerate(self.input_blocks): - h = module(h, emb, context=context, batch_size=b) + transformer_options["block"] = ("input", id) + #h = module(h, emb, context=context, batch_size=b) + h = forward_timestep_embed( + module, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) + h = apply_control(h, control, 'input', cond_idx) + + if "input_block_patch" in transformer_patches: + patch = transformer_patches["input_block_patch"] + for p in patch: + h = p(h, transformer_options) + if id ==0 and self.addition_attention: - h = self.init_attn(h, emb, context=context, batch_size=b) + h = forward_timestep_embed( + self.init_attn, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) ## plug-in adapter features if ((id+1)%3 == 0) and features_adapter is not None: h = h + features_adapter[adapter_idx] adapter_idx += 1 hs.append(h) + if "input_block_patch_after_skip" in transformer_patches: + patch = transformer_patches["input_block_patch_after_skip"] + for p in patch: + h = p(h, transformer_options) if features_adapter is not None: assert len(features_adapter)==adapter_idx, 'Wrong features_adapter' + transformer_options["block"] = ("middle", 0) + h = forward_timestep_embed( + self.middle_block, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) + h = apply_control(h, control, 'middle', cond_idx) + for id, module in enumerate(self.output_blocks): + transformer_options["block"] = ("output", id) + hsp = hs.pop() + hsp = apply_control(hsp, control, 'output', cond_idx) - h = self.middle_block(h, emb, context=context, batch_size=b) - for module in self.output_blocks: - h = torch.cat([h, hs.pop()], dim=1) - h = module(h, emb, context=context, batch_size=b) + if "output_block_patch" in transformer_patches: + patch = transformer_patches["output_block_patch"] + for p in patch: + h, hsp = p(h, hsp, transformer_options) + + h = torch.cat([h, hsp], dim=1) + del hsp + h = forward_timestep_embed( + module, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) h = h.type(x.dtype) h = self.out(h) From 3f3763486bab8cb181739312823c32bf20b2f574 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 15:24:40 -0700 Subject: [PATCH 2/6] Add updated transformer and attention --- modules/attention.py | 676 +++++++++++++++++++ modules/encoders/condition.py | 389 +++++++++++ modules/encoders/resampler.py | 145 ++++ modules/networks/ae_modules.py | 1026 +++++++++++++++++++++++++++++ modules/networks/openaimodel3d.py | 822 +++++++++++++++++++++++ modules/x_transformer.py | 639 ++++++++++++++++++ 6 files changed, 3697 insertions(+) create mode 100644 modules/attention.py create mode 100644 modules/encoders/condition.py create mode 100644 modules/encoders/resampler.py create mode 100644 modules/networks/ae_modules.py create mode 100644 modules/networks/openaimodel3d.py create mode 100644 modules/x_transformer.py diff --git a/modules/attention.py b/modules/attention.py new file mode 100644 index 0000000..7537aa5 --- /dev/null +++ b/modules/attention.py @@ -0,0 +1,676 @@ +import torch +from torch import nn, einsum +import torch.nn.functional as F +from einops import rearrange, repeat +from functools import partial +from ..common import ( + checkpoint, + exists, + default, +) +from ..basics import zero_module +import comfy.ops +ops = comfy.ops.disable_weight_init +from comfy import model_management +from comfy.ldm.modules.attention import optimized_attention, optimized_attention_masked + +if model_management.xformers_enabled(): + import xformers + import xformers.ops + XFORMERS_IS_AVAILBLE = True +else: + XFORMERS_IS_AVAILBLE = False + +class RelativePosition(nn.Module): + """ https://github.com/evelinehong/Transformer_Relative_Position_PyTorch/blob/master/relative_position.py """ + + def __init__(self, num_units, max_relative_position): + super().__init__() + self.num_units = num_units + self.max_relative_position = max_relative_position + self.embeddings_table = nn.Parameter(torch.Tensor(max_relative_position * 2 + 1, num_units)) + nn.init.xavier_uniform_(self.embeddings_table) + + def forward(self, length_q, length_k): + device = self.embeddings_table.device + range_vec_q = torch.arange(length_q, device=device) + range_vec_k = torch.arange(length_k, device=device) + distance_mat = range_vec_k[None, :] - range_vec_q[:, None] + distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position) + final_mat = distance_mat_clipped + self.max_relative_position + final_mat = final_mat.long() + embeddings = self.embeddings_table[final_mat] + return embeddings + + +# TODO Add native Comfy optimized attention. +class CrossAttention(nn.Module): + + def __init__( + self, + query_dim, + context_dim=None, + heads=8, + dim_head=64, + dropout=0., + relative_position=False, + temporal_length=None, + video_length=None, + image_cross_attention=False, + image_cross_attention_scale=1.0, + image_cross_attention_scale_learnable=False, + text_context_len=77, + device=None, + dtype=None, + operations=ops + ): + super().__init__() + inner_dim = dim_head * heads + context_dim = default(context_dim, query_dim) + self.scale = dim_head**-0.5 + self.heads = heads + self.dim_head = dim_head + self.to_q = operations.Linear(query_dim, inner_dim, bias=False, device=device, dtype=dtype) + self.to_k = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) + self.to_v = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) + + self.to_out = nn.Sequential( + operations.Linear(inner_dim, query_dim, device=device, dtype=dtype), + nn.Dropout(dropout) + ) + + self.relative_position = relative_position + if self.relative_position: + assert(temporal_length is not None) + self.relative_position_k = RelativePosition(num_units=dim_head, max_relative_position=temporal_length) + self.relative_position_v = RelativePosition(num_units=dim_head, max_relative_position=temporal_length) + else: + ## only used for spatial attention, while NOT for temporal attention + if XFORMERS_IS_AVAILBLE and temporal_length is None: + self.forward = self.efficient_forward + else: + self.forward = self.comfy_efficient_forward + + self.video_length = video_length + self.image_cross_attention = image_cross_attention + self.image_cross_attention_scale = image_cross_attention_scale + self.text_context_len = text_context_len + self.image_cross_attention_scale_learnable = image_cross_attention_scale_learnable + if self.image_cross_attention: + self.to_k_ip = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) + self.to_v_ip = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) + if image_cross_attention_scale_learnable: + self.register_parameter('alpha', nn.Parameter(torch.tensor(0.)) ) + + def comfy_efficient_forward(self, x, context=None, mask=None): + spatial_self_attn = (context is None) + k_ip, v_ip, out_ip = None, None, None + + h = self.heads + q = self.to_q(x) + context = default(context, x) + + if self.image_cross_attention and not spatial_self_attn: + context, context_image = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:] + k = self.to_k(context) + v = self.to_v(context) + k_ip = self.to_k_ip(context_image) + v_ip = self.to_v_ip(context_image) + else: + if not spatial_self_attn: + context = context[:,:self.text_context_len,:] + k = self.to_k(context) + v = self.to_v(context) + + out = optimized_attention(q, k, v, h) + + if exists(mask): + ## feasible for causal attention mask only + out = optimized_attention_masked(q, k, v, h) + + ## for image cross-attention + if k_ip is not None: + q = rearrange(q, 'b n (h d) -> (b h) n d', h=h) + k_ip, v_ip = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (k_ip, v_ip)) + sim_ip = torch.einsum('b i d, b j d -> b i j', q, k_ip) * self.scale + del k_ip + sim_ip = sim_ip.softmax(dim=-1) + out_ip = torch.einsum('b i j, b j d -> b i d', sim_ip, v_ip) + out_ip = rearrange(out_ip, '(b h) n d -> b n (h d)', h=h) + + if out_ip is not None: + if self.image_cross_attention_scale_learnable: + out = out + self.image_cross_attention_scale * out_ip * (torch.tanh(self.alpha)+1) + else: + out = out + self.image_cross_attention_scale * out_ip + + return self.to_out(out) + + def forward(self, x, context=None, mask=None): + spatial_self_attn = (context is None) + k_ip, v_ip, out_ip = None, None, None + + h = self.heads + q = self.to_q(x) + context = default(context, x) + + if self.image_cross_attention and not spatial_self_attn: + context, context_image = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:] + k = self.to_k(context) + v = self.to_v(context) + k_ip = self.to_k_ip(context_image) + v_ip = self.to_v_ip(context_image) + else: + + # Assumed Spatial Attention (b c h w) + if not spatial_self_attn: + context = context[:,:self.text_context_len,:] + k = self.to_k(context) + v = self.to_v(context) + + + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) + + sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale + if self.relative_position: + len_q, len_k, len_v = q.shape[1], k.shape[1], v.shape[1] + k2 = self.relative_position_k(len_q, len_k) + sim2 = einsum('b t d, t s d -> b t s', q, k2) * self.scale # TODO check + sim += sim2 + del k + + if exists(mask): + ## feasible for causal attention mask only + max_neg_value = -torch.finfo(sim.dtype).max + mask = repeat(mask, 'b i j -> (b h) i j', h=h) + sim.masked_fill_(~(mask>0.5), max_neg_value) + + # attention, what we cannot get enough of + sim = sim.softmax(dim=-1) + + out = torch.einsum('b i j, b j d -> b i d', sim, v) + if self.relative_position: + v2 = self.relative_position_v(len_q, len_v) + out2 = einsum('b t s, t s d -> b t d', sim, v2) # TODO check + out += out2 + out = rearrange(out, '(b h) n d -> b n (h d)', h=h) + + + ## for image cross-attention + if k_ip is not None: + k_ip, v_ip = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (k_ip, v_ip)) + sim_ip = torch.einsum('b i d, b j d -> b i j', q, k_ip) * self.scale + del k_ip + sim_ip = sim_ip.softmax(dim=-1) + out_ip = torch.einsum('b i j, b j d -> b i d', sim_ip, v_ip) + out_ip = rearrange(out_ip, '(b h) n d -> b n (h d)', h=h) + + + if out_ip is not None: + if self.image_cross_attention_scale_learnable: + out = out + self.image_cross_attention_scale * out_ip * (torch.tanh(self.alpha)+1) + else: + out = out + self.image_cross_attention_scale * out_ip + + return self.to_out(out) + + def efficient_forward(self, x, context=None, mask=None): + spatial_self_attn = (context is None) + k_ip, v_ip, out_ip = None, None, None + + q = self.to_q(x) + context = default(context, x) + + if self.image_cross_attention and not spatial_self_attn: + context, context_image = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:] + k = self.to_k(context) + v = self.to_v(context) + k_ip = self.to_k_ip(context_image) + v_ip = self.to_v_ip(context_image) + else: + if not spatial_self_attn: + context = context[:,:self.text_context_len,:] + k = self.to_k(context) + v = self.to_v(context) + + b, _, _ = q.shape + q, k, v = map( + lambda t: t.unsqueeze(3) + .reshape(b, t.shape[1], self.heads, self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b * self.heads, t.shape[1], self.dim_head) + .contiguous(), + (q, k, v), + ) + # actually compute the attention, what we cannot get enough of + out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=None) + + ## for image cross-attention + if k_ip is not None: + k_ip, v_ip = map( + lambda t: t.unsqueeze(3) + .reshape(b, t.shape[1], self.heads, self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b * self.heads, t.shape[1], self.dim_head) + .contiguous(), + (k_ip, v_ip), + ) + out_ip = xformers.ops.memory_efficient_attention(q, k_ip, v_ip, attn_bias=None, op=None) + out_ip = ( + out_ip.unsqueeze(0) + .reshape(b, self.heads, out.shape[1], self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b, out.shape[1], self.heads * self.dim_head) + ) + + if exists(mask): + raise NotImplementedError + out = ( + out.unsqueeze(0) + .reshape(b, self.heads, out.shape[1], self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b, out.shape[1], self.heads * self.dim_head) + ) + if out_ip is not None: + if self.image_cross_attention_scale_learnable: + out = out + self.image_cross_attention_scale * out_ip * (torch.tanh(self.alpha)+1) + else: + out = out + self.image_cross_attention_scale * out_ip + + return self.to_out(out) + + +class BasicTransformerBlock(nn.Module): + + def __init__( + self, + dim, + n_heads, + d_head, + dropout=0., + context_dim=None, + gated_ff=True, + checkpoint=True, + disable_self_attn=False, + attention_cls=None, + video_length=None, + image_cross_attention=False, + image_cross_attention_scale=1.0, + image_cross_attention_scale_learnable=False, + text_context_len=77, + device=None, + dtype=None, + operations=ops + ): + super().__init__() + attn_cls = CrossAttention if attention_cls is None else attention_cls + self.disable_self_attn = disable_self_attn + self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, + context_dim=None, device=device, dtype=dtype if self.disable_self_attn else None) + self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff, device=device, dtype=dtype) + self.attn2 = attn_cls( + query_dim=dim, + context_dim=context_dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + video_length=video_length, + image_cross_attention=image_cross_attention, + image_cross_attention_scale=image_cross_attention_scale, + image_cross_attention_scale_learnable=image_cross_attention_scale_learnable, + text_context_len=text_context_len, + device=device, + dtype=dtype + ) + self.image_cross_attention = image_cross_attention + + self.norm1 = operations.LayerNorm(dim, device=device, dtype=dtype) + self.norm2 = operations.LayerNorm(dim, device=device, dtype=dtype) + self.norm3 = operations.LayerNorm(dim, device=device, dtype=dtype) + self.checkpoint = checkpoint + + def forward(self, x, context=None, mask=None, **kwargs): + ## implementation tricks: because checkpointing doesn't support non-tensor (e.g. None or scalar) arguments + input_tuple = (x,) ## should not be (x), otherwise *input_tuple will decouple x into multiple arguments + if context is not None: + input_tuple = (x, context) + if mask is not None: + forward_mask = partial(self._forward, mask=mask) + return checkpoint(forward_mask, (x,), self.parameters(), self.checkpoint) + return checkpoint(self._forward, input_tuple, self.parameters(), self.checkpoint) + + + def _forward(self, x, context=None, mask=None): + x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None, mask=mask) + x + x = self.attn2(self.norm2(x), context=context, mask=mask) + x + x = self.ff(self.norm3(x)) + x + return x + + +class SpatialTransformer(nn.Module): + """ + Transformer block for image-like data in spatial axis. + First, project the input (aka embedding) + and reshape to b, t, d. + Then apply standard transformer action. + Finally, reshape to image + NEW: use_linear for more efficiency instead of the 1x1 convs + """ + + def __init__( + self, + in_channels, + n_heads, + d_head, + depth=1, + dropout=0., + context_dim=None, + use_checkpoint=True, + disable_self_attn=False, + use_linear=False, + video_length=None, + image_cross_attention=False, + image_cross_attention_scale_learnable=False, + device=None, + dtype=None, + operations=ops + ): + super().__init__() + self.in_channels = in_channels + inner_dim = n_heads * d_head + self.norm = operations.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, device=device, dtype=dtype) + if not use_linear: + self.proj_in = opeations.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0, device=device, dtype=dtype) + else: + self.proj_in = operations.Linear(in_channels, inner_dim, device=device, dtype=dtype) + + attention_cls = None + self.transformer_blocks = nn.ModuleList([ + BasicTransformerBlock( + inner_dim, + n_heads, + d_head, + dropout=dropout, + context_dim=context_dim, + disable_self_attn=disable_self_attn, + checkpoint=use_checkpoint, + attention_cls=attention_cls, + video_length=video_length, + image_cross_attention=image_cross_attention, + image_cross_attention_scale_learnable=image_cross_attention_scale_learnable, + device=device, + dtype=dtype + ) for d in range(depth) + ]) + if not use_linear: + self.proj_out = zero_module(operations.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0, device=device, dtype=dtype)) + else: + self.proj_out = zero_module(operations.Linear(inner_dim, in_channels, device=device, dtype=dtype)) + self.use_linear = use_linear + + def forward(self, x, context=None, **kwargs): + b, c, h, w = x.shape + x_in = x + x = self.norm(x) + if not self.use_linear: + x = self.proj_in(x) + x = rearrange(x, 'b c h w -> b (h w) c').contiguous() + if self.use_linear: + x = self.proj_in(x) + for i, block in enumerate(self.transformer_blocks): + x = block(x, context=context, **kwargs) + if self.use_linear: + x = self.proj_out(x) + x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous() + if not self.use_linear: + x = self.proj_out(x) + return x + x_in + + +class TemporalTransformer(nn.Module): + """ + Transformer block for image-like data in temporal axis. + First, reshape to b, t, d. + Then apply standard transformer action. + Finally, reshape to image + """ + def __init__( + self, + in_channels, + n_heads, + d_head, + depth=1, + dropout=0., + context_dim=None, + use_checkpoint=True, + use_linear=False, + only_self_att=True, + causal_attention=False, + causal_block_size=1, + relative_position=False, + temporal_length=None, + device=None, + dtype=None, + operations=ops + ): + super().__init__() + self.only_self_att = only_self_att + self.relative_position = relative_position + self.causal_attention = causal_attention + self.causal_block_size = causal_block_size + + if only_self_att: + context_dim = None + + self.in_channels = in_channels + inner_dim = n_heads * d_head + self.norm = operations.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, device=device, dtype=dtype) + self.proj_in = nn.Conv1d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0).to(device, dtype) + if not use_linear: + self.proj_in = nn.Conv1d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0).to(device, dtype) + else: + self.proj_in = operations.Linear(in_channels, inner_dim, device=device, dtype=dtype) + + if relative_position: + assert(temporal_length is not None) + attention_cls = partial(CrossAttention, relative_position=True, temporal_length=temporal_length, device=device, dtype=dtype) + else: + attention_cls = partial(CrossAttention, temporal_length=temporal_length, device=device, dtype=dtype) + if self.causal_attention: + assert(temporal_length is not None) + self.mask = torch.tril(torch.ones([1, temporal_length, temporal_length])) + + if self.only_self_att: + context_dim = None + self.transformer_blocks = nn.ModuleList([ + BasicTransformerBlock( + inner_dim, + n_heads, + d_head, + dropout=dropout, + context_dim=context_dim, + attention_cls=attention_cls, + checkpoint=use_checkpoint, + device=device, + dtype=dtype + ) for d in range(depth) + ]) + if not use_linear: + self.proj_out = zero_module(nn.Conv1d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0).to(device, dtype)) + else: + self.proj_out = zero_module(operations.Linear(inner_dim, in_channels, device=device, dtype=dtype)) + self.use_linear = use_linear + + def forward(self, x, context=None): + b, c, t, h, w = x.shape + x_in = x + x = self.norm(x) + x = rearrange(x, 'b c t h w -> (b h w) c t').contiguous() + if not self.use_linear: + x = self.proj_in(x) + x = rearrange(x, 'bhw c t -> bhw t c').contiguous() + if self.use_linear: + x = self.proj_in(x) + + temp_mask = None + if self.causal_attention: + # slice the from mask map + temp_mask = self.mask[:,:t,:t].to(x.device) + + if temp_mask is not None: + mask = temp_mask.to(x.device) + mask = repeat(mask, 'l i j -> (l bhw) i j', bhw=b*h*w) + else: + mask = None + + if self.only_self_att: + ## note: if no context is given, cross-attention defaults to self-attention + for i, block in enumerate(self.transformer_blocks): + x = block(x, mask=mask) + x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous() + else: + x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous() + context = rearrange(context, '(b t) l con -> b t l con', t=t).contiguous() + for i, block in enumerate(self.transformer_blocks): + # calculate each batch one by one (since number in shape could not greater then 65,535 for some package) + for j in range(b): + context_j = repeat( + context[j], + 't l con -> (t r) l con', r=(h * w) // t, t=t).contiguous() + ## note: causal mask will not applied in cross-attention case + x[j] = block(x[j], context=context_j) + + if self.use_linear: + x = self.proj_out(x) + x = rearrange(x, 'b (h w) t c -> b c t h w', h=h, w=w).contiguous() + if not self.use_linear: + x = rearrange(x, 'b hw t c -> (b hw) c t').contiguous() + x = self.proj_out(x) + x = rearrange(x, '(b h w) c t -> b c t h w', b=b, h=h, w=w).contiguous() + + return x + x_in + + +class GEGLU(nn.Module): + def __init__(self, dim_in, dim_out, device=None, dtype=None, operations=ops): + super().__init__() + self.proj = operations.Linear(dim_in, dim_out * 2, device=device, dtype=dtype) + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0., device=None, dtype=None, operations=ops): + super().__init__() + inner_dim = int(dim * mult) + dim_out = default(dim_out, dim) + project_in = nn.Sequential( + operations.Linear(dim, inner_dim, device=device, dtype=dtype), + nn.GELU() + ) if not glu else GEGLU(dim, inner_dim) + + self.net = nn.Sequential( + project_in, + nn.Dropout(dropout), + operations.Linear(inner_dim, dim_out, device=device, dtype=dtype) + ) + + def forward(self, x): + return self.net(x) + + +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32, device=None, dtype=None, operations=ops): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = operations.Conv2d(dim, hidden_dim * 3, 1, bias = False, device=device, dtype=dtype) + self.to_out = operations.Conv2d(hidden_dim, dim, 1, device=device, dtype=dtype) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3) + k = k.softmax(dim=-1) + context = torch.einsum('bhdn,bhen->bhde', k, v) + out = torch.einsum('bhde,bhdn->bhen', context, q) + out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) + return self.to_out(out) + + +class SpatialSelfAttention(nn.Module): + def __init__(self, in_channels, device=None, dtype=None, operations=ops): + super().__init__() + self.in_channels = in_channels + + self.norm = operations.GroupNorm( + num_groups=32, + num_channels=in_channels, + eps=1e-6, + affine=True, + device=device, + dtype=dtype + ) + self.q = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + self.k = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + self.v = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + self.proj_out = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b,c,h,w = q.shape + q = rearrange(q, 'b c h w -> b (h w) c') + k = rearrange(k, 'b c h w -> b c (h w)') + w_ = torch.einsum('bij,bjk->bik', q, k) + + w_ = w_ * (int(c)**(-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = rearrange(v, 'b c h w -> b c (h w)') + w_ = rearrange(w_, 'b i j -> b j i') + h_ = torch.einsum('bij,bjk->bik', v, w_) + h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h) + h_ = self.proj_out(h_) + + return x+h_ diff --git a/modules/encoders/condition.py b/modules/encoders/condition.py new file mode 100644 index 0000000..610322b --- /dev/null +++ b/modules/encoders/condition.py @@ -0,0 +1,389 @@ +import torch +import torch.nn as nn +import kornia +import open_clip +from torch.utils.checkpoint import checkpoint +from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel +from ..common import autocast +from utils.utils import count_params + + +class AbstractEncoder(nn.Module): + def __init__(self): + super().__init__() + + def encode(self, *args, **kwargs): + raise NotImplementedError + + +class IdentityEncoder(AbstractEncoder): + def encode(self, x): + return x + + +class ClassEmbedder(nn.Module): + def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1): + super().__init__() + self.key = key + self.embedding = nn.Embedding(n_classes, embed_dim) + self.n_classes = n_classes + self.ucg_rate = ucg_rate + + def forward(self, batch, key=None, disable_dropout=False): + if key is None: + key = self.key + # this is for use in crossattn + c = batch[key][:, None] + if self.ucg_rate > 0. and not disable_dropout: + mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate) + c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1) + c = c.long() + c = self.embedding(c) + return c + + def get_unconditional_conditioning(self, bs, device="cuda"): + uc_class = self.n_classes - 1 # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000) + uc = torch.ones((bs,), device=device) * uc_class + uc = {self.key: uc} + return uc + + +def disabled_train(self, mode=True): + """Overwrite model.train with this function to make sure train/eval mode + does not change anymore.""" + return self + + +class FrozenT5Embedder(AbstractEncoder): + """Uses the T5 transformer encoder for text""" + + def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77, + freeze=True): # others are google/t5-v1_1-xl and google/t5-v1_1-xxl + super().__init__() + self.tokenizer = T5Tokenizer.from_pretrained(version) + self.transformer = T5EncoderModel.from_pretrained(version) + self.device = device + self.max_length = max_length # TODO: typical value? + if freeze: + self.freeze() + + def freeze(self): + self.transformer = self.transformer.eval() + # self.train = disabled_train + for param in self.parameters(): + param.requires_grad = False + + def forward(self, text): + batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True, + return_overflowing_tokens=False, padding="max_length", return_tensors="pt") + tokens = batch_encoding["input_ids"].to(self.device) + outputs = self.transformer(input_ids=tokens) + + z = outputs.last_hidden_state + return z + + def encode(self, text): + return self(text) + + +class FrozenCLIPEmbedder(AbstractEncoder): + """Uses the CLIP transformer encoder for text (from huggingface)""" + LAYERS = [ + "last", + "pooled", + "hidden" + ] + + def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77, + freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32 + super().__init__() + assert layer in self.LAYERS + self.tokenizer = CLIPTokenizer.from_pretrained(version) + self.transformer = CLIPTextModel.from_pretrained(version) + self.device = device + self.max_length = max_length + if freeze: + self.freeze() + self.layer = layer + self.layer_idx = layer_idx + if layer == "hidden": + assert layer_idx is not None + assert 0 <= abs(layer_idx) <= 12 + + def freeze(self): + self.transformer = self.transformer.eval() + # self.train = disabled_train + for param in self.parameters(): + param.requires_grad = False + + def forward(self, text): + batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True, + return_overflowing_tokens=False, padding="max_length", return_tensors="pt") + tokens = batch_encoding["input_ids"].to(self.device) + outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden") + if self.layer == "last": + z = outputs.last_hidden_state + elif self.layer == "pooled": + z = outputs.pooler_output[:, None, :] + else: + z = outputs.hidden_states[self.layer_idx] + return z + + def encode(self, text): + return self(text) + + +class ClipImageEmbedder(nn.Module): + def __init__( + self, + model, + jit=False, + device='cuda' if torch.cuda.is_available() else 'cpu', + antialias=True, + ucg_rate=0. + ): + super().__init__() + from clip import load as load_clip + self.model, _ = load_clip(name=model, device=device, jit=jit) + + self.antialias = antialias + + self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) + self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) + self.ucg_rate = ucg_rate + + def preprocess(self, x): + # normalize to [0,1] + x = kornia.geometry.resize(x, (224, 224), + interpolation='bicubic', align_corners=True, + antialias=self.antialias) + x = (x + 1.) / 2. + # re-normalize according to clip + x = kornia.enhance.normalize(x, self.mean, self.std) + return x + + def forward(self, x, no_dropout=False): + # x is assumed to be in range [-1,1] + out = self.model.encode_image(self.preprocess(x)) + out = out.to(x.dtype) + if self.ucg_rate > 0. and not no_dropout: + out = torch.bernoulli((1. - self.ucg_rate) * torch.ones(out.shape[0], device=out.device))[:, None] * out + return out + + +class FrozenOpenCLIPEmbedder(AbstractEncoder): + """ + Uses the OpenCLIP transformer encoder for text + """ + LAYERS = [ + # "pooled", + "last", + "penultimate" + ] + + def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, + freeze=True, layer="last"): + super().__init__() + assert layer in self.LAYERS + model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version) + del model.visual + self.model = model + + self.device = device + self.max_length = max_length + if freeze: + self.freeze() + self.layer = layer + if self.layer == "last": + self.layer_idx = 0 + elif self.layer == "penultimate": + self.layer_idx = 1 + else: + raise NotImplementedError() + + def freeze(self): + self.model = self.model.eval() + for param in self.parameters(): + param.requires_grad = False + + def forward(self, text): + tokens = open_clip.tokenize(text) ## all clip models use 77 as context length + z = self.encode_with_transformer(tokens.to(self.device)) + return z + + def encode_with_transformer(self, text): + x = self.model.token_embedding(text) # [batch_size, n_ctx, d_model] + x = x + self.model.positional_embedding + x = x.permute(1, 0, 2) # NLD -> LND + x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask) + x = x.permute(1, 0, 2) # LND -> NLD + x = self.model.ln_final(x) + return x + + def text_transformer_forward(self, x: torch.Tensor, attn_mask=None): + for i, r in enumerate(self.model.transformer.resblocks): + if i == len(self.model.transformer.resblocks) - self.layer_idx: + break + if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(): + x = checkpoint(r, x, attn_mask) + else: + x = r(x, attn_mask=attn_mask) + return x + + def encode(self, text): + return self(text) + + +class FrozenOpenCLIPImageEmbedder(AbstractEncoder): + """ + Uses the OpenCLIP vision transformer encoder for images + """ + + def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, + freeze=True, layer="pooled", antialias=True, ucg_rate=0.): + super().__init__() + model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), + pretrained=version, ) + del model.transformer + self.model = model + # self.mapper = torch.nn.Linear(1280, 1024) + self.device = device + self.max_length = max_length + if freeze: + self.freeze() + self.layer = layer + if self.layer == "penultimate": + raise NotImplementedError() + self.layer_idx = 1 + + self.antialias = antialias + + self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) + self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) + self.ucg_rate = ucg_rate + + def preprocess(self, x): + # normalize to [0,1] + x = kornia.geometry.resize(x, (224, 224), + interpolation='bicubic', align_corners=True, + antialias=self.antialias) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, self.mean, self.std) + return x + + def freeze(self): + self.model = self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + @autocast + def forward(self, image, no_dropout=False): + z = self.encode_with_vision_transformer(image) + if self.ucg_rate > 0. and not no_dropout: + z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z + return z + + def encode_with_vision_transformer(self, img): + img = self.preprocess(img) + x = self.model.visual(img) + return x + + def encode(self, text): + return self(text) + +class FrozenOpenCLIPImageEmbedderV2(AbstractEncoder): + """ + Uses the OpenCLIP vision transformer encoder for images + """ + + def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", + freeze=True, layer="pooled", antialias=True): + super().__init__() + model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), + pretrained=version, ) + del model.transformer + self.model = model + self.device = device + + if freeze: + self.freeze() + self.layer = layer + if self.layer == "penultimate": + raise NotImplementedError() + self.layer_idx = 1 + + self.antialias = antialias + + self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) + self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) + + + def preprocess(self, x): + # normalize to [0,1] + x = kornia.geometry.resize(x, (224, 224), + interpolation='bicubic', align_corners=True, + antialias=self.antialias) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, self.mean, self.std) + return x + + def freeze(self): + self.model = self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + + def forward(self, image, no_dropout=False): + ## image: b c h w + z = self.encode_with_vision_transformer(image) + return z + + def encode_with_vision_transformer(self, x): + x = self.preprocess(x) + + # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1 + if self.model.visual.input_patchnorm: + # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)') + x = x.reshape(x.shape[0], x.shape[1], self.model.visual.grid_size[0], self.model.visual.patch_size[0], self.model.visual.grid_size[1], self.model.visual.patch_size[1]) + x = x.permute(0, 2, 4, 1, 3, 5) + x = x.reshape(x.shape[0], self.model.visual.grid_size[0] * self.model.visual.grid_size[1], -1) + x = self.model.visual.patchnorm_pre_ln(x) + x = self.model.visual.conv1(x) + else: + x = self.model.visual.conv1(x) # shape = [*, width, grid, grid] + x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] + x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] + + # class embeddings and positional embeddings + x = torch.cat( + [self.model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), + x], dim=1) # shape = [*, grid ** 2 + 1, width] + x = x + self.model.visual.positional_embedding.to(x.dtype) + + # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in + x = self.model.visual.patch_dropout(x) + x = self.model.visual.ln_pre(x) + + x = x.permute(1, 0, 2) # NLD -> LND + x = self.model.visual.transformer(x) + x = x.permute(1, 0, 2) # LND -> NLD + + return x + +class FrozenCLIPT5Encoder(AbstractEncoder): + def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cuda", + clip_max_length=77, t5_max_length=77): + super().__init__() + self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length) + self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length) + print(f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder) * 1.e-6:.2f} M parameters, " + f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder) * 1.e-6:.2f} M params.") + + def encode(self, text): + return self(text) + + def forward(self, text): + clip_z = self.clip_encoder.encode(text) + t5_z = self.t5_encoder.encode(text) + return [clip_z, t5_z] diff --git a/modules/encoders/resampler.py b/modules/encoders/resampler.py new file mode 100644 index 0000000..d9abd23 --- /dev/null +++ b/modules/encoders/resampler.py @@ -0,0 +1,145 @@ +# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +# and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py +# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/resampler.py +import math +import torch +import torch.nn as nn + + +class ImageProjModel(nn.Module): + """Projection Model""" + def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): + super().__init__() + self.cross_attention_dim = cross_attention_dim + self.clip_extra_context_tokens = clip_extra_context_tokens + self.proj = nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim) + self.norm = nn.LayerNorm(cross_attention_dim) + + def forward(self, image_embeds): + #embeds = image_embeds + embeds = image_embeds.type(list(self.proj.parameters())[0].dtype) + clip_extra_context_tokens = self.proj(embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim) + clip_extra_context_tokens = self.norm(clip_extra_context_tokens) + return clip_extra_context_tokens + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class Resampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + video_length=None, # using frame-wise version or not + ): + super().__init__() + ## queries for a single frame / image + self.num_queries = num_queries + self.video_length = video_length + + ## queries for each frame + if video_length is not None: + num_queries = num_queries * video_length + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + self.proj_in = nn.Linear(embedding_dim, dim) + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + latents = self.latents.repeat(x.size(0), 1, 1) ## B (T L) C + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + latents = self.norm_out(latents) # B L C or B (T L) C + + return latents \ No newline at end of file diff --git a/modules/networks/ae_modules.py b/modules/networks/ae_modules.py new file mode 100644 index 0000000..35b6817 --- /dev/null +++ b/modules/networks/ae_modules.py @@ -0,0 +1,1026 @@ +# pytorch_diffusion + derived encoder decoder +import math +import torch +import numpy as np +import torch.nn as nn +from einops import rearrange +from utils.utils import instantiate_from_config +from ...modules.attention import LinearAttention + +import comfy.ops +ops = comfy.ops.disable_weight_init + +def nonlinearity(x): + # swish + return x*torch.sigmoid(x) + + +def Normalize(in_channels, num_groups=32, device=None, dtype=None, operations=ops): + return operations.GroupNorm( + num_groups=num_groups, + num_channels=in_channels, + eps=1e-6, + affine=True, + device=device, + dtype=dtype + ) + +class LinAttnBlock(LinearAttention): + """to match AttnBlock usage""" + def __init__(self, in_channels, device=None, dtype=None): + super().__init__(dim=in_channels, heads=1, dim_head=in_channels, device=device, dtype=dtype) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels, device=None, dtype=None, operations=ops): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels, device=device, dtype=dtype) + self.q = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + self.k = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + self.v = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + self.proj_out = operations.Conv2d( + in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b,c,h,w = q.shape + q = q.reshape(b,c,h*w) # bcl + q = q.permute(0,2,1) # bcl -> blc l=hw + k = k.reshape(b,c,h*w) # bcl + + w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c)**(-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b,c,h*w) + w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v,w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b,c,h,w) + + h_ = self.proj_out(h_) + + return x+h_ + +def make_attn(in_channels, attn_type="vanilla", device=None, dtype=None): + assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown' + #print(f"making attention of type '{attn_type}' with {in_channels} in_channels") + if attn_type == "vanilla": + return AttnBlock(in_channels, device=device, dtype=dtype) + elif attn_type == "none": + return nn.Identity(in_channels) + else: + return LinAttnBlock(in_channels, device=device, dtype=dtype) + +class Downsample(nn.Module): + def __init__(self, in_channels, with_conv, device=None, dtype=None, operations=ops): + super().__init__() + self.with_conv = with_conv + self.in_channels = in_channels + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = operations.Conv2d( + in_channels, + in_channels, + kernel_size=3, + stride=2, + padding=0, + device=device, + dtype=dtype + ) + def forward(self, x): + if self.with_conv: + pad = (0,1,0,1) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + return x + +class Upsample(nn.Module): + def __init__(self, in_channels, with_conv, device=None, dtype=None, operations=ops): + super().__init__() + self.with_conv = with_conv + self.in_channels = in_channels + if self.with_conv: + self.conv = operations.Conv2d( + in_channels, + in_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype + ) + + def forward(self, x): + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + return x + +def get_timestep_embedding(timesteps, embedding_dim): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: + From Fairseq. + Build sinusoidal embeddings. + This matches the implementation in tensor2tensor, but differs slightly + from the description in Section 3.5 of "Attention Is All You Need". + """ + assert len(timesteps.shape) == 1 + + half_dim = embedding_dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) + emb = emb.to(device=timesteps.device) + emb = timesteps.float()[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: # zero pad + emb = torch.nn.functional.pad(emb, (0,1,0,0)) + return emb + + + +class ResnetBlock(nn.Module): + def __init__( + self, + *, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout, + temb_channels=512, + device=None, + dtype=None, + operations=ops + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels, device=device, dtype=dtype) + self.conv1 = operations.Conv2d( + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype + ) + if temb_channels > 0: + self.temb_proj = operations.Linear( + temb_channels, + out_channels, + device=device, + dtype=dtype + ) + self.norm2 = Normalize(out_channels, device=device, dtype=dtype) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = operations.Conv2d( + out_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype + ) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = operations.Conv2d( + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=device + ) + else: + self.nin_shortcut = operations.Conv2d( + in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0, + device=device, + dtype=dtype + ) + + def forward(self, x, temb): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None] + + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x+h + +class Model(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1,2,4,8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + use_timestep=True, + use_linear_attn=False, + attn_type="vanilla", + device=None, + dtype=None, + operations=ops + ): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = self.ch*4 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + self.use_timestep = use_timestep + if self.use_timestep: + # timestep embedding + self.temb = nn.Module() + self.temb.dense = nn.ModuleList([ + operations.Linear( + self.ch, + self.temb_ch, + device=device, + dtype=dtype + ), + operations.Linear( + self.temb_ch, + self.temb_ch, + device=device, + dtype=dtype + ), + ]) + + # downsampling + self.conv_in = operations.Conv2d( + in_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype + ) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv, device=device, dtype=dtype) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + skip_in = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + if i_block == self.num_res_blocks: + skip_in = ch*in_ch_mult[i_level] + block.append(ResnetBlock(in_channels=block_in+skip_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv, device=device, dtype=dtype) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in, device=device, dtype=device) + self.conv_out = torch.nn.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + def forward(self, x, t=None, context=None): + #assert x.shape[2] == x.shape[3] == self.resolution + if context is not None: + # assume aligned context, cat along channel axis + x = torch.cat((x, context), dim=1) + if self.use_timestep: + # timestep embedding + assert t is not None + temb = get_timestep_embedding(t, self.ch) + temb = self.temb.dense[0](temb) + temb = nonlinearity(temb) + temb = self.temb.dense[1](temb) + else: + temb = None + + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks+1): + h = self.up[i_level].block[i_block]( + torch.cat([h, hs.pop()], dim=1), temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + def get_last_layer(self): + return self.conv_out.weight + + +class Encoder(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1,2,4,8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + z_channels, + double_z=True, + use_linear_attn=False, + attn_type="vanilla", + device=None, + dtype=None, + operations=ops, + **ignore_kwargs + ): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + # downsampling + self.conv_in = operations.Conv2d(in_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv, device=device, dtype=dtype) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype) + + # end + self.norm_out = Normalize(block_in, device=device, dtype=dtype) + self.conv_out = operations.Conv2d(block_in, + 2*z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + def forward(self, x): + # timestep embedding + temb = None + + # print(f'encoder-input={x.shape}') + # downsampling + hs = [self.conv_in(x)] + # print(f'encoder-conv in feat={hs[0].shape}') + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1], temb) + # print(f'encoder-down feat={h.shape}') + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + # print(f'encoder-downsample (input)={hs[-1].shape}') + hs.append(self.down[i_level].downsample(hs[-1])) + # print(f'encoder-downsample (output)={hs[-1].shape}') + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + # print(f'encoder-mid1 feat={h.shape}') + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + # print(f'encoder-mid2 feat={h.shape}') + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + # print(f'end feat={h.shape}') + return h + + +class Decoder(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1,2,4,8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + z_channels, + give_pre_end=False, + tanh_out=False, + use_linear_attn=False, + attn_type="vanilla", + device=None, + dtype=None, + operations=ops, + **ignorekwargs + ): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.give_pre_end = give_pre_end + self.tanh_out = tanh_out + + # compute in_ch_mult, block_in and curr_res at lowest res + in_ch_mult = (1,)+tuple(ch_mult) + block_in = ch*ch_mult[self.num_resolutions-1] + curr_res = resolution // 2**(self.num_resolutions-1) + self.z_shape = (1,z_channels,curr_res,curr_res) + print("AE working on z of shape {} = {} dimensions.".format( + self.z_shape, np.prod(self.z_shape))) + + # z to block_in + self.conv_in = torch.nn.Conv2d(z_channels, + block_in, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv, device=device, dtype=dtype) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in, device=device, dtype=dtype) + self.conv_out = operations.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + def forward(self, z): + #assert z.shape[1:] == self.z_shape[1:] + self.last_z_shape = z.shape + + # print(f'decoder-input={z.shape}') + # timestep embedding + temb = None + + # z to block_in + h = self.conv_in(z) + # print(f'decoder-conv in feat={h.shape}') + + # middle + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + # print(f'decoder-mid feat={h.shape}') + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks+1): + h = self.up[i_level].block[i_block](h, temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + # print(f'decoder-up feat={h.shape}') + if i_level != 0: + h = self.up[i_level].upsample(h) + # print(f'decoder-upsample feat={h.shape}') + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + # print(f'decoder-conv_out feat={h.shape}') + if self.tanh_out: + h = torch.tanh(h) + return h + + +class SimpleDecoder(nn.Module): + def __init__(self, in_channels, out_channels, device=None, dtype=None, operations=ops, *args, **kwargs): + super().__init__() + self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1), + ResnetBlock(in_channels=in_channels, + out_channels=2 * in_channels, + temb_channels=0, dropout=0.0, + device=device, + dtype=dtype), + ResnetBlock(in_channels=2 * in_channels, + out_channels=4 * in_channels, + temb_channels=0, dropout=0.0, + device=device, + dtype=dtype), + ResnetBlock(in_channels=4 * in_channels, + out_channels=2 * in_channels, + temb_channels=0, dropout=0.0, + device=device, + dtype=dtype), + operations.Conv2d(2*in_channels, in_channels, 1), + Upsample(in_channels, with_conv=True, device=device, dtype=dtype)]) + # end + self.norm_out = Normalize(in_channels, device=device, dtype=dtype) + self.conv_out = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + def forward(self, x): + for i, layer in enumerate(self.model): + if i in [1,2,3]: + x = layer(x, None) + else: + x = layer(x) + + h = self.norm_out(x) + h = nonlinearity(h) + x = self.conv_out(h) + return x + + +class UpsampleDecoder(nn.Module): + def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution, + ch_mult=(2,2), dropout=0.0, device=None, dtype=None, operations=ops): + super().__init__() + # upsampling + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + block_in = in_channels + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.res_blocks = nn.ModuleList() + self.upsample_blocks = nn.ModuleList() + for i_level in range(self.num_resolutions): + res_block = [] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + res_block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + device=device, + dtype=dtype)) + block_in = block_out + self.res_blocks.append(nn.ModuleList(res_block)) + if i_level != self.num_resolutions - 1: + self.upsample_blocks.append(Upsample(block_in, True, device=device, dtype=dtype)) + curr_res = curr_res * 2 + + # end + self.norm_out = Normalize(block_in, device=device, dtype=dtype) + self.conv_out = torch.nn.Conv2d(block_in, + out_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + + def forward(self, x): + # upsampling + h = x + for k, i_level in enumerate(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.res_blocks[i_level][i_block](h, None) + if i_level != self.num_resolutions - 1: + h = self.upsample_blocks[k](h) + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class LatentRescaler(nn.Module): + def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2, device=None, dtype=None, operations=ops): + super().__init__() + # residual block, interpolate, residual block + self.factor = factor + self.conv_in = operations.Conv2d(in_channels, + mid_channels, + kernel_size=3, + stride=1, + padding=1, + device=device, + dtype=dtype) + self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, + out_channels=mid_channels, + temb_channels=0, + dropout=0.0, + device=device, + dtype=dtype) for _ in range(depth)]) + self.attn = AttnBlock(mid_channels, device=device, dtype=dtype) + self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, + out_channels=mid_channels, + temb_channels=0, + dropout=0.0, + device=device, + dtype=dtype) for _ in range(depth)]) + + self.conv_out = operations.Conv2d(mid_channels, + out_channels, + kernel_size=1, + device=device, + dtype=dtype + ) + + def forward(self, x): + x = self.conv_in(x) + for block in self.res_block1: + x = block(x, None) + x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor)))) + x = self.attn(x) + for block in self.res_block2: + x = block(x, None) + x = self.conv_out(x) + return x + + +class MergedRescaleEncoder(nn.Module): + def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, + ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1, device=None, dtype=None, operations=ops): + super().__init__() + intermediate_chn = ch * ch_mult[-1] + self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult, + z_channels=intermediate_chn, double_z=False, resolution=resolution, + attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv, + out_ch=None, device=device, dtype=dtype) + self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn, + mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth, + device=device, dtype=dtype) + + def forward(self, x): + x = self.encoder(x) + x = self.rescaler(x) + return x + + +class MergedRescaleDecoder(nn.Module): + def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8), + dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1, + device=None, dtype=None, operations=ops): + super().__init__() + tmp_chn = z_channels*ch_mult[-1] + self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout, + resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks, + ch_mult=ch_mult, resolution=resolution, ch=ch, device=device, operations=ops) + self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn, + out_channels=tmp_chn, depth=rescale_module_depth, device=device, operations=ops) + + def forward(self, x): + x = self.rescaler(x) + x = self.decoder(x) + return x + + +class Upsampler(nn.Module): + def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2, device=None, dtype=None, operations=ops): + super().__init__() + assert out_size >= in_size + num_blocks = int(np.log2(out_size//in_size))+1 + factor_up = 1.+ (out_size % in_size) + print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}") + self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels, + out_channels=in_channels, device=device, dtype=dtype) + self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2, + attn_resolutions=[], in_channels=None, ch=in_channels, device=device, dtype=dtype, + ch_mult=[ch_mult for _ in range(num_blocks)]) + + def forward(self, x): + x = self.rescaler(x) + x = self.decoder(x) + return x + + +class Resize(nn.Module): + def __init__(self, in_channels=None, learned=False, mode="bilinear", device=None, dtype=None, operations=ops): + super().__init__() + self.with_conv = learned + self.mode = mode + if self.with_conv: + print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode") + raise NotImplementedError() + assert in_channels is not None + # no asymmetric padding in torch conv, must do it ourselves + self.conv = operations.Conv2d(in_channels, + in_channels, + kernel_size=4, + stride=2, + padding=1, + device=device, + dtype=dtype) + + def forward(self, x, scale_factor=1.0): + if scale_factor==1.0: + return x + else: + x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor) + return x + +class FirstStagePostProcessor(nn.Module): + + def __init__(self, ch_mult:list, in_channels, + pretrained_model:nn.Module=None, + reshape=False, + n_channels=None, + dropout=0., + pretrained_config=None, + device=None, + dtype=None, + operations=ops): + super().__init__() + if pretrained_config is None: + assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' + self.pretrained_model = pretrained_model + else: + assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' + self.instantiate_pretrained(pretrained_config) + + self.do_reshape = reshape + + if n_channels is None: + n_channels = self.pretrained_model.encoder.ch + + self.proj_norm = Normalize(in_channels,num_groups=in_channels//2, device=device, dtype=dtype) + self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3, + stride=1,padding=1, device=device, dtype=dtype) + + blocks = [] + downs = [] + ch_in = n_channels + for m in ch_mult: + blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout, device=device, dtype=dtype)) + ch_in = m * n_channels + downs.append(Downsample(ch_in, with_conv=False, device=device, dtype=dtype)) + + self.model = nn.ModuleList(blocks) + self.downsampler = nn.ModuleList(downs) + + + def instantiate_pretrained(self, config): + model = instantiate_from_config(config) + self.pretrained_model = model.eval() + # self.pretrained_model.train = False + for param in self.pretrained_model.parameters(): + param.requires_grad = False + + + @torch.no_grad() + def encode_with_pretrained(self,x): + c = self.pretrained_model.encode(x) + if isinstance(c, DiagonalGaussianDistribution): + c = c.mode() + return c + + def forward(self,x): + z_fs = self.encode_with_pretrained(x) + z = self.proj_norm(z_fs) + z = self.proj(z) + z = nonlinearity(z) + + for submodel, downmodel in zip(self.model,self.downsampler): + z = submodel(z,temb=None) + z = downmodel(z) + + if self.do_reshape: + z = rearrange(z,'b c h w -> b (h w) c') + return z \ No newline at end of file diff --git a/modules/networks/openaimodel3d.py b/modules/networks/openaimodel3d.py new file mode 100644 index 0000000..838b6fe --- /dev/null +++ b/modules/networks/openaimodel3d.py @@ -0,0 +1,822 @@ +from functools import partial +from abc import abstractmethod +import torch +import torch.nn as nn +from einops import rearrange +import torch.nn.functional as F +from ...models.utils_diffusion import timestep_embedding +from ...common import checkpoint +from ...basics import ( + zero_module, + conv_nd, + linear, + avg_pool_nd, + normalization +) +from ...modules.attention import SpatialTransformer, TemporalTransformer +import comfy.ops +import logging + +ops = comfy.ops.disable_weight_init + +class TimestepBlock(nn.Module): + """ + Any module where forward() takes timestep embeddings as a second argument. + """ + @abstractmethod + def forward(self, x, emb): + """ + Apply the module to `x` given `emb` timestep embeddings. + """ + +#This is needed because accelerate makes a copy of transformer_options which breaks "transformer_index" +def forward_timestep_embed(ts, x, emb, context=None, batch_size=None, transformer_options={}): + for layer in ts: + if isinstance(layer, TimestepBlock): + x = layer(x, emb, batch_size=batch_size) + elif isinstance(layer, SpatialTransformer): + x = layer(x, context) + if "transformer_index" in transformer_options: + transformer_options["transformer_index"] += 1 + elif isinstance(layer, TemporalTransformer): + x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size) + x = layer(x, context) + if "transformer_index" in transformer_options: + transformer_options["transformer_index"] += 1 + x = rearrange(x, 'b c f h w -> (b f) c h w') + else: + x = layer(x) + return x + +class TimestepEmbedSequential(nn.Sequential, TimestepBlock): + """ + A sequential module that passes timestep embeddings to the children that + support it as an extra input. + """ + + def forward(self, *args, **kwargs): + return forward_timestep_embed(self, *args, **kwargs) + +class Downsample(nn.Module): + """ + A downsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + downsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=ops): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + stride = 2 if dims != 3 else (1, 2, 2) + if use_conv: + self.op = operations.conv_nd( + dims, self.channels, self.out_channels, 3, stride=stride, padding=padding + ) + else: + assert self.channels == self.out_channels + self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) + + def forward(self, x): + assert x.shape[1] == self.channels + return self.op(x) + +class Upsample(nn.Module): + """ + An upsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + upsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=ops): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + if use_conv: + self.conv = operations.conv_nd(dims, self.channels, self.out_channels, 3, padding=padding, dtype=dtype, device=device) + + def forward(self, x): + assert x.shape[1] == self.channels + if self.dims == 3: + x = F.interpolate(x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode='nearest') + else: + x = F.interpolate(x, scale_factor=2, mode='nearest') + if self.use_conv: + x = self.conv(x) + return x + +class ResBlock(TimestepBlock): + """ + A residual block that can optionally change the number of channels. + :param channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param up: if True, use this block for upsampling. + :param down: if True, use this block for downsampling. + :param use_temporal_conv: if True, use the temporal convolution. + :param use_image_dataset: if True, the temporal parameters will not be optimized. + """ + + def __init__( + self, + channels, + emb_channels, + dropout, + out_channels=None, + use_scale_shift_norm=False, + dims=2, + use_checkpoint=False, + use_conv=False, + up=False, + down=False, + kernel_size=3, + use_temporal_conv=False, + tempspatial_aware=False, + dtype=None, + device=None, + operations=ops + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.dropout = dropout + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_checkpoint = use_checkpoint + self.use_scale_shift_norm = use_scale_shift_norm + self.use_temporal_conv = use_temporal_conv + + if isinstance(kernel_size, list): + padding =[k // 2 for k in kernel_size] + else: + padding = kernel_size // 2 + + # operations used in normalization function + self.in_layers = nn.Sequential( + normalization(channels, dtype=dtype, device=device), + nn.SiLU(), + operations.conv_nd(dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device), + ) + + self.updown = up or down + + if up: + self.h_upd = Upsample(channels, False, dims, dtype=dtype, device=device) + self.x_upd = Upsample(channels, False, dims, dtype=dtype, device=device) + elif down: + self.h_upd = Downsample(channels, False, dims, dtype=dtype, device=device) + self.x_upd = Downsample(channels, False, dims, dtype=dtype, device=device) + else: + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential( + nn.SiLU(), + operations.Linear( + emb_channels, + 2 * self.out_channels if use_scale_shift_norm else self.out_channels, + dtype=dtype, + device=device + ), + ) + self.out_layers = nn.Sequential( + normalization(self.out_channels, dtype=dtype, device=device), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(operations.Conv2d(self.out_channels, self.out_channels, 3, padding=1, dtype=dtype, device=device)), + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device) + else: + self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 1, dtype=dtype, device=device) + + if self.use_temporal_conv: + self.temopral_conv = TemporalConvBlock( + self.out_channels, + self.out_channels, + dropout=0.1, + spatial_aware=tempspatial_aware, + dtype=dtype, + device=device + ) + + def forward(self, x, emb, batch_size=None): + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + :param x: an [N x C x ...] Tensor of features. + :param emb: an [N x emb_channels] Tensor of timestep embeddings. + :return: an [N x C x ...] Tensor of outputs. + """ + input_tuple = (x, emb) + if batch_size: + forward_batchsize = partial(self._forward, batch_size=batch_size) + return checkpoint(forward_batchsize, input_tuple, self.parameters(), self.use_checkpoint) + return checkpoint(self._forward, input_tuple, self.parameters(), self.use_checkpoint) + + def _forward(self, x, emb, batch_size=None): + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + emb_out = self.emb_layers(emb).type(h.dtype) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + if self.use_scale_shift_norm: + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1 + scale) + shift + h = out_rest(h) + else: + h = h + emb_out + h = self.out_layers(h) + h = self.skip_connection(x) + h + + if self.use_temporal_conv and batch_size: + h = rearrange(h, '(b t) c h w -> b c t h w', b=batch_size) + h = self.temopral_conv(h) + h = rearrange(h, 'b c t h w -> (b t) c h w') + return h + +class TemporalConvBlock(nn.Module): + """ + Adapted from modelscope: https://github.com/modelscope/modelscope/blob/master/modelscope/models/multi_modal/video_synthesis/unet_sd.py + """ + def __init__( + self, + in_channels, + out_channels=None, + dropout=0.0, + spatial_aware=False, + dtype=None, + device=None, + operations=ops + ): + super(TemporalConvBlock, self).__init__() + if out_channels is None: + out_channels = in_channels + self.in_channels = in_channels + self.out_channels = out_channels + th_kernel_shape = (3, 1, 1) if not spatial_aware else (3, 3, 1) + th_padding_shape = (1, 0, 0) if not spatial_aware else (1, 1, 0) + tw_kernel_shape = (3, 1, 1) if not spatial_aware else (3, 1, 3) + tw_padding_shape = (1, 0, 0) if not spatial_aware else (1, 0, 1) + + # conv layers + self.conv1 = nn.Sequential( + operations.GroupNorm(32, in_channels, device=device, dtype=dtype), nn.SiLU(), + operations.Conv3d(in_channels, out_channels, th_kernel_shape, padding=th_padding_shape, device=device, dtype=dtype)) + self.conv2 = nn.Sequential( + operations.GroupNorm(32, out_channels, device=device, dtype=dtype), nn.SiLU(), nn.Dropout(dropout), + operations.Conv3d(out_channels, in_channels, tw_kernel_shape, padding=tw_padding_shape, device=device, dtype=dtype)) + self.conv3 = nn.Sequential( + operations.GroupNorm(32, out_channels, device=device, dtype=dtype), nn.SiLU(), nn.Dropout(dropout), + operations.Conv3d(out_channels, in_channels, th_kernel_shape, padding=th_padding_shape, device=device, dtype=dtype)) + self.conv4 = nn.Sequential( + operations.GroupNorm(32, out_channels, device=device, dtype=dtype), nn.SiLU(), nn.Dropout(dropout), + operations.Conv3d(out_channels, in_channels, tw_kernel_shape, padding=tw_padding_shape, device=device, dtype=dtype)) + + # zero out the last layer params,so the conv block is identity + nn.init.zeros_(self.conv4[-1].weight) + nn.init.zeros_(self.conv4[-1].bias) + + def forward(self, x): + identity = x + x = self.conv1(x) + x = self.conv2(x) + x = self.conv3(x) + x = self.conv4(x) + + return identity + x + +def context_processor(context, t, img_emb=None, temporal_size=16, concat_only=False, disable_concat=False): + if disable_concat: + return context + + ## repeat t times for context [(b t) 77 768] & time embedding + ## check if we use per-frame image conditioning + + if img_emb is not None: + context = torch.cat([context, img_emb.to(context.device, context.dtype)], dim=1) + + if concat_only: + return context + + b, l_context, _ = context.shape + if l_context == 77 + t * temporal_size: + context_text, context_img = context[:,:77,:], context[:,77:,:] + context_text = context_text.repeat_interleave(repeats=t, dim=0) + context_img = rearrange(context_img, 'b (t l) c -> (b t) l c', t=t) + context = torch.cat([context_text, context_img], dim=1) + else: + context = context.repeat_interleave(repeats=t, dim=0) + + return context + +def apply_control(h, control, name, cond_idx=None): + if control is not None and name in control and len(control[name]) > 0: + frames = h.shape[0] + ctrl = control[name].pop() + if ctrl is not None: + try: + if cond_idx is not None and ctrl.shape[0] > frames: + ctrl_frames_list = list(range(ctrl.shape[0])) + ctrl_frames = len(ctrl_frames_list) + + idxs = ( + ctrl_frames_list[:ctrl_frames // 2] if cond_idx == 0 else \ + ctrl_frames_list[ctrl_frames // 2:] + ) + + ctrl = ctrl[idxs] + + h += ctrl + except Exception as e: + if h.shape != ctrl.shape: + logging.warning( + "warning control could not be applied {} {}".format(h.shape, ctrl.shape) + ) + logging.warning(e) + return h + +class UNetModel(nn.Module): + """ + The full UNet model with attention and timestep embedding. + :param in_channels: in_channels in the input Tensor. + :param model_channels: base channel count for the model. + :param out_channels: channels in the output Tensor. + :param num_res_blocks: number of residual blocks per downsample. + :param attention_resolutions: a collection of downsample rates at which + attention will take place. May be a set, list, or tuple. + For example, if this contains 4, then at 4x downsampling, attention + will be used. + :param dropout: the dropout probability. + :param channel_mult: channel multiplier for each level of the UNet. + :param conv_resample: if True, use learned convolutions for upsampling and + downsampling. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param num_classes: if specified (as an int), then this model will be + class-conditional with `num_classes` classes. + :param use_checkpoint: use gradient checkpointing to reduce memory usage. + :param num_heads: the number of attention heads in each attention layer. + :param num_heads_channels: if specified, ignore num_heads and instead use + a fixed channel width per attention head. + :param num_heads_upsample: works with num_heads to set a different number + of heads for upsampling. Deprecated. + :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. + :param resblock_updown: use residual blocks for up/downsampling. + :param use_new_attention_order: use a different attention pattern for potentially + increased efficiency. + """ + + def __init__(self, + in_channels, + model_channels, + out_channels, + num_res_blocks, + attention_resolutions, + dropout=0.0, + channel_mult=(1, 2, 4, 8), + conv_resample=True, + dims=2, + context_dim=None, + use_scale_shift_norm=False, + resblock_updown=False, + num_heads=-1, + num_head_channels=-1, + transformer_depth=1, + use_linear=False, + use_checkpoint=False, + temporal_conv=False, + tempspatial_aware=False, + temporal_attention=True, + use_relative_position=True, + use_causal_attention=False, + temporal_length=None, + use_fp16=False, + addition_attention=False, + temporal_selfatt_only=True, + image_cross_attention=False, + image_cross_attention_scale_learnable=False, + default_fs=4, + fs_condition=False, + device=None, + dtype=torch.float16, + operations=ops + ): + super(UNetModel, self).__init__() + if num_heads == -1: + assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' + if num_head_channels == -1: + assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' + + self.in_channels = in_channels + self.model_channels = model_channels + self.out_channels = out_channels + self.num_res_blocks = num_res_blocks + self.attention_resolutions = attention_resolutions + self.dropout = dropout + self.channel_mult = channel_mult + self.conv_resample = conv_resample + self.temporal_attention = temporal_attention + time_embed_dim = model_channels * 4 + self.use_checkpoint = use_checkpoint + temporal_self_att_only = True + self.addition_attention = addition_attention + self.temporal_length = temporal_length + self.image_cross_attention = image_cross_attention + self.image_cross_attention_scale_learnable = image_cross_attention_scale_learnable + self.default_fs = default_fs + self.fs_condition = fs_condition + self.device = device + #self.dtype = dtype + self.dtype = torch.float32 + + ## Time embedding blocks + self.time_embed = nn.Sequential( + linear(model_channels, time_embed_dim, device=device, dtype=self.dtype), + nn.SiLU(), + linear(time_embed_dim, time_embed_dim, device=device, dtype=self.dtype), + ) + if fs_condition: + self.fps_embedding = nn.Sequential( + linear(model_channels, time_embed_dim, device=device, dtype=self.dtype), + nn.SiLU(), + linear(time_embed_dim, time_embed_dim, device=device, dtype=self.dtype), + ) + nn.init.zeros_(self.fps_embedding[-1].weight) + nn.init.zeros_(self.fps_embedding[-1].bias) + ## Input Block + self.input_blocks = nn.ModuleList( + [ + TimestepEmbedSequential( + operations.conv_nd( + dims, + in_channels, + model_channels, + 3, + padding=1, + device=device, + dtype=self.dtype + )) + ] + ) + if self.addition_attention: + self.init_attn=TimestepEmbedSequential( + TemporalTransformer( + model_channels, + n_heads=8, + d_head=num_head_channels, + depth=transformer_depth, + context_dim=context_dim, + use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only, + causal_attention=False, relative_position=use_relative_position, + temporal_length=temporal_length, + device=device, + dtype=self.dtype + )) + + input_block_chans = [model_channels] + ch = model_channels + ds = 1 + for level, mult in enumerate(channel_mult): + for _ in range(num_res_blocks): + layers = [ + ResBlock(ch, time_embed_dim, dropout, + out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, + use_temporal_conv=temporal_conv, + device=device, + dtype=self.dtype + ) + ] + ch = mult * model_channels + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + layers.append( + SpatialTransformer(ch, num_heads, dim_head, + depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, + use_checkpoint=use_checkpoint, disable_self_attn=False, + video_length=temporal_length, image_cross_attention=self.image_cross_attention, + image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable, + device=device, + dtype=self.dtype + ) + ) + if self.temporal_attention: + layers.append( + TemporalTransformer(ch, num_heads, dim_head, + depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, + use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only, + causal_attention=use_causal_attention, relative_position=use_relative_position, + temporal_length=temporal_length, + device=device, + dtype=self.dtype + ) + ) + self.input_blocks.append(TimestepEmbedSequential(*layers)) + input_block_chans.append(ch) + if level != len(channel_mult) - 1: + out_ch = ch + self.input_blocks.append( + TimestepEmbedSequential( + ResBlock(ch, time_embed_dim, dropout, + out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + down=True, + device=device, + dtype=self.dtype + ) + if resblock_updown + else Downsample( + ch, + conv_resample, + dims=dims, + out_channels=out_ch, + device=device, + dtype=self.dtype + ) + ) + ) + ch = out_ch + input_block_chans.append(ch) + ds *= 2 + + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + layers = [ + ResBlock(ch, time_embed_dim, dropout, + dims=dims, use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, + use_temporal_conv=temporal_conv, + device=device, + dtype=self.dtype + ), + SpatialTransformer(ch, num_heads, dim_head, + depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, + use_checkpoint=use_checkpoint, disable_self_attn=False, video_length=temporal_length, + image_cross_attention=self.image_cross_attention,image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable, + device=device, + dtype=self.dtype + ) + ] + if self.temporal_attention: + layers.append( + TemporalTransformer(ch, num_heads, dim_head, + depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, + use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only, + causal_attention=use_causal_attention, relative_position=use_relative_position, + temporal_length=temporal_length, + device=device, + dtype=self.dtype + ) + ) + layers.append( + ResBlock(ch, time_embed_dim, dropout, + dims=dims, use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, + use_temporal_conv=temporal_conv, + device=device, + dtype=self.dtype + ) + ) + + ## Middle Block + self.middle_block = TimestepEmbedSequential(*layers) + + ## Output Block + self.output_blocks = nn.ModuleList([]) + for level, mult in list(enumerate(channel_mult))[::-1]: + for i in range(num_res_blocks + 1): + ich = input_block_chans.pop() + layers = [ + ResBlock(ch + ich, time_embed_dim, dropout, + out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, + use_temporal_conv=temporal_conv, + device=device, + dtype=self.dtype + ) + ] + ch = model_channels * mult + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + layers.append( + SpatialTransformer(ch, num_heads, dim_head, + depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, + use_checkpoint=use_checkpoint, disable_self_attn=False, video_length=temporal_length, + image_cross_attention=self.image_cross_attention,image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable, + device=device, + dtype=self.dtype + ) + ) + if self.temporal_attention: + layers.append( + TemporalTransformer(ch, num_heads, dim_head, + depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, + use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only, + causal_attention=use_causal_attention, relative_position=use_relative_position, + temporal_length=temporal_length, + device=device, + dtype=self.dtype + ) + ) + if level and i == num_res_blocks: + out_ch = ch + layers.append( + ResBlock(ch, time_embed_dim, dropout, + out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + up=True, + device=device, + dtype=self.dtype + ) + if resblock_updown + else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) + ) + ds //= 2 + self.output_blocks.append(TimestepEmbedSequential(*layers)) + + self.out = nn.Sequential( + normalization(ch, device=device, dtype=self.dtype), + nn.SiLU(), + zero_module( + operations.conv_nd( + dims, + model_channels, + out_channels, + 3, + padding=1, + device=device, + dtype=self.dtype + ) + ), + ) + + # TODO Add Transformer options to leverage the usage of patches. + def forward( + self, + x, + timesteps, + context=None, + context_in=None, + cc_concat=None, + num_video_frames=16, + features_adapter=None, + fs=None, + img_emb=None, + control=None, + transformer_options={}, + cond_idx=None, + **kwargs + ): + + if any([fs is None, img_emb is None, cc_concat is None]): + raise ValueError("One or more of the required inputs for UNet Forward is None.") + + cond_idx = transformer_options.get("cond_idx", None) + transformer_options['original_shape'] = list(x.shape) + transformer_options['transformer_index'] = 0 + transformer_patches = transformer_options.get("patches", {}) + + # In ComfyUI, the frames are always with the batch, so we deconstruct it here. + # This is mandatory as this is a video based model. + # We usually denote "f" as frames, but will use "t" (time) to be consistent with DynamiCrafter. + b,_,t,_,_ = x.shape + + context = context_in + cc_concat = cc_concat.to(x.device, x.dtype) + x = torch.cat([x, cc_concat], dim=1) + + fs = fs.to(x.device, x.dtype) + + timestep = timesteps + context = context_processor(context, num_video_frames, img_emb=img_emb) + + t_emb = timestep_embedding(timestep, self.model_channels, repeat_only=False, dtype=self.dtype) + emb = self.time_embed(t_emb) + emb = emb.repeat_interleave(repeats=t, dim=0) + + ## always in shape (b t) c h w, except for temporal layer + x = rearrange(x, 'b c t h w -> (b t) c h w') + + ## combine emb + if self.fs_condition: + if fs is None: + fs = torch.tensor( + [self.default_fs] * b, dtype=torch.long, device=x.device) + fs_emb = timestep_embedding(fs, self.model_channels, repeat_only=False, dtype=self.dtype).type(x.dtype) + + fs_embed = self.fps_embedding(fs_emb) + fs_embed = fs_embed.repeat_interleave(repeats=t, dim=0) + + emb = emb + fs_embed + + h = x.type(self.dtype) + adapter_idx = 0 + hs = [] + + for id, module in enumerate(self.input_blocks): + transformer_options["block"] = ("input", id) + #h = module(h, emb, context=context, batch_size=b) + h = forward_timestep_embed( + module, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) + h = apply_control(h, control, 'input', cond_idx) + + if "input_block_patch" in transformer_patches: + patch = transformer_patches["input_block_patch"] + for p in patch: + h = p(h, transformer_options) + + if id ==0 and self.addition_attention: + h = forward_timestep_embed( + self.init_attn, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) + ## plug-in adapter features + if ((id+1)%3 == 0) and features_adapter is not None: + h = h + features_adapter[adapter_idx] + adapter_idx += 1 + hs.append(h) + if "input_block_patch_after_skip" in transformer_patches: + patch = transformer_patches["input_block_patch_after_skip"] + for p in patch: + h = p(h, transformer_options) + if features_adapter is not None: + assert len(features_adapter)==adapter_idx, 'Wrong features_adapter' + transformer_options["block"] = ("middle", 0) + h = forward_timestep_embed( + self.middle_block, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) + h = apply_control(h, control, 'middle', cond_idx) + for id, module in enumerate(self.output_blocks): + transformer_options["block"] = ("output", id) + hsp = hs.pop() + hsp = apply_control(hsp, control, 'output', cond_idx) + + if "output_block_patch" in transformer_patches: + patch = transformer_patches["output_block_patch"] + for p in patch: + h, hsp = p(h, hsp, transformer_options) + + h = torch.cat([h, hsp], dim=1) + del hsp + h = forward_timestep_embed( + module, + h, + emb, + context=context, + batch_size=b, + transformer_options=transformer_options + ) + h = h.type(x.dtype) + h = self.out(h) + + # We output with the tensor unfolded framewise, then reshape them to batched using ComfyUI nodes. + h = rearrange(h, '(b t) c h w -> b c t h w', t=num_video_frames) + + return h \ No newline at end of file diff --git a/modules/x_transformer.py b/modules/x_transformer.py new file mode 100644 index 0000000..5321012 --- /dev/null +++ b/modules/x_transformer.py @@ -0,0 +1,639 @@ +"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers""" +from functools import partial +from inspect import isfunction +from collections import namedtuple +from einops import rearrange, repeat +import torch +from torch import nn, einsum +import torch.nn.functional as F + +# constants +DEFAULT_DIM_HEAD = 64 + +Intermediates = namedtuple('Intermediates', [ + 'pre_softmax_attn', + 'post_softmax_attn' +]) + +LayerIntermediates = namedtuple('Intermediates', [ + 'hiddens', + 'attn_intermediates' +]) + + +class AbsolutePositionalEmbedding(nn.Module): + def __init__(self, dim, max_seq_len): + super().__init__() + self.emb = nn.Embedding(max_seq_len, dim) + self.init_() + + def init_(self): + nn.init.normal_(self.emb.weight, std=0.02) + + def forward(self, x): + n = torch.arange(x.shape[1], device=x.device) + return self.emb(n)[None, :, :] + + +class FixedPositionalEmbedding(nn.Module): + def __init__(self, dim): + super().__init__() + inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer('inv_freq', inv_freq) + + def forward(self, x, seq_dim=1, offset=0): + t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset + sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq) + emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1) + return emb[None, :, :] + + +# helpers + +def exists(val): + return val is not None + + +def default(val, d): + if exists(val): + return val + return d() if isfunction(d) else d + + +def always(val): + def inner(*args, **kwargs): + return val + return inner + + +def not_equals(val): + def inner(x): + return x != val + return inner + + +def equals(val): + def inner(x): + return x == val + return inner + + +def max_neg_value(tensor): + return -torch.finfo(tensor.dtype).max + + +# keyword argument helpers + +def pick_and_pop(keys, d): + values = list(map(lambda key: d.pop(key), keys)) + return dict(zip(keys, values)) + + +def group_dict_by_key(cond, d): + return_val = [dict(), dict()] + for key in d.keys(): + match = bool(cond(key)) + ind = int(not match) + return_val[ind][key] = d[key] + return (*return_val,) + + +def string_begins_with(prefix, str): + return str.startswith(prefix) + + +def group_by_key_prefix(prefix, d): + return group_dict_by_key(partial(string_begins_with, prefix), d) + + +def groupby_prefix_and_trim(prefix, d): + kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d) + kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items()))) + return kwargs_without_prefix, kwargs + + +# classes +class Scale(nn.Module): + def __init__(self, value, fn): + super().__init__() + self.value = value + self.fn = fn + + def forward(self, x, **kwargs): + x, *rest = self.fn(x, **kwargs) + return (x * self.value, *rest) + + +class Rezero(nn.Module): + def __init__(self, fn): + super().__init__() + self.fn = fn + self.g = nn.Parameter(torch.zeros(1)) + + def forward(self, x, **kwargs): + x, *rest = self.fn(x, **kwargs) + return (x * self.g, *rest) + + +class ScaleNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.scale = dim ** -0.5 + self.eps = eps + self.g = nn.Parameter(torch.ones(1)) + + def forward(self, x): + norm = torch.norm(x, dim=-1, keepdim=True) * self.scale + return x / norm.clamp(min=self.eps) * self.g + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-8): + super().__init__() + self.scale = dim ** -0.5 + self.eps = eps + self.g = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + norm = torch.norm(x, dim=-1, keepdim=True) * self.scale + return x / norm.clamp(min=self.eps) * self.g + + +class Residual(nn.Module): + def forward(self, x, residual): + return x + residual + + +class GRUGating(nn.Module): + def __init__(self, dim): + super().__init__() + self.gru = nn.GRUCell(dim, dim) + + def forward(self, x, residual): + gated_output = self.gru( + rearrange(x, 'b n d -> (b n) d'), + rearrange(residual, 'b n d -> (b n) d') + ) + + return gated_output.reshape_as(x) + + +# feedforward + +class GEGLU(nn.Module): + def __init__(self, dim_in, dim_out): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2) + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.): + super().__init__() + inner_dim = int(dim * mult) + dim_out = default(dim_out, dim) + project_in = nn.Sequential( + nn.Linear(dim, inner_dim), + nn.GELU() + ) if not glu else GEGLU(dim, inner_dim) + + self.net = nn.Sequential( + project_in, + nn.Dropout(dropout), + nn.Linear(inner_dim, dim_out) + ) + + def forward(self, x): + return self.net(x) + + +# attention. +class Attention(nn.Module): + def __init__( + self, + dim, + dim_head=DEFAULT_DIM_HEAD, + heads=8, + causal=False, + mask=None, + talking_heads=False, + sparse_topk=None, + use_entmax15=False, + num_mem_kv=0, + dropout=0., + on_attn=False + ): + super().__init__() + if use_entmax15: + raise NotImplementedError("Check out entmax activation instead of softmax activation!") + self.scale = dim_head ** -0.5 + self.heads = heads + self.causal = causal + self.mask = mask + + inner_dim = dim_head * heads + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_k = nn.Linear(dim, inner_dim, bias=False) + self.to_v = nn.Linear(dim, inner_dim, bias=False) + self.dropout = nn.Dropout(dropout) + + # talking heads + self.talking_heads = talking_heads + if talking_heads: + self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads)) + self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads)) + + # explicit topk sparse attention + self.sparse_topk = sparse_topk + + # entmax + #self.attn_fn = entmax15 if use_entmax15 else F.softmax + self.attn_fn = F.softmax + + # add memory key / values + self.num_mem_kv = num_mem_kv + if num_mem_kv > 0: + self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head)) + self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head)) + + # attention on attention + self.attn_on_attn = on_attn + self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim) + + def forward( + self, + x, + context=None, + mask=None, + context_mask=None, + rel_pos=None, + sinusoidal_emb=None, + prev_attn=None, + mem=None + ): + b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device + kv_input = default(context, x) + + q_input = x + k_input = kv_input + v_input = kv_input + + if exists(mem): + k_input = torch.cat((mem, k_input), dim=-2) + v_input = torch.cat((mem, v_input), dim=-2) + + if exists(sinusoidal_emb): + # in shortformer, the query would start at a position offset depending on the past cached memory + offset = k_input.shape[-2] - q_input.shape[-2] + q_input = q_input + sinusoidal_emb(q_input, offset=offset) + k_input = k_input + sinusoidal_emb(k_input) + + q = self.to_q(q_input) + k = self.to_k(k_input) + v = self.to_v(v_input) + + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v)) + + input_mask = None + if any(map(exists, (mask, context_mask))): + q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool()) + k_mask = q_mask if not exists(context) else context_mask + k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool()) + q_mask = rearrange(q_mask, 'b i -> b () i ()') + k_mask = rearrange(k_mask, 'b j -> b () () j') + input_mask = q_mask * k_mask + + if self.num_mem_kv > 0: + mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v)) + k = torch.cat((mem_k, k), dim=-2) + v = torch.cat((mem_v, v), dim=-2) + if exists(input_mask): + input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True) + + dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale + mask_value = max_neg_value(dots) + + if exists(prev_attn): + dots = dots + prev_attn + + pre_softmax_attn = dots + + if talking_heads: + dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous() + + if exists(rel_pos): + dots = rel_pos(dots) + + if exists(input_mask): + dots.masked_fill_(~input_mask, mask_value) + del input_mask + + if self.causal: + i, j = dots.shape[-2:] + r = torch.arange(i, device=device) + mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j') + mask = F.pad(mask, (j - i, 0), value=False) + dots.masked_fill_(mask, mask_value) + del mask + + if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]: + top, _ = dots.topk(self.sparse_topk, dim=-1) + vk = top[..., -1].unsqueeze(-1).expand_as(dots) + mask = dots < vk + dots.masked_fill_(mask, mask_value) + del mask + + attn = self.attn_fn(dots, dim=-1) + post_softmax_attn = attn + + attn = self.dropout(attn) + + if talking_heads: + attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous() + + out = einsum('b h i j, b h j d -> b h i d', attn, v) + out = rearrange(out, 'b h n d -> b n (h d)') + + intermediates = Intermediates( + pre_softmax_attn=pre_softmax_attn, + post_softmax_attn=post_softmax_attn + ) + + return self.to_out(out), intermediates + + +class AttentionLayers(nn.Module): + def __init__( + self, + dim, + depth, + heads=8, + causal=False, + cross_attend=False, + only_cross=False, + use_scalenorm=False, + use_rmsnorm=False, + use_rezero=False, + rel_pos_num_buckets=32, + rel_pos_max_distance=128, + position_infused_attn=False, + custom_layers=None, + sandwich_coef=None, + par_ratio=None, + residual_attn=False, + cross_residual_attn=False, + macaron=False, + pre_norm=True, + gate_residual=False, + **kwargs + ): + super().__init__() + ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs) + attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs) + + dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD) + + self.dim = dim + self.depth = depth + self.layers = nn.ModuleList([]) + + self.has_pos_emb = position_infused_attn + self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None + self.rotary_pos_emb = always(None) + + assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance' + self.rel_pos = None + + self.pre_norm = pre_norm + + self.residual_attn = residual_attn + self.cross_residual_attn = cross_residual_attn + + norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm + norm_class = RMSNorm if use_rmsnorm else norm_class + norm_fn = partial(norm_class, dim) + + norm_fn = nn.Identity if use_rezero else norm_fn + branch_fn = Rezero if use_rezero else None + + if cross_attend and not only_cross: + default_block = ('a', 'c', 'f') + elif cross_attend and only_cross: + default_block = ('c', 'f') + else: + default_block = ('a', 'f') + + if macaron: + default_block = ('f',) + default_block + + if exists(custom_layers): + layer_types = custom_layers + elif exists(par_ratio): + par_depth = depth * len(default_block) + assert 1 < par_ratio <= par_depth, 'par ratio out of range' + default_block = tuple(filter(not_equals('f'), default_block)) + par_attn = par_depth // par_ratio + depth_cut = par_depth * 2 // 3 # 2 / 3 attention layer cutoff suggested by PAR paper + par_width = (depth_cut + depth_cut // par_attn) // par_attn + assert len(default_block) <= par_width, 'default block is too large for par_ratio' + par_block = default_block + ('f',) * (par_width - len(default_block)) + par_head = par_block * par_attn + layer_types = par_head + ('f',) * (par_depth - len(par_head)) + elif exists(sandwich_coef): + assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth' + layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef + else: + layer_types = default_block * depth + + self.layer_types = layer_types + self.num_attn_layers = len(list(filter(equals('a'), layer_types))) + + for layer_type in self.layer_types: + if layer_type == 'a': + layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs) + elif layer_type == 'c': + layer = Attention(dim, heads=heads, **attn_kwargs) + elif layer_type == 'f': + layer = FeedForward(dim, **ff_kwargs) + layer = layer if not macaron else Scale(0.5, layer) + else: + raise Exception(f'invalid layer type {layer_type}') + + if isinstance(layer, Attention) and exists(branch_fn): + layer = branch_fn(layer) + + if gate_residual: + residual_fn = GRUGating(dim) + else: + residual_fn = Residual() + + self.layers.append(nn.ModuleList([ + norm_fn(), + layer, + residual_fn + ])) + + def forward( + self, + x, + context=None, + mask=None, + context_mask=None, + mems=None, + return_hiddens=False + ): + hiddens = [] + intermediates = [] + prev_attn = None + prev_cross_attn = None + + mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers + + for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)): + is_last = ind == (len(self.layers) - 1) + + if layer_type == 'a': + hiddens.append(x) + layer_mem = mems.pop(0) + + residual = x + + if self.pre_norm: + x = norm(x) + + if layer_type == 'a': + out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos, + prev_attn=prev_attn, mem=layer_mem) + elif layer_type == 'c': + out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn) + elif layer_type == 'f': + out = block(x) + + x = residual_fn(out, residual) + + if layer_type in ('a', 'c'): + intermediates.append(inter) + + if layer_type == 'a' and self.residual_attn: + prev_attn = inter.pre_softmax_attn + elif layer_type == 'c' and self.cross_residual_attn: + prev_cross_attn = inter.pre_softmax_attn + + if not self.pre_norm and not is_last: + x = norm(x) + + if return_hiddens: + intermediates = LayerIntermediates( + hiddens=hiddens, + attn_intermediates=intermediates + ) + + return x, intermediates + + return x + + +class Encoder(AttentionLayers): + def __init__(self, **kwargs): + assert 'causal' not in kwargs, 'cannot set causality on encoder' + super().__init__(causal=False, **kwargs) + + + +class TransformerWrapper(nn.Module): + def __init__( + self, + *, + num_tokens, + max_seq_len, + attn_layers, + emb_dim=None, + max_mem_len=0., + emb_dropout=0., + num_memory_tokens=None, + tie_embedding=False, + use_pos_emb=True + ): + super().__init__() + assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder' + + dim = attn_layers.dim + emb_dim = default(emb_dim, dim) + + self.max_seq_len = max_seq_len + self.max_mem_len = max_mem_len + self.num_tokens = num_tokens + + self.token_emb = nn.Embedding(num_tokens, emb_dim) + self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if ( + use_pos_emb and not attn_layers.has_pos_emb) else always(0) + self.emb_dropout = nn.Dropout(emb_dropout) + + self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity() + self.attn_layers = attn_layers + self.norm = nn.LayerNorm(dim) + + self.init_() + + self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t() + + # memory tokens (like [cls]) from Memory Transformers paper + num_memory_tokens = default(num_memory_tokens, 0) + self.num_memory_tokens = num_memory_tokens + if num_memory_tokens > 0: + self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim)) + + # let funnel encoder know number of memory tokens, if specified + if hasattr(attn_layers, 'num_memory_tokens'): + attn_layers.num_memory_tokens = num_memory_tokens + + def init_(self): + nn.init.normal_(self.token_emb.weight, std=0.02) + + def forward( + self, + x, + return_embeddings=False, + mask=None, + return_mems=False, + return_attn=False, + mems=None, + **kwargs + ): + b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens + x = self.token_emb(x) + x += self.pos_emb(x) + x = self.emb_dropout(x) + + x = self.project_emb(x) + + if num_mem > 0: + mem = repeat(self.memory_tokens, 'n d -> b n d', b=b) + x = torch.cat((mem, x), dim=1) + + # auto-handle masking after appending memory tokens + if exists(mask): + mask = F.pad(mask, (num_mem, 0), value=True) + + x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs) + x = self.norm(x) + + mem, x = x[:, :num_mem], x[:, num_mem:] + + out = self.to_logits(x) if not return_embeddings else x + + if return_mems: + hiddens = intermediates.hiddens + new_mems = list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens))) if exists(mems) else hiddens + new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems)) + return out, new_mems + + if return_attn: + attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates)) + return out, attn_maps + + return out \ No newline at end of file From 37ce3a22ff43a200e11fb1f0dea4f1fde5977821 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 15:25:31 -0700 Subject: [PATCH 3/6] Delete modules directory --- modules/attention.py | 676 ------------------- modules/encoders/condition.py | 389 ----------- modules/encoders/resampler.py | 145 ---- modules/networks/ae_modules.py | 1026 ----------------------------- modules/networks/openaimodel3d.py | 822 ----------------------- modules/x_transformer.py | 639 ------------------ 6 files changed, 3697 deletions(-) delete mode 100644 modules/attention.py delete mode 100644 modules/encoders/condition.py delete mode 100644 modules/encoders/resampler.py delete mode 100644 modules/networks/ae_modules.py delete mode 100644 modules/networks/openaimodel3d.py delete mode 100644 modules/x_transformer.py diff --git a/modules/attention.py b/modules/attention.py deleted file mode 100644 index 7537aa5..0000000 --- a/modules/attention.py +++ /dev/null @@ -1,676 +0,0 @@ -import torch -from torch import nn, einsum -import torch.nn.functional as F -from einops import rearrange, repeat -from functools import partial -from ..common import ( - checkpoint, - exists, - default, -) -from ..basics import zero_module -import comfy.ops -ops = comfy.ops.disable_weight_init -from comfy import model_management -from comfy.ldm.modules.attention import optimized_attention, optimized_attention_masked - -if model_management.xformers_enabled(): - import xformers - import xformers.ops - XFORMERS_IS_AVAILBLE = True -else: - XFORMERS_IS_AVAILBLE = False - -class RelativePosition(nn.Module): - """ https://github.com/evelinehong/Transformer_Relative_Position_PyTorch/blob/master/relative_position.py """ - - def __init__(self, num_units, max_relative_position): - super().__init__() - self.num_units = num_units - self.max_relative_position = max_relative_position - self.embeddings_table = nn.Parameter(torch.Tensor(max_relative_position * 2 + 1, num_units)) - nn.init.xavier_uniform_(self.embeddings_table) - - def forward(self, length_q, length_k): - device = self.embeddings_table.device - range_vec_q = torch.arange(length_q, device=device) - range_vec_k = torch.arange(length_k, device=device) - distance_mat = range_vec_k[None, :] - range_vec_q[:, None] - distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position) - final_mat = distance_mat_clipped + self.max_relative_position - final_mat = final_mat.long() - embeddings = self.embeddings_table[final_mat] - return embeddings - - -# TODO Add native Comfy optimized attention. -class CrossAttention(nn.Module): - - def __init__( - self, - query_dim, - context_dim=None, - heads=8, - dim_head=64, - dropout=0., - relative_position=False, - temporal_length=None, - video_length=None, - image_cross_attention=False, - image_cross_attention_scale=1.0, - image_cross_attention_scale_learnable=False, - text_context_len=77, - device=None, - dtype=None, - operations=ops - ): - super().__init__() - inner_dim = dim_head * heads - context_dim = default(context_dim, query_dim) - self.scale = dim_head**-0.5 - self.heads = heads - self.dim_head = dim_head - self.to_q = operations.Linear(query_dim, inner_dim, bias=False, device=device, dtype=dtype) - self.to_k = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) - self.to_v = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) - - self.to_out = nn.Sequential( - operations.Linear(inner_dim, query_dim, device=device, dtype=dtype), - nn.Dropout(dropout) - ) - - self.relative_position = relative_position - if self.relative_position: - assert(temporal_length is not None) - self.relative_position_k = RelativePosition(num_units=dim_head, max_relative_position=temporal_length) - self.relative_position_v = RelativePosition(num_units=dim_head, max_relative_position=temporal_length) - else: - ## only used for spatial attention, while NOT for temporal attention - if XFORMERS_IS_AVAILBLE and temporal_length is None: - self.forward = self.efficient_forward - else: - self.forward = self.comfy_efficient_forward - - self.video_length = video_length - self.image_cross_attention = image_cross_attention - self.image_cross_attention_scale = image_cross_attention_scale - self.text_context_len = text_context_len - self.image_cross_attention_scale_learnable = image_cross_attention_scale_learnable - if self.image_cross_attention: - self.to_k_ip = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) - self.to_v_ip = operations.Linear(context_dim, inner_dim, bias=False, device=device, dtype=dtype) - if image_cross_attention_scale_learnable: - self.register_parameter('alpha', nn.Parameter(torch.tensor(0.)) ) - - def comfy_efficient_forward(self, x, context=None, mask=None): - spatial_self_attn = (context is None) - k_ip, v_ip, out_ip = None, None, None - - h = self.heads - q = self.to_q(x) - context = default(context, x) - - if self.image_cross_attention and not spatial_self_attn: - context, context_image = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:] - k = self.to_k(context) - v = self.to_v(context) - k_ip = self.to_k_ip(context_image) - v_ip = self.to_v_ip(context_image) - else: - if not spatial_self_attn: - context = context[:,:self.text_context_len,:] - k = self.to_k(context) - v = self.to_v(context) - - out = optimized_attention(q, k, v, h) - - if exists(mask): - ## feasible for causal attention mask only - out = optimized_attention_masked(q, k, v, h) - - ## for image cross-attention - if k_ip is not None: - q = rearrange(q, 'b n (h d) -> (b h) n d', h=h) - k_ip, v_ip = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (k_ip, v_ip)) - sim_ip = torch.einsum('b i d, b j d -> b i j', q, k_ip) * self.scale - del k_ip - sim_ip = sim_ip.softmax(dim=-1) - out_ip = torch.einsum('b i j, b j d -> b i d', sim_ip, v_ip) - out_ip = rearrange(out_ip, '(b h) n d -> b n (h d)', h=h) - - if out_ip is not None: - if self.image_cross_attention_scale_learnable: - out = out + self.image_cross_attention_scale * out_ip * (torch.tanh(self.alpha)+1) - else: - out = out + self.image_cross_attention_scale * out_ip - - return self.to_out(out) - - def forward(self, x, context=None, mask=None): - spatial_self_attn = (context is None) - k_ip, v_ip, out_ip = None, None, None - - h = self.heads - q = self.to_q(x) - context = default(context, x) - - if self.image_cross_attention and not spatial_self_attn: - context, context_image = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:] - k = self.to_k(context) - v = self.to_v(context) - k_ip = self.to_k_ip(context_image) - v_ip = self.to_v_ip(context_image) - else: - - # Assumed Spatial Attention (b c h w) - if not spatial_self_attn: - context = context[:,:self.text_context_len,:] - k = self.to_k(context) - v = self.to_v(context) - - - q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) - - sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale - if self.relative_position: - len_q, len_k, len_v = q.shape[1], k.shape[1], v.shape[1] - k2 = self.relative_position_k(len_q, len_k) - sim2 = einsum('b t d, t s d -> b t s', q, k2) * self.scale # TODO check - sim += sim2 - del k - - if exists(mask): - ## feasible for causal attention mask only - max_neg_value = -torch.finfo(sim.dtype).max - mask = repeat(mask, 'b i j -> (b h) i j', h=h) - sim.masked_fill_(~(mask>0.5), max_neg_value) - - # attention, what we cannot get enough of - sim = sim.softmax(dim=-1) - - out = torch.einsum('b i j, b j d -> b i d', sim, v) - if self.relative_position: - v2 = self.relative_position_v(len_q, len_v) - out2 = einsum('b t s, t s d -> b t d', sim, v2) # TODO check - out += out2 - out = rearrange(out, '(b h) n d -> b n (h d)', h=h) - - - ## for image cross-attention - if k_ip is not None: - k_ip, v_ip = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (k_ip, v_ip)) - sim_ip = torch.einsum('b i d, b j d -> b i j', q, k_ip) * self.scale - del k_ip - sim_ip = sim_ip.softmax(dim=-1) - out_ip = torch.einsum('b i j, b j d -> b i d', sim_ip, v_ip) - out_ip = rearrange(out_ip, '(b h) n d -> b n (h d)', h=h) - - - if out_ip is not None: - if self.image_cross_attention_scale_learnable: - out = out + self.image_cross_attention_scale * out_ip * (torch.tanh(self.alpha)+1) - else: - out = out + self.image_cross_attention_scale * out_ip - - return self.to_out(out) - - def efficient_forward(self, x, context=None, mask=None): - spatial_self_attn = (context is None) - k_ip, v_ip, out_ip = None, None, None - - q = self.to_q(x) - context = default(context, x) - - if self.image_cross_attention and not spatial_self_attn: - context, context_image = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:] - k = self.to_k(context) - v = self.to_v(context) - k_ip = self.to_k_ip(context_image) - v_ip = self.to_v_ip(context_image) - else: - if not spatial_self_attn: - context = context[:,:self.text_context_len,:] - k = self.to_k(context) - v = self.to_v(context) - - b, _, _ = q.shape - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, t.shape[1], self.heads, self.dim_head) - .permute(0, 2, 1, 3) - .reshape(b * self.heads, t.shape[1], self.dim_head) - .contiguous(), - (q, k, v), - ) - # actually compute the attention, what we cannot get enough of - out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=None) - - ## for image cross-attention - if k_ip is not None: - k_ip, v_ip = map( - lambda t: t.unsqueeze(3) - .reshape(b, t.shape[1], self.heads, self.dim_head) - .permute(0, 2, 1, 3) - .reshape(b * self.heads, t.shape[1], self.dim_head) - .contiguous(), - (k_ip, v_ip), - ) - out_ip = xformers.ops.memory_efficient_attention(q, k_ip, v_ip, attn_bias=None, op=None) - out_ip = ( - out_ip.unsqueeze(0) - .reshape(b, self.heads, out.shape[1], self.dim_head) - .permute(0, 2, 1, 3) - .reshape(b, out.shape[1], self.heads * self.dim_head) - ) - - if exists(mask): - raise NotImplementedError - out = ( - out.unsqueeze(0) - .reshape(b, self.heads, out.shape[1], self.dim_head) - .permute(0, 2, 1, 3) - .reshape(b, out.shape[1], self.heads * self.dim_head) - ) - if out_ip is not None: - if self.image_cross_attention_scale_learnable: - out = out + self.image_cross_attention_scale * out_ip * (torch.tanh(self.alpha)+1) - else: - out = out + self.image_cross_attention_scale * out_ip - - return self.to_out(out) - - -class BasicTransformerBlock(nn.Module): - - def __init__( - self, - dim, - n_heads, - d_head, - dropout=0., - context_dim=None, - gated_ff=True, - checkpoint=True, - disable_self_attn=False, - attention_cls=None, - video_length=None, - image_cross_attention=False, - image_cross_attention_scale=1.0, - image_cross_attention_scale_learnable=False, - text_context_len=77, - device=None, - dtype=None, - operations=ops - ): - super().__init__() - attn_cls = CrossAttention if attention_cls is None else attention_cls - self.disable_self_attn = disable_self_attn - self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, - context_dim=None, device=device, dtype=dtype if self.disable_self_attn else None) - self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff, device=device, dtype=dtype) - self.attn2 = attn_cls( - query_dim=dim, - context_dim=context_dim, - heads=n_heads, - dim_head=d_head, - dropout=dropout, - video_length=video_length, - image_cross_attention=image_cross_attention, - image_cross_attention_scale=image_cross_attention_scale, - image_cross_attention_scale_learnable=image_cross_attention_scale_learnable, - text_context_len=text_context_len, - device=device, - dtype=dtype - ) - self.image_cross_attention = image_cross_attention - - self.norm1 = operations.LayerNorm(dim, device=device, dtype=dtype) - self.norm2 = operations.LayerNorm(dim, device=device, dtype=dtype) - self.norm3 = operations.LayerNorm(dim, device=device, dtype=dtype) - self.checkpoint = checkpoint - - def forward(self, x, context=None, mask=None, **kwargs): - ## implementation tricks: because checkpointing doesn't support non-tensor (e.g. None or scalar) arguments - input_tuple = (x,) ## should not be (x), otherwise *input_tuple will decouple x into multiple arguments - if context is not None: - input_tuple = (x, context) - if mask is not None: - forward_mask = partial(self._forward, mask=mask) - return checkpoint(forward_mask, (x,), self.parameters(), self.checkpoint) - return checkpoint(self._forward, input_tuple, self.parameters(), self.checkpoint) - - - def _forward(self, x, context=None, mask=None): - x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None, mask=mask) + x - x = self.attn2(self.norm2(x), context=context, mask=mask) + x - x = self.ff(self.norm3(x)) + x - return x - - -class SpatialTransformer(nn.Module): - """ - Transformer block for image-like data in spatial axis. - First, project the input (aka embedding) - and reshape to b, t, d. - Then apply standard transformer action. - Finally, reshape to image - NEW: use_linear for more efficiency instead of the 1x1 convs - """ - - def __init__( - self, - in_channels, - n_heads, - d_head, - depth=1, - dropout=0., - context_dim=None, - use_checkpoint=True, - disable_self_attn=False, - use_linear=False, - video_length=None, - image_cross_attention=False, - image_cross_attention_scale_learnable=False, - device=None, - dtype=None, - operations=ops - ): - super().__init__() - self.in_channels = in_channels - inner_dim = n_heads * d_head - self.norm = operations.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, device=device, dtype=dtype) - if not use_linear: - self.proj_in = opeations.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0, device=device, dtype=dtype) - else: - self.proj_in = operations.Linear(in_channels, inner_dim, device=device, dtype=dtype) - - attention_cls = None - self.transformer_blocks = nn.ModuleList([ - BasicTransformerBlock( - inner_dim, - n_heads, - d_head, - dropout=dropout, - context_dim=context_dim, - disable_self_attn=disable_self_attn, - checkpoint=use_checkpoint, - attention_cls=attention_cls, - video_length=video_length, - image_cross_attention=image_cross_attention, - image_cross_attention_scale_learnable=image_cross_attention_scale_learnable, - device=device, - dtype=dtype - ) for d in range(depth) - ]) - if not use_linear: - self.proj_out = zero_module(operations.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0, device=device, dtype=dtype)) - else: - self.proj_out = zero_module(operations.Linear(inner_dim, in_channels, device=device, dtype=dtype)) - self.use_linear = use_linear - - def forward(self, x, context=None, **kwargs): - b, c, h, w = x.shape - x_in = x - x = self.norm(x) - if not self.use_linear: - x = self.proj_in(x) - x = rearrange(x, 'b c h w -> b (h w) c').contiguous() - if self.use_linear: - x = self.proj_in(x) - for i, block in enumerate(self.transformer_blocks): - x = block(x, context=context, **kwargs) - if self.use_linear: - x = self.proj_out(x) - x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous() - if not self.use_linear: - x = self.proj_out(x) - return x + x_in - - -class TemporalTransformer(nn.Module): - """ - Transformer block for image-like data in temporal axis. - First, reshape to b, t, d. - Then apply standard transformer action. - Finally, reshape to image - """ - def __init__( - self, - in_channels, - n_heads, - d_head, - depth=1, - dropout=0., - context_dim=None, - use_checkpoint=True, - use_linear=False, - only_self_att=True, - causal_attention=False, - causal_block_size=1, - relative_position=False, - temporal_length=None, - device=None, - dtype=None, - operations=ops - ): - super().__init__() - self.only_self_att = only_self_att - self.relative_position = relative_position - self.causal_attention = causal_attention - self.causal_block_size = causal_block_size - - if only_self_att: - context_dim = None - - self.in_channels = in_channels - inner_dim = n_heads * d_head - self.norm = operations.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, device=device, dtype=dtype) - self.proj_in = nn.Conv1d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0).to(device, dtype) - if not use_linear: - self.proj_in = nn.Conv1d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0).to(device, dtype) - else: - self.proj_in = operations.Linear(in_channels, inner_dim, device=device, dtype=dtype) - - if relative_position: - assert(temporal_length is not None) - attention_cls = partial(CrossAttention, relative_position=True, temporal_length=temporal_length, device=device, dtype=dtype) - else: - attention_cls = partial(CrossAttention, temporal_length=temporal_length, device=device, dtype=dtype) - if self.causal_attention: - assert(temporal_length is not None) - self.mask = torch.tril(torch.ones([1, temporal_length, temporal_length])) - - if self.only_self_att: - context_dim = None - self.transformer_blocks = nn.ModuleList([ - BasicTransformerBlock( - inner_dim, - n_heads, - d_head, - dropout=dropout, - context_dim=context_dim, - attention_cls=attention_cls, - checkpoint=use_checkpoint, - device=device, - dtype=dtype - ) for d in range(depth) - ]) - if not use_linear: - self.proj_out = zero_module(nn.Conv1d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0).to(device, dtype)) - else: - self.proj_out = zero_module(operations.Linear(inner_dim, in_channels, device=device, dtype=dtype)) - self.use_linear = use_linear - - def forward(self, x, context=None): - b, c, t, h, w = x.shape - x_in = x - x = self.norm(x) - x = rearrange(x, 'b c t h w -> (b h w) c t').contiguous() - if not self.use_linear: - x = self.proj_in(x) - x = rearrange(x, 'bhw c t -> bhw t c').contiguous() - if self.use_linear: - x = self.proj_in(x) - - temp_mask = None - if self.causal_attention: - # slice the from mask map - temp_mask = self.mask[:,:t,:t].to(x.device) - - if temp_mask is not None: - mask = temp_mask.to(x.device) - mask = repeat(mask, 'l i j -> (l bhw) i j', bhw=b*h*w) - else: - mask = None - - if self.only_self_att: - ## note: if no context is given, cross-attention defaults to self-attention - for i, block in enumerate(self.transformer_blocks): - x = block(x, mask=mask) - x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous() - else: - x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous() - context = rearrange(context, '(b t) l con -> b t l con', t=t).contiguous() - for i, block in enumerate(self.transformer_blocks): - # calculate each batch one by one (since number in shape could not greater then 65,535 for some package) - for j in range(b): - context_j = repeat( - context[j], - 't l con -> (t r) l con', r=(h * w) // t, t=t).contiguous() - ## note: causal mask will not applied in cross-attention case - x[j] = block(x[j], context=context_j) - - if self.use_linear: - x = self.proj_out(x) - x = rearrange(x, 'b (h w) t c -> b c t h w', h=h, w=w).contiguous() - if not self.use_linear: - x = rearrange(x, 'b hw t c -> (b hw) c t').contiguous() - x = self.proj_out(x) - x = rearrange(x, '(b h w) c t -> b c t h w', b=b, h=h, w=w).contiguous() - - return x + x_in - - -class GEGLU(nn.Module): - def __init__(self, dim_in, dim_out, device=None, dtype=None, operations=ops): - super().__init__() - self.proj = operations.Linear(dim_in, dim_out * 2, device=device, dtype=dtype) - - def forward(self, x): - x, gate = self.proj(x).chunk(2, dim=-1) - return x * F.gelu(gate) - - -class FeedForward(nn.Module): - def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0., device=None, dtype=None, operations=ops): - super().__init__() - inner_dim = int(dim * mult) - dim_out = default(dim_out, dim) - project_in = nn.Sequential( - operations.Linear(dim, inner_dim, device=device, dtype=dtype), - nn.GELU() - ) if not glu else GEGLU(dim, inner_dim) - - self.net = nn.Sequential( - project_in, - nn.Dropout(dropout), - operations.Linear(inner_dim, dim_out, device=device, dtype=dtype) - ) - - def forward(self, x): - return self.net(x) - - -class LinearAttention(nn.Module): - def __init__(self, dim, heads=4, dim_head=32, device=None, dtype=None, operations=ops): - super().__init__() - self.heads = heads - hidden_dim = dim_head * heads - self.to_qkv = operations.Conv2d(dim, hidden_dim * 3, 1, bias = False, device=device, dtype=dtype) - self.to_out = operations.Conv2d(hidden_dim, dim, 1, device=device, dtype=dtype) - - def forward(self, x): - b, c, h, w = x.shape - qkv = self.to_qkv(x) - q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3) - k = k.softmax(dim=-1) - context = torch.einsum('bhdn,bhen->bhde', k, v) - out = torch.einsum('bhde,bhdn->bhen', context, q) - out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) - return self.to_out(out) - - -class SpatialSelfAttention(nn.Module): - def __init__(self, in_channels, device=None, dtype=None, operations=ops): - super().__init__() - self.in_channels = in_channels - - self.norm = operations.GroupNorm( - num_groups=32, - num_channels=in_channels, - eps=1e-6, - affine=True, - device=device, - dtype=dtype - ) - self.q = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - self.k = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - self.v = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - self.proj_out = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - - def forward(self, x): - h_ = x - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) - - # compute attention - b,c,h,w = q.shape - q = rearrange(q, 'b c h w -> b (h w) c') - k = rearrange(k, 'b c h w -> b c (h w)') - w_ = torch.einsum('bij,bjk->bik', q, k) - - w_ = w_ * (int(c)**(-0.5)) - w_ = torch.nn.functional.softmax(w_, dim=2) - - # attend to values - v = rearrange(v, 'b c h w -> b c (h w)') - w_ = rearrange(w_, 'b i j -> b j i') - h_ = torch.einsum('bij,bjk->bik', v, w_) - h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h) - h_ = self.proj_out(h_) - - return x+h_ diff --git a/modules/encoders/condition.py b/modules/encoders/condition.py deleted file mode 100644 index 610322b..0000000 --- a/modules/encoders/condition.py +++ /dev/null @@ -1,389 +0,0 @@ -import torch -import torch.nn as nn -import kornia -import open_clip -from torch.utils.checkpoint import checkpoint -from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel -from ..common import autocast -from utils.utils import count_params - - -class AbstractEncoder(nn.Module): - def __init__(self): - super().__init__() - - def encode(self, *args, **kwargs): - raise NotImplementedError - - -class IdentityEncoder(AbstractEncoder): - def encode(self, x): - return x - - -class ClassEmbedder(nn.Module): - def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1): - super().__init__() - self.key = key - self.embedding = nn.Embedding(n_classes, embed_dim) - self.n_classes = n_classes - self.ucg_rate = ucg_rate - - def forward(self, batch, key=None, disable_dropout=False): - if key is None: - key = self.key - # this is for use in crossattn - c = batch[key][:, None] - if self.ucg_rate > 0. and not disable_dropout: - mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate) - c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1) - c = c.long() - c = self.embedding(c) - return c - - def get_unconditional_conditioning(self, bs, device="cuda"): - uc_class = self.n_classes - 1 # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000) - uc = torch.ones((bs,), device=device) * uc_class - uc = {self.key: uc} - return uc - - -def disabled_train(self, mode=True): - """Overwrite model.train with this function to make sure train/eval mode - does not change anymore.""" - return self - - -class FrozenT5Embedder(AbstractEncoder): - """Uses the T5 transformer encoder for text""" - - def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77, - freeze=True): # others are google/t5-v1_1-xl and google/t5-v1_1-xxl - super().__init__() - self.tokenizer = T5Tokenizer.from_pretrained(version) - self.transformer = T5EncoderModel.from_pretrained(version) - self.device = device - self.max_length = max_length # TODO: typical value? - if freeze: - self.freeze() - - def freeze(self): - self.transformer = self.transformer.eval() - # self.train = disabled_train - for param in self.parameters(): - param.requires_grad = False - - def forward(self, text): - batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True, - return_overflowing_tokens=False, padding="max_length", return_tensors="pt") - tokens = batch_encoding["input_ids"].to(self.device) - outputs = self.transformer(input_ids=tokens) - - z = outputs.last_hidden_state - return z - - def encode(self, text): - return self(text) - - -class FrozenCLIPEmbedder(AbstractEncoder): - """Uses the CLIP transformer encoder for text (from huggingface)""" - LAYERS = [ - "last", - "pooled", - "hidden" - ] - - def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77, - freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32 - super().__init__() - assert layer in self.LAYERS - self.tokenizer = CLIPTokenizer.from_pretrained(version) - self.transformer = CLIPTextModel.from_pretrained(version) - self.device = device - self.max_length = max_length - if freeze: - self.freeze() - self.layer = layer - self.layer_idx = layer_idx - if layer == "hidden": - assert layer_idx is not None - assert 0 <= abs(layer_idx) <= 12 - - def freeze(self): - self.transformer = self.transformer.eval() - # self.train = disabled_train - for param in self.parameters(): - param.requires_grad = False - - def forward(self, text): - batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True, - return_overflowing_tokens=False, padding="max_length", return_tensors="pt") - tokens = batch_encoding["input_ids"].to(self.device) - outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden") - if self.layer == "last": - z = outputs.last_hidden_state - elif self.layer == "pooled": - z = outputs.pooler_output[:, None, :] - else: - z = outputs.hidden_states[self.layer_idx] - return z - - def encode(self, text): - return self(text) - - -class ClipImageEmbedder(nn.Module): - def __init__( - self, - model, - jit=False, - device='cuda' if torch.cuda.is_available() else 'cpu', - antialias=True, - ucg_rate=0. - ): - super().__init__() - from clip import load as load_clip - self.model, _ = load_clip(name=model, device=device, jit=jit) - - self.antialias = antialias - - self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) - self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) - self.ucg_rate = ucg_rate - - def preprocess(self, x): - # normalize to [0,1] - x = kornia.geometry.resize(x, (224, 224), - interpolation='bicubic', align_corners=True, - antialias=self.antialias) - x = (x + 1.) / 2. - # re-normalize according to clip - x = kornia.enhance.normalize(x, self.mean, self.std) - return x - - def forward(self, x, no_dropout=False): - # x is assumed to be in range [-1,1] - out = self.model.encode_image(self.preprocess(x)) - out = out.to(x.dtype) - if self.ucg_rate > 0. and not no_dropout: - out = torch.bernoulli((1. - self.ucg_rate) * torch.ones(out.shape[0], device=out.device))[:, None] * out - return out - - -class FrozenOpenCLIPEmbedder(AbstractEncoder): - """ - Uses the OpenCLIP transformer encoder for text - """ - LAYERS = [ - # "pooled", - "last", - "penultimate" - ] - - def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, - freeze=True, layer="last"): - super().__init__() - assert layer in self.LAYERS - model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version) - del model.visual - self.model = model - - self.device = device - self.max_length = max_length - if freeze: - self.freeze() - self.layer = layer - if self.layer == "last": - self.layer_idx = 0 - elif self.layer == "penultimate": - self.layer_idx = 1 - else: - raise NotImplementedError() - - def freeze(self): - self.model = self.model.eval() - for param in self.parameters(): - param.requires_grad = False - - def forward(self, text): - tokens = open_clip.tokenize(text) ## all clip models use 77 as context length - z = self.encode_with_transformer(tokens.to(self.device)) - return z - - def encode_with_transformer(self, text): - x = self.model.token_embedding(text) # [batch_size, n_ctx, d_model] - x = x + self.model.positional_embedding - x = x.permute(1, 0, 2) # NLD -> LND - x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask) - x = x.permute(1, 0, 2) # LND -> NLD - x = self.model.ln_final(x) - return x - - def text_transformer_forward(self, x: torch.Tensor, attn_mask=None): - for i, r in enumerate(self.model.transformer.resblocks): - if i == len(self.model.transformer.resblocks) - self.layer_idx: - break - if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting(): - x = checkpoint(r, x, attn_mask) - else: - x = r(x, attn_mask=attn_mask) - return x - - def encode(self, text): - return self(text) - - -class FrozenOpenCLIPImageEmbedder(AbstractEncoder): - """ - Uses the OpenCLIP vision transformer encoder for images - """ - - def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, - freeze=True, layer="pooled", antialias=True, ucg_rate=0.): - super().__init__() - model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), - pretrained=version, ) - del model.transformer - self.model = model - # self.mapper = torch.nn.Linear(1280, 1024) - self.device = device - self.max_length = max_length - if freeze: - self.freeze() - self.layer = layer - if self.layer == "penultimate": - raise NotImplementedError() - self.layer_idx = 1 - - self.antialias = antialias - - self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) - self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) - self.ucg_rate = ucg_rate - - def preprocess(self, x): - # normalize to [0,1] - x = kornia.geometry.resize(x, (224, 224), - interpolation='bicubic', align_corners=True, - antialias=self.antialias) - x = (x + 1.) / 2. - # renormalize according to clip - x = kornia.enhance.normalize(x, self.mean, self.std) - return x - - def freeze(self): - self.model = self.model.eval() - for param in self.model.parameters(): - param.requires_grad = False - - @autocast - def forward(self, image, no_dropout=False): - z = self.encode_with_vision_transformer(image) - if self.ucg_rate > 0. and not no_dropout: - z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z - return z - - def encode_with_vision_transformer(self, img): - img = self.preprocess(img) - x = self.model.visual(img) - return x - - def encode(self, text): - return self(text) - -class FrozenOpenCLIPImageEmbedderV2(AbstractEncoder): - """ - Uses the OpenCLIP vision transformer encoder for images - """ - - def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", - freeze=True, layer="pooled", antialias=True): - super().__init__() - model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), - pretrained=version, ) - del model.transformer - self.model = model - self.device = device - - if freeze: - self.freeze() - self.layer = layer - if self.layer == "penultimate": - raise NotImplementedError() - self.layer_idx = 1 - - self.antialias = antialias - - self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False) - self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False) - - - def preprocess(self, x): - # normalize to [0,1] - x = kornia.geometry.resize(x, (224, 224), - interpolation='bicubic', align_corners=True, - antialias=self.antialias) - x = (x + 1.) / 2. - # renormalize according to clip - x = kornia.enhance.normalize(x, self.mean, self.std) - return x - - def freeze(self): - self.model = self.model.eval() - for param in self.model.parameters(): - param.requires_grad = False - - def forward(self, image, no_dropout=False): - ## image: b c h w - z = self.encode_with_vision_transformer(image) - return z - - def encode_with_vision_transformer(self, x): - x = self.preprocess(x) - - # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1 - if self.model.visual.input_patchnorm: - # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)') - x = x.reshape(x.shape[0], x.shape[1], self.model.visual.grid_size[0], self.model.visual.patch_size[0], self.model.visual.grid_size[1], self.model.visual.patch_size[1]) - x = x.permute(0, 2, 4, 1, 3, 5) - x = x.reshape(x.shape[0], self.model.visual.grid_size[0] * self.model.visual.grid_size[1], -1) - x = self.model.visual.patchnorm_pre_ln(x) - x = self.model.visual.conv1(x) - else: - x = self.model.visual.conv1(x) # shape = [*, width, grid, grid] - x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] - x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] - - # class embeddings and positional embeddings - x = torch.cat( - [self.model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), - x], dim=1) # shape = [*, grid ** 2 + 1, width] - x = x + self.model.visual.positional_embedding.to(x.dtype) - - # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in - x = self.model.visual.patch_dropout(x) - x = self.model.visual.ln_pre(x) - - x = x.permute(1, 0, 2) # NLD -> LND - x = self.model.visual.transformer(x) - x = x.permute(1, 0, 2) # LND -> NLD - - return x - -class FrozenCLIPT5Encoder(AbstractEncoder): - def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cuda", - clip_max_length=77, t5_max_length=77): - super().__init__() - self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length) - self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length) - print(f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder) * 1.e-6:.2f} M parameters, " - f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder) * 1.e-6:.2f} M params.") - - def encode(self, text): - return self(text) - - def forward(self, text): - clip_z = self.clip_encoder.encode(text) - t5_z = self.t5_encoder.encode(text) - return [clip_z, t5_z] diff --git a/modules/encoders/resampler.py b/modules/encoders/resampler.py deleted file mode 100644 index d9abd23..0000000 --- a/modules/encoders/resampler.py +++ /dev/null @@ -1,145 +0,0 @@ -# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py -# and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py -# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/resampler.py -import math -import torch -import torch.nn as nn - - -class ImageProjModel(nn.Module): - """Projection Model""" - def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): - super().__init__() - self.cross_attention_dim = cross_attention_dim - self.clip_extra_context_tokens = clip_extra_context_tokens - self.proj = nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim) - self.norm = nn.LayerNorm(cross_attention_dim) - - def forward(self, image_embeds): - #embeds = image_embeds - embeds = image_embeds.type(list(self.proj.parameters())[0].dtype) - clip_extra_context_tokens = self.proj(embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim) - clip_extra_context_tokens = self.norm(clip_extra_context_tokens) - return clip_extra_context_tokens - - -# FFN -def FeedForward(dim, mult=4): - inner_dim = int(dim * mult) - return nn.Sequential( - nn.LayerNorm(dim), - nn.Linear(dim, inner_dim, bias=False), - nn.GELU(), - nn.Linear(inner_dim, dim, bias=False), - ) - - -def reshape_tensor(x, heads): - bs, length, width = x.shape - #(bs, length, width) --> (bs, length, n_heads, dim_per_head) - x = x.view(bs, length, heads, -1) - # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) - x = x.transpose(1, 2) - # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) - x = x.reshape(bs, heads, length, -1) - return x - - -class PerceiverAttention(nn.Module): - def __init__(self, *, dim, dim_head=64, heads=8): - super().__init__() - self.scale = dim_head**-0.5 - self.dim_head = dim_head - self.heads = heads - inner_dim = dim_head * heads - - self.norm1 = nn.LayerNorm(dim) - self.norm2 = nn.LayerNorm(dim) - - self.to_q = nn.Linear(dim, inner_dim, bias=False) - self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) - self.to_out = nn.Linear(inner_dim, dim, bias=False) - - - def forward(self, x, latents): - """ - Args: - x (torch.Tensor): image features - shape (b, n1, D) - latent (torch.Tensor): latent features - shape (b, n2, D) - """ - x = self.norm1(x) - latents = self.norm2(latents) - - b, l, _ = latents.shape - - q = self.to_q(latents) - kv_input = torch.cat((x, latents), dim=-2) - k, v = self.to_kv(kv_input).chunk(2, dim=-1) - - q = reshape_tensor(q, self.heads) - k = reshape_tensor(k, self.heads) - v = reshape_tensor(v, self.heads) - - # attention - scale = 1 / math.sqrt(math.sqrt(self.dim_head)) - weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards - weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) - out = weight @ v - - out = out.permute(0, 2, 1, 3).reshape(b, l, -1) - - return self.to_out(out) - - -class Resampler(nn.Module): - def __init__( - self, - dim=1024, - depth=8, - dim_head=64, - heads=16, - num_queries=8, - embedding_dim=768, - output_dim=1024, - ff_mult=4, - video_length=None, # using frame-wise version or not - ): - super().__init__() - ## queries for a single frame / image - self.num_queries = num_queries - self.video_length = video_length - - ## queries for each frame - if video_length is not None: - num_queries = num_queries * video_length - - self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) - self.proj_in = nn.Linear(embedding_dim, dim) - self.proj_out = nn.Linear(dim, output_dim) - self.norm_out = nn.LayerNorm(output_dim) - - self.layers = nn.ModuleList([]) - for _ in range(depth): - self.layers.append( - nn.ModuleList( - [ - PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), - FeedForward(dim=dim, mult=ff_mult), - ] - ) - ) - - def forward(self, x): - latents = self.latents.repeat(x.size(0), 1, 1) ## B (T L) C - x = self.proj_in(x) - - for attn, ff in self.layers: - latents = attn(x, latents) + latents - latents = ff(latents) + latents - - latents = self.proj_out(latents) - latents = self.norm_out(latents) # B L C or B (T L) C - - return latents \ No newline at end of file diff --git a/modules/networks/ae_modules.py b/modules/networks/ae_modules.py deleted file mode 100644 index 35b6817..0000000 --- a/modules/networks/ae_modules.py +++ /dev/null @@ -1,1026 +0,0 @@ -# pytorch_diffusion + derived encoder decoder -import math -import torch -import numpy as np -import torch.nn as nn -from einops import rearrange -from utils.utils import instantiate_from_config -from ...modules.attention import LinearAttention - -import comfy.ops -ops = comfy.ops.disable_weight_init - -def nonlinearity(x): - # swish - return x*torch.sigmoid(x) - - -def Normalize(in_channels, num_groups=32, device=None, dtype=None, operations=ops): - return operations.GroupNorm( - num_groups=num_groups, - num_channels=in_channels, - eps=1e-6, - affine=True, - device=device, - dtype=dtype - ) - -class LinAttnBlock(LinearAttention): - """to match AttnBlock usage""" - def __init__(self, in_channels, device=None, dtype=None): - super().__init__(dim=in_channels, heads=1, dim_head=in_channels, device=device, dtype=dtype) - - -class AttnBlock(nn.Module): - def __init__(self, in_channels, device=None, dtype=None, operations=ops): - super().__init__() - self.in_channels = in_channels - - self.norm = Normalize(in_channels, device=device, dtype=dtype) - self.q = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - self.k = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - self.v = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - self.proj_out = operations.Conv2d( - in_channels, - in_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - - def forward(self, x): - h_ = x - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) - - # compute attention - b,c,h,w = q.shape - q = q.reshape(b,c,h*w) # bcl - q = q.permute(0,2,1) # bcl -> blc l=hw - k = k.reshape(b,c,h*w) # bcl - - w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] - w_ = w_ * (int(c)**(-0.5)) - w_ = torch.nn.functional.softmax(w_, dim=2) - - # attend to values - v = v.reshape(b,c,h*w) - w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q) - h_ = torch.bmm(v,w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] - h_ = h_.reshape(b,c,h,w) - - h_ = self.proj_out(h_) - - return x+h_ - -def make_attn(in_channels, attn_type="vanilla", device=None, dtype=None): - assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown' - #print(f"making attention of type '{attn_type}' with {in_channels} in_channels") - if attn_type == "vanilla": - return AttnBlock(in_channels, device=device, dtype=dtype) - elif attn_type == "none": - return nn.Identity(in_channels) - else: - return LinAttnBlock(in_channels, device=device, dtype=dtype) - -class Downsample(nn.Module): - def __init__(self, in_channels, with_conv, device=None, dtype=None, operations=ops): - super().__init__() - self.with_conv = with_conv - self.in_channels = in_channels - if self.with_conv: - # no asymmetric padding in torch conv, must do it ourselves - self.conv = operations.Conv2d( - in_channels, - in_channels, - kernel_size=3, - stride=2, - padding=0, - device=device, - dtype=dtype - ) - def forward(self, x): - if self.with_conv: - pad = (0,1,0,1) - x = torch.nn.functional.pad(x, pad, mode="constant", value=0) - x = self.conv(x) - else: - x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) - return x - -class Upsample(nn.Module): - def __init__(self, in_channels, with_conv, device=None, dtype=None, operations=ops): - super().__init__() - self.with_conv = with_conv - self.in_channels = in_channels - if self.with_conv: - self.conv = operations.Conv2d( - in_channels, - in_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype - ) - - def forward(self, x): - x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") - if self.with_conv: - x = self.conv(x) - return x - -def get_timestep_embedding(timesteps, embedding_dim): - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: - From Fairseq. - Build sinusoidal embeddings. - This matches the implementation in tensor2tensor, but differs slightly - from the description in Section 3.5 of "Attention Is All You Need". - """ - assert len(timesteps.shape) == 1 - - half_dim = embedding_dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) - emb = emb.to(device=timesteps.device) - emb = timesteps.float()[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0,1,0,0)) - return emb - - - -class ResnetBlock(nn.Module): - def __init__( - self, - *, - in_channels, - out_channels=None, - conv_shortcut=False, - dropout, - temb_channels=512, - device=None, - dtype=None, - operations=ops - ): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.use_conv_shortcut = conv_shortcut - - self.norm1 = Normalize(in_channels, device=device, dtype=dtype) - self.conv1 = operations.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype - ) - if temb_channels > 0: - self.temb_proj = operations.Linear( - temb_channels, - out_channels, - device=device, - dtype=dtype - ) - self.norm2 = Normalize(out_channels, device=device, dtype=dtype) - self.dropout = torch.nn.Dropout(dropout) - self.conv2 = operations.Conv2d( - out_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype - ) - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = operations.Conv2d( - in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=device - ) - else: - self.nin_shortcut = operations.Conv2d( - in_channels, - out_channels, - kernel_size=1, - stride=1, - padding=0, - device=device, - dtype=dtype - ) - - def forward(self, x, temb): - h = x - h = self.norm1(h) - h = nonlinearity(h) - h = self.conv1(h) - - if temb is not None: - h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None] - - h = self.norm2(h) - h = nonlinearity(h) - h = self.dropout(h) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - x = self.conv_shortcut(x) - else: - x = self.nin_shortcut(x) - - return x+h - -class Model(nn.Module): - def __init__( - self, - *, - ch, - out_ch, - ch_mult=(1,2,4,8), - num_res_blocks, - attn_resolutions, - dropout=0.0, - resamp_with_conv=True, - in_channels, - resolution, - use_timestep=True, - use_linear_attn=False, - attn_type="vanilla", - device=None, - dtype=None, - operations=ops - ): - super().__init__() - if use_linear_attn: attn_type = "linear" - self.ch = ch - self.temb_ch = self.ch*4 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - - self.use_timestep = use_timestep - if self.use_timestep: - # timestep embedding - self.temb = nn.Module() - self.temb.dense = nn.ModuleList([ - operations.Linear( - self.ch, - self.temb_ch, - device=device, - dtype=dtype - ), - operations.Linear( - self.temb_ch, - self.temb_ch, - device=device, - dtype=dtype - ), - ]) - - # downsampling - self.conv_in = operations.Conv2d( - in_channels, - self.ch, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype - ) - - curr_res = resolution - in_ch_mult = (1,)+tuple(ch_mult) - self.down = nn.ModuleList() - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = ch*in_ch_mult[i_level] - block_out = ch*ch_mult[i_level] - for i_block in range(self.num_res_blocks): - block.append(ResnetBlock(in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype)) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions-1: - down.downsample = Downsample(block_in, resamp_with_conv, device=device, dtype=dtype) - curr_res = curr_res // 2 - self.down.append(down) - - # middle - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock(in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype) - self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype) - self.mid.block_2 = ResnetBlock(in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype) - - # upsampling - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch*ch_mult[i_level] - skip_in = ch*ch_mult[i_level] - for i_block in range(self.num_res_blocks+1): - if i_block == self.num_res_blocks: - skip_in = ch*in_ch_mult[i_level] - block.append(ResnetBlock(in_channels=block_in+skip_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype)) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv, device=device, dtype=dtype) - curr_res = curr_res * 2 - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = Normalize(block_in, device=device, dtype=device) - self.conv_out = torch.nn.Conv2d(block_in, - out_ch, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - def forward(self, x, t=None, context=None): - #assert x.shape[2] == x.shape[3] == self.resolution - if context is not None: - # assume aligned context, cat along channel axis - x = torch.cat((x, context), dim=1) - if self.use_timestep: - # timestep embedding - assert t is not None - temb = get_timestep_embedding(t, self.ch) - temb = self.temb.dense[0](temb) - temb = nonlinearity(temb) - temb = self.temb.dense[1](temb) - else: - temb = None - - # downsampling - hs = [self.conv_in(x)] - for i_level in range(self.num_resolutions): - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](hs[-1], temb) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - hs.append(h) - if i_level != self.num_resolutions-1: - hs.append(self.down[i_level].downsample(hs[-1])) - - # middle - h = hs[-1] - h = self.mid.block_1(h, temb) - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb) - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks+1): - h = self.up[i_level].block[i_block]( - torch.cat([h, hs.pop()], dim=1), temb) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - if i_level != 0: - h = self.up[i_level].upsample(h) - - # end - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - return h - - def get_last_layer(self): - return self.conv_out.weight - - -class Encoder(nn.Module): - def __init__( - self, - *, - ch, - out_ch, - ch_mult=(1,2,4,8), - num_res_blocks, - attn_resolutions, - dropout=0.0, - resamp_with_conv=True, - in_channels, - resolution, - z_channels, - double_z=True, - use_linear_attn=False, - attn_type="vanilla", - device=None, - dtype=None, - operations=ops, - **ignore_kwargs - ): - super().__init__() - if use_linear_attn: attn_type = "linear" - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - - # downsampling - self.conv_in = operations.Conv2d(in_channels, - self.ch, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - curr_res = resolution - in_ch_mult = (1,)+tuple(ch_mult) - self.in_ch_mult = in_ch_mult - self.down = nn.ModuleList() - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = ch*in_ch_mult[i_level] - block_out = ch*ch_mult[i_level] - for i_block in range(self.num_res_blocks): - block.append(ResnetBlock(in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype)) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions-1: - down.downsample = Downsample(block_in, resamp_with_conv, device=device, dtype=dtype) - curr_res = curr_res // 2 - self.down.append(down) - - # middle - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock(in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype) - self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype) - self.mid.block_2 = ResnetBlock(in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype) - - # end - self.norm_out = Normalize(block_in, device=device, dtype=dtype) - self.conv_out = operations.Conv2d(block_in, - 2*z_channels if double_z else z_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - def forward(self, x): - # timestep embedding - temb = None - - # print(f'encoder-input={x.shape}') - # downsampling - hs = [self.conv_in(x)] - # print(f'encoder-conv in feat={hs[0].shape}') - for i_level in range(self.num_resolutions): - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](hs[-1], temb) - # print(f'encoder-down feat={h.shape}') - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - hs.append(h) - if i_level != self.num_resolutions-1: - # print(f'encoder-downsample (input)={hs[-1].shape}') - hs.append(self.down[i_level].downsample(hs[-1])) - # print(f'encoder-downsample (output)={hs[-1].shape}') - - # middle - h = hs[-1] - h = self.mid.block_1(h, temb) - # print(f'encoder-mid1 feat={h.shape}') - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb) - # print(f'encoder-mid2 feat={h.shape}') - - # end - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - # print(f'end feat={h.shape}') - return h - - -class Decoder(nn.Module): - def __init__( - self, - *, - ch, - out_ch, - ch_mult=(1,2,4,8), - num_res_blocks, - attn_resolutions, - dropout=0.0, - resamp_with_conv=True, - in_channels, - resolution, - z_channels, - give_pre_end=False, - tanh_out=False, - use_linear_attn=False, - attn_type="vanilla", - device=None, - dtype=None, - operations=ops, - **ignorekwargs - ): - super().__init__() - if use_linear_attn: attn_type = "linear" - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.give_pre_end = give_pre_end - self.tanh_out = tanh_out - - # compute in_ch_mult, block_in and curr_res at lowest res - in_ch_mult = (1,)+tuple(ch_mult) - block_in = ch*ch_mult[self.num_resolutions-1] - curr_res = resolution // 2**(self.num_resolutions-1) - self.z_shape = (1,z_channels,curr_res,curr_res) - print("AE working on z of shape {} = {} dimensions.".format( - self.z_shape, np.prod(self.z_shape))) - - # z to block_in - self.conv_in = torch.nn.Conv2d(z_channels, - block_in, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - # middle - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock(in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype) - self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype) - self.mid.block_2 = ResnetBlock(in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype) - - # upsampling - self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = ch*ch_mult[i_level] - for i_block in range(self.num_res_blocks+1): - block.append(ResnetBlock(in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype)) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(make_attn(block_in, attn_type=attn_type, device=device, dtype=dtype)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv, device=device, dtype=dtype) - curr_res = curr_res * 2 - self.up.insert(0, up) # prepend to get consistent order - - # end - self.norm_out = Normalize(block_in, device=device, dtype=dtype) - self.conv_out = operations.Conv2d(block_in, - out_ch, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - def forward(self, z): - #assert z.shape[1:] == self.z_shape[1:] - self.last_z_shape = z.shape - - # print(f'decoder-input={z.shape}') - # timestep embedding - temb = None - - # z to block_in - h = self.conv_in(z) - # print(f'decoder-conv in feat={h.shape}') - - # middle - h = self.mid.block_1(h, temb) - h = self.mid.attn_1(h) - h = self.mid.block_2(h, temb) - # print(f'decoder-mid feat={h.shape}') - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks+1): - h = self.up[i_level].block[i_block](h, temb) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - # print(f'decoder-up feat={h.shape}') - if i_level != 0: - h = self.up[i_level].upsample(h) - # print(f'decoder-upsample feat={h.shape}') - - # end - if self.give_pre_end: - return h - - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - # print(f'decoder-conv_out feat={h.shape}') - if self.tanh_out: - h = torch.tanh(h) - return h - - -class SimpleDecoder(nn.Module): - def __init__(self, in_channels, out_channels, device=None, dtype=None, operations=ops, *args, **kwargs): - super().__init__() - self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1), - ResnetBlock(in_channels=in_channels, - out_channels=2 * in_channels, - temb_channels=0, dropout=0.0, - device=device, - dtype=dtype), - ResnetBlock(in_channels=2 * in_channels, - out_channels=4 * in_channels, - temb_channels=0, dropout=0.0, - device=device, - dtype=dtype), - ResnetBlock(in_channels=4 * in_channels, - out_channels=2 * in_channels, - temb_channels=0, dropout=0.0, - device=device, - dtype=dtype), - operations.Conv2d(2*in_channels, in_channels, 1), - Upsample(in_channels, with_conv=True, device=device, dtype=dtype)]) - # end - self.norm_out = Normalize(in_channels, device=device, dtype=dtype) - self.conv_out = torch.nn.Conv2d(in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - def forward(self, x): - for i, layer in enumerate(self.model): - if i in [1,2,3]: - x = layer(x, None) - else: - x = layer(x) - - h = self.norm_out(x) - h = nonlinearity(h) - x = self.conv_out(h) - return x - - -class UpsampleDecoder(nn.Module): - def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution, - ch_mult=(2,2), dropout=0.0, device=None, dtype=None, operations=ops): - super().__init__() - # upsampling - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - block_in = in_channels - curr_res = resolution // 2 ** (self.num_resolutions - 1) - self.res_blocks = nn.ModuleList() - self.upsample_blocks = nn.ModuleList() - for i_level in range(self.num_resolutions): - res_block = [] - block_out = ch * ch_mult[i_level] - for i_block in range(self.num_res_blocks + 1): - res_block.append(ResnetBlock(in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - device=device, - dtype=dtype)) - block_in = block_out - self.res_blocks.append(nn.ModuleList(res_block)) - if i_level != self.num_resolutions - 1: - self.upsample_blocks.append(Upsample(block_in, True, device=device, dtype=dtype)) - curr_res = curr_res * 2 - - # end - self.norm_out = Normalize(block_in, device=device, dtype=dtype) - self.conv_out = torch.nn.Conv2d(block_in, - out_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - - def forward(self, x): - # upsampling - h = x - for k, i_level in enumerate(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - h = self.res_blocks[i_level][i_block](h, None) - if i_level != self.num_resolutions - 1: - h = self.upsample_blocks[k](h) - h = self.norm_out(h) - h = nonlinearity(h) - h = self.conv_out(h) - return h - - -class LatentRescaler(nn.Module): - def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2, device=None, dtype=None, operations=ops): - super().__init__() - # residual block, interpolate, residual block - self.factor = factor - self.conv_in = operations.Conv2d(in_channels, - mid_channels, - kernel_size=3, - stride=1, - padding=1, - device=device, - dtype=dtype) - self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, - out_channels=mid_channels, - temb_channels=0, - dropout=0.0, - device=device, - dtype=dtype) for _ in range(depth)]) - self.attn = AttnBlock(mid_channels, device=device, dtype=dtype) - self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, - out_channels=mid_channels, - temb_channels=0, - dropout=0.0, - device=device, - dtype=dtype) for _ in range(depth)]) - - self.conv_out = operations.Conv2d(mid_channels, - out_channels, - kernel_size=1, - device=device, - dtype=dtype - ) - - def forward(self, x): - x = self.conv_in(x) - for block in self.res_block1: - x = block(x, None) - x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor)))) - x = self.attn(x) - for block in self.res_block2: - x = block(x, None) - x = self.conv_out(x) - return x - - -class MergedRescaleEncoder(nn.Module): - def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks, - attn_resolutions, dropout=0.0, resamp_with_conv=True, - ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1, device=None, dtype=None, operations=ops): - super().__init__() - intermediate_chn = ch * ch_mult[-1] - self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult, - z_channels=intermediate_chn, double_z=False, resolution=resolution, - attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv, - out_ch=None, device=device, dtype=dtype) - self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn, - mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth, - device=device, dtype=dtype) - - def forward(self, x): - x = self.encoder(x) - x = self.rescaler(x) - return x - - -class MergedRescaleDecoder(nn.Module): - def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8), - dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1, - device=None, dtype=None, operations=ops): - super().__init__() - tmp_chn = z_channels*ch_mult[-1] - self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout, - resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks, - ch_mult=ch_mult, resolution=resolution, ch=ch, device=device, operations=ops) - self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn, - out_channels=tmp_chn, depth=rescale_module_depth, device=device, operations=ops) - - def forward(self, x): - x = self.rescaler(x) - x = self.decoder(x) - return x - - -class Upsampler(nn.Module): - def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2, device=None, dtype=None, operations=ops): - super().__init__() - assert out_size >= in_size - num_blocks = int(np.log2(out_size//in_size))+1 - factor_up = 1.+ (out_size % in_size) - print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}") - self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels, - out_channels=in_channels, device=device, dtype=dtype) - self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2, - attn_resolutions=[], in_channels=None, ch=in_channels, device=device, dtype=dtype, - ch_mult=[ch_mult for _ in range(num_blocks)]) - - def forward(self, x): - x = self.rescaler(x) - x = self.decoder(x) - return x - - -class Resize(nn.Module): - def __init__(self, in_channels=None, learned=False, mode="bilinear", device=None, dtype=None, operations=ops): - super().__init__() - self.with_conv = learned - self.mode = mode - if self.with_conv: - print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode") - raise NotImplementedError() - assert in_channels is not None - # no asymmetric padding in torch conv, must do it ourselves - self.conv = operations.Conv2d(in_channels, - in_channels, - kernel_size=4, - stride=2, - padding=1, - device=device, - dtype=dtype) - - def forward(self, x, scale_factor=1.0): - if scale_factor==1.0: - return x - else: - x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor) - return x - -class FirstStagePostProcessor(nn.Module): - - def __init__(self, ch_mult:list, in_channels, - pretrained_model:nn.Module=None, - reshape=False, - n_channels=None, - dropout=0., - pretrained_config=None, - device=None, - dtype=None, - operations=ops): - super().__init__() - if pretrained_config is None: - assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' - self.pretrained_model = pretrained_model - else: - assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' - self.instantiate_pretrained(pretrained_config) - - self.do_reshape = reshape - - if n_channels is None: - n_channels = self.pretrained_model.encoder.ch - - self.proj_norm = Normalize(in_channels,num_groups=in_channels//2, device=device, dtype=dtype) - self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3, - stride=1,padding=1, device=device, dtype=dtype) - - blocks = [] - downs = [] - ch_in = n_channels - for m in ch_mult: - blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout, device=device, dtype=dtype)) - ch_in = m * n_channels - downs.append(Downsample(ch_in, with_conv=False, device=device, dtype=dtype)) - - self.model = nn.ModuleList(blocks) - self.downsampler = nn.ModuleList(downs) - - - def instantiate_pretrained(self, config): - model = instantiate_from_config(config) - self.pretrained_model = model.eval() - # self.pretrained_model.train = False - for param in self.pretrained_model.parameters(): - param.requires_grad = False - - - @torch.no_grad() - def encode_with_pretrained(self,x): - c = self.pretrained_model.encode(x) - if isinstance(c, DiagonalGaussianDistribution): - c = c.mode() - return c - - def forward(self,x): - z_fs = self.encode_with_pretrained(x) - z = self.proj_norm(z_fs) - z = self.proj(z) - z = nonlinearity(z) - - for submodel, downmodel in zip(self.model,self.downsampler): - z = submodel(z,temb=None) - z = downmodel(z) - - if self.do_reshape: - z = rearrange(z,'b c h w -> b (h w) c') - return z \ No newline at end of file diff --git a/modules/networks/openaimodel3d.py b/modules/networks/openaimodel3d.py deleted file mode 100644 index 838b6fe..0000000 --- a/modules/networks/openaimodel3d.py +++ /dev/null @@ -1,822 +0,0 @@ -from functools import partial -from abc import abstractmethod -import torch -import torch.nn as nn -from einops import rearrange -import torch.nn.functional as F -from ...models.utils_diffusion import timestep_embedding -from ...common import checkpoint -from ...basics import ( - zero_module, - conv_nd, - linear, - avg_pool_nd, - normalization -) -from ...modules.attention import SpatialTransformer, TemporalTransformer -import comfy.ops -import logging - -ops = comfy.ops.disable_weight_init - -class TimestepBlock(nn.Module): - """ - Any module where forward() takes timestep embeddings as a second argument. - """ - @abstractmethod - def forward(self, x, emb): - """ - Apply the module to `x` given `emb` timestep embeddings. - """ - -#This is needed because accelerate makes a copy of transformer_options which breaks "transformer_index" -def forward_timestep_embed(ts, x, emb, context=None, batch_size=None, transformer_options={}): - for layer in ts: - if isinstance(layer, TimestepBlock): - x = layer(x, emb, batch_size=batch_size) - elif isinstance(layer, SpatialTransformer): - x = layer(x, context) - if "transformer_index" in transformer_options: - transformer_options["transformer_index"] += 1 - elif isinstance(layer, TemporalTransformer): - x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size) - x = layer(x, context) - if "transformer_index" in transformer_options: - transformer_options["transformer_index"] += 1 - x = rearrange(x, 'b c f h w -> (b f) c h w') - else: - x = layer(x) - return x - -class TimestepEmbedSequential(nn.Sequential, TimestepBlock): - """ - A sequential module that passes timestep embeddings to the children that - support it as an extra input. - """ - - def forward(self, *args, **kwargs): - return forward_timestep_embed(self, *args, **kwargs) - -class Downsample(nn.Module): - """ - A downsampling layer with an optional convolution. - :param channels: channels in the inputs and outputs. - :param use_conv: a bool determining if a convolution is applied. - :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then - downsampling occurs in the inner-two dimensions. - """ - - def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=ops): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.dims = dims - stride = 2 if dims != 3 else (1, 2, 2) - if use_conv: - self.op = operations.conv_nd( - dims, self.channels, self.out_channels, 3, stride=stride, padding=padding - ) - else: - assert self.channels == self.out_channels - self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) - - def forward(self, x): - assert x.shape[1] == self.channels - return self.op(x) - -class Upsample(nn.Module): - """ - An upsampling layer with an optional convolution. - :param channels: channels in the inputs and outputs. - :param use_conv: a bool determining if a convolution is applied. - :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then - upsampling occurs in the inner-two dimensions. - """ - - def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=ops): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.dims = dims - if use_conv: - self.conv = operations.conv_nd(dims, self.channels, self.out_channels, 3, padding=padding, dtype=dtype, device=device) - - def forward(self, x): - assert x.shape[1] == self.channels - if self.dims == 3: - x = F.interpolate(x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode='nearest') - else: - x = F.interpolate(x, scale_factor=2, mode='nearest') - if self.use_conv: - x = self.conv(x) - return x - -class ResBlock(TimestepBlock): - """ - A residual block that can optionally change the number of channels. - :param channels: the number of input channels. - :param emb_channels: the number of timestep embedding channels. - :param dropout: the rate of dropout. - :param out_channels: if specified, the number of out channels. - :param use_conv: if True and out_channels is specified, use a spatial - convolution instead of a smaller 1x1 convolution to change the - channels in the skip connection. - :param dims: determines if the signal is 1D, 2D, or 3D. - :param up: if True, use this block for upsampling. - :param down: if True, use this block for downsampling. - :param use_temporal_conv: if True, use the temporal convolution. - :param use_image_dataset: if True, the temporal parameters will not be optimized. - """ - - def __init__( - self, - channels, - emb_channels, - dropout, - out_channels=None, - use_scale_shift_norm=False, - dims=2, - use_checkpoint=False, - use_conv=False, - up=False, - down=False, - kernel_size=3, - use_temporal_conv=False, - tempspatial_aware=False, - dtype=None, - device=None, - operations=ops - ): - super().__init__() - self.channels = channels - self.emb_channels = emb_channels - self.dropout = dropout - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.use_checkpoint = use_checkpoint - self.use_scale_shift_norm = use_scale_shift_norm - self.use_temporal_conv = use_temporal_conv - - if isinstance(kernel_size, list): - padding =[k // 2 for k in kernel_size] - else: - padding = kernel_size // 2 - - # operations used in normalization function - self.in_layers = nn.Sequential( - normalization(channels, dtype=dtype, device=device), - nn.SiLU(), - operations.conv_nd(dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device), - ) - - self.updown = up or down - - if up: - self.h_upd = Upsample(channels, False, dims, dtype=dtype, device=device) - self.x_upd = Upsample(channels, False, dims, dtype=dtype, device=device) - elif down: - self.h_upd = Downsample(channels, False, dims, dtype=dtype, device=device) - self.x_upd = Downsample(channels, False, dims, dtype=dtype, device=device) - else: - self.h_upd = self.x_upd = nn.Identity() - - self.emb_layers = nn.Sequential( - nn.SiLU(), - operations.Linear( - emb_channels, - 2 * self.out_channels if use_scale_shift_norm else self.out_channels, - dtype=dtype, - device=device - ), - ) - self.out_layers = nn.Sequential( - normalization(self.out_channels, dtype=dtype, device=device), - nn.SiLU(), - nn.Dropout(p=dropout), - zero_module(operations.Conv2d(self.out_channels, self.out_channels, 3, padding=1, dtype=dtype, device=device)), - ) - - if self.out_channels == channels: - self.skip_connection = nn.Identity() - elif use_conv: - self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 3, padding=1, dtype=dtype, device=device) - else: - self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 1, dtype=dtype, device=device) - - if self.use_temporal_conv: - self.temopral_conv = TemporalConvBlock( - self.out_channels, - self.out_channels, - dropout=0.1, - spatial_aware=tempspatial_aware, - dtype=dtype, - device=device - ) - - def forward(self, x, emb, batch_size=None): - """ - Apply the block to a Tensor, conditioned on a timestep embedding. - :param x: an [N x C x ...] Tensor of features. - :param emb: an [N x emb_channels] Tensor of timestep embeddings. - :return: an [N x C x ...] Tensor of outputs. - """ - input_tuple = (x, emb) - if batch_size: - forward_batchsize = partial(self._forward, batch_size=batch_size) - return checkpoint(forward_batchsize, input_tuple, self.parameters(), self.use_checkpoint) - return checkpoint(self._forward, input_tuple, self.parameters(), self.use_checkpoint) - - def _forward(self, x, emb, batch_size=None): - if self.updown: - in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] - h = in_rest(x) - h = self.h_upd(h) - x = self.x_upd(x) - h = in_conv(h) - else: - h = self.in_layers(x) - emb_out = self.emb_layers(emb).type(h.dtype) - while len(emb_out.shape) < len(h.shape): - emb_out = emb_out[..., None] - if self.use_scale_shift_norm: - out_norm, out_rest = self.out_layers[0], self.out_layers[1:] - scale, shift = torch.chunk(emb_out, 2, dim=1) - h = out_norm(h) * (1 + scale) + shift - h = out_rest(h) - else: - h = h + emb_out - h = self.out_layers(h) - h = self.skip_connection(x) + h - - if self.use_temporal_conv and batch_size: - h = rearrange(h, '(b t) c h w -> b c t h w', b=batch_size) - h = self.temopral_conv(h) - h = rearrange(h, 'b c t h w -> (b t) c h w') - return h - -class TemporalConvBlock(nn.Module): - """ - Adapted from modelscope: https://github.com/modelscope/modelscope/blob/master/modelscope/models/multi_modal/video_synthesis/unet_sd.py - """ - def __init__( - self, - in_channels, - out_channels=None, - dropout=0.0, - spatial_aware=False, - dtype=None, - device=None, - operations=ops - ): - super(TemporalConvBlock, self).__init__() - if out_channels is None: - out_channels = in_channels - self.in_channels = in_channels - self.out_channels = out_channels - th_kernel_shape = (3, 1, 1) if not spatial_aware else (3, 3, 1) - th_padding_shape = (1, 0, 0) if not spatial_aware else (1, 1, 0) - tw_kernel_shape = (3, 1, 1) if not spatial_aware else (3, 1, 3) - tw_padding_shape = (1, 0, 0) if not spatial_aware else (1, 0, 1) - - # conv layers - self.conv1 = nn.Sequential( - operations.GroupNorm(32, in_channels, device=device, dtype=dtype), nn.SiLU(), - operations.Conv3d(in_channels, out_channels, th_kernel_shape, padding=th_padding_shape, device=device, dtype=dtype)) - self.conv2 = nn.Sequential( - operations.GroupNorm(32, out_channels, device=device, dtype=dtype), nn.SiLU(), nn.Dropout(dropout), - operations.Conv3d(out_channels, in_channels, tw_kernel_shape, padding=tw_padding_shape, device=device, dtype=dtype)) - self.conv3 = nn.Sequential( - operations.GroupNorm(32, out_channels, device=device, dtype=dtype), nn.SiLU(), nn.Dropout(dropout), - operations.Conv3d(out_channels, in_channels, th_kernel_shape, padding=th_padding_shape, device=device, dtype=dtype)) - self.conv4 = nn.Sequential( - operations.GroupNorm(32, out_channels, device=device, dtype=dtype), nn.SiLU(), nn.Dropout(dropout), - operations.Conv3d(out_channels, in_channels, tw_kernel_shape, padding=tw_padding_shape, device=device, dtype=dtype)) - - # zero out the last layer params,so the conv block is identity - nn.init.zeros_(self.conv4[-1].weight) - nn.init.zeros_(self.conv4[-1].bias) - - def forward(self, x): - identity = x - x = self.conv1(x) - x = self.conv2(x) - x = self.conv3(x) - x = self.conv4(x) - - return identity + x - -def context_processor(context, t, img_emb=None, temporal_size=16, concat_only=False, disable_concat=False): - if disable_concat: - return context - - ## repeat t times for context [(b t) 77 768] & time embedding - ## check if we use per-frame image conditioning - - if img_emb is not None: - context = torch.cat([context, img_emb.to(context.device, context.dtype)], dim=1) - - if concat_only: - return context - - b, l_context, _ = context.shape - if l_context == 77 + t * temporal_size: - context_text, context_img = context[:,:77,:], context[:,77:,:] - context_text = context_text.repeat_interleave(repeats=t, dim=0) - context_img = rearrange(context_img, 'b (t l) c -> (b t) l c', t=t) - context = torch.cat([context_text, context_img], dim=1) - else: - context = context.repeat_interleave(repeats=t, dim=0) - - return context - -def apply_control(h, control, name, cond_idx=None): - if control is not None and name in control and len(control[name]) > 0: - frames = h.shape[0] - ctrl = control[name].pop() - if ctrl is not None: - try: - if cond_idx is not None and ctrl.shape[0] > frames: - ctrl_frames_list = list(range(ctrl.shape[0])) - ctrl_frames = len(ctrl_frames_list) - - idxs = ( - ctrl_frames_list[:ctrl_frames // 2] if cond_idx == 0 else \ - ctrl_frames_list[ctrl_frames // 2:] - ) - - ctrl = ctrl[idxs] - - h += ctrl - except Exception as e: - if h.shape != ctrl.shape: - logging.warning( - "warning control could not be applied {} {}".format(h.shape, ctrl.shape) - ) - logging.warning(e) - return h - -class UNetModel(nn.Module): - """ - The full UNet model with attention and timestep embedding. - :param in_channels: in_channels in the input Tensor. - :param model_channels: base channel count for the model. - :param out_channels: channels in the output Tensor. - :param num_res_blocks: number of residual blocks per downsample. - :param attention_resolutions: a collection of downsample rates at which - attention will take place. May be a set, list, or tuple. - For example, if this contains 4, then at 4x downsampling, attention - will be used. - :param dropout: the dropout probability. - :param channel_mult: channel multiplier for each level of the UNet. - :param conv_resample: if True, use learned convolutions for upsampling and - downsampling. - :param dims: determines if the signal is 1D, 2D, or 3D. - :param num_classes: if specified (as an int), then this model will be - class-conditional with `num_classes` classes. - :param use_checkpoint: use gradient checkpointing to reduce memory usage. - :param num_heads: the number of attention heads in each attention layer. - :param num_heads_channels: if specified, ignore num_heads and instead use - a fixed channel width per attention head. - :param num_heads_upsample: works with num_heads to set a different number - of heads for upsampling. Deprecated. - :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. - :param resblock_updown: use residual blocks for up/downsampling. - :param use_new_attention_order: use a different attention pattern for potentially - increased efficiency. - """ - - def __init__(self, - in_channels, - model_channels, - out_channels, - num_res_blocks, - attention_resolutions, - dropout=0.0, - channel_mult=(1, 2, 4, 8), - conv_resample=True, - dims=2, - context_dim=None, - use_scale_shift_norm=False, - resblock_updown=False, - num_heads=-1, - num_head_channels=-1, - transformer_depth=1, - use_linear=False, - use_checkpoint=False, - temporal_conv=False, - tempspatial_aware=False, - temporal_attention=True, - use_relative_position=True, - use_causal_attention=False, - temporal_length=None, - use_fp16=False, - addition_attention=False, - temporal_selfatt_only=True, - image_cross_attention=False, - image_cross_attention_scale_learnable=False, - default_fs=4, - fs_condition=False, - device=None, - dtype=torch.float16, - operations=ops - ): - super(UNetModel, self).__init__() - if num_heads == -1: - assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' - if num_head_channels == -1: - assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' - - self.in_channels = in_channels - self.model_channels = model_channels - self.out_channels = out_channels - self.num_res_blocks = num_res_blocks - self.attention_resolutions = attention_resolutions - self.dropout = dropout - self.channel_mult = channel_mult - self.conv_resample = conv_resample - self.temporal_attention = temporal_attention - time_embed_dim = model_channels * 4 - self.use_checkpoint = use_checkpoint - temporal_self_att_only = True - self.addition_attention = addition_attention - self.temporal_length = temporal_length - self.image_cross_attention = image_cross_attention - self.image_cross_attention_scale_learnable = image_cross_attention_scale_learnable - self.default_fs = default_fs - self.fs_condition = fs_condition - self.device = device - #self.dtype = dtype - self.dtype = torch.float32 - - ## Time embedding blocks - self.time_embed = nn.Sequential( - linear(model_channels, time_embed_dim, device=device, dtype=self.dtype), - nn.SiLU(), - linear(time_embed_dim, time_embed_dim, device=device, dtype=self.dtype), - ) - if fs_condition: - self.fps_embedding = nn.Sequential( - linear(model_channels, time_embed_dim, device=device, dtype=self.dtype), - nn.SiLU(), - linear(time_embed_dim, time_embed_dim, device=device, dtype=self.dtype), - ) - nn.init.zeros_(self.fps_embedding[-1].weight) - nn.init.zeros_(self.fps_embedding[-1].bias) - ## Input Block - self.input_blocks = nn.ModuleList( - [ - TimestepEmbedSequential( - operations.conv_nd( - dims, - in_channels, - model_channels, - 3, - padding=1, - device=device, - dtype=self.dtype - )) - ] - ) - if self.addition_attention: - self.init_attn=TimestepEmbedSequential( - TemporalTransformer( - model_channels, - n_heads=8, - d_head=num_head_channels, - depth=transformer_depth, - context_dim=context_dim, - use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only, - causal_attention=False, relative_position=use_relative_position, - temporal_length=temporal_length, - device=device, - dtype=self.dtype - )) - - input_block_chans = [model_channels] - ch = model_channels - ds = 1 - for level, mult in enumerate(channel_mult): - for _ in range(num_res_blocks): - layers = [ - ResBlock(ch, time_embed_dim, dropout, - out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint, - use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, - use_temporal_conv=temporal_conv, - device=device, - dtype=self.dtype - ) - ] - ch = mult * model_channels - if ds in attention_resolutions: - if num_head_channels == -1: - dim_head = ch // num_heads - else: - num_heads = ch // num_head_channels - dim_head = num_head_channels - layers.append( - SpatialTransformer(ch, num_heads, dim_head, - depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, - use_checkpoint=use_checkpoint, disable_self_attn=False, - video_length=temporal_length, image_cross_attention=self.image_cross_attention, - image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable, - device=device, - dtype=self.dtype - ) - ) - if self.temporal_attention: - layers.append( - TemporalTransformer(ch, num_heads, dim_head, - depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, - use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only, - causal_attention=use_causal_attention, relative_position=use_relative_position, - temporal_length=temporal_length, - device=device, - dtype=self.dtype - ) - ) - self.input_blocks.append(TimestepEmbedSequential(*layers)) - input_block_chans.append(ch) - if level != len(channel_mult) - 1: - out_ch = ch - self.input_blocks.append( - TimestepEmbedSequential( - ResBlock(ch, time_embed_dim, dropout, - out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, - use_scale_shift_norm=use_scale_shift_norm, - down=True, - device=device, - dtype=self.dtype - ) - if resblock_updown - else Downsample( - ch, - conv_resample, - dims=dims, - out_channels=out_ch, - device=device, - dtype=self.dtype - ) - ) - ) - ch = out_ch - input_block_chans.append(ch) - ds *= 2 - - if num_head_channels == -1: - dim_head = ch // num_heads - else: - num_heads = ch // num_head_channels - dim_head = num_head_channels - layers = [ - ResBlock(ch, time_embed_dim, dropout, - dims=dims, use_checkpoint=use_checkpoint, - use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, - use_temporal_conv=temporal_conv, - device=device, - dtype=self.dtype - ), - SpatialTransformer(ch, num_heads, dim_head, - depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, - use_checkpoint=use_checkpoint, disable_self_attn=False, video_length=temporal_length, - image_cross_attention=self.image_cross_attention,image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable, - device=device, - dtype=self.dtype - ) - ] - if self.temporal_attention: - layers.append( - TemporalTransformer(ch, num_heads, dim_head, - depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, - use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only, - causal_attention=use_causal_attention, relative_position=use_relative_position, - temporal_length=temporal_length, - device=device, - dtype=self.dtype - ) - ) - layers.append( - ResBlock(ch, time_embed_dim, dropout, - dims=dims, use_checkpoint=use_checkpoint, - use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, - use_temporal_conv=temporal_conv, - device=device, - dtype=self.dtype - ) - ) - - ## Middle Block - self.middle_block = TimestepEmbedSequential(*layers) - - ## Output Block - self.output_blocks = nn.ModuleList([]) - for level, mult in list(enumerate(channel_mult))[::-1]: - for i in range(num_res_blocks + 1): - ich = input_block_chans.pop() - layers = [ - ResBlock(ch + ich, time_embed_dim, dropout, - out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint, - use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware, - use_temporal_conv=temporal_conv, - device=device, - dtype=self.dtype - ) - ] - ch = model_channels * mult - if ds in attention_resolutions: - if num_head_channels == -1: - dim_head = ch // num_heads - else: - num_heads = ch // num_head_channels - dim_head = num_head_channels - layers.append( - SpatialTransformer(ch, num_heads, dim_head, - depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, - use_checkpoint=use_checkpoint, disable_self_attn=False, video_length=temporal_length, - image_cross_attention=self.image_cross_attention,image_cross_attention_scale_learnable=self.image_cross_attention_scale_learnable, - device=device, - dtype=self.dtype - ) - ) - if self.temporal_attention: - layers.append( - TemporalTransformer(ch, num_heads, dim_head, - depth=transformer_depth, context_dim=context_dim, use_linear=use_linear, - use_checkpoint=use_checkpoint, only_self_att=temporal_self_att_only, - causal_attention=use_causal_attention, relative_position=use_relative_position, - temporal_length=temporal_length, - device=device, - dtype=self.dtype - ) - ) - if level and i == num_res_blocks: - out_ch = ch - layers.append( - ResBlock(ch, time_embed_dim, dropout, - out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint, - use_scale_shift_norm=use_scale_shift_norm, - up=True, - device=device, - dtype=self.dtype - ) - if resblock_updown - else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) - ) - ds //= 2 - self.output_blocks.append(TimestepEmbedSequential(*layers)) - - self.out = nn.Sequential( - normalization(ch, device=device, dtype=self.dtype), - nn.SiLU(), - zero_module( - operations.conv_nd( - dims, - model_channels, - out_channels, - 3, - padding=1, - device=device, - dtype=self.dtype - ) - ), - ) - - # TODO Add Transformer options to leverage the usage of patches. - def forward( - self, - x, - timesteps, - context=None, - context_in=None, - cc_concat=None, - num_video_frames=16, - features_adapter=None, - fs=None, - img_emb=None, - control=None, - transformer_options={}, - cond_idx=None, - **kwargs - ): - - if any([fs is None, img_emb is None, cc_concat is None]): - raise ValueError("One or more of the required inputs for UNet Forward is None.") - - cond_idx = transformer_options.get("cond_idx", None) - transformer_options['original_shape'] = list(x.shape) - transformer_options['transformer_index'] = 0 - transformer_patches = transformer_options.get("patches", {}) - - # In ComfyUI, the frames are always with the batch, so we deconstruct it here. - # This is mandatory as this is a video based model. - # We usually denote "f" as frames, but will use "t" (time) to be consistent with DynamiCrafter. - b,_,t,_,_ = x.shape - - context = context_in - cc_concat = cc_concat.to(x.device, x.dtype) - x = torch.cat([x, cc_concat], dim=1) - - fs = fs.to(x.device, x.dtype) - - timestep = timesteps - context = context_processor(context, num_video_frames, img_emb=img_emb) - - t_emb = timestep_embedding(timestep, self.model_channels, repeat_only=False, dtype=self.dtype) - emb = self.time_embed(t_emb) - emb = emb.repeat_interleave(repeats=t, dim=0) - - ## always in shape (b t) c h w, except for temporal layer - x = rearrange(x, 'b c t h w -> (b t) c h w') - - ## combine emb - if self.fs_condition: - if fs is None: - fs = torch.tensor( - [self.default_fs] * b, dtype=torch.long, device=x.device) - fs_emb = timestep_embedding(fs, self.model_channels, repeat_only=False, dtype=self.dtype).type(x.dtype) - - fs_embed = self.fps_embedding(fs_emb) - fs_embed = fs_embed.repeat_interleave(repeats=t, dim=0) - - emb = emb + fs_embed - - h = x.type(self.dtype) - adapter_idx = 0 - hs = [] - - for id, module in enumerate(self.input_blocks): - transformer_options["block"] = ("input", id) - #h = module(h, emb, context=context, batch_size=b) - h = forward_timestep_embed( - module, - h, - emb, - context=context, - batch_size=b, - transformer_options=transformer_options - ) - h = apply_control(h, control, 'input', cond_idx) - - if "input_block_patch" in transformer_patches: - patch = transformer_patches["input_block_patch"] - for p in patch: - h = p(h, transformer_options) - - if id ==0 and self.addition_attention: - h = forward_timestep_embed( - self.init_attn, - h, - emb, - context=context, - batch_size=b, - transformer_options=transformer_options - ) - ## plug-in adapter features - if ((id+1)%3 == 0) and features_adapter is not None: - h = h + features_adapter[adapter_idx] - adapter_idx += 1 - hs.append(h) - if "input_block_patch_after_skip" in transformer_patches: - patch = transformer_patches["input_block_patch_after_skip"] - for p in patch: - h = p(h, transformer_options) - if features_adapter is not None: - assert len(features_adapter)==adapter_idx, 'Wrong features_adapter' - transformer_options["block"] = ("middle", 0) - h = forward_timestep_embed( - self.middle_block, - h, - emb, - context=context, - batch_size=b, - transformer_options=transformer_options - ) - h = apply_control(h, control, 'middle', cond_idx) - for id, module in enumerate(self.output_blocks): - transformer_options["block"] = ("output", id) - hsp = hs.pop() - hsp = apply_control(hsp, control, 'output', cond_idx) - - if "output_block_patch" in transformer_patches: - patch = transformer_patches["output_block_patch"] - for p in patch: - h, hsp = p(h, hsp, transformer_options) - - h = torch.cat([h, hsp], dim=1) - del hsp - h = forward_timestep_embed( - module, - h, - emb, - context=context, - batch_size=b, - transformer_options=transformer_options - ) - h = h.type(x.dtype) - h = self.out(h) - - # We output with the tensor unfolded framewise, then reshape them to batched using ComfyUI nodes. - h = rearrange(h, '(b t) c h w -> b c t h w', t=num_video_frames) - - return h \ No newline at end of file diff --git a/modules/x_transformer.py b/modules/x_transformer.py deleted file mode 100644 index 5321012..0000000 --- a/modules/x_transformer.py +++ /dev/null @@ -1,639 +0,0 @@ -"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers""" -from functools import partial -from inspect import isfunction -from collections import namedtuple -from einops import rearrange, repeat -import torch -from torch import nn, einsum -import torch.nn.functional as F - -# constants -DEFAULT_DIM_HEAD = 64 - -Intermediates = namedtuple('Intermediates', [ - 'pre_softmax_attn', - 'post_softmax_attn' -]) - -LayerIntermediates = namedtuple('Intermediates', [ - 'hiddens', - 'attn_intermediates' -]) - - -class AbsolutePositionalEmbedding(nn.Module): - def __init__(self, dim, max_seq_len): - super().__init__() - self.emb = nn.Embedding(max_seq_len, dim) - self.init_() - - def init_(self): - nn.init.normal_(self.emb.weight, std=0.02) - - def forward(self, x): - n = torch.arange(x.shape[1], device=x.device) - return self.emb(n)[None, :, :] - - -class FixedPositionalEmbedding(nn.Module): - def __init__(self, dim): - super().__init__() - inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim)) - self.register_buffer('inv_freq', inv_freq) - - def forward(self, x, seq_dim=1, offset=0): - t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset - sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq) - emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1) - return emb[None, :, :] - - -# helpers - -def exists(val): - return val is not None - - -def default(val, d): - if exists(val): - return val - return d() if isfunction(d) else d - - -def always(val): - def inner(*args, **kwargs): - return val - return inner - - -def not_equals(val): - def inner(x): - return x != val - return inner - - -def equals(val): - def inner(x): - return x == val - return inner - - -def max_neg_value(tensor): - return -torch.finfo(tensor.dtype).max - - -# keyword argument helpers - -def pick_and_pop(keys, d): - values = list(map(lambda key: d.pop(key), keys)) - return dict(zip(keys, values)) - - -def group_dict_by_key(cond, d): - return_val = [dict(), dict()] - for key in d.keys(): - match = bool(cond(key)) - ind = int(not match) - return_val[ind][key] = d[key] - return (*return_val,) - - -def string_begins_with(prefix, str): - return str.startswith(prefix) - - -def group_by_key_prefix(prefix, d): - return group_dict_by_key(partial(string_begins_with, prefix), d) - - -def groupby_prefix_and_trim(prefix, d): - kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d) - kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items()))) - return kwargs_without_prefix, kwargs - - -# classes -class Scale(nn.Module): - def __init__(self, value, fn): - super().__init__() - self.value = value - self.fn = fn - - def forward(self, x, **kwargs): - x, *rest = self.fn(x, **kwargs) - return (x * self.value, *rest) - - -class Rezero(nn.Module): - def __init__(self, fn): - super().__init__() - self.fn = fn - self.g = nn.Parameter(torch.zeros(1)) - - def forward(self, x, **kwargs): - x, *rest = self.fn(x, **kwargs) - return (x * self.g, *rest) - - -class ScaleNorm(nn.Module): - def __init__(self, dim, eps=1e-5): - super().__init__() - self.scale = dim ** -0.5 - self.eps = eps - self.g = nn.Parameter(torch.ones(1)) - - def forward(self, x): - norm = torch.norm(x, dim=-1, keepdim=True) * self.scale - return x / norm.clamp(min=self.eps) * self.g - - -class RMSNorm(nn.Module): - def __init__(self, dim, eps=1e-8): - super().__init__() - self.scale = dim ** -0.5 - self.eps = eps - self.g = nn.Parameter(torch.ones(dim)) - - def forward(self, x): - norm = torch.norm(x, dim=-1, keepdim=True) * self.scale - return x / norm.clamp(min=self.eps) * self.g - - -class Residual(nn.Module): - def forward(self, x, residual): - return x + residual - - -class GRUGating(nn.Module): - def __init__(self, dim): - super().__init__() - self.gru = nn.GRUCell(dim, dim) - - def forward(self, x, residual): - gated_output = self.gru( - rearrange(x, 'b n d -> (b n) d'), - rearrange(residual, 'b n d -> (b n) d') - ) - - return gated_output.reshape_as(x) - - -# feedforward - -class GEGLU(nn.Module): - def __init__(self, dim_in, dim_out): - super().__init__() - self.proj = nn.Linear(dim_in, dim_out * 2) - - def forward(self, x): - x, gate = self.proj(x).chunk(2, dim=-1) - return x * F.gelu(gate) - - -class FeedForward(nn.Module): - def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.): - super().__init__() - inner_dim = int(dim * mult) - dim_out = default(dim_out, dim) - project_in = nn.Sequential( - nn.Linear(dim, inner_dim), - nn.GELU() - ) if not glu else GEGLU(dim, inner_dim) - - self.net = nn.Sequential( - project_in, - nn.Dropout(dropout), - nn.Linear(inner_dim, dim_out) - ) - - def forward(self, x): - return self.net(x) - - -# attention. -class Attention(nn.Module): - def __init__( - self, - dim, - dim_head=DEFAULT_DIM_HEAD, - heads=8, - causal=False, - mask=None, - talking_heads=False, - sparse_topk=None, - use_entmax15=False, - num_mem_kv=0, - dropout=0., - on_attn=False - ): - super().__init__() - if use_entmax15: - raise NotImplementedError("Check out entmax activation instead of softmax activation!") - self.scale = dim_head ** -0.5 - self.heads = heads - self.causal = causal - self.mask = mask - - inner_dim = dim_head * heads - - self.to_q = nn.Linear(dim, inner_dim, bias=False) - self.to_k = nn.Linear(dim, inner_dim, bias=False) - self.to_v = nn.Linear(dim, inner_dim, bias=False) - self.dropout = nn.Dropout(dropout) - - # talking heads - self.talking_heads = talking_heads - if talking_heads: - self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads)) - self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads)) - - # explicit topk sparse attention - self.sparse_topk = sparse_topk - - # entmax - #self.attn_fn = entmax15 if use_entmax15 else F.softmax - self.attn_fn = F.softmax - - # add memory key / values - self.num_mem_kv = num_mem_kv - if num_mem_kv > 0: - self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head)) - self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head)) - - # attention on attention - self.attn_on_attn = on_attn - self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim) - - def forward( - self, - x, - context=None, - mask=None, - context_mask=None, - rel_pos=None, - sinusoidal_emb=None, - prev_attn=None, - mem=None - ): - b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device - kv_input = default(context, x) - - q_input = x - k_input = kv_input - v_input = kv_input - - if exists(mem): - k_input = torch.cat((mem, k_input), dim=-2) - v_input = torch.cat((mem, v_input), dim=-2) - - if exists(sinusoidal_emb): - # in shortformer, the query would start at a position offset depending on the past cached memory - offset = k_input.shape[-2] - q_input.shape[-2] - q_input = q_input + sinusoidal_emb(q_input, offset=offset) - k_input = k_input + sinusoidal_emb(k_input) - - q = self.to_q(q_input) - k = self.to_k(k_input) - v = self.to_v(v_input) - - q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v)) - - input_mask = None - if any(map(exists, (mask, context_mask))): - q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool()) - k_mask = q_mask if not exists(context) else context_mask - k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool()) - q_mask = rearrange(q_mask, 'b i -> b () i ()') - k_mask = rearrange(k_mask, 'b j -> b () () j') - input_mask = q_mask * k_mask - - if self.num_mem_kv > 0: - mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v)) - k = torch.cat((mem_k, k), dim=-2) - v = torch.cat((mem_v, v), dim=-2) - if exists(input_mask): - input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True) - - dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale - mask_value = max_neg_value(dots) - - if exists(prev_attn): - dots = dots + prev_attn - - pre_softmax_attn = dots - - if talking_heads: - dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous() - - if exists(rel_pos): - dots = rel_pos(dots) - - if exists(input_mask): - dots.masked_fill_(~input_mask, mask_value) - del input_mask - - if self.causal: - i, j = dots.shape[-2:] - r = torch.arange(i, device=device) - mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j') - mask = F.pad(mask, (j - i, 0), value=False) - dots.masked_fill_(mask, mask_value) - del mask - - if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]: - top, _ = dots.topk(self.sparse_topk, dim=-1) - vk = top[..., -1].unsqueeze(-1).expand_as(dots) - mask = dots < vk - dots.masked_fill_(mask, mask_value) - del mask - - attn = self.attn_fn(dots, dim=-1) - post_softmax_attn = attn - - attn = self.dropout(attn) - - if talking_heads: - attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous() - - out = einsum('b h i j, b h j d -> b h i d', attn, v) - out = rearrange(out, 'b h n d -> b n (h d)') - - intermediates = Intermediates( - pre_softmax_attn=pre_softmax_attn, - post_softmax_attn=post_softmax_attn - ) - - return self.to_out(out), intermediates - - -class AttentionLayers(nn.Module): - def __init__( - self, - dim, - depth, - heads=8, - causal=False, - cross_attend=False, - only_cross=False, - use_scalenorm=False, - use_rmsnorm=False, - use_rezero=False, - rel_pos_num_buckets=32, - rel_pos_max_distance=128, - position_infused_attn=False, - custom_layers=None, - sandwich_coef=None, - par_ratio=None, - residual_attn=False, - cross_residual_attn=False, - macaron=False, - pre_norm=True, - gate_residual=False, - **kwargs - ): - super().__init__() - ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs) - attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs) - - dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD) - - self.dim = dim - self.depth = depth - self.layers = nn.ModuleList([]) - - self.has_pos_emb = position_infused_attn - self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None - self.rotary_pos_emb = always(None) - - assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance' - self.rel_pos = None - - self.pre_norm = pre_norm - - self.residual_attn = residual_attn - self.cross_residual_attn = cross_residual_attn - - norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm - norm_class = RMSNorm if use_rmsnorm else norm_class - norm_fn = partial(norm_class, dim) - - norm_fn = nn.Identity if use_rezero else norm_fn - branch_fn = Rezero if use_rezero else None - - if cross_attend and not only_cross: - default_block = ('a', 'c', 'f') - elif cross_attend and only_cross: - default_block = ('c', 'f') - else: - default_block = ('a', 'f') - - if macaron: - default_block = ('f',) + default_block - - if exists(custom_layers): - layer_types = custom_layers - elif exists(par_ratio): - par_depth = depth * len(default_block) - assert 1 < par_ratio <= par_depth, 'par ratio out of range' - default_block = tuple(filter(not_equals('f'), default_block)) - par_attn = par_depth // par_ratio - depth_cut = par_depth * 2 // 3 # 2 / 3 attention layer cutoff suggested by PAR paper - par_width = (depth_cut + depth_cut // par_attn) // par_attn - assert len(default_block) <= par_width, 'default block is too large for par_ratio' - par_block = default_block + ('f',) * (par_width - len(default_block)) - par_head = par_block * par_attn - layer_types = par_head + ('f',) * (par_depth - len(par_head)) - elif exists(sandwich_coef): - assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth' - layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef - else: - layer_types = default_block * depth - - self.layer_types = layer_types - self.num_attn_layers = len(list(filter(equals('a'), layer_types))) - - for layer_type in self.layer_types: - if layer_type == 'a': - layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs) - elif layer_type == 'c': - layer = Attention(dim, heads=heads, **attn_kwargs) - elif layer_type == 'f': - layer = FeedForward(dim, **ff_kwargs) - layer = layer if not macaron else Scale(0.5, layer) - else: - raise Exception(f'invalid layer type {layer_type}') - - if isinstance(layer, Attention) and exists(branch_fn): - layer = branch_fn(layer) - - if gate_residual: - residual_fn = GRUGating(dim) - else: - residual_fn = Residual() - - self.layers.append(nn.ModuleList([ - norm_fn(), - layer, - residual_fn - ])) - - def forward( - self, - x, - context=None, - mask=None, - context_mask=None, - mems=None, - return_hiddens=False - ): - hiddens = [] - intermediates = [] - prev_attn = None - prev_cross_attn = None - - mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers - - for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)): - is_last = ind == (len(self.layers) - 1) - - if layer_type == 'a': - hiddens.append(x) - layer_mem = mems.pop(0) - - residual = x - - if self.pre_norm: - x = norm(x) - - if layer_type == 'a': - out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos, - prev_attn=prev_attn, mem=layer_mem) - elif layer_type == 'c': - out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn) - elif layer_type == 'f': - out = block(x) - - x = residual_fn(out, residual) - - if layer_type in ('a', 'c'): - intermediates.append(inter) - - if layer_type == 'a' and self.residual_attn: - prev_attn = inter.pre_softmax_attn - elif layer_type == 'c' and self.cross_residual_attn: - prev_cross_attn = inter.pre_softmax_attn - - if not self.pre_norm and not is_last: - x = norm(x) - - if return_hiddens: - intermediates = LayerIntermediates( - hiddens=hiddens, - attn_intermediates=intermediates - ) - - return x, intermediates - - return x - - -class Encoder(AttentionLayers): - def __init__(self, **kwargs): - assert 'causal' not in kwargs, 'cannot set causality on encoder' - super().__init__(causal=False, **kwargs) - - - -class TransformerWrapper(nn.Module): - def __init__( - self, - *, - num_tokens, - max_seq_len, - attn_layers, - emb_dim=None, - max_mem_len=0., - emb_dropout=0., - num_memory_tokens=None, - tie_embedding=False, - use_pos_emb=True - ): - super().__init__() - assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder' - - dim = attn_layers.dim - emb_dim = default(emb_dim, dim) - - self.max_seq_len = max_seq_len - self.max_mem_len = max_mem_len - self.num_tokens = num_tokens - - self.token_emb = nn.Embedding(num_tokens, emb_dim) - self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if ( - use_pos_emb and not attn_layers.has_pos_emb) else always(0) - self.emb_dropout = nn.Dropout(emb_dropout) - - self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity() - self.attn_layers = attn_layers - self.norm = nn.LayerNorm(dim) - - self.init_() - - self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t() - - # memory tokens (like [cls]) from Memory Transformers paper - num_memory_tokens = default(num_memory_tokens, 0) - self.num_memory_tokens = num_memory_tokens - if num_memory_tokens > 0: - self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim)) - - # let funnel encoder know number of memory tokens, if specified - if hasattr(attn_layers, 'num_memory_tokens'): - attn_layers.num_memory_tokens = num_memory_tokens - - def init_(self): - nn.init.normal_(self.token_emb.weight, std=0.02) - - def forward( - self, - x, - return_embeddings=False, - mask=None, - return_mems=False, - return_attn=False, - mems=None, - **kwargs - ): - b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens - x = self.token_emb(x) - x += self.pos_emb(x) - x = self.emb_dropout(x) - - x = self.project_emb(x) - - if num_mem > 0: - mem = repeat(self.memory_tokens, 'n d -> b n d', b=b) - x = torch.cat((mem, x), dim=1) - - # auto-handle masking after appending memory tokens - if exists(mask): - mask = F.pad(mask, (num_mem, 0), value=True) - - x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs) - x = self.norm(x) - - mem, x = x[:, :num_mem], x[:, num_mem:] - - out = self.to_logits(x) if not return_embeddings else x - - if return_mems: - hiddens = intermediates.hiddens - new_mems = list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens))) if exists(mems) else hiddens - new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems)) - return out, new_mems - - if return_attn: - attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates)) - return out, attn_maps - - return out \ No newline at end of file From ccd9611e2b0ae9ac3cf87ac1dc4d829a87f4bfbb Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 15:30:37 -0700 Subject: [PATCH 4/6] Update attention.py --- lvdm/modules/attention.py | 140 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 136 insertions(+), 4 deletions(-) diff --git a/lvdm/modules/attention.py b/lvdm/modules/attention.py index 7537aa5..f71c89d 100644 --- a/lvdm/modules/attention.py +++ b/lvdm/modules/attention.py @@ -294,9 +294,11 @@ class BasicTransformerBlock(nn.Module): disable_self_attn=False, attention_cls=None, video_length=None, + inner_dim=None, image_cross_attention=False, image_cross_attention_scale=1.0, image_cross_attention_scale_learnable=False, + switch_temporal_ca_to_sa=False, text_context_len=77, device=None, dtype=None, @@ -304,6 +306,23 @@ class BasicTransformerBlock(nn.Module): ): super().__init__() attn_cls = CrossAttention if attention_cls is None else attention_cls + + self.ff_in = ff_in or inner_dim is not None + if self.ff_in: + self.norm_in = operations.LayerNorm(dim, dtype=dtype, device=device) + self.ff_in = FeedForward( + dim, + dim_out=inner_dim, + dropout=dropout, + glu=gated_ff, + dtype=dtype, + device=device, + operations=operations + ) + if inner_dim is None: + inner_dim = dim + + self.is_res = inner_dim == dim self.disable_self_attn = disable_self_attn self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, context_dim=None, device=device, dtype=dtype if self.disable_self_attn else None) @@ -327,7 +346,11 @@ class BasicTransformerBlock(nn.Module): self.norm1 = operations.LayerNorm(dim, device=device, dtype=dtype) self.norm2 = operations.LayerNorm(dim, device=device, dtype=dtype) self.norm3 = operations.LayerNorm(dim, device=device, dtype=dtype) + + self.n_heads = n_heads + self.d_head = d_head self.checkpoint = checkpoint + self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa def forward(self, x, context=None, mask=None, **kwargs): ## implementation tricks: because checkpointing doesn't support non-tensor (e.g. None or scalar) arguments @@ -341,9 +364,117 @@ class BasicTransformerBlock(nn.Module): def _forward(self, x, context=None, mask=None): - x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None, mask=mask) + x - x = self.attn2(self.norm2(x), context=context, mask=mask) + x - x = self.ff(self.norm3(x)) + x + extra_options = {} + block = transformer_options.get("block", None) + block_index = transformer_options.get("block_index", 0) + transformer_patches = {} + transformer_patches_replace = {} + + for k in transformer_options: + if k == "patches": + transformer_patches = transformer_options[k] + elif k == "patches_replace": + transformer_patches_replace = transformer_options[k] + else: + extra_options[k] = transformer_options[k] + + extra_options["n_heads"] = self.n_heads + extra_options["dim_head"] = self.d_head + + if self.ff_in: + x_skip = x + x = self.ff_in(self.norm_in(x)) + if self.is_res: + x += x_skip + + n = self.norm1(x) + if self.disable_self_attn: + context_attn1 = context + else: + context_attn1 = None + value_attn1 = None + + if "attn1_patch" in transformer_patches: + patch = transformer_patches["attn1_patch"] + if context_attn1 is None: + context_attn1 = n + value_attn1 = context_attn1 + for p in patch: + n, context_attn1, value_attn1 = p(n, context_attn1, value_attn1, extra_options) + + if block is not None: + transformer_block = (block[0], block[1], block_index) + else: + transformer_block = None + attn1_replace_patch = transformer_patches_replace.get("attn1", {}) + block_attn1 = transformer_block + if block_attn1 not in attn1_replace_patch: + block_attn1 = block + + if block_attn1 in attn1_replace_patch: + if context_attn1 is None: + context_attn1 = n + value_attn1 = n + n = self.attn1.to_q(n) + context_attn1 = self.attn1.to_k(context_attn1) + value_attn1 = self.attn1.to_v(value_attn1) + n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options) + n = self.attn1.to_out(n) + else: + n = self.attn1(n, context=context_attn1, value=value_attn1) + + if "attn1_output_patch" in transformer_patches: + patch = transformer_patches["attn1_output_patch"] + for p in patch: + n = p(n, extra_options) + + x += n + if "middle_patch" in transformer_patches: + patch = transformer_patches["middle_patch"] + for p in patch: + x = p(x, extra_options) + + if self.attn2 is not None: + n = self.norm2(x) + if self.switch_temporal_ca_to_sa: + context_attn2 = n + else: + context_attn2 = context + value_attn2 = None + if "attn2_patch" in transformer_patches: + patch = transformer_patches["attn2_patch"] + value_attn2 = context_attn2 + for p in patch: + n, context_attn2, value_attn2 = p(n, context_attn2, value_attn2, extra_options) + + attn2_replace_patch = transformer_patches_replace.get("attn2", {}) + block_attn2 = transformer_block + if block_attn2 not in attn2_replace_patch: + block_attn2 = block + + if block_attn2 in attn2_replace_patch: + if value_attn2 is None: + value_attn2 = context_attn2 + n = self.attn2.to_q(n) + context_attn2 = self.attn2.to_k(context_attn2) + value_attn2 = self.attn2.to_v(value_attn2) + n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options) + n = self.attn2.to_out(n) + else: + n = self.attn2(n, context=context_attn2, value=value_attn2) + + if "attn2_output_patch" in transformer_patches: + patch = transformer_patches["attn2_output_patch"] + for p in patch: + n = p(n, extra_options) + + x += n + if self.is_res: + x_skip = x + x = self.ff(self.norm3(x)) + if self.is_res: + x += x_skip + return x @@ -408,7 +539,7 @@ class SpatialTransformer(nn.Module): self.proj_out = zero_module(operations.Linear(inner_dim, in_channels, device=device, dtype=dtype)) self.use_linear = use_linear - def forward(self, x, context=None, **kwargs): + def forward(self, x, context=None, transformer_options={}, **kwargs): b, c, h, w = x.shape x_in = x x = self.norm(x) @@ -418,6 +549,7 @@ class SpatialTransformer(nn.Module): if self.use_linear: x = self.proj_in(x) for i, block in enumerate(self.transformer_blocks): + transformer_options['block_index'] = i x = block(x, context=context, **kwargs) if self.use_linear: x = self.proj_out(x) From f7ad172ae3fcc5dc06342d9b50d195c17cc90b08 Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 15:31:18 -0700 Subject: [PATCH 5/6] Complete ControlNet Support --- nodes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nodes.py b/nodes.py index c0b239d..6055180 100644 --- a/nodes.py +++ b/nodes.py @@ -124,6 +124,7 @@ class DynamiCrafterProcessor: outs = [] for i in range(batch_size): + model_in_kwargs['transformer_options']['cond_idx'] = i x_out = apply_model( x_in[[i]], t=torch.cat([t[:1]]), From d61c5910ae05c47919295317398d97ae6551974f Mon Sep 17 00:00:00 2001 From: ExponentialML <59846140+ExponentialML@users.noreply.github.com> Date: Fri, 22 Mar 2024 15:37:31 -0700 Subject: [PATCH 6/6] Update README.md --- README.md | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 2c2c136..df4f7ed 100644 --- a/README.md +++ b/README.md @@ -56,15 +56,32 @@ If you are downloading the CLIP and VAE models separately, place them under thei - **empty_latent**: An empty latent with the same size and frames as the processed ones. - **latent_img**: If you're doing Img2Img based workflows, this is the necessary one to use. - + +# ControlNet Support + +You can now use DynamiCrafter by applying ControlNet to the Spatial (image) portion to guide video generations in various ways. +The ControlNets are based on 2.1, so you must download them at the link below (Thanks @thibaud !) . + +**ControlNet 2.1**: https://huggingface.co/thibaud/controlnet-sd21 + +After you download them, you can use them as you would with any other workflow. + +# Tips + > [!TIP] > You don't have to use the latent outputs. As long as you use the same frame length (as your batch size) and same height and with as your image inputs, you can use your own latents. > This means that you can experiment with inpainting and so on. +> [!TIP] +> You can choose which frame you use as init by using VAE Encode Inpaint or Set Latent Noise Mask. You set the beginning batch mask to full black, while the rest are at full white. This also means you can do interpolation with regular models. +> As these workflows are more advanced, examples will arrive at a future date. + # TODO - [x] Add various workflows. -- [ ] Add support for Spatial Transformer options. -- [ ] Ensure attention optimizations are working properly. +- [ ] Add advanced workflows. +- [x] Add support for Spatial Transformer options. +- [x] Add ControlNet support. +- [x] Ensure attention optimizations are working properly. - [ ] Add autoregressive nodes (this may be a separate repository) - [x] Add examples. (For more, [check here](https://github.com/Doubiiu/DynamiCrafter?tab=readme-ov-file#11-showcases-576x1024)).