@echo off
setlocal enabledelayedexpansion

echo.
echo ============================
echo  ComfyUI Auto Installer
echo ============================
echo.

REM Define paths
set "comfyPath=%CD%\ComfyUI"
set "customNodesPath=%comfyPath%\custom_nodes"
set "venvPath=%comfyPath%\venv"
set "pythonPath=%venvPath%\Scripts\python.exe"
set "activateScript=%venvPath%\Scripts\activate.bat"

REM -------------------------------
REM Check if Python is available
REM -------------------------------
where python >nul 2>&1
if %errorlevel% NEQ 0 (
    echo Python is not found in PATH.
    echo Would you like to install Python 3.12 now? [Y/N]
    set /p "PYINSTALL=Your choice [Y/N]: "
    if /i "!PYINSTALL!"=="Y" (
        echo Downloading Python 3.12...
        %SystemRoot%\System32\curl.exe -L -o python-installer.exe https://www.python.org/ftp/python/3.12.3/python-3.12.3-amd64.exe
        if not exist python-installer.exe (
            echo Failed to download Python installer.
            pause
            exit /b 1
        )
        echo Installing Python 3.12...
        start /wait python-installer.exe /quiet InstallAllUsers=1 PrependPath=1 Include_test=0
        del /f python-installer.exe
    ) else (
        echo Cannot continue without Python. Exiting...
        pause
        exit /b 1
    )
)

REM Get Python version from stderr
for /f "tokens=2 delims= " %%A in ('python -V 2^>^&1') do set PYVERSION=%%A

REM Parse major and minor
for /f "tokens=1,2 delims=." %%B in ("%PYVERSION%") do (
    set "PYMAJOR=%%B"
    set "PYMINOR=%%C"
)

echo Detected Python version: %PYVERSION%
if not "%PYMAJOR%.%PYMINOR%"=="3.12" (
    echo Your current Python version is %PYVERSION%. It may not be supported.
    echo Would you like to:
    echo   [Y] Install Python 3.12
    echo   [N] Continue using current version
    set /p "PYCHOICE=Choose [Y/N]: "
    if /i "!PYCHOICE!"=="Y" (
        echo Downloading Python 3.12...
        %SystemRoot%\System32\curl.exe -L -o python-installer.exe https://www.python.org/ftp/python/3.12.3/python-3.12.3-amd64.exe
        if not exist python-installer.exe (
            echo Failed to download Python installer.
            pause
            exit /b 1
        )
        echo Installing Python 3.12...
        start /wait python-installer.exe /quiet InstallAllUsers=1 PrependPath=1 Include_test=0
        del /f python-installer.exe
    )
)

REM -------------------------------
REM Check Git
REM -------------------------------
git --version >nul 2>&1
if %errorlevel% NEQ 0 (
    echo Git not found. Installing Git...
    powershell -Command "& {Invoke-WebRequest -Uri 'https://github.com/git-for-windows/git/releases/download/v2.41.0.windows.3/Git-2.41.0.3-64-bit.exe' -OutFile 'git_installer.exe'}"
    if %errorlevel% NEQ 0 (
        echo Failed to download Git.
        pause
        exit /b 1
    )
    start /wait git_installer.exe /VERYSILENT
    del /f git_installer.exe
    echo Git installed successfully.
)

REM -------------------------------
REM Clone ComfyUI
REM -------------------------------
if not exist "%comfyPath%" (
    echo Cloning ComfyUI...
    git clone https://github.com/comfyanonymous/ComfyUI.git "%comfyPath%"
) else (
    echo ComfyUI already exists. Skipping clone.
)

REM -------------------------------
REM Create venv
REM -------------------------------
if not exist "%venvPath%" (
    echo Creating virtual environment...
    python -m venv "%venvPath%"
)

REM -------------------------------
REM Activate venv and install deps
REM -------------------------------
call "%activateScript%"
echo Installing CUDA-enabled torch manually...
"%pythonPath%" -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu128

echo Installing dependencies...
"%pythonPath%" -m pip install --upgrade pip
"%pythonPath%" -m pip install -r "%comfyPath%\requirements.txt"

REM -------------------------------
REM Clone custom nodes
REM -------------------------------
if not exist "%customNodesPath%" (
    mkdir "%customNodesPath%"
)

echo Cloning custom nodes...

