@echo off
setlocal enabledelayedexpansion

set "basePath=%CD%"

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

:check_folder
if exist "%basePath%\ComfyUI" (
    echo ComfyUI folder detected
	set "modelsPath=%basePath%\ComfyUI\models"
) else if exist "%basePath%\ComfyUI_windows_portable" (
    echo ComfyUI folder detected
	set "modelsPath=%basePath%\ComfyUI_windows_portable\ComfyUI\models"
) else (
	echo ComfyUI folder not detected, give ComfyUI folder path:
    set /p "comfyPath=Path: "
    set "modelsPath=!comfyPath!\models"
    echo comfyPath = !comfyPath!
    echo modelsPath = !modelsPath!
    pause
)

:CHOOSE_OPTION
REM Ask user for installation type
echo Do you want to download FLUX base models?
echo A) Base
echo B) fp8
echo C) All
echo D) No
set /p "CHOICE=Enter your choice (A,B,C or D) and press Enter: "
	if /i "%CHOICE%"=="A" (
	set "DOWNLOAD=yes"
) else if /i "%CHOICE%"=="B" (
	set "DOWNLOAD=yes"
) else if /i "%CHOICE%"=="C" (
	set "DOWNLOAD=yes"
) else if /i "%CHOICE%"=="D" (
	set "DOWNLOAD=yes"
) else (
	echo Invalid choice. Please enter A,B,C or D.
	goto CHOOSE_OPTION
)

:CHOOSE_FLUX_GGUF
REM Ask user if they FLUXt to download FLUX GGUF Model
echo Do you FLUXt to download FLUX GGUF models?
echo A) Q8_0 + T5_Q8 (24GB Vram)
echo B) Q5_K_S + T5_Q5_K_M (16GB Vram)
echo C) Q4_K_S + T5_Q3_K_L (less than 12GB Vram)
echo D) All
echo E) No
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 Invalid choice. Please enter A,B,C,D or E.
	goto CHOOSE_FLUX_GGUF
)

:CHOOSE_FLUX_SCHNELL
REM Ask user if they want to download FLUX SCHNELL Model
echo Do you want to download the FLUX SCHNELL Model?
echo A) Yes
echo B) No
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 Invalid choice. Please enter A or B.
    goto CHOOSE_FLUX_SCHNELL
)

:CHOOSE_FLUX_CONTROLNET
REM Ask user if they want to download FLUX SCHNELL Model
echo Do you want to download FLUX ControlNet Models?
echo A) Yes
echo B) No
set /p "FLUX_CONTROLNET_CHOICE=Enter your choice (A or B) and press Enter: "

if /i "%FLUX_CONTROLNET_CHOICE%"=="A" (
    set "DOWNLOAD_FLUX_CONTROLNET=yes"
) else if /i "%FLUX_CONTROLNET_CHOICE%"=="B" (
    set "DOWNLOAD_FLUX_CONTROLNET=no"
) else (
echo Invalid choice. Please enter A or B.
    goto CHOOSE_FLUX_CONTROLNET
)

:CHOOSE_FLUX_PULID
REM Ask user if they want to download FLUX SCHNELL Model
echo Do you want to download FLUX PuLID and REDUX Models?
echo A) Yes
echo B) No
set /p "FLUX_PULID_CHOICE=Enter your choice (A or B) and press Enter: "

if /i "%FLUX_PULID_CHOICE%"=="A" (
    set "DOWNLOAD_FLUX_PULID=yes"
) else if /i "%FLUX_PULID_CHOICE%"=="B" (
    set "DOWNLOAD_FLUX_PULID=no"
) else (
echo Invalid choice. Please enter A or B.
    goto CHOOSE_FLUX_PULID
)

:CHOOSE_FLUX_LORA
REM Ask user if they want to download FLUX LORA
echo Do you want to download UmeAiRT LoRAs?
echo A) Yes
echo B) No
set /p "FLUX_LORA_CHOICE=Enter your choice (A or B) and press Enter: "

