@echo off
setlocal enabledelayedexpansion

:CHOOSE_OPTION
REM Ask user for installation type
echo [33mChoose your preferred installation:[0m
echo [32mA) Fast-Lowvram install[0m
echo [32mB) Unoptimized normal model[0m
set /p "CHOICE=Enter your choice (A or B) and press Enter: "

if /i "%CHOICE%"=="A" (
    set "INSTALL_TYPE=fast-lowvram"
) else if /i "%CHOICE%"=="B" (
    set "INSTALL_TYPE=unoptimized"
) else (
echo [31mInvalid choice. Please enter A or B.[0m
    goto CHOOSE_OPTION
)

:CHOOSE_FLUX_SCHNELL
REM Ask user if they want to download FLUX SCHNELL Model
echo [33mDo you want to download SD3.5 Turbo model?[0m
echo [32mA) Yes[0m
echo [32mB) No[0m
set /p "FLUX_SCHNELL_CHOICE=Enter your choice (A or B) and press Enter: "

if /i "%FLUX_SCHNELL_CHOICE%"=="A" (
    set "DOWNLOAD_FLUX_SCHNELL=yes"
) else if /i "%FLUX_SCHNELL_CHOICE%"=="B" (
    set "DOWNLOAD_FLUX_SCHNELL=no"
) else (
echo [31mInvalid choice. Please enter A or B.[0m
    goto CHOOSE_FLUX_SCHNELL
)

:CHOOSE_FLUX_GGUF
REM Ask user if they want to download FLUX GGUF Model
echo [33mDo you want to download SD3.5 GGUF models?[0m
echo [32mA) Q8_0 (24GB Vram)[0m
echo [32mB) Q5_0 (16GB Vram)[0m
echo [32mC) Q4_0 (less than 12GB Vram)[0m
echo [32mD) All[0m
echo [32mE) No[0m
set /p "FLUX_GGUF_CHOICE=Enter your choice (A,B,C,D or E) and press Enter: "

if /i "%FLUX_GGUF_CHOICE%"=="A" (
    set "DOWNLOAD_GGUF=yes"
) else if /i "%FLUX_GGUF_CHOICE%"=="B" (
    set "DOWNLOAD_GGUF=yes"
) else if /i "%FLUX_GGUF_CHOICE%"=="C" (
    set "DOWNLOAD_GGUF=yes"
) else if /i "%FLUX_GGUF_CHOICE%"=="D" (
    set "DOWNLOAD_GGUF=yes"
) else if /i "%FLUX_GGUF_CHOICE%"=="E" (
    set "DOWNLOAD_GGUF=no"
) else (
echo [31mInvalid choice. Please enter A,B,C,D or E.[0m
    goto CHOOSE_FLUX_GGUF
)

REM Check if 7-Zip is installed and get its path
for %%I in (7z.exe) do set "SEVEN_ZIP_PATH=%%~$PATH:I"
if not defined SEVEN_ZIP_PATH (
    if exist "%ProgramFiles%\7-Zip\7z.exe" (
        set "SEVEN_ZIP_PATH=%ProgramFiles%\7-Zip\7z.exe"
    ) else if exist "%ProgramFiles(x86)%\7-Zip\7z.exe" (
        set "SEVEN_ZIP_PATH=%ProgramFiles(x86)%\7-Zip\7z.exe"
    ) else (
        echo 7-Zip is not installed. Downloading and installing...
        curl -L -o 7z-installer.exe https://www.7-zip.org/a/7z2201-x64.exe
        7z-installer.exe /S
        set "SEVEN_ZIP_PATH=%ProgramFiles%\7-Zip\7z.exe"
        if not exist "%SEVEN_ZIP_PATH%" (
            echo Installation of 7-Zip failed. Please install it manually and try again.
            exit /b 1
        )
        del 7z-installer.exe
    )
)

REM Check and install Git
git --version > NUL 2>&1
if %errorlevel% NEQ 0 (
    echo 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-2.41.0.3-64-bit.exe'; if ($LASTEXITCODE -ne 0) { exit 1 }}"
    if %errorlevel% NEQ 0 (
        echo Failed to download Git installer.
        exit /b
    )
    start /wait Git-2.41.0.3-64-bit.exe /VERYSILENT
    del Git-2.41.0.3-64-bit.exe
) else (
    echo Git already installed.
)

REM Download ComfyUI
echo [33mDownloading ComfyUI...[0m
curl -L -o ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z https://github.com/comfyanonymous/ComfyUI/releases/download/latest/ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z

REM Extract ComfyUI
echo Extracting ComfyUI...
"%SEVEN_ZIP_PATH%" x ComfyUI_windows_portable_nvidia_cu121_or_cpu.7z -o"%CD%" -y

REM Check if extraction was successful
if not exist "ComfyUI_windows_portable" (
    echo Extraction failed. Please check the downloaded file and try again.
    exit /b 1
)

REM Navigate to custom_nodes folder
cd ComfyUI_windows_portable\ComfyUI\custom_nodes

REM Clone ComfyUI-Manager
echo [33mCloning ComfyUI-Manager...[0m
git clone https://github.com/ltdrdata/ComfyUI-Manager.git

REM Navigate to models folder
cd ..\models

REM Download VAE file
echo [33mDownloading VAE file...[0m
cd vae
curl -L -o diffusion_pytorch_model.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/vae/diffusion_pytorch_model.safetensors?download=true

cd ..

REM Download CLIP files
echo [33mDownloading CLIP files...[0m
cd clip
curl -L -o clip_l.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/clip_l.safetensors?download=true
curl -L -o clip_g.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/clip_g.safetensors?download=true
curl -L -o t5xxl_fp8_e4m3fn.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5xxl_fp8_e4m3fn.safetensors?download=true
curl -L -o t5xxl_fp16.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5xxl_fp16.safetensors?download=true

cd ..

REM Download UNET file based on user choice
echo [33mDownloading checkpoints file...[0m
cd checkpoints
curl -L -o sd3.5_medium.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/checkpoints/sd3.5_medium.safetensors?download=true
if "%INSTALL_TYPE%"=="fast-lowvram" (
    curl -L -o sd3.5_large_fp8_scaled.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/checkpoints/sd3.5_large_fp8_scaled.safetensors?download=true
) else (
    curl -L -o sd3.5_large.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/checkpoints/sd3.5_large.safetensors?download=true
)

if "%DOWNLOAD_GGUF%"=="yes" (
    echo [33mDownloading SD3.5 GGUF Model...[0m
    if /i "%FLUX_GGUF_CHOICE%"=="A" (
	curl -L -o sd3.5_large-Q8_0.gguf https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/resolve/main/sd3.5_large-Q8_0.gguf?download=true
	) else if /i "%FLUX_GGUF_CHOICE%"=="B" (
	curl -L -o sd3.5_large-Q5_0.gguf https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/resolve/main/sd3.5_large-Q5_0.gguf?download=true
	) else if /i "%FLUX_GGUF_CHOICE%"=="C" (
	curl -L -o sd3.5_large-Q4_0.gguf https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/resolve/main/sd3.5_large-Q4_0.gguf?download=true
	) else if /i "%FLUX_GGUF_CHOICE%"=="D" (
	curl -L -o sd3.5_large-Q8_0.gguf https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/resolve/main/sd3.5_large-Q8_0.gguf?download=true
	curl -L -o sd3.5_large-Q5_0.gguf https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/resolve/main/sd3.5_large-Q5_0.gguf?download=true
	curl -L -o sd3.5_large-Q4_0.gguf https://huggingface.co/city96/stable-diffusion-3.5-large-gguf/resolve/main/sd3.5_large-Q4_0.gguf?download=true
	)
)

REM Download FLUX SCHNELL Model if user chose to
if "%DOWNLOAD_FLUX_SCHNELL%"=="yes" (
    echo [33mDownloading SD3.5 Turbo model...[0m
    curl -L -o sd3.5_large_turbo.safetensors https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/checkpoints/sd3.5_large_turbo.safetensors?download=true
)

cd ..

REM Download upscale model
echo [33mDownloading upscale models...[0m
cd upscale_models
curl -L -o 4x_NMKD-Siax_200k.pth https://huggingface.co/uwg/upscaler/resolve/main/ESRGAN/4x_NMKD-Siax_200k.pth?download=true
curl -L -o 4x-ClearRealityV1.pth https://huggingface.co/skbhadra/ClearRealityV1/resolve/main/4x-ClearRealityV1.pth?download=true

cd ..\..

mkdir .\user\default
echo [33mDownloading comfy settings...[0m
cd user\default
curl -L -o comfy.settings.json https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/comfy.settings.json?download=true
echo [33mDownloading comfy workflow...[0m
mkdir .\workflows
cd workflows
curl -L -o "UmeAiRT - Flux workflow.7z" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/workflows/UmeAiRT-Flux_workflow.7z?download=true
"%SEVEN_ZIP_PATH%" x "UmeAiRT - Flux workflow.7z" -o"%CD%" -y   >nul 2>&1
del /f "UmeAiRT - Flux workflow.7z" -force   >nul 2>&1
cd ..\..\..\..

cd ComfyUI
cd custom_nodes
mkdir ComfyUI-NAI-styler
cd ComfyUI-NAI-styler
mkdir CSV
cd CSV
curl -L -o "naifilters.csv" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/naifilters.csv?download=true
curl -L -o "naistyles.csv" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/naistyles.csv?download=true
curl -L -o "naitypes.csv" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/naitypes.csv?download=true

cd ..\..\..\..

REM Final steps based on user choice
if "%INSTALL_TYPE%"=="fast-lowvram" (
    echo Downloading special run file for fast-lowvram...
    curl -L -o "run_nvidia_gpu-LOWVRAM.bat" "https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/scripts/run_nvidia_gpu-LOWVRAM.bat?download=true"
    echo ComfyUI and FLUX installed. Running ComfyUI...
    call "run_nvidia_gpu-LOWVRAM.bat"
) else (
    echo ComfyUI and FLUX installed. HAVE FUN ;)
)