Update app.py

app.py

@@ -1,205 +1,185 @@
-import gradio as gr
-import torch
-import os
-import sys
-import tempfile
-import shutil
-import subprocess
-import spaces
-
-# --- Configuration ---
-# Path to the cloned UniRig repository directory within the Space
-# IMPORTANT: You must clone the UniRig repository into this directory in your Hugging Face Space.
-UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
-
-# Check if UniRig directory exists
-if not os.path.isdir(UNIRIG_REPO_DIR):
-    # This message will appear in logs, Gradio app might fail to start fully.
-    print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
-    # Optionally, raise an error to make it more visible if the app starts
-    # raise RuntimeError(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
-
-
-# Determine processing device (CUDA if available, otherwise CPU)
-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}")
-    # Note: UniRig scripts might manage device internally or via Hydra configs.
-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  # Use the current python interpreter
-    cmd = [python_exe, "-m"] + command_args
-    
-    print(f"Running {step_name}: {' '.join(cmd)}")
-    
-    # UniRig scripts often expect to be run from the root of the UniRig repository
-    # because they use Hydra and relative paths for configs (e.g., conf/config.yaml)
-    process_env = os.environ.copy()
-    
-    # Add UniRig's parent directory to PYTHONPATH so `import unirig` works if needed,
-    # and the UniRig directory itself so its internal imports work.
-    # However, `python -m` typically handles package discovery well if CWD is correct.
-    # process_env["PYTHONPATH"] = f"{UNIRIG_REPO_DIR}{os.pathsep}{process_env.get('PYTHONPATH', '')}"
-
-    try:
-        # Execute the command
-        # It's crucial to set `cwd=UNIRIG_REPO_DIR` for Hydra to find its configs.
-        result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
-        print(f"{step_name} output:\n{result.stdout}")
-        if result.stderr:
-            print(f"{step_name} errors (non-fatal if check=True passed):\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}")
-        raise gr.Error(f"Error in UniRig {step_name}: {e.stderr[:500]}") # Show first 500 chars of error
-    except FileNotFoundError:
-        # This can happen if UNIRIG_REPO_DIR is not populated correctly or python_exe is wrong
-        print(f"ERROR: Could not find executable or script for {step_name}. Is UniRig cloned correctly in {UNIRIG_REPO_DIR}?")
-        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs.")
-
-
-# --- Core Rigging Function ---
-@spaces.GPU
-def rig_glb_mesh_multistep(input_glb_file_obj):
-    """
-    Takes an input GLB file object, 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.")
-
-    if input_glb_file_obj is None:
-        raise gr.Error("No input file provided. Please upload a .glb mesh.")
-
-    input_glb_path = input_glb_file_obj.name  # Path to the temporary uploaded file
-
-    # Create a dedicated temporary directory for all intermediate and final files for this run
-    # This helps in organization and cleanup.
-    processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
-    print(f"Using temporary processing directory: {processing_temp_dir}")
-
-    try:
-        # Define paths for intermediate and final files within the processing_temp_dir
-        # UniRig scripts expect output paths.
-        base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
-        
-        # Step 1: Skeleton Prediction
-        # Output is typically an FBX file for the skeleton
-        temp_skeleton_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
-        print("Step 1: Predicting Skeleton...")
-        # Command: python -m unirig.predict_skeleton +input_path=<input_glb_path> +output_path=<temp_skeleton_path>
-        # Note: UniRig's scripts might have default output locations or require specific Hydra overrides.
-        # The `+` syntax is for Hydra overrides.
-        # Check UniRig's `conf/predict_skeleton.yaml` for default config values.
-        run_unirig_command([
-            "unirig.predict_skeleton",
-            f"input.path={input_glb_path}", # Use dot notation for Hydra parameters
-            f"output.path={temp_skeleton_path}",
-            # Add other necessary overrides, e.g., for device if not auto-detected well
-            # f"device={str(DEVICE)}" # If UniRig's script accepts this override
-        ], "Skeleton Prediction")
-        print(f"Skeleton predicted at: {temp_skeleton_path}")
-        if not os.path.exists(temp_skeleton_path):
-            raise gr.Error("Skeleton prediction failed to produce an output file.")
-
-        # Step 2: Skinning Weight Prediction
-        # Input: skeleton FBX and original GLB. Output: skinned FBX (or other format)
-        temp_skin_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
-        print("Step 2: Predicting Skinning Weights...")
-        # Command: python -m unirig.predict_skin +input_path=<temp_skeleton_path> +output_path=<temp_skin_path> +source_mesh_path=<input_glb_path>
-        run_unirig_command([
-            "unirig.predict_skin",
-            f"input.skeleton_path={temp_skeleton_path}", # Check exact Hydra param name in UniRig
-            f"input.source_mesh_path={input_glb_path}", # Check exact Hydra param name
-            f"output.path={temp_skin_path}",
-        ], "Skinning Prediction")
-        print(f"Skinning predicted at: {temp_skin_path}")
-        if not os.path.exists(temp_skin_path):
-            raise gr.Error("Skinning prediction failed to produce an output file.")
-
-        # Step 3: Merge Skeleton/Skin with Original Mesh
-        # Input: original GLB and the skin FBX (which contains skeleton + weights). Output: final rigged GLB
-        final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
-        print("Step 3: Merging Results...")
-        # Command: python -m unirig.merge_skeleton_skin +source_path=<temp_skin_path> +target_path=<input_glb_path> +output_path=<final_rigged_glb_path>
-        run_unirig_command([
-            "unirig.merge_skeleton_skin",
-            f"input.source_rig_path={temp_skin_path}", # Path to the file with skeleton and skin weights
-            f"input.target_mesh_path={input_glb_path}", # Path to the original mesh
-            f"output.path={final_rigged_glb_path}",
-        ], "Merging")
-        print(f"Final rigged mesh at: {final_rigged_glb_path}")
-        if not os.path.exists(final_rigged_glb_path):
-            raise gr.Error("Merging process failed to produce the final rigged GLB file.")
-
-        # The final_rigged_glb_path needs to be accessible by Gradio to serve it.
-        # Gradio usually copies temp files it creates, but here we created it.
-        # We return the path, and Gradio should handle it.
-        # The processing_temp_dir will be cleaned up by Gradio if input_glb_file_obj is from gr.File
-        # or we can clean it up if we copy the final file to a Gradio managed temp location.
-        # For gr.Model3D, returning a path is fine.
-        return final_rigged_glb_path
-
-    except gr.Error: # Re-raise Gradio errors directly
-        raise
-    except Exception as e:
-        print(f"An unexpected error occurred: {e}")
-        # Clean up the processing directory in case of an unhandled error
-        if os.path.exists(processing_temp_dir):
-            shutil.rmtree(processing_temp_dir)
-        raise gr.Error(f"An unexpected error occurred during processing: {str(e)}")
-    # Note: Do not clean up processing_temp_dir in a `finally` block here if returning path from it,
-    # as Gradio needs the file to exist to serve it. Gradio's gr.File output type handles temp file cleanup.
-    # If outputting gr.File, copy the final file to a new tempfile managed by Gradio.
-    # For gr.Model3D, path is fine.
-
-# --- 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"],
-)
-
-iface = gr.Interface(
-    fn=rig_glb_mesh_multistep,
-    inputs=gr.File(label="Upload .glb Mesh File", type="file"),
-    outputs=gr.Model3D(
-        label="Rigged 3D Model (.glb)",
-        clear_color=[0.8, 0.8, 0.8, 1.0],
-        # Note: Model3D might have issues with complex GLBs or certain rigging structures.
-        # A gr.File output for download might be a safer fallback.
-        # outputs=[gr.Model3D(...), gr.File(label="Download Rigged GLB")]
-    ),
-    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: '{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__":
-    # Perform a quick check for UniRig directory on launch
-    if not os.path.isdir(UNIRIG_REPO_DIR):
-        print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will likely fail.")
-        # You could display this error in the Gradio interface itself using a dummy function or Markdown.
-    
-    # For local testing, you might need to set PYTHONPATH or ensure UniRig is installed.
-    # Example: os.environ["PYTHONPATH"] = f"{UNIRIG_REPO_DIR}{os.pathsep}{os.environ.get('PYTHONPATH', '')}"
-    
-    iface.launch()
+   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):
+            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.")
+       
+       iface.launch()
+