if /i "%FLUX_LORA_CHOICE%"=="A" (
    set "DOWNLOAD_FLUX_LORA=yes"
) else if /i "%FLUX_LORA_CHOICE%"=="B" (
    set "DOWNLOAD_FLUX_LORA=no"
) else (
echo Invalid choice. Please enter A or B.
    goto CHOOSE_FLUX_LORA
)

if /i "%CHOICE%"=="A" (
    echo Downloading unet models file...
	powershell -NoProfile -Command "$p = Join-Path $env:modelsPath 'diffusion_models\FLUX'; if (-Not (Test-Path $p)) { New-Item -ItemType Directory -Path $p | Out-Null }"
	curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-dev.sft" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-dev.sft?download=true
) else if /i "%CHOICE%"=="B" (
    echo Downloading unet models file...
	powershell -NoProfile -Command "$p = Join-Path $env:modelsPath 'diffusion_models\FLUX'; if (-Not (Test-Path $p)) { New-Item -ItemType Directory -Path $p | Out-Null }"
	curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-dev-fp8.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-dev-fp8.safetensors?download=true
) else if /i "%CHOICE%"=="C" (
    echo Downloading unet models file...
	powershell -NoProfile -Command "$p = Join-Path $env:modelsPath 'diffusion_models\FLUX'; if (-Not (Test-Path $p)) { New-Item -ItemType Directory -Path $p | Out-Null }"
	curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-dev.sft" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-dev.sft?download=true
    curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-dev-fp8.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-dev-fp8.safetensors?download=true
) else if /i "%CHOICE%"=="D" (
    set "INSTALL_TYPE=optimized"
) else (
echo Invalid choice. Please enter A,B,C or D.
    goto CHOOSE_OPTION
)

REM Download VAE file
echo Downloading VAE file...
curl -L -o "%modelsPath%/vae/ae.safetensors" "https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/vae/ae.safetensors?download=true"

REM Download CLIP files
echo Downloading CLIP files...
curl -L -o "%modelsPath%/clip/clip_l.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/clip_l.safetensors?download=true
curl -L -o "%modelsPath%/clip/longclip-L.pt" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/longclip-L.pt?download=true
curl -L -o "%modelsPath%/clip/t5xxl_fp8_e4m3fn.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5xxl_fp8_e4m3fn.safetensors?download=true
curl -L -o "%modelsPath%/clip/t5xxl_fp16.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5xxl_fp16.safetensors?download=true
curl -L -o "%modelsPath%/clip/ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/ViT-L-14-TEXT-detail-improved-hiT-GmP-TE-only-HF.safetensors?download=true

if "%DOWNLOAD_GGUF%"=="yes" (
    echo Downloading FLUX GGUF Quant Model...
	powershell -NoProfile -Command "$p = Join-Path $env:modelsPath 'unet\FLUX'; if (-Not (Test-Path $p)) { New-Item -ItemType Directory -Path $p | Out-Null }"
    if /i "%FLUX_GGUF_CHOICE%"=="A" (
	curl -L -o "%modelsPath%/clip/t5-v1_1-xxl-encoder-Q8_0.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5-v1_1-xxl-encoder-Q8_0.gguf?download=true
	curl -L -o "%modelsPath%/unet/FLUX/flux1-dev-Q8_0.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-dev-Q8_0.gguf?download=true
	) else if /i "%FLUX_GGUF_CHOICE%"=="B" (
	curl -L -o "%modelsPath%/clip/t5-v1_1-xxl-encoder-Q5_K_M.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5-v1_1-xxl-encoder-Q5_K_M.gguf?download=true
	curl -L -o "%modelsPath%/unet/FLUX/flux1-dev-Q5_K_S.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-dev-Q5_K_S.gguf?download=true
	) else if /i "%FLUX_GGUF_CHOICE%"=="C" (
	curl -L -o "%modelsPath%/clip/t5-v1_1-xxl-encoder-Q3_K_S.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5-v1_1-xxl-encoder-Q3_K_S.gguf?download=true
	curl -L -o "%modelsPath%/unet/FLUX/flux1-dev-Q4_K_S.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/flux1-dev-Q4_K_S.gguf?download=true
	) else if /i "%FLUX_GGUF_CHOICE%"=="D" (
	curl -L -o "%modelsPath%/clip/t5-v1_1-xxl-encoder-Q8_0.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5-v1_1-xxl-encoder-Q8_0.gguf?download=true
	curl -L -o "%modelsPath%/clip/t5-v1_1-xxl-encoder-Q5_K_M.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5-v1_1-xxl-encoder-Q5_K_M.gguf?download=true
	curl -L -o "%modelsPath%/clip/t5-v1_1-xxl-encoder-Q3_K_S.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip/t5-v1_1-xxl-encoder-Q3_K_S.gguf?download=true
	curl -L -o "%modelsPath%/unet/FLUX/flux1-dev-Q8_0.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-dev-Q8_0.gguf?download=true
	curl -L -o "%modelsPath%/unet/FLUX/flux1-dev-Q5_K_S.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-dev-Q5_K_S.gguf?download=true
	curl -L -o "%modelsPath%/unet/FLUX/flux1-dev-Q4_K_S.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/flux1-dev-Q4_K_S.gguf?download=true
	)
)
REM Download clip vision
echo Downloading clip vision file...
curl -L -o "%modelsPath%/clip_vision/clip_vision_h.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip_vision/clip_vision_h.safetensors?download=true

