Update unet/mv_unet.py

unet/mv_unet.py

@@ -204,11 +204,14 @@ class SPADUnetModel(UNetModel, ModelMixin, ConfigMixin):
         timesteps = rearrange(timesteps, "n v -> (n v)")
         t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
         time = self.time_embed(t_emb)
+        print("old time: ", time.shape)
         time = rearrange(time, "(n v) d -> n v d", n=n_objects, v=n_views)
+        # 2, 4, 1280
 
         # extract txt and cam embedding (absolute) from context
         if len(context) == 2:
             txt, cam = context
+            print("txt shape", txt.shape)
         elif len(context) == 3:
             txt, cam, epi_mask = context
             txt = (txt, epi_mask)
@@ -219,13 +222,19 @@ class SPADUnetModel(UNetModel, ModelMixin, ConfigMixin):
         if x.shape[2] > 4:
             plucker, x = x[:, :, 4:], x[:, :, :4]
             txt = (*txt, plucker) if isinstance(txt, tuple) else (txt, plucker)
+            print("extracted")
 
+        # print("txt shape: ", txt.shape)
         # combine timestep and camera embedding (resnet)
-        time_cam = time + cam
+        time_cam = time # add + cam later
         del time, cam
 
         # encode
+        
         h = x.type(self.dtype)
+        print("h: ", h.shape)
+        print("time_cam: ", time_cam.shape)
+        # print("txt: ", txt.shape)
         hs = self.encode(h, time_cam, txt, self.input_blocks)
 
         # middle block