init

modules/attn_wrapper.py

@@ -0,0 +1,63 @@
+import torch
+
+from .ref_controller import RefMode
+
+
+def get_attn_wrapper(cls):
+    class RefAttention(cls):
+        ref_mode = RefMode.OFF
+        ref_bank = None
+        is_cfg = True
+        style_fidelity = 0.5
+        ref_attn_weight = 0.5
+        attention_bound = 1.0
+        hidden_states = None
+
+        def forward(self, *args, **kwargs):
+            if self.ref_mode == RefMode.OFF:
+                return super().forward(*args, **kwargs)
+            
+            return self._ref_forward(*args, **kwargs)
+        
+        def _ref_forward(self, norm_hidden_states, context, **kwargs):
+            if self.ref_mode == RefMode.WRITE:
+                self.ref_bank = norm_hidden_states.detach().clone()
+                return super().forward(norm_hidden_states, context, **kwargs)
+            if self.ref_mode == RefMode.READ:
+                if self.attention_bound > self.ref_attn_weight:
+                    if len(self.ref_bank) == 2:
+                        batch_size = len(norm_hidden_states) // len(self.ref_bank)
+                        uc_bank = self.ref_bank[0].unsqueeze(0)
+                        uc_bank = torch.cat([uc_bank] * batch_size)
+                        c_bank = self.ref_bank[1].unsqueeze(0)
+                        c_bank = torch.cat([c_bank] * batch_size)
+                        bank = torch.cat([uc_bank, c_bank])
+                    else:
+                        bank = self.ref_bank
+                    bank_hidden_states = torch.cat([bank, norm_hidden_states], dim=1)
+                    attn_output_uc = (
+                        super().forward(
+                            norm_hidden_states,
+                            bank_hidden_states,
+                            **kwargs
+                        )
+                    )
+                    attn_output_c = attn_output_uc.clone()
+                    if self.is_cfg and self.style_fidelity > 0:
+                        uc_mask =  torch.Tensor(
+                                [1] * (norm_hidden_states.shape[0] // 2)
+                                + [0] * (norm_hidden_states.shape[0] // 2)
+                            ).to(norm_hidden_states.device).bool()
+                        attn_output_c[uc_mask] = (
+                            super().forward(
+                                norm_hidden_states[uc_mask],
+                                norm_hidden_states[uc_mask],
+                                **kwargs
+                            )
+                        )
+                    attn_output = self.style_fidelity * attn_output_c + (1.0 - self.style_fidelity) * attn_output_uc
+                else:
+                    attn_output = super().forward(norm_hidden_states, context, **kwargs)
+                return attn_output
+            
+    return RefAttention

modules/ref_controller.py

@@ -0,0 +1,54 @@
+import enum
+
+
+class RefMode(enum.Enum):
+    OFF = 1
+    WRITE = 2
+    READ = 3
+
+
+class RefController:
+    modules = []
+    temp_modules = []
+    default_video_length = 16
+    default_full_length = 16
+
+    def add_module(self, module):
+        self.modules.append(module)
+
+    def set_mode(self, mode: RefMode):
+        for module in self.modules:
+            module.ref_mode = mode
+
+    def clear_modules(self):
+        for module in self.modules:
+            module.ref_bank = None
+        self.set_mode(RefMode.OFF)
+
+    def set_cfg(self, cfg: bool):
+        for module in self.modules:
+            module.cfg = cfg
+
+    def set_attention_bound(self, attention_bound: float):
+        for module in self.modules:
+            module.attention_bound = attention_bound
+
+    def set_style_fidelity(self, style_fidelity: float): 
+        for module in self.modules:
+            module.style_fidelity = style_fidelity
+
+    def add_temporal_transformer(self, module):
+        self.temp_modules.append(module)
+
+    def get_default_temporal_values(self):
+        if len(self.temp_modules) > 0:
+            self.default_full_length = self.temp_modules[0].full_length
+            self.default_video_length = self.temp_modules[0].video_length
+
+    def set_temporal_values_to(self, full_length, video_length):
+        for module in self.temp_modules:
+            module.full_length = full_length
+            module.video_length = video_length
+
+    def set_temporal_values_to_default(self):
+        self.set_temporal_values_to(self.default_full_length, self.default_video_length)