app.py

import shutil
import subprocess
import time
from pathlib import Path
from typing import Tuple

import gradio as gr
import lightning as L
import spaces
import torch
import yaml

subprocess.run('pip install flash-attn --no-build-isolation', shell=True)

# Get the PyTorch and CUDA versions
torch_version = torch.__version__.split("+")[0]  # Strips any "+cuXXX" suffix
cuda_version = torch.version.cuda
spconv_version = "-cu121" if cuda_version else "" 

# Format CUDA version to match the URL convention (e.g., "cu118" for CUDA 11.8)
if cuda_version:
    cuda_version = f"cu{cuda_version.replace('.', '')}"
else:
    cuda_version = "cpu"  # Fallback in case CUDA is not available


subprocess.run(f'pip install spconv{spconv_version}', shell=True)
subprocess.run(f'pip install torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-{torch_version}+{cuda_version}.html --no-cache-dir', shell=True)


# Helper functions
def validate_input_file(file_path: str) -> bool:
    """Validate if the input file format is supported."""
    supported_formats = ['.obj', '.fbx', '.glb']
    if not file_path or not Path(file_path).exists():
        return False
    
    file_ext = Path(file_path).suffix.lower()
    return file_ext in supported_formats

def extract_mesh_python(input_file: str, output_dir: str) -> str:
    """
    Extract mesh data from 3D model using Python (replaces extract.sh)
    Returns path to generated .npz file
    """
    # Import required modules
    from src.data.extract import get_files, extract_builtin
    
    # Create extraction parameters
    files = get_files(
        data_name="raw_data.npz",
        inputs=str(input_file),
        input_dataset_dir=None,
        output_dataset_dir=output_dir,
        force_override=True,
        warning=False,
    )
    
    if not files:
        raise RuntimeError("No files to extract")
    
    # Run the actual extraction
    timestamp = str(int(time.time()))
    extract_builtin(
        output_folder=output_dir,
        target_count=50000,
        num_runs=1,
        id=0,
        time=timestamp,
        files=files,
    )
    
    # Return the directory path where raw_data.npz was created
    # The dataset expects to find raw_data.npz in this directory
    expected_npz_dir = files[0][1]  # This is the output directory
    expected_npz_file = Path(expected_npz_dir) / "raw_data.npz"
    
    if not expected_npz_file.exists():
        raise RuntimeError(f"Extraction failed: {expected_npz_file} not found")
    
    return expected_npz_dir  # Return the directory containing raw_data.npz

