#!/bin/bash
 
# Set up variables
TITLE="ComfyUI + Model Installer"
COMFYUI_URL_PART1=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v8.0/ComfyUI_Windows_portable_cu124.7z.001
COMFYUI_URL_PART2=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v8.0/ComfyUI_Windows_portable_cu124.7z.002
MODELS_URL=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v8.0/models.zip.001
HUNYUAN_SAFETENSORS_URL=https://huggingface.co/Kijai/Hunyuan3D-2_safetensors/resolve/main/hunyuan3d-dit-v2-0-fp16.safetensors
HUNYUAN_CKPT_URL=https://huggingface.co/tencent/Hunyuan3D-2/raw/main/hunyuan3d-dit-v2-0/model.ckpt
 
# Set installation paths
INSTALL_DIR="$HOME/ComfyUI_Windows_portable"
CUSTOM_NODES_DIR="$INSTALL_DIR/ComfyUI/custom_nodes"
MODELS_DIR="$INSTALL_DIR/ComfyUI/models"
CHECKPOINTS_DIR="$MODELS_DIR/checkpoints"
PYTHON_DIR="$INSTALL_DIR/python_standalone"
 
echo "$TITLE"
echo "Installing ComfyUI and Models..."
 
# Check system RAM
RAM_GB=$(awk '/MemTotal/{printf "%.0f\n", $2/1024/1024}' /proc/meminfo)
echo "System RAM: ${RAM_GB} GB"
 
# Check GPU
GPU_NAME=$(lspci | grep -i nvidia)
if [ -z "$GPU_NAME" ]; then
    echo "ERROR: No NVIDIA GPU detected."
    exit 1
fi
echo "Detected GPU: $GPU_NAME"
 
# Create directories
mkdir -p "$INSTALL_DIR"
mkdir -p "$MODELS_DIR"
mkdir -p "$CHECKPOINTS_DIR"
 
echo
echo "[Step 1] Downloading ComfyUI (Split Archive)..."
 
# Check for existing ComfyUI files in temp
if [ -f "/tmp/ComfyUI_Windows_portable_cu124.7z.001" ]; then
    echo "Found existing ComfyUI part 1 download, skipping download..."
else
    # Download ComfyUI Part 1
    echo "Downloading ComfyUI part 1..."
    curl -L --retry 3 --retry-delay 2 \
        -o "/tmp/ComfyUI_Windows_portable_cu124.7z.001" \
        "$COMFYUI_URL_PART1"
fi
 
if [ -f "/tmp/ComfyUI_Windows_portable_cu124.7z.002" ]; then
    echo "Found existing ComfyUI part 2 download, skipping download..."
else
    # Download ComfyUI Part 2
    echo "Downloading ComfyUI part 2..."
    curl -L --retry 3 --retry-delay 2 \
        -o "/tmp/ComfyUI_Windows_portable_cu124.7z.002" \
        "$COMFYUI_URL_PART2"
fi
 
echo
echo "[Step 2] Installing p7zip for extraction..."
if ! command -v 7z &> /dev/null; then
    sudo apt-get update && sudo apt-get install -y p7zip-full
else
    echo "p7zip already installed."
fi
 
echo "[Step 3] Extracting ComfyUI..."
 
# First, delete existing directory if it exists
if [ -d "$INSTALL_DIR" ]; then
    echo "Removing existing installation..."
    rm -rf "$INSTALL_DIR"
fi
mkdir "$INSTALL_DIR"
 
# Change to temp directory for extraction
cd "/tmp"
 
echo "Extracting to: $INSTALL_DIR"
7z x "/tmp/ComfyUI_Windows_portable_cu124.7z.001" -o"$INSTALL_DIR" -y
 
# Verify the extraction and adjust directories
if [ ! -f "$PYTHON_DIR/python" ]; then
    echo "ERROR: Python not found in expected location"
    echo "Available files in installation directory:"
    ls "$INSTALL_DIR"
    exit 1
fi
 
