Update app.py

app.py

@@ -102,15 +102,13 @@ def preprocess_image(image_array, state):
     if image_array is None:
         return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=True)
 
-    session_id = request.session_hash
     base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
     os.makedirs(base_dir, exist_ok=True)
-    state.update({"session_id": session_id, "base_dir": base_dir})
+    state.update({"base_dir": base_dir})
 
     img = Image.fromarray(image_array)
     saved_image_path = os.path.join(base_dir, f"{session_id}.png")
     img.save(saved_image_path)
-    state["image_path"] = saved_image_path
     
     try:
         p = subprocess.run([
@@ -218,7 +216,11 @@ def step4_track(state):
 
 # New: run all steps sequentially
 @spaces.GPU()
-def run_pipeline(image_array, state):
+def run_pipeline(image_array, state, request: gr.Request):
+
+    session_id = request.session_hash
+    state.update({"session_id": session_id, "base_dir": base_dir})
+
     # Step 1
     status1, crop_img, state, _, _ = preprocess_image(image_array, state)
     if "❌" in status1:
@@ -245,9 +247,15 @@ def cleanup(request: gr.Request):
         shutil.rmtree(d1, ignore_errors=True)
         shutil.rmtree(d2, ignore_errors=True)
 
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1024px;
+}
+"""
+
 # Build Gradio UI
-demo = gr.Blocks()
-with demo:
+with gr.Blocks(css=css) as demo:
     gr.Markdown("## Pixel3dmm [Image Mode]")
     gr.Markdown("Versatile Screen-Space Priors for Single-Image 3D Face Reconstruction.")
     gr.HTML(
@@ -260,38 +268,40 @@ with demo:
         </div>
         """
     )
+
+    with gr.Column(elem_id="col-container"):
+        
+        with gr.Row():
+            with gr.Column():
+                image_in = gr.Image(label="Upload Image", type="numpy", height=512)
+                status = gr.Textbox(label="Status", lines=6, interactive=True, value="Upload an image to start.")
+                state = gr.State({})
+            with gr.Column():
+                with gr.Row():
+                    crop_img = gr.Image(label="Preprocessed", height=256)
+                    normals_img = gr.Image(label="Normals", height=256)
+                with gr.Row():
+                    uv_img = gr.Image(label="UV Map", height=256)
+                    track_img = gr.Image(label="Tracking", height=256)
+                mesh_file = gr.Model3D(label="3D Model Preview")
     
-    with gr.Row():
-        with gr.Column():
-            image_in = gr.Image(label="Upload Image", type="numpy", height=512)
-            status = gr.Textbox(label="Status", lines=6, interactive=True, value="Upload an image to start.")
-            state = gr.State({})
-        with gr.Column():
-            with gr.Row():
-                crop_img = gr.Image(label="Preprocessed", height=256)
-                normals_img = gr.Image(label="Normals", height=256)
-            with gr.Row():
-                uv_img = gr.Image(label="UV Map", height=256)
-                track_img = gr.Image(label="Tracking", height=256)
-            mesh_file = gr.Model3D(label="3D Model Preview")
-
-            examples = gr.Examples(
-                examples=["example_images/jim_carrey.png", "example_images/margaret_qualley.png"],
-                inputs=image_in,
-                outputs=[],  # outputs are specified via the callback
-                fn=run_pipeline,
-                cache_examples=True  # cache results of examples for faster loading
-            )
-
-    run_btn = gr.Button("Run Pipeline")
-
-    run_btn.click(
-        fn=run_pipeline,
-        inputs=[image_in, state],
-        outputs=[status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
-    )
-    examples.outputs = [status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
-    image_in.upload(fn=reset_all, inputs=None, outputs=[crop_img, normals_img, uv_img, track_img, mesh_file, status, state, run_btn])
+                examples = gr.Examples(
+                    examples=["example_images/jim_carrey.png", "example_images/margaret_qualley.png"],
+                    inputs=image_in,
+                    outputs=[],  # outputs are specified via the callback
+                    fn=run_pipeline,
+                    cache_examples=True  # cache results of examples for faster loading
+                )
+    
+        run_btn = gr.Button("Run Pipeline")
+    
+        run_btn.click(
+            fn=run_pipeline,
+            inputs=[image_in, state, Request],
+            outputs=[status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
+        )
+        examples.outputs = [status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
+        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(cleanup)