Update gradio_app.py

gradio_app.py

@@ -53,35 +53,32 @@ def create_rgba_image(rgb_image: Image.Image, mask: np.ndarray = None) -> Image.
     return rgba_image
 
 def create_batch(input_image: Image.Image) -> dict[str, Any]:
-    """Prepare image batch for model input."""
-    # 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": 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
+   """Prepare image batch for model input."""
+   # 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 
+   mask_cond = img_array[..., 3:4]  # Alpha channel
+   # Blend RGB with background based on alpha
+   rgb_cond = np.clip(
+       img_array[..., :3] * mask_cond + BACKGROUND_COLOR * (1 - mask_cond), 
+       0, 
+       1
+   )
+   
+   batch = {
+       "rgb_cond": torch.from_numpy(rgb_cond).unsqueeze(0),
+       "mask_cond": torch.from_numpy(mask_cond).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 | None, Image.Image | None]:
     """Generate image from prompt and convert to 3D model."""