@echo off
setlocal EnableDelayedExpansion

title ComfyUI + Model Installer
echo Installing ComfyUI and Models...

:: Set download URLs
set "COMFYUI_URL_PART1=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v8.0/ComfyUI_Windows_portable_cu124.7z.001"
set "COMFYUI_URL_PART2=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v8.0/ComfyUI_Windows_portable_cu124.7z.002"

rem set "COMFYUI_URL_PART1=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v7.2/ComfyUI_Windows_portable_cu124.7z.001"
rem set "COMFYUI_URL_PART2=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v7.2/ComfyUI_Windows_portable_cu124.7z.002"
rem set "COMFYUI_URL_PART3=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v7.2/ComfyUI_Windows_portable_cu124.7z.003"
set "MODELS_URL=https://github.com/YanWenKun/ComfyUI-Windows-Portable/releases/download/v8.0/models.zip.001"
set "SEVENZIP_URL=https://www.7-zip.org/a/7z2301-x64.exe"
set "HUNYUAN_SAFETENSORS_URL=https://huggingface.co/Kijai/Hunyuan3D-2_safetensors/resolve/main/hunyuan3d-dit-v2-0-fp16.safetensors"
set "HUNYUAN_CKPT_URL=https://huggingface.co/tencent/Hunyuan3D-2/raw/main/hunyuan3d-dit-v2-0/model.ckpt"

:: Set installation paths
set "INSTALL_DIR=%USERPROFILE%\ComfyUI_Windows_portable"
set "CUSTOM_NODES_DIR=%INSTALL_DIR%\ComfyUI\custom_nodes"
set "MODELS_DIR=%INSTALL_DIR%\ComfyUI\models"
set "CHECKPOINTS_DIR=%MODELS_DIR%\checkpoints"
set "PYTHON_DIR=%INSTALL_DIR%\python_standalone"

:: Check system RAM
for /f "tokens=*" %%a in ('powershell -command "$ram = (Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory/1GB; Write-Host $ram"') do set "RAM_GB=%%a"
echo System RAM: %RAM_GB% GB

:: Check GPU
echo Checking NVIDIA GPU...
for /f "tokens=*" %%a in ('powershell -command "$gpu = Get-WmiObject Win32_VideoController | Where-Object { $_.Name -like '*NVIDIA*' }; if($gpu) { Write-Host $gpu.Name }"') do set "GPU_NAME=%%a"
if not defined GPU_NAME (
    echo ERROR: No NVIDIA GPU detected.
    pause
    exit /b 1
)
echo Detected GPU: %GPU_NAME%

:: Create directories
echo Creating directories...
mkdir "%INSTALL_DIR%" 2>nul
mkdir "%MODELS_DIR%" 2>nul
mkdir "%CHECKPOINTS_DIR%" 2>nul

echo.
echo [Step 1] Downloading ComfyUI (Split Archive)...

:: Check for existing ComfyUI files in temp
if exist "%TEMP%\ComfyUI_Windows_portable_cu124.7z.001" (
    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 "%TEMP%\ComfyUI_Windows_portable_cu124.7z.001" ^
        "%COMFYUI_URL_PART1%"
)

if exist "%TEMP%\ComfyUI_Windows_portable_cu124.7z.002" (
    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 "%TEMP%\ComfyUI_Windows_portable_cu124.7z.002" ^
        "%COMFYUI_URL_PART2%"
)

rem if exist "%TEMP%\ComfyUI_Windows_portable_cu124.7z.003" (
rem    echo Found existing ComfyUI part 3 download, skipping download...
rem ) else (
rem    :: Download ComfyUI Part 3
rem    echo Downloading ComfyUI part 3...
rem    curl -L --retry 3 --retry-delay 2 ^
rem        -o "%TEMP%\ComfyUI_Windows_portable_cu124.7z.003" ^
rem        "%COMFYUI_URL_PART3%"
rem )

echo.
echo [Step 2] Installing 7-Zip...
if not exist "C:\Program Files\7-Zip\7z.exe" (
    curl -L -o "%TEMP%\7z-installer.exe" "%SEVENZIP_URL%"
    "%TEMP%\7z-installer.exe" /S
    timeout /t 3 /nobreak >nul
)

echo [Step 3] Extracting ComfyUI...
:: First, delete existing directory if it exists
if exist "%INSTALL_DIR%" (
    echo Removing existing installation...
    rd /s /q "%INSTALL_DIR%"
)
mkdir "%INSTALL_DIR%"

:: Change to TEMP directory for extraction
cd /d "%TEMP%"

echo Extracting to: %INSTALL_DIR%
"C:\Program Files\7-Zip\7z.exe" x "%TEMP%\ComfyUI_Windows_portable_cu124.7z.001" -o"%USERPROFILE%" -y

:: Verify the extraction and adjust directories
if exist "%PYTHON_DIR%\python.exe" (
    echo Python installation found at expected location...
) else (
    echo ERROR: Python not found in expected location
    echo Available files in installation directory:
    dir "%INSTALL_DIR%"
    pause
    exit /b 1
)

echo.
echo [Step 4] Downloading Hunyuan3D Models...
:: Ensure the diffusion_models directory exists
if not exist "%INSTALL_DIR%\ComfyUI\models\diffusion_models" (
    mkdir "%INSTALL_DIR%\ComfyUI\models\diffusion_models"
)

:: Download safetensors model
if not exist "%INSTALL_DIR%\ComfyUI\models\diffusion_models\hunyuan3d-dit-v2-0-fp16.safetensors" (
    echo Downloading Hunyuan3D safetensors model...
    curl -L -o "%INSTALL_DIR%\ComfyUI\models\diffusion_models\hunyuan3d-dit-v2-0-fp16.safetensors" ^
    "%HUNYUAN_SAFETENSORS_URL%"
    
    if !ERRORLEVEL! NEQ 0 (
        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.
    )
) else (
    echo Hunyuan3D safetensors model already exists. Skipping download.
)

