Update app.py

app.py

@@ -5,9 +5,7 @@ import sys
 import tempfile
 import shutil
 import subprocess
-import spaces
-# from huggingface_hub import HfApi, snapshot_download # For future model management if needed
-# import spaces # For @spaces.GPU decorator if you add it
+import spaces # Ensure spaces is imported if @spaces.GPU is used
 
 # --- Configuration ---
 # Path to the cloned UniRig repository directory within the Space
@@ -25,27 +23,40 @@ if DEVICE.type == 'cuda':
 else:
     print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
 
-@spaces.GPU
-def run_unirig_command(command_args, step_name):
-    """Helper function to run UniRig commands using subprocess."""
-    python_exe = sys.executable
-    # Ensure the command starts with the python executable and '-m' for module execution
-    cmd = [python_exe, "-m"] + command_args
+@spaces.GPU # Decorator for ZeroGPU
+def run_unirig_command(command_list, step_name):
+    """
+    Helper function to run UniRig commands (now expecting bash scripts) using subprocess.
+    command_list: The full command and its arguments, e.g., ["bash", "script.sh", "--arg", "value"]
+    """
+    # The command_list is now expected to be the full command, e.g., starting with "bash"
+    cmd = command_list
     
     print(f"Running {step_name}: {' '.join(cmd)}")
     
     process_env = os.environ.copy()
     
-    # Explicitly add UNIRIG_REPO_DIR to PYTHONPATH for the subprocess.
-    # This ensures that Python can find the 'unirig' package located within UNIRIG_REPO_DIR.
-    # UNIRIG_REPO_DIR itself is the directory containing the 'unirig' package folder.
+    # Determine the path to the 'src' directory within UniRig, where the 'unirig' package resides.
+    unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
+
+    # Explicitly add UNIRIG_REPO_DIR/src to PYTHONPATH for the subprocess.
+    # The bash scripts will internally call Python, which needs to find the 'unirig' package.
+    # Also, keep UNIRIG_REPO_DIR itself in case some scripts or modules there are run directly
+    # or expect the project root to be in PYTHONPATH.
     existing_pythonpath = process_env.get('PYTHONPATH', '')
-    process_env["PYTHONPATH"] = f"{UNIRIG_REPO_DIR}{os.pathsep}{existing_pythonpath}"
+    new_pythonpath_parts = [unirig_src_dir, UNIRIG_REPO_DIR] # UniRig/src first, then UniRig/
+    if existing_pythonpath:
+        # Prepend our paths to existing PYTHONPATH
+        new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
+    
+    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts)) # filter(None,...) handles empty existing_pythonpath
     print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")
 
 
     try:
-        # Execute the command from the UniRig directory for Hydra to find configs
+        # Execute the command from the UniRig directory (UNIRIG_REPO_DIR)
+        # This is crucial for the bash scripts to find their relative paths (e.g., to Python scripts)
+        # and for any underlying Python/Hydra calls to find configurations (e.g., in UniRig/configs/)
         result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
         print(f"{step_name} STDOUT:\n{result.stdout}")
         if result.stderr:
@@ -60,87 +71,92 @@ def run_unirig_command(command_args, step_name):
         error_summary = e.stderr.splitlines()[-5:] # Last 5 lines of stderr
         raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
     except FileNotFoundError:
-        print(f"ERROR: Could not find executable or script for {step_name}. Is UniRig cloned correctly in {UNIRIG_REPO_DIR} and Python environment setup?")
-        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs and UniRig directory structure.")
+        # This error means the executable (e.g., "bash" or the script itself) was not found.
+        print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}. Is UniRig cloned correctly and 'bash' available?")
+        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs, UniRig directory structure, and script paths.")
     except Exception as e_general:
         print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
         raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}")
 
 
-# @spaces.GPU # You can specify type like @spaces.GPU(type="t4") or count
-@spaces.GPU
+@spaces.GPU # Decorator for ZeroGPU
 def rig_glb_mesh_multistep(input_glb_file_obj):
     """
     Takes an input GLB file object (from gr.File with type="filepath"),
-    rigs it using the new UniRig multi-step process,
+    rigs it using the new UniRig multi-step process by calling its bash scripts,
     and returns the path to the final rigged GLB file.
     """
     if not os.path.isdir(UNIRIG_REPO_DIR):
          raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Cannot proceed. Please check Space setup.")
 
     if input_glb_file_obj is None:
-        # This case should ideally be handled by Gradio's input validation if `allow_none=False` (default)
         raise gr.Error("No input file provided. Please upload a .glb mesh.")
 
-    # When type="filepath", input_glb_file_obj is the path string directly
-    input_glb_path = input_glb_file_obj
+    input_glb_path = input_glb_file_obj # This is the absolute path from gr.File(type="filepath")
     print(f"Input GLB path received: {input_glb_path}")
 
     # Create a dedicated temporary directory for all intermediate and final files
+    # The output paths for UniRig scripts will point into this directory.
     processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
     print(f"Using temporary processing directory: {processing_temp_dir}")
 
     try:
         base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
         
-        # Step 1: Skeleton Prediction
-        temp_skeleton_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
+        # Define absolute paths for intermediate files within the processing_temp_dir
+        abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
+        abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
+        abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
+
+        # Step 1: Skeleton Prediction using generate_skeleton.sh
         print("Step 1: Predicting Skeleton...")
