@echo off
setlocal enabledelayedexpansion

set "basePath=%CD%"
set "comfyPath=%basePath%\ComfyUI"
set "pythonPath=%basePath%\python_embeded"
set "customNodesPath=%comfyPath%\custom_nodes"

if not exist "%basePath%\logs" mkdir "%basePath%\logs"

curl -L -o banner.txt https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/banner.txt?download=true >> "%basePath%\logs\SageAttention.txt" 2>&1
echo -------------------------------------------------------------------------------
type banner.txt
echo -------------------------------------------------------------------------------
echo                   ComfyUI - WAN2.1 - SageAttention installer                          
echo                                                                     V1.3          
echo -------------------------------------------------------------------------------
del /f banner.txt

:check_folder
if exist "%basePath%\ComfyUI" (
    echo ComfyUI folder detected
	set "comfyPath=%basePath%\ComfyUI"
	if exist "%basePath%\python_embeded" (
		echo Python folder detected
		set "pythonPath=%basePath%\python_embeded"
	) else (
		echo Python folder not detected, give Python folder path :
	pause
	)
) else if exist "%basePath%\ComfyUI_windows_portable" (
	set "comfyPath=%basePath%\ComfyUI_windows_portable\ComfyUI"
	set "pythonPath=%basePath%\ComfyUI_windows_portable\python_embeded"
    echo ComfyUI folder detected
) else (
	echo ComfyUI folder not detected, give ComfyUi folder path :
	set /p "comfyPath=Path: "
	echo Python folder not detected, give Python folder path :
	set /p "pythonPath=Path: "
)

:CHOOSE_CLEAN
REM Ask user if they want to download FLUX SCHNELL Model
echo Do you want to do a clean install? (old triton and sageattention will be deleted)
echo A) Yes
echo B) No
set /p "CHOOSE_CLEAN=Enter your choice (A or B) and press Enter: "

if /i "%CHOOSE_CLEAN%"=="A" (
	echo Uninstalling Triton and SageAttention...
	"%pythonPath%\python.exe" -m pip uninstall -y triton-windows  >> "%basePath%\logs\SageAttention.txt" 2>&1
	"%pythonPath%\python.exe" -m pip uninstall -y triton  >> "%basePath%\logs\SageAttention.txt" 2>&1
	"%pythonPath%\python.exe" -m pip uninstall -y sageattention  >> "%basePath%\logs\SageAttention.txt" 2>&1

	echo Removing SageAttention build files...
	rmdir /s /q "SageAttention"  >> "%basePath%\logs\SageAttention.txt" 2>&1
	rmdir /s /q "%pythonPath%\libs"  >> "%basePath%\logs\SageAttention.txt" 2>&1
	rmdir /s /q "%pythonPath%\include"  >> "%basePath%\logs\SageAttention.txt" 2>&1
) else if /i "%CHOOSE_CLEAN%"=="B" (
    set "CHOOSE_CLEANL=no"
) else (
echo Invalid choice. Please enter A or B.
    goto CHOOSE_CLEAN
)

echo Installing Visual Studio Build Tools...
winget install --id Microsoft.VisualStudio.2022.BuildTools -e --source winget --override "--quiet --wait --norestart --add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 --add Microsoft.VisualStudio.Component.Windows10SDK.20348"

echo Installing Triton...
curl -L -o "%pythonPath%\triton-3.3.0-py3-none-any.whl" https://github.com/woct0rdho/triton-windows/releases/download/empty/triton-3.3.0-py3-none-any.whl >> "%basePath%\logs\SageAttention.txt" 2>&1
"%pythonPath%\python.exe" -m pip install "%pythonPath%\triton-3.3.0-py3-none-any.whl" >> "%basePath%\logs\SageAttention.txt" 2>&1
"%pythonPath%\python.exe" -s -m pip install triton-windows >> "%basePath%\logs\SageAttention.txt" 2>&1

echo Downloading Python include/libs...
curl -L -o "%pythonPath%\python_3.12.9_include_libs.zip" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/others/python_3.12.9_include_libs.zip?download=true >> "%basePath%\logs\SageAttention.txt" 2>&1
tar -xf "%pythonPath%\python_3.12.9_include_libs.zip" -C "%pythonPath%" >> "%basePath%\logs\SageAttention.txt" 2>&1

echo Downloading SageAttention...
git clone https://github.com/thu-ml/SageAttention.git >> "%basePath%\logs\SageAttention.txt" 2>&1

echo Installing SageAttention...
"%pythonPath%\python.exe" -s -m pip install -e SageAttention >> "%basePath%\logs\SageAttention.txt" 2>&1
rmdir /s /q "SageAttention"  >> "%basePath%\logs\SageAttention.txt" 2>&1

echo Installation complete
pause