@echo off
chcp 65001 >nul 2>&1
setlocal
title ACE-Step
pushd "%~dp0"

REM ============================ НАСТРОЙКИ ============================
REM Каталог виртуального окружения
set "VENV_DIR=venv"
REM Порт веб-интерфейса
set "PORT=7865"
REM Адрес, на котором слушает сервер (0.0.0.0 - доступ из локальной сети)
set "SERVER_NAME=127.0.0.1"
REM Номер видеокарты
set "DEVICE_ID=0"
REM Путь к весам модели. Пусто - скачать в ~/.cache/ace-step/checkpoints
set "CHECKPOINT_PATH="
REM Сборка PyTorch (cu126 / cu124 / cu121 - под вашу версию CUDA)
set "TORCH_INDEX=https://download.pytorch.org/whl/cu126"
REM ==================================================================

REM Имя скрипта надо запомнить до shift: shift сдвигает и %0
set "SCRIPT_NAME=%~nx0"

set "OPT_LOWVRAM=0"
set "OPT_CPU=0"
set "OPT_SHARE=0"
set "DO_SETUP_ONLY=0"
set "DO_REINSTALL=0"
set "DO_UPDATE=0"

set "FROM_EXPLORER=0"
echo %cmdcmdline% | find /i "%~nx0" >nul 2>&1 && set "FROM_EXPLORER=1"

REM ------------------------- разбор аргументов -------------------------
:parse_args
if "%~1"=="" goto args_done
if /i "%~1"=="--help"      goto usage
if /i "%~1"=="-h"          goto usage
if /i "%~1"=="/?"          goto usage
if /i "%~1"=="--lowvram"   ( set "OPT_LOWVRAM=1"  & shift & goto parse_args )
if /i "%~1"=="--cpu"       ( set "OPT_CPU=1"      & shift & goto parse_args )
if /i "%~1"=="--share"     ( set "OPT_SHARE=1"    & shift & goto parse_args )
if /i "%~1"=="--setup"     ( set "DO_SETUP_ONLY=1" & shift & goto parse_args )
if /i "%~1"=="--reinstall" ( set "DO_REINSTALL=1" & shift & goto parse_args )
if /i "%~1"=="--update"    ( set "DO_UPDATE=1"    & shift & goto parse_args )
if /i "%~1"=="--listen"    ( set "SERVER_NAME=0.0.0.0" & shift & goto parse_args )
if /i "%~1"=="--port"      ( set "PORT=%~2"       & shift & shift & goto parse_args )
if /i "%~1"=="--device"    ( set "DEVICE_ID=%~2"  & shift & shift & goto parse_args )
echo [ОШИБКА] Неизвестный аргумент: %~1
echo Запустите "%SCRIPT_NAME% --help" для справки.
goto fail
:args_done

echo.
echo ==========================================
echo   ACE-Step - генерация музыки
echo ==========================================
echo.

set "VENV_PY=%CD%\%VENV_DIR%\Scripts\python.exe"

if "%DO_REINSTALL%"=="1" (
    if exist "%VENV_DIR%\" (
        echo [1/4] Удаляю старое окружение "%VENV_DIR%"...
        rmdir /s /q "%VENV_DIR%"
    )
)

REM ------------------------- поиск Python -------------------------
if exist "%VENV_PY%" goto venv_ready

echo [1/4] Ищу подходящий Python...
set "SYS_PY="
for %%V in (3.12 3.11 3.10 3.13) do (
    if not defined SYS_PY (
        py -%%V -c "import sys" >nul 2>&1 && set "SYS_PY=py -%%V"
    )
)
if not defined SYS_PY (
    python -c "import sys; sys.exit(0 if sys.version_info >= (3,10) else 1)" >nul 2>&1 && set "SYS_PY=python"
)
if not defined SYS_PY (
    echo.
    echo [ОШИБКА] Не найден Python 3.10 или новее.
    echo Установите его с https://www.python.org/downloads/
    echo и обязательно отметьте галочку "Add Python to PATH".
    goto fail
)
for /f "delims=" %%O in ('%SYS_PY% -c "import sys;print(sys.version.split()[0])" 2^>nul') do set "PYVER=%%O"
echo       Найден Python %PYVER% (%SYS_PY%)

echo [1/4] Создаю виртуальное окружение в "%VENV_DIR%"...
%SYS_PY% -m venv "%VENV_DIR%"
if errorlevel 1 (
    echo [ОШИБКА] Не удалось создать виртуальное окружение.
    goto fail
)
if not exist "%VENV_PY%" (
    echo [ОШИБКА] Виртуальное окружение создано некорректно: нет "%VENV_PY%".
    goto fail
)
set "FRESH_VENV=1"

:venv_ready
if not defined FRESH_VENV echo [1/4] Виртуальное окружение найдено: %VENV_DIR%

REM ------------------------- зависимости -------------------------
echo [2/4] Проверяю зависимости...

set "NEED_TORCH=0"
"%VENV_PY%" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('torch') else 1)" >nul 2>&1 || set "NEED_TORCH=1"
REM Проверяем именно установку пакета: сам каталог acestep лежит рядом со скриптом,
REM поэтому find_spec('acestep') сработал бы даже без установленных зависимостей.
set "NEED_ACESTEP=0"
"%VENV_PY%" -c "import importlib.metadata as md, importlib.util as u, sys; md.version('ace_step'); sys.exit(0 if u.find_spec('click') and u.find_spec('gradio') else 1)" >nul 2>&1 || set "NEED_ACESTEP=1"

if "%DO_UPDATE%"=="1" set "NEED_TORCH=1"
if "%DO_UPDATE%"=="1" set "NEED_ACESTEP=1"

