Update app_gradio.py

app_gradio.py

@@ -67,12 +67,23 @@ def prepare_latents(pipe, x_aug):
 
 @torch.no_grad()
 def invert(pipe, inv, load_name, device="cuda", dtype=torch.bfloat16):
-    input_img = [load_image(load_name, 256).to(device, dtype=dtype).unsqueeze(1)] * 5
-    input_img = torch.cat(input_img, dim=1)
-    latents = prepare_latents(pipe, input_img).to(torch.bfloat16)
+    # Load and process the image
+    input_img = load_image(load_name, target_size=256).to(device, dtype=torch.float32)  # Shape: (1, C, H, W)
+    input_img = input_img.unsqueeze(1).repeat(1, 5, 1, 1, 1)  # Add time dimension and repeat for T=5
+    # Shape: (B=1, T=5, C=3, H=256, W=256)
+    
+    # Convert image to latent space
+    latents = prepare_latents(pipe, input_img).to(dtype)  # Shape: (B, latent_dim, T, H/8, W/8)
+
+    # Configure the inversion process
     inv.set_timesteps(25)
-    id_latents = dd_inversion(pipe, inv, video_latent=latents, num_inv_steps=25, prompt="")[-1].to(dtype)
-    return torch.mean(id_latents, dim=2, keepdim=True)
+
+    # Perform inversion and extract final latent representation
+    id_latents = dd_inversion(pipe, inv, video_latent=latents, num_inv_steps=25, prompt="")[-1]
+    id_latents = id_latents.to(dtype)  # Ensure correct dtype
+    id_latents = torch.mean(id_latents, dim=2, keepdim=True)  # Shape: (B, latent_dim, 1, H/8, W/8)
+
+    return id_latents
 
 def load_primary_models(pretrained_model_path):
     return (