update image_to_3d to support multiple images and handle numpy arrays

app.py

@@ -86,7 +86,8 @@ def get_image_gen_pipeline():
             device = "cuda" if torch.cuda.is_available() else "cpu"
             dtype = torch.bfloat16
             _image_gen_pipeline = DiffusionPipeline.from_pretrained(
-                "black-forest-labs/FLUX.1-schnell",
+                # "black-forest-labs/FLUX.1-schnell",
+                "black-forest-labs/FLUX.1-dev",
                 torch_dtype=dtype,
             ).to(device)
             
@@ -320,25 +321,29 @@ def image_to_3d(
     slat_sampling_steps: int,
 ) -> Tuple[dict, str]:
     try:
-        # Use a fixed temp directory instead of user-specific
-        temp_dir = os.path.join(TMP_DIR, "temp_output")
-        os.makedirs(temp_dir, exist_ok=True)
+        if isinstance(image, dict) and "image" in image:
+            image = image["image"]
         
-        # Get the pipeline using the getter function
+        # 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]
+
         pipeline = get_trellis_pipeline()
         if pipeline is None:
             return None, "Trellis pipeline is unavailable."
-        
-        # Call cuda() here in the GPU worker process
         pipeline.cuda()
-        
-        # Convert image to the right format if needed
-        if isinstance(image, np.ndarray):
-            image = Image.fromarray(image.astype('uint8'))
-            
-        # Make sure we have a list of images as expected by the pipeline
-        input_image = [image]
-        
+
         outputs = pipeline.run(
             input_image,
             seed=seed,
@@ -627,4 +632,3 @@ if __name__ == "__main__":
 
     demo = create_interface()
     demo.launch(debug=True)
-