@echo off
rem onwx-Image (Windows): Python 3.13, the IRON / mlir-aie toolchain (pinned in requirements.txt),
rem the XRT SDK, the Z-Image-Turbo weights and the NPU kernels for every picture size.
rem Requires the AMD NPU driver and Visual Studio C++ Build Tools.
setlocal
cd /d "%~dp0"
set PYTHONUTF8=1

if not exist ".venv\Scripts\python.exe" goto create
fc /b requirements.txt .venv\requirements.done >nul 2>&1 && goto model
goto install
:create
call "%~dp0find_python.bat"
if not defined PY (
    pause
    exit /b 1
)
echo [setup] creating Python environment in .venv ...
%PY% -m venv .venv
if errorlevel 1 (
    echo [setup] could not create the Python environment.
    pause
    exit /b 1
)
:install
".venv\Scripts\python.exe" -m pip install --disable-pip-version-check -r requirements.txt "triton-xdna==3.6.0.2026102404+a414baf" ^
  --find-links https://github.com/amd/Triton-XDNA/releases/expanded_assets/latest-wheels ^
  --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-no-rtti-2 ^
  --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly ^
  --find-links https://github.com/Xilinx/mlir-air/releases/expanded_assets/latest-air-wheels-no-rtti
if not errorlevel 1 goto installed
if defined RETRIED (
    echo [setup] package installation failed.
    pause
    exit /b 1
)
echo [setup] installation failed; rebuilding the Python environment once ...
set RETRIED=1
rmdir /s /q .venv
if not defined PY call "%~dp0find_python.bat"
%PY% -m venv .venv
goto install
:installed
copy /y requirements.txt .venv\requirements.done >nul
".venv\Scripts\python.exe" -m onwx_image.xdna_setup
if errorlevel 1 (
    echo [setup] XRT SDK installation failed.
    pause
    exit /b 1
)

:model
if not exist "C:\Windows\System32\AMD\xrt-smi.exe" (
    echo [setup] note: the AMD NPU driver and XRT are required for NPU inference.
)
echo [setup] Z-Image-Turbo model ...
".venv\Scripts\python.exe" -m onwx_image.fetch_model
if errorlevel 1 (
    echo [setup] the model download failed - run setup.bat again to resume.
    pause
    exit /b 1
)
echo [setup] NPU kernels for every picture size (about 30 minutes the first time; then a minute) ...
".venv\Scripts\python.exe" -m onwx_image.prebuild
if errorlevel 1 (
    echo [setup] building the NPU kernels failed: the app builds them at each size's first picture instead.
)

rem Start menu entry: pythonw -m onwx_image (no console window); the entry of versions before v1.0.0 goes.
set "PYW=%CD%\.venv\Scripts\pythonw.exe"
powershell -NoProfile -Command ^
  "Remove-Item -ErrorAction SilentlyContinue ([Environment]::GetFolderPath('Programs') + '\onw-Image.lnk');" ^
  "$s = (New-Object -ComObject WScript.Shell).CreateShortcut([Environment]::GetFolderPath('Programs') + '\onwx-Image.lnk');" ^
  "$s.TargetPath = '%PYW%'; $s.Arguments = '-m onwx_image'; $s.WorkingDirectory = '%CD%';" ^
  "$s.Description = 'onwx-Image - text-to-image on the AMD Ryzen AI NPU'; $s.IconLocation = '%CD%\onwx_image\web\icon.ico,0'; $s.Save()"
echo [setup] Start menu: onwx-Image

echo.
echo [start] onwx-Image is starting on the AMD Ryzen AI NPU.
start "" "%PYW%" -m onwx_image
timeout /t 8 >nul