v5b and matryoshka-liked style transfer

hellomeme/model_config.json

@@ -22,5 +22,6 @@
       "[preset]Drawing": ["songkey/loras_sd_1_5", "Drawing.safetensors"]
     }
   },
-  "prompt": "(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture"
+  "prompt": "(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture",
+  "prompt_new": "best quality, highly detailed, masterpiece, person, portrait, headshot, looking at the viewer, sharp focus."
 }

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__(

hellomeme/pipelines/pipline_hm5_image.py

@@ -20,10 +20,8 @@ from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
 from diffusers import DPMSolverMultistepScheduler
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
-from ..models import (HM3Denoising3D,
-                     HMPipeline, HM5ReferenceAdapter,
-                     HM5ControlNetBase,
-                     HM5SD15ControlProj)
+from ..models import (HM3Denoising3D, HMPipeline, HM5ReferenceAdapter,
+                      HM5bReferenceAdapter, HM5ControlNetBase, HM5SD15ControlProj)
 
 class HM5ImagePipeline(HMPipeline):
     def caryomitosis(self, **kwargs):
@@ -46,20 +44,28 @@ class HM5ImagePipeline(HMPipeline):
         self.safety_checker.cpu()
 
     def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False):
-
         self.version = version
         if modelscope:
             from modelscope import snapshot_download
-            hm_reference_dir = snapshot_download('songkey/hm5_reference')
+            if version == 'v5b':
+                hm_reference_dir = snapshot_download('songkey/hm5b_reference')
+            else:
+                hm_reference_dir = snapshot_download('songkey/hm5_reference')
             hm_control_dir = snapshot_download('songkey/hm5_control_base')
             hm_control_proj_dir = snapshot_download('songkey/hm5_control_proj')
         else:
-            hm_reference_dir = 'songkey/hm5_reference'
+            if version == 'v5b':
+                hm_reference_dir = 'songkey/hm5b_reference'
+            else:
+                hm_reference_dir = 'songkey/hm5_reference'
             hm_control_dir = 'songkey/hm5_control_base'
             hm_control_proj_dir = 'songkey/hm5_control_proj'
 
         if isinstance(self.unet, HM3Denoising3D):
-            hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir)
+            if version == 'v5b':
+                hm_adapter = HM5bReferenceAdapter.from_pretrained(hm_reference_dir)
+            else:
+                hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir)
 
             self.unet.insert_reference_adapter(hm_adapter)
             self.unet.to(device='cpu', dtype=dtype).eval()

hellomeme/pipelines/pipline_hm_image.py

@@ -45,6 +45,7 @@ class HMImagePipeline(HMPipeline):
         self.safety_checker.cpu()
 
     def insert_hm_modules(self, version, dtype, modelscope=False):
+        self.version = version
         if modelscope:
             from modelscope import snapshot_download
             hm_reference_dir = snapshot_download('songkey/hm_reference')