add option n_view

unet/config.json

@@ -64,7 +64,7 @@
   ],
   "upcast_attention": false,
   "use_linear_projection": false,
-
+  "n_view": 4, 
   "multiview_attn_position": "attn1", 
   "num_modalities": 1,
   "latent_size": 64,

unet/mv_unet.py

@@ -142,6 +142,7 @@ class UnifieldWrappedUNet(UNet2DConditionModel):
             cross_attention_norm: Optional[str] = None,
             addition_embed_type_num_heads: int = 64,
             multiview_attn_position: str = "attn1", 
+            n_views: int = 4,
             num_modalities: int = 1, 
             latent_size: int = 64, 
             multiview_chain_pose: str = "parralle", 
@@ -149,7 +150,7 @@ class UnifieldWrappedUNet(UNet2DConditionModel):
         ):   
         super().__init__(**{
                 k: v for k, v in locals().items() if k not in 
-                ["self", "kwargs", "__class__", "multiview_attn_position", "num_modalities", "latent_size", "multiview_chain_pose"]
+                ["self", "kwargs", "__class__", "multiview_attn_position", "n_views", "num_modalities", "latent_size", "multiview_chain_pose"]
             })
         
         add_multiview_processor(
@@ -170,7 +171,6 @@ class UnifieldWrappedUNet(UNet2DConditionModel):
         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,