REM Download upscale model
echo Downloading upscale models...
curl -L -o "%modelsPath%/upscale_models/RealESRGAN_x4plus.pth" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/upscale_models/RealESRGAN_x4plus.pth?download=true
curl -L -o "%modelsPath%/upscale_models/RealESRGAN_x4plus_anime_6B.pth" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/upscale_models/RealESRGAN_x4plus_anime_6B.pth?download=true

REM Download FLUX SCHNELL Model if user chose to
if "%DOWNLOAD_FLUX_SCHNELL%"=="yes" (
    echo Downloading FLUX SCHNELL Model...
    curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-schnell-fp8.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-schnell-fp8.safetensors?download=true
)

mkdir "%modelsPath%/LLM"
mkdir "%modelsPath%/xlabs/controlnets"

if "%DOWNLOAD_FLUX_CONTROLNET%"=="yes" (
    echo Downloading FLUX ControlNet Models...
	powershell -NoProfile -Command "$p = Join-Path $env:modelsPath 'unet\FLUX'; if (-Not (Test-Path $p)) { New-Item -ItemType Directory -Path $p | Out-Null }"
    curl -L -o "%modelsPath%/xlabs/controlnets/flux-canny-controlnet-v3.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/xlabs/controlnets/flux-canny-controlnet-v3.safetensors?download=true
    curl -L -o "%modelsPath%/xlabs/controlnets/flux-depth-controlnet-v3.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/xlabs/controlnets/flux-depth-controlnet-v3.safetensors?download=true
    curl -L -o "%modelsPath%/xlabs/controlnets/flux-hed-controlnet-v3.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/xlabs/controlnets/flux-hed-controlnet-v3.safetensors?download=true
	curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-depth-dev-fp8.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-depth-dev-fp8.safetensors?download=true
	curl -L -o "%modelsPath%/diffusion_models/FLUX/flux1-canny-dev-fp8.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/diffusion_models/FLUX/flux1-canny-dev-fp8.safetensors?download=true
	if "%DOWNLOAD_GGUF%"=="yes" (
		echo Downloading ControlNet GGUF Model...
		if /i "%FLUX_GGUF_CHOICE%"=="A" (
		curl -L -o "%modelsPath%/unet/FLUX/flux1-depth-dev-fp16-Q8_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-depth-dev-fp16-Q8_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-canny-dev-fp16-Q8_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-canny-dev-fp16-Q8_0-GGUF.gguf?download=true
		) else if /i "%FLUX_GGUF_CHOICE%"=="B" (
		curl -L -o "%modelsPath%/unet/FLUX/flux1-depth-dev-fp16-Q5_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-depth-dev-fp16-Q5_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-canny-dev-fp16-Q5_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-canny-dev-fp16-Q5_0-GGUF.gguf?download=true
		) else if /i "%FLUX_GGUF_CHOICE%"=="C" (
		curl -L -o "%modelsPath%/unet/FLUX/flux1-depth-dev-fp16-Q4_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-depth-dev-fp16-Q4_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-canny-dev-fp16-Q4_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-canny-dev-fp16-Q4_0-GGUF.gguf?download=true
		) else if /i "%FLUX_GGUF_CHOICE%"=="D" (
		curl -L -o "%modelsPath%/unet/FLUX/flux1-depth-dev-fp16-Q8_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-depth-dev-fp16-Q8_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-depth-dev-fp16-Q5_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-depth-dev-fp16-Q5_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-depth-dev-fp16-Q4_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-depth-dev-fp16-Q4_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-canny-dev-fp16-Q8_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-canny-dev-fp16-Q8_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-canny-dev-fp16-Q5_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-canny-dev-fp16-Q5_0-GGUF.gguf?download=true
		curl -L -o "%modelsPath%/unet/FLUX/flux1-canny-dev-fp16-Q4_0-GGUF.gguf" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/unet/FLUX/flux1-canny-dev-fp16-Q4_0-GGUF.gguf?download=true
		)
	)
	curl -L -o "%modelsPath%/controlnet/diffusion_pytorch_model_promax.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/controlnet/diffusion_pytorch_model_promax.safetensors?download=true
	curl -L -o "%modelsPath%/controlnet/Shakker-LabsFLUX1-dev-ControlNet-Union-Pro.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/controlnet/Shakker-LabsFLUX1-dev-ControlNet-Union-Pro.safetensors?download=true
)

