Update gradio_app.py

gradio_app.py

@@ -43,14 +43,14 @@ intrinsic, intrinsic_normed_cond = spar3d_utils.create_intrinsic_from_fov_rad(
     COND_FOVY, COND_HEIGHT, COND_WIDTH
 )
 
-def create_rgba_image(rgb_image: Image.Image, alpha: np.ndarray = None) -> Image.Image:
-    """Create an RGBA image from RGB image and optional alpha channel."""
-    if alpha is None:
-        alpha = np.full(rgb_image.size[::-1], 255, dtype=np.uint8)
-    rgba = Image.new('RGBA', rgb_image.size)
-    rgba.paste(rgb_image)
-    rgba.putalpha(Image.fromarray(alpha))
-    return rgba
+def create_rgba_image(rgb_image: Image.Image, mask: np.ndarray = None) -> Image.Image:
+    """Create an RGBA image from RGB image and optional mask."""
+    rgba_image = rgb_image.convert('RGBA')
+    if mask is not None:
+        # Convert mask to alpha channel format
+        alpha = Image.fromarray((mask * 255).astype(np.uint8))
+        rgba_image.putalpha(alpha)
+    return rgba_image
 
 def create_batch(input_image: Image.Image) -> dict[str, Any]:
     """Prepare image batch for model input."""
@@ -100,12 +100,15 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
         # Process the generated image
         rgb_image = generated_image.convert('RGB')
         
-        # Remove background and get mask
-        mask = bg_remover.process(rgb_image)
-        mask_uint8 = (mask * 255).astype(np.uint8)
+        # Remove background
+        no_bg_image = bg_remover.process(rgb_image)
+        
+        # Convert to numpy array to extract mask
+        no_bg_array = np.array(no_bg_image)
+        mask = (no_bg_array.sum(axis=2) > 0).astype(np.float32)
         
         # Create RGBA image
-        rgba_image = create_rgba_image(rgb_image, mask_uint8)
+        rgba_image = create_rgba_image(rgb_image, mask)
         
         # Auto crop with foreground
         processed_image = spar3d_utils.foreground_crop(