trying to tickle the dimension problem of video output tensor

serve/vote_utils.py

@@ -122,6 +122,9 @@ def vote_last_response_vg(state, vote_type, model_selector, request: gr.Request)
     else:
         print("======== video shape: ========")
         print(state.output.shape)
+        # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]
+        if state.output.shape[-1] != 3:
+            state.output = state.output.permute(0, 2, 3, 1)
         imageio.mimwrite(output_file, state.output, fps=8, quality=9)
     save_video_file_on_log_server(output_file)
 
@@ -148,6 +151,9 @@ def vote_last_response_vgm(states, vote_type, model_selectors, request: gr.Reque
         elif isinstance(state.output, torch.Tensor):
             print("======== video shape: ========")
             print(state.output.shape)
+            # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]
+            if state.output.shape[-1] != 3:
+                state.output = state.output.permute(0, 2, 3, 1)
             imageio.mimwrite(output_file, state.output, fps=8, quality=9)
         else:
             r = requests.get(state.output)
@@ -1228,6 +1234,9 @@ def generate_vg(gen_func, state, text, model_name, request: gr.Request):
     else:
         print("======== video shape: ========")
         print(state.output.shape)
+        # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]
+        if state.output.shape[-1] != 3:
+            state.output = state.output.permute(0, 2, 3, 1)
         imageio.mimwrite(output_file, state.output, fps=8, quality=9)
 
     save_video_file_on_log_server(output_file)
@@ -1343,6 +1352,9 @@ def generate_vgm(gen_func, state0, state1, text, model_name0, model_name1, reque
             print("======== video shape: ========")
             print(state.output)
             print(state.output.shape)
+            # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]
+            if state.output.shape[-1] != 3:
+                state.output = state.output.permute(0, 2, 3, 1)
             imageio.mimwrite(output_file, state.output, fps=8, quality=9)
         save_video_file_on_log_server(output_file)
     yield state0, state1, f'{VIDEO_DIR}/generation/{state0.conv_id}.mp4', f'{VIDEO_DIR}/generation/{state1.conv_id}.mp4'
@@ -1472,6 +1484,9 @@ def generate_vgm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
         else:
             print("======== video shape: ========")
             print(state.output.shape)
+            # Assuming state.output has to be a tensor with shape [num_frames, height, width, num_channels]
+            if state.output.shape[-1] != 3:
+                state.output = state.output.permute(0, 2, 3, 1)
             imageio.mimwrite(output_file, state.output, fps=8, quality=9)
         save_video_file_on_log_server(output_file)