REM Download FLUX SCHNELL Model if user chose to
if "%DOWNLOAD_FLUX_LORA%"=="yes" (
    echo Downloading UmeAiRT LoRAs...
	powershell -NoProfile -Command "$p = Join-Path $env:modelsPath 'loras\FLUX'; if (-Not (Test-Path $p)) { New-Item -ItemType Directory -Path $p | Out-Null }"
    curl -L -o "%modelsPath%/loras/FLUX/ume_sky_v2.safetensors" https://huggingface.co/UmeAiRT/FLUX.1-dev-LoRA-Ume_Sky/resolve/main/ume_sky_v2.safetensors?download=true
    curl -L -o "%modelsPath%/loras/FLUX/ume_modern_pixelart.safetensors" https://huggingface.co/UmeAiRT/FLUX.1-dev-LoRA-Modern_Pixel_art/resolve/main/ume_modern_pixelart.safetensors?download=true
    curl -L -o "%modelsPath%/loras/FLUX/ume_classic_Romanticism.safetensors" https://huggingface.co/UmeAiRT/FLUX.1-dev-LoRA-Romanticism/resolve/main/ume_classic_Romanticism.safetensors?download=true
    curl -L -o "%modelsPath%/loras/FLUX/ume_classic_impressionist.safetensors" https://huggingface.co/UmeAiRT/FLUX.1-dev-LoRA-Impressionism/resolve/main/ume_classic_impressionist.safetensors?download=true
    curl -L -o "%modelsPath%/loras/FLUX/ume_the-little-newspaper.safetensors" https://huggingface.co/UmeAiRT/FLUX.1-dev-LoRA-Ume_J1900/resolve/main/umej1900.safetensors?download=true
    curl -L -o "%modelsPath%/loras/FLUX/ume_knight.safetensors" https://huggingface.co/UmeAiRT/FLUX.1-dev-LoRA-Ume_Knight/resolve/main/ume_gachaak.safetensors?download=true
)

REM Download FLUX SCHNELL Model if user chose to
if "%DOWNLOAD_FLUX_PuLID%"=="yes" (
    echo Downloading FLUX PuLID...
	mkdir "%modelsPath%/pulid"
    curl -L -o "%modelsPath%/pulid/pulid_flux_v0.9.0.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/pulid/pulid_flux_v0.9.0.safetensors?download=true
    curl -L -o "%modelsPath%/clip_vision/sigclip_vision_patch14_384.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/clip_vision/sigclip_vision_patch14_384.safetensors?download=true
    curl -L -o "%modelsPath%/style_models/flux1-redux-dev.safetensors" https://huggingface.co/UmeAiRT/ComfyUI-Auto_installer/resolve/main/models/style_models/flux1-redux-dev.safetensors?download=true
)

echo Models downloaded.
pause