Upload app.py

app.py

@@ -11,9 +11,10 @@ import numpy as np
 import imageio
 from easydict import EasyDict as edict
 from PIL import Image
+from Amodal3R.pipelines import Amodal3RImageTo3DPipeline
 from trellis.pipelines import TrellisImageTo3DPipeline
-from trellis.representations import Gaussian, MeshExtractResult
-from trellis.utils import render_utils, postprocessing_utils
+from Amodal3R.representations import Gaussian, MeshExtractResult
+from Amodal3R.utils import render_utils, postprocessing_utils
 
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -31,16 +32,16 @@ def end_session(req: gr.Request):
     shutil.rmtree(user_dir)
 
 
-def preprocess_image(image: Image.Image) -> Image.Image:
-    """
-    Preprocess the input image.
-    Args:
-        image (Image.Image): The input image.
-    Returns:
-        Image.Image: The preprocessed image.
-    """
-    processed_image = pipeline.preprocess_image(image)
-    return processed_image
+# def preprocess_image(image: Image.Image) -> Image.Image:
+#     """
+#     Preprocess the input image.
+#     Args:
+#         image (Image.Image): The input image.
+#     Returns:
+#         Image.Image: The preprocessed image.
+#     """
+#     processed_image = pipeline.preprocess_image(image)
+#     return processed_image
 
 
 def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
@@ -108,6 +109,7 @@ def get_seed(randomize_seed: bool, seed: int) -> int:
 @spaces.GPU
 def image_to_3d(
     image: Image.Image,
+    mask: Image.Image,
     seed: int,
     ss_guidance_strength: float,
     ss_sampling_steps: int,
@@ -135,6 +137,7 @@ def image_to_3d(
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     outputs = pipeline.run_multi_image(
         [image],
+        [mask],
         seed=seed,
         formats=["gaussian", "mesh"],
         preprocess_image=False,
@@ -283,8 +286,8 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     with gr.Row():
         with gr.Column():
             with gr.Tabs() as input_tabs:
-                with gr.Tab(label="Single Image", id=0) as single_image_input_tab:
-                    image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
+                image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
+                mask_prompt = gr.Image(label="Mask Prompt", format="png", image_mode="L", type="pil", height=300)
             
             with gr.Accordion(label="Generation Settings", open=False):
                 seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
@@ -328,11 +331,12 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     demo.unload(end_session)
 
     
-    image_prompt.upload(
-        preprocess_image,
-        inputs=[image_prompt],
-        outputs=[image_prompt],
-    )
+    # image_prompt.upload(
+    #     lambda x:x,
+    #     inputs=[image_prompt],
+    #     outputs=[image_prompt],
+    # )
+
 
     generate_btn.click(
         get_seed,
@@ -340,7 +344,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[seed],
     ).then(
         image_to_3d,
-        inputs=[image_prompt, is_multiimage, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
+        inputs=[image_prompt, mask_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
         outputs=[output_buf, video_output],
     ).then(
         lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
@@ -378,7 +382,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+    pipeline = pipeline = Amodal3RImageTo3DPipeline.from_pretrained("Sm0kyWu/Amodal3R")
     pipeline.cuda()
     try:
         pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg