cahnge image type to PIL

app.py

@@ -62,7 +62,7 @@ def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shutil.rmtree(user_dir)
 
-def preprocess_image(image):
+def preprocess_image(image: Image.Image) -> Image.Image:
     trellis = get_trellis_pipeline()
     if trellis is None:
         # If the pipeline is not loaded, just return the original image
@@ -313,7 +313,7 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
 
 @spaces.GPU
 def image_to_3d(
-    image,
+    image: Image.Image,
     seed: int,
     ss_guidance_strength: float,
     ss_sampling_steps: int,
@@ -321,31 +321,18 @@ def image_to_3d(
     slat_sampling_steps: int,
 ) -> Tuple[dict, str]:
     try:
-        if isinstance(image, dict) and "image" in image:
-            image = image["image"]
-        
-        # If user passed multiple images
-        if isinstance(image, list):
-            input_image = []
-            for img in image:
-                if isinstance(img, dict) and "image" in img:
-                    img = img["image"]
-                if isinstance(img, np.ndarray):
-                    img = Image.fromarray(img.astype("uint8"))
-                input_image.append(img)
-        else:
-            # Single image
-            if isinstance(image, np.ndarray):
-                image = Image.fromarray(image.astype("uint8"))
-            input_image = [image]
-
+        # Load the Trellis pipeline
         pipeline = get_trellis_pipeline()
         if pipeline is None:
             return None, "Trellis pipeline is unavailable."
         pipeline.cuda()
 
+        # Preprocess image
+        image = preprocess_image(image)
+
+        # Run the pipeline
         outputs = pipeline.run(
-            input_image,
+            image,
             seed=seed,
             formats=["gaussian", "mesh"],
             preprocess_image=False,
@@ -473,7 +460,7 @@ def create_interface():
                 max_length=2048,
             )
             visual_button = gr.Button("Create visual with Flux")
-            generated_image = gr.Image(show_label=False)
+            generated_image = gr.Image(show_label=False, format="png", image_mode="RGBA", type="pil", height=300)
 
             preprocessed_button = gr.Button("Preprocess image")
             preprocessed_image = gr.Image(show_label=False)
@@ -586,7 +573,7 @@ def create_interface():
         gr.on(
             triggers=[gen3d_button.click],
             fn=image_to_3d,
-            inputs=[preprocessed_image, trellis_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
+            inputs=[generated_image, trellis_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
             outputs=[output_state, video_output],
         )
         # .then(