reference unet

modules/attn_wrapper.py

@@ -9,7 +9,7 @@ def get_attn_wrapper(cls):
         ref_bank = None
         is_cfg = True
         style_fidelity = 0.5
-        ref_attn_weight = 0.5
+        ref_attn_weight = 1.0
         attention_bound = 1.0
         hidden_states = None
 
@@ -25,18 +25,10 @@ def get_attn_wrapper(cls):
                 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])
-                    elif len(self.ref_bank) == 1:
-                        bank = torch.cat([self.ref_bank] * len(norm_hidden_states))
-                    else:
-                        bank = self.ref_bank
-                    bank_hidden_states = torch.cat([bank, norm_hidden_states], dim=1)
+                    bank = self.ref_bank
+
+                    self.ref_bank = None
+                    bank_hidden_states = torch.cat([norm_hidden_states, bank], dim=1)
                     attn_output_uc = (
                         super().forward(
                             norm_hidden_states,
@@ -52,7 +44,7 @@ def get_attn_wrapper(cls):
                                 + [0] * (norm_hidden_states.shape[0] // 2)
                             ).to(norm_hidden_states.device).bool()
                         attn_sub =  super().forward(
-                                norm_hidden_states[uc_mask],
+                                norm_hidden_states[uc_mask].clone(),
                                 norm_hidden_states[uc_mask].clone(),
                                 **kwargs
                             )

modules/ref_controller.py

@@ -17,23 +17,28 @@ class RefController:
         self.modules.append(module)
 
     def set_mode(self, mode: RefMode):
+        print('------- CTRL SETTING MODE TO', mode, 'FOR', len(self.modules))
         for module in self.modules:
             module.ref_mode = mode
 
     def clear_modules(self):
+        print('------- CTRL CLEARING MODULES FOR', len(self.modules))
         for module in self.modules:
             module.ref_bank = None
         self.set_mode(RefMode.OFF)
 
     def set_cfg(self, cfg: bool):
+        print('------- CTRL SETTING CFG TO', cfg, 'FOR', len(self.modules))
         for module in self.modules:
             module.is_cfg = cfg
 
     def set_attention_bound(self, attention_bound: float):
+        print('------- CTRL SETTING attention bound  TO', attention_bound, 'FOR', len(self.modules))
         for module in self.modules:
             module.attention_bound = attention_bound
 
-    def set_style_fidelity(self, style_fidelity: float):
+    def set_style_fidelity(self, style_fidelity: float): 
+        print('------- CTRL SETTING style fidelity  TO', style_fidelity, 'FOR', len(self.modules))
         for module in self.modules:
             module.style_fidelity = style_fidelity
 
@@ -51,5 +56,4 @@ class RefController:
             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)
+        self.set_temporal_values_to(self.default_full_length, self.default_video_length)

modules/unet_wrapper.py

@@ -0,0 +1,59 @@
+import torch
+
+from .ref_controller import RefController, RefMode
+from ..utils.ref_config import RefConfig
+from ..utils.noise_utils import add_noise
+
+def get_unet_wrapper(cls, ref_controller: RefController):
+    class RefUNet(cls):
+        is_ref = True
+        def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+            if 'ref_config' not in transformer_options:
+                return super().forward(x, 
+                                timesteps=timesteps,
+                                context=context,
+                                y=y,
+                                control=control,
+                                transformer_options=transformer_options,
+                                **kwargs)
+            
+            ref_config: RefConfig = transformer_options['ref_config']
+            ref_latent = ref_config.ref_latent
+            ref_latent = torch.cat([ref_latent]*2)
+
+            sigma = ref_config.sampling.sigma(timesteps)
+
+            ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0]).to(x.device).to(x.dtype)
+
+            ref_controller.set_cfg(True)  #TODO
+            ref_controller.set_attention_bound(ref_config.attention_bound)
+            ref_controller.set_style_fidelity(ref_config.style_fidelity)
+
+            ref_controller.clear_modules()
+            try:
+                ref_controller.set_mode(RefMode.WRITE)
+                super().forward(ref_latent_noised, 
+                                timesteps=timesteps,
+                                context=context,
+                                y=y,
+                                control=None,
+                                transformer_options=transformer_options,
+                                **kwargs)
+                
+                ref_controller.set_mode(RefMode.READ)
+                output = super().forward(x, 
+                                timesteps=timesteps,
+                                context=context,
+                                y=y,
+                                control=control,
+                                transformer_options=transformer_options,
+                                **kwargs)
+                ref_controller.set_mode(RefMode.OFF)
+                return output
+            # except Exception as e:
+            #     print('issue with ref')
+            #     raise e
+            finally:
+                ref_controller.clear_modules()
+        
+    return RefUNet