update app.py

app.py

@@ -15,22 +15,28 @@ from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-
+# Constants
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
+# Initialize pipeline at the module level
+pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+pipeline.cuda()
+try:
+    # Preload rembg
+    pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
+except:
+    pass
 
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
     
-    
 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: Image.Image) -> Image.Image:
     """
     Preprocess the input image.
@@ -44,7 +50,6 @@ def preprocess_image(image: Image.Image) -> Image.Image:
     processed_image = pipeline.preprocess_image(image)
     return processed_image
 
-
 def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
     """
     Preprocess a list of input images.
@@ -59,7 +64,6 @@ def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image
     processed_images = [pipeline.preprocess_image(image) for image in images]
     return processed_images
 
-
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     return {
         'gaussian': {
@@ -76,7 +80,6 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
         },
     }
     
-    
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
@@ -99,14 +102,12 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     
     return gs, mesh
 
-
 def get_seed(randomize_seed: bool, seed: int) -> int:
     """
     Get the random seed.
     """
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
-
 @spaces.GPU
 def image_to_3d(
     image: Image.Image,
@@ -122,21 +123,6 @@ def image_to_3d(
 ) -> Tuple[dict, str]:
     """
     Convert an image to a 3D model.
-
-    Args:
-        image (Image.Image): The input image.
-        multiimages (List[Tuple[Image.Image, str]]): The input images in multi-image mode.
-        is_multiimage (bool): Whether is in multi-image mode.
-        seed (int): The random seed.
-        ss_guidance_strength (float): The guidance strength for sparse structure generation.
-        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
-        slat_guidance_strength (float): The guidance strength for structured latent generation.
-        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
-        multiimage_algo (Literal["multidiffusion", "stochastic"]): The algorithm for multi-image generation.
-
-    Returns:
-        dict: The information of the generated 3D model.
-        str: The path to the video of the 3D model.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     if not is_multiimage:
@@ -179,7 +165,6 @@ def image_to_3d(
     torch.cuda.empty_cache()
     return state, video_path
 
-
 @spaces.GPU(duration=90)
 def extract_glb(
     state: dict,
@@ -189,14 +174,6 @@ def extract_glb(
 ) -> Tuple[str, str]:
     """
     Extract a GLB file from the 3D model.
-
-    Args:
-        state (dict): The state of the generated 3D model.
-        mesh_simplify (float): The mesh simplification factor.
-        texture_size (int): The texture resolution.
-
-    Returns:
-        str: The path to the extracted GLB file.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh = unpack_state(state)
@@ -206,17 +183,10 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-
 @spaces.GPU
 def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     """
     Extract a Gaussian file from the 3D model.
-
-    Args:
-        state (dict): The state of the generated 3D model.
-
-    Returns:
-        str: The path to the extracted Gaussian file.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, _ = unpack_state(state)
@@ -225,7 +195,6 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     torch.cuda.empty_cache()
     return gaussian_path, gaussian_path
 
-
 def prepare_multi_example() -> List[Image.Image]:
     multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
     images = []
@@ -239,7 +208,6 @@ def prepare_multi_example() -> List[Image.Image]:
         images.append(Image.fromarray(np.concatenate(_images, axis=1)))
     return images
 
-
 def split_image(image: Image.Image) -> List[Image.Image]:
     """
     Split an image into multiple views.
@@ -254,7 +222,6 @@ def split_image(image: Image.Image) -> List[Image.Image]:
         images.append(Image.fromarray(image[:, s:e+1]))
     return [preprocess_image(image) for image in images]
 
-
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
@@ -401,14 +368,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         lambda: gr.Button(interactive=False),
         outputs=[download_glb],
     )
-    
 
 # Launch the Gradio app
 if __name__ == "__main__":
-    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()
+    demo.launch()