Update gradio_app.py

gradio_app.py

@@ -142,6 +142,16 @@ save_fps = 8
 print("resolution:", resolution)
 print("init done.")
 
+def transpose_if_needed(tensor):
+    h = tensor.shape[-2]
+    w = tensor.shape[-1]
+    if h > w:
+        tensor = tensor.permute(0, 2, 1)
+    return tensor
+
+def untranspose(tensor):
+    ndim = tensor.ndim
+    return tensor.transpose(ndim-1, ndim-2)
 
 @spaces.GPU(duration=200)
 def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, image2=None, frame_guides=None,control_scale=0.6):
@@ -172,6 +182,7 @@ def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, i
     noise_shape = [batch_size, channels, frames, h, w]
 
     # text cond
+    transposed = False 
     with torch.no_grad(), torch.cuda.amp.autocast():
         text_emb = model.get_learned_conditioning([prompt])
         print("before control")
@@ -185,6 +196,7 @@ def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, i
                 
                 #cn_tensor = (cn_tensor / 255. - 0.5) * 2
                 cn_tensor = ( cn_tensor/255.0 )
+                cn_tensor = transpose_if_needed(cn_tensor)
                 cn_tensor_resized = transform(cn_tensor) #3,h,w
 
                 cn_video = cn_tensor_resized.unsqueeze(0).unsqueeze(2) # bc1hw
@@ -211,8 +223,10 @@ def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, i
 
         # img cond
         img_tensor = torch.from_numpy(image).permute(2, 0, 1).float().to(model.device)
+        input_h, input_w = img_tensor.shape[1:]
         img_tensor = (img_tensor / 255. - 0.5) * 2
-
+        img_tensor = transpose_if_needed(img_tensor)
+        
         image_tensor_resized = transform(img_tensor) #3,h,w
         videos = image_tensor_resized.unsqueeze(0).unsqueeze(2) # bc1hw
         print("get latent z")
@@ -222,6 +236,7 @@ def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, i
         if image2 is not None:
             img_tensor2 = torch.from_numpy(image2).permute(2, 0, 1).float().to(model.device)
             img_tensor2 = (img_tensor2 / 255. - 0.5) * 2
+            img_tensor2 = transpose_if_needed(img_tensor2)
             image_tensor_resized2 = transform(img_tensor2) #3,h,w
             videos2 = image_tensor_resized2.unsqueeze(0).unsqueeze(2) # bchw
             videos2 = repeat(videos2, 'b c t h w -> b c (repeat t) h w', repeat=frames//2)
@@ -263,6 +278,9 @@ def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, i
 
     global result_dir
     global save_fps
+    if input_h > input_w:
+        batch_samples = untranspose(batch_samples)
+        
     save_videos(batch_samples, result_dir, filenames=[prompt_str], fps=save_fps)
     print(f"Saved in {prompt_str}. Time used: {(time.time() - start):.2f} seconds")
     model = model.cpu()