Fix

Image3DProcessor.py

@@ -40,13 +40,13 @@ class Image3DProcessor:
         self.model.load_state_dict(ckpt_loaded["model_state_dict"])
         self.model.to(self.device)
         self.model.eval()
-        
+
+    @torch.no_grad()
     def preprocess(self, input_image, preprocess_background=True, foreground_ratio=0.65):
-        # Convert bytes to PIL Image if necessary
-        if isinstance(input_image, bytes):
-            input_image = Image.open(BytesIO(input_image))
-        
+        # Create a new Rembg session
         rembg_session = rembg.new_session()
+
+        # Preprocess input image
         if preprocess_background:
             image = input_image.convert("RGB")
             image = remove_background(image, rembg_session)
@@ -56,15 +56,19 @@ class Image3DProcessor:
             image = input_image
             if image.mode == "RGBA":
                 image = set_white_background(image)
+
         image = resize_to_128(image)
+
         return image
     
     @torch.no_grad()
     def reconstruct_and_export(self, image):
-        # Convert PIL Image to NumPy array if needed
+        # Ensure the input image is a NumPy array after preprocessing
         if isinstance(image, Image.Image):
             image = np.array(image)
-        
+        elif isinstance(image, bytes):
+            image = np.array(Image.open(BytesIO(image)))
+
         image_tensor = to_tensor(image).to(self.device)
         view_to_world_source, rot_transform_quats = get_source_camera_v2w_rmo_and_quats()
         view_to_world_source = view_to_world_source.to(self.device)
@@ -113,4 +117,4 @@ class Image3DProcessor:
         with open(mesh_path, "rb") as mesh_file:
             mesh_data = mesh_file.read()
 
-        return mesh_data, video_data
+        return mesh_data, video_data

app.py

@@ -17,8 +17,10 @@ class VideoGenerator:
         self.processor = Image3DProcessor(model_cfg_path, model_repo_id, model_filename)
 
     def generate_3d_video(self, image):
-        # Process the image and create a 3D video and mesh
-        mesh_data, video_data = self.processor.reconstruct_and_export(image)
+        # Preprocess the image first
+        processed_image = self.processor.preprocess(image)
+        # Then pass it to reconstruct_and_export
+        mesh_data, video_data = self.processor.reconstruct_and_export(processed_image)
         return mesh_data, video_data
 
 class GradioApp:
@@ -64,4 +66,4 @@ class GradioApp:
 
 if __name__ == "__main__":
     app = GradioApp()
-    app.launch()
+    app.launch()