echo
echo "[Step 4] Downloading Hunyuan3D Models..."
# Ensure the diffusion_models directory exists
mkdir -p "$INSTALL_DIR/ComfyUI/models/diffusion_models"
 
# Download safetensors model
if [ ! -f "$INSTALL_DIR/ComfyUI/models/diffusion_models/hunyuan3d-dit-v2-0-fp16.safetensors" ]; then
    echo "Downloading Hunyuan3D safetensors model..."
    curl -L -o "$INSTALL_DIR/ComfyUI/models/diffusion_models/hunyuan3d-dit-v2-0-fp16.safetensors" \
    "$HUNYUAN_SAFETENSORS_URL"
 
    if [ $? -ne 0 ]; then
        echo "Failed to download Hunyuan3D safetensors model."
        echo "Please download manually from:"
        echo "https://huggingface.co/Kijai/Hunyuan3D-2_safetensors/blob/main/hunyuan3d-dit-v2-0-fp16.safetensors"
        echo "And place it in: $INSTALL_DIR/ComfyUI/models/diffusion_models"
    else
        echo "Successfully downloaded Hunyuan3D safetensors model."
    fi
else
    echo "Hunyuan3D safetensors model already exists. Skipping download."
fi
 
# Download .ckpt model
if [ ! -f "$INSTALL_DIR/ComfyUI/models/diffusion_models/hunyuan3d-dit-v2-0.ckpt" ]; then
    echo "Downloading Hunyuan3D .ckpt model..."
    curl -L -o "$INSTALL_DIR/ComfyUI/models/diffusion_models/hunyuan3d-dit-v2-0.ckpt" \
    "$HUNYUAN_CKPT_URL"
 
    if [ $? -ne 0 ]; then
        echo "Failed to download Hunyuan3D .ckpt model."
        echo "Please download manually from:"
        echo "https://huggingface.co/tencent/Hunyuan3D-2/raw/main/hunyuan3d-dit-v2-0/model.ckpt"
        echo "And place it in: $INSTALL_DIR/ComfyUI/models/diffusion_models"
    else
        echo "Successfully downloaded Hunyuan3D .ckpt model."
    fi
else
    echo "Hunyuan3D .ckpt model already exists. Skipping download."
fi
 
echo
echo "[Step 5] Downloading Base Models..."
if [ ! -d "$CHECKPOINTS_DIR" ]; then
    mkdir "$CHECKPOINTS_DIR"
fi
 
curl -L --retry 3 --retry-delay 2 -o "$CHECKPOINTS_DIR/models.zip.001" "$MODELS_URL"
 
echo.
echo "[Step 6] Setting up Python Environment..."
cd "$INSTALL_DIR"
echo "Initializing Python..."
"$PYTHON_DIR/python3" -m ensurepip
if [ $? -ne 0 ]; then
    echo "Failed to initialize pip"
    exit 1
fi
"$PYTHON_DIR/python3" -m pip install --upgrade pip
 
echo.
echo "[Step 7] Creating Custom Nodes Directory..."
mkdir "$CUSTOM_NODES_DIR" 2>/dev/null
 
echo.
echo "[Step 8] Installing and Initializing Custom Nodes..."
 
# Define an array of custom node repositories
CUSTOM_NODES=(
    "https://github.com/Kijai/ComfyUI-Hunyuan3DWrapper.git"
    "https://github.com/WASasquatch/was-node-suite-comfyui.git"
    "https://github.com/rgthree/rgthree-comfy.git"
    "https://github.com/Mikubill/sd-webui-controlnet.git"
    "https://github.com/ltdrdata/ComfyUI-Manager.git"
    "https://github.com/storyicon/comfyui_controlnet_aux.git"
    "https://github.com/MrForExample/ComfyUI-Inspire-Pack.git"
    "https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved.git"
    "https://github.com/ZHO-ZHO-ZHO/ComfyUI-Image-Selector.git"
    "https://github.com/pythongosssss/ComfyUI-Custom-Scripts.git"
)
 
