Update app.py

app.py

@@ -169,8 +169,6 @@ def run_triposg(image_path: str,
         if mesh is None:
             outputs[i] = trimesh.Trimesh(vertices=[[0,0,0]], faces=[[0,0,0]])
 
-    # Merge and color
-    merged = get_colored_mesh_composition(outputs)
 
     export_dir = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], session_id)
     os.makedirs(export_dir, exist_ok=True)
@@ -181,13 +179,23 @@ def run_triposg(image_path: str,
         part = os.path.join(export_dir, f"part_{idx:02}.glb")
         mesh.export(part)
         parts.append(part)
+        
+    zip_path = os.path.join(os.environ["PARTCRAFTER_PROCESSED"], f"{session_id}.zip")
+    
+    # shutil.make_archive wants the base name without extension:
+    base_name = zip_path[:-4]  # strip off '.zip'
+    shutil.make_archive(base_name, 'zip', export_dir)
+    
+    # Merge and color
+    merged = get_colored_mesh_composition(outputs)
+
 
     glb_path = os.path.join(export_dir, "object.glb")
     merged.export(glb_path)
 
     mesh_file = first_file_from_dir(export_dir, "glb")
     
-    return mesh_file, export_dir
+    return mesh_file, export_dir, zip_path
 
 def cleanup(request: gr.Request):
 
@@ -254,6 +262,7 @@ def build_demo():
                     )
                     output_model = gr.Model3D(label="Merged 3D Object")
                     output_dir = gr.Textbox(label="Export Directory", visible=False)
+                    download_zip = gr.File(label="Download All Parts (zip)")
                     examples = gr.Examples(
                         
                         examples=[
@@ -262,8 +271,8 @@ def build_demo():
                                 5,
                             ], 
                             [
-                                "assets/images/np5_bird.png", 
-                                5,
+                                "assets/images/np7_1c004909dedb4ebe8db69b4d7b077434.png", 
+                                7,
                             ], 
                             [
                                 "assets/images/np4_7bd5d25aa77b4fb18e780d7a4c97d342.png", 
@@ -271,7 +280,7 @@ def build_demo():
                             ], 
                         ],
                         inputs=[input_image, num_parts],
-                        outputs=[output_model, output_dir],
+                        outputs=[output_model, output_dir, download_zip],
                         fn=run_triposg,
                         cache_examples=True,
                     )
@@ -279,7 +288,7 @@ def build_demo():
             run_button.click(fn=run_triposg,
                              inputs=[input_image, num_parts, seed, num_tokens, num_steps,
                                      guidance, flash_decoder, remove_bg, session_state],
-                             outputs=[output_model, output_dir])
+                             outputs=[output_model, output_dir, download_zip])
         return demo
 
 if __name__ == "__main__":