fix cfg1 and test noise

modules/attn_wrapper.py

@@ -32,6 +32,8 @@ def get_attn_wrapper(cls):
                         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)
@@ -42,20 +44,20 @@ def get_attn_wrapper(cls):
                             **kwargs
                         )
                     )
-                    attn_output_c = attn_output_uc.clone()
+                    attn_output = attn_output_uc
                     if self.is_cfg and self.style_fidelity > 0:
+                        attn_output_c = attn_output_uc.clone()
                         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],
+                        attn_sub =  super().forward(
                                 norm_hidden_states[uc_mask],
+                                norm_hidden_states[uc_mask].clone(),
                                 **kwargs
                             )
-                        )
-                    attn_output = self.style_fidelity * attn_output_c + (1.0 - self.style_fidelity) * attn_output_uc
+                        attn_output_c[uc_mask] = attn_sub
+                        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

modules/ref_controller.py

@@ -27,13 +27,13 @@ class RefController:
 
     def set_cfg(self, cfg: bool):
         for module in self.modules:
-            module.cfg = cfg
+            module.is_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): 
+    def set_style_fidelity(self, style_fidelity: float):
         for module in self.modules:
             module.style_fidelity = style_fidelity
 
@@ -51,4 +51,5 @@ 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)