def run_skeleton_inference_python(input_file: str, output_file: str, seed: int = 12345) -> str:
    """
    Run skeleton inference using Python (replaces skeleton part of generate_skeleton.sh)
    Returns path to skeleton FBX file
    """
    from box import Box

    from src.data.datapath import Datapath
    from src.data.dataset import DatasetConfig, UniRigDatasetModule
    from src.data.transform import TransformConfig
    from src.inference.download import download
    from src.model.parse import get_model
    from src.system.parse import get_system, get_writer
    from src.tokenizer.parse import get_tokenizer
    from src.tokenizer.spec import TokenizerConfig
    
    # Set random seed
    L.seed_everything(seed, workers=True)
    
    # Load task configuration
    task_config_path = "configs/task/quick_inference_skeleton_articulationxl_ar_256.yaml"
    if not Path(task_config_path).exists():
        raise FileNotFoundError(f"Task configuration file not found: {task_config_path}")
    
    # Load the task configuration
    with open(task_config_path, 'r') as f:
        task = Box(yaml.safe_load(f))
    
    # Create temporary npz directory
    npz_dir = Path(output_file).parent / "npz"
    npz_dir.mkdir(exist_ok=True)
    
    # Extract mesh data
    npz_data_dir = extract_mesh_python(input_file, npz_dir)
    
    # Setup datapath with the directory containing raw_data.npz
    datapath = Datapath(files=[npz_data_dir], cls=None)
    
    # Load configurations
    data_config = Box(yaml.safe_load(open("configs/data/quick_inference.yaml", 'r')))
    transform_config = Box(yaml.safe_load(open("configs/transform/inference_ar_transform.yaml", 'r')))
    
    # Get tokenizer
    tokenizer_config = TokenizerConfig.parse(config=Box(yaml.safe_load(open("configs/tokenizer/tokenizer_parts_articulationxl_256.yaml", 'r'))))
    tokenizer = get_tokenizer(config=tokenizer_config)
    
    # Get model
    model_config = Box(yaml.safe_load(open("configs/model/unirig_ar_350m_1024_81920_float32.yaml", 'r')))
    model = get_model(tokenizer=tokenizer, **model_config)
    
    # Setup datasets and transforms
    predict_dataset_config = DatasetConfig.parse(config=data_config.predict_dataset_config).split_by_cls()
    predict_transform_config = TransformConfig.parse(config=transform_config.predict_transform_config)
    
    # Create data module
    data = UniRigDatasetModule(
        process_fn=model._process_fn,
        predict_dataset_config=predict_dataset_config,
        predict_transform_config=predict_transform_config,
        tokenizer_config=tokenizer_config,
        debug=False,
        data_name="raw_data.npz",
        datapath=datapath,
        cls=None,
    )
    
    # Setup callbacks and writer
    callbacks = []
    writer_config = task.writer.copy()
    writer_config['npz_dir'] = str(npz_dir)
    writer_config['output_dir'] = str(Path(output_file).parent)
    writer_config['output_name'] = Path(output_file).name
    writer_config['user_mode'] = False  # Set to False to enable NPZ export
    print(f"Writer config: {writer_config}")
    # But we want the FBX to go to our specified location when in user mode for FBX
    callbacks.append(get_writer(**writer_config, order_config=predict_transform_config.order_config))
    
    # Get system
    system_config = Box(yaml.safe_load(open("configs/system/ar_inference_articulationxl.yaml", 'r')))
    system = get_system(**system_config, model=model, steps_per_epoch=1)
    
    # Setup trainer
    trainer_config = task.trainer
    resume_from_checkpoint = download(task.resume_from_checkpoint)
    
    trainer = L.Trainer(callbacks=callbacks, logger=None, **trainer_config)
    
    # Run prediction
    trainer.predict(system, datamodule=data, ckpt_path=resume_from_checkpoint, return_predictions=False)
    
    # The actual output file will be in a subdirectory named after the input file
    # Look for the generated skeleton.fbx file
    input_name_stem = Path(input_file).stem
    actual_output_dir = Path(output_file).parent / input_name_stem
    actual_output_file = actual_output_dir / "skeleton.fbx"
    
    if not actual_output_file.exists():
        # Try alternative locations - look for any skeleton.fbx file in the output directory
        alt_files = list(Path(output_file).parent.rglob("skeleton.fbx"))
        if alt_files:
            actual_output_file = alt_files[0]
            print(f"Found skeleton at alternative location: {actual_output_file}")
        else:
            # List all files for debugging
            all_files = list(Path(output_file).parent.rglob("*"))
            print(f"Available files: {[str(f) for f in all_files]}")
            raise RuntimeError(f"Skeleton FBX file not found. Expected at: {actual_output_file}")
    
    # Copy to the expected output location
    if actual_output_file != Path(output_file):
        shutil.copy2(actual_output_file, output_file)
        print(f"Copied skeleton from {actual_output_file} to {output_file}")
    
    print(f"Generated skeleton at: {output_file}")
    return str(output_file)

