diff --git a/lgm/mvdream/mv_unet.py b/lgm/mvdream/mv_unet.py index 7d9ad4d..2071d2a 100644 --- a/lgm/mvdream/mv_unet.py +++ b/lgm/mvdream/mv_unet.py @@ -11,19 +11,39 @@ from einops import rearrange, repeat from diffusers.configuration_utils import ConfigMixin from diffusers.models.modeling_utils import ModelMixin -# require xformers! -import xformers -import xformers.ops - from kiui.cam import orbit_camera +def memory_efficient_attention(q, k, v, attn_bias=None): + """ + Implements a memory-efficient attention mechanism. + + Parameters: + - q, k, v: Query, Key, Value tensors of shape (batch_size, seq_len_q, dim) + - attn_bias: Optional bias tensor of shape (seq_len_k, seq_len_v) + + Returns: + - output tensor of shape (batch_size, seq_len_q, dim) + """ + # Calculate attention scores + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(k.size(-1)) + if attn_bias is not None: + scores += attn_bias + + # Apply softmax to get attention weights + attention_weights = torch.nn.functional.softmax(scores, dim=-1) + + # Compute the weighted sum of values + output = torch.matmul(attention_weights, v) + + return output + def get_camera( num_frames, elevation=0, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False, ): angle_gap = azimuth_span / num_frames cameras = [] for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap): - + pose = orbit_camera(elevation, azimuth, radius=1) # [4, 4] # opengl to blender @@ -140,17 +160,17 @@ class FeedForward(nn.Module): class MemoryEfficientCrossAttention(nn.Module): # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 def __init__( - self, - query_dim, - context_dim=None, - heads=8, - dim_head=64, + self, + query_dim, + context_dim=None, + heads=8, + dim_head=64, dropout=0.0, ip_dim=0, ip_weight=1, ): super().__init__() - + inner_dim = dim_head * heads context_dim = default(context_dim, query_dim) @@ -199,7 +219,7 @@ class MemoryEfficientCrossAttention(nn.Module): ) # actually compute the attention, what we cannot get enough of - out = xformers.ops.memory_efficient_attention( + out = memory_efficient_attention( q, k, v, attn_bias=None, op=self.attention_op ) @@ -213,7 +233,7 @@ class MemoryEfficientCrossAttention(nn.Module): (k_ip, v_ip), ) # actually compute the attention, what we cannot get enough of - out_ip = xformers.ops.memory_efficient_attention( + out_ip = memory_efficient_attention( q, k_ip, v_ip, attn_bias=None, op=self.attention_op ) out = out + self.ip_weight * out_ip @@ -228,7 +248,7 @@ class MemoryEfficientCrossAttention(nn.Module): class BasicTransformerBlock3D(nn.Module): - + def __init__( self, dim, @@ -259,7 +279,7 @@ class BasicTransformerBlock3D(nn.Module): # ip only applies to cross-attention ip_dim=ip_dim, ip_weight=ip_weight, - ) + ) self.norm1 = nn.LayerNorm(dim) self.norm2 = nn.LayerNorm(dim) self.norm3 = nn.LayerNorm(dim) @@ -311,9 +331,9 @@ class SpatialTransformer3D(nn.Module): for d in range(depth) ] ) - + self.proj_out = zero_module(nn.Linear(in_channels, inner_dim)) - + def forward(self, x, context=None, num_frames=1): # note: if no context is given, cross-attention defaults to self-attention @@ -328,7 +348,7 @@ class SpatialTransformer3D(nn.Module): x = block(x, context=context[i], num_frames=num_frames) x = self.proj_out(x) x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous() - + return x + x_in @@ -672,7 +692,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin): ): super().__init__() assert context_dim is not None - + if num_heads_upsample == -1: num_heads_upsample = num_heads @@ -699,7 +719,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin): "as a list/tuple (per-level) with the same length as channel_mult" ) self.num_res_blocks = num_res_blocks - + if num_attention_blocks is not None: assert len(num_attention_blocks) == len(self.num_res_blocks) assert all( @@ -848,7 +868,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin): else: num_heads = ch // num_head_channels dim_head = num_head_channels - + self.middle_block = CondSequential( ResBlock( ch, @@ -865,7 +885,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin): depth=transformer_depth, ip_dim=self.ip_dim, ip_weight=self.ip_weight, - ), + ), ResBlock( ch, time_embed_dim, @@ -983,7 +1003,7 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin): # Add camera embeddings if camera is not None: emb = emb + self.camera_embed(camera) - + # imagedream variant if self.ip_dim > 0: x[(num_frames - 1) :: num_frames, :, :, :] = ip_img # place at [4, 9] @@ -1002,4 +1022,4 @@ class MultiViewUNetModel(ModelMixin, ConfigMixin): if self.predict_codebook_ids: return self.id_predictor(h) else: - return self.out(h) \ No newline at end of file + return self.out(h)