Update gradio_app.py

gradio_app.py

@@ -71,68 +71,77 @@ def create_batch(input_image: Image.Image) -> dict[str, Any]:
         rgb = img_array
         mask = np.ones((*img_array.shape[:2], 1), dtype=np.float32)
     
-    # Convert to tensors and keep in channel-last format initially
+    # Convert to tensors while keeping channel-last format
     rgb = torch.from_numpy(rgb).float()  # [H, W, 3]
     mask = torch.from_numpy(mask).float()  # [H, W, 1]
     print("[debug] rgb tensor shape:", rgb.shape)
     print("[debug] mask tensor shape:", mask.shape)
     
-    # Create background blend
-    bg_tensor = torch.tensor(BACKGROUND_COLOR)  # [3]
+    # Create background blend (match channel-last format)
+    bg_tensor = torch.tensor(BACKGROUND_COLOR).view(1, 1, 3)  # [1, 1, 3]
     print("[debug] bg_tensor shape:", bg_tensor.shape)
     
-    # Blend RGB with background using mask
-    rgb_cond = torch.lerp(
-        bg_tensor.view(1, 1, 3),  # [1, 1, 3]
-        rgb,  # [H, W, 3]
-        mask  # [H, W, 1]
-    )
+    # Blend RGB with background using mask (all in channel-last format)
+    rgb_cond = torch.lerp(bg_tensor, rgb, mask)  # [H, W, 3]
     print("[debug] rgb_cond shape after blend:", rgb_cond.shape)
 
-    # Permute the tensors to [B, C, H, W] format at the end
-    rgb_cond = rgb_cond.permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]
-    mask = mask.permute(2, 0, 1).unsqueeze(0)  # [1, 1, H, W]
+    # Move channels to correct dimension and add batch dimension
+    # Important: For SPAR3D image tokenizer, we need [B, H, W, C] format
+    rgb_cond = rgb_cond.unsqueeze(0)  # [1, H, W, 3]
+    mask = mask.unsqueeze(0)  # [1, H, W, 1]
     
     print("[debug] rgb_cond final shape:", rgb_cond.shape)
     print("[debug] mask final shape:", mask.shape)
 
+    # Create the batch dictionary
     batch = {
-        "rgb_cond": rgb_cond,
-        "mask_cond": mask,
-        "c2w_cond": c2w_cond.unsqueeze(0),
-        "intrinsic_cond": intrinsic.unsqueeze(0),
-        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
+        "rgb_cond": rgb_cond,  # [1, H, W, 3]
+        "mask_cond": mask,  # [1, H, W, 1]
+        "c2w_cond": c2w_cond.unsqueeze(0),  # [1, 4, 4]
+        "intrinsic_cond": intrinsic.unsqueeze(0),  # [1, 3, 3]
+        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),  # [1, 3, 3]
     }
     
-    # Final shapes check
+    print("\nFinal batch shapes:")
     for k, v in batch.items():
         print(f"[debug] {k} final shape:", v.shape)
+    print("\nrgb_cond max:", batch["rgb_cond"].max())
+    print("rgb_cond min:", batch["rgb_cond"].min())
+    print("mask_cond unique values:", torch.unique(batch["mask_cond"]))
         
     return batch
 
 def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
     """Process batch through model and generate point cloud."""
-    print("[debug] Starting forward_model")
+    print("\n[debug] Starting forward_model")
     print("[debug] Input rgb_cond shape:", batch["rgb_cond"].shape)
-    
-    # Ensure input is in correct format [B, C, H, W]
-    if batch["rgb_cond"].shape[1] != 3:
-        batch["rgb_cond"] = batch["rgb_cond"].permute(0, 3, 1, 2)
-    if batch["mask_cond"].shape[1] != 1:
-        batch["mask_cond"] = batch["mask_cond"].permute(0, 3, 1, 2)
-        
-    print("[debug] Processed rgb_cond shape:", batch["rgb_cond"].shape)
-    print("[debug] Processed mask_cond shape:", batch["mask_cond"].shape)
+    print("[debug] Input mask_cond shape:", batch["mask_cond"].shape)
     
     batch_size = batch["rgb_cond"].shape[0]
+    assert batch_size == 1, f"Expected batch size 1, got {batch_size}"
+    
+    # Print value ranges for debugging
+    print("\nValue ranges:")
+    print("rgb_cond max:", batch["rgb_cond"].max())
+    print("rgb_cond min:", batch["rgb_cond"].min())
+    print("mask_cond unique values:", torch.unique(batch["mask_cond"]))
     
     # Generate point cloud tokens
-    print("[debug] Generating point cloud tokens")
-    cond_tokens = system.forward_pdiff_cond(batch)
-    print("[debug] cond_tokens shape:", cond_tokens.shape)
+    print("\n[debug] Generating point cloud tokens")
+    try:
+        cond_tokens = system.forward_pdiff_cond(batch)
+        print("[debug] cond_tokens shape:", cond_tokens.shape)
+    except Exception as e:
+        print("\n[ERROR] Failed in forward_pdiff_cond:")
+        print(e)
+        print("\nInput tensor properties:")
+        print("rgb_cond dtype:", batch["rgb_cond"].dtype)
+        print("rgb_cond device:", batch["rgb_cond"].device)
+        print("rgb_cond requires_grad:", batch["rgb_cond"].requires_grad)
+        raise
     
     # Sample points
-    print("[debug] Sampling points")
+    print("\n[debug] Sampling points")
     sample_iter = system.sampler.sample_batch_progressive(
         batch_size,
         cond_tokens,
@@ -145,7 +154,6 @@ def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
         samples = x["xstart"]
     
     print("[debug] samples shape before permute:", samples.shape)
-    # Convert samples to point cloud format
     pc_cond = samples.permute(0, 2, 1).float()
     print("[debug] pc_cond shape after permute:", pc_cond.shape)