Remove temp dir and gr.State

app.py

@@ -1,8 +1,8 @@
 import os
 import pathlib
 import shlex
-import shutil
 import subprocess
+import tempfile
 
 os.environ["SPCONV_ALGO"] = "native"
 
@@ -29,25 +29,12 @@ from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import postprocessing_utils, render_utils
 
 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)
-
 
 pipeline = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
 pipeline.cuda()
 pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))  # Preload rembg
 
 
-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.
 
@@ -73,24 +60,26 @@ def preprocess_images(images: list[tuple[Image.Image, str]]) -> list[Image.Image
     return [pipeline.preprocess_image(image) for image in images]
 
 
-def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
-    return {
+def save_state_to_file(gs: Gaussian, mesh: MeshExtractResult, output_path: str) -> None:
+    state = {
         "gaussian": {
             **gs.init_params,
-            "_xyz": gs._xyz.cpu().numpy(),
-            "_features_dc": gs._features_dc.cpu().numpy(),
-            "_scaling": gs._scaling.cpu().numpy(),
-            "_rotation": gs._rotation.cpu().numpy(),
-            "_opacity": gs._opacity.cpu().numpy(),
+            "_xyz": gs._xyz,
+            "_features_dc": gs._features_dc,
+            "_scaling": gs._scaling,
+            "_rotation": gs._rotation,
+            "_opacity": gs._opacity,
         },
         "mesh": {
-            "vertices": mesh.vertices.cpu().numpy(),
-            "faces": mesh.faces.cpu().numpy(),
+            "vertices": mesh.vertices,
+            "faces": mesh.faces,
         },
     }
+    torch.save(state, output_path)
 
 
-def unpack_state(state: dict) -> tuple[Gaussian, EasyDict, str]:
+def load_state_from_file(state_path: str) -> tuple[Gaussian, EasyDict]:
+    state = torch.load(state_path)
     gs = Gaussian(
         aabb=state["gaussian"]["aabb"],
         sh_degree=state["gaussian"]["sh_degree"],
@@ -99,15 +88,15 @@ def unpack_state(state: dict) -> tuple[Gaussian, EasyDict, str]:
         opacity_bias=state["gaussian"]["opacity_bias"],
         scaling_activation=state["gaussian"]["scaling_activation"],
     )
-    gs._xyz = torch.tensor(state["gaussian"]["_xyz"], device="cuda")
-    gs._features_dc = torch.tensor(state["gaussian"]["_features_dc"], device="cuda")
-    gs._scaling = torch.tensor(state["gaussian"]["_scaling"], device="cuda")
-    gs._rotation = torch.tensor(state["gaussian"]["_rotation"], device="cuda")
-    gs._opacity = torch.tensor(state["gaussian"]["_opacity"], device="cuda")
+    gs._xyz = state["gaussian"]["_xyz"]
+    gs._features_dc = state["gaussian"]["_features_dc"]
+    gs._scaling = state["gaussian"]["_scaling"]
+    gs._rotation = state["gaussian"]["_rotation"]
+    gs._opacity = state["gaussian"]["_opacity"]
 
     mesh = EasyDict(
-        vertices=torch.tensor(state["mesh"]["vertices"], device="cuda"),
-        faces=torch.tensor(state["mesh"]["faces"], device="cuda"),
+        vertices=state["mesh"]["vertices"],
+        faces=state["mesh"]["faces"],
     )
 
     return gs, mesh
@@ -126,8 +115,7 @@ def image_to_3d(
     ss_sampling_steps: int,
     slat_guidance_strength: float,
     slat_sampling_steps: int,
-    req: gr.Request,
-) -> tuple[dict, str]:
+) -> tuple[str, str]:
     """Convert an image to a 3D model.
 
     Args:
@@ -139,11 +127,9 @@ def image_to_3d(
         slat_sampling_steps (int): The number of sampling steps for structured latent generation.
 
     Returns:
-        dict: The information of the generated 3D model.
+        str: The path to the pickle file that contains the state 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))
-
     outputs = pipeline.run(
         image,
         seed=seed,
@@ -162,69 +148,55 @@ def image_to_3d(
     video = render_utils.render_video(outputs["gaussian"][0], num_frames=120)["color"]
     video_geo = render_utils.render_video(outputs["mesh"][0], num_frames=120)["normal"]
     video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-    video_path = os.path.join(user_dir, "sample.mp4")
-    imageio.mimsave(video_path, video, fps=15)
-    state = pack_state(outputs["gaussian"][0], outputs["mesh"][0])
-    torch.cuda.empty_cache()
-    return state, video_path
+
+    with (
+        tempfile.NamedTemporaryFile(suffix=".pth", delete=False) as state_file,
+        tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as video_file,
+    ):
+        save_state_to_file(outputs["gaussian"][0], outputs["mesh"][0], state_file.name)
+        torch.cuda.empty_cache()
+        imageio.mimsave(video_file.name, video, fps=15)
+        return state_file.name, video_file.name
 
 
 @spaces.GPU(duration=90)
 def extract_glb(
-    state: dict,
+    state_path: str,
     mesh_simplify: float,
     texture_size: int,
-    req: gr.Request,
-) -> tuple[str, str]:
+) -> str:
     """Extract a GLB file from the 3D model.
 
     Args:
-        state (dict): The state of the generated 3D model.
+        state_path (str): The path to the pickle file that contains 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)
+    gs, mesh = load_state_from_file(state_path)
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(user_dir, "sample.glb")
-    glb.export(glb_path)
     torch.cuda.empty_cache()
