@echo off
rem usage: install-comfyui-depends-qwen-image-21.cmd [--force] [--delete] [comfyui_root]
rem   --force      refetch everything, ignore size match, never prompt
rem   --delete     remove the model files and exit, no network access
rem   ASSUME_YES=1 answer yes to mismatch prompts when not on a terminal
rem   HF_TOKEN=hf_...   bearer token, forwarded across the CDN redirect
rem   DEL=1         alias for --force

setlocal enabledelayedexpansion

set "BASE=https://huggingface.co/Comfy-Org/Qwen-Image-2.1/resolve/main"
set "FORCE=0"
set "MODE=fetch"
set "ROOT=C:\AI\ComfyUI"

:parse
if "%~1"=="" goto done_parse
set "arg=%~1"
if /i "!arg!"=="--force" set "FORCE=1" & shift & goto parse
if /i "!arg!"=="-f" set "FORCE=1" & shift & goto parse
if /i "!arg!"=="--delete" set "MODE=delete" & shift & goto parse
if /i "!arg!"=="-d" set "MODE=delete" & shift & goto parse
if "!arg:~0,1!"=="-" echo unknown option: !arg! & exit /b 2
set "ROOT=!arg!"
shift
goto parse
:done_parse

if defined DEL if "!DEL!"=="1" set "FORCE=1"

cd /d "%ROOT%" || exit /b 1

set "N_NEW=0"
set "N_SKIP=0"
set "N_UPD=0"
set "N_BAD=0"
set "N_DEL=0"
set "FAILED=0"

call :all "%BASE%/vae/qwen_image_2.1_vae_bf16.safetensors" models/vae || set "FAILED=1"
call :all "%BASE%/text_encoders/qwen3vl_8b_int8_convrot.safetensors" models/text_encoders || set "FAILED=1"
call :all "%BASE%/text_encoders/qwen3.5_9b_qwen_image_2.1_pe_t2i.int8_convrot.safetensors" models/text_encoders || set "FAILED=1"
call :all "%BASE%/text_encoders/qwen3.5_9b_qwen_image_2.1_pe_i2i.int8_convrot.safetensors" models/text_encoders || set "FAILED=1"
call :all "%BASE%/diffusion_models/qwen_image_2.1_int8_convrot.safetensors" models/diffusion_models || set "FAILED=1"

echo.
if "%MODE%"=="delete" (
    echo summary: !N_DEL! deleted
    exit /b 0
)
echo summary: !N_NEW! installed, !N_UPD! updated, !N_SKIP! already present, !N_BAD! bad
if "%FAILED%"=="1" exit /b 1
exit /b 0

rem ---------- subroutines ----------

:all
set "URL=%~1"
set "DIR=%~2"
set "FILENAME=%URL:/=\%"
for %%F in ("!FILENAME!") do set "FILENAME=%%~nxF"
set "PATH_FULL=%DIR%\%FILENAME%"
set "PART=%PATH_FULL%.part"

if "%MODE%"=="delete" goto :do_remove
goto :do_fetch

:do_remove
echo == !PATH_FULL!
if exist "!PATH_FULL!" (
    for %%A in ("!PATH_FULL!") do echo    deleting !PATH_FULL! (%%~zA bytes)
    del /f /q "!PATH_FULL!" >nul 2>&1
    set /a N_DEL+=1
) else (
    echo    nothing to delete
)
if exist "!PART!" (
    echo    deleting !PART!
    del /f /q "!PART!" >nul 2>&1
)
exit /b 0

:do_fetch
echo == !PATH_FULL!
echo    url:    !URL!

rem Get expected size from server
set "WANT="
for /f "tokens=2 delims=:" %%a in ('curl -sIL "!URL!" 2^>nul ^| findstr /i /c:"x-linked-size:"') do set "WANT=%%a"
if defined WANT (
    set "WANT=!WANT: =!"
    echo    expect: !WANT! bytes
) else (
    echo    expect: unknown, server published no size
)

set "ACTION=install"
if exist "!PATH_FULL!" (
    for %%A in ("!PATH_FULL!") do set "HAVE=%%~zA"
    echo    have:   !HAVE! bytes
    if "!FORCE!"=="1" (
        echo    --force, removing and refetching
        del /f /q "!PATH_FULL!" >nul 2>&1
        set "ACTION=update"
    ) else if not defined WANT (
        echo    warn:   no size to compare against, cannot verify
        set /a N_SKIP+=1
        exit /b 0
    ) else if "!HAVE!"=="!WANT!" (
        echo    ok:     size matches, skipping download
        set /a N_SKIP+=1
        exit /b 0
    ) else (
        if not defined ASSUME_YES (
            echo    size mismatch. Remove and refetch !WANT! bytes? [y/N]
            set /p "ANS="
        ) else if "!ASSUME_YES!"=="1" (
            set "ANS=y"
        ) else (
            set "ANS="
        )
        if /i "!ANS!"=="y" (
            echo    BAD:    size mismatch, removing and refetching
            del /f /q "!PATH_FULL!" >nul 2>&1
            set "ACTION=update"
        ) else (
            echo    keep:   left !PATH_FULL! unchanged
            set /a N_BAD+=1
            exit /b 1
        )
    )
)

echo    fetching to !PART!
if not exist "%DIR%" mkdir "%DIR%" >nul 2>&1
del /f /q "!PART!" >nul 2>&1

set "CURL_OPTS=-fL --progress-bar --retry 5 --retry-delay 3 --retry-all-errors --connect-timeout 30"
if defined HF_TOKEN set "CURL_OPTS=!CURL_OPTS! --location-trusted -H "Authorization: Bearer !HF_TOKEN!""

curl !CURL_OPTS! -o "!PART!" "!URL!" >nul 2>&1
if errorlevel 1 (
    echo    error:  download failed, discarded !PART!
    del /f /q "!PART!" >nul 2>&1
    set /a N_BAD+=1
    exit /b 1
)

for %%A in ("!PART!") do set "HAVE=%%~zA"
echo    have:   !HAVE! bytes
if defined WANT if not "!HAVE!"=="!WANT!" (
    echo    BAD:    size mismatch, discarded !PART!
    del /f /q "!PART!" >nul 2>&1
    set /a N_BAD+=1
    exit /b 1
)

move /y "!PART!" "!PATH_FULL!" >nul 2>&1
if "!ACTION!"=="update" (
    echo    done:   updated !PATH_FULL!
    set /a N_UPD+=1
) else (
    echo    done:   installed !PATH_FULL!
    set /a N_NEW+=1
)
exit /b 0