fix: handle non-RGB images in BiRefNet (CMYK, grayscale, LA modes)

Always convert to RGB before the transform pipeline to prevent
shape mismatch errors with non-standard image modes.
Fixed in both utils_birefnet.py and trellis2/pipelines/rembg/BiRefNet.py.

trellis2/pipelines/rembg/BiRefNet.py

@@ -30,13 +30,18 @@ class BiRefNet:
         
     def __call__(self, image: Image.Image) -> Image.Image:
         image_size = image.size
-        input_images = self.transform_image(image).unsqueeze(0).to("cuda")
+        # Always convert to RGB for the transform (handles RGBA, L, LA, CMYK, P, etc.)
+        rgb_image = image.convert('RGB')
+
+        input_images = self.transform_image(rgb_image).unsqueeze(0).to("cuda")
         # Prediction
         with torch.no_grad():
             preds = self.model(input_images)[-1].sigmoid().cpu()
         pred = preds[0].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(image_size)
-        image.putalpha(mask)
-        return image
+        # Convert to RGBA so putalpha works regardless of the original mode
+        rgba_image = rgb_image.convert('RGBA')
+        rgba_image.putalpha(mask)
+        return rgba_image
     

utils_birefnet.py

@@ -29,16 +29,17 @@ class BiRefNet:
         
     def __call__(self, image: Image.Image) -> Image.Image:
         image_size = image.size
-        # Handle alpha channel if present
-        if image.mode == 'RGBA':
-            image = image.convert('RGB')
-            
-        input_images = self.transform_image(image).unsqueeze(0).to("cuda")
+        # Always convert to RGB for the transform (handles RGBA, L, LA, CMYK, P, etc.)
+        rgb_image = image.convert('RGB')
+
+        input_images = self.transform_image(rgb_image).unsqueeze(0).to("cuda")
         # Prediction
         with torch.no_grad():
             preds = self.model(input_images)[-1].sigmoid().cpu()
         pred = preds[0].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(image_size)
-        image.putalpha(mask)
-        return image
+        # Convert to RGBA so putalpha works regardless of the original mode
+        rgba_image = rgb_image.convert('RGBA')
+        rgba_image.putalpha(mask)
+        return rgba_image