:: Download .ckpt model
if not exist "%INSTALL_DIR%\ComfyUI\models\diffusion_models\hunyuan3d-dit-v2-0.ckpt" (
    echo Downloading Hunyuan3D .ckpt model...
    curl -L -o "%INSTALL_DIR%\ComfyUI\models\diffusion_models\hunyuan3d-dit-v2-0.ckpt" ^
    "%HUNYUAN_CKPT_URL%"
    
    if !ERRORLEVEL! NEQ 0 (
        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.
    )
) else (
    echo Hunyuan3D .ckpt model already exists. Skipping download.
)

echo.
echo [Step 5] Downloading Base Models...
if not exist "%CHECKPOINTS_DIR%" mkdir "%CHECKPOINTS_DIR%"

curl -L --retry 3 --retry-delay 2 -o "%CHECKPOINTS_DIR%\models.zip.001" "%MODELS_URL%"
)

rem echo.
rem echo [Step 5] Extracting Models...
rem "C:\Program Files\7-Zip\7z.exe" x "%TEMP%\models.zip.001" -o"%MODELS_DIR%" -y

echo.
echo [Step 6] Setting up Python Environment...
cd /d "%INSTALL_DIR%"
echo Initializing Python...
"%PYTHON_DIR%\python.exe" -m ensurepip
if !ERRORLEVEL! NEQ 0 (
    echo Failed to initialize pip
    exit /b 1
)
"%PYTHON_DIR%\python.exe" -m pip install --upgrade pip

echo.
echo [Step 7] Creating Custom Nodes Directory...
mkdir "%CUSTOM_NODES_DIR%" 2>nul
cd /d "%CUSTOM_NODES_DIR%"

echo.
echo [Step 8] Installing and Initializing Custom Nodes...

:: Define an array of custom node repositories
set "CUSTOM_NODES[0]=https://github.com/Kijai/ComfyUI-Hunyuan3DWrapper.git"
set "CUSTOM_NODES[1]=https://github.com/WASasquatch/was-node-suite-comfyui.git"
set "CUSTOM_NODES[2]=https://github.com/rgthree/rgthree-comfy.git"
set "CUSTOM_NODES[3]=https://github.com/Mikubill/sd-webui-controlnet.git"
set "CUSTOM_NODES[4]=https://github.com/ltdrdata/ComfyUI-Manager.git"
set "CUSTOM_NODES[5]=https://github.com/storyicon/comfyui_controlnet_aux.git"
set "CUSTOM_NODES[6]=https://github.com/MrForExample/ComfyUI-Inspire-Pack.git"
set "CUSTOM_NODES[7]=https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved.git"
set "CUSTOM_NODES[8]=https://github.com/ZHO-ZHO-ZHO/ComfyUI-Image-Selector.git"
set "CUSTOM_NODES[9]=https://github.com/pythongosssss/ComfyUI-Custom-Scripts.git"

:: Ensure custom nodes directory exists
if not exist "%CUSTOM_NODES_DIR%" mkdir "%CUSTOM_NODES_DIR%"
cd /d "%CUSTOM_NODES_DIR%"

:: Clone or update repositories
set "idx=0"
:clone_loop
if defined CUSTOM_NODES[%idx%] (
    for /f "tokens=*" %%a in ("!CUSTOM_NODES[%idx%]!") do (
        set "repo=%%a"
        set "repo_name=!repo:~38,-4!"
        
        if not exist "%CUSTOM_NODES_DIR%\!repo_name!" (
            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 ..
        )
    )
    set /a "idx+=1"
    goto :clone_loop
)

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 %%D in ("%CUSTOM_NODES_DIR%\*") do (
    if exist "%%D\.git" (
        echo Repository already initialized: %%~nxD
    ) else (
        echo Initializing Git for: %%~nxD
        cd "%%D"
        git init
        git remote add origin https://github.com/%%~nxD/%%~nxD.git
        cd ..
    )
)

echo.
echo [Step 8b] Installing Hunyuan3D Wrapper Dependencies...
if exist "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper" (
    echo Installing Hunyuan3D dependencies...
    "%PYTHON_DIR%\python.exe" -m pip install -r "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper\requirements.txt"

    :: Install precompiled rasterizer
    echo Installing precompiled rasterizer...
    "%PYTHON_DIR%\python.exe" -m pip install "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper\wheels\custom_rasterizer-0.1-cp312-cp312-win_amd64.whl"
)

echo.
echo [Step 9] Compiling Components...
if exist "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper\hy3dgen\texgen\custom_rasterizer" (
    cd /d "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper\hy3dgen\texgen\custom_rasterizer"
    "%PYTHON_DIR%\python.exe" setup.py install
)

if exist "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper\hy3dgen\texgen\differentiable_renderer" (
    cd /d "%CUSTOM_NODES_DIR%\ComfyUI-Hunyuan3DWrapper\hy3dgen\texgen\differentiable_renderer"
    "%PYTHON_DIR%\python.exe" setup.py build_ext --inplace
)

:: Cleanup
del "%TEMP%\ComfyUI_Windows_portable_cu124.7z.*" 2>nul
del "%TEMP%\models.zip.*" 2>nul
del "%TEMP%\7z-installer.exe" 2>nul

echo.
echo [Step 10] Creating Desktop Shortcut...
(
echo @echo off
echo cd /d "%INSTALL_DIR%"
echo start run_nvidia_gpu.bat
) > "%USERPROFILE%\Desktop\Start_ComfyUI.bat"

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.bat shortcut on your desktop.
pause