# Ensure custom nodes directory exists
if [ ! -d "$CUSTOM_NODES_DIR" ]; then
    mkdir "$CUSTOM_NODES_DIR"
fi
 
cd "$CUSTOM_NODES_DIR"
 
# Clone or update repositories
for repo in "${CUSTOM_NODES[@]}"; do
    repo_name=$(basename "$repo" .git)
    if [ ! -d "$CUSTOM_NODES_DIR/$repo_name" ]; then
        echo "Cloning: $repo_name"
        git clone "$repo" "$CUSTOM_NODES_DIR/$repo_name"
    else
        echo "Updating existing repository: $repo_name"
        cd "$repo_name"
        git fetch origin
        git reset --hard origin/main
        cd ..
    fi
done
 
echo.
echo "[Step 8a] Initializing Git for Existing Nodes..."
# Reinitialize Git for any existing nodes that might have been extracted without Git metadata
for d in "$CUSTOM_NODES_DIR"/*/; do
    if [ -d "$d.git" ]; then
        echo "Repository already initialized: $(basename "$d")"
    else
        echo "Initializing Git for: $(basename "$d")"
        cd "$(basename "$d")"
        git init
        git remote add origin "https://github.com/$(basename "$d")/$(basename "$d").git"
        cd ..
    fi
done
 
echo.
echo "[Step 8b] Installing Hunyuan3D Wrapper Dependencies..."
if [ -d "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper" ]; then
    echo "Installing Hunyuan3D dependencies..."
    "$PYTHON_DIR/python3" -m pip install -r "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper/requirements.txt"
 
    # Install precompiled rasterizer
    echo "Installing precompiled rasterizer..."
    "$PYTHON_DIR/python3" -m pip install "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper/wheels/custom_rasterizer-0.1-cp312-cp312-linux_x86_64.whl"
fi
 
echo.
echo "[Step 9] Compiling Components..."
if [ -d "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper/hy3dgen/texgen/custom_rasterizer" ]; then
    cd "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper/hy3dgen/texgen/custom_rasterizer"
    "$PYTHON_DIR/python3" setup.py install
fi
 
if [ -d "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper/hy3dgen/texgen/differentiable_renderer" ]; then
    cd "$CUSTOM_NODES_DIR/ComfyUI-Hunyuan3DWrapper/hy3dgen/texgen/differentiable_renderer"
    "$PYTHON_DIR/python3" setup.py build_ext --inplace
fi
 
# Cleanup
rm -f "$TEMP/ComfyUI_Windows_portable_cu124.7z.*" 2>/dev/null
rm -f "$TEMP/models.zip.*" 2>/dev/null
 
echo.
echo "[Step 10] Creating Desktop Shortcut..."
(
echo #!/bin/bash
echo "cd $INSTALL_DIR"
echo "./run_nvidia_gpu.sh"
) > "$HOME/Desktop/Start_ComfyUI.sh"
 
chmod +x "$HOME/Desktop/Start_ComfyUI.sh"
 
echo.
echo Installation Complete!
echo.
echo Important Directories:
echo - ComfyUI Installation: $INSTALL_DIR
echo - Models Location: $CHECKPOINTS_DIR
echo.
echo You can launch ComfyUI using the Start_ComfyUI.sh shortcut on your desktop.
 
# End of script
 
### Notes:
1. **7-Zip**: Since 7-Zip is not commonly used in Linux, we will use `unzip` or another tool like `p7zip` to handle `.7z` files.
2. **Python Wheel**: The wheel file name for Linux was changed from `win_amd64` to `linux_x86_64`.
3. **Desktop Shortcut**: A Bash script is created instead of a Windows batch file.
4. **System RAM and GPU Checks**: These checks are omitted as they are not directly translatable without additional tools or commands specific to Linux.
 
Ensure you have the necessary permissions and dependencies installed before running the script, such as `unzip`, `p7zip-full`, and any other required packages for the custom nodes and ComfyUI.