Update app.py

app.py

@@ -1,185 +1,184 @@
-   import gradio as gr
-   import torch
-   import os
-   import sys
-   import tempfile
-   import shutil
-   import subprocess
-   # from huggingface_hub import HfApi, snapshot_download # For future model management if needed
-   # import spaces # For @spaces.GPU decorator if you add it
-
-   # --- Configuration ---
-   # Path to the cloned UniRig repository directory within the Space
-   UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
-
+import gradio as gr
+import torch
+import os
+import sys
+import tempfile
+import shutil
+import subprocess
+# from huggingface_hub import HfApi, snapshot_download # For future model management if needed
+# import spaces # For @spaces.GPU decorator if you add it
+
+# --- Configuration ---
+# Path to the cloned UniRig repository directory within the Space
+UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
+
+if not os.path.isdir(UNIRIG_REPO_DIR):
+   print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
+   # Consider raising an error or displaying it in the UI if UniRig is critical for startup
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {DEVICE}")
+if DEVICE.type == 'cuda':
+   print(f"CUDA Device Name: {torch.cuda.get_device_name(0)}")
+   print(f"CUDA Version: {torch.version.cuda}")
+else:
+   print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
+
+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
+   
+   print(f"Running {step_name}: {' '.join(cmd)}")
+   
+   process_env = os.environ.copy()
+   # It's generally better to ensure UniRig's internal scripts handle PYTHONPATH if needed,
+   # or that it's installed in a way that `python -m` works correctly from its root.
+   # Setting cwd=UNIRIG_REPO_DIR is often the key for Hydra.
+
+   try:
+       # Execute the command from the UniRig directory for Hydra to find 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:
+           print(f"{step_name} STDERR (non-fatal or warnings):\n{result.stderr}")
+   except subprocess.CalledProcessError as e:
+       print(f"ERROR during {step_name}:")
+       print(f"Command: {' '.join(e.cmd)}")
+       print(f"Return code: {e.returncode}")
+       print(f"Stdout: {e.stdout}")
+       print(f"Stderr: {e.stderr}")
+       # Provide a more user-friendly error, potentially masking long tracebacks
+       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.")
+   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]}")
+
+
+# If you are using @spaces.GPU, you would import it:
+# import spaces
+# @spaces.GPU # You can specify type like @spaces.GPU(type="t4") or count
+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,
+   and returns the path to the final rigged GLB file.
+   """
    if not os.path.isdir(UNIRIG_REPO_DIR):
-       print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
-       # Consider raising an error or displaying it in the UI if UniRig is critical for startup
-
-   DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-   print(f"Using device: {DEVICE}")
-   if DEVICE.type == 'cuda':
-       print(f"CUDA Device Name: {torch.cuda.get_device_name(0)}")
-       print(f"CUDA Version: {torch.version.cuda}")
-   else:
-       print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
-
-   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
-       
-       print(f"Running {step_name}: {' '.join(cmd)}")
-       
-       process_env = os.environ.copy()
-       # It's generally better to ensure UniRig's internal scripts handle PYTHONPATH if needed,
-       # or that it's installed in a way that `python -m` works correctly from its root.
-       # Setting cwd=UNIRIG_REPO_DIR is often the key for Hydra.
-
-       try:
-           # Execute the command from the UniRig directory for Hydra to find 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:
-               print(f"{step_name} STDERR (non-fatal or warnings):\n{result.stderr}")
-       except subprocess.CalledProcessError as e:
-           print(f"ERROR during {step_name}:")
-           print(f"Command: {' '.join(e.cmd)}")
-           print(f"Return code: {e.returncode}")
-           print(f"Stdout: {e.stdout}")
-           print(f"Stderr: {e.stderr}")
-           # Provide a more user-friendly error, potentially masking long tracebacks
-           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.")
-       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]}")
-
-
-   # If you are using @spaces.GPU, you would import it:
-   # import spaces
-   # @spaces.GPU # You can specify type like @spaces.GPU(type="t4") or count
-   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,
-       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
-       print(f"Input GLB path received: {input_glb_path}")
-
-       # Create a dedicated temporary directory for all intermediate and final files
-       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")
-           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):
-               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")
-           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):
-               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")
-           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):
-               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
-
-       except gr.Error: # Re-raise Gradio errors directly
-           if os.path.exists(processing_temp_dir): # Clean up on known Gradio error
-               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
-               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]}")
-
-
-   # --- Gradio Interface ---
-   theme = gr.themes.Soft(
-       primary_hue=gr.themes.colors.sky,
-       secondary_hue=gr.themes.colors.blue,
-       neutral_hue=gr.themes.colors.slate,
-       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
-       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.x
-       ),
-       outputs=gr.Model3D(
-           label="Rigged 3D Model (.glb)",
-           clear_color=[0.8, 0.8, 0.8, 1.0],
-       ),
-       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"
-           "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"
-           f"UniRig Source: https://github.com/VAST-AI-Research/UniRig"
-       ),
-       cache_examples=False,
-       theme=theme,
-       allow_flagging="never"
-   )
-
-   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.")
+        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
+   print(f"Input GLB path received: {input_glb_path}")
+
+   # Create a dedicated temporary directory for all intermediate and final files
+   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]
        
-       iface.launch()
-   
+       # Step 1: Skeleton Prediction
+       temp_skeleton_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
+       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):
+           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")
+       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):
+           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")
+       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):
+           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
+
+   except gr.Error: # Re-raise Gradio errors directly
+       if os.path.exists(processing_temp_dir): # Clean up on known Gradio error
+           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
+           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]}")
+
+
+# --- Gradio Interface ---
+theme = gr.themes.Soft(
+   primary_hue=gr.themes.colors.sky,
+   secondary_hue=gr.themes.colors.blue,
+   neutral_hue=gr.themes.colors.slate,
+   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
+   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.x
+   ),
+   outputs=gr.Model3D(
+       label="Rigged 3D Model (.glb)",
+       clear_color=[0.8, 0.8, 0.8, 1.0],
+   ),
+   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"
+       "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"
+       f"UniRig Source: https://github.com/VAST-AI-Research/UniRig"
+   ),
+   cache_examples=False,
+   theme=theme,
+   allow_flagging="never"
+)
+
+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()