set repos[0]=https://github.com/ltdrdata/ComfyUI-Impact-Pack
set repos[1]=https://github.com/ltdrdata/ComfyUI-Impact-Subpack
set repos[2]=https://github.com/Smirnov75/ComfyUI-mxToolkit
set repos[3]=https://github.com/yolain/ComfyUI-Easy-Use
set repos[4]=https://github.com/pythongosssss/ComfyUI-Custom-Scripts
set repos[5]=https://github.com/crystian/ComfyUI-Crystools
set repos[6]=https://github.com/alexopus/ComfyUI-Image-Saver
set repos[7]=https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes
set repos[8]=https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet
set repos[9]=https://github.com/kijai/ComfyUI-KJNodes
set repos[10]=https://github.com/Fannovel16/comfyui_controlnet_aux
set repos[11]=https://github.com/vslinx/ComfyUI-vslinx-nodes.git
set repos[12]=https://github.com/chrisgoringe/cg-image-filter
set repos[13]=https://github.com/rgthree/rgthree-comfy
set repos[14]=https://github.com/cubiq/ComfyUI_IPAdapter_plus
set repos[15]=https://github.com/pythongosssss/ComfyUI-WD14-Tagger.git
set repos[16]=https://github.com/Comfy-Org/ComfyUI-Manager
set repos[17]=https://github.com/wallish77/wlsh_nodes.git
set repos[18]=https://github.com/lquesada/ComfyUI-Inpaint-CropAndStitch

for /L %%i in (0,1,18) do (
    set "repo=!repos[%%i]!"
    for %%A in (!repo!) do (
        set "folderName=%%~nxA"
        set "targetPath=%customNodesPath%\!folderName!"

        if not exist "!targetPath!" (
            echo   - Cloning !folderName!...
            git clone !repo! "!targetPath!"
        ) else (
            echo   - !folderName! already exists. Skipping.
        )

        REM Install requirements if available
        if exist "!targetPath!\requirements.txt" (
            echo     Installing requirements for !folderName!...
            "%pythonPath%" -s -m pip install -r "!targetPath!\requirements.txt"
        )
    )
)

REM -------------------------------
REM Download JSON workflow files
REM -------------------------------
set "workflowFolder=%comfyPath%\user\default\workflows"
if not exist "%workflowFolder%" (
    mkdir "%workflowFolder%"
)

REM -------------------------------
REM Download SAM model (Segment Anything)
REM -------------------------------
set "samFolder=%comfyPath%\models\sams"
if not exist "%samFolder%" (
    mkdir "%samFolder%"
)

set "samFile=%samFolder%\sam_vit_b_01ec64.pth"
if exist "%samFile%" (
    echo   - sam_vit_b_01ec64.pth already exists. Skipping download.
) else (
    echo Downloading SAM model...
    %SystemRoot%\System32\curl.exe -L -o "%samFile%" ^
        https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth
)

echo Checking workflow files...

REM Define workflow filenames and URLs
set workflow[0]=TXT2IMG-ADetailer-v5.0-vslinx.json
set url[0]=https://huggingface.co/vslinx/ComfyUIDetailerWorkflow-vslinx/resolve/main/workflows/TXT2IMG/v5.0/TXT2IMG-ADetailer-v5.0-vslinx.json

set workflow[1]=IMG2IMG-ADetailer-v5.0-vslinx.json
set url[1]=https://huggingface.co/vslinx/ComfyUIDetailerWorkflow-vslinx/resolve/main/workflows/IMG2IMG/v5.0/IMG2IMG-ADetailer-v5.0-vslinx.json

REM Loop through and download if not already present
for /L %%i in (0,1,1) do (
    call set "file=%%workflow[%%i]%%"
    call set "link=%%url[%%i]%%"
    set "path=%workflowFolder%\!file!"

    if exist "!path!" (
        echo   - !file! already exists. Skipping download.
    ) else (
        echo   ↓ Downloading !file!...
        %SystemRoot%\System32\curl.exe -L -o "!path!" "!link!"
    )
)

echo.
echo Would you like to download required models (ControlNet, IPAdapter, Upscale Model, etc.)?
echo   [Y] Yes
echo   [N] No
set /p "download_models=Choose [Y/N]: "

