Update gradio_app.py

gradio_app.py

@@ -212,14 +212,17 @@ def get_image(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, i
         # z = get_latent_z(model, videos) #bc,1,hw
         videos = repeat(videos, 'b c t h w -> b c (repeat t) h w', repeat=frames//2)
 
-        img_tensor2 = torch.from_numpy(image2).permute(2, 0, 1).float().to(model.device)
-        img_tensor2 = (img_tensor2 / 255. - 0.5) * 2
-        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)
-        
-        
-        videos = torch.cat([videos, videos2], dim=2)
+        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
+            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)
+            
+            videos = torch.cat([videos, videos2], dim=2)
+        else:
+            videos = torch.cat([videos, videos], dim=2)
+            
         z, hs = get_latent_z_with_hidden_states(model, videos)
 
         img_tensor_repeat = torch.zeros_like(z)