-    return glb_path, glb_path
+    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as glb_file:
+        glb.export(glb_file.name)
+        return glb_file.name
 
 
 @spaces.GPU
-def extract_gaussian(state: dict, req: gr.Request) -> tuple[str, str]:
+def extract_gaussian(state_path: str) -> str:
     """Extract a Gaussian file from the 3D model.
 
     Args:
-        state (dict): The state of the generated 3D model.
+        state_path (str): The path to the pickle file that contains 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)
-    gaussian_path = os.path.join(user_dir, "sample.ply")
-    gs.save_ply(gaussian_path)
-    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 = []
-    for case in multi_case:
-        _images = []
-        for i in range(1, 4):
-            img = Image.open(f"assets/example_multi_image/{case}_{i}.png")
-            W, H = img.size
-            img = img.resize((int(W / H * 512), 512))
-            _images.append(np.array(img))
-        images.append(Image.fromarray(np.concatenate(_images, axis=1)))
-    return images
+    gs, _ = load_state_from_file(state_path)
+    with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as gaussian_file:
+        gs.save_ply(gaussian_file.name)
+        return gaussian_file.name
 
 
 with gr.Blocks(delete_cache=(600, 600)) as demo:
@@ -279,11 +251,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
             model_output = gr.Model3D(label="Extracted GLB/Gaussian", height=300)
 
-            with gr.Row():
-                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)
-
-    output_buf = gr.State()
+    state_file_path = gr.Textbox(visible=False)
 
     examples = gr.Examples(
         examples=sorted(pathlib.Path("assets/example_image").glob("*.png")),
@@ -294,10 +262,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         examples_per_page=64,
     )
 
-    # Handlers
-    demo.load(start_session)
-    demo.unload(end_session)
-
     image_prompt.upload(
         fn=preprocess_image,
         inputs=image_prompt,
@@ -318,40 +282,21 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             slat_guidance_strength,
             slat_sampling_steps,
         ],
-        outputs=[output_buf, video_output],
+        outputs=[state_file_path, video_output],
     ).then(
         fn=lambda: (gr.Button(interactive=True), gr.Button(interactive=True)),
         outputs=[extract_glb_btn, extract_gs_btn],
+        api_name=False,
     )
 
     video_output.clear(
         fn=lambda: (gr.Button(interactive=False), gr.Button(interactive=False)),
         outputs=[extract_glb_btn, extract_gs_btn],
+        api_name=False,
     )
 
-    extract_glb_btn.click(
-        fn=extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size],
-        outputs=[model_output, download_glb],
-    ).then(
-        fn=lambda: gr.Button(interactive=True),
-        outputs=[download_glb],
-    )
-
-    extract_gs_btn.click(
-        fn=extract_gaussian,
-        inputs=[output_buf],
-        outputs=[model_output, download_gs],
-    ).then(
-        fn=lambda: gr.Button(interactive=True),
-        outputs=[download_gs],
-    )
-
-    model_output.clear(
-        fn=lambda: gr.Button(interactive=False),
-        outputs=[download_glb],
-    )
-
+    extract_glb_btn.click(fn=extract_glb, inputs=[state_file_path, mesh_simplify, texture_size], outputs=model_output)
+    extract_gs_btn.click(fn=extract_gaussian, inputs=state_file_path, outputs=model_output)
 
 if __name__ == "__main__":
     demo.launch(mcp_server=True)