@echo off
chcp 65001 >nul 2>&1
setlocal EnableDelayedExpansion

title SDXL Long Context — Learned Sinusoidal

set "PYTHON=..\python_embeded\python.exe"
set "SCRIPT=convert_learned.py"
set "CKPT_DIR=..\ComfyUI\models\checkpoints"

:: ============================================================
::  SETTINGS
:: ============================================================
set "CKPT_NAME=FernflowerAI-2D.safetensors"

:: Method: learned_sinusoidal / vaswani_fitted / hybrid
::
:: learned_sinusoidal — fit N sinusoids to each dimension
:: vaswani_fitted     — fit original Transformer PE formula
:: hybrid             — best of both + spectral fallback (recommended)
set "METHOD=hybrid"

:: Number of sinusoidal components (for learned_sinusoidal)
set "N_COMPONENTS=3"

:: Blend width at boundary
set "BLEND_WIDTH=12"

:: Target length
set "TARGET_LEN=248"
:: ============================================================

set "INPUT=%CKPT_DIR%\%CKPT_NAME%"
for %%A in ("%CKPT_NAME%") do set "BASE=%%~nA"
set "OUTPUT=%CKPT_DIR%\%BASE%_long%TARGET_LEN%.safetensors"

echo.
echo  ================================================================
echo   SDXL Long Context — Learned Sinusoidal Extrapolation
echo   Extend CLIP context: 77 -^> %TARGET_LEN% tokens
echo   Method: %METHOD%
echo   Uses the same math as CLIP training initialization
echo  ================================================================
echo.

if not exist "%PYTHON%" (
    echo  [X] Python not found: %PYTHON%
    echo      Run this from the scripts folder
    goto :fail
)
if not exist "%SCRIPT%" (
    echo  [X] Script not found: %SCRIPT%
    goto :fail
)
if not exist "%INPUT%" (
    echo  [X] Checkpoint not found: %CKPT_NAME%
    echo.
    echo      Available checkpoints:
    for %%F in ("%CKPT_DIR%\*.safetensors") do echo        %%~nxF
    echo.
    echo      Edit CKPT_NAME in this bat file
    goto :fail
)

%PYTHON% -c "import safetensors" >nul 2>&1
if errorlevel 1 (
    echo  Installing safetensors...
    %PYTHON% -m pip install safetensors
)
%PYTHON% -c "import numpy" >nul 2>&1
if errorlevel 1 (
    echo  Installing numpy...
    %PYTHON% -m pip install numpy
)
%PYTHON% -c "import scipy" >nul 2>&1
if errorlevel 1 (
    echo  Installing scipy...
    %PYTHON% -m pip install scipy
)

echo  Input:      %CKPT_NAME%
echo  Output:     %BASE%_long%TARGET_LEN%.safetensors
echo  Method:     %METHOD%
echo.

%PYTHON% %SCRIPT% ^
    --input "%INPUT%" ^
    --output "%OUTPUT%" ^
    --target_len %TARGET_LEN% ^
    --method %METHOD% ^
    --n_components %N_COMPONENTS% ^
    --blend_width %BLEND_WIDTH% ^
    --analyze

if errorlevel 1 goto :fail

echo.
echo  ================================================================
echo   DONE!
echo.
echo   Output: %BASE%_long%TARGET_LEN%.safetensors
echo.
echo   Method: %METHOD%
echo   - Fits the mathematical model used in CLIP training
echo   - Original 77 positions preserved
echo   - Extended positions follow learned frequency patterns
echo.
echo   Use "Load Checkpoint (SDXL Long 248)" node in ComfyUI.
echo  ================================================================
echo.
pause
exit /b 0

:fail
echo.
echo  [X] Error! Check output above.
echo.
pause
exit /b 1