@echo off
setlocal enabledelayedexpansion

echo FramePack Installer
echo =================
echo.

REM Check if Python is installed
python --version >nul 2>&1
if %errorlevel% neq 0 (
    echo Python is not installed or not in PATH.
    echo Please install Python 3.10+ from https://www.python.org/downloads/
    echo Make sure to check "Add Python to PATH" during installation.
    echo.
    pause
    exit /b 1
)

REM Check Python version (need 3.10+)
for /f "tokens=2" %%I in ('python --version 2^>^&1') do (
    set pyver=%%I
)
for /f "tokens=1,2 delims=." %%a in ("%pyver%") do (
    set major=%%a
    set minor=%%b
)
if %major% LSS 3 (
    echo Python version %pyver% is too old. FramePack requires Python 3.10+
    pause
    exit /b 1
)
if %major% EQU 3 if %minor% LSS 10 (
    echo Python version %pyver% is too old. FramePack requires Python 3.10+
    pause
    exit /b 1
)

REM Set up repo directory
set REPO_DIR=%~dp0FramePack
if not exist "%REPO_DIR%" (
    echo Cloning FramePack repository...
    git clone https://github.com/lllyasviel/FramePack.git "%REPO_DIR%"
    if %errorlevel% neq 0 (
        echo Failed to clone repository. Make sure git is installed.
        echo Install git from https://git-scm.com/downloads
        pause
        exit /b 1
    )
) else (
    echo FramePack repository already exists.
    echo Updating repository...
    cd "%REPO_DIR%"
    git pull
    cd ..
)

REM Create virtual environment
set VENV_DIR=%~dp0framepack_venv
if not exist "%VENV_DIR%" (
    echo Creating virtual environment...
    python -m venv "%VENV_DIR%"
    if %errorlevel% neq 0 (
        echo Failed to create virtual environment.
        pause
        exit /b 1
    )
)

REM Activate virtual environment and install dependencies
echo Installing dependencies...
call "%VENV_DIR%\Scripts\activate.bat"

REM Install PyTorch with CUDA support (using available versions)
echo Installing PyTorch and dependencies...
python -m pip install --upgrade pip

echo Installing PyTorch packages (latest available for CUDA 12.6)...
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

echo Installing common dependencies...
pip install numpy pillow

REM Check if torchvision is installed correctly
python -c "import torchvision; print('Torchvision version:', torchvision.__version__)"
if %errorlevel% neq 0 (
    echo Warning: Torchvision installation verification failed.
    echo Attempting alternative installation...
    pip install torchvision --force-reinstall --no-deps --index-url https://download.pytorch.org/whl/cu126
    python -c "import torchvision"
    if %errorlevel% neq 0 (
        echo Failed to install torchvision. Please check your Python environment.
        pause
        exit /b 1
    )
)

REM Install requirements
cd "%REPO_DIR%"

echo Checking for requirements.txt...
if not exist requirements.txt (
    echo Warning: requirements.txt not found in the repository!
    echo Installing commonly needed packages instead...
) else (
    echo Installing requirements from requirements.txt...
    pip install -r requirements.txt
    if %errorlevel% neq 0 (
        echo Warning: Some packages in requirements.txt could not be installed.
        echo Continuing with installation of key packages...
    )
)

REM Install common dependencies that might be needed
echo Installing additional dependencies...
pip install --upgrade gradio diffusers transformers safetensors
pip install --upgrade accelerate einops kornia opencv-python
pip install --upgrade huggingface_hub
pip install --upgrade scipy scikit-image scikit-learn
pip install --upgrade ftfy

REM Try to install xformers if possible
echo Attempting to install xformers...
pip install --upgrade xformers --index-url https://download.pytorch.org/whl/cu126
if %errorlevel% neq 0 (
    echo xformers installation failed, but this is optional.
    echo FramePack will still work without it.
)


REM Create a launcher batch file
set LAUNCHER=%~dp0run_framepack.bat

echo @echo off > "%LAUNCHER%"
echo setlocal enabledelayedexpansion >> "%LAUNCHER%"
echo title FramePack >> "%LAUNCHER%"
echo. >> "%LAUNCHER%"
echo echo ========================================== >> "%LAUNCHER%"
echo echo             FramePack Launcher            >> "%LAUNCHER%"
echo echo ========================================== >> "%LAUNCHER%"
echo echo. >> "%LAUNCHER%"

echo REM Activate virtual environment >> "%LAUNCHER%"
echo call "%VENV_DIR%\Scripts\activate.bat" >> "%LAUNCHER%"
echo. >> "%LAUNCHER%"

echo REM Basic dependency check >> "%LAUNCHER%"
echo echo Checking dependencies... >> "%LAUNCHER%"
echo python -c "import torch" 2^>nul >> "%LAUNCHER%"
echo if %%errorlevel%% neq 0 ( >> "%LAUNCHER%"
echo   echo Error: PyTorch is not installed correctly. >> "%LAUNCHER%"
echo   echo Please run install_framepack.bat again to fix dependencies. >> "%LAUNCHER%"
echo   pause >> "%LAUNCHER%"
echo   exit /b 1 >> "%LAUNCHER%"
echo ) >> "%LAUNCHER%"
echo. >> "%LAUNCHER%"