if "%NEED_TORCH%%NEED_ACESTEP%"=="00" goto deps_ready

echo       Обновляю pip...
"%VENV_PY%" -m pip install --upgrade pip setuptools wheel --quiet
if errorlevel 1 (
    echo [ОШИБКА] Не удалось обновить pip. Проверьте подключение к интернету.
    goto fail
)

if "%NEED_TORCH%"=="1" (
    if "%OPT_CPU%"=="1" (
        echo       Устанавливаю PyTorch для CPU. Это займёт несколько минут...
        "%VENV_PY%" -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
    ) else (
        echo       Устанавливаю PyTorch с поддержкой CUDA. Это займёт несколько минут...
        "%VENV_PY%" -m pip install --upgrade torch torchvision torchaudio --index-url %TORCH_INDEX%
    )
    if errorlevel 1 (
        echo [ОШИБКА] Не удалось установить PyTorch.
        echo Если у вас другая версия CUDA, поменяйте TORCH_INDEX в начале этого файла.
        goto fail
    )
)

if "%NEED_ACESTEP%"=="1" (
    echo       Устанавливаю ACE-Step и зависимости...
    "%VENV_PY%" -m pip install -e .
    if errorlevel 1 (
        echo [ОШИБКА] Не удалось установить ACE-Step.
        goto fail
    )
)

:deps_ready

if "%OPT_LOWVRAM%"=="1" (
    "%VENV_PY%" -c "import importlib.util,sys; sys.exit(0 if importlib.util.find_spec('triton') else 1)" >nul 2>&1 || (
        echo       Режим экономии VRAM: устанавливаю triton-windows...
        "%VENV_PY%" -m pip install triton-windows --quiet
        if errorlevel 1 echo [ВНИМАНИЕ] triton-windows не установился, --torch_compile может не заработать.
    )
)

REM ------------------------- проверка GPU -------------------------
echo [3/4] Проверяю видеокарту...
if "%OPT_CPU%"=="1" goto gpu_cpu_mode

set "BF16=true"
set "GPUINFO="
set "GPUTMP=%TEMP%\acestep_gpu_%RANDOM%.txt"
"%VENV_PY%" -c "import torch; print((torch.cuda.get_device_name(0) + ' - ' + str(round(torch.cuda.get_device_properties(0).total_memory/1073741824, 1)) + ' GB VRAM') if torch.cuda.is_available() else 'NO_CUDA')" > "%GPUTMP%" 2>nul
if exist "%GPUTMP%" set /p GPUINFO=<"%GPUTMP%"
del "%GPUTMP%" >nul 2>&1
if not defined GPUINFO set "GPUINFO=NO_CUDA"
if not "%GPUINFO%"=="NO_CUDA" goto gpu_ok
echo.
echo [ВНИМАНИЕ] CUDA недоступна - модель будет работать на процессоре ^(очень медленно^).
echo Если у вас есть видеокарта NVIDIA, обновите драйвер и переустановите PyTorch:
echo     "%SCRIPT_NAME%" --update
echo.
goto gpu_done

:gpu_ok
echo       %GPUINFO%
goto gpu_done

:gpu_cpu_mode
echo       Принудительный режим CPU. Генерация будет очень медленной.
set "DEVICE_ID=-1"
set "BF16=false"

:gpu_done

if "%DO_SETUP_ONLY%"=="1" (
    echo.
    echo Установка завершена. Для запуска выполните "%SCRIPT_NAME%".
    goto success
)

REM ------------------------- запуск -------------------------
set "ARGS=--port %PORT% --server_name %SERVER_NAME% --device_id=%DEVICE_ID% --bf16 %BF16%"
if not "%CHECKPOINT_PATH%"=="" set ARGS=%ARGS% --checkpoint_path "%CHECKPOINT_PATH%"
if "%OPT_SHARE%"=="1"   set "ARGS=%ARGS% --share true"
if "%OPT_LOWVRAM%"=="1" set "ARGS=%ARGS% --cpu_offload true --overlapped_decode true --torch_compile true"

echo [4/4] Запускаю веб-интерфейс...
echo.
echo       Адрес: http://%SERVER_NAME%:%PORT%
if "%OPT_LOWVRAM%"=="1" echo       Режим: экономия видеопамяти
echo       Первый запуск скачивает веса модели (~8 ГБ) - наберитесь терпения.
echo       Остановить: Ctrl+C в этом окне.
echo.

"%VENV_PY%" -m acestep.gui %ARGS%
if errorlevel 1 (
    echo.
    echo [ОШИБКА] ACE-Step завершился с ошибкой. Текст ошибки - выше.
    goto fail
)
goto success

REM ------------------------- справка -------------------------
:usage
echo.
echo Использование: %SCRIPT_NAME% [флаги]
echo.
echo   --lowvram      Режим экономии видеопамяти (~8 ГБ VRAM)
echo   --cpu          Запуск на процессоре, без CUDA (очень медленно)
echo   --share        Публичная ссылка Gradio
echo   --port ^<N^>     Порт веб-интерфейса (по умолчанию %PORT%)
echo   --device ^<N^>   Номер видеокарты (по умолчанию %DEVICE_ID%)
echo   --listen       Слушать 0.0.0.0 (доступ из локальной сети)
echo   --reinstall    Пересоздать виртуальное окружение с нуля
echo   --update       Обновить зависимости
echo   --setup        Только установка, без запуска
echo   --help         Эта справка
echo.
echo Настройки по умолчанию - в блоке НАСТРОЙКИ в начале файла.
echo.
goto success

REM ------------------------- завершение -------------------------
:fail
echo.
pause
popd
endlocal
exit /b 1

:success
if "%FROM_EXPLORER%"=="1" pause
popd
endlocal
exit /b 0