.

main.py

@@ -323,7 +323,7 @@ def predict_next_frame(previous_frames: List[np.ndarray], previous_actions: List
     # Convert the generated frame to the correct format
     new_frame = new_frame.transpose(1, 2, 0)
     print (new_frame.max(), new_frame.min())
-    new_frame = new_frame * data_std + data_mean
+    #new_frame = new_frame * data_std + data_mean
     new_frame_denormalized = denormalize_image(new_frame, source_range=(-1, 1))
     
     # Draw the trace of previous actions

utils.py

@@ -8,7 +8,7 @@ from huggingface_hub import hf_hub_download
 import json
 import os
 import time
-DEBUG = True
+DEBUG = False
 
 def load_model_from_config(config_path, model_name, device='cuda'):
     # Load the config file
@@ -38,11 +38,13 @@ def sample_frame(model: LatentDiffusion, prompt: str, image_sequence: torch.Tens
         #uc = model.enc_concat_seq(uc, u_dict, 'c_concat')
         
         c_dict = {'c_crossattn': prompt, 'c_concat': image_sequence}
+        
         model.eval()
-        c = model.get_learned_conditioning(c_dict)
-        print (c['c_crossattn'].shape)
-        print (c['c_crossattn'][0])
+        #c = model.get_learned_conditioning(c_dict)
+        #print (c['c_crossattn'].shape)
+        #print (c['c_crossattn'][0])
         print (prompt)
+        c = {}
         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)
@@ -91,7 +93,12 @@ def sample_frame(model: LatentDiffusion, prompt: str, image_sequence: torch.Tens
             #x_samples_ddim = torch.zeros((1, 3, 384, 512))
             #x_samples_ddim[:, :, 128:128+48, 160:160+64] = samples_ddim[:, :3]
         else:
-            x_samples_ddim = model.decode_first_stage(samples_ddim)
+            data_mean = -0.54
+            data_std = 6.78
+            data_min = -27.681446075439453
+            data_max = 30.854148864746094
+            x_samples_ddim = x_samples_ddim * data_std + data_mean
+            x_samples_ddim = model.decode_first_stage(x_samples_ddim)
         #x_samples_ddim = pos_map.to(c['c_concat'].device).unsqueeze(0).expand(-1, 3, -1, -1)
         #x_samples_ddim = model.decode_first_stage(x_samples_ddim)
         #x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)