Update gradio_app.py

gradio_app.py

@@ -43,31 +43,47 @@ intrinsic, intrinsic_normed_cond = spar3d_utils.create_intrinsic_from_fov_rad(
     COND_FOVY, COND_HEIGHT, COND_WIDTH
 )
 
-def create_batch(input_image: Image) -> dict[str, Any]:
+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_batch(input_image: Image.Image) -> dict[str, Any]:
     """Prepare image batch for model input."""
-    img_cond = (
-        torch.from_numpy(
-            np.asarray(input_image.resize((COND_WIDTH, COND_HEIGHT))).astype(np.float32)
-            / 255.0
-        )
-        .float()
-        .clip(0, 1)
-    )
-    mask_cond = img_cond[:, :, -1:]
-    rgb_cond = torch.lerp(
-        torch.tensor(BACKGROUND_COLOR)[None, None, :], img_cond[:, :, :3], mask_cond
-    )
+    # Ensure input is RGBA
+    if input_image.mode != 'RGBA':
+        input_image = input_image.convert('RGBA')
+    
+    # Resize and convert to numpy array
+    resized_image = input_image.resize((COND_WIDTH, COND_HEIGHT))
+    img_array = np.array(resized_image).astype(np.float32) / 255.0
+    
+    # Split into RGB and alpha
+    rgb = img_array[..., :3]
+    alpha = img_array[..., 3:4]
+    
+    # Convert to tensors
+    rgb_tensor = torch.from_numpy(rgb).float()
+    alpha_tensor = torch.from_numpy(alpha).float()
+    
+    # Create background blend
+    bg_tensor = torch.tensor(BACKGROUND_COLOR)[None, None, :]
+    rgb_cond = torch.lerp(bg_tensor, rgb_tensor, alpha_tensor)
 
     batch = {
         "rgb_cond": rgb_cond.unsqueeze(0),
-        "mask_cond": mask_cond.unsqueeze(0),
+        "mask_cond": alpha_tensor.unsqueeze(0),
         "c2w_cond": c2w_cond.unsqueeze(0),
         "intrinsic_cond": intrinsic.unsqueeze(0),
         "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
     }
     return batch
 
-def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, height: int = 1024) -> tuple[str, Image.Image]:
+def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, height: int = 1024) -> tuple[str | None, Image.Image | None]:
     """Generate image from prompt and convert to 3D model."""
     try:
         # Generate image using FLUX
@@ -81,23 +97,26 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
             guidance_scale=0.0
         ).images[0]
         
-        # Convert PIL image to RGBA
-        input_image = generated_image.convert("RGBA")
+        # Process the generated image
+        rgb_image = generated_image.convert('RGB')
+        
+        # Remove background and get mask
+        mask = bg_remover.process_image(rgb_image)
+        mask_uint8 = (mask * 255).astype(np.uint8)
         
-        # Remove background
-        rgba_image = bg_remover.process(input_image.convert("RGB"))
-        rgba_image.putalpha(255)  # Add alpha channel
+        # Create RGBA image
+        rgba_image = create_rgba_image(rgb_image, mask_uint8)
         
-        # Auto crop
-        input_image = spar3d_utils.foreground_crop(
+        # Auto crop with foreground
+        processed_image = spar3d_utils.foreground_crop(
             rgba_image,
             crop_ratio=1.3,
             newsize=(COND_WIDTH, COND_HEIGHT),
             no_crop=False
         )
 
-        # Prepare batch
-        batch = create_batch(input_image)
+        # Prepare batch for 3D generation
+        batch = create_batch(processed_image)
         batch = {k: v.to(device) for k, v in batch.items()}
 
         # Generate mesh
@@ -120,7 +139,7 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
         return output_path, generated_image
         
     except Exception as e:
-        print(f"Error: {str(e)}")
+        print(f"Error during generation: {str(e)}")
         return None, None
 
 # Create Gradio interface