clean up

app.py

@@ -235,6 +235,15 @@ def run_pipeline(image_array, state):
     final_status = "\n".join([status1, status2, status3, status4])
     return final_status, crop_img, normals_img, uv_img, track_img, mesh_file, state
 
+# Cleanup on unload
+def cleanup(state):
+    sid = state.get("session_id")
+    if sid:
+        d1 = os.path.join(env["PIXEL3DMM_PREPROCESSED_DATA"], sid)
+        d2 = os.path.join(env["PIXEL3DMM_TRACKING_OUTPUT"], sid)
+        shutil.rmtree(d1, ignore_errors=True)
+        shutil.rmtree(d2, ignore_errors=True)
+
 # Build Gradio UI
 demo = gr.Blocks()
 with demo:
@@ -263,6 +272,7 @@ with demo:
 
     image_in.upload(fn=reset_all, inputs=None, outputs=[crop_img, normals_img, uv_img, track_img, mesh_file, status, state, run_btn])
 
+demo.unload(fn=cleanup, inputs=[state], outputs=[])
 
 demo.queue()
 demo.launch(share=True, ssr_mode=False)