.

main.py

@@ -223,7 +223,7 @@ def predict_next_frame(previous_frames, previous_actions: List[Tuple[str, List[i
     # Convert the image sequence to a tensor and concatenate in the channel dimension
     #image_sequence_tensor = torch.from_numpy(normalize_images(image_sequence_list, target_range=(-1, 1)))
     #image_sequence_tensor = image_sequence_tensor.to(device)
-    image_sequence_tensor = torch.cat(image_sequence, dim=1)
+    image_sequence_tensor = torch.cat(image_sequence, dim=-1)
     
     #image_sequence_tensor = (image_sequence_tensor - data_mean) / data_std
     

utils.py

@@ -45,6 +45,7 @@ def sample_frame(model: LatentDiffusion, prompt: str, image_sequence: torch.Tens
         #print (c['c_crossattn'][0])
         print (prompt)
         c = {'c_concat': image_sequence.unsqueeze(0)}
+        print (image_sequence.shape, c['c_concat'].shape)
         #c = model.enc_concat_seq(c, c_dict, 'c_concat')
         # Zero out the corresponding subtensors in c_concat for padding images
         #padding_mask = torch.isclose(image_sequence, torch.tensor(-1.0), rtol=1e-5, atol=1e-5).all(dim=(1, 2, 3)).unsqueeze(0)