-        run_unirig_command([
-            "unirig.predict_skeleton",
-            f"input.path={os.path.abspath(input_glb_path)}", # Use absolute path for robustness
-            f"output.path={os.path.abspath(temp_skeleton_path)}",
-            # f"device={str(DEVICE)}" # If UniRig's script accepts this override and handles it
-        ], "Skeleton Prediction")
-        if not os.path.exists(temp_skeleton_path):
+        skeleton_cmd = [
+            "bash", "launch/inference/generate_skeleton.sh",
+            "--input", input_glb_path, # Input is the original GLB
+            "--output", abs_skeleton_output_path
+        ]
+        run_unirig_command(skeleton_cmd, "Skeleton Prediction")
+        if not os.path.exists(abs_skeleton_output_path):
             raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")
 
-        # Step 2: Skinning Weight Prediction
-        temp_skin_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
+        # Step 2: Skinning Weight Prediction using generate_skin.sh
         print("Step 2: Predicting Skinning Weights...")
-        run_unirig_command([
-            "unirig.predict_skin",
-            f"input.skeleton_path={os.path.abspath(temp_skeleton_path)}",
-            f"input.source_mesh_path={os.path.abspath(input_glb_path)}",
-            f"output.path={os.path.abspath(temp_skin_path)}",
-        ], "Skinning Prediction")
-        if not os.path.exists(temp_skin_path):
+        # generate_skin.sh requires the skeleton from step 1 as --input,
+        # and the original mesh as --source.
+        skin_cmd = [
+            "bash", "launch/inference/generate_skin.sh",
+            "--input", abs_skeleton_output_path,    # Input is the skeleton FBX from previous step
+            "--source", input_glb_path,             # Source is the original GLB mesh
+            "--output", abs_skin_output_path
+        ]
+        run_unirig_command(skin_cmd, "Skinning Prediction")
+        if not os.path.exists(abs_skin_output_path):
             raise gr.Error("Skinning prediction failed to produce an output file. Check logs for UniRig errors.")
 
-        # Step 3: Merge Skeleton/Skin with Original Mesh
-        final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
+        # Step 3: Merge Skeleton/Skin with Original Mesh using merge.sh
         print("Step 3: Merging Results...")
-        run_unirig_command([
-            "unirig.merge_skeleton_skin",
-            f"input.source_rig_path={os.path.abspath(temp_skin_path)}",
-            f"input.target_mesh_path={os.path.abspath(input_glb_path)}",
-            f"output.path={os.path.abspath(final_rigged_glb_path)}",
-        ], "Merging")
-        if not os.path.exists(final_rigged_glb_path):
+        # merge.sh requires the skinned FBX as --source (which contains skeleton and weights)
+        # and the original GLB as --target.
+        merge_cmd = [
+            "bash", "launch/inference/merge.sh",
+            "--source", abs_skin_output_path,      # Source is the skinned FBX from previous step
+            "--target", input_glb_path,            # Target is the original GLB mesh
+            "--output", abs_final_rigged_glb_path
+        ]
+        run_unirig_command(merge_cmd, "Merging")
+        if not os.path.exists(abs_final_rigged_glb_path):
             raise gr.Error("Merging process failed to produce the final rigged GLB file. Check logs for UniRig errors.")
 
-        # final_rigged_glb_path is in processing_temp_dir.
-        # Gradio's gr.Model3D output component will handle serving this file.
-        return final_rigged_glb_path
+        return abs_final_rigged_glb_path
 
-    except gr.Error: # Re-raise Gradio errors directly
-        if os.path.exists(processing_temp_dir): # Clean up on known Gradio error
+    except gr.Error: 
+        if os.path.exists(processing_temp_dir): 
             shutil.rmtree(processing_temp_dir)
             print(f"Cleaned up temporary directory: {processing_temp_dir}")
         raise
     except Exception as e:
         print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
-        if os.path.exists(processing_temp_dir): # Clean up on unexpected error
+        if os.path.exists(processing_temp_dir): 
             shutil.rmtree(processing_temp_dir)
             print(f"Cleaned up temporary directory: {processing_temp_dir}")
         raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")
@@ -154,17 +170,14 @@ theme = gr.themes.Soft(
     font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
 )
 
-# Ensure UNIRIG_REPO_DIR check happens before interface is built if it's critical
-if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__": # Check only if running as main script
+if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__": 
     print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will not work.")
 
-# Define the interface
-# Note: The @spaces.GPU decorator would go above the function `rig_glb_mesh_multistep`
 iface = gr.Interface(
     fn=rig_glb_mesh_multistep, 
     inputs=gr.File(
         label="Upload .glb Mesh File",
-        type="filepath"  # Corrected type for Gradio 5.x
+        type="filepath" 
     ),
     outputs=gr.Model3D(
         label="Rigged 3D Model (.glb)",
@@ -172,7 +185,7 @@ iface = gr.Interface(
     ),
     title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
     description=(
-        "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh.\n"
+        "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh by calling its provided bash scripts.\n"
         "The process involves: 1. Skeleton Prediction, 2. Skinning Weight Prediction, 3. Merging.\n"
         "This may take several minutes. Ensure your GLB has clean geometry.\n"
         f"Running on: {str(DEVICE).upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
@@ -180,12 +193,10 @@ iface = gr.Interface(
     ),
     cache_examples=False,
     theme=theme
-    # allow_flagging="never" # Removed as it's deprecated in Gradio 4.x and default behavior is usually no flagging.
-                           # If specific flagging control is needed, use `flagging_options` or similar.
 )
 
 if __name__ == "__main__":
     if not os.path.isdir(UNIRIG_REPO_DIR):
         print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}. Ensure it's cloned in the Space's root.")
     
-    iface.launch()
+    iface.launch()