def run_skin_inference_python(skeleton_file: str, output_file: str) -> str:
    """
    Run skin inference using Python (replaces skin part of generate_skin.sh)
    Returns path to skin FBX file
    """
    from box import Box

    from src.data.datapath import Datapath
    from src.data.dataset import DatasetConfig, UniRigDatasetModule
    from src.data.transform import TransformConfig
    from src.inference.download import download
    from src.model.parse import get_model
    from src.system.parse import get_system, get_writer
    
    # Load task configuration
    task_config_path = "configs/task/quick_inference_unirig_skin.yaml"
    with open(task_config_path, 'r') as f:
        task = Box(yaml.safe_load(f))
            
    # Look for files matching predict_skeleton.npz pattern recursively
    skeleton_work_dir = Path(skeleton_file).parent
    all_npz_files = list(skeleton_work_dir.rglob("**/*.npz"))
    
    # Setup datapath - need to pass the directory containing the NPZ file
    skeleton_npz_dir = all_npz_files[0].parent
    datapath = Datapath(files=[str(skeleton_npz_dir)], cls=None)
    
    # Load configurations
    data_config = Box(yaml.safe_load(open("configs/data/quick_inference.yaml", 'r')))
    transform_config = Box(yaml.safe_load(open("configs/transform/inference_skin_transform.yaml", 'r')))
    
    # Get model
    model_config = Box(yaml.safe_load(open("configs/model/unirig_skin.yaml", 'r')))
    model = get_model(tokenizer=None, **model_config)
    
    # Setup datasets and transforms
    predict_dataset_config = DatasetConfig.parse(config=data_config.predict_dataset_config).split_by_cls()
    predict_transform_config = TransformConfig.parse(config=transform_config.predict_transform_config)
    
    # Create data module
    data = UniRigDatasetModule(
        process_fn=model._process_fn,
        predict_dataset_config=predict_dataset_config,
        predict_transform_config=predict_transform_config,
        tokenizer_config=None,
        debug=False,
        data_name="predict_skeleton.npz",
        datapath=datapath,
        cls=None,
    )
    
    # Setup callbacks and writer
    callbacks = []
    writer_config = task.writer.copy()
    writer_config['npz_dir'] = str(skeleton_npz_dir)
    writer_config['output_name'] = str(output_file)
    writer_config['user_mode'] = True
    writer_config['export_fbx'] = True  # Enable FBX export
    callbacks.append(get_writer(**writer_config, order_config=predict_transform_config.order_config))
    
    # Get system
    system_config = Box(yaml.safe_load(open("configs/system/skin.yaml", 'r')))
    system = get_system(**system_config, model=model, steps_per_epoch=1)
    
    # Setup trainer
    trainer_config = task.trainer
    resume_from_checkpoint = download(task.resume_from_checkpoint)
    
    trainer = L.Trainer(callbacks=callbacks, logger=None, **trainer_config)
    
    # Run prediction
    trainer.predict(system, datamodule=data, ckpt_path=resume_from_checkpoint, return_predictions=False)
    
    # The skin FBX file should be generated with the specified output name
    # Since user_mode is True and export_fbx is True, it should create the file directly
    if not Path(output_file).exists():
        # Look for generated skin FBX files in the output directory
        skin_files = list(Path(output_file).parent.rglob("*skin*.fbx"))
        if skin_files:
            actual_output_file = skin_files[0]
            # Copy/move to the expected location
            shutil.copy2(actual_output_file, output_file)
        else:
            raise RuntimeError(f"Skin FBX file not found. Expected at: {output_file}")
    
    return str(output_file)

def merge_results_python(source_file: str, target_file: str, output_file: str) -> str:
    """
    Merge results using Python (replaces merge.sh)
    Returns path to merged file
    """
    from src.inference.merge import transfer
    
    # Validate input paths
    if not Path(source_file).exists():
        raise ValueError(f"Source file does not exist: {source_file}")
    if not Path(target_file).exists():
        raise ValueError(f"Target file does not exist: {target_file}")
    
    # Ensure output directory exists
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Use the transfer function directly
    transfer(source=str(source_file), target=str(target_file), output=str(output_path), add_root=False)
    
    # Validate that the output file was created and is a valid file
    if not output_path.exists():
        raise RuntimeError(f"Merge failed: Output file not created at {output_path}")
    
    if not output_path.is_file():
        raise RuntimeError(f"Merge failed: Output path is not a valid file: {output_path}")
    
    return str(output_path.resolve())

