update

app.py

@@ -35,22 +35,6 @@ TMP_DIR = os.path.join(
 )
 os.makedirs(TMP_DIR, exist_ok=True)
 
-RBG_REMOVER = RembgRemover()
-SAM_PREDICTOR = SAMPredictor(model_type="vit_h")
-# DELIGHT = DelightingModel()
-# IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
-# PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
-#     "JeffreyXiang/TRELLIS-image-large"
-# )
-# # PIPELINE.cuda()
-
-IMAGE_BUFFER = {}
-# SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
-# GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
-# AESTHETIC_CHECKER = ImageAestheticChecker()
-# CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
-# URDF_CONVERTOR = URDFGenerator(GPT_CLIENT, render_view_num=4)
-
 
 def start_session(req: gr.Request) -> None:
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -233,11 +217,7 @@ with gr.Blocks(
                         )
                     ],
                     inputs=[image_prompt],
-                    fn=partial(
-                        preprocess_image_fn,
-                        model=RBG_REMOVER,
-                        buffer=IMAGE_BUFFER,
-                    ),
+                    fn=preprocess_image_fn,
                     outputs=[image_prompt],
                     run_on_click=True,
                     examples_per_page=32,
@@ -305,7 +285,7 @@ with gr.Blocks(
     )
 
     image_prompt.upload(
-        partial(preprocess_image_fn, model=gr.State(value=RBG_REMOVER), buffer=gr.State(value=IMAGE_BUFFER)),
+        preprocess_image_fn,
         inputs=[image_prompt],
         outputs=[image_prompt],
     )
@@ -431,7 +411,7 @@ with gr.Blocks(
             slat_guidance_strength,
             slat_sampling_steps,
             gr.State(lambda: IMAGE_BUFFER),
-            # gr.State(lambda: PIPELINE),
+            gr.State(lambda: PIPELINE),
             gr.State(lambda: TMP_DIR),
             image_seg_sam,
             is_samimage,
@@ -448,8 +428,8 @@ with gr.Blocks(
             output_buf,
             project_delight,
             gr.State(lambda: TMP_DIR),
-            # gr.State(lambda: DELIGHT),
-            # gr.State(lambda: IMAGESR_MODEL),
+            gr.State(lambda: DELIGHT),
+            gr.State(lambda: IMAGESR_MODEL),
         ],
         outputs=[
             model_output_mesh,
@@ -472,9 +452,9 @@ with gr.Blocks(
             mass_range_text,
             asset_version_text,
             gr.State(lambda: TMP_DIR),
-            # gr.State(lambda: URDF_CONVERTOR),
+            gr.State(lambda: URDF_CONVERTOR),
             gr.State(lambda: IMAGE_BUFFER),
-            # gr.State(lambda: CHECKERS),
+            gr.State(lambda: CHECKERS),
         ],
         outputs=[
             download_urdf,

common.py

@@ -44,6 +44,7 @@ from asset3d_gen.validators.quality_checkers import (
     MeshGeoChecker,
 )
 from asset3d_gen.validators.urdf_convertor import URDFGenerator, zip_files
+from asset3d_gen.utils.gpt_clients import GPT_CLIENT
 
 current_file_path = os.path.abspath(__file__)
 current_dir = os.path.dirname(current_file_path)
@@ -70,6 +71,23 @@ MAX_SEED = 100000
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "false"
 
 
+RBG_REMOVER = RembgRemover()
+SAM_PREDICTOR = SAMPredictor(model_type="vit_h")
+DELIGHT = DelightingModel()
+IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
+PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+    "JeffreyXiang/TRELLIS-image-large"
+)
+# PIPELINE.cuda()
+
+IMAGE_BUFFER = {}
+SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
+GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
+AESTHETIC_CHECKER = ImageAestheticChecker()
+CHECKERS = [GEO_CHECKER, SEG_CHECKER, AESTHETIC_CHECKER]
+URDF_CONVERTOR = URDFGenerator(GPT_CLIENT, render_view_num=4)
+
+
 @spaces.GPU
 def render_mesh(sample, extrinsics, intrinsics, options={}, **kwargs):
     renderer = MeshRenderer()
@@ -126,21 +144,16 @@ def render_video(
 @spaces.GPU
 def preprocess_image_fn(
     image: str | np.ndarray | Image.Image,
-    model: DelightingModel | RembgRemover,
-    buffer: dict = None,
 ) -> Image.Image:
     if isinstance(image, str):
         image = Image.open(image)
     elif isinstance(image, np.ndarray):
         image = Image.fromarray(image)
 
-    if buffer is not None:
-        buffer["raw_image"] = image
+    if IMAGE_BUFFER is not None:
+        IMAGE_BUFFER["raw_image"] = image
 
-    if isinstance(model, DelightingModel):
-        image = model(image, preprocess=True, target_wh=(512, 512))
-    elif isinstance(model, RembgRemover):
-        image = model(image)
+    image = RBG_REMOVER(image)
     image = trellis_preprocess(image)
 
     return image