Update app.py

still working on the high quality download

app.py

@@ -158,7 +158,7 @@ def extract_glb(
     mesh_simplify: float,
     texture_size: int,
     req: gr.Request,
-) -> Tuple[str, str]:
+) -> Tuple[dict, bytes]:
     """
     Extract a GLB file from the 3D model.
 
@@ -169,22 +169,25 @@ def extract_glb(
         req (gr.Request): Gradio request object.
 
     Returns:
-        Tuple[str, str]: The path to the extracted GLB file.
+        Tuple[dict, bytes]: The model state and the GLB file bytes.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh, trial_id = unpack_state(state)
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = os.path.join(user_dir, f"{trial_id}.glb")
     glb.export(glb_path)
+    # Read the GLB file as bytes
+    with open(glb_path, "rb") as f:
+        glb_bytes = f.read()
     torch.cuda.empty_cache()
-    return glb_path, glb_path
+    return state, glb_bytes
 
 # New High-Quality GLB Extraction Function
 @spaces.GPU
 def extract_glb_high_quality(
     state: dict,
     req: gr.Request,
-) -> Tuple[str, str]:
+) -> Tuple[dict, bytes]:
     """
     Extract a high-quality GLB file from the 3D model without polygon reduction.
 
@@ -193,7 +196,7 @@ def extract_glb_high_quality(
         req (gr.Request): Gradio request object.
 
     Returns:
-        Tuple[str, str]: The path to the high-quality GLB file.
+        Tuple[dict, bytes]: The model state and the high-quality GLB file bytes.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh, trial_id = unpack_state(state)
@@ -202,8 +205,11 @@ def extract_glb_high_quality(
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=0.0, texture_size=2048, verbose=False)
     glb_path = os.path.join(user_dir, f"{trial_id}_high_quality.glb")
     glb.export(glb_path)
+    # Read the GLB file as bytes
+    with open(glb_path, "rb") as f:
+        glb_bytes = f.read()
     torch.cuda.empty_cache()
-    return glb_path, glb_path
+    return state, glb_bytes
 
 # Gradio Interface Definition
 with gr.Blocks(delete_cache=(600, 600)) as demo:
@@ -324,8 +330,8 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
     # State Variables
     output_buf = gr.State()
-    glb_path_state = gr.State()  # For standard GLB
-    glb_high_quality_path_state = gr.State()  # For high-quality GLB
+    glb_bytes_state = gr.State()  # For standard GLB
+    glb_high_quality_bytes_state = gr.State()  # For high-quality GLB
 
     # Example Images
     with gr.Row():
@@ -377,7 +383,11 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     extract_glb_btn.click(
         extract_glb,
         inputs=[output_buf, mesh_simplify, texture_size],
-        outputs=[model_output, download_glb],
+        outputs=[model_output, glb_bytes_state],
+    ).then(
+        lambda glb_bytes: (glb_bytes, ),
+        inputs=[glb_bytes_state],
+        outputs=[download_glb],
     ).then(
         lambda: gr.DownloadButton.update(interactive=True),
         outputs=[download_glb],
@@ -387,10 +397,10 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     extract_glb_high_quality_btn.click(
         extract_glb_high_quality,
         inputs=[output_buf],
-        outputs=[model_output, glb_high_quality_path_state],
+        outputs=[model_output, glb_high_quality_bytes_state],
     ).then(
-        lambda glb_path: glb_path,  # Pass the file path directly
-        inputs=[glb_high_quality_path_state],
+        lambda glb_bytes: (glb_bytes, ),
+        inputs=[glb_high_quality_bytes_state],
         outputs=[download_high_quality_glb],
     ).then(
         lambda: gr.DownloadButton.update(interactive=True),