debugging

app.py

@@ -21,7 +21,7 @@ TMP_DIR = "/tmp/Trellis-demo"
 
 os.makedirs(TMP_DIR, exist_ok=True)
 
-
+@spaces.GPU
 def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     """
     Preprocess the input image.
@@ -34,6 +34,7 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
         Image.Image: The preprocessed image.
     """
     trial_id = str(uuid.uuid4())
+    preload()
     processed_image = pipeline.preprocess_image(image)
     processed_image.save(f"{TMP_DIR}/{trial_id}.png")
     return trial_id, processed_image
@@ -100,6 +101,9 @@ def image_to_3d(trial_id: str, seed: int, randomize_seed: bool, ss_guidance_stre
     """
     if randomize_seed:
         seed = np.random.randint(0, MAX_SEED)
+
+
+    preload()
     outputs = pipeline.run(
         Image.open(f"{TMP_DIR}/{trial_id}.png"),
         seed=seed,
@@ -202,19 +206,19 @@ with gr.Blocks() as demo:
     trial_id = gr.Textbox(visible=False)
     output_buf = gr.State()
 
-    # Example images at the bottom of the page
-    with gr.Row():
-        examples = gr.Examples(
-            examples=[
-                f'assets/example_image/{image}'
-                for image in os.listdir("assets/example_image")
-            ],
-            inputs=[image_prompt],
-            fn=preprocess_image,
-            outputs=[trial_id, image_prompt],
-            run_on_click=True,
-            examples_per_page=64,
-        )
+    # # Example images at the bottom of the page
+    # with gr.Row():
+    #     examples = gr.Examples(
+    #         examples=[
+    #             f'assets/example_image/{image}'
+    #             for image in os.listdir("assets/example_image")
+    #         ],
+    #         inputs=[image_prompt],
+    #         fn=preprocess_image,
+    #         outputs=[trial_id, image_prompt],
+    #         run_on_click=True,
+    #         examples_per_page=64,
+    #     )
 
     # Handlers
     image_prompt.upload(
@@ -273,8 +277,13 @@ def cleanup_tmp_dir():
 cleanup_thread = threading.Thread(target=cleanup_tmp_dir)
 cleanup_thread.start()
 
+
 @spaces.GPU
 def preload():
+    global preloaded
+    if preloaded:
+        return
+    preloaded = True
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
     try:
@@ -282,9 +291,13 @@ def preload():
     except:
         pass
 
-
 # Launch the Gradio app
 if __name__ == "__main__":
     global pipeline
-    preload()
+    # pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+    # pipeline.cuda()
+    # try:
+    #     pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
+    # except:
+    #     pass
     demo.launch()