update unet

unet/mv_unet.py

@@ -2,6 +2,7 @@ import torch
 from typing import Optional, Tuple, Union
 from diffusers import UNet2DConditionModel
 from diffusers.models.attention_processor import Attention
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
 
 
 def switch_multiview_processor(model, enable_filter=lambda x:True):
@@ -160,4 +161,17 @@ class UnifieldWrappedUNet(UNet2DConditionModel):
             ) 
         
         switch_multiview_processor(self, enable_filter=lambda name: name.endswith(f"{multiview_attn_position}.processor"))
-        
+
+    def __call__(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        condition_latens: torch.Tensor = None,
+        class_labels: Optional[torch.Tensor] = None,  
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        
+        sample = torch.cat([sample, condition_latens], dim=1)
+        return self.forward( 
+            sample, timestep, encoder_hidden_states, class_labels=class_labels,  
+        )