@spaces.GPU()
def complete_pipeline(input_file: str, seed: int = 12345) -> Tuple[str, list]:
    """
    Run the complete rigging pipeline: skeleton generation → skinning → merge.
    
    Args:
        input_file: Path to the input 3D model file
        seed: Random seed for reproducible results
        
    Returns:
        Tuple of (final_file_path, list_of_intermediate_files)
    """
    # Create temp directory
    base_dir = Path(__file__).parent
    temp_dir = base_dir / "tmp"
    temp_dir.mkdir(exist_ok=True)
    
    # Supported file formats
    supported_formats = ['.obj', '.fbx', '.glb']
    
    # Validate input file
    if not validate_input_file(input_file):
        raise gr.Error(f"Error: Invalid or unsupported file format. Supported formats: {', '.join(supported_formats)}")
    
    # Create working directory
    file_stem = Path(input_file).stem
    input_model_dir = temp_dir / f"{file_stem}_{seed}"
    input_model_dir.mkdir(exist_ok=True)

    # Copy input file to working directory
    input_file = Path(input_file)
    shutil.copy2(input_file, input_model_dir / input_file.name)
    input_file = input_model_dir / input_file.name
    print(f"New input file path: {input_file}")
    
    # Step 1: Generate skeleton        
    output_skeleton_file = input_model_dir / f"{file_stem}_skeleton.fbx"
    run_skeleton_inference_python(input_file, output_skeleton_file, seed)        

    # Step 2: Generate skinning
    output_skin_file = input_model_dir / f"{file_stem}_skin.fbx"
    run_skin_inference_python(output_skeleton_file, output_skin_file)
    
    # Step 3: Merge results
    final_file = input_model_dir / f"{file_stem}_rigged.glb"
    merge_results_python(output_skin_file, input_file, final_file)

    return str(final_file), [str(output_skeleton_file), str(output_skin_file), str(final_file)]


def create_app():
    """Create and configure the Gradio interface."""
    
    with gr.Blocks(title="UniRig - 3D Model Rigging Demo") as interface:
        
        # Header
        gr.HTML("""
        <div class="title" style="text-align: center">
            <h1>🎯 UniRig: Automated 3D Model Rigging</h1>
            <p style="font-size: 1.1em; color: #6b7280;">
                Leverage deep learning to automatically generate skeletons and skinning weights for your 3D models
            </p>
        </div>
        """)
        
        # Usage Instructions Section
        gr.Markdown("""
        ## 📋 How to Use ?
        1. **Upload your 3D model** - Drop your .obj, .fbx, or .glb file in the upload area
        2. **Set random seed** (optional) - Use the same seed for reproducible results
        3. **Click "Start Complete Pipeline"** - The AI will automatically rig your model
        4. **Download results** - `_skeleton.fbx` is the base model with skeleton, `_skin.fbx` is the base model with armature/skeleton and skinning weights, and `_rigged.*` is the final rigged model ready for use.

        **Supported File Formats:** .obj, .fbx, .glb
        **Note:** The process may take a few minutes depending on the model complexity and server load.
        """)
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                input_3d_model = gr.Model3D(label="Upload 3D Model")
                
                with gr.Row(equal_height=True):
                    seed = gr.Number(
                        value=12345,
                        label="Random Seed (for reproducible results)",
                        scale=4,
                    )
                    random_btn = gr.Button("🔄 Random Seed", variant="secondary", scale=1)
                
                pipeline_btn = gr.Button("🎯 Start Complete Pipeline", variant="primary", size="lg")
            
            with gr.Column():
                pipeline_skeleton_out = gr.Model3D(label="Final Rigged Model", scale=4)
                files_to_download = gr.Files(label="Download Files", scale=1)
                   
        random_btn.click(
            fn=lambda: int(torch.randint(0, 100000, (1,)).item()),
            outputs=seed
        )
        
        pipeline_btn.click(
            fn=complete_pipeline,
            inputs=[input_3d_model, seed],
            outputs=[pipeline_skeleton_out, files_to_download]
        )
        
        # Footer
        gr.HTML("""
        <div style="text-align: center; margin-top: 2em; padding: 1em; border-radius: 8px;">
            <p style="color: #6b7280;">
                🔬 <strong>UniRig</strong> - Research by Tsinghua University & Tripo<br>
                📄 <a href="https://arxiv.org/abs/2504.12451" target="_blank">Paper</a> | 
                🏠 <a href="https://zjp-shadow.github.io/works/UniRig/" target="_blank">Project Page</a> | 
                🤗 <a href="https://huggingface.co/VAST-AI/UniRig" target="_blank">Models</a>
            </p>
        </div>
        """)
    
    return interface

if __name__ == "__main__":
    # Create and launch the interface
    app = create_app()
    
    # Launch configuration
    app.queue().launch()