if /i "%download_models%"=="Y" (
    echo Downloading required models...

    REM Create folders if they don't exist
    if not exist "%comfyPath%\models\vae" mkdir "%comfyPath%\models\vae"
    if not exist "%comfyPath%\models\upscale_models" mkdir "%comfyPath%\models\upscale_models"
    if not exist "%comfyPath%\models\ipadapter" mkdir "%comfyPath%\models\ipadapter"
    if not exist "%comfyPath%\models\clip_vision" mkdir "%comfyPath%\models\clip_vision"
    if not exist "%comfyPath%\models\controlnet" mkdir "%comfyPath%\models\controlnet"

    REM Download files
    echo Downloading VAE...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\vae\sdxl_vae.safetensors" ^
        https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors

    echo Downloading Upscale Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\upscale_models\4x_foolhardy_Remacri.pth" ^
        https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth

    echo Downloading IPAdapter model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\ipadapter\noobIPAMARK1_mark1.safetensors" ^
        https://huggingface.co/vslinx/ComfyUIDetailerWorkflow-vslinx/resolve/main/models/IPAdapter/noobIPAMARK1_mark1.safetensors

    echo Downloading CLIP Vision Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\clip_vision\CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors" ^
        https://huggingface.co/XuminYu/example_safetensors/resolve/4b89d7ebd99a9913f0abbec4bf0f54932b11d243/CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors

    echo Downloading ControlNet Canny Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsCanny.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-canny/resolve/main/noob_sdxl_controlnet_canny.safetensors   

    echo Downloading ControlNet DepthMidas Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsDepthMidasv1-1.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-depth_midas-v1-1/resolve/main/noob-sdxl-controlnet-depth-midas-v1-1.safetensors      

    echo Downloading ControlNet Lineart-Anime Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsLineartAnime.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-lineart_anime/resolve/main/noob-sdxl-controlnet-lineart_anime.safetensors      

    echo Downloading ControlNet Lineart-Realistic Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsLineartRealistic.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-lineart_realistic/resolve/main/noob-sdxl-controlnet-lineart_realistic.safetensors  

    echo Downloading ControlNet Manga line Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsMangaLine.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-manga_line/resolve/main/noob-sdxl-controlnet-manga-line.safetensors    

    echo Downloading ControlNet OpenPose Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsOpenPose.safetensors" ^
        https://huggingface.co/Laxhar/noob_openpose/resolve/main/openpose_pre.safetensors

    echo Downloading ControlNet Softedge Hed Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsSoftedgeHed.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-softedge_hed/resolve/main/noob-sdxl-controlnet-softedge_hed.safetensors

    echo Downloading ControlNet Depth Model...
    %SystemRoot%\System32\curl.exe -L -o "%comfyPath%\models\controlnet\noobaiXLControlnet_epsDepth.safetensors" ^
        https://huggingface.co/Eugeoter/noob-sdxl-controlnet-depth/resolve/main/noob_sdxl_controlnet_depth.safetensors                                

    echo All required models have been downloaded.
) else (
    echo Skipping model downloads.
)

REM -------------------------------
REM Setup Complete
REM -------------------------------
echo.
echo ============================
echo   Setup Complete!
echo ============================

REM -------------------------------
REM Offer to create Desktop shortcut
REM -------------------------------
echo Would you like to create a desktop shortcut to start ComfyUI?
echo   [Y] Yes
echo   [N] No
set /p "MAKE_SHORTCUT=Choose [Y/N]: "

if /i "!MAKE_SHORTCUT!"=="Y" (
    set "shortcutBat=%USERPROFILE%\Desktop\Start_ComfyUI.bat"
    echo @echo off > "!shortcutBat!"
    echo cd /d "%comfyPath%" >> "!shortcutBat!"
    echo call "%venvPath%\Scripts\activate.bat" >> "!shortcutBat!"
    echo python main.py >> "!shortcutBat!"
    echo pause >> "!shortcutBat!"
    echo Shortcut created on Desktop as Start_ComfyUI.bat
)

REM -------------------------------
REM Ask to start ComfyUI
REM -------------------------------
echo Would you like to start ComfyUI now?
echo   [Y] Yes
echo   [N] No
set /p "STARTNOW=Choose [Y/N]: "
if /i "!STARTNOW!"=="Y" (
    echo Starting ComfyUI in a new shell...
    start "" cmd /k ^
        "cd /d "%comfyPath%" && call "%venvPath%\Scripts\activate.bat" && python main.py"
)

pause