#!/bin/bash
set -e

echo "Starting RunPod Stable-Makeup Serverless"

# Check if workspace volume is properly mounted
if [ ! -d "/workspace" ]; then
    echo "ERROR: /workspace volume is not mounted. Please create a RunPod with a Network Volume mounted at /workspace"
    exit 1
fi

# Check available disk space and display
df -h /workspace
df -h /

# Make sure our directories exist
mkdir -p /workspace/app /workspace/cache /workspace/models /workspace/models/stablemakeup

# Clone the repository with LFS since we have sufficient network storage (40GB)
if [ ! -d "/workspace/app/stable-makeup" ]; then
    echo "Cloning repository with LFS from Hugging Face..."
    
    # Setup git-lfs
    git lfs install
    
    # Clone the repository with LFS
    git clone https://huggingface.co/spaces/edgarhnd/Stable-Makeup-unofficial /workspace/app/stable-makeup
    
    # Copy our handler file
    cp /rp_handler.py /workspace/app/stable-makeup/
    
    # Modify the inference_utils.py file to use the workspace volume for checkpoints
    sed -i 's|"./checkpoints/stablemakeup"|"/workspace/models/stablemakeup"|g' /workspace/app/stable-makeup/inference_utils.py
    
    # Install required dependencies
    cd /workspace/app/stable-makeup
    echo "Installing dependencies..."
    pip install --no-cache-dir -r requirements.txt
    pip install --no-cache-dir opencv-python-headless
    
    # Create checkpoints directory and copy model files if they exist in the repo
    mkdir -p /workspace/models/stablemakeup
    if [ -d "/workspace/app/stable-makeup/checkpoints/stablemakeup" ]; then
        echo "Copying model files from cloned repository to workspace models directory..."
        cp -v /workspace/app/stable-makeup/checkpoints/stablemakeup/* /workspace/models/stablemakeup/
    fi
fi

# If model files are still missing, download them directly
if [ ! -f "/workspace/models/stablemakeup/pytorch_model.bin" ]; then
    echo "Model files not found in cloned repository. Downloading directly..."
    
    echo "Downloading pytorch_model.bin..."
    curl -L -o /workspace/models/stablemakeup/pytorch_model.bin https://huggingface.co/spaces/edgarhnd/Stable-Makeup-unofficial/resolve/main/checkpoints/stablemakeup/pytorch_model.bin
    
    echo "Downloading pytorch_model_1.bin..."
    curl -L -o /workspace/models/stablemakeup/pytorch_model_1.bin https://huggingface.co/spaces/edgarhnd/Stable-Makeup-unofficial/resolve/main/checkpoints/stablemakeup/pytorch_model_1.bin
    
    echo "Downloading pytorch_model_2.bin..."
    curl -L -o /workspace/models/stablemakeup/pytorch_model_2.bin https://huggingface.co/spaces/edgarhnd/Stable-Makeup-unofficial/resolve/main/checkpoints/stablemakeup/pytorch_model_2.bin
fi

# Start the handler
echo "Starting serverless handler..."
cd /workspace/app/stable-makeup
exec python -u rp_handler.py