v5b and matryoshka-liked style transfer

hellomeme/models/__init__.py

@@ -12,7 +12,7 @@ from .hm_denoising_motion import HMDenoisingMotion
 from .hm_control import (HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2,
                          HMV3ControlNet, HMControlNetBase, HM5ControlNetBase,
                          HM4SD15ControlProj, HM5SD15ControlProj)
-from .hm_adapters import (HMReferenceAdapter, HM3ReferenceAdapter, HM5ReferenceAdapter,
+from .hm_adapters import (HMReferenceAdapter, HM3ReferenceAdapter, HM5ReferenceAdapter, HM5bReferenceAdapter,
                           HM3MotionAdapter, HM5MotionAdapter, HMPipeline)
 from .hm_denoising_3d import HMDenoising3D
 from .hm3_denoising_3d import HM3Denoising3D

hellomeme/models/hm_adapters.py

@@ -41,6 +41,7 @@ from diffusers.models.lora import adjust_lora_scale_text_encoder
 
 from .hm_blocks import (SKReferenceAttentionV3,
                         SKReferenceAttentionV5,
+                        STKReferenceModule,
                         SKReferenceAttention,
                         SKMotionModule,
                         SKMotionModuleV5)
@@ -136,6 +137,24 @@ class HM3ReferenceAdapter(ModelMixin, ConfigMixin):
             )
 
 
+class HM5bReferenceAdapter(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(self, block_down_channels: Tuple[int] = (320, 640, 1280, 1280),
+                     block_up_channels: Tuple[int] = (1280, 1280, 1280, 640),
+                     num_attention_heads: int = 8,
+                     use_3d: bool = False):
+        super().__init__()
+
+        self.reference_modules_up = nn.ModuleList([])
+        for i, in_channels in enumerate(block_up_channels):
+            self.reference_modules_up.append(
+                STKReferenceModule(
+                    in_channels=in_channels,
+                    num_attention_heads=num_attention_heads,
+                )
+            )
+
+
 class HM5ReferenceAdapter(ModelMixin, ConfigMixin):
     @register_to_config
     def __init__(self, block_down_channels: Tuple[int] = (320, 640, 1280, 1280),

hellomeme/models/hm_blocks.py

@@ -191,6 +191,38 @@ class SKReferenceAttentionV5(nn.Module):
 
         return hidden_states + self.proj(res2.contiguous())
 
+class STKReferenceModule(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_attention_heads: int = 1,
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.attn_ref = STKAttentionV5(in_channels=in_channels,
+                                     mid_channels=in_channels,
+                                     num_attention_heads=num_attention_heads,
+                                     time_embed_dim=None,
+                                     norm_elementwise_affine=norm_elementwise_affine,
+                                     norm_eps=norm_eps,
+                                     temporal_attn=False,)
+        self.attn = STKAttentionV5(in_channels=in_channels,
+                                     mid_channels=in_channels,
+                                     num_attention_heads=num_attention_heads,
+                                     time_embed_dim=None,
+                                     norm_elementwise_affine=norm_elementwise_affine,
+                                     norm_eps=norm_eps,
+                                     temporal_attn=False,)
+    def forward(self, hidden_states, ref_stats, num_frames):
+        if ref_stats.shape[0] != hidden_states.shape[0]:
+            ref_stats = ref_stats.repeat_interleave(num_frames, dim=0)
+        hidden_states = rearrange(hidden_states, "(b f) c h w -> b f h w c", f=num_frames)
+        ref_stats = rearrange(ref_stats, "(b f) c h w -> b f h w c", f=num_frames)
+        ref_stats = self.attn_ref(ref_stats, residual=False)
+        hidden_states = self.attn(hidden_states, ref_stats, residual=True)
+        hidden_states = rearrange(hidden_states, "b f h w c -> (b f) c h w")
+        return hidden_states
 
 class SmallUnet(nn.Module):
     def __init__(self, in_channels: int = 3,
@@ -380,7 +412,7 @@ class STKAttentionV5(nn.Module):
         self.norm = nn.LayerNorm(mid_channels, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
         self.proj = zero_module(nn.Conv2d(mid_channels, in_channels, kernel_size=3, padding=1))
 
-    def forward(self, hidden_states, ref_stats=None, emb=None):
+    def forward(self, hidden_states, ref_stats=None, emb=None, residual=True):
         b, f, h, w, _ = hidden_states.shape
 
         hidden_states_in = rearrange(hidden_states, "b f h w c -> (b f) c h w")
@@ -433,7 +465,9 @@ class STKAttentionV5(nn.Module):
         else:
             rest = rearrange(resy, "b f h w c -> (b f) c h w")
 
-        return hidden_states + rearrange(self.proj(rest), "(b f) c h w -> b f h w c", f=f)
+        rest = rearrange(self.proj(rest), "(b f) c h w -> b f h w c", f=f)
+
+        return hidden_states + rest if residual else rest
 
 class STKCrossAttentionV5(nn.Module):
     def __init__(