echo REM Check for CUDA availability >> "%LAUNCHER%"
echo python -c "import torch; print('CUDA Available:', torch.cuda.is_available()); print('CUDA Version:', torch.version.cuda if torch.cuda.is_available() else 'N/A'); print('GPU Model:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'No GPU detected')" >> "%LAUNCHER%"
echo if %%errorlevel%% neq 0 ( >> "%LAUNCHER%"
echo   echo Warning: Could not verify CUDA availability. >> "%LAUNCHER%"
echo ) >> "%LAUNCHER%"
echo. >> "%LAUNCHER%"

echo cd "%REPO_DIR%" >> "%LAUNCHER%"
echo echo. >> "%LAUNCHER%"
echo echo Starting FramePack... >> "%LAUNCHER%"
echo echo. >> "%LAUNCHER%"
echo echo If the application fails to start, check these common issues: >> "%LAUNCHER%"
echo echo  - Make sure you have an NVIDIA GPU with CUDA support >> "%LAUNCHER%"
echo echo  - At least 6GB of VRAM is recommended >> "%LAUNCHER%"
echo echo  - Try running the installer again if modules are missing >> "%LAUNCHER%"
echo echo. >> "%LAUNCHER%"

echo python demo_gradio.py %%* >> "%LAUNCHER%"
echo. >> "%LAUNCHER%"
echo if %%errorlevel%% neq 0 ( >> "%LAUNCHER%"
echo   echo. >> "%LAUNCHER%"
echo   echo FramePack exited with an error. >> "%LAUNCHER%"
echo   echo. >> "%LAUNCHER%"
echo   echo Possible solutions: >> "%LAUNCHER%"
echo   echo  1. Try reinstalling by running install_framepack.bat >> "%LAUNCHER%"
echo   echo  2. Check if your GPU meets the requirements >> "%LAUNCHER%"
echo   echo  3. For more help, visit: https://github.com/lllyasviel/FramePack >> "%LAUNCHER%"
echo ) else ( >> "%LAUNCHER%"
echo   echo. >> "%LAUNCHER%"
echo   echo FramePack closed successfully. >> "%LAUNCHER%"
echo ) >> "%LAUNCHER%"
echo. >> "%LAUNCHER%"
echo pause >> "%LAUNCHER%"

REM Create a minimal test script to verify imports
set TEST_SCRIPT=%TEMP%\test_imports.py
echo import torch > "%TEST_SCRIPT%"
echo import torchvision >> "%TEST_SCRIPT%"
echo print("PyTorch version:", torch.__version__) >> "%TEST_SCRIPT%"
echo print("Torchvision version:", torchvision.__version__) >> "%TEST_SCRIPT%"
echo try: >> "%TEST_SCRIPT%"
echo     import diffusers >> "%TEST_SCRIPT%"
echo     print("Diffusers version:", diffusers.__version__) >> "%TEST_SCRIPT%"
echo except ImportError: >> "%TEST_SCRIPT%"
echo     print("WARNING: diffusers not installed") >> "%TEST_SCRIPT%"
echo try: >> "%TEST_SCRIPT%"
echo     import transformers >> "%TEST_SCRIPT%"
echo     print("Transformers version:", transformers.__version__) >> "%TEST_SCRIPT%"
echo except ImportError: >> "%TEST_SCRIPT%"
echo     print("WARNING: transformers not installed") >> "%TEST_SCRIPT%"
echo try: >> "%TEST_SCRIPT%"
echo     import gradio >> "%TEST_SCRIPT%"
echo     print("Gradio version:", gradio.__version__) >> "%TEST_SCRIPT%"
echo except ImportError: >> "%TEST_SCRIPT%"
echo     print("WARNING: gradio not installed") >> "%TEST_SCRIPT%"

echo.
echo Verifying installation...
python "%TEST_SCRIPT%"
set VERIFY_RESULT=%errorlevel%

echo.
if %VERIFY_RESULT% neq 0 (
    echo ⚠️ Warning: Some dependencies may not be installed correctly.
    echo.
    echo If you encounter issues running FramePack, try these steps:
    echo  1. Make sure you have an NVIDIA GPU with CUDA support
    echo  2. Try running the installer again
    echo  3. If problems persist, install packages manually:
    echo     pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126
    echo     pip install gradio diffusers transformers
    echo.
) else (
    echo ✅ Core dependencies verified successfully!
)

echo Installation complete!
echo.
echo To run FramePack, use the run_framepack.bat file or run:
echo   %LAUNCHER%
echo.
echo You can pass additional parameters like --share for public link:
echo   %LAUNCHER% --share
echo.
echo If you encounter any issues, please check the FramePack GitHub repository:
echo   https://github.com/lllyasviel/FramePack
echo.
pause