Update gradio_app.py

gradio_app.py

@@ -58,13 +58,22 @@ def create_rgba_image(rgb_image: Image.Image, mask: np.ndarray = None) -> Image.
 
 def create_batch(input_image: Image.Image) -> dict[str, Any]:
     """Prepare image batch for model input."""
-    # Convert input image to numpy array and normalize
-    img_array = np.array(input_image.resize((COND_WIDTH, COND_HEIGHT))).astype(np.float32) / 255.0
+    # Resize and convert input image to numpy array
+    resized_image = input_image.resize((COND_WIDTH, COND_HEIGHT))
+    img_array = np.array(resized_image).astype(np.float32) / 255.0
     print("[debug] img_array shape:", img_array.shape)
     
     # Extract RGB and alpha channels
-    rgb = torch.from_numpy(img_array[..., :3]).float()
-    mask = torch.from_numpy(img_array[..., 3:4]).float()
+    if img_array.shape[-1] == 4:  # RGBA
+        rgb = img_array[..., :3]
+        mask = img_array[..., 3:4]
+    else:  # RGB
+        rgb = img_array
+        mask = np.ones((*img_array.shape[:2], 1), dtype=np.float32)
+    
+    # Convert to tensors
+    rgb = torch.from_numpy(rgb).float()
+    mask = torch.from_numpy(mask).float()
     print("[debug] rgb tensor shape:", rgb.shape)
     print("[debug] mask tensor shape:", mask.shape)
     
@@ -76,15 +85,16 @@ def create_batch(input_image: Image.Image) -> dict[str, Any]:
     rgb_cond = torch.lerp(bg_tensor, rgb, mask)
     print("[debug] rgb_cond shape:", rgb_cond.shape)
 
-    # Note: We need to permute the tensors to match the expected shape
-    rgb_cond = rgb_cond.permute(2, 0, 1)  # Change from [H, W, C] to [C, H, W]
-    mask = mask.permute(2, 0, 1)  # Change from [H, W, 1] to [1, H, W]
+    # Permute the tensors to match the expected shape [B, C, H, W]
+    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]
+    
     print("[debug] rgb_cond after permute shape:", rgb_cond.shape)
     print("[debug] mask after permute shape:", mask.shape)
 
     batch = {
-        "rgb_cond": rgb_cond.unsqueeze(0),  # Add batch dimension
-        "mask_cond": mask.unsqueeze(0),
+        "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),
@@ -112,25 +122,23 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
             guidance_scale=0.0
         ).images[0]
         
-        # Process the generated image
         print("[debug] converting the image to rgb")
         rgb_image = generated_image.convert('RGB')
         
-        # Remove background
         print("[debug] removing the background by calling bg_remover.process(rgb_image)")
         no_bg_image = bg_remover.process(rgb_image)
         
-        # Convert to numpy array to extract mask
         print("[debug] converting to numpy array to extract the mask")
         no_bg_array = np.array(no_bg_image)
-        mask = (no_bg_array.sum(axis=2) > 0).astype(np.float32)
         
-        # Create RGBA image
+        # Create mask based on RGB values
+        mask = ((no_bg_array > 0).any(axis=2)).astype(np.float32)
+        mask = np.expand_dims(mask, axis=2)  # Add channel dimension
+        
         print("[debug] creating the RGBA image using create_rgba_image(rgb_image, mask)")
         rgba_image = create_rgba_image(rgb_image, mask)
         
-        # Auto crop with foreground
-        print(f"[debug] auto-cropping the rgba_image using spar3d_utils.foreground_crop(...). newsize=(COND_WIDTH, COND_HEIGHT) = ({COND_WIDTH}, {COND_HEIGHT})")
+        print(f"[debug] auto-cropping the rgba_image using spar3d_utils.foreground_crop(...)")
         processed_image = spar3d_utils.foreground_crop(
             rgba_image,
             crop_ratio=1.3,
@@ -138,8 +146,8 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
             no_crop=False
         )
 
+        # Forward pass through SPAR3D
         print("[debug] preparing the batch by calling create_batch(processed_image)")
-        # Prepare batch for 3D generation
         batch = create_batch(processed_image)
         batch = {k: v.to(device) for k, v in batch.items()}
 
@@ -147,6 +155,24 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
         with torch.no_grad():
             print("[debug] calling torch.autocast(....) to generate the mesh")
             with torch.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu', dtype=torch.bfloat16):
+                # Add point cloud conditioning to match expected input
+                if "pc_cond" not in batch:
+                    # Sample tokens from model's diffusion process
+                    cond_tokens = spar3d_model.forward_pdiff_cond(batch)
+                    sample_iter = spar3d_model.sampler.sample_batch_progressive(
+                        1,  # batch size
+                        cond_tokens,
+                        guidance_scale=3.0,
+                        device=device,
+                    )
+                    for x in sample_iter:
+                        samples = x["xstart"]
+                    # Add point cloud to batch
+                    batch["pc_cond"] = samples.permute(0, 2, 1).float()
+                    batch["pc_cond"] = spar3d_utils.normalize_pc_bbox(batch["pc_cond"])
+                    # Subsample to 512 points
+                    batch["pc_cond"] = batch["pc_cond"][:, torch.randperm(batch["pc_cond"].shape[1])[:512]]
+
                 trimesh_mesh, _ = spar3d_model.generate_mesh(
                     batch,
                     1024,  # texture_resolution