diff --git a/.gitignore b/.gitignore
index 4631ea00..b2db6d29 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,7 @@ queue_images/
modules/toolbox/model_esrgan/
modules/toolbox/model_rife/
.framepack/
+.claude/
modules/toolbox/data/
modules/toolbox/bin
queue.json
diff --git a/INSTALLATION_NOTES.md b/INSTALLATION_NOTES.md
new file mode 100644
index 00000000..ae95a2fb
--- /dev/null
+++ b/INSTALLATION_NOTES.md
@@ -0,0 +1,139 @@
+# FramePack Studio Installation Notes
+
+## Installation Summary
+
+**Installation Method:** Fresh venv created and packages installed via automated process
+**Date:** December 29, 2025
+**Python Version:** 3.13
+
+### Installed Packages
+
+**Core AI/ML Packages:**
+- PyTorch: 2.8.0+cu128
+- TorchVision: 0.23.0+cu128
+- TorchAudio: 2.8.0+cu128
+- Diffusers: 0.36.0
+- Transformers: 4.57.3
+- Accelerate: 1.12.0
+- PEFT: 0.18.0
+
+**UI & Web:**
+- Gradio: 6.2.0
+
+**Media Processing:**
+- OpenCV Contrib Python: 4.12.0.88
+- av: 16.0.1
+- imageio: 2.37.2
+- imageio-ffmpeg: 0.6.0
+- decord: 0.6.0
+
+**Upscaling/Enhancement:**
+- facexlib: 0.3.0
+- gfpgan: 1.3.8
+- realesrgan: 0.3.0
+
+**Utilities:**
+- sentencepiece: 0.2.1
+- torchsde: 0.2.6
+- scipy: 1.16.3
+- numpy: 2.2.6
+
+### CUDA Configuration
+- CUDA Version: 12.8
+- Optimized for RTX 40xx and 50xx GPUs
+
+### Known Issues
+
+#### BasicSR Not Installed
+BasicSR has compatibility issues with Python 3.13 and could not be installed.
+
+**Impact:**
+- Some advanced toolbox features (ESRGAN, GFPGAN, RealESRGAN) may have limited functionality
+- Main FramePack Studio features should work normally
+
+**Workaround:**
+1. Use Python 3.10, 3.11, or 3.12 if BasicSR features are critical
+2. Wait for BasicSR to release Python 3.13 compatible wheels
+
+**Alternative:** You can try installing BasicSR manually later if needed:
+```bash
+venv\Scripts\activate
+pip install basicsr --no-build-isolation
+```
+
+#### NumPy Version
+- Requirements specify numpy==1.26.2 but numpy 2.2.6 was installed
+- This is a newer version with better Python 3.13 support
+- Most packages are compatible with NumPy 2.x
+
+### Optional Acceleration Packages
+
+The install script offers optional acceleration packages:
+
+**Sage Attention:**
+- Requires: triton-windows<3.4
+- Significant speed improvements for RTX 40xx/50xx
+- Pre-built wheels available for Python 3.10-3.12
+
+**Flash Attention:**
+- Alternative acceleration method
+- Pre-built wheels available
+
+**Note:** These may not be available for Python 3.13 yet.
+
+### Scripts Created
+
+1. **install_40xx_50xx.bat** - Full installation script for RTX 40xx/50xx GPUs
+2. **start.bat** - Launches FramePack Studio with venv activation
+3. **activate_venv.bat** - Activates venv for manual commands
+
+### Running the Application
+
+```bash
+# Option 1: Use the start script
+start.bat
+
+# Option 2: Use the existing run script
+run.bat
+
+# Option 3: Manual
+venv\Scripts\activate
+python studio.py
+```
+
+### Verifying Installation
+
+```bash
+venv\Scripts\activate
+python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA: {torch.cuda.is_available()}')"
+```
+
+Expected output:
+```
+PyTorch: 2.8.0+cu128
+CUDA: True
+```
+
+## Troubleshooting
+
+### If PyTorch CUDA is not detected:
+1. Ensure NVIDIA drivers are up to date
+2. Verify nvidia-smi shows your GPU
+3. Reinstall PyTorch with CUDA 12.8
+
+### If packages are missing:
+```bash
+venv\Scripts\activate
+pip install -r requirements.txt
+```
+
+### Clean Reinstall:
+1. Delete the `venv` folder
+2. Run `install_40xx_50xx.bat` again
+
+## Python Version Recommendation
+
+**Recommended:** Python 3.10, 3.11, or 3.12
+**Current:** Python 3.13 (newer, some packages may have limited compatibility)
+
+If you experience issues, consider using Python 3.12 for maximum compatibility.
diff --git a/activate_venv.bat b/activate_venv.bat
new file mode 100644
index 00000000..be8e8353
--- /dev/null
+++ b/activate_venv.bat
@@ -0,0 +1,24 @@
+@echo off
+echo ============================================
+echo Activating FramePack-Studio Virtual Environment
+echo ============================================
+echo.
+
+REM Check if venv exists
+if not exist "%cd%\venv\Scripts\activate.bat" (
+ echo Error: Virtual environment not found!
+ echo Please run install_40xx_50xx.bat first to set up the environment.
+ echo.
+ pause
+ exit /b 1
+)
+
+echo Virtual environment activated.
+echo You can now run Python commands within this environment.
+echo.
+echo To deactivate, type: deactivate
+echo To run FramePack-Studio, type: python studio.py
+echo.
+
+REM Activate and keep command prompt open
+cmd /k "%cd%\venv\Scripts\activate.bat"
diff --git a/diffusers_helper/models/hunyuan_video_packed.py b/diffusers_helper/models/hunyuan_video_packed.py
index 29695620..3cefec0d 100644
--- a/diffusers_helper/models/hunyuan_video_packed.py
+++ b/diffusers_helper/models/hunyuan_video_packed.py
@@ -63,7 +63,7 @@
has_xformers = xformers_attn_func is not None
if has_sage:
- print("✅ Using SAGE Attention (highest performance).")
+ print("[OK] Using SAGE Attention (highest performance).")
ignored = []
if has_flash:
ignored.append("Flash Attention")
@@ -72,16 +72,16 @@
if ignored:
print(f" - Ignoring other installed attention libraries: {', '.join(ignored)}")
elif has_flash:
- print("✅ Using Flash Attention (high performance).")
+ print("[OK] Using Flash Attention (high performance).")
if has_xformers:
print(" - Consider installing SAGE Attention for highest performance.")
print(" - Ignoring other installed attention library: xFormers")
elif has_xformers:
- print("✅ Using xFormers.")
+ print("[OK] Using xFormers.")
print(" - Consider installing SAGE Attention for highest performance.")
print(" - or Consider installing Flash Attention for high performance.")
else:
- print("⚠️ No attention library found. Using native PyTorch Scaled Dot Product Attention.")
+ print("[WARNING] No attention library found. Using native PyTorch Scaled Dot Product Attention.")
print(" - For better performance, consider installing one of:")
print(" SAGE Attention (highest performance), Flash Attention (high performance), or xFormers.")
print("-------------------------------\n")
diff --git a/install.bat b/install.bat
index e4936ba7..e678e0c3 100644
--- a/install.bat
+++ b/install.bat
@@ -1,208 +1,208 @@
-@echo off
-echo FramePack-Studio Setup Script
-setlocal enabledelayedexpansion
-
-REM Check if Python is installed (basic check)
-where python >nul 2>&1
-if %errorlevel% neq 0 (
- echo Error: Python is not installed or not in your PATH. Please install Python and try again.
- goto end
-)
-
-if exist "%cd%/venv" (
-echo Virtual Environment already exists.
-set /p choice= "Do you want to reinstall packages?[Y/N]: "
-
-if "!choice!" == "y" (goto checkgpu)
-if "!choice!"=="Y" (goto checkgpu)
-
-goto end
-)
-
-REM Check the python version
-echo Python versions 3.10-3.12 have been confirmed to work. Other versions are currently not supported. You currently have:
-python -V
-set choice=
-set /p choice= "Do you want to continue?[Y/N]: "
-
-
-if "!choice!" == "y" (goto makevenv)
-if "!choice!"=="Y" (goto makevenv)
-
-goto end
-
-:makevenv
-REM This creates a virtual environment in the folder
-echo Creating a Virtual Environment...
-python -m venv venv
-echo Upgrading pip in Virtual Environment to lower chance of error...
-"%cd%/venv/Scripts/python.exe" -m pip install --upgrade pip
-
-:checkgpu
-REM ask Windows for GPU
-where nvidia-smi >nul 2>&1
-if %errorlevel% neq 0 (
- echo Error: Nvidia GPU doesn't exist or drivers installed incorrectly. Please confirm your drivers are installed.
- goto end
-)
-
-echo Checking your GPU...
-
-for /F "tokens=* skip=1" %%n in ('nvidia-smi --query-gpu=name') do set GPU_NAME=%%n && goto gpuchecked
-
-:gpuchecked
-echo Detected %GPU_NAME%
-set "GPU_SERIES=%GPU_NAME:*RTX =%"
-set "GPU_SERIES=%GPU_SERIES:~0,2%00"
-
-REM This gets the shortened Python version for later use. e.g. 3.10.13 becomes 310.
-for /f "delims=" %%A in ('python -V') do set "pyv=%%A"
-for /f "tokens=2 delims= " %%A in ("%pyv%") do (
- set pyv=%%A
-)
-set pyv=%pyv:.=%
-set pyv=%pyv:~0,3%
-
-echo Installing torch...
-
-if !GPU_SERIES! geq 5000 (
- goto torch270
-) else (
- goto torch260
-)
-
-REM RTX 5000 Series
-:torch270
-"%cd%/venv/Scripts/pip.exe" install torch==2.7.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 --force-reinstall
-REM Check if pip installation was successful
-if %errorlevel% neq 0 (
- echo Warning: Failed to install dependencies. You may need to install them manually.
- goto end
-)
-
-REM Ask if user wants Sage Attention
-set choice=
-echo Do you want to install any of the following? They speed up generation.
-echo 1) Sage Attention
-echo 2) Flash Attention
-echo 3) BOTH!
-echo 4) No
-set /p choice= "Input Selection: "
-
-set both="N"
-
-if "!choice!" == "1" (goto triton270)
-if "!choice!"== "2" (goto flash270)
-if "!choice!"== "3" (set both="Y"
-goto triton270
-)
-
-goto requirements
-
-:triton270
-REM Sage Attention and Triton for Torch 2.7.0
-"%cd%/venv/Scripts/pip.exe" install "triton-windows<3.4" --force-reinstall
-"%cd%/venv/Scripts/pip.exe" install "https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu128torch2.7.0-cp%pyv%-cp%pyv%-win_amd64.whl" --force-reinstall
-echo Finishing up installing triton-windows. This requires extraction of libraries into Python Folder...
-
-REM Check for python version and download the triton-windows required libs accordingly
-if %pyv% == 310 (
- powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.10.11_include_libs.zip', 'triton-lib.zip')"
-)
-
-if %pyv% == 311 (
- powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.11.9_include_libs.zip', 'triton-lib.zip')"
-)
-
-if %pyv% == 312 (
- powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.12.7_include_libs.zip', 'triton-lib.zip')"
-)
-
-REM Extract the zip into the Python Folder and Delete zip
-powershell Expand-Archive -Path '%cd%\triton-lib.zip' -DestinationPath '%cd%\venv\Scripts\' -force
-del triton-lib.zip
-if %both% == "Y" (goto flash270)
-
-goto requirements
-
-:flash270
-REM Install flash-attn.
-"%cd%/venv/Scripts/pip.exe" install "https://huggingface.co/lldacing/flash-attention-windows-wheel/resolve/main/flash_attn-2.7.4.post1%%2Bcu128torch2.7.0cxx11abiFALSE-cp%pyv%-cp%pyv%-win_amd64.whl?download=true"
-goto requirements
-
-
-REM RTX 4000 Series and below
-:torch260
-"%cd%/venv/Scripts/pip.exe" install torch==2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 --force-reinstall
-REM Check if pip installation was successful
-if %errorlevel% neq 0 (
- echo Warning: Failed to install dependencies. You may need to install them manually.
- goto end
-)
-
-REM Ask if user wants Sage Attention
-set choice=
-echo Do you want to install any of the following? They speed up generation.
-echo 1) Sage Attention
-echo 2) Flash Attention
-echo 3) BOTH!
-echo 4) No
-set /p choice= "Input Selection: "
-
-set both="N"
-
-if "!choice!" == "1" (goto triton260)
-if "!choice!"== "2" (goto flash260)
-if "!choice!"== "3" (set both="Y"
-goto triton260)
-
-goto requirements
-
-:triton260
-REM Sage Attention and Triton for Torch 2.6.0
-"%cd%/venv/Scripts/pip.exe" install "triton-windows<3.3.0" --force-reinstall
-"%cd%/venv/Scripts/pip.exe" install https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu126torch2.6.0-cp%pyv%-cp%pyv%-win_amd64.whl --force-reinstall
-
-echo Finishing up installing triton-windows. This requires extraction of libraries into Python Folder...
-
-REM Check for python version and download the triton-windows required libs accordingly
-if %pyv% == 310 (
- powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.10.11_include_libs.zip', 'triton-lib.zip')"
-)
-
-if %pyv% == 311 (
- powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.11.9_include_libs.zip', 'triton-lib.zip')"
-)
-
-if %pyv% == 312 (
- powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.12.7_include_libs.zip', 'triton-lib.zip')"
-)
-
-REM Extract the zip into the Python Folder and Delete zip
-powershell Expand-Archive -Path '%cd%\triton-lib.zip' -DestinationPath '%cd%\venv\Scripts\' -force
-del triton-lib.zip
-
-if %both% == "Y" (goto flash260)
-
-goto requirements
-
-:flash260
-REM Install flash-attn.
-"%cd%/venv/Scripts/pip.exe" install "https://huggingface.co/lldacing/flash-attention-windows-wheel/resolve/main/flash_attn-2.7.4%%2Bcu126torch2.6.0cxx11abiFALSE-cp%pyv%-cp%pyv%-win_amd64.whl?download=true"
-
-:requirements
-echo Installing remaining required packages through pip...
-REM This assumes there's a requirements.txt file in the root
-"%cd%/venv/Scripts/pip.exe" install -r requirements.txt
-
-REM Check if pip installation was successful
-if %errorlevel% neq 0 (
- echo Warning: Failed to install dependencies. You may need to install them manually.
- goto end
-)
-
-echo Setup complete.
-
-:end
-echo Exiting setup script.
-pause
+@echo off
+echo FramePack-Studio Setup Script
+setlocal enabledelayedexpansion
+
+REM Check if Python is installed (basic check)
+where python >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Error: Python is not installed or not in your PATH. Please install Python and try again.
+ goto end
+)
+
+if exist "%cd%/venv" (
+echo Virtual Environment already exists.
+set /p choice= "Do you want to reinstall packages?[Y/N]: "
+
+if "!choice!" == "y" (goto checkgpu)
+if "!choice!"=="Y" (goto checkgpu)
+
+goto end
+)
+
+REM Check the python version
+echo Python versions 3.10-3.12 have been confirmed to work. Other versions are currently not supported. You currently have:
+python -V
+set choice=
+set /p choice= "Do you want to continue?[Y/N]: "
+
+
+if "!choice!" == "y" (goto makevenv)
+if "!choice!"=="Y" (goto makevenv)
+
+goto end
+
+:makevenv
+REM This creates a virtual environment in the folder
+echo Creating a Virtual Environment...
+python -m venv venv
+echo Upgrading pip in Virtual Environment to lower chance of error...
+"%cd%/venv/Scripts/python.exe" -m pip install --upgrade pip
+
+:checkgpu
+REM ask Windows for GPU
+where nvidia-smi >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Error: Nvidia GPU doesn't exist or drivers installed incorrectly. Please confirm your drivers are installed.
+ goto end
+)
+
+echo Checking your GPU...
+
+for /F "tokens=* skip=1" %%n in ('nvidia-smi --query-gpu=name') do set GPU_NAME=%%n && goto gpuchecked
+
+:gpuchecked
+echo Detected %GPU_NAME%
+set "GPU_SERIES=%GPU_NAME:*RTX =%"
+set "GPU_SERIES=%GPU_SERIES:~0,2%00"
+
+REM This gets the shortened Python version for later use. e.g. 3.10.13 becomes 310.
+for /f "delims=" %%A in ('python -V') do set "pyv=%%A"
+for /f "tokens=2 delims= " %%A in ("%pyv%") do (
+ set pyv=%%A
+)
+set pyv=%pyv:.=%
+set pyv=%pyv:~0,3%
+
+echo Installing torch...
+
+if !GPU_SERIES! geq 5000 (
+ goto torch270
+) else (
+ goto torch260
+)
+
+REM RTX 5000 Series
+:torch270
+"%cd%/venv/Scripts/pip.exe" install torch==2.7.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 --force-reinstall
+REM Check if pip installation was successful
+if %errorlevel% neq 0 (
+ echo Warning: Failed to install dependencies. You may need to install them manually.
+ goto end
+)
+
+REM Ask if user wants Sage Attention
+set choice=
+echo Do you want to install any of the following? They speed up generation.
+echo 1) Sage Attention
+echo 2) Flash Attention
+echo 3) BOTH!
+echo 4) No
+set /p choice= "Input Selection: "
+
+set both="N"
+
+if "!choice!" == "1" (goto triton270)
+if "!choice!"== "2" (goto flash270)
+if "!choice!"== "3" (set both="Y"
+goto triton270
+)
+
+goto requirements
+
+:triton270
+REM Sage Attention and Triton for Torch 2.7.0
+"%cd%/venv/Scripts/pip.exe" install "triton-windows<3.4" --force-reinstall
+"%cd%/venv/Scripts/pip.exe" install "https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu128torch2.7.0-cp%pyv%-cp%pyv%-win_amd64.whl" --force-reinstall
+echo Finishing up installing triton-windows. This requires extraction of libraries into Python Folder...
+
+REM Check for python version and download the triton-windows required libs accordingly
+if %pyv% == 310 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.10.11_include_libs.zip', 'triton-lib.zip')"
+)
+
+if %pyv% == 311 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.11.9_include_libs.zip', 'triton-lib.zip')"
+)
+
+if %pyv% == 312 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.12.7_include_libs.zip', 'triton-lib.zip')"
+)
+
+REM Extract the zip into the Python Folder and Delete zip
+powershell Expand-Archive -Path '%cd%\triton-lib.zip' -DestinationPath '%cd%\venv\Scripts\' -force
+del triton-lib.zip
+if %both% == "Y" (goto flash270)
+
+goto requirements
+
+:flash270
+REM Install flash-attn.
+"%cd%/venv/Scripts/pip.exe" install "https://huggingface.co/lldacing/flash-attention-windows-wheel/resolve/main/flash_attn-2.7.4.post1%%2Bcu128torch2.7.0cxx11abiFALSE-cp%pyv%-cp%pyv%-win_amd64.whl?download=true"
+goto requirements
+
+
+REM RTX 4000 Series and below
+:torch260
+"%cd%/venv/Scripts/pip.exe" install torch==2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 --force-reinstall
+REM Check if pip installation was successful
+if %errorlevel% neq 0 (
+ echo Warning: Failed to install dependencies. You may need to install them manually.
+ goto end
+)
+
+REM Ask if user wants Sage Attention
+set choice=
+echo Do you want to install any of the following? They speed up generation.
+echo 1) Sage Attention
+echo 2) Flash Attention
+echo 3) BOTH!
+echo 4) No
+set /p choice= "Input Selection: "
+
+set both="N"
+
+if "!choice!" == "1" (goto triton260)
+if "!choice!"== "2" (goto flash260)
+if "!choice!"== "3" (set both="Y"
+goto triton260)
+
+goto requirements
+
+:triton260
+REM Sage Attention and Triton for Torch 2.6.0
+"%cd%/venv/Scripts/pip.exe" install "triton-windows<3.3.0" --force-reinstall
+"%cd%/venv/Scripts/pip.exe" install https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu126torch2.6.0-cp%pyv%-cp%pyv%-win_amd64.whl --force-reinstall
+
+echo Finishing up installing triton-windows. This requires extraction of libraries into Python Folder...
+
+REM Check for python version and download the triton-windows required libs accordingly
+if %pyv% == 310 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.10.11_include_libs.zip', 'triton-lib.zip')"
+)
+
+if %pyv% == 311 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.11.9_include_libs.zip', 'triton-lib.zip')"
+)
+
+if %pyv% == 312 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.12.7_include_libs.zip', 'triton-lib.zip')"
+)
+
+REM Extract the zip into the Python Folder and Delete zip
+powershell Expand-Archive -Path '%cd%\triton-lib.zip' -DestinationPath '%cd%\venv\Scripts\' -force
+del triton-lib.zip
+
+if %both% == "Y" (goto flash260)
+
+goto requirements
+
+:flash260
+REM Install flash-attn.
+"%cd%/venv/Scripts/pip.exe" install "https://huggingface.co/lldacing/flash-attention-windows-wheel/resolve/main/flash_attn-2.7.4%%2Bcu126torch2.6.0cxx11abiFALSE-cp%pyv%-cp%pyv%-win_amd64.whl?download=true"
+
+:requirements
+echo Installing remaining required packages through pip...
+REM This assumes there's a requirements.txt file in the root
+"%cd%/venv/Scripts/pip.exe" install -r requirements.txt
+
+REM Check if pip installation was successful
+if %errorlevel% neq 0 (
+ echo Warning: Failed to install dependencies. You may need to install them manually.
+ goto end
+)
+
+echo Setup complete.
+
+:end
+echo Exiting setup script.
+pause
diff --git a/install_40xx_50xx.bat b/install_40xx_50xx.bat
new file mode 100644
index 00000000..4cfa287d
--- /dev/null
+++ b/install_40xx_50xx.bat
@@ -0,0 +1,172 @@
+@echo off
+echo FramePack-Studio Setup Script for RTX 40xx/50xx GPUs
+echo Using PyTorch 2.8.0 with CUDA 12.8
+echo ============================================
+setlocal enabledelayedexpansion
+
+REM Check if Python is installed
+where python >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Error: Python is not installed or not in your PATH. Please install Python and try again.
+ goto end
+)
+
+REM Check Python version
+echo Checking Python version...
+echo Python versions 3.10-3.13 are supported. You currently have:
+python -V
+echo.
+
+REM Check for existing venv
+if exist "%cd%\venv" (
+ echo Virtual Environment already exists.
+ set /p choice="Do you want to delete and recreate it? [Y/N]: "
+ if /i "!choice!" == "Y" (
+ echo Deleting existing venv...
+ rmdir /s /q "%cd%\venv"
+ goto makevenv
+ ) else (
+ echo Keeping existing venv. Installing packages...
+ goto checkgpu
+ )
+) else (
+ goto makevenv
+)
+
+:makevenv
+echo.
+echo Creating Virtual Environment...
+python -m venv venv
+if %errorlevel% neq 0 (
+ echo Error: Failed to create virtual environment.
+ goto end
+)
+
+echo Upgrading pip in Virtual Environment...
+"%cd%\venv\Scripts\python.exe" -m pip install --upgrade pip
+
+:checkgpu
+REM Check for Nvidia GPU
+where nvidia-smi >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Warning: nvidia-smi not found. Cannot detect GPU.
+ echo Continuing with installation anyway...
+ goto installpytorch
+)
+
+echo.
+echo Detecting GPU...
+for /F "tokens=* skip=1" %%n in ('nvidia-smi --query-gpu=name') do set GPU_NAME=%%n && goto gpuchecked
+
+:gpuchecked
+echo Detected: %GPU_NAME%
+echo.
+
+:installpytorch
+echo Installing PyTorch 2.8.0 with CUDA 12.8...
+echo This may take several minutes depending on your internet connection...
+"%cd%\venv\Scripts\pip.exe" install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
+
+if %errorlevel% neq 0 (
+ echo Error: Failed to install PyTorch. Please check your internet connection and try again.
+ goto end
+)
+
+echo.
+echo PyTorch 2.8.0 installed successfully!
+echo.
+
+REM Get Python version for acceleration packages
+for /f "delims=" %%A in ('python -V') do set "pyv=%%A"
+for /f "tokens=2 delims= " %%A in ("%pyv%") do set pyv=%%A
+set pyv=%pyv:.=%
+set pyv=%pyv:~0,3%
+
+REM Ask about acceleration packages
+echo ============================================
+echo Optional: Install acceleration packages?
+echo These can significantly speed up generation:
+echo.
+echo 1) Sage Attention (recommended for RTX 40xx/50xx)
+echo 2) Flash Attention (alternative acceleration)
+echo 3) BOTH (maximum performance)
+echo 4) Skip (no acceleration packages)
+echo.
+set /p accel_choice="Enter your choice [1-4]: "
+
+if "!accel_choice!" == "1" goto install_sage
+if "!accel_choice!" == "2" goto install_flash
+if "!accel_choice!" == "3" (
+ set install_both=Y
+ goto install_sage
+)
+goto requirements
+
+:install_sage
+echo.
+echo Installing Sage Attention...
+echo Installing triton-windows...
+"%cd%\venv\Scripts\pip.exe" install "triton-windows<3.4" --force-reinstall
+
+echo Downloading triton libraries for Python %pyv%...
+if %pyv% == 310 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.10.11_include_libs.zip', 'triton-lib.zip')"
+)
+if %pyv% == 311 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.11.9_include_libs.zip', 'triton-lib.zip')"
+)
+if %pyv% == 312 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.12.7_include_libs.zip', 'triton-lib.zip')"
+)
+if %pyv% == 313 (
+ powershell -Command "(New-Object System.Net.WebClient).DownloadFile('https://github.com/woct0rdho/triton-windows/releases/download/v3.0.0-windows.post1/python_3.12.7_include_libs.zip', 'triton-lib.zip')"
+)
+
+if exist "triton-lib.zip" (
+ echo Extracting triton libraries...
+ powershell Expand-Archive -Path '%cd%\triton-lib.zip' -DestinationPath '%cd%\venv\Scripts\' -force
+ del triton-lib.zip
+ echo Triton libraries installed successfully.
+) else (
+ echo Warning: Could not download triton libraries for Python %pyv%.
+)
+
+echo Installing Sage Attention for PyTorch 2.8.0...
+"%cd%\venv\Scripts\pip.exe" install "https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu128torch2.7.0-cp%pyv%-cp%pyv%-win_amd64.whl" --force-reinstall
+
+if "!install_both!" == "Y" goto install_flash
+goto requirements
+
+:install_flash
+echo.
+echo Installing Flash Attention...
+"%cd%\venv\Scripts\pip.exe" install "https://huggingface.co/lldacing/flash-attention-windows-wheel/resolve/main/flash_attn-2.7.4.post1%%2Bcu128torch2.7.0cxx11abiFALSE-cp%pyv%-cp%pyv%-win_amd64.whl?download=true"
+
+:requirements
+echo.
+echo ============================================
+echo Installing remaining required packages...
+if exist "%cd%\requirements.txt" (
+ "%cd%\venv\Scripts\pip.exe" install -r requirements.txt
+ if %errorlevel% neq 0 (
+ echo Warning: Some packages failed to install. Check the output above.
+ )
+) else (
+ echo Warning: requirements.txt not found. Skipping additional packages.
+)
+
+echo.
+echo ============================================
+echo Setup complete!
+echo.
+echo To start FramePack-Studio:
+echo - Run: start.bat (or run.bat)
+echo.
+
+REM Verify PyTorch installation
+echo Verifying PyTorch installation...
+"%cd%\venv\Scripts\python.exe" -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA Available: {torch.cuda.is_available()}'); print(f'CUDA Version: {torch.version.cuda if torch.cuda.is_available() else \"N/A\"}')"
+
+:end
+echo.
+pause
diff --git a/modules/__init__.py b/modules/__init__.py
index 42856011..f5cd0387 100644
--- a/modules/__init__.py
+++ b/modules/__init__.py
@@ -1,4 +1,4 @@
-# modules/__init__.py
-
-# Workaround for the single lora bug. Must not be an empty string.
-DUMMY_LORA_NAME = " "
+# modules/__init__.py
+
+# Workaround for the single lora bug. Must not be an empty string.
+DUMMY_LORA_NAME = " "
diff --git a/modules/interface.py b/modules/interface.py
index 2fe30194..6bcd1e59 100644
--- a/modules/interface.py
+++ b/modules/interface.py
@@ -1,2472 +1,2481 @@
-import gradio as gr
-import time
-import datetime
-import random
-import json
-import os
-import shutil
-from typing import List, Dict, Any, Optional
-from PIL import Image, ImageDraw, ImageFont
-import numpy as np
-import base64
-import io
-import functools
-
-from modules.version import APP_VERSION, APP_VERSION_DISPLAY
-
-import subprocess
-import itertools
-import re
-from collections import defaultdict
-import imageio
-import imageio.plugins.ffmpeg
-import ffmpeg
-from diffusers_helper.utils import generate_timestamp
-
-from modules.video_queue import JobStatus, Job, JobType
-from modules.prompt_handler import get_section_boundaries, get_quick_prompts, parse_timestamped_prompt
-from modules.llm_enhancer import enhance_prompt
-from modules.llm_captioner import caption_image
-from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
-from diffusers_helper.bucket_tools import find_nearest_bucket
-from modules.pipelines.metadata_utils import create_metadata
-from modules import DUMMY_LORA_NAME # Import the constant
-
-from modules.toolbox_app import tb_processor
-from modules.toolbox_app import tb_create_video_toolbox_ui, tb_get_formatted_toolbar_stats
-from modules.xy_plot_ui import create_xy_plot_ui, xy_plot_process
-
-# Define the dummy LoRA name as a constant
-
-def create_interface(
- process_fn,
- monitor_fn,
- end_process_fn,
- update_queue_status_fn,
- load_lora_file_fn,
- job_queue,
- settings,
- default_prompt: str = '[1s: The person waves hello] [3s: The person jumps up and down] [5s: The person does a dance]',
- lora_names: list = [],
- lora_values: list = []
-):
- """
- Create the Gradio interface for the video generation application
-
- Args:
- process_fn: Function to process a new job
- monitor_fn: Function to monitor an existing job
- end_process_fn: Function to cancel the current job
- update_queue_status_fn: Function to update the queue status display
- default_prompt: Default prompt text
- lora_names: List of loaded LoRA names
-
- Returns:
- Gradio Blocks interface
- """
- def is_video_model(model_type_value):
- return model_type_value in ["Video", "Video with Endframe", "Video F1"]
-
- # Add near the top of create_interface function, after the initial setup
- def get_latents_display_top():
- """Get current latents display preference - centralized access point"""
- return settings.get("latents_display_top", False)
-
- def create_latents_layout_update():
- """Create a standardized layout update based on current setting"""
- display_top = get_latents_display_top()
- if display_top:
- return (
- gr.update(visible=True), # top_preview_row
- gr.update(visible=False, value=None) # preview_image (right column)
- )
- else:
- return (
- gr.update(visible=False), # top_preview_row
- gr.update(visible=True) # preview_image (right column)
- )
-
-
-
- # Get section boundaries and quick prompts
- section_boundaries = get_section_boundaries()
- quick_prompts = get_quick_prompts()
-
- # --- Function to update queue stats (Moved earlier to resolve UnboundLocalError) ---
- def update_stats(*args): # Accept any arguments and ignore them
- # Get queue status data
- queue_status_data = update_queue_status_fn()
-
- # Get queue statistics for the toolbar display
- jobs = job_queue.get_all_jobs()
-
- # Count jobs by status
- pending_count = 0
- running_count = 0
- completed_count = 0
-
- for job in jobs:
- if hasattr(job, 'status'):
- status = str(job.status)
- if status == "JobStatus.PENDING":
- pending_count += 1
- elif status == "JobStatus.RUNNING":
- running_count += 1
- elif status == "JobStatus.COMPLETED":
- completed_count += 1
-
- # Format the queue stats display text
- queue_stats_text = f"Queue: {pending_count} | Running: {running_count} | Completed: {completed_count}
"
-
- return queue_status_data, queue_stats_text
-
- # --- Preset System Functions ---
- PRESET_FILE = os.path.join(".framepack", "generation_presets.json")
-
- def load_presets(model_type):
- if not os.path.exists(PRESET_FILE):
- return []
- with open(PRESET_FILE, 'r') as f:
- data = json.load(f)
- return list(data.get(model_type, {}).keys())
-
- # Create the interface
- css = make_progress_bar_css()
- css += """
-
- .short-import-box, .short-import-box > div {
- min-height: 40px !important;
- height: 40px !important;
- }
- /* Image container styling - more aggressive approach */
- .contain-image, .contain-image > div, .contain-image > div > img {
- object-fit: contain !important;
- }
-
- #non-mirrored-video {
- transform: scaleX(-1) !important;
- }
-
- /* Target all images in the contain-image class and its children */
- .contain-image img,
- .contain-image > div > img,
- .contain-image * img {
- object-fit: contain !important;
- width: 100% !important;
- height: 60vh !important;
- max-height: 100% !important;
- max-width: 100% !important;
- }
-
- /* Additional selectors to override Gradio defaults */
- .gradio-container img,
- .gradio-container .svelte-1b5oq5x,
- .gradio-container [data-testid="image"] img {
- object-fit: contain !important;
- }
-
- /* Toolbar styling */
- #fixed-toolbar {
- position: fixed;
- top: 0;
- left: 0;
- width: 100vw;
- z-index: 1000;
- background: #333;
- color: #fff;
- padding: 0px 10px; /* Reduced top/bottom padding */
- display: flex;
- align-items: center;
- gap: 8px;
- box-shadow: 0 2px 8px rgba(0,0,0,0.1);
- }
-
- /* Responsive toolbar title */
- .toolbar-title {
- font-size: 1.4rem;
- margin: 0;
- color: white;
- white-space: nowrap;
- overflow: hidden;
- text-overflow: ellipsis;
- }
-
- /* Toolbar Patreon link */
- .toolbar-patreon {
- margin: 0 0 0 20px;
- color: white;
- font-size: 0.9rem;
- white-space: nowrap;
- display: inline-block;
- }
- .toolbar-patreon a {
- color: white;
- text-decoration: none;
- }
- .toolbar-patreon a:hover {
- text-decoration: underline;
- }
-
- /* Toolbar Version number */
- .toolbar-version {
- margin: 0 15px; /* Space around version */
- color: white;
- font-size: 0.8rem;
- white-space: nowrap;
- display: inline-block;
- }
-
- /* Responsive design for screens */
- @media (max-width: 1147px) {
- .toolbar-patreon, .toolbar-version { /* Hide both on smaller screens */
- display: none;
- }
- .footer-patreon, .footer-version { /* Show both in footer on smaller screens */
- display: inline-block !important; /* Ensure they are shown */
- }
- #fixed-toolbar {
- gap: 4px !important; /* Reduce gap for screens <= 1024px */
- }
- #fixed-toolbar > div:first-child { /* Target the first gr.Column (Title) */
- min-width: fit-content !important; /* Override Python-set min-width */
- flex-shrink: 0 !important; /* Prevent title column from shrinking too much */
- }
- }
-
- @media (min-width: 1148px) {
- .footer-patreon, .footer-version { /* Hide both in footer on larger screens */
- display: none !important;
- }
- }
-
- @media (max-width: 768px) {
- .toolbar-title {
- font-size: 1.1rem;
- max-width: 150px;
- }
- #fixed-toolbar {
- padding: 3px 6px;
- gap: 4px;
- }
- .toolbar-text {
- font-size: 0.75rem;
- }
- }
-
- @media (max-width: 510px) {
- #toolbar-ram-col, #toolbar-vram-col, #toolbar-gpu-col {
- display: none !important;
- }
- }
-
- @media (max-width: 480px) {
- .toolbar-title {
- font-size: 1rem;
- max-width: 120px;
- }
- #fixed-toolbar {
- padding: 2px 4px;
- gap: 2px;
- }
- .toolbar-text {
- font-size: 0.7rem;
- }
- }
-
- /* Button styling */
- #toolbar-add-to-queue-btn button {
- font-size: 14px !important;
- padding: 4px 16px !important;
- height: 32px !important;
- min-width: 80px !important;
- }
- .narrow-button {
- min-width: 40px !important;
- width: 40px !important;
- padding: 0 !important;
- margin: 0 !important;
- }
- .gr-button-primary {
- color: white;
- }
-
- /* Layout adjustments */
- body, .gradio-container {
- padding-top: 42px !important; /* Adjusted for new toolbar height (36px - 10px) */
- }
-
- @media (max-width: 848px) {
- body, .gradio-container {
- padding-top: 48px !important;
- }
- }
-
- @media (max-width: 768px) {
- body, .gradio-container {
- padding-top: 22px !important; /* Adjusted for new toolbar height (32px - 10px) */
- }
- }
-
- @media (max-width: 480px) {
- body, .gradio-container {
- padding-top: 18px !important; /* Adjusted for new toolbar height (28px - 10px) */
- }
- }
-
- /* hide the gr.Video source selection bar for tb_input_video_component */
- #toolbox-video-player .source-selection {
- display: none !important;
- }
- /* control sizing for gr.Video components */
- .video-size video {
- max-height: 60vh;
- min-height: 300px !important;
- object-fit: contain;
- }
- /* NEW: Closes the gap between input tabs and the pipeline accordion below them */
- #pipeline-controls-wrapper {
- margin-top: -15px !important; /* Adjust this value to get the perfect "snug" fit */
- }
- /* --- NEW CSS RULE FOR GALLERY SCROLLING --- */
- #gallery-scroll-wrapper {
- max-height: 600px; /* Set your desired fixed height */
- overflow-y: auto; /* Add a scrollbar only when needed */
- }
- #toolbox-start-pipeline-btn {
- margin-top: -14px !important; /* Adjust this value to get the perfect alignment */
- }
-
- .control-group {
- border-top: 1px solid #ccc;
- border-bottom: 1px solid #ccc;
- margin: 12px 0;
- }
- """
-
- # Get the theme from settings
- current_theme = settings.get("gradio_theme", "default") # Use default if not found
- block = gr.Blocks(css=css, title="FramePack Studio", theme=current_theme).queue()
-
- with block:
- with gr.Row(elem_id="fixed-toolbar"):
- with gr.Column(scale=0, min_width=400): # Title/Version/Patreon
- gr.HTML(f"""
-
- """)
- # REMOVED: refresh_stats_btn - Toolbar refresh button is no longer needed
- # with gr.Column(scale=0, min_width=40):
- # refresh_stats_btn = gr.Button("⟳", elem_id="refresh-stats-btn", elem_classes="narrow-button")
- with gr.Column(scale=1, min_width=180): # Queue Stats
- queue_stats_display = gr.Markdown("Queue: 0 | Running: 0 | Completed: 0
")
-
- # --- System Stats Display - Single gr.Textbox per stat ---
- with gr.Column(scale=0, min_width=173, elem_id="toolbar-ram-col"): # RAM Column
- toolbar_ram_display_component = gr.Textbox(
- value="RAM: N/A",
- interactive=False,
- lines=1,
- max_lines=1,
- show_label=False,
- container=False,
- elem_id="toolbar-ram-stat",
- elem_classes="toolbar-stat-textbox"
- )
- with gr.Column(scale=0, min_width=138, elem_id="toolbar-vram-col"): # VRAM Column
- toolbar_vram_display_component = gr.Textbox(
- value="VRAM: N/A",
- interactive=False,
- lines=1,
- max_lines=1,
- show_label=False,
- container=False,
- elem_id="toolbar-vram-stat",
- elem_classes="toolbar-stat-textbox"
- # Visibility controlled by tb_get_formatted_toolbar_stats
- )
- with gr.Column(scale=0, min_width=130, elem_id="toolbar-gpu-col"): # GPU Column
- toolbar_gpu_display_component = gr.Textbox(
- value="GPU: N/A",
- interactive=False,
- lines=1,
- max_lines=1,
- show_label=False,
- container=False,
- elem_id="toolbar-gpu-stat",
- elem_classes="toolbar-stat-textbox"
- # Visibility controlled by tb_get_formatted_toolbar_stats
- )
- # --- End of System Stats Display ---
-
- # Removed old version_display column
- # --- End of Toolbar ---
-
- # Essential to capture main_tabs_component for later use by send_to_toolbox_btn
- with gr.Tabs(elem_id="main_tabs") as main_tabs_component:
- with gr.Tab("Generate", id="generate_tab"):
- # NEW: Top preview area for latents display
- with gr.Row(visible=get_latents_display_top()) as top_preview_row:
- top_preview_image = gr.Image(
- label="Next Latents (Top Display)",
- height=150,
- visible=True,
- type="numpy",
- interactive=False,
- elem_classes="contain-image",
- image_mode="RGB"
- )
-
- with gr.Row():
- with gr.Column(scale=2):
- model_type = gr.Radio(
- choices=[("Original", "Original"), ("Original with Endframe", "Original with Endframe"), ("F1", "F1"), ("Video", "Video"), ("Video with Endframe", "Video with Endframe"), ("Video F1", "Video F1")],
- value="Original",
- label="Generation Type"
- )
- with gr.Accordion("Original Presets", open=False, visible=True) as preset_accordion:
- with gr.Row():
- preset_dropdown = gr.Dropdown(label="Select Preset", choices=load_presets("Original"), interactive=True, scale=2)
- delete_preset_button = gr.Button("🗑️ Delete", variant="stop", scale=1)
- with gr.Row():
- preset_name_textbox = gr.Textbox(label="Preset Name", placeholder="Enter a name for your preset", scale=2)
- save_preset_button = gr.Button("💾 Save", variant="primary", scale=1)
- with gr.Row(visible=False) as confirm_delete_row:
- gr.Markdown("### Are you sure you want to delete this preset?")
- confirm_delete_yes_btn = gr.Button("🗑️ Yes, Delete", variant="stop")
- confirm_delete_no_btn = gr.Button("↩️ No, Go Back")
- with gr.Accordion("Basic Parameters", open=True, visible=True) as basic_parameters_accordion:
- with gr.Group():
- total_second_length = gr.Slider(label="Video Length (Seconds)", minimum=1, maximum=120, value=6, step=0.1)
- with gr.Row("Resolution"):
- resolutionW = gr.Slider(
- label="Width", minimum=128, maximum=768, value=640, step=32,
- info="Nearest valid width will be used."
- )
- resolutionH = gr.Slider(
- label="Height", minimum=128, maximum=768, value=640, step=32,
- info="Nearest valid height will be used."
- )
- resolution_text = gr.Markdown(value="Selected bucket for resolution: 640 x 640
", label="", show_label=False)
-
- # --- START OF REFACTORED XY PLOT SECTION ---
- xy_plot_components = create_xy_plot_ui(
- lora_names=lora_names,
- default_prompt=default_prompt,
- DUMMY_LORA_NAME=DUMMY_LORA_NAME,
- )
- xy_group = xy_plot_components["group"]
- xy_plot_status = xy_plot_components["status"]
- xy_plot_output = xy_plot_components["output"]
- # --- END OF REFACTORED XY PLOT SECTION ---
-
- with gr.Group(visible=True) as standard_generation_group: # Default visibility: True because "Original" model is not "Video"
- with gr.Group(visible=True) as image_input_group: # This group now only contains the start frame image
- with gr.Row():
- with gr.Column(scale=1): # Start Frame Image Column
- input_image = gr.Image(
- sources='upload',
- type="numpy",
- label="Start Frame (optional)",
- elem_classes="contain-image",
- image_mode="RGB",
- show_download_button=False,
- show_label=True, # Keep label for clarity
- container=True
- )
-
- with gr.Group(visible=False) as video_input_group:
- input_video = gr.Video(
- sources='upload',
- label="Video Input",
- height=420,
- show_label=True
- )
- combine_with_source = gr.Checkbox(
- label="Combine with source video",
- value=True,
- info="If checked, the source video will be combined with the generated video",
- interactive=True
- )
- num_cleaned_frames = gr.Slider(label="Number of Context Frames (Adherence to Video)", minimum=2, maximum=10, value=5, step=1, interactive=True, info="Expensive. Retain more video details. Reduce if memory issues or motion too restricted (jumpcut, ignoring prompt, still).")
-
-
- # End Frame Image Input
- # Initial visibility is False, controlled by update_input_visibility
- with gr.Column(scale=1, visible=False) as end_frame_group_original:
- end_frame_image_original = gr.Image(
- sources='upload',
- type="numpy",
- label="End Frame (Optional)",
- elem_classes="contain-image",
- image_mode="RGB",
- show_download_button=False,
- show_label=True,
- container=True
- )
-
- # End Frame Influence slider
- # Initial visibility is False, controlled by update_input_visibility
- with gr.Group(visible=False) as end_frame_slider_group:
- end_frame_strength_original = gr.Slider(
- label="End Frame Influence",
- minimum=0.05,
- maximum=1.0,
- value=1.0,
- step=0.05,
- info="Controls how strongly the end frame guides the generation. 1.0 is full influence."
- )
-
-
-
- with gr.Row():
- prompt = gr.Textbox(label="Prompt", value=default_prompt, scale=10)
- with gr.Row():
- enhance_prompt_btn = gr.Button("✨ Enhance", scale=1)
- caption_btn = gr.Button("✨ Caption", scale=1)
-
- with gr.Accordion("Prompt Parameters", open=False):
- n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=True) # Make visible for both models
-
- blend_sections = gr.Slider(
- minimum=0, maximum=10, value=4, step=1,
- label="Number of sections to blend between prompts"
- )
- with gr.Accordion("Batch Input", open=False):
- batch_input_images = gr.File(
- label="Batch Images (Upload one or more)",
- file_count="multiple",
- file_types=["image"],
- type="filepath"
- )
- batch_input_gallery = gr.Gallery(
- label="Selected Batch Images",
- visible=False,
- columns=5,
- object_fit="contain",
- height="auto"
- )
- add_batch_to_queue_btn = gr.Button("🚀 Add Batch to Queue", variant="primary")
- with gr.Accordion("Generation Parameters", open=True):
- with gr.Row():
- steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1)
- def on_input_image_change(img):
- if img is not None:
- return gr.update(info="Nearest valid bucket size will be used. Height will be adjusted automatically."), gr.update(visible=False)
- else:
- return gr.update(info="Nearest valid width will be used."), gr.update(visible=True)
- input_image.change(fn=on_input_image_change, inputs=[input_image], outputs=[resolutionW, resolutionH])
- def on_resolution_change(img, resolutionW, resolutionH):
- out_bucket_resH, out_bucket_resW = [640, 640]
- if img is not None:
- H, W, _ = img.shape
- out_bucket_resH, out_bucket_resW = find_nearest_bucket(H, W, resolution=resolutionW)
- else:
- out_bucket_resH, out_bucket_resW = find_nearest_bucket(resolutionH, resolutionW, (resolutionW+resolutionH)/2) # if resolutionW > resolutionH else resolutionH
- return gr.update(value=f"Selected bucket for resolution: {out_bucket_resW} x {out_bucket_resH}
")
- resolutionW.change(fn=on_resolution_change, inputs=[input_image, resolutionW, resolutionH], outputs=[resolution_text], show_progress="hidden")
- resolutionH.change(fn=on_resolution_change, inputs=[input_image, resolutionW, resolutionH], outputs=[resolution_text], show_progress="hidden")
-
- with gr.Row():
- seed = gr.Number(label="Seed", value=2500, precision=0)
- randomize_seed = gr.Checkbox(label="Randomize", value=True, info="Generate a new random seed for each job")
- with gr.Accordion("LoRAs", open=False):
- with gr.Row():
- lora_selector = gr.Dropdown(
- choices=lora_names,
- label="Select LoRAs to Load",
- multiselect=True,
- value=[],
- info="Select one or more LoRAs to use for this job"
- )
- lora_names_states = gr.State(lora_names)
- lora_sliders = {}
- for lora in lora_names:
- lora_sliders[lora] = gr.Slider(
- minimum=0.0, maximum=2.0, value=1.0, step=0.01,
- label=f"{lora} Weight", visible=False, interactive=True
- )
- with gr.Accordion("Latent Image Options", open=False):
- latent_type = gr.Dropdown(
- ["Noise", "White", "Black", "Green Screen"], label="Latent Image", value="Noise", info="Used as a starting point if no image is provided"
- )
- with gr.Accordion("Advanced Parameters", open=False):
- gr.Markdown("#### Motion Model")
- gr.Markdown("Settings for precise control of the motion model")
-
- with gr.Group(elem_classes="control-group"):
- latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, info='Change at your own risk, very experimental') # Should not change
- gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=10.0, step=0.5)
-
- gr.Markdown("#### CFG Scale")
- gr.Markdown("Much better prompt following. Warning: Modifying these values from their defaults will almost double generation time. ⚠️")
-
- with gr.Group(elem_classes="control-group"):
- cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=3.0, value=1.0, step=0.1)
- rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.05)
-
- gr.Markdown("#### Cache Options")
- gr.Markdown("Using a cache will speed up generation. May affect quality, fine or even coarse details, and may change or inhibit motion. You can choose at most one.")
-
- with gr.Group(elem_classes="control-group"):
- with gr.Row():
- cache_type = gr.Radio(["MagCache", "TeaCache", "None"], value='MagCache', label="Caching strategy", info="Which cache implementation to use, if any")
-
- with gr.Row(): # MagCache now first
- magcache_threshold = gr.Slider(label="MagCache Threshold", minimum=0.01, maximum=1.0, step=0.01, value=0.1, visible=True, info='[⬇️ **Faster**] Error tolerance. Lower = more estimated steps')
- magcache_max_consecutive_skips = gr.Slider(label="MagCache Max Consecutive Skips", minimum=1, maximum=5, step=1, value=2, visible=True, info='[⬆️ **Faster**] Allow multiple estimated steps in a row')
- magcache_retention_ratio = gr.Slider(label="MagCache Retention Ratio", minimum=0.0, maximum=1.0, step=0.01, value=0.25, visible=True, info='[⬇️ **Faster**] Disallow estimation in critical early steps')
-
- with gr.Row():
- teacache_num_steps = gr.Slider(label="TeaCache steps", minimum=1, maximum=50, step=1, value=25, visible=False, info='How many intermediate sections to keep in the cache')
- teacache_rel_l1_thresh = gr.Slider(label="TeaCache rel_l1_thresh", minimum=0.01, maximum=1.0, step=0.01, value=0.15, visible=False, info='[⬇️ **Faster**] Relative L1 Threshold')
-
- def update_cache_type(cache_type: str):
- enable_magcache = False
- enable_teacache = False
-
- if cache_type == 'MagCache':
- enable_magcache = True
- elif cache_type == 'TeaCache':
- enable_teacache = True
-
- magcache_threshold_update = gr.update(visible=enable_magcache)
- magcache_max_consecutive_skips_update = gr.update(visible=enable_magcache)
- magcache_retention_ratio_update = gr.update(visible=enable_magcache)
-
- teacache_num_steps_update = gr.update(visible=enable_teacache)
- teacache_rel_l1_thresh_update = gr.update(visible=enable_teacache)
-
- return [
- magcache_threshold_update,
- magcache_max_consecutive_skips_update,
- magcache_retention_ratio_update,
- teacache_num_steps_update,
- teacache_rel_l1_thresh_update
- ]
-
-
- cache_type.change(fn=update_cache_type, inputs=cache_type, outputs=[
- magcache_threshold,
- magcache_max_consecutive_skips,
- magcache_retention_ratio,
- teacache_num_steps,
- teacache_rel_l1_thresh
- ])
-
- with gr.Row("Metadata"):
- json_upload = gr.File(
- label="Upload Metadata JSON (optional)",
- file_types=[".json"],
- type="filepath",
- height=140,
- )
-
- with gr.Column():
- preview_image = gr.Image(
- label="Next Latents",
- height=150,
- visible=not get_latents_display_top(),
- type="numpy",
- interactive=False,
- elem_classes="contain-image",
- image_mode="RGB"
- )
- result_video = gr.Video(label="Finished Frames", autoplay=True, show_share_button=False, height=256, loop=True)
- progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
- progress_bar = gr.HTML('', elem_classes='no-generating-animation')
- with gr.Row():
- current_job_id = gr.Textbox(label="Current Job ID", value="", visible=True, interactive=True)
- start_button = gr.Button(value="🚀 Add to Queue", variant="primary", elem_id="toolbar-add-to-queue-btn")
- xy_plot_process_btn = gr.Button("🚀 Submit XY Plot", visible=False)
- video_input_required_message = gr.Markdown(
- "Input video required
", visible=False
- )
- end_button = gr.Button(value="❌ Cancel Current Job", interactive=True, visible=False)
-
-
-
- with gr.Tab("Queue"):
- with gr.Row():
- with gr.Column():
- with gr.Row() as queue_controls_row:
- refresh_button = gr.Button("🔄 Refresh Queue")
- load_queue_button = gr.Button("▶️ Resume Queue")
- queue_export_button = gr.Button("📦 Export Queue")
- clear_complete_button = gr.Button("🧹 Clear Completed Jobs", variant="secondary")
- clear_queue_button = gr.Button("❌ Cancel Queued Jobs", variant="stop")
- with gr.Row():
- import_queue_file = gr.File(
- label="Import Queue",
- file_types=[".json", ".zip"],
- type="filepath",
- visible=True,
- elem_classes="short-import-box"
- )
-
- with gr.Row(visible=False) as confirm_cancel_row:
- gr.Markdown("### Are you sure you want to cancel all pending jobs?")
- confirm_cancel_yes_btn = gr.Button("❌ Yes, Cancel All", variant="stop")
- confirm_cancel_no_btn = gr.Button("↩️ No, Go Back")
-
- with gr.Row():
- queue_status = gr.DataFrame(
- headers=["Job ID", "Type", "Status", "Created", "Started", "Completed", "Elapsed", "Preview"],
- datatype=["str", "str", "str", "str", "str", "str", "str", "html"],
- label="Job Queue"
- )
-
- with gr.Accordion("Queue Documentation", open=False):
- gr.Markdown("""
- ## Queue Tab Guide
-
- This tab is for managing your generation jobs.
-
- - **Refresh Queue**: Update the job list.
- - **Cancel Queue**: Stop all pending jobs.
- - **Clear Complete**: Remove finished, failed, or cancelled jobs from the list.
- - **Load Queue**: Load jobs from the default `queue.json`.
- - **Export Queue**: Save the current job list and its images to a zip file.
- - **Import Queue**: Load a queue from a `.json` or `.zip` file.
- """)
-
- # --- Event Handlers for Queue Tab ---
-
- # Function to clear all jobs in the queue
- def clear_all_jobs():
- try:
- cancelled_count = job_queue.clear_queue()
- print(f"Cleared {cancelled_count} jobs from the queue")
- return update_stats()
- except Exception as e:
- import traceback
- print(f"Error in clear_all_jobs: {e}")
- traceback.print_exc()
- return [], ""
-
- # Function to clear completed and cancelled jobs
- def clear_completed_jobs():
- try:
- removed_count = job_queue.clear_completed_jobs()
- print(f"Removed {removed_count} completed/cancelled jobs from the queue")
- return update_stats()
- except Exception as e:
- import traceback
- print(f"Error in clear_completed_jobs: {e}")
- traceback.print_exc()
- return [], ""
-
- # Function to load queue from queue.json
- def load_queue_from_json():
- try:
- loaded_count = job_queue.load_queue_from_json()
- print(f"Loaded {loaded_count} jobs from queue.json")
- return update_stats()
- except Exception as e:
- import traceback
- print(f"Error loading queue from JSON: {e}")
- traceback.print_exc()
- return [], ""
-
- # Function to import queue from a custom JSON file
- def import_queue_from_file(file_path):
- if not file_path:
- return update_stats()
- try:
- loaded_count = job_queue.load_queue_from_json(file_path)
- print(f"Loaded {loaded_count} jobs from {file_path}")
- return update_stats()
- except Exception as e:
- import traceback
- print(f"Error importing queue from file: {e}")
- traceback.print_exc()
- return [], ""
-
- # Function to export queue to a zip file
- def export_queue_to_zip():
- try:
- zip_path = job_queue.export_queue_to_zip()
- if zip_path and os.path.exists(zip_path):
- print(f"Queue exported to {zip_path}")
- else:
- print("Failed to export queue to zip")
- return update_stats()
- except Exception as e:
- import traceback
- print(f"Error exporting queue to zip: {e}")
- traceback.print_exc()
- return [], ""
-
- # --- Connect Buttons ---
- refresh_button.click(fn=update_stats, inputs=[], outputs=[queue_status, queue_stats_display])
-
- # Confirmation logic for Cancel Queue
- def show_cancel_confirmation():
- return gr.update(visible=False), gr.update(visible=True)
-
- def hide_cancel_confirmation():
- return gr.update(visible=True), gr.update(visible=False)
-
- def confirmed_clear_all_jobs():
- qs_data, qs_text = clear_all_jobs()
- return qs_data, qs_text, gr.update(visible=True), gr.update(visible=False)
-
- clear_queue_button.click(fn=show_cancel_confirmation, inputs=None, outputs=[queue_controls_row, confirm_cancel_row])
- confirm_cancel_no_btn.click(fn=hide_cancel_confirmation, inputs=None, outputs=[queue_controls_row, confirm_cancel_row])
- confirm_cancel_yes_btn.click(fn=confirmed_clear_all_jobs, inputs=None, outputs=[queue_status, queue_stats_display, queue_controls_row, confirm_cancel_row])
-
- clear_complete_button.click(fn=clear_completed_jobs, inputs=[], outputs=[queue_status, queue_stats_display])
- queue_export_button.click(fn=export_queue_to_zip, inputs=[], outputs=[queue_status, queue_stats_display])
-
- # Create a container for thumbnails (kept for potential future use, though not displayed in DataFrame)
- with gr.Row():
- thumbnail_container = gr.Column()
- thumbnail_container.elem_classes = ["thumbnail-container"]
-
- # Add CSS for thumbnails
-
- with gr.Tab("Outputs", id="outputs_tab"): # Ensure 'id' is present for tab switching
- outputDirectory_video = settings.get("output_dir", settings.default_settings['output_dir'])
- outputDirectory_metadata = settings.get("metadata_dir", settings.default_settings['metadata_dir'])
- def get_gallery_items():
- items = []
- for f in os.listdir(outputDirectory_metadata):
- if f.endswith(".png"):
- prefix = os.path.splitext(f)[0]
- latest_video = get_latest_video_version(prefix)
- if latest_video:
- video_path = os.path.join(outputDirectory_video, latest_video)
- mtime = os.path.getmtime(video_path)
- preview_path = os.path.join(outputDirectory_metadata, f)
- items.append((preview_path, prefix, mtime))
- items.sort(key=lambda x: x[2], reverse=True)
- return [(i[0], i[1]) for i in items]
- def get_latest_video_version(prefix):
- max_number = -1
- selected_file = None
- for f in os.listdir(outputDirectory_video):
- if f.startswith(prefix + "_") and f.endswith(".mp4"):
- # Skip files that include "combined" in their name
- if "combined" in f:
- continue
- try:
- num = int(f.replace(prefix + "_", '').replace(".mp4", ''))
- if num > max_number:
- max_number = num
- selected_file = f
- except ValueError:
- # Ignore files that do not have a valid number in their name
- continue
- return selected_file
- # load_video_and_info_from_prefix now also returns button visibility
- def load_video_and_info_from_prefix(prefix):
- video_file = get_latest_video_version(prefix)
- json_path = os.path.join(outputDirectory_metadata, prefix) + ".json"
-
- if not video_file or not os.path.exists(os.path.join(outputDirectory_video, video_file)) or not os.path.exists(json_path):
- # If video or info not found, button should be hidden
- return None, "Video or JSON not found.", gr.update(visible=False)
-
- video_path = os.path.join(outputDirectory_video, video_file)
- info_content = {"description": "no info"}
- if os.path.exists(json_path):
- with open(json_path, "r", encoding="utf-8") as f:
- info_content = json.load(f)
- # If video and info found, button should be visible
- return video_path, json.dumps(info_content, indent=2, ensure_ascii=False), gr.update(visible=True)
-
- gallery_items_state = gr.State(get_gallery_items())
- selected_original_video_path_state = gr.State(None) # Holds the ORIGINAL, UNPROCESSED path
- with gr.Row():
- with gr.Column(scale=2):
- thumbs = gr.Gallery(
- # value=[i[0] for i in get_gallery_items()],
- columns=[4],
- allow_preview=False,
- object_fit="cover",
- height="auto"
- )
- refresh_button = gr.Button("🔄 Update Gallery")
- with gr.Column(scale=5):
- video_out = gr.Video(sources=[], autoplay=True, loop=True, visible=False)
- with gr.Column(scale=1):
- info_out = gr.Textbox(label="Generation info", visible=False)
- send_to_toolbox_btn = gr.Button("➡️ Send to Post-processing", visible=False) # Added new send_to_toolbox_btn
- def refresh_gallery():
- new_items = get_gallery_items()
- return gr.update(value=[i[0] for i in new_items]), new_items
- refresh_button.click(fn=refresh_gallery, outputs=[thumbs, gallery_items_state])
-
- # MODIFIED: on_select now handles visibility of the new button
- def on_select(evt: gr.SelectData, gallery_items):
- if evt.index is None or not gallery_items or evt.index >= len(gallery_items):
- return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
-
- prefix = gallery_items[evt.index][1]
- # original_video_path is e.g., "outputs/my_actual_video.mp4"
- original_video_path, info_string, button_visibility_update = load_video_and_info_from_prefix(prefix)
-
- # Determine visibility for video and info based on whether video_path was found
- video_out_update = gr.update(value=original_video_path, visible=bool(original_video_path))
- info_out_update = gr.update(value=info_string, visible=bool(original_video_path))
-
- # IMPORTANT: Store the ORIGINAL, UNPROCESSED path in the gr.State
- return video_out_update, info_out_update, button_visibility_update, original_video_path
-
- thumbs.select(
- fn=on_select,
- inputs=[gallery_items_state],
- outputs=[video_out, info_out, send_to_toolbox_btn, selected_original_video_path_state] # Output original path to State
- )
- with gr.Tab("Post-processing", id="toolbox_tab"):
- # Call the function from toolbox_app.py to build the Toolbox UI
- # The toolbox_ui_layout (e.g., a gr.Column) is automatically placed here.
- toolbox_ui_layout, tb_target_video_input = tb_create_video_toolbox_ui()
-
- with gr.Tab("Settings"):
- with gr.Row():
- with gr.Column():
- save_metadata = gr.Checkbox(
- label="Save Metadata",
- info="Save to JSON file",
- value=settings.get("save_metadata", 6),
- )
- gpu_memory_preservation = gr.Slider(
- label="Memory Buffer for Stability (VRAM GB)",
- minimum=1,
- maximum=128,
- step=0.1,
- value=settings.get("gpu_memory_preservation", 6),
- info="Increase reserve if you see computer freezes, stagnant generation, or super slow sampling steps (try 1G at a time).\
- Otherwise smaller buffer is faster. Some models and lora need more buffer than others. \
- (5.5 - 8.5 is a common range)"
- )
- mp4_crf = gr.Slider(
- label="MP4 Compression",
- minimum=0,
- maximum=100,
- step=1,
- value=settings.get("mp4_crf", 16),
- info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs."
- )
- clean_up_videos = gr.Checkbox(
- label="Clean up video files",
- value=settings.get("clean_up_videos", True),
- info="If checked, only the final video will be kept after generation."
- )
- auto_cleanup_on_startup = gr.Checkbox(
- label="Automatically clean up temp folders on startup",
- value=settings.get("auto_cleanup_on_startup", False),
- info="If checked, temporary files (inc. post-processing) will be cleaned up when the application starts."
- )
-
- latents_display_top = gr.Checkbox(
- label="Display Next Latents across top of interface",
- value=get_latents_display_top(),
- info="If checked, the Next Latents preview will be displayed across the top of the interface instead of in the right column."
- )
-
- # gr.Markdown("---")
- # gr.Markdown("### Startup Settings")
- gr.Markdown("")
- # Initial values for startup preset dropdown
- # Ensure settings and load_presets are available in this scope
- initial_startup_model_val = settings.get("startup_model_type", "None")
- initial_startup_presets_choices_val = []
- initial_startup_preset_value_val = None
-
- if initial_startup_model_val and initial_startup_model_val != "None":
- # load_presets is defined further down in create_interface
- initial_startup_presets_choices_val = load_presets(initial_startup_model_val)
- saved_preset_for_initial_model_val = settings.get("startup_preset_name")
- if saved_preset_for_initial_model_val in initial_startup_presets_choices_val:
- initial_startup_preset_value_val = saved_preset_for_initial_model_val
-
- startup_model_type_dropdown = gr.Dropdown(
- label="Startup Model Type",
- choices=["None"] + [choice[0] for choice in model_type.choices if choice[0] != "XY Plot"], # model_type is the Radio on Generate tab
- value=initial_startup_model_val,
- info="Select a model type to load on startup. 'None' to disable."
- )
- startup_preset_name_dropdown = gr.Dropdown(
- label="Startup Preset",
- choices=initial_startup_presets_choices_val,
- value=initial_startup_preset_value_val,
- info="Select a preset for the startup model. Updates when Startup Model Type changes.",
- interactive=True # Must be interactive to be updated by another component
- )
-
- with gr.Accordion("System Prompt", open=False):
- with gr.Row(equal_height=True): # New Row to contain checkbox and reset button
- override_system_prompt = gr.Checkbox(
- label="Override System Prompt",
- value=settings.get("override_system_prompt", False),
- info="If checked, the system prompt template below will be used instead of the default one.",
- scale=1 # Give checkbox some scale
- )
- reset_system_prompt_btn = gr.Button(
- "🔄 Reset",
- scale=0
- )
- system_prompt_template = gr.Textbox(
- label="System Prompt Template",
- value=settings.get("system_prompt_template", "{\"template\": \"<|start_header_id|>system<|end_header_id|>\\n\\nDescribe the video by detailing the following aspects: 1. The main content and theme of the video.2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.4. background environment, light, style and atmosphere.5. camera angles, movements, and transitions used in the video:<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{}<|eot_id|>\", \"crop_start\": 95}"),
- lines=10,
- info="System prompt template used for video generation. Must be a valid JSON or Python dictionary string with 'template' and 'crop_start' keys. Example: {\"template\": \"your template here\", \"crop_start\": 95}"
- )
- # The reset_system_prompt_btn is now defined above within the Row
-
- # --- Settings Tab Event Handlers ---
-
- output_dir = gr.Textbox(
- label="Output Directory",
- value=settings.get("output_dir"),
- placeholder="Path to save generated videos"
- )
- metadata_dir = gr.Textbox(
- label="Metadata Directory",
- value=settings.get("metadata_dir"),
- placeholder="Path to save metadata files"
- )
- lora_dir = gr.Textbox(
- label="LoRA Directory",
- value=settings.get("lora_dir"),
- placeholder="Path to LoRA models"
- )
- gradio_temp_dir = gr.Textbox(label="Gradio Temporary Directory", value=settings.get("gradio_temp_dir"))
- auto_save = gr.Checkbox(
- label="Auto-save settings",
- value=settings.get("auto_save_settings", True)
- )
- # Add Gradio Theme Dropdown
- gradio_themes = ["default", "base", "soft", "glass", "mono", "origin", "citrus", "monochrome", "ocean", "NoCrypt/miku", "earneleh/paris", "gstaff/xkcd"]
- theme_dropdown = gr.Dropdown(
- label="Theme",
- choices=gradio_themes,
- value=settings.get("gradio_theme", "default"),
- info="Select the Gradio UI theme. Requires restart."
- )
- save_btn = gr.Button("💾 Save Settings")
- cleanup_btn = gr.Button("🗑️ Clean Up Temporary Files")
- status = gr.HTML("")
- cleanup_output = gr.Textbox(label="Cleanup Status", interactive=False)
-
- def save_settings(save_metadata, gpu_memory_preservation, mp4_crf, clean_up_videos, auto_cleanup_on_startup_val, latents_display_top_val, override_system_prompt_value, system_prompt_template_value, output_dir, metadata_dir, lora_dir, gradio_temp_dir, auto_save, selected_theme, startup_model_type_val, startup_preset_name_val):
- """Handles the manual 'Save Settings' button click."""
- # This function is for the manual save button.
- # It collects all current UI values and saves them.
- # The auto-save logic is handled by individual .change() and .blur() handlers
- # calling settings.set().
-
- # First, update the settings object with all current values from the UI
- try:
- # Save the system prompt template as is, without trying to parse it
- # The hunyuan.py file will handle parsing it when needed
- processed_template = system_prompt_template_value
-
- settings.save_settings(
- save_metadata=save_metadata,
- gpu_memory_preservation=gpu_memory_preservation,
- mp4_crf=mp4_crf,
- clean_up_videos=clean_up_videos,
- auto_cleanup_on_startup=auto_cleanup_on_startup_val, # ADDED
- latents_display_top=latents_display_top_val, # NEW: Added latents display position setting
- override_system_prompt=override_system_prompt_value,
- system_prompt_template=processed_template,
- output_dir=output_dir,
- metadata_dir=metadata_dir,
- lora_dir=lora_dir,
- gradio_temp_dir=gradio_temp_dir,
- auto_save_settings=auto_save,
- gradio_theme=selected_theme,
- startup_model_type=startup_model_type_val,
- startup_preset_name=startup_preset_name_val
- )
- # settings.save_settings() is called inside settings.save_settings if auto_save is true,
- # but for the manual button, we ensure it saves regardless of the auto_save flag's previous state.
- # The call above to settings.save_settings already handles writing to disk.
- return "Settings saved successfully! Restart required for theme change.
"
- except Exception as e:
- return f"Error saving settings: {str(e)}
"
-
- def handle_individual_setting_change(key, value, setting_name_for_ui):
- """Called by .change() and .submit() events of individual setting components."""
- if key == "auto_save_settings":
- # For the "auto_save_settings" checkbox itself:
- # 1. Update its value directly in the settings object in memory.
- # This bypasses the conditional save logic within settings.set() for this specific action.
- settings.settings[key] = value
- # 2. Force a save of all settings to disk. This will be correct because either:
- # - auto_save_settings is turning True: so all changes already in memory need to be saved now.
- # - auto_save_settings turning False from True: prior changes already saved so only auto_save_settings will be saved.
- settings.save_settings()
- # 3. Provide feedback.
- if value is True:
- return f"'{setting_name_for_ui}' setting is now ON and saved.
"
- else:
- return f"'{setting_name_for_ui}' setting is now OFF and saved.
"
- else:
- # For all other settings:
- # Let settings.set() handle the auto-save logic based on the current "auto_save_settings" value.
- settings.set(key, value) # settings.set() will call save_settings() if auto_save is True
- if settings.get("auto_save_settings"): # Check the current state of auto_save
- return f"'{setting_name_for_ui}' setting auto-saved.
"
- else:
- return f"'{setting_name_for_ui}' setting changed (auto-save is off, click 'Save Settings').
"
-
- # REMOVE `cleanup_temp_folder` from the `inputs` list
- save_btn.click(
- fn=save_settings,
- inputs=[save_metadata, gpu_memory_preservation, mp4_crf, clean_up_videos, auto_cleanup_on_startup, latents_display_top, override_system_prompt, system_prompt_template, output_dir, metadata_dir, lora_dir, gradio_temp_dir, auto_save, theme_dropdown, startup_model_type_dropdown, startup_preset_name_dropdown],
- outputs=[status]
- ).then(
- # NEW: Update latents display layout after manual save
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
- def reset_system_prompt_template_value():
- return settings.default_settings["system_prompt_template"], False
-
- reset_system_prompt_btn.click(
- fn=reset_system_prompt_template_value,
- outputs=[system_prompt_template, override_system_prompt]
- ).then( # Trigger auto-save for the reset values if auto-save is on
- lambda val_template, val_override: handle_individual_setting_change("system_prompt_template", val_template, "System Prompt Template") or handle_individual_setting_change("override_system_prompt", val_override, "Override System Prompt"),
- inputs=[system_prompt_template, override_system_prompt], outputs=[status])
-
- def manual_cleanup_handler():
- """UI handler for the manual cleanup button."""
- # This directly calls the toolbox_processor's cleanup method and returns the summary string.
- summary = tb_processor.tb_clear_temporary_files()
- return summary
-
- cleanup_btn.click(
- fn=manual_cleanup_handler,
- inputs=None,
- outputs=[cleanup_output]
- )
-
- # Add .change handlers for auto-saving individual settings
- save_metadata.change(lambda v: handle_individual_setting_change("save_metadata", v, "Save Metadata"), inputs=[save_metadata], outputs=[status])
- gpu_memory_preservation.change(lambda v: handle_individual_setting_change("gpu_memory_preservation", v, "GPU Memory Preservation"), inputs=[gpu_memory_preservation], outputs=[status])
- mp4_crf.change(lambda v: handle_individual_setting_change("mp4_crf", v, "MP4 Compression"), inputs=[mp4_crf], outputs=[status])
- clean_up_videos.change(lambda v: handle_individual_setting_change("clean_up_videos", v, "Clean Up Videos"), inputs=[clean_up_videos], outputs=[status])
-
- # NEW: auto-cleanup temp files on startup checkbox
- auto_cleanup_on_startup.change(lambda v: handle_individual_setting_change("auto_cleanup_on_startup", v, "Auto Cleanup on Startup"), inputs=[auto_cleanup_on_startup], outputs=[status])
-
- # NEW: latents display position setting
- latents_display_top.change(lambda v: handle_individual_setting_change("latents_display_top", v, "Latents Display Position"), inputs=[latents_display_top], outputs=[status])
-
-
-
- # Connect the latents display setting to layout updates
- def update_latents_display_layout_from_checkbox(display_top):
- """Update layout when checkbox changes - uses the checkbox value directly"""
- if display_top:
- return (
- gr.update(visible=True), # top_preview_row
- gr.update(visible=False, value=None) # preview_image (right column)
- )
- else:
- return (
- gr.update(visible=False), # top_preview_row
- gr.update(visible=True) # preview_image (right column)
- )
-
- latents_display_top.change(
- fn=update_latents_display_layout_from_checkbox,
- inputs=[latents_display_top],
- outputs=[top_preview_row, preview_image]
- )
-
- override_system_prompt.change(lambda v: handle_individual_setting_change("override_system_prompt", v, "Override System Prompt"), inputs=[override_system_prompt], outputs=[status])
- # Using .blur for text changes so they are processed after the user finishes, not on every keystroke
- system_prompt_template.blur(lambda v: handle_individual_setting_change("system_prompt_template", v, "System Prompt Template"), inputs=[system_prompt_template], outputs=[status])
- # reset_system_prompt_btn # is handled separately above, on click
-
- # Using .blur for text changes so they are processed after the user finishes, not on every keystroke
- output_dir.blur(lambda v: handle_individual_setting_change("output_dir", v, "Output Directory"), inputs=[output_dir], outputs=[status])
- metadata_dir.blur(lambda v: handle_individual_setting_change("metadata_dir", v, "Metadata Directory"), inputs=[metadata_dir], outputs=[status])
- lora_dir.blur(lambda v: handle_individual_setting_change("lora_dir", v, "LoRA Directory"), inputs=[lora_dir], outputs=[status])
- gradio_temp_dir.blur(lambda v: handle_individual_setting_change("gradio_temp_dir", v, "Gradio Temporary Directory"), inputs=[gradio_temp_dir], outputs=[status])
-
- auto_save.change(lambda v: handle_individual_setting_change("auto_save_settings", v, "Auto-save Settings"), inputs=[auto_save], outputs=[status])
- theme_dropdown.change(lambda v: handle_individual_setting_change("gradio_theme", v, "Theme"), inputs=[theme_dropdown], outputs=[status])
-
- # Event handlers for startup settings
- def update_startup_preset_dropdown_choices(selected_startup_model_type_from_ui):
- if not selected_startup_model_type_from_ui or selected_startup_model_type_from_ui == "None":
- return gr.update(choices=[], value=None)
-
- loaded_presets_for_model = load_presets(selected_startup_model_type_from_ui)
-
- # Get the preset name that was saved for the *previous* model type
- current_saved_startup_preset = settings.get("startup_preset_name")
-
- # Default to None
- value_to_select = None
- # If the previously saved preset name exists for the new model, select it
- if current_saved_startup_preset and current_saved_startup_preset in loaded_presets_for_model:
- value_to_select = current_saved_startup_preset
-
- return gr.update(choices=loaded_presets_for_model, value=value_to_select)
-
- startup_model_type_dropdown.change(
- fn=lambda v: handle_individual_setting_change("startup_model_type", v, "Startup Model Type"),
- inputs=[startup_model_type_dropdown], outputs=[status]
- ).then( # Chain the update to the preset dropdown
- fn=update_startup_preset_dropdown_choices, inputs=[startup_model_type_dropdown], outputs=[startup_preset_name_dropdown])
- startup_preset_name_dropdown.change(lambda v: handle_individual_setting_change("startup_preset_name", v, "Startup Preset Name"), inputs=[startup_preset_name_dropdown], outputs=[status])
-
- # --- Event Handlers and Connections (Now correctly indented) ---
-
- # --- Connect Monitoring ---
- # Auto-check for current job on page load and job change
- def check_for_current_job():
- # This function will be called when the interface loads
- # It will check if there's a current job in the queue and update the UI
- with job_queue.lock:
- current_job = job_queue.current_job
- if current_job:
- # Return all the necessary information to update the preview windows
- job_id = current_job.id
- result = current_job.result
- preview = current_job.progress_data.get('preview') if current_job.progress_data else None
- desc = current_job.progress_data.get('desc', '') if current_job.progress_data else ''
- html = current_job.progress_data.get('html', '') if current_job.progress_data else ''
-
- # Also trigger the monitor_job function to start monitoring this job
- print(f"Auto-check found current job {job_id}, triggering monitor_job")
- return job_id, result, preview, preview, desc, html
- return None, None, None, None, '', ''
-
- # Auto-check for current job on page load and handle handoff between jobs.
- def check_for_current_job_and_monitor():
- # This function is now the key to the handoff.
- # It finds the current job and returns its ID, which will trigger the monitor.
- job_id, result, preview, top_preview, desc, html = check_for_current_job()
- # We also need to get fresh stats at the same time.
- queue_status_data, queue_stats_text = update_stats()
- # Return everything needed to update the UI atomically.
- return job_id, result, preview, top_preview, desc, html, queue_status_data, queue_stats_text
-
- # Connect the main process function (wrapper for adding to queue)
- def process_with_queue_update(model_type_arg, *args):
- # Call update_stats to get both queue_status_data and queue_stats_text
- queue_status_data, queue_stats_text = update_stats() # MODIFIED
-
- # Extract all arguments (ensure order matches inputs lists)
- # The order here MUST match the order in the `ips` list.
- # RT_BORG: Global settings gpu_memory_preservation, mp4_crf, save_metadata removed from direct args.
- (input_image_arg,
- input_video_arg,
- end_frame_image_original_arg,
- end_frame_strength_original_arg,
- prompt_text_arg,
- n_prompt_arg,
- seed_arg, # the seed value
- randomize_seed_arg, # the boolean value of the checkbox
- total_second_length_arg,
- latent_window_size_arg,
- steps_arg,
- cfg_arg,
- gs_arg,
- rs_arg,
- cache_type_arg,
- teacache_num_steps_arg,
- teacache_rel_l1_thresh_arg,
- magcache_threshold_arg,
- magcache_max_consecutive_skips_arg,
- magcache_retention_ratio_arg,
- blend_sections_arg,
- latent_type_arg,
- clean_up_videos_arg, # UI checkbox from Generate tab
- selected_loras_arg,
- resolutionW_arg, resolutionH_arg,
- combine_with_source_arg,
- num_cleaned_frames_arg,
- lora_names_states_arg, # This is from lora_names_states (gr.State)
- *lora_slider_values_tuple # Remaining args are LoRA slider values
- ) = args
- # DO NOT parse the prompt here. Parsing happens once in the worker.
-
- # Determine the model type to send to the backend
- backend_model_type = model_type_arg # model_type_arg is the UI selection
- if model_type_arg == "Video with Endframe":
- backend_model_type = "Video" # The backend "Video" model_type handles with and without endframe
-
- # Use the appropriate input based on model type
- is_ui_video_model = is_video_model(model_type_arg)
- input_data = input_video_arg if is_ui_video_model else input_image_arg
-
- # Define actual end_frame params to pass to backend
- actual_end_frame_image_for_backend = None
- actual_end_frame_strength_for_backend = 1.0 # Default strength
-
- if model_type_arg == "Original with Endframe" or model_type_arg == "F1 with Endframe" or model_type_arg == "Video with Endframe":
- actual_end_frame_image_for_backend = end_frame_image_original_arg
- actual_end_frame_strength_for_backend = end_frame_strength_original_arg
-
- # Get the input video path for Video model
- input_image_path = None
- if is_ui_video_model and input_video_arg is not None:
- # For Video models, input_video contains the path to the video file
- input_image_path = input_video_arg
-
- # Use the current seed value as is for this job
- # Call the process function with all arguments
- # Pass the backend_model_type and the ORIGINAL prompt_text string to the backend process function
- result = process_fn(backend_model_type, input_data, actual_end_frame_image_for_backend, actual_end_frame_strength_for_backend,
- prompt_text_arg, n_prompt_arg, seed_arg, total_second_length_arg,
- latent_window_size_arg, steps_arg, cfg_arg, gs_arg, rs_arg,
- cache_type_arg == 'TeaCache', teacache_num_steps_arg, teacache_rel_l1_thresh_arg,
- cache_type_arg == 'MagCache', magcache_threshold_arg, magcache_max_consecutive_skips_arg, magcache_retention_ratio_arg,
- blend_sections_arg, latent_type_arg, clean_up_videos_arg, # clean_up_videos_arg is from UI
- selected_loras_arg, resolutionW_arg, resolutionH_arg,
- input_image_path,
- combine_with_source_arg,
- num_cleaned_frames_arg,
- lora_names_states_arg,
- *lora_slider_values_tuple
- )
- # If randomize_seed is checked, generate a new random seed for the next job
- new_seed_value = None
- if randomize_seed_arg:
- new_seed_value = random.randint(0, 21474)
- print(f"Generated new seed for next job: {new_seed_value}")
-
- # Create the button update for start_button WITHOUT interactive=True.
- # The interactivity will be set by update_start_button_state later in the chain.
- start_button_update_after_add = gr.update(value="🚀 Add to Queue")
-
- # If a job ID was created, automatically start monitoring it and update queue
- if result and result[1]: # Check if job_id exists in results
- job_id = result[1]
- # queue_status_data = update_queue_status_fn() # OLD: update_stats now called earlier
- # Call update_stats again AFTER the job is added to get the freshest stats
- queue_status_data, queue_stats_text = update_stats()
-
-
- # Add the new seed value to the results if randomize is checked
- if new_seed_value is not None:
- # Use result[6] directly for end_button to preserve its value. Add gr.update() for video_input_required_message.
- return [result[0], job_id, result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, new_seed_value, gr.update()]
- else:
- # Use result[6] directly for end_button to preserve its value. Add gr.update() for video_input_required_message.
- return [result[0], job_id, result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, gr.update(), gr.update()]
-
- # If no job ID was created, still return the new seed if randomize is checked
- # Also, ensure we return the latest stats even if no job was created (e.g., error during param validation)
- queue_status_data, queue_stats_text = update_stats()
- if new_seed_value is not None:
- # Make sure to preserve the end_button update from result[6]
- return [result[0], result[1], result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, new_seed_value, gr.update()]
- else:
- # Make sure to preserve the end_button update from result[6]
- return [result[0], result[1], result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, gr.update(), gr.update()]
-
- # Custom end process function that ensures the queue is updated and changes button text
- def end_process_with_update():
- _ = end_process_fn() # Call the original end_process_fn
- # Now, get fresh stats for both queue table and toolbar
- queue_status_data, queue_stats_text = update_stats()
-
- # Don't try to get the new job ID immediately after cancellation
- # The monitor_job function will handle the transition to the next job
-
- # Change the cancel button text to "Cancelling..." and make it non-interactive
- # This ensures the button stays in this state until the job is fully cancelled
- return queue_status_data, queue_stats_text, gr.update(value="Cancelling...", interactive=False), gr.update(value=None)
-
- # MODIFIED handle_send_video_to_toolbox:
- def handle_send_video_to_toolbox(original_path_from_state): # Input is now the original path from gr.State
- print(f"Sending selected Outputs' video to Post-processing: {original_path_from_state}")
-
- if original_path_from_state and isinstance(original_path_from_state, str) and os.path.exists(original_path_from_state):
- # tb_target_video_input will now process the ORIGINAL path (e.g., "outputs/my_actual_video.mp4").
- return gr.update(value=original_path_from_state), gr.update(selected="toolbox_tab")
- else:
- print(f"No valid video path (from State) found to send. Path: {original_path_from_state}")
- return gr.update(), gr.update()
-
- send_to_toolbox_btn.click(
- fn=handle_send_video_to_toolbox,
- inputs=[selected_original_video_path_state], # INPUT IS NOW THE gr.State holding the ORIGINAL path
- outputs=[
- tb_target_video_input, # This is tb_input_video_component from toolbox_app.py
- main_tabs_component
- ]
- )
-
- # --- Inputs Lists ---
- # --- Inputs for all models ---
- ips = [
- input_image, # Corresponds to input_image_arg
- input_video, # Corresponds to input_video_arg
- end_frame_image_original, # Corresponds to end_frame_image_original_arg
- end_frame_strength_original,# Corresponds to end_frame_strength_original_arg
- prompt, # Corresponds to prompt_text_arg
- n_prompt, # Corresponds to n_prompt_arg
- seed, # Corresponds to seed_arg
- randomize_seed, # Corresponds to randomize_seed_arg
- total_second_length, # Corresponds to total_second_length_arg
- latent_window_size, # Corresponds to latent_window_size_arg
- steps, # Corresponds to steps_arg
- cfg, # Corresponds to cfg_arg
- gs, # Corresponds to gs_arg
- rs, # Corresponds to rs_arg
- cache_type, # Corresponds to cache_type_arg
- teacache_num_steps, # Corresponds to teacache_num_steps_arg
- teacache_rel_l1_thresh, # Corresponds to teacache_rel_l1_thresh_arg
- magcache_threshold, # Corresponds to magcache_threshold_arg
- magcache_max_consecutive_skips, # Corresponds to magcache_max_consecutive_skips_arg
- magcache_retention_ratio, # Corresponds to magcache_retention_ratio_arg
- blend_sections, # Corresponds to blend_sections_arg
- latent_type, # Corresponds to latent_type_arg
- clean_up_videos, # Corresponds to clean_up_videos_arg (UI checkbox)
- lora_selector, # Corresponds to selected_loras_arg
- resolutionW, # Corresponds to resolutionW_arg
- resolutionH, # Corresponds to resolutionH_arg
- combine_with_source, # Corresponds to combine_with_source_arg
- num_cleaned_frames, # Corresponds to num_cleaned_frames_arg
- lora_names_states # Corresponds to lora_names_states_arg
- ]
- # Add LoRA sliders to the input list
- ips.extend([lora_sliders[lora] for lora in lora_names])
-
-
- # --- Connect Buttons ---
- def handle_start_button(selected_model, *args):
- # For other model types, use the regular process function
- return process_with_queue_update(selected_model, *args)
-
- def handle_batch_add_to_queue(*args):
- # The last argument will be the list of files from batch_input_images
- batch_files = args[-1]
- if not batch_files or not isinstance(batch_files, list):
- print("No batch images provided.")
- return
-
- print(f"Starting batch processing for {len(batch_files)} images.")
-
- # Reconstruct the arguments for the single process function, excluding the batch files list
- single_job_args = list(args[:-1])
-
- # The first argument to process_with_queue_update is model_type
- model_type_arg = single_job_args.pop(0)
-
- # Keep track of the seed
- current_seed = single_job_args[6] # seed is the 7th element in the ips list
- randomize_seed_arg = single_job_args[7] # randomize_seed is the 8th
-
- for image_path in batch_files:
- # --- FIX IS HERE ---
- # Load the image from the path into a NumPy array
- try:
- pil_image = Image.open(image_path).convert("RGB")
- numpy_image = np.array(pil_image)
- except Exception as e:
- print(f"Error loading batch image {image_path}: {e}. Skipping.")
- continue
- # --- END OF FIX ---
-
- # Replace the single input_image argument with the loaded NumPy image
- current_job_args = single_job_args[:]
- current_job_args[0] = numpy_image # Use the loaded numpy_image
- current_job_args[6] = current_seed # Set the seed for the current job
-
- # Call the original processing function with the modified arguments
- process_with_queue_update(model_type_arg, *current_job_args)
-
- # If randomize seed is checked, generate a new one for the next image
- if randomize_seed_arg:
- current_seed = random.randint(0, 21474)
-
- print("Batch processing complete. All jobs added to the queue.")
-
- # Validation ensures the start button is only enabled when appropriate
- def update_start_button_state(*args):
- """
- Validation fails if a video model is selected and no input video is provided.
- Updates the start button interactivity and validation message visibility.
- Handles variable inputs from different Gradio event chains.
- """
- # The required values are the last two arguments provided by the Gradio event
- if len(args) >= 2:
- selected_model = args[-2]
- input_video_value = args[-1]
- else:
- # Fallback or error handling if not enough arguments are received
- # This might happen if the event is triggered in an unexpected way
- print(f"Warning: update_start_button_state received {len(args)} args, expected at least 2.")
- # Default to a safe state (button disabled)
- return gr.Button(value="❌ Error", interactive=False), gr.update(visible=True)
-
- video_provided = input_video_value is not None
-
- if is_video_model(selected_model) and not video_provided:
- # Video model selected, but no video provided
- return gr.Button(value="❌ Missing Video", interactive=False), gr.update(visible=True)
- else:
- # Either not a video model, or video model selected and video provided
- return gr.update(value="🚀 Add to Queue", interactive=True), gr.update(visible=False)
- # Function to update button state before processing
- def update_button_before_processing(selected_model, *args):
- # First update the button to show "Adding..." and disable it
- # Also return current stats so they don't get blanked out during the "Adding..." phase
- qs_data, qs_text = update_stats()
- return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(value="⏳ Adding...", interactive=False), gr.update(), qs_data, qs_text, gr.update(), gr.update() # Added update for video_input_required_message
-
- # Connect the start button to first update its state
- start_button.click(
- fn=update_button_before_processing,
- inputs=[model_type] + ips,
- outputs=[result_video, current_job_id, preview_image, top_preview_image, progress_desc, progress_bar, start_button, end_button, queue_status, queue_stats_display, seed, video_input_required_message]
- ).then(
- # Then process the job
- fn=handle_start_button,
- inputs=[model_type] + ips,
- outputs=[result_video, current_job_id, preview_image, progress_desc, progress_bar, start_button, end_button, queue_status, queue_stats_display, seed, video_input_required_message] # Added video_input_required_message
- ).then( # Ensure validation is re-checked after job processing completes
- fn=update_start_button_state,
- inputs=[model_type, input_video], # Current values of model_type and input_video
- outputs=[start_button, video_input_required_message]
- )
-
- def show_batch_gallery(files):
- return gr.update(value=files, visible=True) if files else gr.update(visible=False)
-
- batch_input_images.change(
- fn=show_batch_gallery,
- inputs=[batch_input_images],
- outputs=[batch_input_gallery]
- )
-
- # We need to gather all the same inputs as the single 'Add to Queue' button, plus the new file input
- batch_ips = [model_type] + ips + [batch_input_images]
-
- add_batch_to_queue_btn.click(
- fn=handle_batch_add_to_queue,
- inputs=batch_ips,
- outputs=None # No direct output updates from this button
- ).then(
- fn=update_stats, # Refresh the queue stats in the UI
- inputs=None,
- outputs=[queue_status, queue_stats_display]
- ).then(
- # This new block checks for a running job and updates the monitor UI
- fn=check_for_current_job,
- inputs=None,
- outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
- ).then(
- # NEW: Update latents display layout after loading queue to ensure correct visibility
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
- # --- START OF REFACTORED XY PLOT EVENT WIRING ---
- # Get the process button from the created components
- xy_plot_process_btn = xy_plot_components["process_btn"]
-
- # Prepare the process function with its static dependencies (job_queue, settings)
- fn_xy_process_with_deps = functools.partial(xy_plot_process, job_queue, settings)
-
- # Construct the full list of inputs for the click handler in the correct order
- c = xy_plot_components
- xy_plot_input_components = [
- c["model_type"], c["input_image"], c["end_frame_image_original"],
- c["end_frame_strength_original"], c["latent_type"], c["prompt"],
- c["blend_sections"], c["steps"], c["total_second_length"],
- resolutionW, resolutionH, # The components from the main UI
- c["seed"], c["randomize_seed"],
- c["use_teacache"], c["teacache_num_steps"], c["teacache_rel_l1_thresh"],
- c["use_magcache"], c["magcache_threshold"], c["magcache_max_consecutive_skips"], c["magcache_retention_ratio"],
- c["latent_window_size"], c["cfg"], c["gs"], c["rs"],
- c["gpu_memory_preservation"], c["mp4_crf"],
- c["axis_x_switch"], c["axis_x_value_text"], c["axis_x_value_dropdown"],
- c["axis_y_switch"], c["axis_y_value_text"], c["axis_y_value_dropdown"],
- c["axis_z_switch"], c["axis_z_value_text"], c["axis_z_value_dropdown"],
- c["lora_selector"]
- ]
- # LoRA sliders are in a dictionary, so we add their values to the list
- xy_plot_input_components.extend(c["lora_sliders"].values())
-
- # Wire the click handler for the XY Plot button
- xy_plot_process_btn.click(
- fn=fn_xy_process_with_deps,
- inputs=xy_plot_input_components,
- outputs=[xy_plot_status, xy_plot_output]
- ).then(
- fn=update_stats,
- inputs=None,
- outputs=[queue_status, queue_stats_display]
- ).then(
- fn=check_for_current_job,
- inputs=None,
- outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
- ).then(
- # NEW: Update latents display layout after XY plot to ensure correct visibility
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
- # --- END OF REFACTORED XY PLOT EVENT WIRING ---
-
-
-
- # MODIFIED: on_model_type_change to handle new "XY Plot" option
- def on_model_type_change(selected_model):
- is_xy_plot = selected_model == "XY Plot"
- is_ui_video_model_flag = is_video_model(selected_model)
- shows_end_frame = selected_model in ["Original with Endframe", "Video with Endframe"]
-
- return (
- gr.update(visible=not is_xy_plot), # standard_generation_group
- gr.update(visible=is_xy_plot), # xy_group
- gr.update(visible=not is_xy_plot and not is_ui_video_model_flag), # image_input_group
- gr.update(visible=not is_xy_plot and is_ui_video_model_flag), # video_input_group
- gr.update(visible=not is_xy_plot and shows_end_frame), # end_frame_group_original
- gr.update(visible=not is_xy_plot and shows_end_frame), # end_frame_slider_group
- gr.update(visible=not is_xy_plot), # start_button
- gr.update(visible=is_xy_plot) # xy_plot_process_btn
- )
-
- # Model change listener
- model_type.change(
- fn=on_model_type_change,
- inputs=model_type,
- outputs=[
- standard_generation_group,
- xy_group,
- image_input_group,
- video_input_group,
- end_frame_group_original,
- end_frame_slider_group,
- start_button,
- xy_plot_process_btn # This is the button returned from the dictionary
- ]
- ).then( # Also trigger validation after model type changes
- fn=update_start_button_state,
- inputs=[model_type, input_video],
- outputs=[start_button, video_input_required_message]
- )
-
- # Connect input_video change to the validation function
- input_video.change(
- fn=update_start_button_state,
- inputs=[model_type, input_video],
- outputs=[start_button, video_input_required_message]
- )
- # Also trigger validation when video is cleared
- input_video.clear(
- fn=update_start_button_state,
- inputs=[model_type, input_video],
- outputs=[start_button, video_input_required_message]
- )
-
-
-
- # Auto-monitor the current job when job_id changes
- current_job_id.change(
- fn=monitor_fn,
- inputs=[current_job_id],
- outputs=[result_video, preview_image, top_preview_image, progress_desc, progress_bar, start_button, end_button]
- ).then(
- fn=update_stats, # When a monitor finishes, always update the stats.
- inputs=None,
- outputs=[queue_status, queue_stats_display]
- ).then( # re-validate button state
- fn=update_start_button_state,
- inputs=[model_type, input_video],
- outputs=[start_button, video_input_required_message]
- ).then(
- # NEW: Update latents display layout after monitoring to ensure correct visibility
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
- # The "end_button" (Cancel Job) is the trigger for the next job's monitor.
- # When a job is cancelled, we check for the next one.
- end_button.click(
- fn=end_process_with_update,
- outputs=[queue_status, queue_stats_display, end_button, current_job_id]
- ).then(
- fn=check_for_current_job_and_monitor,
- inputs=[],
- outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar, queue_status, queue_stats_display]
- ).then(
- # NEW: Update latents display layout after job handoff to ensure correct visibility
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
- load_queue_button.click(
- fn=load_queue_from_json,
- inputs=[],
- outputs=[queue_status, queue_stats_display]
- ).then( # ADD THIS .then() CLAUSE
- fn=check_for_current_job,
- inputs=[],
- outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
- ).then(
- # NEW: Update latents display layout after loading queue to ensure correct visibility
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
- import_queue_file.change(
- fn=import_queue_from_file,
- inputs=[import_queue_file],
- outputs=[queue_status, queue_stats_display]
- ).then( # ADD THIS .then() CLAUSE
- fn=check_for_current_job,
- inputs=[],
- outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
- ).then(
- # NEW: Update latents display layout after importing queue to ensure correct visibility
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
-
- # --- Connect Queue Refresh ---
- # The update_stats function is now defined much earlier.
-
- # REMOVED: refresh_stats_btn.click - Toolbar refresh button is no longer needed
- # refresh_stats_btn.click(
- # fn=update_stats,
- # inputs=None,
- # outputs=[queue_status, queue_stats_display]
- # )
-
- # Set up auto-refresh for queue status
- # Instead of using a timer with 'every' parameter, we'll use the queue refresh button
- # and rely on manual refreshes. The user can click the refresh button in the toolbar
- # to update the stats.
-
- # --- Connect LoRA UI ---
- # Function to update slider visibility based on selection
- def update_lora_sliders(selected_loras):
- updates = []
- # Suppress dummy LoRA from workaround for the single lora bug.
- # Filter out the dummy LoRA for display purposes in the dropdown
- actual_selected_loras_for_display = [lora for lora in selected_loras if lora != DUMMY_LORA_NAME]
- updates.append(gr.update(value=actual_selected_loras_for_display)) # First update is for the dropdown itself
-
- # Need to handle potential missing keys if lora_names changes dynamically
- # lora_names is from the create_interface scope
- for lora_name_key in lora_names: # Iterate using lora_names to maintain order
- if lora_name_key == DUMMY_LORA_NAME: # Check for dummy LoRA
- updates.append(gr.update(visible=False))
- else:
- # Visibility of sliders should be based on actual_selected_loras_for_display
- updates.append(gr.update(visible=(lora_name_key in actual_selected_loras_for_display)))
- return updates # This list will be correctly ordered
-
- # Connect the dropdown to the sliders
- lora_selector.change(
- fn=update_lora_sliders,
- inputs=[lora_selector],
- outputs=[lora_selector] + [lora_sliders[lora] for lora in lora_names if lora in lora_sliders]
- )
-
- def apply_preset(preset_name, model_type):
- if not preset_name:
- # Create a list of empty updates matching the number of components
- return [gr.update()] * len(ui_components)
-
- with open(PRESET_FILE, 'r') as f:
- data = json.load(f)
- preset = data.get(model_type, {}).get(preset_name, {})
-
- # Initialize updates for all components
- updates = {key: gr.update() for key in ui_components.keys()}
-
- # Update components based on the preset
- for key, value in preset.items():
- if key in updates:
- updates[key] = gr.update(value=value)
-
- # Handle LoRA sliders specifically
- if 'lora_values' in preset and isinstance(preset['lora_values'], dict):
- lora_values_dict = preset['lora_values']
- for lora_name, lora_value in lora_values_dict.items():
- if lora_name in updates:
- updates[lora_name] = gr.update(value=lora_value)
-
- # Convert the dictionary of updates to a list in the correct order
- return [updates[key] for key in ui_components.keys()]
-
- def save_preset(preset_name, model_type, *args):
- if not preset_name:
- return gr.update()
-
- # Ensure the directory exists
- os.makedirs(os.path.dirname(PRESET_FILE), exist_ok=True)
-
- if not os.path.exists(PRESET_FILE):
- with open(PRESET_FILE, 'w') as f:
- json.dump({}, f)
-
- with open(PRESET_FILE, 'r') as f:
- data = json.load(f)
-
- if model_type not in data:
- data[model_type] = {}
-
- keys = list(ui_components.keys())
-
- # Create a dictionary from the passed arguments
- args_dict = {keys[i]: args[i] for i in range(len(keys))}
-
- # Build the preset data from the arguments dictionary
- preset_data = {key: args_dict[key] for key in ui_components.keys() if key not in lora_sliders}
-
- # Handle LoRA values separately
- selected_loras = args_dict.get("lora_selector", [])
- lora_values = {}
- for lora_name in selected_loras:
- if lora_name in args_dict:
- lora_values[lora_name] = args_dict[lora_name]
-
- preset_data['lora_values'] = lora_values
-
- # Remove individual lora sliders from the top-level preset data
- for lora_name in lora_sliders:
- if lora_name in preset_data:
- del preset_data[lora_name]
-
- data[model_type][preset_name] = preset_data
-
- with open(PRESET_FILE, 'w') as f:
- json.dump(data, f, indent=2)
-
- return gr.update(choices=load_presets(model_type), value=preset_name)
-
- def delete_preset(preset_name, model_type):
- if not preset_name:
- return gr.update(), gr.update(visible=True), gr.update(visible=False)
-
- with open(PRESET_FILE, 'r') as f:
- data = json.load(f)
-
- if model_type in data and preset_name in data[model_type]:
- del data[model_type][preset_name]
-
- with open(PRESET_FILE, 'w') as f:
- json.dump(data, f, indent=2)
-
- return gr.update(choices=load_presets(model_type), value=None), gr.update(visible=True), gr.update(visible=False)
-
- # --- Connect Preset UI ---
- # Without this refresh, if you define a new preset for the Startup Model Type, and then try to select it in settings, it won't show up.
- def refresh_settings_tab_startup_presets_if_needed(generate_tab_model_type_value, settings_tab_startup_model_type_value):
- # generate_tab_model_type_value is the model for which a preset was just saved
- # settings_tab_startup_model_type_value is the current selection in the startup model dropdown on settings tab
- if generate_tab_model_type_value == settings_tab_startup_model_type_value and settings_tab_startup_model_type_value != "None":
- return update_startup_preset_dropdown_choices(settings_tab_startup_model_type_value)
- return gr.update()
-
- ui_components = {
- # Prompts
- "prompt": prompt,
- "n_prompt": n_prompt,
- "blend_sections": blend_sections,
- # Basic Params
- "steps": steps,
- "total_second_length": total_second_length,
- "resolutionW": resolutionW,
- "resolutionH": resolutionH,
- "seed": seed,
- "randomize_seed": randomize_seed,
- # Advanced Params
- "gs": gs,
- "cfg": cfg,
- "rs": rs,
- "latent_window_size": latent_window_size,
- # Cache type (Mag/Tea/None)
- "cache_type": cache_type,
- # TeaCache
- "teacache_num_steps": teacache_num_steps,
- "teacache_rel_l1_thresh": teacache_rel_l1_thresh,
- # MagCache
- "magcache_threshold": magcache_threshold,
- "magcache_max_consecutive_skips": magcache_max_consecutive_skips,
- "magcache_retention_ratio": magcache_retention_ratio,
- # Input Options
- "latent_type": latent_type,
- "end_frame_strength_original": end_frame_strength_original,
- # Video Specific
- "combine_with_source": combine_with_source,
- "num_cleaned_frames": num_cleaned_frames,
- # LoRAs
- "lora_selector": lora_selector,
- **lora_sliders
- }
-
- model_type.change(
- fn=lambda mt: (gr.update(choices=load_presets(mt)), gr.update(label=f"{mt} Presets")),
- inputs=[model_type],
- outputs=[preset_dropdown, preset_accordion]
- )
-
- preset_dropdown.select(
- fn=apply_preset,
- inputs=[preset_dropdown, model_type],
- outputs=list(ui_components.values())
- ).then(
- lambda name: name,
- inputs=[preset_dropdown],
- outputs=[preset_name_textbox]
- )
-
- save_preset_button.click(
- fn=save_preset,
- inputs=[preset_name_textbox, model_type, *list(ui_components.values())],
- outputs=[preset_dropdown] # preset_dropdown is on Generate tab
- ).then(
- fn=refresh_settings_tab_startup_presets_if_needed,
- inputs=[model_type, startup_model_type_dropdown], # model_type (Generate tab), startup_model_type_dropdown (Settings tab)
- outputs=[startup_preset_name_dropdown] # startup_preset_name_dropdown (Settings tab)
- )
-
- def show_delete_confirmation():
- return gr.update(visible=False), gr.update(visible=True)
-
- def hide_delete_confirmation():
- return gr.update(visible=True), gr.update(visible=False)
-
- delete_preset_button.click(
- fn=show_delete_confirmation,
- outputs=[save_preset_button, confirm_delete_row]
- )
-
- confirm_delete_no_btn.click(
- fn=hide_delete_confirmation,
- outputs=[save_preset_button, confirm_delete_row]
- )
-
- confirm_delete_yes_btn.click(
- fn=delete_preset,
- inputs=[preset_dropdown, model_type],
- outputs=[preset_dropdown, save_preset_button, confirm_delete_row]
- )
-
- # --- Definition of apply_startup_settings (AFTER ui_components and apply_preset are defined) ---
- # This function needs access to `settings`, `model_type` (Generate tab Radio),
- # `preset_dropdown` (Generate tab Dropdown), `preset_name_textbox` (Generate tab Textbox),
- # `ui_components` (dict of all other UI elements), `load_presets`, and `apply_preset`.
- # All these are available in the scope of `create_interface`.
- def apply_startup_settings():
- startup_model_val = settings.get("startup_model_type", "None")
- startup_preset_val = settings.get("startup_preset_name", None)
-
- # Default updates (no change)
- model_type_update = gr.update()
- preset_dropdown_update = gr.update()
- preset_name_textbox_update = gr.update()
-
- # ui_components is now defined
- ui_components_updates_list = [gr.update() for _ in ui_components]
-
- if startup_model_val and startup_model_val != "None":
- model_type_update = gr.update(value=startup_model_val)
-
- presets_for_startup_model = load_presets(startup_model_val) # load_presets is defined earlier
- preset_dropdown_update = gr.update(choices=presets_for_startup_model)
- preset_name_textbox_update = gr.update(value="")
-
- if startup_preset_val and startup_preset_val in presets_for_startup_model:
- preset_dropdown_update = gr.update(choices=presets_for_startup_model, value=startup_preset_val)
- preset_name_textbox_update = gr.update(value=startup_preset_val)
-
- # apply_preset is now defined
- ui_components_updates_list = apply_preset(startup_preset_val, startup_model_val)
-
- # NEW: Ensure latents_display_top checkbox reflects the current setting
- latents_display_top_update = gr.update(value=get_latents_display_top())
-
- return tuple([model_type_update, preset_dropdown_update, preset_name_textbox_update] + ui_components_updates_list + [latents_display_top_update])
-
-
- # --- Auto-refresh for Toolbar System Stats Monitor (Timer) ---
- main_toolbar_system_stats_timer = gr.Timer(2, active=True)
-
- main_toolbar_system_stats_timer.tick(
- fn=tb_get_formatted_toolbar_stats, # Function imported from toolbox_app.py
- inputs=None,
- outputs=[ # Target the Textbox components
- toolbar_ram_display_component,
- toolbar_vram_display_component,
- toolbar_gpu_display_component
- ]
- )
-
- # --- Connect Metadata Loading ---
- # Function to load metadata from JSON file
- def load_metadata_from_json(json_path):
- # Define the total number of output components to handle errors gracefully
- num_outputs = 20 + len(lora_sliders)
-
- if not json_path:
- # Return empty updates for all components if no file is provided
- return [gr.update()] * num_outputs
-
- try:
- with open(json_path, 'r') as f:
- metadata = json.load(f)
-
- # Extract values from metadata with defaults
- prompt_val = metadata.get('prompt')
- n_prompt_val = metadata.get('negative_prompt')
- seed_val = metadata.get('seed')
- steps_val = metadata.get('steps')
- total_second_length_val = metadata.get('total_second_length')
- end_frame_strength_val = metadata.get('end_frame_strength')
- model_type_val = metadata.get('model_type')
- lora_weights = metadata.get('loras', {})
- latent_window_size_val = metadata.get('latent_window_size')
- resolutionW_val = metadata.get('resolutionW')
- resolutionH_val = metadata.get('resolutionH')
- blend_sections_val = metadata.get('blend_sections')
- # Determine cache_type from metadata, with fallback for older formats
- cache_type_val = metadata.get('cache_type')
- if cache_type_val is None:
- use_magcache = metadata.get('use_magcache', False)
- use_teacache = metadata.get('use_teacache', False)
- if use_magcache:
- cache_type_val = "MagCache"
- elif use_teacache:
- cache_type_val = "TeaCache"
- else:
- cache_type_val = "None"
- magcache_threshold_val = metadata.get('magcache_threshold')
- magcache_max_consecutive_skips_val = metadata.get('magcache_max_consecutive_skips')
- magcache_retention_ratio_val = metadata.get('magcache_retention_ratio')
- teacache_num_steps_val = metadata.get('teacache_num_steps')
- teacache_rel_l1_thresh_val = metadata.get('teacache_rel_l1_thresh')
- latent_type_val = metadata.get('latent_type')
- combine_with_source_val = metadata.get('combine_with_source')
-
- # Get the names of the selected LoRAs from the metadata
- selected_lora_names = list(lora_weights.keys())
-
- print(f"Loaded metadata from JSON: {json_path}")
- print(f"Model Type: {model_type_val}, Prompt: {prompt_val}, Seed: {seed_val}, LoRAs: {selected_lora_names}")
-
- # Create a list of UI updates
- updates = [
- gr.update(value=prompt_val) if prompt_val is not None else gr.update(),
- gr.update(value=n_prompt_val) if n_prompt_val is not None else gr.update(),
- gr.update(value=seed_val) if seed_val is not None else gr.update(),
- gr.update(value=steps_val) if steps_val is not None else gr.update(),
- gr.update(value=total_second_length_val) if total_second_length_val is not None else gr.update(),
- gr.update(value=end_frame_strength_val) if end_frame_strength_val is not None else gr.update(),
- gr.update(value=model_type_val) if model_type_val else gr.update(),
- gr.update(value=selected_lora_names) if selected_lora_names else gr.update(),
- gr.update(value=latent_window_size_val) if latent_window_size_val is not None else gr.update(),
- gr.update(value=resolutionW_val) if resolutionW_val is not None else gr.update(),
- gr.update(value=resolutionH_val) if resolutionH_val is not None else gr.update(),
- gr.update(value=blend_sections_val) if blend_sections_val is not None else gr.update(),
- gr.update(value=cache_type_val),
- gr.update(value=magcache_threshold_val),
- gr.update(value=magcache_max_consecutive_skips_val),
- gr.update(value=magcache_retention_ratio_val),
- gr.update(value=teacache_num_steps_val) if teacache_num_steps_val is not None else gr.update(),
- gr.update(value=teacache_rel_l1_thresh_val) if teacache_rel_l1_thresh_val is not None else gr.update(),
- gr.update(value=latent_type_val) if latent_type_val else gr.update(),
- gr.update(value=combine_with_source_val) if combine_with_source_val else gr.update(),
- ]
-
- # Update LoRA sliders based on loaded weights
- for lora in lora_names:
- if lora in lora_weights:
- updates.append(gr.update(value=lora_weights[lora], visible=True))
- else:
- # Hide sliders for LoRAs not in the metadata
- updates.append(gr.update(visible=False))
-
- return updates
-
- except Exception as e:
- print(f"Error loading metadata: {e}")
- import traceback
- traceback.print_exc()
- # Return empty updates for all components on error
- return [gr.update()] * num_outputs
-
-
- # Connect JSON metadata loader for Original tab
- json_upload.change(
- fn=load_metadata_from_json,
- inputs=[json_upload],
- outputs=[
- prompt,
- n_prompt,
- seed,
- steps,
- total_second_length,
- end_frame_strength_original,
- model_type,
- lora_selector,
- latent_window_size,
- resolutionW,
- resolutionH,
- blend_sections,
- cache_type,
- magcache_threshold,
- magcache_max_consecutive_skips,
- magcache_retention_ratio,
- teacache_num_steps,
- teacache_rel_l1_thresh,
- latent_type,
- combine_with_source
- ] + [lora_sliders[lora] for lora in lora_names]
- )
-
-
- # --- Helper Functions (defined within create_interface scope if needed by handlers) ---
- # Function to get queue statistics
- def get_queue_stats():
- try:
- # Get all jobs from the queue
- jobs = job_queue.get_all_jobs()
-
- # Count jobs by status
- status_counts = {
- "QUEUED": 0,
- "RUNNING": 0,
- "COMPLETED": 0,
- "FAILED": 0,
- "CANCELLED": 0
- }
-
- for job in jobs:
- if hasattr(job, 'status'):
- status = str(job.status) # Use str() for safety
- if status in status_counts:
- status_counts[status] += 1
-
- # Format the display text
- stats_text = f"Queue: {status_counts['QUEUED']} | Running: {status_counts['RUNNING']} | Completed: {status_counts['COMPLETED']} | Failed: {status_counts['FAILED']} | Cancelled: {status_counts['CANCELLED']}"
-
- return f"{stats_text}
"
-
- except Exception as e:
- print(f"Error getting queue stats: {e}")
- return "Error loading queue stats
"
-
- # Add footer with social links
- with gr.Row(elem_id="footer"):
- with gr.Column(scale=1):
- gr.HTML(f"""
-
- """)
-
- # Add CSS for footer
-
- # gr.HTML("""
- #
- # """)
-
- # --- Function to update latents display layout on interface load ---
- def update_latents_layout_on_load():
- """Update latents display layout based on saved setting when interface loads"""
- return create_latents_layout_update()
-
- # Connect the auto-check function to the interface load event
- block.load(
- fn=check_for_current_job_and_monitor, # Use the new combined function
- inputs=[],
- outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar, queue_status, queue_stats_display]
-
- ).then(
- fn=apply_startup_settings, # apply_startup_settings is now defined
- inputs=None,
- outputs=[model_type, preset_dropdown, preset_name_textbox] + list(ui_components.values()) + [latents_display_top] # ui_components is now defined
- ).then(
- fn=update_start_button_state, # Ensure button state is correct after startup settings
- inputs=[model_type, input_video],
- outputs=[start_button, video_input_required_message]
- ).then(
- # NEW: Update latents display layout based on saved setting
- fn=create_latents_layout_update,
- inputs=None,
- outputs=[top_preview_row, preview_image]
- )
-
- # --- Prompt Enhancer Connection ---
- def handle_enhance_prompt(current_prompt_text):
- """Calls the LLM enhancer and returns the updated text."""
- if not current_prompt_text:
- return ""
- print("UI: Enhance button clicked. Sending prompt to enhancer.")
- enhanced_text = enhance_prompt(current_prompt_text)
- print(f"UI: Received enhanced prompt: {enhanced_text}")
- return gr.update(value=enhanced_text)
-
- enhance_prompt_btn.click(
- fn=handle_enhance_prompt,
- inputs=[prompt],
- outputs=[prompt]
- )
-
- # --- Captioner Connection ---
- def handle_caption(input_image, prompt):
- """Calls the LLM enhancer and returns the updated text."""
- if input_image is None:
- return prompt # Return current prompt if no image is provided
- caption_text = caption_image(input_image)
- print(f"UI: Received caption: {caption_text}")
- return gr.update(value=caption_text)
-
- caption_btn.click(
- fn=handle_caption,
- inputs=[input_image, prompt],
- outputs=[prompt]
- )
-
- return block
-
-# --- Top-level Helper Functions (Used by Gradio callbacks, must be defined outside create_interface) ---
-
-def format_queue_status(jobs):
- """Format job data for display in the queue status table"""
- rows = []
- for job in jobs:
- created = time.strftime('%H:%M:%S', time.localtime(job.created_at)) if job.created_at else ""
- started = time.strftime('%H:%M:%S', time.localtime(job.started_at)) if job.started_at else ""
- completed = time.strftime('%H:%M:%S', time.localtime(job.completed_at)) if job.completed_at else ""
-
- # Calculate elapsed time
- elapsed_time = ""
- if job.started_at:
- if job.completed_at:
- start_datetime = datetime.datetime.fromtimestamp(job.started_at)
- complete_datetime = datetime.datetime.fromtimestamp(job.completed_at)
- elapsed_seconds = (complete_datetime - start_datetime).total_seconds()
- elapsed_time = f"{elapsed_seconds:.2f}s"
- else:
- # For running jobs, calculate elapsed time from now
- start_datetime = datetime.datetime.fromtimestamp(job.started_at)
- current_datetime = datetime.datetime.now()
- elapsed_seconds = (current_datetime - start_datetime).total_seconds()
- elapsed_time = f"{elapsed_seconds:.2f}s (running)"
-
- # Get generation type from job data
- generation_type = getattr(job, 'generation_type', 'Original')
-
- # Get thumbnail from job data and format it as HTML for display
- thumbnail = getattr(job, 'thumbnail', None)
- thumbnail_html = f'
' if thumbnail else ""
-
- rows.append([
- job.id[:6] + '...',
- generation_type,
- job.status.value,
- created,
- started,
- completed,
- elapsed_time,
- thumbnail_html # Add formatted thumbnail HTML to row data
- ])
- return rows
-
-# Create the queue status update function (wrapper around format_queue_status)
-def update_queue_status_with_thumbnails(): # Function name is now slightly misleading, but keep for now to avoid breaking clicks
- # This function is likely called by the refresh button and potentially the timer
- # It needs access to the job_queue object
- # Assuming job_queue is accessible globally or passed appropriately
- # For now, let's assume it's globally accessible as defined in studio.py
- # If not, this needs adjustment based on how job_queue is managed.
- try:
- # Need access to the global job_queue instance from studio.py
- # This might require restructuring or passing job_queue differently.
- # For now, assuming it's accessible (this might fail if run standalone)
- from __main__ import job_queue # Attempt to import from main script scope
-
- jobs = job_queue.get_all_jobs()
- for job in jobs:
- if job.status == JobStatus.PENDING:
- job.queue_position = job_queue.get_queue_position(job.id)
-
- if job_queue.current_job:
- job_queue.current_job.status = JobStatus.RUNNING
-
- return format_queue_status(jobs)
- except ImportError:
- print("Error: Could not import job_queue. Queue status update might fail.")
- return [] # Return empty list on error
- except Exception as e:
- print(f"Error updating queue status: {e}")
- return []
+import gradio as gr
+import time
+import datetime
+import random
+import json
+import os
+import shutil
+from typing import List, Dict, Any, Optional
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import base64
+import io
+import functools
+
+from modules.version import APP_VERSION, APP_VERSION_DISPLAY
+
+import subprocess
+import itertools
+import re
+from collections import defaultdict
+import imageio
+import imageio.plugins.ffmpeg
+import ffmpeg
+from diffusers_helper.utils import generate_timestamp
+
+from modules.video_queue import JobStatus, Job, JobType
+from modules.prompt_handler import get_section_boundaries, get_quick_prompts, parse_timestamped_prompt
+from modules.llm_enhancer import enhance_prompt
+from modules.llm_captioner import caption_image
+from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
+from diffusers_helper.bucket_tools import find_nearest_bucket
+from modules.pipelines.metadata_utils import create_metadata
+from modules import DUMMY_LORA_NAME # Import the constant
+
+from modules.toolbox_app import tb_processor
+from modules.toolbox_app import tb_create_video_toolbox_ui, tb_get_formatted_toolbar_stats
+from modules.xy_plot_ui import create_xy_plot_ui, xy_plot_process
+
+# Define the dummy LoRA name as a constant
+
+def create_interface(
+ process_fn,
+ monitor_fn,
+ end_process_fn,
+ update_queue_status_fn,
+ load_lora_file_fn,
+ job_queue,
+ settings,
+ default_prompt: str = '[1s: The person waves hello] [3s: The person jumps up and down] [5s: The person does a dance]',
+ lora_names: list = [],
+ lora_values: list = []
+):
+ """
+ Create the Gradio interface for the video generation application
+
+ Args:
+ process_fn: Function to process a new job
+ monitor_fn: Function to monitor an existing job
+ end_process_fn: Function to cancel the current job
+ update_queue_status_fn: Function to update the queue status display
+ default_prompt: Default prompt text
+ lora_names: List of loaded LoRA names
+
+ Returns:
+ Gradio Blocks interface
+ """
+ def is_video_model(model_type_value):
+ return model_type_value in ["Video", "Video with Endframe", "Video F1"]
+
+ # Add near the top of create_interface function, after the initial setup
+ def get_latents_display_top():
+ """Get current latents display preference - centralized access point"""
+ return settings.get("latents_display_top", False)
+
+ def create_latents_layout_update():
+ """Create a standardized layout update based on current setting"""
+ display_top = get_latents_display_top()
+ if display_top:
+ return (
+ gr.update(visible=True), # top_preview_row
+ gr.update(visible=False, value=None) # preview_image (right column)
+ )
+ else:
+ return (
+ gr.update(visible=False), # top_preview_row
+ gr.update(visible=True) # preview_image (right column)
+ )
+
+
+
+ # Get section boundaries and quick prompts
+ section_boundaries = get_section_boundaries()
+ quick_prompts = get_quick_prompts()
+
+ # --- Function to update queue stats (Moved earlier to resolve UnboundLocalError) ---
+ def update_stats(*args): # Accept any arguments and ignore them
+ # Get queue status data
+ queue_status_data = update_queue_status_fn()
+
+ # Get queue statistics for the toolbar display
+ jobs = job_queue.get_all_jobs()
+
+ # Count jobs by status
+ pending_count = 0
+ running_count = 0
+ completed_count = 0
+
+ for job in jobs:
+ if hasattr(job, 'status'):
+ status = str(job.status)
+ if status == "JobStatus.PENDING":
+ pending_count += 1
+ elif status == "JobStatus.RUNNING":
+ running_count += 1
+ elif status == "JobStatus.COMPLETED":
+ completed_count += 1
+
+ # Format the queue stats display text
+ queue_stats_text = f"Queue: {pending_count} | Running: {running_count} | Completed: {completed_count}
"
+
+ return queue_status_data, queue_stats_text
+
+ # --- Preset System Functions ---
+ PRESET_FILE = os.path.join(".framepack", "generation_presets.json")
+
+ def load_presets(model_type):
+ if not os.path.exists(PRESET_FILE):
+ return []
+ with open(PRESET_FILE, 'r') as f:
+ data = json.load(f)
+ return list(data.get(model_type, {}).keys())
+
+ # Create the interface
+ css = make_progress_bar_css()
+ css += """
+
+ .short-import-box, .short-import-box > div {
+ min-height: 40px !important;
+ height: 40px !important;
+ }
+ /* Image container styling - more aggressive approach */
+ .contain-image, .contain-image > div, .contain-image > div > img {
+ object-fit: contain !important;
+ }
+
+ #non-mirrored-video {
+ transform: scaleX(-1) !important;
+ }
+
+ /* Target all images in the contain-image class and its children */
+ .contain-image img,
+ .contain-image > div > img,
+ .contain-image * img {
+ object-fit: contain !important;
+ width: 100% !important;
+ height: 60vh !important;
+ max-height: 100% !important;
+ max-width: 100% !important;
+ }
+
+ /* Additional selectors to override Gradio defaults */
+ .gradio-container img,
+ .gradio-container .svelte-1b5oq5x,
+ .gradio-container [data-testid="image"] img {
+ object-fit: contain !important;
+ }
+
+ /* Toolbar styling */
+ #fixed-toolbar {
+ position: fixed;
+ top: 0;
+ left: 0;
+ width: 100vw;
+ z-index: 1000;
+ background: #333;
+ color: #fff;
+ padding: 0px 10px; /* Reduced top/bottom padding */
+ display: flex;
+ align-items: center;
+ gap: 8px;
+ box-shadow: 0 2px 8px rgba(0,0,0,0.1);
+ }
+
+ /* Responsive toolbar title */
+ .toolbar-title {
+ font-size: 1.4rem;
+ margin: 0;
+ color: white;
+ white-space: nowrap;
+ overflow: hidden;
+ text-overflow: ellipsis;
+ }
+
+ /* Toolbar Patreon link */
+ .toolbar-patreon {
+ margin: 0 0 0 20px;
+ color: white;
+ font-size: 0.9rem;
+ white-space: nowrap;
+ display: inline-block;
+ }
+ .toolbar-patreon a {
+ color: white;
+ text-decoration: none;
+ }
+ .toolbar-patreon a:hover {
+ text-decoration: underline;
+ }
+
+ /* Toolbar Version number */
+ .toolbar-version {
+ margin: 0 15px; /* Space around version */
+ color: white;
+ font-size: 0.8rem;
+ white-space: nowrap;
+ display: inline-block;
+ }
+
+ /* System Monitor Textboxes - Medium gray background for visibility */
+ .toolbar-stat-textbox input[type="text"],
+ .toolbar-stat-textbox .scroll-hide {
+ background-color: #666 !important; /* Medium gray */
+ color: #fff !important; /* White text */
+ border: 1px solid #555 !important;
+ font-weight: 500 !important;
+ }
+
+ /* Responsive design for screens */
+ @media (max-width: 1147px) {
+ .toolbar-patreon, .toolbar-version { /* Hide both on smaller screens */
+ display: none;
+ }
+ .footer-patreon, .footer-version { /* Show both in footer on smaller screens */
+ display: inline-block !important; /* Ensure they are shown */
+ }
+ #fixed-toolbar {
+ gap: 4px !important; /* Reduce gap for screens <= 1024px */
+ }
+ #fixed-toolbar > div:first-child { /* Target the first gr.Column (Title) */
+ min-width: fit-content !important; /* Override Python-set min-width */
+ flex-shrink: 0 !important; /* Prevent title column from shrinking too much */
+ }
+ }
+
+ @media (min-width: 1148px) {
+ .footer-patreon, .footer-version { /* Hide both in footer on larger screens */
+ display: none !important;
+ }
+ }
+
+ @media (max-width: 768px) {
+ .toolbar-title {
+ font-size: 1.1rem;
+ max-width: 150px;
+ }
+ #fixed-toolbar {
+ padding: 3px 6px;
+ gap: 4px;
+ }
+ .toolbar-text {
+ font-size: 0.75rem;
+ }
+ }
+
+ @media (max-width: 510px) {
+ #toolbar-ram-col, #toolbar-vram-col, #toolbar-gpu-col {
+ display: none !important;
+ }
+ }
+
+ @media (max-width: 480px) {
+ .toolbar-title {
+ font-size: 1rem;
+ max-width: 120px;
+ }
+ #fixed-toolbar {
+ padding: 2px 4px;
+ gap: 2px;
+ }
+ .toolbar-text {
+ font-size: 0.7rem;
+ }
+ }
+
+ /* Button styling */
+ #toolbar-add-to-queue-btn button {
+ font-size: 14px !important;
+ padding: 4px 16px !important;
+ height: 32px !important;
+ min-width: 80px !important;
+ }
+ .narrow-button {
+ min-width: 40px !important;
+ width: 40px !important;
+ padding: 0 !important;
+ margin: 0 !important;
+ }
+ .gr-button-primary {
+ color: white;
+ }
+
+ /* Layout adjustments */
+ body, .gradio-container {
+ padding-top: 42px !important; /* Adjusted for new toolbar height (36px - 10px) */
+ }
+
+ @media (max-width: 848px) {
+ body, .gradio-container {
+ padding-top: 48px !important;
+ }
+ }
+
+ @media (max-width: 768px) {
+ body, .gradio-container {
+ padding-top: 22px !important; /* Adjusted for new toolbar height (32px - 10px) */
+ }
+ }
+
+ @media (max-width: 480px) {
+ body, .gradio-container {
+ padding-top: 18px !important; /* Adjusted for new toolbar height (28px - 10px) */
+ }
+ }
+
+ /* hide the gr.Video source selection bar for tb_input_video_component */
+ #toolbox-video-player .source-selection {
+ display: none !important;
+ }
+ /* control sizing for gr.Video components */
+ .video-size video {
+ max-height: 60vh;
+ min-height: 300px !important;
+ object-fit: contain;
+ }
+ /* NEW: Closes the gap between input tabs and the pipeline accordion below them */
+ #pipeline-controls-wrapper {
+ margin-top: -15px !important; /* Adjust this value to get the perfect "snug" fit */
+ }
+ /* --- NEW CSS RULE FOR GALLERY SCROLLING --- */
+ #gallery-scroll-wrapper {
+ max-height: 600px; /* Set your desired fixed height */
+ overflow-y: auto; /* Add a scrollbar only when needed */
+ }
+ #toolbox-start-pipeline-btn {
+ margin-top: -14px !important; /* Adjust this value to get the perfect alignment */
+ }
+
+ .control-group {
+ border-top: 1px solid #ccc;
+ border-bottom: 1px solid #ccc;
+ margin: 12px 0;
+ }
+ """
+
+ # Get the theme from settings
+ current_theme = settings.get("gradio_theme", "default") # Use default if not found
+ block = gr.Blocks(theme=current_theme).queue()
+ block.css = css # Store for later use in launch()
+ block.title = "FramePack Studio" # Store for later use in launch()
+
+ with block:
+ with gr.Row(elem_id="fixed-toolbar"):
+ with gr.Column(scale=0, min_width=400): # Title/Version/Patreon
+ gr.HTML(f"""
+
+ """)
+ # REMOVED: refresh_stats_btn - Toolbar refresh button is no longer needed
+ # with gr.Column(scale=0, min_width=40):
+ # refresh_stats_btn = gr.Button("⟳", elem_id="refresh-stats-btn", elem_classes="narrow-button")
+ with gr.Column(scale=1, min_width=180): # Queue Stats
+ queue_stats_display = gr.Markdown("Queue: 0 | Running: 0 | Completed: 0
")
+
+ # --- System Stats Display - Single gr.Textbox per stat ---
+ with gr.Column(scale=0, min_width=173, elem_id="toolbar-ram-col"): # RAM Column
+ toolbar_ram_display_component = gr.Textbox(
+ value="RAM: N/A",
+ interactive=False,
+ lines=1,
+ max_lines=1,
+ show_label=False,
+ container=False,
+ elem_id="toolbar-ram-stat",
+ elem_classes="toolbar-stat-textbox"
+ )
+ with gr.Column(scale=0, min_width=138, elem_id="toolbar-vram-col"): # VRAM Column
+ toolbar_vram_display_component = gr.Textbox(
+ value="VRAM: N/A",
+ interactive=False,
+ lines=1,
+ max_lines=1,
+ show_label=False,
+ container=False,
+ elem_id="toolbar-vram-stat",
+ elem_classes="toolbar-stat-textbox"
+ # Visibility controlled by tb_get_formatted_toolbar_stats
+ )
+ with gr.Column(scale=0, min_width=130, elem_id="toolbar-gpu-col"): # GPU Column
+ toolbar_gpu_display_component = gr.Textbox(
+ value="GPU: N/A",
+ interactive=False,
+ lines=1,
+ max_lines=1,
+ show_label=False,
+ container=False,
+ elem_id="toolbar-gpu-stat",
+ elem_classes="toolbar-stat-textbox"
+ # Visibility controlled by tb_get_formatted_toolbar_stats
+ )
+ # --- End of System Stats Display ---
+
+ # Removed old version_display column
+ # --- End of Toolbar ---
+
+ # Essential to capture main_tabs_component for later use by send_to_toolbox_btn
+ with gr.Tabs(elem_id="main_tabs") as main_tabs_component:
+ with gr.Tab("Generate", id="generate_tab"):
+ # NEW: Top preview area for latents display
+ with gr.Row(visible=get_latents_display_top()) as top_preview_row:
+ top_preview_image = gr.Image(
+ label="Next Latents (Top Display)",
+ height=150,
+ visible=True,
+ type="numpy",
+ interactive=False,
+ elem_classes="contain-image",
+ image_mode="RGB"
+ )
+
+ with gr.Row():
+ with gr.Column(scale=2):
+ model_type = gr.Radio(
+ choices=[("Original", "Original"), ("Original with Endframe", "Original with Endframe"), ("F1", "F1"), ("Video", "Video"), ("Video with Endframe", "Video with Endframe"), ("Video F1", "Video F1")],
+ value="Original",
+ label="Generation Type"
+ )
+ with gr.Accordion("Original Presets", open=False, visible=True) as preset_accordion:
+ with gr.Row():
+ preset_dropdown = gr.Dropdown(label="Select Preset", choices=load_presets("Original"), interactive=True, scale=2)
+ delete_preset_button = gr.Button("🗑️ Delete", variant="stop", scale=1)
+ with gr.Row():
+ preset_name_textbox = gr.Textbox(label="Preset Name", placeholder="Enter a name for your preset", scale=2)
+ save_preset_button = gr.Button("💾 Save", variant="primary", scale=1)
+ with gr.Row(visible=False) as confirm_delete_row:
+ gr.Markdown("### Are you sure you want to delete this preset?")
+ confirm_delete_yes_btn = gr.Button("🗑️ Yes, Delete", variant="stop")
+ confirm_delete_no_btn = gr.Button("↩️ No, Go Back")
+ with gr.Accordion("Basic Parameters", open=True, visible=True) as basic_parameters_accordion:
+ with gr.Group():
+ total_second_length = gr.Slider(label="Video Length (Seconds)", minimum=1, maximum=120, value=6, step=0.1)
+ with gr.Row("Resolution"):
+ resolutionW = gr.Slider(
+ label="Width", minimum=128, maximum=768, value=640, step=32,
+ info="Nearest valid width will be used."
+ )
+ resolutionH = gr.Slider(
+ label="Height", minimum=128, maximum=768, value=640, step=32,
+ info="Nearest valid height will be used."
+ )
+ resolution_text = gr.Markdown(value="Selected bucket for resolution: 640 x 640
", label="", show_label=False)
+
+ # --- START OF REFACTORED XY PLOT SECTION ---
+ xy_plot_components = create_xy_plot_ui(
+ lora_names=lora_names,
+ default_prompt=default_prompt,
+ DUMMY_LORA_NAME=DUMMY_LORA_NAME,
+ )
+ xy_group = xy_plot_components["group"]
+ xy_plot_status = xy_plot_components["status"]
+ xy_plot_output = xy_plot_components["output"]
+ # --- END OF REFACTORED XY PLOT SECTION ---
+
+ with gr.Group(visible=True) as standard_generation_group: # Default visibility: True because "Original" model is not "Video"
+ with gr.Group(visible=True) as image_input_group: # This group now only contains the start frame image
+ with gr.Row():
+ with gr.Column(scale=1): # Start Frame Image Column
+ input_image = gr.Image(
+ sources='upload',
+ type="numpy",
+ label="Start Frame (optional)",
+ elem_classes="contain-image",
+ image_mode="RGB",
+ show_label=True, # Keep label for clarity
+ container=True
+ )
+
+ with gr.Group(visible=False) as video_input_group:
+ input_video = gr.Video(
+ sources='upload',
+ label="Video Input",
+ height=420,
+ show_label=True
+ )
+ combine_with_source = gr.Checkbox(
+ label="Combine with source video",
+ value=True,
+ info="If checked, the source video will be combined with the generated video",
+ interactive=True
+ )
+ num_cleaned_frames = gr.Slider(label="Number of Context Frames (Adherence to Video)", minimum=2, maximum=10, value=5, step=1, interactive=True, info="Expensive. Retain more video details. Reduce if memory issues or motion too restricted (jumpcut, ignoring prompt, still).")
+
+
+ # End Frame Image Input
+ # Initial visibility is False, controlled by update_input_visibility
+ with gr.Column(scale=1, visible=False) as end_frame_group_original:
+ end_frame_image_original = gr.Image(
+ sources='upload',
+ type="numpy",
+ label="End Frame (Optional)",
+ elem_classes="contain-image",
+ image_mode="RGB",
+ show_label=True,
+ container=True
+ )
+
+ # End Frame Influence slider
+ # Initial visibility is False, controlled by update_input_visibility
+ with gr.Group(visible=False) as end_frame_slider_group:
+ end_frame_strength_original = gr.Slider(
+ label="End Frame Influence",
+ minimum=0.05,
+ maximum=1.0,
+ value=1.0,
+ step=0.05,
+ info="Controls how strongly the end frame guides the generation. 1.0 is full influence."
+ )
+
+
+
+ with gr.Row():
+ prompt = gr.Textbox(label="Prompt", value=default_prompt, scale=10)
+ with gr.Row():
+ enhance_prompt_btn = gr.Button("✨ Enhance", scale=1)
+ caption_btn = gr.Button("✨ Caption", scale=1)
+
+ with gr.Accordion("Prompt Parameters", open=False):
+ n_prompt = gr.Textbox(label="Negative Prompt", value="", visible=True) # Make visible for both models
+
+ blend_sections = gr.Slider(
+ minimum=0, maximum=10, value=4, step=1,
+ label="Number of sections to blend between prompts"
+ )
+ with gr.Accordion("Batch Input", open=False):
+ batch_input_images = gr.File(
+ label="Batch Images (Upload one or more)",
+ file_count="multiple",
+ file_types=["image"],
+ type="filepath"
+ )
+ batch_input_gallery = gr.Gallery(
+ label="Selected Batch Images",
+ visible=False,
+ columns=5,
+ object_fit="contain",
+ height="auto"
+ )
+ add_batch_to_queue_btn = gr.Button("🚀 Add Batch to Queue", variant="primary")
+ with gr.Accordion("Generation Parameters", open=True):
+ with gr.Row():
+ steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=25, step=1)
+ def on_input_image_change(img):
+ if img is not None:
+ return gr.update(info="Nearest valid bucket size will be used. Height will be adjusted automatically."), gr.update(visible=False)
+ else:
+ return gr.update(info="Nearest valid width will be used."), gr.update(visible=True)
+ input_image.change(fn=on_input_image_change, inputs=[input_image], outputs=[resolutionW, resolutionH])
+ def on_resolution_change(img, resolutionW, resolutionH):
+ out_bucket_resH, out_bucket_resW = [640, 640]
+ if img is not None:
+ H, W, _ = img.shape
+ out_bucket_resH, out_bucket_resW = find_nearest_bucket(H, W, resolution=resolutionW)
+ else:
+ out_bucket_resH, out_bucket_resW = find_nearest_bucket(resolutionH, resolutionW, (resolutionW+resolutionH)/2) # if resolutionW > resolutionH else resolutionH
+ return gr.update(value=f"Selected bucket for resolution: {out_bucket_resW} x {out_bucket_resH}
")
+ resolutionW.change(fn=on_resolution_change, inputs=[input_image, resolutionW, resolutionH], outputs=[resolution_text], show_progress="hidden")
+ resolutionH.change(fn=on_resolution_change, inputs=[input_image, resolutionW, resolutionH], outputs=[resolution_text], show_progress="hidden")
+
+ with gr.Row():
+ seed = gr.Number(label="Seed", value=2500, precision=0)
+ randomize_seed = gr.Checkbox(label="Randomize", value=True, info="Generate a new random seed for each job")
+ with gr.Accordion("LoRAs", open=False):
+ with gr.Row():
+ lora_selector = gr.Dropdown(
+ choices=lora_names,
+ label="Select LoRAs to Load",
+ multiselect=True,
+ value=[],
+ info="Select one or more LoRAs to use for this job"
+ )
+ lora_names_states = gr.State(lora_names)
+ lora_sliders = {}
+ for lora in lora_names:
+ lora_sliders[lora] = gr.Slider(
+ minimum=0.0, maximum=2.0, value=1.0, step=0.01,
+ label=f"{lora} Weight", visible=False, interactive=True
+ )
+ with gr.Accordion("Latent Image Options", open=False):
+ latent_type = gr.Dropdown(
+ ["Noise", "White", "Black", "Green Screen"], label="Latent Image", value="Noise", info="Used as a starting point if no image is provided"
+ )
+ with gr.Accordion("Advanced Parameters", open=False):
+ gr.Markdown("#### Motion Model")
+ gr.Markdown("Settings for precise control of the motion model")
+
+ with gr.Group(elem_classes="control-group"):
+ latent_window_size = gr.Slider(label="Latent Window Size", minimum=1, maximum=33, value=9, step=1, info='Change at your own risk, very experimental') # Should not change
+ gs = gr.Slider(label="Distilled CFG Scale", minimum=1.0, maximum=32.0, value=10.0, step=0.5)
+
+ gr.Markdown("#### CFG Scale")
+ gr.Markdown("Much better prompt following. Warning: Modifying these values from their defaults will almost double generation time. ⚠️")
+
+ with gr.Group(elem_classes="control-group"):
+ cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=3.0, value=1.0, step=0.1)
+ rs = gr.Slider(label="CFG Re-Scale", minimum=0.0, maximum=1.0, value=0.0, step=0.05)
+
+ gr.Markdown("#### Cache Options")
+ gr.Markdown("Using a cache will speed up generation. May affect quality, fine or even coarse details, and may change or inhibit motion. You can choose at most one.")
+
+ with gr.Group(elem_classes="control-group"):
+ with gr.Row():
+ cache_type = gr.Radio(["MagCache", "TeaCache", "None"], value='MagCache', label="Caching strategy", info="Which cache implementation to use, if any")
+
+ with gr.Row(): # MagCache now first
+ magcache_threshold = gr.Slider(label="MagCache Threshold", minimum=0.01, maximum=1.0, step=0.01, value=0.1, visible=True, info='[⬇️ **Faster**] Error tolerance. Lower = more estimated steps')
+ magcache_max_consecutive_skips = gr.Slider(label="MagCache Max Consecutive Skips", minimum=1, maximum=5, step=1, value=2, visible=True, info='[⬆️ **Faster**] Allow multiple estimated steps in a row')
+ magcache_retention_ratio = gr.Slider(label="MagCache Retention Ratio", minimum=0.0, maximum=1.0, step=0.01, value=0.25, visible=True, info='[⬇️ **Faster**] Disallow estimation in critical early steps')
+
+ with gr.Row():
+ teacache_num_steps = gr.Slider(label="TeaCache steps", minimum=1, maximum=50, step=1, value=25, visible=False, info='How many intermediate sections to keep in the cache')
+ teacache_rel_l1_thresh = gr.Slider(label="TeaCache rel_l1_thresh", minimum=0.01, maximum=1.0, step=0.01, value=0.15, visible=False, info='[⬇️ **Faster**] Relative L1 Threshold')
+
+ def update_cache_type(cache_type: str):
+ enable_magcache = False
+ enable_teacache = False
+
+ if cache_type == 'MagCache':
+ enable_magcache = True
+ elif cache_type == 'TeaCache':
+ enable_teacache = True
+
+ magcache_threshold_update = gr.update(visible=enable_magcache)
+ magcache_max_consecutive_skips_update = gr.update(visible=enable_magcache)
+ magcache_retention_ratio_update = gr.update(visible=enable_magcache)
+
+ teacache_num_steps_update = gr.update(visible=enable_teacache)
+ teacache_rel_l1_thresh_update = gr.update(visible=enable_teacache)
+
+ return [
+ magcache_threshold_update,
+ magcache_max_consecutive_skips_update,
+ magcache_retention_ratio_update,
+ teacache_num_steps_update,
+ teacache_rel_l1_thresh_update
+ ]
+
+
+ cache_type.change(fn=update_cache_type, inputs=cache_type, outputs=[
+ magcache_threshold,
+ magcache_max_consecutive_skips,
+ magcache_retention_ratio,
+ teacache_num_steps,
+ teacache_rel_l1_thresh
+ ])
+
+ with gr.Row("Metadata"):
+ json_upload = gr.File(
+ label="Upload Metadata JSON (optional)",
+ file_types=[".json"],
+ type="filepath",
+ height=140,
+ )
+
+ with gr.Column():
+ preview_image = gr.Image(
+ label="Next Latents",
+ height=150,
+ visible=not get_latents_display_top(),
+ type="numpy",
+ interactive=False,
+ elem_classes="contain-image",
+ image_mode="RGB"
+ )
+ result_video = gr.Video(label="Finished Frames", autoplay=True, height=256, loop=True)
+ progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
+ progress_bar = gr.HTML('', elem_classes='no-generating-animation')
+ with gr.Row():
+ current_job_id = gr.Textbox(label="Current Job ID", value="", visible=True, interactive=True)
+ start_button = gr.Button(value="🚀 Add to Queue", variant="primary", elem_id="toolbar-add-to-queue-btn")
+ xy_plot_process_btn = gr.Button("🚀 Submit XY Plot", visible=False)
+ video_input_required_message = gr.Markdown(
+ "Input video required
", visible=False
+ )
+ end_button = gr.Button(value="❌ Cancel Current Job", interactive=True, visible=False)
+
+
+
+ with gr.Tab("Queue"):
+ with gr.Row():
+ with gr.Column():
+ with gr.Row() as queue_controls_row:
+ refresh_button = gr.Button("🔄 Refresh Queue")
+ load_queue_button = gr.Button("▶️ Resume Queue")
+ queue_export_button = gr.Button("📦 Export Queue")
+ clear_complete_button = gr.Button("🧹 Clear Completed Jobs", variant="secondary")
+ clear_queue_button = gr.Button("❌ Cancel Queued Jobs", variant="stop")
+ with gr.Row():
+ import_queue_file = gr.File(
+ label="Import Queue",
+ file_types=[".json", ".zip"],
+ type="filepath",
+ visible=True,
+ elem_classes="short-import-box"
+ )
+
+ with gr.Row(visible=False) as confirm_cancel_row:
+ gr.Markdown("### Are you sure you want to cancel all pending jobs?")
+ confirm_cancel_yes_btn = gr.Button("❌ Yes, Cancel All", variant="stop")
+ confirm_cancel_no_btn = gr.Button("↩️ No, Go Back")
+
+ with gr.Row():
+ queue_status = gr.DataFrame(
+ headers=["Job ID", "Type", "Status", "Created", "Started", "Completed", "Elapsed", "Preview"],
+ datatype=["str", "str", "str", "str", "str", "str", "str", "html"],
+ label="Job Queue"
+ )
+
+ with gr.Accordion("Queue Documentation", open=False):
+ gr.Markdown("""
+ ## Queue Tab Guide
+
+ This tab is for managing your generation jobs.
+
+ - **Refresh Queue**: Update the job list.
+ - **Cancel Queue**: Stop all pending jobs.
+ - **Clear Complete**: Remove finished, failed, or cancelled jobs from the list.
+ - **Load Queue**: Load jobs from the default `queue.json`.
+ - **Export Queue**: Save the current job list and its images to a zip file.
+ - **Import Queue**: Load a queue from a `.json` or `.zip` file.
+ """)
+
+ # --- Event Handlers for Queue Tab ---
+
+ # Function to clear all jobs in the queue
+ def clear_all_jobs():
+ try:
+ cancelled_count = job_queue.clear_queue()
+ print(f"Cleared {cancelled_count} jobs from the queue")
+ return update_stats()
+ except Exception as e:
+ import traceback
+ print(f"Error in clear_all_jobs: {e}")
+ traceback.print_exc()
+ return [], ""
+
+ # Function to clear completed and cancelled jobs
+ def clear_completed_jobs():
+ try:
+ removed_count = job_queue.clear_completed_jobs()
+ print(f"Removed {removed_count} completed/cancelled jobs from the queue")
+ return update_stats()
+ except Exception as e:
+ import traceback
+ print(f"Error in clear_completed_jobs: {e}")
+ traceback.print_exc()
+ return [], ""
+
+ # Function to load queue from queue.json
+ def load_queue_from_json():
+ try:
+ loaded_count = job_queue.load_queue_from_json()
+ print(f"Loaded {loaded_count} jobs from queue.json")
+ return update_stats()
+ except Exception as e:
+ import traceback
+ print(f"Error loading queue from JSON: {e}")
+ traceback.print_exc()
+ return [], ""
+
+ # Function to import queue from a custom JSON file
+ def import_queue_from_file(file_path):
+ if not file_path:
+ return update_stats()
+ try:
+ loaded_count = job_queue.load_queue_from_json(file_path)
+ print(f"Loaded {loaded_count} jobs from {file_path}")
+ return update_stats()
+ except Exception as e:
+ import traceback
+ print(f"Error importing queue from file: {e}")
+ traceback.print_exc()
+ return [], ""
+
+ # Function to export queue to a zip file
+ def export_queue_to_zip():
+ try:
+ zip_path = job_queue.export_queue_to_zip()
+ if zip_path and os.path.exists(zip_path):
+ print(f"Queue exported to {zip_path}")
+ else:
+ print("Failed to export queue to zip")
+ return update_stats()
+ except Exception as e:
+ import traceback
+ print(f"Error exporting queue to zip: {e}")
+ traceback.print_exc()
+ return [], ""
+
+ # --- Connect Buttons ---
+ refresh_button.click(fn=update_stats, inputs=[], outputs=[queue_status, queue_stats_display])
+
+ # Confirmation logic for Cancel Queue
+ def show_cancel_confirmation():
+ return gr.update(visible=False), gr.update(visible=True)
+
+ def hide_cancel_confirmation():
+ return gr.update(visible=True), gr.update(visible=False)
+
+ def confirmed_clear_all_jobs():
+ qs_data, qs_text = clear_all_jobs()
+ return qs_data, qs_text, gr.update(visible=True), gr.update(visible=False)
+
+ clear_queue_button.click(fn=show_cancel_confirmation, inputs=None, outputs=[queue_controls_row, confirm_cancel_row])
+ confirm_cancel_no_btn.click(fn=hide_cancel_confirmation, inputs=None, outputs=[queue_controls_row, confirm_cancel_row])
+ confirm_cancel_yes_btn.click(fn=confirmed_clear_all_jobs, inputs=None, outputs=[queue_status, queue_stats_display, queue_controls_row, confirm_cancel_row])
+
+ clear_complete_button.click(fn=clear_completed_jobs, inputs=[], outputs=[queue_status, queue_stats_display])
+ queue_export_button.click(fn=export_queue_to_zip, inputs=[], outputs=[queue_status, queue_stats_display])
+
+ # Create a container for thumbnails (kept for potential future use, though not displayed in DataFrame)
+ with gr.Row():
+ thumbnail_container = gr.Column()
+ thumbnail_container.elem_classes = ["thumbnail-container"]
+
+ # Add CSS for thumbnails
+
+ with gr.Tab("Outputs", id="outputs_tab"): # Ensure 'id' is present for tab switching
+ outputDirectory_video = settings.get("output_dir", settings.default_settings['output_dir'])
+ outputDirectory_metadata = settings.get("metadata_dir", settings.default_settings['metadata_dir'])
+ def get_gallery_items():
+ items = []
+ for f in os.listdir(outputDirectory_metadata):
+ if f.endswith(".png"):
+ prefix = os.path.splitext(f)[0]
+ latest_video = get_latest_video_version(prefix)
+ if latest_video:
+ video_path = os.path.join(outputDirectory_video, latest_video)
+ mtime = os.path.getmtime(video_path)
+ preview_path = os.path.join(outputDirectory_metadata, f)
+ items.append((preview_path, prefix, mtime))
+ items.sort(key=lambda x: x[2], reverse=True)
+ return [(i[0], i[1]) for i in items]
+ def get_latest_video_version(prefix):
+ max_number = -1
+ selected_file = None
+ for f in os.listdir(outputDirectory_video):
+ if f.startswith(prefix + "_") and f.endswith(".mp4"):
+ # Skip files that include "combined" in their name
+ if "combined" in f:
+ continue
+ try:
+ num = int(f.replace(prefix + "_", '').replace(".mp4", ''))
+ if num > max_number:
+ max_number = num
+ selected_file = f
+ except ValueError:
+ # Ignore files that do not have a valid number in their name
+ continue
+ return selected_file
+ # load_video_and_info_from_prefix now also returns button visibility
+ def load_video_and_info_from_prefix(prefix):
+ video_file = get_latest_video_version(prefix)
+ json_path = os.path.join(outputDirectory_metadata, prefix) + ".json"
+
+ if not video_file or not os.path.exists(os.path.join(outputDirectory_video, video_file)) or not os.path.exists(json_path):
+ # If video or info not found, button should be hidden
+ return None, "Video or JSON not found.", gr.update(visible=False)
+
+ video_path = os.path.join(outputDirectory_video, video_file)
+ info_content = {"description": "no info"}
+ if os.path.exists(json_path):
+ with open(json_path, "r", encoding="utf-8") as f:
+ info_content = json.load(f)
+ # If video and info found, button should be visible
+ return video_path, json.dumps(info_content, indent=2, ensure_ascii=False), gr.update(visible=True)
+
+ gallery_items_state = gr.State(get_gallery_items())
+ selected_original_video_path_state = gr.State(None) # Holds the ORIGINAL, UNPROCESSED path
+ with gr.Row():
+ with gr.Column(scale=2):
+ thumbs = gr.Gallery(
+ # value=[i[0] for i in get_gallery_items()],
+ columns=[4],
+ allow_preview=False,
+ object_fit="cover",
+ height="auto"
+ )
+ refresh_button = gr.Button("🔄 Update Gallery")
+ with gr.Column(scale=5):
+ video_out = gr.Video(sources=[], autoplay=True, loop=True, visible=False)
+ with gr.Column(scale=1):
+ info_out = gr.Textbox(label="Generation info", visible=False)
+ send_to_toolbox_btn = gr.Button("➡️ Send to Post-processing", visible=False) # Added new send_to_toolbox_btn
+ def refresh_gallery():
+ new_items = get_gallery_items()
+ return gr.update(value=[i[0] for i in new_items]), new_items
+ refresh_button.click(fn=refresh_gallery, outputs=[thumbs, gallery_items_state])
+
+ # MODIFIED: on_select now handles visibility of the new button
+ def on_select(evt: gr.SelectData, gallery_items):
+ if evt.index is None or not gallery_items or evt.index >= len(gallery_items):
+ return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
+
+ prefix = gallery_items[evt.index][1]
+ # original_video_path is e.g., "outputs/my_actual_video.mp4"
+ original_video_path, info_string, button_visibility_update = load_video_and_info_from_prefix(prefix)
+
+ # Determine visibility for video and info based on whether video_path was found
+ video_out_update = gr.update(value=original_video_path, visible=bool(original_video_path))
+ info_out_update = gr.update(value=info_string, visible=bool(original_video_path))
+
+ # IMPORTANT: Store the ORIGINAL, UNPROCESSED path in the gr.State
+ return video_out_update, info_out_update, button_visibility_update, original_video_path
+
+ thumbs.select(
+ fn=on_select,
+ inputs=[gallery_items_state],
+ outputs=[video_out, info_out, send_to_toolbox_btn, selected_original_video_path_state] # Output original path to State
+ )
+ with gr.Tab("Post-processing", id="toolbox_tab"):
+ # Call the function from toolbox_app.py to build the Toolbox UI
+ # The toolbox_ui_layout (e.g., a gr.Column) is automatically placed here.
+ toolbox_ui_layout, tb_target_video_input = tb_create_video_toolbox_ui()
+
+ with gr.Tab("Settings"):
+ with gr.Row():
+ with gr.Column():
+ save_metadata = gr.Checkbox(
+ label="Save Metadata",
+ info="Save to JSON file",
+ value=settings.get("save_metadata", 6),
+ )
+ gpu_memory_preservation = gr.Slider(
+ label="Memory Buffer for Stability (VRAM GB)",
+ minimum=1,
+ maximum=128,
+ step=0.1,
+ value=settings.get("gpu_memory_preservation", 6),
+ info="Increase reserve if you see computer freezes, stagnant generation, or super slow sampling steps (try 1G at a time).\
+ Otherwise smaller buffer is faster. Some models and lora need more buffer than others. \
+ (5.5 - 8.5 is a common range)"
+ )
+ mp4_crf = gr.Slider(
+ label="MP4 Compression",
+ minimum=0,
+ maximum=100,
+ step=1,
+ value=settings.get("mp4_crf", 16),
+ info="Lower means better quality. 0 is uncompressed. Change to 16 if you get black outputs."
+ )
+ clean_up_videos = gr.Checkbox(
+ label="Clean up video files",
+ value=settings.get("clean_up_videos", True),
+ info="If checked, only the final video will be kept after generation."
+ )
+ auto_cleanup_on_startup = gr.Checkbox(
+ label="Automatically clean up temp folders on startup",
+ value=settings.get("auto_cleanup_on_startup", False),
+ info="If checked, temporary files (inc. post-processing) will be cleaned up when the application starts."
+ )
+
+ latents_display_top = gr.Checkbox(
+ label="Display Next Latents across top of interface",
+ value=get_latents_display_top(),
+ info="If checked, the Next Latents preview will be displayed across the top of the interface instead of in the right column."
+ )
+
+ # gr.Markdown("---")
+ # gr.Markdown("### Startup Settings")
+ gr.Markdown("")
+ # Initial values for startup preset dropdown
+ # Ensure settings and load_presets are available in this scope
+ initial_startup_model_val = settings.get("startup_model_type", "None")
+ initial_startup_presets_choices_val = []
+ initial_startup_preset_value_val = None
+
+ if initial_startup_model_val and initial_startup_model_val != "None":
+ # load_presets is defined further down in create_interface
+ initial_startup_presets_choices_val = load_presets(initial_startup_model_val)
+ saved_preset_for_initial_model_val = settings.get("startup_preset_name")
+ if saved_preset_for_initial_model_val in initial_startup_presets_choices_val:
+ initial_startup_preset_value_val = saved_preset_for_initial_model_val
+
+ startup_model_type_dropdown = gr.Dropdown(
+ label="Startup Model Type",
+ choices=["None"] + [choice[0] for choice in model_type.choices if choice[0] != "XY Plot"], # model_type is the Radio on Generate tab
+ value=initial_startup_model_val,
+ info="Select a model type to load on startup. 'None' to disable."
+ )
+ startup_preset_name_dropdown = gr.Dropdown(
+ label="Startup Preset",
+ choices=initial_startup_presets_choices_val,
+ value=initial_startup_preset_value_val,
+ info="Select a preset for the startup model. Updates when Startup Model Type changes.",
+ interactive=True # Must be interactive to be updated by another component
+ )
+
+ with gr.Accordion("System Prompt", open=False):
+ with gr.Row(equal_height=True): # New Row to contain checkbox and reset button
+ override_system_prompt = gr.Checkbox(
+ label="Override System Prompt",
+ value=settings.get("override_system_prompt", False),
+ info="If checked, the system prompt template below will be used instead of the default one.",
+ scale=1 # Give checkbox some scale
+ )
+ reset_system_prompt_btn = gr.Button(
+ "🔄 Reset",
+ scale=0
+ )
+ system_prompt_template = gr.Textbox(
+ label="System Prompt Template",
+ value=settings.get("system_prompt_template", "{\"template\": \"<|start_header_id|>system<|end_header_id|>\\n\\nDescribe the video by detailing the following aspects: 1. The main content and theme of the video.2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.4. background environment, light, style and atmosphere.5. camera angles, movements, and transitions used in the video:<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{}<|eot_id|>\", \"crop_start\": 95}"),
+ lines=10,
+ info="System prompt template used for video generation. Must be a valid JSON or Python dictionary string with 'template' and 'crop_start' keys. Example: {\"template\": \"your template here\", \"crop_start\": 95}"
+ )
+ # The reset_system_prompt_btn is now defined above within the Row
+
+ # --- Settings Tab Event Handlers ---
+
+ output_dir = gr.Textbox(
+ label="Output Directory",
+ value=settings.get("output_dir"),
+ placeholder="Path to save generated videos"
+ )
+ metadata_dir = gr.Textbox(
+ label="Metadata Directory",
+ value=settings.get("metadata_dir"),
+ placeholder="Path to save metadata files"
+ )
+ lora_dir = gr.Textbox(
+ label="LoRA Directory",
+ value=settings.get("lora_dir"),
+ placeholder="Path to LoRA models"
+ )
+ gradio_temp_dir = gr.Textbox(label="Gradio Temporary Directory", value=settings.get("gradio_temp_dir"))
+ auto_save = gr.Checkbox(
+ label="Auto-save settings",
+ value=settings.get("auto_save_settings", True)
+ )
+ # Add Gradio Theme Dropdown
+ gradio_themes = ["default", "base", "soft", "glass", "mono", "origin", "citrus", "monochrome", "ocean", "NoCrypt/miku", "earneleh/paris", "gstaff/xkcd"]
+ theme_dropdown = gr.Dropdown(
+ label="Theme",
+ choices=gradio_themes,
+ value=settings.get("gradio_theme", "default"),
+ info="Select the Gradio UI theme. Requires restart."
+ )
+ save_btn = gr.Button("💾 Save Settings")
+ cleanup_btn = gr.Button("🗑️ Clean Up Temporary Files")
+ status = gr.HTML("")
+ cleanup_output = gr.Textbox(label="Cleanup Status", interactive=False)
+
+ def save_settings(save_metadata, gpu_memory_preservation, mp4_crf, clean_up_videos, auto_cleanup_on_startup_val, latents_display_top_val, override_system_prompt_value, system_prompt_template_value, output_dir, metadata_dir, lora_dir, gradio_temp_dir, auto_save, selected_theme, startup_model_type_val, startup_preset_name_val):
+ """Handles the manual 'Save Settings' button click."""
+ # This function is for the manual save button.
+ # It collects all current UI values and saves them.
+ # The auto-save logic is handled by individual .change() and .blur() handlers
+ # calling settings.set().
+
+ # First, update the settings object with all current values from the UI
+ try:
+ # Save the system prompt template as is, without trying to parse it
+ # The hunyuan.py file will handle parsing it when needed
+ processed_template = system_prompt_template_value
+
+ settings.save_settings(
+ save_metadata=save_metadata,
+ gpu_memory_preservation=gpu_memory_preservation,
+ mp4_crf=mp4_crf,
+ clean_up_videos=clean_up_videos,
+ auto_cleanup_on_startup=auto_cleanup_on_startup_val, # ADDED
+ latents_display_top=latents_display_top_val, # NEW: Added latents display position setting
+ override_system_prompt=override_system_prompt_value,
+ system_prompt_template=processed_template,
+ output_dir=output_dir,
+ metadata_dir=metadata_dir,
+ lora_dir=lora_dir,
+ gradio_temp_dir=gradio_temp_dir,
+ auto_save_settings=auto_save,
+ gradio_theme=selected_theme,
+ startup_model_type=startup_model_type_val,
+ startup_preset_name=startup_preset_name_val
+ )
+ # settings.save_settings() is called inside settings.save_settings if auto_save is true,
+ # but for the manual button, we ensure it saves regardless of the auto_save flag's previous state.
+ # The call above to settings.save_settings already handles writing to disk.
+ return "Settings saved successfully! Restart required for theme change.
"
+ except Exception as e:
+ return f"Error saving settings: {str(e)}
"
+
+ def handle_individual_setting_change(key, value, setting_name_for_ui):
+ """Called by .change() and .submit() events of individual setting components."""
+ if key == "auto_save_settings":
+ # For the "auto_save_settings" checkbox itself:
+ # 1. Update its value directly in the settings object in memory.
+ # This bypasses the conditional save logic within settings.set() for this specific action.
+ settings.settings[key] = value
+ # 2. Force a save of all settings to disk. This will be correct because either:
+ # - auto_save_settings is turning True: so all changes already in memory need to be saved now.
+ # - auto_save_settings turning False from True: prior changes already saved so only auto_save_settings will be saved.
+ settings.save_settings()
+ # 3. Provide feedback.
+ if value is True:
+ return f"'{setting_name_for_ui}' setting is now ON and saved.
"
+ else:
+ return f"'{setting_name_for_ui}' setting is now OFF and saved.
"
+ else:
+ # For all other settings:
+ # Let settings.set() handle the auto-save logic based on the current "auto_save_settings" value.
+ settings.set(key, value) # settings.set() will call save_settings() if auto_save is True
+ if settings.get("auto_save_settings"): # Check the current state of auto_save
+ return f"'{setting_name_for_ui}' setting auto-saved.
"
+ else:
+ return f"'{setting_name_for_ui}' setting changed (auto-save is off, click 'Save Settings').
"
+
+ # REMOVE `cleanup_temp_folder` from the `inputs` list
+ save_btn.click(
+ fn=save_settings,
+ inputs=[save_metadata, gpu_memory_preservation, mp4_crf, clean_up_videos, auto_cleanup_on_startup, latents_display_top, override_system_prompt, system_prompt_template, output_dir, metadata_dir, lora_dir, gradio_temp_dir, auto_save, theme_dropdown, startup_model_type_dropdown, startup_preset_name_dropdown],
+ outputs=[status]
+ ).then(
+ # NEW: Update latents display layout after manual save
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+ def reset_system_prompt_template_value():
+ return settings.default_settings["system_prompt_template"], False
+
+ reset_system_prompt_btn.click(
+ fn=reset_system_prompt_template_value,
+ outputs=[system_prompt_template, override_system_prompt]
+ ).then( # Trigger auto-save for the reset values if auto-save is on
+ lambda val_template, val_override: handle_individual_setting_change("system_prompt_template", val_template, "System Prompt Template") or handle_individual_setting_change("override_system_prompt", val_override, "Override System Prompt"),
+ inputs=[system_prompt_template, override_system_prompt], outputs=[status])
+
+ def manual_cleanup_handler():
+ """UI handler for the manual cleanup button."""
+ # This directly calls the toolbox_processor's cleanup method and returns the summary string.
+ summary = tb_processor.tb_clear_temporary_files()
+ return summary
+
+ cleanup_btn.click(
+ fn=manual_cleanup_handler,
+ inputs=None,
+ outputs=[cleanup_output]
+ )
+
+ # Add .change handlers for auto-saving individual settings
+ save_metadata.change(lambda v: handle_individual_setting_change("save_metadata", v, "Save Metadata"), inputs=[save_metadata], outputs=[status])
+ gpu_memory_preservation.change(lambda v: handle_individual_setting_change("gpu_memory_preservation", v, "GPU Memory Preservation"), inputs=[gpu_memory_preservation], outputs=[status])
+ mp4_crf.change(lambda v: handle_individual_setting_change("mp4_crf", v, "MP4 Compression"), inputs=[mp4_crf], outputs=[status])
+ clean_up_videos.change(lambda v: handle_individual_setting_change("clean_up_videos", v, "Clean Up Videos"), inputs=[clean_up_videos], outputs=[status])
+
+ # NEW: auto-cleanup temp files on startup checkbox
+ auto_cleanup_on_startup.change(lambda v: handle_individual_setting_change("auto_cleanup_on_startup", v, "Auto Cleanup on Startup"), inputs=[auto_cleanup_on_startup], outputs=[status])
+
+ # NEW: latents display position setting
+ latents_display_top.change(lambda v: handle_individual_setting_change("latents_display_top", v, "Latents Display Position"), inputs=[latents_display_top], outputs=[status])
+
+
+
+ # Connect the latents display setting to layout updates
+ def update_latents_display_layout_from_checkbox(display_top):
+ """Update layout when checkbox changes - uses the checkbox value directly"""
+ if display_top:
+ return (
+ gr.update(visible=True), # top_preview_row
+ gr.update(visible=False, value=None) # preview_image (right column)
+ )
+ else:
+ return (
+ gr.update(visible=False), # top_preview_row
+ gr.update(visible=True) # preview_image (right column)
+ )
+
+ latents_display_top.change(
+ fn=update_latents_display_layout_from_checkbox,
+ inputs=[latents_display_top],
+ outputs=[top_preview_row, preview_image]
+ )
+
+ override_system_prompt.change(lambda v: handle_individual_setting_change("override_system_prompt", v, "Override System Prompt"), inputs=[override_system_prompt], outputs=[status])
+ # Using .blur for text changes so they are processed after the user finishes, not on every keystroke
+ system_prompt_template.blur(lambda v: handle_individual_setting_change("system_prompt_template", v, "System Prompt Template"), inputs=[system_prompt_template], outputs=[status])
+ # reset_system_prompt_btn # is handled separately above, on click
+
+ # Using .blur for text changes so they are processed after the user finishes, not on every keystroke
+ output_dir.blur(lambda v: handle_individual_setting_change("output_dir", v, "Output Directory"), inputs=[output_dir], outputs=[status])
+ metadata_dir.blur(lambda v: handle_individual_setting_change("metadata_dir", v, "Metadata Directory"), inputs=[metadata_dir], outputs=[status])
+ lora_dir.blur(lambda v: handle_individual_setting_change("lora_dir", v, "LoRA Directory"), inputs=[lora_dir], outputs=[status])
+ gradio_temp_dir.blur(lambda v: handle_individual_setting_change("gradio_temp_dir", v, "Gradio Temporary Directory"), inputs=[gradio_temp_dir], outputs=[status])
+
+ auto_save.change(lambda v: handle_individual_setting_change("auto_save_settings", v, "Auto-save Settings"), inputs=[auto_save], outputs=[status])
+ theme_dropdown.change(lambda v: handle_individual_setting_change("gradio_theme", v, "Theme"), inputs=[theme_dropdown], outputs=[status])
+
+ # Event handlers for startup settings
+ def update_startup_preset_dropdown_choices(selected_startup_model_type_from_ui):
+ if not selected_startup_model_type_from_ui or selected_startup_model_type_from_ui == "None":
+ return gr.update(choices=[], value=None)
+
+ loaded_presets_for_model = load_presets(selected_startup_model_type_from_ui)
+
+ # Get the preset name that was saved for the *previous* model type
+ current_saved_startup_preset = settings.get("startup_preset_name")
+
+ # Default to None
+ value_to_select = None
+ # If the previously saved preset name exists for the new model, select it
+ if current_saved_startup_preset and current_saved_startup_preset in loaded_presets_for_model:
+ value_to_select = current_saved_startup_preset
+
+ return gr.update(choices=loaded_presets_for_model, value=value_to_select)
+
+ startup_model_type_dropdown.change(
+ fn=lambda v: handle_individual_setting_change("startup_model_type", v, "Startup Model Type"),
+ inputs=[startup_model_type_dropdown], outputs=[status]
+ ).then( # Chain the update to the preset dropdown
+ fn=update_startup_preset_dropdown_choices, inputs=[startup_model_type_dropdown], outputs=[startup_preset_name_dropdown])
+ startup_preset_name_dropdown.change(lambda v: handle_individual_setting_change("startup_preset_name", v, "Startup Preset Name"), inputs=[startup_preset_name_dropdown], outputs=[status])
+
+ # --- Event Handlers and Connections (Now correctly indented) ---
+
+ # --- Connect Monitoring ---
+ # Auto-check for current job on page load and job change
+ def check_for_current_job():
+ # This function will be called when the interface loads
+ # It will check if there's a current job in the queue and update the UI
+ with job_queue.lock:
+ current_job = job_queue.current_job
+ if current_job:
+ # Return all the necessary information to update the preview windows
+ job_id = current_job.id
+ result = current_job.result
+ preview = current_job.progress_data.get('preview') if current_job.progress_data else None
+ desc = current_job.progress_data.get('desc', '') if current_job.progress_data else ''
+ html = current_job.progress_data.get('html', '') if current_job.progress_data else ''
+
+ # Also trigger the monitor_job function to start monitoring this job
+ print(f"Auto-check found current job {job_id}, triggering monitor_job")
+ return job_id, result, preview, preview, desc, html
+ return None, None, None, None, '', ''
+
+ # Auto-check for current job on page load and handle handoff between jobs.
+ def check_for_current_job_and_monitor():
+ # This function is now the key to the handoff.
+ # It finds the current job and returns its ID, which will trigger the monitor.
+ job_id, result, preview, top_preview, desc, html = check_for_current_job()
+ # We also need to get fresh stats at the same time.
+ queue_status_data, queue_stats_text = update_stats()
+ # Return everything needed to update the UI atomically.
+ return job_id, result, preview, top_preview, desc, html, queue_status_data, queue_stats_text
+
+ # Connect the main process function (wrapper for adding to queue)
+ def process_with_queue_update(model_type_arg, *args):
+ # Call update_stats to get both queue_status_data and queue_stats_text
+ queue_status_data, queue_stats_text = update_stats() # MODIFIED
+
+ # Extract all arguments (ensure order matches inputs lists)
+ # The order here MUST match the order in the `ips` list.
+ # RT_BORG: Global settings gpu_memory_preservation, mp4_crf, save_metadata removed from direct args.
+ (input_image_arg,
+ input_video_arg,
+ end_frame_image_original_arg,
+ end_frame_strength_original_arg,
+ prompt_text_arg,
+ n_prompt_arg,
+ seed_arg, # the seed value
+ randomize_seed_arg, # the boolean value of the checkbox
+ total_second_length_arg,
+ latent_window_size_arg,
+ steps_arg,
+ cfg_arg,
+ gs_arg,
+ rs_arg,
+ cache_type_arg,
+ teacache_num_steps_arg,
+ teacache_rel_l1_thresh_arg,
+ magcache_threshold_arg,
+ magcache_max_consecutive_skips_arg,
+ magcache_retention_ratio_arg,
+ blend_sections_arg,
+ latent_type_arg,
+ clean_up_videos_arg, # UI checkbox from Generate tab
+ selected_loras_arg,
+ resolutionW_arg, resolutionH_arg,
+ combine_with_source_arg,
+ num_cleaned_frames_arg,
+ lora_names_states_arg, # This is from lora_names_states (gr.State)
+ *lora_slider_values_tuple # Remaining args are LoRA slider values
+ ) = args
+ # DO NOT parse the prompt here. Parsing happens once in the worker.
+
+ # Determine the model type to send to the backend
+ backend_model_type = model_type_arg # model_type_arg is the UI selection
+ if model_type_arg == "Video with Endframe":
+ backend_model_type = "Video" # The backend "Video" model_type handles with and without endframe
+
+ # Use the appropriate input based on model type
+ is_ui_video_model = is_video_model(model_type_arg)
+ input_data = input_video_arg if is_ui_video_model else input_image_arg
+
+ # Define actual end_frame params to pass to backend
+ actual_end_frame_image_for_backend = None
+ actual_end_frame_strength_for_backend = 1.0 # Default strength
+
+ if model_type_arg == "Original with Endframe" or model_type_arg == "F1 with Endframe" or model_type_arg == "Video with Endframe":
+ actual_end_frame_image_for_backend = end_frame_image_original_arg
+ actual_end_frame_strength_for_backend = end_frame_strength_original_arg
+
+ # Get the input video path for Video model
+ input_image_path = None
+ if is_ui_video_model and input_video_arg is not None:
+ # For Video models, input_video contains the path to the video file
+ input_image_path = input_video_arg
+
+ # Use the current seed value as is for this job
+ # Call the process function with all arguments
+ # Pass the backend_model_type and the ORIGINAL prompt_text string to the backend process function
+ result = process_fn(backend_model_type, input_data, actual_end_frame_image_for_backend, actual_end_frame_strength_for_backend,
+ prompt_text_arg, n_prompt_arg, seed_arg, total_second_length_arg,
+ latent_window_size_arg, steps_arg, cfg_arg, gs_arg, rs_arg,
+ cache_type_arg == 'TeaCache', teacache_num_steps_arg, teacache_rel_l1_thresh_arg,
+ cache_type_arg == 'MagCache', magcache_threshold_arg, magcache_max_consecutive_skips_arg, magcache_retention_ratio_arg,
+ blend_sections_arg, latent_type_arg, clean_up_videos_arg, # clean_up_videos_arg is from UI
+ selected_loras_arg, resolutionW_arg, resolutionH_arg,
+ input_image_path,
+ combine_with_source_arg,
+ num_cleaned_frames_arg,
+ lora_names_states_arg,
+ *lora_slider_values_tuple
+ )
+ # If randomize_seed is checked, generate a new random seed for the next job
+ new_seed_value = None
+ if randomize_seed_arg:
+ new_seed_value = random.randint(0, 21474)
+ print(f"Generated new seed for next job: {new_seed_value}")
+
+ # Create the button update for start_button WITHOUT interactive=True.
+ # The interactivity will be set by update_start_button_state later in the chain.
+ start_button_update_after_add = gr.update(value="🚀 Add to Queue")
+
+ # If a job ID was created, automatically start monitoring it and update queue
+ if result and result[1]: # Check if job_id exists in results
+ job_id = result[1]
+ # queue_status_data = update_queue_status_fn() # OLD: update_stats now called earlier
+ # Call update_stats again AFTER the job is added to get the freshest stats
+ queue_status_data, queue_stats_text = update_stats()
+
+
+ # Add the new seed value to the results if randomize is checked
+ if new_seed_value is not None:
+ # Use result[6] directly for end_button to preserve its value. Add gr.update() for video_input_required_message.
+ return [result[0], job_id, result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, new_seed_value, gr.update()]
+ else:
+ # Use result[6] directly for end_button to preserve its value. Add gr.update() for video_input_required_message.
+ return [result[0], job_id, result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, gr.update(), gr.update()]
+
+ # If no job ID was created, still return the new seed if randomize is checked
+ # Also, ensure we return the latest stats even if no job was created (e.g., error during param validation)
+ queue_status_data, queue_stats_text = update_stats()
+ if new_seed_value is not None:
+ # Make sure to preserve the end_button update from result[6]
+ return [result[0], result[1], result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, new_seed_value, gr.update()]
+ else:
+ # Make sure to preserve the end_button update from result[6]
+ return [result[0], result[1], result[2], result[3], result[4], start_button_update_after_add, result[6], queue_status_data, queue_stats_text, gr.update(), gr.update()]
+
+ # Custom end process function that ensures the queue is updated and changes button text
+ def end_process_with_update():
+ _ = end_process_fn() # Call the original end_process_fn
+ # Now, get fresh stats for both queue table and toolbar
+ queue_status_data, queue_stats_text = update_stats()
+
+ # Don't try to get the new job ID immediately after cancellation
+ # The monitor_job function will handle the transition to the next job
+
+ # Change the cancel button text to "Cancelling..." and make it non-interactive
+ # This ensures the button stays in this state until the job is fully cancelled
+ return queue_status_data, queue_stats_text, gr.update(value="Cancelling...", interactive=False), gr.update(value=None)
+
+ # MODIFIED handle_send_video_to_toolbox:
+ def handle_send_video_to_toolbox(original_path_from_state): # Input is now the original path from gr.State
+ print(f"Sending selected Outputs' video to Post-processing: {original_path_from_state}")
+
+ if original_path_from_state and isinstance(original_path_from_state, str) and os.path.exists(original_path_from_state):
+ # tb_target_video_input will now process the ORIGINAL path (e.g., "outputs/my_actual_video.mp4").
+ return gr.update(value=original_path_from_state), gr.update(selected="toolbox_tab")
+ else:
+ print(f"No valid video path (from State) found to send. Path: {original_path_from_state}")
+ return gr.update(), gr.update()
+
+ send_to_toolbox_btn.click(
+ fn=handle_send_video_to_toolbox,
+ inputs=[selected_original_video_path_state], # INPUT IS NOW THE gr.State holding the ORIGINAL path
+ outputs=[
+ tb_target_video_input, # This is tb_input_video_component from toolbox_app.py
+ main_tabs_component
+ ]
+ )
+
+ # --- Inputs Lists ---
+ # --- Inputs for all models ---
+ ips = [
+ input_image, # Corresponds to input_image_arg
+ input_video, # Corresponds to input_video_arg
+ end_frame_image_original, # Corresponds to end_frame_image_original_arg
+ end_frame_strength_original,# Corresponds to end_frame_strength_original_arg
+ prompt, # Corresponds to prompt_text_arg
+ n_prompt, # Corresponds to n_prompt_arg
+ seed, # Corresponds to seed_arg
+ randomize_seed, # Corresponds to randomize_seed_arg
+ total_second_length, # Corresponds to total_second_length_arg
+ latent_window_size, # Corresponds to latent_window_size_arg
+ steps, # Corresponds to steps_arg
+ cfg, # Corresponds to cfg_arg
+ gs, # Corresponds to gs_arg
+ rs, # Corresponds to rs_arg
+ cache_type, # Corresponds to cache_type_arg
+ teacache_num_steps, # Corresponds to teacache_num_steps_arg
+ teacache_rel_l1_thresh, # Corresponds to teacache_rel_l1_thresh_arg
+ magcache_threshold, # Corresponds to magcache_threshold_arg
+ magcache_max_consecutive_skips, # Corresponds to magcache_max_consecutive_skips_arg
+ magcache_retention_ratio, # Corresponds to magcache_retention_ratio_arg
+ blend_sections, # Corresponds to blend_sections_arg
+ latent_type, # Corresponds to latent_type_arg
+ clean_up_videos, # Corresponds to clean_up_videos_arg (UI checkbox)
+ lora_selector, # Corresponds to selected_loras_arg
+ resolutionW, # Corresponds to resolutionW_arg
+ resolutionH, # Corresponds to resolutionH_arg
+ combine_with_source, # Corresponds to combine_with_source_arg
+ num_cleaned_frames, # Corresponds to num_cleaned_frames_arg
+ lora_names_states # Corresponds to lora_names_states_arg
+ ]
+ # Add LoRA sliders to the input list
+ ips.extend([lora_sliders[lora] for lora in lora_names])
+
+
+ # --- Connect Buttons ---
+ def handle_start_button(selected_model, *args):
+ # For other model types, use the regular process function
+ return process_with_queue_update(selected_model, *args)
+
+ def handle_batch_add_to_queue(*args):
+ # The last argument will be the list of files from batch_input_images
+ batch_files = args[-1]
+ if not batch_files or not isinstance(batch_files, list):
+ print("No batch images provided.")
+ return
+
+ print(f"Starting batch processing for {len(batch_files)} images.")
+
+ # Reconstruct the arguments for the single process function, excluding the batch files list
+ single_job_args = list(args[:-1])
+
+ # The first argument to process_with_queue_update is model_type
+ model_type_arg = single_job_args.pop(0)
+
+ # Keep track of the seed
+ current_seed = single_job_args[6] # seed is the 7th element in the ips list
+ randomize_seed_arg = single_job_args[7] # randomize_seed is the 8th
+
+ for image_path in batch_files:
+ # --- FIX IS HERE ---
+ # Load the image from the path into a NumPy array
+ try:
+ pil_image = Image.open(image_path).convert("RGB")
+ numpy_image = np.array(pil_image)
+ except Exception as e:
+ print(f"Error loading batch image {image_path}: {e}. Skipping.")
+ continue
+ # --- END OF FIX ---
+
+ # Replace the single input_image argument with the loaded NumPy image
+ current_job_args = single_job_args[:]
+ current_job_args[0] = numpy_image # Use the loaded numpy_image
+ current_job_args[6] = current_seed # Set the seed for the current job
+
+ # Call the original processing function with the modified arguments
+ process_with_queue_update(model_type_arg, *current_job_args)
+
+ # If randomize seed is checked, generate a new one for the next image
+ if randomize_seed_arg:
+ current_seed = random.randint(0, 21474)
+
+ print("Batch processing complete. All jobs added to the queue.")
+
+ # Validation ensures the start button is only enabled when appropriate
+ def update_start_button_state(*args):
+ """
+ Validation fails if a video model is selected and no input video is provided.
+ Updates the start button interactivity and validation message visibility.
+ Handles variable inputs from different Gradio event chains.
+ """
+ # The required values are the last two arguments provided by the Gradio event
+ if len(args) >= 2:
+ selected_model = args[-2]
+ input_video_value = args[-1]
+ else:
+ # Fallback or error handling if not enough arguments are received
+ # This might happen if the event is triggered in an unexpected way
+ print(f"Warning: update_start_button_state received {len(args)} args, expected at least 2.")
+ # Default to a safe state (button disabled)
+ return gr.Button(value="❌ Error", interactive=False), gr.update(visible=True)
+
+ video_provided = input_video_value is not None
+
+ if is_video_model(selected_model) and not video_provided:
+ # Video model selected, but no video provided
+ return gr.Button(value="❌ Missing Video", interactive=False), gr.update(visible=True)
+ else:
+ # Either not a video model, or video model selected and video provided
+ return gr.update(value="🚀 Add to Queue", interactive=True), gr.update(visible=False)
+ # Function to update button state before processing
+ def update_button_before_processing(selected_model, *args):
+ # First update the button to show "Adding..." and disable it
+ # Also return current stats so they don't get blanked out during the "Adding..." phase
+ qs_data, qs_text = update_stats()
+ return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(value="⏳ Adding...", interactive=False), gr.update(), qs_data, qs_text, gr.update(), gr.update() # Added update for video_input_required_message
+
+ # Connect the start button to first update its state
+ start_button.click(
+ fn=update_button_before_processing,
+ inputs=[model_type] + ips,
+ outputs=[result_video, current_job_id, preview_image, top_preview_image, progress_desc, progress_bar, start_button, end_button, queue_status, queue_stats_display, seed, video_input_required_message]
+ ).then(
+ # Then process the job
+ fn=handle_start_button,
+ inputs=[model_type] + ips,
+ outputs=[result_video, current_job_id, preview_image, progress_desc, progress_bar, start_button, end_button, queue_status, queue_stats_display, seed, video_input_required_message] # Added video_input_required_message
+ ).then( # Ensure validation is re-checked after job processing completes
+ fn=update_start_button_state,
+ inputs=[model_type, input_video], # Current values of model_type and input_video
+ outputs=[start_button, video_input_required_message]
+ )
+
+ def show_batch_gallery(files):
+ return gr.update(value=files, visible=True) if files else gr.update(visible=False)
+
+ batch_input_images.change(
+ fn=show_batch_gallery,
+ inputs=[batch_input_images],
+ outputs=[batch_input_gallery]
+ )
+
+ # We need to gather all the same inputs as the single 'Add to Queue' button, plus the new file input
+ batch_ips = [model_type] + ips + [batch_input_images]
+
+ add_batch_to_queue_btn.click(
+ fn=handle_batch_add_to_queue,
+ inputs=batch_ips,
+ outputs=None # No direct output updates from this button
+ ).then(
+ fn=update_stats, # Refresh the queue stats in the UI
+ inputs=None,
+ outputs=[queue_status, queue_stats_display]
+ ).then(
+ # This new block checks for a running job and updates the monitor UI
+ fn=check_for_current_job,
+ inputs=None,
+ outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
+ ).then(
+ # NEW: Update latents display layout after loading queue to ensure correct visibility
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+ # --- START OF REFACTORED XY PLOT EVENT WIRING ---
+ # Get the process button from the created components
+ xy_plot_process_btn = xy_plot_components["process_btn"]
+
+ # Prepare the process function with its static dependencies (job_queue, settings)
+ fn_xy_process_with_deps = functools.partial(xy_plot_process, job_queue, settings)
+
+ # Construct the full list of inputs for the click handler in the correct order
+ c = xy_plot_components
+ xy_plot_input_components = [
+ c["model_type"], c["input_image"], c["end_frame_image_original"],
+ c["end_frame_strength_original"], c["latent_type"], c["prompt"],
+ c["blend_sections"], c["steps"], c["total_second_length"],
+ resolutionW, resolutionH, # The components from the main UI
+ c["seed"], c["randomize_seed"],
+ c["use_teacache"], c["teacache_num_steps"], c["teacache_rel_l1_thresh"],
+ c["use_magcache"], c["magcache_threshold"], c["magcache_max_consecutive_skips"], c["magcache_retention_ratio"],
+ c["latent_window_size"], c["cfg"], c["gs"], c["rs"],
+ c["gpu_memory_preservation"], c["mp4_crf"],
+ c["axis_x_switch"], c["axis_x_value_text"], c["axis_x_value_dropdown"],
+ c["axis_y_switch"], c["axis_y_value_text"], c["axis_y_value_dropdown"],
+ c["axis_z_switch"], c["axis_z_value_text"], c["axis_z_value_dropdown"],
+ c["lora_selector"]
+ ]
+ # LoRA sliders are in a dictionary, so we add their values to the list
+ xy_plot_input_components.extend(c["lora_sliders"].values())
+
+ # Wire the click handler for the XY Plot button
+ xy_plot_process_btn.click(
+ fn=fn_xy_process_with_deps,
+ inputs=xy_plot_input_components,
+ outputs=[xy_plot_status, xy_plot_output]
+ ).then(
+ fn=update_stats,
+ inputs=None,
+ outputs=[queue_status, queue_stats_display]
+ ).then(
+ fn=check_for_current_job,
+ inputs=None,
+ outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
+ ).then(
+ # NEW: Update latents display layout after XY plot to ensure correct visibility
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+ # --- END OF REFACTORED XY PLOT EVENT WIRING ---
+
+
+
+ # MODIFIED: on_model_type_change to handle new "XY Plot" option
+ def on_model_type_change(selected_model):
+ is_xy_plot = selected_model == "XY Plot"
+ is_ui_video_model_flag = is_video_model(selected_model)
+ shows_end_frame = selected_model in ["Original with Endframe", "Video with Endframe"]
+
+ return (
+ gr.update(visible=not is_xy_plot), # standard_generation_group
+ gr.update(visible=is_xy_plot), # xy_group
+ gr.update(visible=not is_xy_plot and not is_ui_video_model_flag), # image_input_group
+ gr.update(visible=not is_xy_plot and is_ui_video_model_flag), # video_input_group
+ gr.update(visible=not is_xy_plot and shows_end_frame), # end_frame_group_original
+ gr.update(visible=not is_xy_plot and shows_end_frame), # end_frame_slider_group
+ gr.update(visible=not is_xy_plot), # start_button
+ gr.update(visible=is_xy_plot) # xy_plot_process_btn
+ )
+
+ # Model change listener
+ model_type.change(
+ fn=on_model_type_change,
+ inputs=model_type,
+ outputs=[
+ standard_generation_group,
+ xy_group,
+ image_input_group,
+ video_input_group,
+ end_frame_group_original,
+ end_frame_slider_group,
+ start_button,
+ xy_plot_process_btn # This is the button returned from the dictionary
+ ]
+ ).then( # Also trigger validation after model type changes
+ fn=update_start_button_state,
+ inputs=[model_type, input_video],
+ outputs=[start_button, video_input_required_message]
+ )
+
+ # Connect input_video change to the validation function
+ input_video.change(
+ fn=update_start_button_state,
+ inputs=[model_type, input_video],
+ outputs=[start_button, video_input_required_message]
+ )
+ # Also trigger validation when video is cleared
+ input_video.clear(
+ fn=update_start_button_state,
+ inputs=[model_type, input_video],
+ outputs=[start_button, video_input_required_message]
+ )
+
+
+
+ # Auto-monitor the current job when job_id changes
+ current_job_id.change(
+ fn=monitor_fn,
+ inputs=[current_job_id],
+ outputs=[result_video, preview_image, top_preview_image, progress_desc, progress_bar, start_button, end_button]
+ ).then(
+ fn=update_stats, # When a monitor finishes, always update the stats.
+ inputs=None,
+ outputs=[queue_status, queue_stats_display]
+ ).then( # re-validate button state
+ fn=update_start_button_state,
+ inputs=[model_type, input_video],
+ outputs=[start_button, video_input_required_message]
+ ).then(
+ # NEW: Update latents display layout after monitoring to ensure correct visibility
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+ # The "end_button" (Cancel Job) is the trigger for the next job's monitor.
+ # When a job is cancelled, we check for the next one.
+ end_button.click(
+ fn=end_process_with_update,
+ outputs=[queue_status, queue_stats_display, end_button, current_job_id]
+ ).then(
+ fn=check_for_current_job_and_monitor,
+ inputs=[],
+ outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar, queue_status, queue_stats_display]
+ ).then(
+ # NEW: Update latents display layout after job handoff to ensure correct visibility
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+ load_queue_button.click(
+ fn=load_queue_from_json,
+ inputs=[],
+ outputs=[queue_status, queue_stats_display]
+ ).then( # ADD THIS .then() CLAUSE
+ fn=check_for_current_job,
+ inputs=[],
+ outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
+ ).then(
+ # NEW: Update latents display layout after loading queue to ensure correct visibility
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+ import_queue_file.change(
+ fn=import_queue_from_file,
+ inputs=[import_queue_file],
+ outputs=[queue_status, queue_stats_display]
+ ).then( # ADD THIS .then() CLAUSE
+ fn=check_for_current_job,
+ inputs=[],
+ outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar]
+ ).then(
+ # NEW: Update latents display layout after importing queue to ensure correct visibility
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+
+ # --- Connect Queue Refresh ---
+ # The update_stats function is now defined much earlier.
+
+ # REMOVED: refresh_stats_btn.click - Toolbar refresh button is no longer needed
+ # refresh_stats_btn.click(
+ # fn=update_stats,
+ # inputs=None,
+ # outputs=[queue_status, queue_stats_display]
+ # )
+
+ # Set up auto-refresh for queue status
+ # Instead of using a timer with 'every' parameter, we'll use the queue refresh button
+ # and rely on manual refreshes. The user can click the refresh button in the toolbar
+ # to update the stats.
+
+ # --- Connect LoRA UI ---
+ # Function to update slider visibility based on selection
+ def update_lora_sliders(selected_loras):
+ updates = []
+ # Suppress dummy LoRA from workaround for the single lora bug.
+ # Filter out the dummy LoRA for display purposes in the dropdown
+ actual_selected_loras_for_display = [lora for lora in selected_loras if lora != DUMMY_LORA_NAME]
+ updates.append(gr.update(value=actual_selected_loras_for_display)) # First update is for the dropdown itself
+
+ # Need to handle potential missing keys if lora_names changes dynamically
+ # lora_names is from the create_interface scope
+ for lora_name_key in lora_names: # Iterate using lora_names to maintain order
+ if lora_name_key == DUMMY_LORA_NAME: # Check for dummy LoRA
+ updates.append(gr.update(visible=False))
+ else:
+ # Visibility of sliders should be based on actual_selected_loras_for_display
+ updates.append(gr.update(visible=(lora_name_key in actual_selected_loras_for_display)))
+ return updates # This list will be correctly ordered
+
+ # Connect the dropdown to the sliders
+ lora_selector.change(
+ fn=update_lora_sliders,
+ inputs=[lora_selector],
+ outputs=[lora_selector] + [lora_sliders[lora] for lora in lora_names if lora in lora_sliders]
+ )
+
+ def apply_preset(preset_name, model_type):
+ if not preset_name:
+ # Create a list of empty updates matching the number of components
+ return [gr.update()] * len(ui_components)
+
+ with open(PRESET_FILE, 'r') as f:
+ data = json.load(f)
+ preset = data.get(model_type, {}).get(preset_name, {})
+
+ # Initialize updates for all components
+ updates = {key: gr.update() for key in ui_components.keys()}
+
+ # Update components based on the preset
+ for key, value in preset.items():
+ if key in updates:
+ updates[key] = gr.update(value=value)
+
+ # Handle LoRA sliders specifically
+ if 'lora_values' in preset and isinstance(preset['lora_values'], dict):
+ lora_values_dict = preset['lora_values']
+ for lora_name, lora_value in lora_values_dict.items():
+ if lora_name in updates:
+ updates[lora_name] = gr.update(value=lora_value)
+
+ # Convert the dictionary of updates to a list in the correct order
+ return [updates[key] for key in ui_components.keys()]
+
+ def save_preset(preset_name, model_type, *args):
+ if not preset_name:
+ return gr.update()
+
+ # Ensure the directory exists
+ os.makedirs(os.path.dirname(PRESET_FILE), exist_ok=True)
+
+ if not os.path.exists(PRESET_FILE):
+ with open(PRESET_FILE, 'w') as f:
+ json.dump({}, f)
+
+ with open(PRESET_FILE, 'r') as f:
+ data = json.load(f)
+
+ if model_type not in data:
+ data[model_type] = {}
+
+ keys = list(ui_components.keys())
+
+ # Create a dictionary from the passed arguments
+ args_dict = {keys[i]: args[i] for i in range(len(keys))}
+
+ # Build the preset data from the arguments dictionary
+ preset_data = {key: args_dict[key] for key in ui_components.keys() if key not in lora_sliders}
+
+ # Handle LoRA values separately
+ selected_loras = args_dict.get("lora_selector", [])
+ lora_values = {}
+ for lora_name in selected_loras:
+ if lora_name in args_dict:
+ lora_values[lora_name] = args_dict[lora_name]
+
+ preset_data['lora_values'] = lora_values
+
+ # Remove individual lora sliders from the top-level preset data
+ for lora_name in lora_sliders:
+ if lora_name in preset_data:
+ del preset_data[lora_name]
+
+ data[model_type][preset_name] = preset_data
+
+ with open(PRESET_FILE, 'w') as f:
+ json.dump(data, f, indent=2)
+
+ return gr.update(choices=load_presets(model_type), value=preset_name)
+
+ def delete_preset(preset_name, model_type):
+ if not preset_name:
+ return gr.update(), gr.update(visible=True), gr.update(visible=False)
+
+ with open(PRESET_FILE, 'r') as f:
+ data = json.load(f)
+
+ if model_type in data and preset_name in data[model_type]:
+ del data[model_type][preset_name]
+
+ with open(PRESET_FILE, 'w') as f:
+ json.dump(data, f, indent=2)
+
+ return gr.update(choices=load_presets(model_type), value=None), gr.update(visible=True), gr.update(visible=False)
+
+ # --- Connect Preset UI ---
+ # Without this refresh, if you define a new preset for the Startup Model Type, and then try to select it in settings, it won't show up.
+ def refresh_settings_tab_startup_presets_if_needed(generate_tab_model_type_value, settings_tab_startup_model_type_value):
+ # generate_tab_model_type_value is the model for which a preset was just saved
+ # settings_tab_startup_model_type_value is the current selection in the startup model dropdown on settings tab
+ if generate_tab_model_type_value == settings_tab_startup_model_type_value and settings_tab_startup_model_type_value != "None":
+ return update_startup_preset_dropdown_choices(settings_tab_startup_model_type_value)
+ return gr.update()
+
+ ui_components = {
+ # Prompts
+ "prompt": prompt,
+ "n_prompt": n_prompt,
+ "blend_sections": blend_sections,
+ # Basic Params
+ "steps": steps,
+ "total_second_length": total_second_length,
+ "resolutionW": resolutionW,
+ "resolutionH": resolutionH,
+ "seed": seed,
+ "randomize_seed": randomize_seed,
+ # Advanced Params
+ "gs": gs,
+ "cfg": cfg,
+ "rs": rs,
+ "latent_window_size": latent_window_size,
+ # Cache type (Mag/Tea/None)
+ "cache_type": cache_type,
+ # TeaCache
+ "teacache_num_steps": teacache_num_steps,
+ "teacache_rel_l1_thresh": teacache_rel_l1_thresh,
+ # MagCache
+ "magcache_threshold": magcache_threshold,
+ "magcache_max_consecutive_skips": magcache_max_consecutive_skips,
+ "magcache_retention_ratio": magcache_retention_ratio,
+ # Input Options
+ "latent_type": latent_type,
+ "end_frame_strength_original": end_frame_strength_original,
+ # Video Specific
+ "combine_with_source": combine_with_source,
+ "num_cleaned_frames": num_cleaned_frames,
+ # LoRAs
+ "lora_selector": lora_selector,
+ **lora_sliders
+ }
+
+ model_type.change(
+ fn=lambda mt: (gr.update(choices=load_presets(mt)), gr.update(label=f"{mt} Presets")),
+ inputs=[model_type],
+ outputs=[preset_dropdown, preset_accordion]
+ )
+
+ preset_dropdown.select(
+ fn=apply_preset,
+ inputs=[preset_dropdown, model_type],
+ outputs=list(ui_components.values())
+ ).then(
+ lambda name: name,
+ inputs=[preset_dropdown],
+ outputs=[preset_name_textbox]
+ )
+
+ save_preset_button.click(
+ fn=save_preset,
+ inputs=[preset_name_textbox, model_type, *list(ui_components.values())],
+ outputs=[preset_dropdown] # preset_dropdown is on Generate tab
+ ).then(
+ fn=refresh_settings_tab_startup_presets_if_needed,
+ inputs=[model_type, startup_model_type_dropdown], # model_type (Generate tab), startup_model_type_dropdown (Settings tab)
+ outputs=[startup_preset_name_dropdown] # startup_preset_name_dropdown (Settings tab)
+ )
+
+ def show_delete_confirmation():
+ return gr.update(visible=False), gr.update(visible=True)
+
+ def hide_delete_confirmation():
+ return gr.update(visible=True), gr.update(visible=False)
+
+ delete_preset_button.click(
+ fn=show_delete_confirmation,
+ outputs=[save_preset_button, confirm_delete_row]
+ )
+
+ confirm_delete_no_btn.click(
+ fn=hide_delete_confirmation,
+ outputs=[save_preset_button, confirm_delete_row]
+ )
+
+ confirm_delete_yes_btn.click(
+ fn=delete_preset,
+ inputs=[preset_dropdown, model_type],
+ outputs=[preset_dropdown, save_preset_button, confirm_delete_row]
+ )
+
+ # --- Definition of apply_startup_settings (AFTER ui_components and apply_preset are defined) ---
+ # This function needs access to `settings`, `model_type` (Generate tab Radio),
+ # `preset_dropdown` (Generate tab Dropdown), `preset_name_textbox` (Generate tab Textbox),
+ # `ui_components` (dict of all other UI elements), `load_presets`, and `apply_preset`.
+ # All these are available in the scope of `create_interface`.
+ def apply_startup_settings():
+ startup_model_val = settings.get("startup_model_type", "None")
+ startup_preset_val = settings.get("startup_preset_name", None)
+
+ # Default updates (no change)
+ model_type_update = gr.update()
+ preset_dropdown_update = gr.update()
+ preset_name_textbox_update = gr.update()
+
+ # ui_components is now defined
+ ui_components_updates_list = [gr.update() for _ in ui_components]
+
+ if startup_model_val and startup_model_val != "None":
+ model_type_update = gr.update(value=startup_model_val)
+
+ presets_for_startup_model = load_presets(startup_model_val) # load_presets is defined earlier
+ preset_dropdown_update = gr.update(choices=presets_for_startup_model)
+ preset_name_textbox_update = gr.update(value="")
+
+ if startup_preset_val and startup_preset_val in presets_for_startup_model:
+ preset_dropdown_update = gr.update(choices=presets_for_startup_model, value=startup_preset_val)
+ preset_name_textbox_update = gr.update(value=startup_preset_val)
+
+ # apply_preset is now defined
+ ui_components_updates_list = apply_preset(startup_preset_val, startup_model_val)
+
+ # NEW: Ensure latents_display_top checkbox reflects the current setting
+ latents_display_top_update = gr.update(value=get_latents_display_top())
+
+ return tuple([model_type_update, preset_dropdown_update, preset_name_textbox_update] + ui_components_updates_list + [latents_display_top_update])
+
+
+ # --- Auto-refresh for Toolbar System Stats Monitor (Timer) ---
+ main_toolbar_system_stats_timer = gr.Timer(2, active=True)
+
+ main_toolbar_system_stats_timer.tick(
+ fn=tb_get_formatted_toolbar_stats, # Function imported from toolbox_app.py
+ inputs=None,
+ outputs=[ # Target the Textbox components
+ toolbar_ram_display_component,
+ toolbar_vram_display_component,
+ toolbar_gpu_display_component
+ ]
+ )
+
+ # --- Connect Metadata Loading ---
+ # Function to load metadata from JSON file
+ def load_metadata_from_json(json_path):
+ # Define the total number of output components to handle errors gracefully
+ num_outputs = 20 + len(lora_sliders)
+
+ if not json_path:
+ # Return empty updates for all components if no file is provided
+ return [gr.update()] * num_outputs
+
+ try:
+ with open(json_path, 'r') as f:
+ metadata = json.load(f)
+
+ # Extract values from metadata with defaults
+ prompt_val = metadata.get('prompt')
+ n_prompt_val = metadata.get('negative_prompt')
+ seed_val = metadata.get('seed')
+ steps_val = metadata.get('steps')
+ total_second_length_val = metadata.get('total_second_length')
+ end_frame_strength_val = metadata.get('end_frame_strength')
+ model_type_val = metadata.get('model_type')
+ lora_weights = metadata.get('loras', {})
+ latent_window_size_val = metadata.get('latent_window_size')
+ resolutionW_val = metadata.get('resolutionW')
+ resolutionH_val = metadata.get('resolutionH')
+ blend_sections_val = metadata.get('blend_sections')
+ # Determine cache_type from metadata, with fallback for older formats
+ cache_type_val = metadata.get('cache_type')
+ if cache_type_val is None:
+ use_magcache = metadata.get('use_magcache', False)
+ use_teacache = metadata.get('use_teacache', False)
+ if use_magcache:
+ cache_type_val = "MagCache"
+ elif use_teacache:
+ cache_type_val = "TeaCache"
+ else:
+ cache_type_val = "None"
+ magcache_threshold_val = metadata.get('magcache_threshold')
+ magcache_max_consecutive_skips_val = metadata.get('magcache_max_consecutive_skips')
+ magcache_retention_ratio_val = metadata.get('magcache_retention_ratio')
+ teacache_num_steps_val = metadata.get('teacache_num_steps')
+ teacache_rel_l1_thresh_val = metadata.get('teacache_rel_l1_thresh')
+ latent_type_val = metadata.get('latent_type')
+ combine_with_source_val = metadata.get('combine_with_source')
+
+ # Get the names of the selected LoRAs from the metadata
+ selected_lora_names = list(lora_weights.keys())
+
+ print(f"Loaded metadata from JSON: {json_path}")
+ print(f"Model Type: {model_type_val}, Prompt: {prompt_val}, Seed: {seed_val}, LoRAs: {selected_lora_names}")
+
+ # Create a list of UI updates
+ updates = [
+ gr.update(value=prompt_val) if prompt_val is not None else gr.update(),
+ gr.update(value=n_prompt_val) if n_prompt_val is not None else gr.update(),
+ gr.update(value=seed_val) if seed_val is not None else gr.update(),
+ gr.update(value=steps_val) if steps_val is not None else gr.update(),
+ gr.update(value=total_second_length_val) if total_second_length_val is not None else gr.update(),
+ gr.update(value=end_frame_strength_val) if end_frame_strength_val is not None else gr.update(),
+ gr.update(value=model_type_val) if model_type_val else gr.update(),
+ gr.update(value=selected_lora_names) if selected_lora_names else gr.update(),
+ gr.update(value=latent_window_size_val) if latent_window_size_val is not None else gr.update(),
+ gr.update(value=resolutionW_val) if resolutionW_val is not None else gr.update(),
+ gr.update(value=resolutionH_val) if resolutionH_val is not None else gr.update(),
+ gr.update(value=blend_sections_val) if blend_sections_val is not None else gr.update(),
+ gr.update(value=cache_type_val),
+ gr.update(value=magcache_threshold_val),
+ gr.update(value=magcache_max_consecutive_skips_val),
+ gr.update(value=magcache_retention_ratio_val),
+ gr.update(value=teacache_num_steps_val) if teacache_num_steps_val is not None else gr.update(),
+ gr.update(value=teacache_rel_l1_thresh_val) if teacache_rel_l1_thresh_val is not None else gr.update(),
+ gr.update(value=latent_type_val) if latent_type_val else gr.update(),
+ gr.update(value=combine_with_source_val) if combine_with_source_val else gr.update(),
+ ]
+
+ # Update LoRA sliders based on loaded weights
+ for lora in lora_names:
+ if lora in lora_weights:
+ updates.append(gr.update(value=lora_weights[lora], visible=True))
+ else:
+ # Hide sliders for LoRAs not in the metadata
+ updates.append(gr.update(visible=False))
+
+ return updates
+
+ except Exception as e:
+ print(f"Error loading metadata: {e}")
+ import traceback
+ traceback.print_exc()
+ # Return empty updates for all components on error
+ return [gr.update()] * num_outputs
+
+
+ # Connect JSON metadata loader for Original tab
+ json_upload.change(
+ fn=load_metadata_from_json,
+ inputs=[json_upload],
+ outputs=[
+ prompt,
+ n_prompt,
+ seed,
+ steps,
+ total_second_length,
+ end_frame_strength_original,
+ model_type,
+ lora_selector,
+ latent_window_size,
+ resolutionW,
+ resolutionH,
+ blend_sections,
+ cache_type,
+ magcache_threshold,
+ magcache_max_consecutive_skips,
+ magcache_retention_ratio,
+ teacache_num_steps,
+ teacache_rel_l1_thresh,
+ latent_type,
+ combine_with_source
+ ] + [lora_sliders[lora] for lora in lora_names]
+ )
+
+
+ # --- Helper Functions (defined within create_interface scope if needed by handlers) ---
+ # Function to get queue statistics
+ def get_queue_stats():
+ try:
+ # Get all jobs from the queue
+ jobs = job_queue.get_all_jobs()
+
+ # Count jobs by status
+ status_counts = {
+ "QUEUED": 0,
+ "RUNNING": 0,
+ "COMPLETED": 0,
+ "FAILED": 0,
+ "CANCELLED": 0
+ }
+
+ for job in jobs:
+ if hasattr(job, 'status'):
+ status = str(job.status) # Use str() for safety
+ if status in status_counts:
+ status_counts[status] += 1
+
+ # Format the display text
+ stats_text = f"Queue: {status_counts['QUEUED']} | Running: {status_counts['RUNNING']} | Completed: {status_counts['COMPLETED']} | Failed: {status_counts['FAILED']} | Cancelled: {status_counts['CANCELLED']}"
+
+ return f"{stats_text}
"
+
+ except Exception as e:
+ print(f"Error getting queue stats: {e}")
+ return "Error loading queue stats
"
+
+ # Add footer with social links
+ with gr.Row(elem_id="footer"):
+ with gr.Column(scale=1):
+ gr.HTML(f"""
+
+ """)
+
+ # Add CSS for footer
+
+ # gr.HTML("""
+ #
+ # """)
+
+ # --- Function to update latents display layout on interface load ---
+ def update_latents_layout_on_load():
+ """Update latents display layout based on saved setting when interface loads"""
+ return create_latents_layout_update()
+
+ # Connect the auto-check function to the interface load event
+ block.load(
+ fn=check_for_current_job_and_monitor, # Use the new combined function
+ inputs=[],
+ outputs=[current_job_id, result_video, preview_image, top_preview_image, progress_desc, progress_bar, queue_status, queue_stats_display]
+
+ ).then(
+ fn=apply_startup_settings, # apply_startup_settings is now defined
+ inputs=None,
+ outputs=[model_type, preset_dropdown, preset_name_textbox] + list(ui_components.values()) + [latents_display_top] # ui_components is now defined
+ ).then(
+ fn=update_start_button_state, # Ensure button state is correct after startup settings
+ inputs=[model_type, input_video],
+ outputs=[start_button, video_input_required_message]
+ ).then(
+ # NEW: Update latents display layout based on saved setting
+ fn=create_latents_layout_update,
+ inputs=None,
+ outputs=[top_preview_row, preview_image]
+ )
+
+ # --- Prompt Enhancer Connection ---
+ def handle_enhance_prompt(current_prompt_text):
+ """Calls the LLM enhancer and returns the updated text."""
+ if not current_prompt_text:
+ return ""
+ print("UI: Enhance button clicked. Sending prompt to enhancer.")
+ enhanced_text = enhance_prompt(current_prompt_text)
+ print(f"UI: Received enhanced prompt: {enhanced_text}")
+ return gr.update(value=enhanced_text)
+
+ enhance_prompt_btn.click(
+ fn=handle_enhance_prompt,
+ inputs=[prompt],
+ outputs=[prompt]
+ )
+
+ # --- Captioner Connection ---
+ def handle_caption(input_image, prompt):
+ """Calls the LLM enhancer and returns the updated text."""
+ if input_image is None:
+ return prompt # Return current prompt if no image is provided
+ caption_text = caption_image(input_image)
+ print(f"UI: Received caption: {caption_text}")
+ return gr.update(value=caption_text)
+
+ caption_btn.click(
+ fn=handle_caption,
+ inputs=[input_image, prompt],
+ outputs=[prompt]
+ )
+
+ return block
+
+# --- Top-level Helper Functions (Used by Gradio callbacks, must be defined outside create_interface) ---
+
+def format_queue_status(jobs):
+ """Format job data for display in the queue status table"""
+ rows = []
+ for job in jobs:
+ created = time.strftime('%H:%M:%S', time.localtime(job.created_at)) if job.created_at else ""
+ started = time.strftime('%H:%M:%S', time.localtime(job.started_at)) if job.started_at else ""
+ completed = time.strftime('%H:%M:%S', time.localtime(job.completed_at)) if job.completed_at else ""
+
+ # Calculate elapsed time
+ elapsed_time = ""
+ if job.started_at:
+ if job.completed_at:
+ start_datetime = datetime.datetime.fromtimestamp(job.started_at)
+ complete_datetime = datetime.datetime.fromtimestamp(job.completed_at)
+ elapsed_seconds = (complete_datetime - start_datetime).total_seconds()
+ elapsed_time = f"{elapsed_seconds:.2f}s"
+ else:
+ # For running jobs, calculate elapsed time from now
+ start_datetime = datetime.datetime.fromtimestamp(job.started_at)
+ current_datetime = datetime.datetime.now()
+ elapsed_seconds = (current_datetime - start_datetime).total_seconds()
+ elapsed_time = f"{elapsed_seconds:.2f}s (running)"
+
+ # Get generation type from job data
+ generation_type = getattr(job, 'generation_type', 'Original')
+
+ # Get thumbnail from job data and format it as HTML for display
+ thumbnail = getattr(job, 'thumbnail', None)
+ thumbnail_html = f'' if thumbnail else ""
+
+ rows.append([
+ job.id[:6] + '...',
+ generation_type,
+ job.status.value,
+ created,
+ started,
+ completed,
+ elapsed_time,
+ thumbnail_html # Add formatted thumbnail HTML to row data
+ ])
+ return rows
+
+# Create the queue status update function (wrapper around format_queue_status)
+def update_queue_status_with_thumbnails(): # Function name is now slightly misleading, but keep for now to avoid breaking clicks
+ # This function is likely called by the refresh button and potentially the timer
+ # It needs access to the job_queue object
+ # Assuming job_queue is accessible globally or passed appropriately
+ # For now, let's assume it's globally accessible as defined in studio.py
+ # If not, this needs adjustment based on how job_queue is managed.
+ try:
+ # Need access to the global job_queue instance from studio.py
+ # This might require restructuring or passing job_queue differently.
+ # For now, assuming it's accessible (this might fail if run standalone)
+ from __main__ import job_queue # Attempt to import from main script scope
+
+ jobs = job_queue.get_all_jobs()
+ for job in jobs:
+ if job.status == JobStatus.PENDING:
+ job.queue_position = job_queue.get_queue_position(job.id)
+
+ if job_queue.current_job:
+ job_queue.current_job.status = JobStatus.RUNNING
+
+ return format_queue_status(jobs)
+ except ImportError:
+ print("Error: Could not import job_queue. Queue status update might fail.")
+ return [] # Return empty list on error
+ except Exception as e:
+ print(f"Error updating queue status: {e}")
+ return []
diff --git a/modules/prompt_handler.py b/modules/prompt_handler.py
index d2ce826b..41396bc8 100644
--- a/modules/prompt_handler.py
+++ b/modules/prompt_handler.py
@@ -1,164 +1,164 @@
-import re
-from dataclasses import dataclass
-from typing import List, Optional
-
-
-@dataclass
-class PromptSection:
- """Represents a section of the prompt with specific timing information"""
- prompt: str
- start_time: float = 0 # in seconds
- end_time: Optional[float] = None # in seconds, None means until the end
-
-
-def snap_to_section_boundaries(prompt_sections: List[PromptSection], latent_window_size: int, fps: int = 30) -> List[PromptSection]:
- """
- Adjust timestamps to align with model's internal section boundaries
-
- Args:
- prompt_sections: List of PromptSection objects
- latent_window_size: Size of the latent window used in the model
- fps: Frames per second (default: 30)
-
- Returns:
- List of PromptSection objects with aligned timestamps
- """
- section_duration = (latent_window_size * 4 - 3) / fps # Duration of one section in seconds
-
- aligned_sections = []
- for section in prompt_sections:
- # Snap start time to nearest section boundary
- aligned_start = round(section.start_time / section_duration) * section_duration
-
- # Snap end time to nearest section boundary
- aligned_end = None
- if section.end_time is not None:
- aligned_end = round(section.end_time / section_duration) * section_duration
-
- # Ensure minimum section length
- if aligned_end is not None and aligned_end <= aligned_start:
- aligned_end = aligned_start + section_duration
-
- aligned_sections.append(PromptSection(
- prompt=section.prompt,
- start_time=aligned_start,
- end_time=aligned_end
- ))
-
- return aligned_sections
-
-
-def parse_timestamped_prompt(prompt_text: str, total_duration: float, latent_window_size: int = 9, generation_type: str = "Original") -> List[PromptSection]:
- """
- Parse a prompt with timestamps in the format [0s-2s: text] or [3s: text]
-
- Args:
- prompt_text: The input prompt text with optional timestamp sections
- total_duration: Total duration of the video in seconds
- latent_window_size: Size of the latent window used in the model
- generation_type: Type of generation ("Original" or "F1")
-
- Returns:
- List of PromptSection objects with timestamps aligned to section boundaries
- and reversed to account for reverse generation (only for Original type)
- """
- # Default prompt for the entire duration if no timestamps are found
- if "[" not in prompt_text or "]" not in prompt_text:
- return [PromptSection(prompt=prompt_text.strip())]
-
- sections = []
- # Find all timestamp sections [time: text]
- timestamp_pattern = r'\[(\d+(?:\.\d+)?s)(?:-(\d+(?:\.\d+)?s))?\s*:\s*(.*?)\]'
- regular_text = prompt_text
-
- for match in re.finditer(timestamp_pattern, prompt_text):
- start_time_str = match.group(1)
- end_time_str = match.group(2)
- section_text = match.group(3).strip()
-
- # Convert time strings to seconds
- start_time = float(start_time_str.rstrip('s'))
- end_time = float(end_time_str.rstrip('s')) if end_time_str else None
-
- sections.append(PromptSection(
- prompt=section_text,
- start_time=start_time,
- end_time=end_time
- ))
-
- # Remove the processed section from regular_text
- regular_text = regular_text.replace(match.group(0), "")
-
- # If there's any text outside of timestamp sections, use it as a default for the entire duration
- regular_text = regular_text.strip()
- if regular_text:
- sections.append(PromptSection(
- prompt=regular_text,
- start_time=0,
- end_time=None
- ))
-
- # Sort sections by start time
- sections.sort(key=lambda x: x.start_time)
-
- # Fill in end times if not specified
- for i in range(len(sections) - 1):
- if sections[i].end_time is None:
- sections[i].end_time = sections[i+1].start_time
-
- # Set the last section's end time to the total duration if not specified
- if sections and sections[-1].end_time is None:
- sections[-1].end_time = total_duration
-
- # Snap timestamps to section boundaries
- sections = snap_to_section_boundaries(sections, latent_window_size)
-
- # Only reverse timestamps for Original generation type
- if generation_type in ("Original", "Original with Endframe", "Video"):
- # Now reverse the timestamps to account for reverse generation
- reversed_sections = []
- for section in sections:
- reversed_start = total_duration - section.end_time if section.end_time is not None else 0
- reversed_end = total_duration - section.start_time
- reversed_sections.append(PromptSection(
- prompt=section.prompt,
- start_time=reversed_start,
- end_time=reversed_end
- ))
-
- # Sort the reversed sections by start time
- reversed_sections.sort(key=lambda x: x.start_time)
- return reversed_sections
-
- return sections
-
-
-def get_section_boundaries(latent_window_size: int = 9, count: int = 10) -> str:
- """
- Calculate and format section boundaries for UI display
-
- Args:
- latent_window_size: Size of the latent window used in the model
- count: Number of boundaries to display
-
- Returns:
- Formatted string of section boundaries
- """
- section_duration = (latent_window_size * 4 - 3) / 30
- return ", ".join([f"{i*section_duration:.1f}s" for i in range(count)])
-
-
-def get_quick_prompts() -> List[List[str]]:
- """
- Get a list of example timestamped prompts
-
- Returns:
- List of example prompts formatted for Gradio Dataset
- """
- prompts = [
- '[0s: The person waves hello] [2s: The person jumps up and down] [4s: The person does a spin]',
- '[0s: The person raises both arms slowly] [2s: The person claps hands enthusiastically]',
- '[0s: Person gives thumbs up] [1.1s: Person smiles and winks] [2.2s: Person shows two thumbs down]',
- '[0s: Person looks surprised] [1.1s: Person raises arms above head] [2.2s-3.3s: Person puts hands on hips]'
- ]
- return [[x] for x in prompts]
+import re
+from dataclasses import dataclass
+from typing import List, Optional
+
+
+@dataclass
+class PromptSection:
+ """Represents a section of the prompt with specific timing information"""
+ prompt: str
+ start_time: float = 0 # in seconds
+ end_time: Optional[float] = None # in seconds, None means until the end
+
+
+def snap_to_section_boundaries(prompt_sections: List[PromptSection], latent_window_size: int, fps: int = 30) -> List[PromptSection]:
+ """
+ Adjust timestamps to align with model's internal section boundaries
+
+ Args:
+ prompt_sections: List of PromptSection objects
+ latent_window_size: Size of the latent window used in the model
+ fps: Frames per second (default: 30)
+
+ Returns:
+ List of PromptSection objects with aligned timestamps
+ """
+ section_duration = (latent_window_size * 4 - 3) / fps # Duration of one section in seconds
+
+ aligned_sections = []
+ for section in prompt_sections:
+ # Snap start time to nearest section boundary
+ aligned_start = round(section.start_time / section_duration) * section_duration
+
+ # Snap end time to nearest section boundary
+ aligned_end = None
+ if section.end_time is not None:
+ aligned_end = round(section.end_time / section_duration) * section_duration
+
+ # Ensure minimum section length
+ if aligned_end is not None and aligned_end <= aligned_start:
+ aligned_end = aligned_start + section_duration
+
+ aligned_sections.append(PromptSection(
+ prompt=section.prompt,
+ start_time=aligned_start,
+ end_time=aligned_end
+ ))
+
+ return aligned_sections
+
+
+def parse_timestamped_prompt(prompt_text: str, total_duration: float, latent_window_size: int = 9, generation_type: str = "Original") -> List[PromptSection]:
+ """
+ Parse a prompt with timestamps in the format [0s-2s: text] or [3s: text]
+
+ Args:
+ prompt_text: The input prompt text with optional timestamp sections
+ total_duration: Total duration of the video in seconds
+ latent_window_size: Size of the latent window used in the model
+ generation_type: Type of generation ("Original" or "F1")
+
+ Returns:
+ List of PromptSection objects with timestamps aligned to section boundaries
+ and reversed to account for reverse generation (only for Original type)
+ """
+ # Default prompt for the entire duration if no timestamps are found
+ if "[" not in prompt_text or "]" not in prompt_text:
+ return [PromptSection(prompt=prompt_text.strip())]
+
+ sections = []
+ # Find all timestamp sections [time: text]
+ timestamp_pattern = r'\[(\d+(?:\.\d+)?s)(?:-(\d+(?:\.\d+)?s))?\s*:\s*(.*?)\]'
+ regular_text = prompt_text
+
+ for match in re.finditer(timestamp_pattern, prompt_text):
+ start_time_str = match.group(1)
+ end_time_str = match.group(2)
+ section_text = match.group(3).strip()
+
+ # Convert time strings to seconds
+ start_time = float(start_time_str.rstrip('s'))
+ end_time = float(end_time_str.rstrip('s')) if end_time_str else None
+
+ sections.append(PromptSection(
+ prompt=section_text,
+ start_time=start_time,
+ end_time=end_time
+ ))
+
+ # Remove the processed section from regular_text
+ regular_text = regular_text.replace(match.group(0), "")
+
+ # If there's any text outside of timestamp sections, use it as a default for the entire duration
+ regular_text = regular_text.strip()
+ if regular_text:
+ sections.append(PromptSection(
+ prompt=regular_text,
+ start_time=0,
+ end_time=None
+ ))
+
+ # Sort sections by start time
+ sections.sort(key=lambda x: x.start_time)
+
+ # Fill in end times if not specified
+ for i in range(len(sections) - 1):
+ if sections[i].end_time is None:
+ sections[i].end_time = sections[i+1].start_time
+
+ # Set the last section's end time to the total duration if not specified
+ if sections and sections[-1].end_time is None:
+ sections[-1].end_time = total_duration
+
+ # Snap timestamps to section boundaries
+ sections = snap_to_section_boundaries(sections, latent_window_size)
+
+ # Only reverse timestamps for Original generation type
+ if generation_type in ("Original", "Original with Endframe", "Video"):
+ # Now reverse the timestamps to account for reverse generation
+ reversed_sections = []
+ for section in sections:
+ reversed_start = total_duration - section.end_time if section.end_time is not None else 0
+ reversed_end = total_duration - section.start_time
+ reversed_sections.append(PromptSection(
+ prompt=section.prompt,
+ start_time=reversed_start,
+ end_time=reversed_end
+ ))
+
+ # Sort the reversed sections by start time
+ reversed_sections.sort(key=lambda x: x.start_time)
+ return reversed_sections
+
+ return sections
+
+
+def get_section_boundaries(latent_window_size: int = 9, count: int = 10) -> str:
+ """
+ Calculate and format section boundaries for UI display
+
+ Args:
+ latent_window_size: Size of the latent window used in the model
+ count: Number of boundaries to display
+
+ Returns:
+ Formatted string of section boundaries
+ """
+ section_duration = (latent_window_size * 4 - 3) / 30
+ return ", ".join([f"{i*section_duration:.1f}s" for i in range(count)])
+
+
+def get_quick_prompts() -> List[List[str]]:
+ """
+ Get a list of example timestamped prompts
+
+ Returns:
+ List of example prompts formatted for Gradio Dataset
+ """
+ prompts = [
+ '[0s: The person waves hello] [2s: The person jumps up and down] [4s: The person does a spin]',
+ '[0s: The person raises both arms slowly] [2s: The person claps hands enthusiastically]',
+ '[0s: Person gives thumbs up] [1.1s: Person smiles and winks] [2.2s: Person shows two thumbs down]',
+ '[0s: Person looks surprised] [1.1s: Person raises arms above head] [2.2s-3.3s: Person puts hands on hips]'
+ ]
+ return [[x] for x in prompts]
diff --git a/modules/settings.py b/modules/settings.py
index e3bd25df..d2dfc28d 100644
--- a/modules/settings.py
+++ b/modules/settings.py
@@ -1,88 +1,88 @@
-import json
-from pathlib import Path
-from typing import Dict, Any, Optional
-import os
-
-class Settings:
- def __init__(self):
- # Get the project root directory (where settings.py is located)
- project_root = Path(__file__).parent.parent
-
- self.settings_file = project_root / ".framepack" / "settings.json"
- self.settings_file.parent.mkdir(parents=True, exist_ok=True)
-
- # Set default paths relative to project root
- self.default_settings = {
- "save_metadata": True,
- "gpu_memory_preservation": 6,
- "output_dir": str(project_root / "outputs"),
- "metadata_dir": str(project_root / "outputs"),
- "lora_dir": str(project_root / "loras"),
- "gradio_temp_dir": str(project_root / "temp"),
- "input_files_dir": str(project_root / "input_files"), # New setting for input files
- "auto_save_settings": True,
- "gradio_theme": "default",
- "mp4_crf": 16,
- "clean_up_videos": True,
- "override_system_prompt": False,
- "auto_cleanup_on_startup": False, # ADDED: New setting for startup cleanup
- "latents_display_top": False, # NEW: Control latents preview position (False = right column, True = top of interface)
- "system_prompt_template": "{\"template\": \"<|start_header_id|>system<|end_header_id|>\\n\\nDescribe the video by detailing the following aspects: 1. The main content and theme of the video.2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.4. background environment, light, style and atmosphere.5. camera angles, movements, and transitions used in the video:<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{}<|eot_id|>\", \"crop_start\": 95}",
- "startup_model_type": "None",
- "startup_preset_name": None,
- "enhancer_prompt_template": """You are a creative assistant for a text-to-video generator. Your task is to take a user's prompt and make it more descriptive, vivid, and detailed. Focus on visual elements. Do not change the core action, but embellish it.
-
-User prompt: "{text_to_enhance}"
-
-Enhanced prompt:"""
- }
- self.settings = self.load_settings()
-
- def load_settings(self) -> Dict[str, Any]:
- """Load settings from file or return defaults"""
- if self.settings_file.exists():
- try:
- with open(self.settings_file, 'r') as f:
- loaded_settings = json.load(f)
- # Merge with defaults to ensure all settings exist
- settings = self.default_settings.copy()
- settings.update(loaded_settings)
- return settings
- except Exception as e:
- print(f"Error loading settings: {e}")
- return self.default_settings.copy()
- return self.default_settings.copy()
-
- def save_settings(self, **kwargs):
- """Save settings to file. Accepts keyword arguments for any settings to update."""
- # Update self.settings with any provided keyword arguments
- self.settings.update(kwargs)
- # Ensure all default fields are present
- for k, v in self.default_settings.items():
- self.settings.setdefault(k, v)
-
- # Ensure directories exist for relevant fields
- for dir_key in ["output_dir", "metadata_dir", "lora_dir", "gradio_temp_dir"]:
- dir_path = self.settings.get(dir_key)
- if dir_path:
- os.makedirs(dir_path, exist_ok=True)
-
- # Save to file
- with open(self.settings_file, 'w') as f:
- json.dump(self.settings, f, indent=4)
-
- def get(self, key: str, default: Any = None) -> Any:
- """Get a setting value"""
- return self.settings.get(key, default)
-
- def set(self, key: str, value: Any) -> None:
- """Set a setting value"""
- self.settings[key] = value
- if self.settings.get("auto_save_settings", True):
- self.save_settings()
-
- def update(self, settings: Dict[str, Any]) -> None:
- """Update multiple settings at once"""
- self.settings.update(settings)
- if self.settings.get("auto_save_settings", True):
- self.save_settings()
+import json
+from pathlib import Path
+from typing import Dict, Any, Optional
+import os
+
+class Settings:
+ def __init__(self):
+ # Get the project root directory (where settings.py is located)
+ project_root = Path(__file__).parent.parent
+
+ self.settings_file = project_root / ".framepack" / "settings.json"
+ self.settings_file.parent.mkdir(parents=True, exist_ok=True)
+
+ # Set default paths relative to project root
+ self.default_settings = {
+ "save_metadata": True,
+ "gpu_memory_preservation": 6,
+ "output_dir": str(project_root / "outputs"),
+ "metadata_dir": str(project_root / "outputs"),
+ "lora_dir": str(project_root / "loras"),
+ "gradio_temp_dir": str(project_root / "temp"),
+ "input_files_dir": str(project_root / "input_files"), # New setting for input files
+ "auto_save_settings": True,
+ "gradio_theme": "default",
+ "mp4_crf": 16,
+ "clean_up_videos": True,
+ "override_system_prompt": False,
+ "auto_cleanup_on_startup": False, # ADDED: New setting for startup cleanup
+ "latents_display_top": False, # NEW: Control latents preview position (False = right column, True = top of interface)
+ "system_prompt_template": "{\"template\": \"<|start_header_id|>system<|end_header_id|>\\n\\nDescribe the video by detailing the following aspects: 1. The main content and theme of the video.2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.4. background environment, light, style and atmosphere.5. camera angles, movements, and transitions used in the video:<|eot_id|><|start_header_id|>user<|end_header_id|>\\n\\n{}<|eot_id|>\", \"crop_start\": 95}",
+ "startup_model_type": "None",
+ "startup_preset_name": None,
+ "enhancer_prompt_template": """You are a creative assistant for a text-to-video generator. Your task is to take a user's prompt and make it more descriptive, vivid, and detailed. Focus on visual elements. Do not change the core action, but embellish it.
+
+User prompt: "{text_to_enhance}"
+
+Enhanced prompt:"""
+ }
+ self.settings = self.load_settings()
+
+ def load_settings(self) -> Dict[str, Any]:
+ """Load settings from file or return defaults"""
+ if self.settings_file.exists():
+ try:
+ with open(self.settings_file, 'r') as f:
+ loaded_settings = json.load(f)
+ # Merge with defaults to ensure all settings exist
+ settings = self.default_settings.copy()
+ settings.update(loaded_settings)
+ return settings
+ except Exception as e:
+ print(f"Error loading settings: {e}")
+ return self.default_settings.copy()
+ return self.default_settings.copy()
+
+ def save_settings(self, **kwargs):
+ """Save settings to file. Accepts keyword arguments for any settings to update."""
+ # Update self.settings with any provided keyword arguments
+ self.settings.update(kwargs)
+ # Ensure all default fields are present
+ for k, v in self.default_settings.items():
+ self.settings.setdefault(k, v)
+
+ # Ensure directories exist for relevant fields
+ for dir_key in ["output_dir", "metadata_dir", "lora_dir", "gradio_temp_dir"]:
+ dir_path = self.settings.get(dir_key)
+ if dir_path:
+ os.makedirs(dir_path, exist_ok=True)
+
+ # Save to file
+ with open(self.settings_file, 'w') as f:
+ json.dump(self.settings, f, indent=4)
+
+ def get(self, key: str, default: Any = None) -> Any:
+ """Get a setting value"""
+ return self.settings.get(key, default)
+
+ def set(self, key: str, value: Any) -> None:
+ """Set a setting value"""
+ self.settings[key] = value
+ if self.settings.get("auto_save_settings", True):
+ self.save_settings()
+
+ def update(self, settings: Dict[str, Any]) -> None:
+ """Update multiple settings at once"""
+ self.settings.update(settings)
+ if self.settings.get("auto_save_settings", True):
+ self.save_settings()
diff --git a/modules/toolbox/esrgan_core.py b/modules/toolbox/esrgan_core.py
index 9055eeb9..cbd23f23 100644
--- a/modules/toolbox/esrgan_core.py
+++ b/modules/toolbox/esrgan_core.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
import os
import torch
import gc
@@ -7,10 +9,18 @@
from pathlib import Path
from huggingface_hub import snapshot_download
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from realesrgan import RealESRGANer
-from realesrgan.archs.srvgg_arch import SRVGGNetCompact
-from basicsr.utils.download_util import load_file_from_url # Import for direct downloads
+
+# Conditional import for BasicSR and RealESRGAN
+try:
+ from basicsr.archs.rrdbnet_arch import RRDBNet
+ from realesrgan import RealESRGANer
+ from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+ from basicsr.utils.download_util import load_file_from_url
+ BASICSR_AVAILABLE = True
+except ImportError:
+ BASICSR_AVAILABLE = False
+ print("[WARNING] BasicSR not available. ESRGAN upscaling features will be disabled.")
+ print(" To enable, install BasicSR: pip install basicsr")
# Conditional import for GFPGAN
try:
@@ -28,13 +38,16 @@
class ESRGANUpscaler:
def __init__(self, message_manager: MessageManager, device: torch.device):
+ if not BASICSR_AVAILABLE:
+ raise ImportError("BasicSR is not available. Cannot initialize ESRGANUpscaler.")
+
self.message_manager = message_manager
self.device = device
self.model_dir = Path(MODEL_ESRGAN_PATH)
self.gfpgan_model_dir = Path(MODEL_GFPGAN_PATH) # GFPGAN model directory
os.makedirs(self.model_dir, exist_ok=True)
os.makedirs(self.gfpgan_model_dir, exist_ok=True) # Ensure GFPGAN model dir exists
-
+
self.supported_models = {
"RealESRGAN_x2plus": {
diff --git a/modules/toolbox/setup_ffmpeg.py b/modules/toolbox/setup_ffmpeg.py
index 08c6502c..3bae305a 100644
--- a/modules/toolbox/setup_ffmpeg.py
+++ b/modules/toolbox/setup_ffmpeg.py
@@ -86,10 +86,10 @@ def setup_ffmpeg():
os.chmod(ffmpeg_path, 0o755)
os.chmod(ffprobe_path, 0o755)
- print(f"✅ FFmpeg setup complete. Binaries are in: {bin_dir}")
+ print(f"[OK] FFmpeg setup complete. Binaries are in: {bin_dir}")
except Exception as e:
- print(f"\n❌ Error setting up FFmpeg: {e}")
+ print(f"\n[ERROR] Error setting up FFmpeg: {e}")
import traceback
traceback.print_exc()
print("\nPlease download FFmpeg manually and place the 'ffmpeg' and 'ffprobe' executables in the 'bin' directory.")
diff --git a/modules/toolbox/toolbox_processor.py b/modules/toolbox/toolbox_processor.py
index 86c9cccd..53a04f62 100644
--- a/modules/toolbox/toolbox_processor.py
+++ b/modules/toolbox/toolbox_processor.py
@@ -21,9 +21,15 @@
from torchvision.transforms.functional import to_tensor, to_pil_image
from modules.toolbox.rife_core import RIFEHandler
-from modules.toolbox.esrgan_core import ESRGANUpscaler
from modules.toolbox.message_manager import MessageManager
+# Conditional import for ESRGAN
+try:
+ from modules.toolbox.esrgan_core import ESRGANUpscaler
+ ESRGAN_AVAILABLE = True
+except ImportError:
+ ESRGAN_AVAILABLE = False
+
device_name_str = devicetorch.get(torch)
VIDEO_QUALITY = 8 # Used by imageio.mimwrite quality/quantizer
@@ -33,7 +39,17 @@ def __init__(self, message_manager: MessageManager, settings):
self.message_manager = message_manager
self.rife_handler = RIFEHandler(message_manager)
self.device_obj = torch.device(device_name_str) # Store device_obj
- self.esrgan_upscaler = ESRGANUpscaler(message_manager, self.device_obj)
+
+ # Only initialize ESRGAN if available
+ if ESRGAN_AVAILABLE:
+ try:
+ self.esrgan_upscaler = ESRGANUpscaler(message_manager, self.device_obj)
+ except (ImportError, NameError) as e:
+ print(f"[WARNING] Could not initialize ESRGAN: {e}")
+ self.esrgan_upscaler = None
+ else:
+ self.esrgan_upscaler = None
+
self.settings = settings
self.project_root = Path(__file__).resolve().parents[2]
@@ -1593,13 +1609,18 @@ def tb_apply_filters(self, video_path, brightness, contrast, saturation, tempera
finally: gc.collect()
- def tb_upscale_video(self, video_path, model_key: str, output_scale_factor_ui: float,
- tile_size: int, enhance_face: bool,
+ def tb_upscale_video(self, video_path, model_key: str, output_scale_factor_ui: float,
+ tile_size: int, enhance_face: bool,
denoise_strength_ui: float | None,
use_streaming: bool, # New parameter from UI
progress=gr.Progress()):
if video_path is None: self.message_manager.add_warning("No input video for upscaling."); return None
-
+
+ # Check if ESRGAN is available
+ if self.esrgan_upscaler is None:
+ self.message_manager.add_error("ESRGAN upscaling is not available. BasicSR module is required.")
+ return None
+
reader = None
writer = None
final_output_path = None
diff --git a/modules/toolbox_app.py b/modules/toolbox_app.py
index 2977470c..6087e644 100644
--- a/modules/toolbox_app.py
+++ b/modules/toolbox_app.py
@@ -59,7 +59,15 @@ def wrapper(self, *args, **kwargs):
# --- Local Application Imports ---
from modules.settings import Settings
-from modules.toolbox.esrgan_core import ESRGANUpscaler
+
+# Conditional import for ESRGAN
+try:
+ from modules.toolbox.esrgan_core import ESRGANUpscaler
+ ESRGAN_AVAILABLE = True
+except ImportError as e:
+ ESRGAN_AVAILABLE = False
+ print(f"[WARNING] ESRGAN upscaling not available: {e}")
+
from modules.toolbox.message_manager import MessageManager
from modules.toolbox.rife_core import RIFEHandler
from modules.toolbox.setup_ffmpeg import setup_ffmpeg
@@ -249,7 +257,7 @@ def _get_default_workflow_params():
# Gets default values from filter settings and adds other op defaults
params = TB_DEFAULT_FILTER_SETTINGS.copy()
params.update({
- "upscale_model": list(tb_processor.esrgan_upscaler.supported_models.keys())[0] if tb_processor.esrgan_upscaler.supported_models else None,
+ "upscale_model": list(tb_processor.esrgan_upscaler.supported_models.keys())[0] if (tb_processor.esrgan_upscaler and tb_processor.esrgan_upscaler.supported_models) else None,
"upscale_factor": 2.0,
"tile_size": 0,
"enhance_face": False,
@@ -1198,13 +1206,17 @@ def tb_create_video_toolbox_ui():
gr.Markdown("Upscale video resolution using Real-ESRGAN.")
with gr.Row():
with gr.Column(scale=2):
+ # Check if ESRGAN is available
+ esrgan_choices = list(tb_processor.esrgan_upscaler.supported_models.keys()) if (tb_processor.esrgan_upscaler and tb_processor.esrgan_upscaler.supported_models) else []
+ esrgan_default = esrgan_choices[0] if esrgan_choices else None
+
tb_upscale_model_select = gr.Dropdown(
- choices=list(tb_processor.esrgan_upscaler.supported_models.keys()),
- value=list(tb_processor.esrgan_upscaler.supported_models.keys())[0] if tb_processor.esrgan_upscaler.supported_models else None,
+ choices=esrgan_choices,
+ value=esrgan_default,
label="ESRGAN Model",
- info="Select the Real-ESRGAN model."
+ info="Select the Real-ESRGAN model." if esrgan_choices else "ESRGAN not available (BasicSR required)"
)
- default_model_key_init = list(tb_processor.esrgan_upscaler.supported_models.keys())[0] if tb_processor.esrgan_upscaler.supported_models else None
+ default_model_key_init = esrgan_default
initial_model_info_gr_val, initial_slider_gr_val, initial_denoise_gr_val = tb_get_model_info_and_update_scale_slider(default_model_key_init)
tb_selected_model_scale_display = gr.Textbox(
diff --git a/modules/video_queue.py b/modules/video_queue.py
index 52b09d1a..f79640d3 100644
--- a/modules/video_queue.py
+++ b/modules/video_queue.py
@@ -1,1655 +1,1655 @@
-import threading
-import time
-import uuid
-import json
-import os
-import zipfile
-import shutil
-from dataclasses import dataclass, field
-from enum import Enum
-from typing import Dict, Any, Optional, List
-import queue as queue_module # Renamed to avoid conflicts
-import io
-import base64
-from PIL import Image
-import numpy as np
-
-from diffusers_helper.thread_utils import AsyncStream
-from modules.pipelines.metadata_utils import create_metadata
-from modules.settings import Settings
-from diffusers_helper.gradio.progress_bar import make_progress_bar_html
-
-
-# Simple LIFO queue implementation to avoid dependency on queue.LifoQueue
-class SimpleLifoQueue:
- def __init__(self):
- self._queue = []
- self._mutex = threading.Lock()
- self._not_empty = threading.Condition(self._mutex)
-
- def put(self, item):
- with self._mutex:
- self._queue.append(item)
- self._not_empty.notify()
-
- def get(self):
- with self._not_empty:
- while not self._queue:
- self._not_empty.wait()
- return self._queue.pop()
-
- def task_done(self):
- pass # For compatibility with queue.Queue
-
-
-class JobStatus(Enum):
- PENDING = "pending"
- RUNNING = "running"
- COMPLETED = "completed"
- FAILED = "failed"
- CANCELLED = "cancelled"
-
-
-class JobType(Enum):
- SINGLE = "single"
- GRID = "grid"
-
-
-@dataclass
-class Job:
- id: str
- params: Dict[str, Any]
- status: JobStatus = JobStatus.PENDING
- job_type: JobType = JobType.SINGLE
- child_job_ids: List[str] = field(default_factory=list)
- parent_job_id: Optional[str] = None
- created_at: float = field(default_factory=time.time)
- started_at: Optional[float] = None
- completed_at: Optional[float] = None
- error: Optional[str] = None
- result: Optional[str] = None
- progress_data: Optional[Dict] = None
- queue_position: Optional[int] = None
- stream: Optional[Any] = None
- input_image: Optional[np.ndarray] = None
- latent_type: Optional[str] = None
- thumbnail: Optional[str] = None
- generation_type: Optional[str] = None # Added generation_type
- input_image_saved: bool = False # Flag to track if input image has been saved
- end_frame_image_saved: bool = False # Flag to track if end frame image has been saved
-
- def __post_init__(self):
- # Store generation type
- self.generation_type = self.params.get('model_type', 'Original') # Initialize generation_type
-
- # Store input image or latent type
- if 'input_image' in self.params and self.params['input_image'] is not None:
- self.input_image = self.params['input_image']
- # Create thumbnail
- if isinstance(self.input_image, np.ndarray):
- # Handle numpy array (image)
- img = Image.fromarray(self.input_image)
- img.thumbnail((100, 100))
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
- elif isinstance(self.input_image, str):
- # Handle string (video path)
- try:
- print(f"Attempting to extract thumbnail from video: {self.input_image}")
- # Try to extract frames from the video using imageio
- import imageio
-
- # Check if the file exists
- if not os.path.exists(self.input_image):
- print(f"Video file not found: {self.input_image}")
- raise FileNotFoundError(f"Video file not found: {self.input_image}")
-
- # Create outputs directory if it doesn't exist
- os.makedirs("outputs", exist_ok=True)
-
- # Try to open the video file
- try:
- reader = imageio.get_reader(self.input_image)
- print(f"Successfully opened video file with imageio")
- except Exception as e:
- print(f"Failed to open video with imageio: {e}")
- raise
-
- # Get the total number of frames
- num_frames = None
- try:
- # Try to get the number of frames from metadata
- meta_data = reader.get_meta_data()
- print(f"Video metadata: {meta_data}")
- num_frames = meta_data.get('nframes')
- if num_frames is None or num_frames == float('inf'):
- print("Number of frames not available in metadata")
- # If not available, try to count frames
- if hasattr(reader, 'count_frames'):
- print("Trying to count frames...")
- num_frames = reader.count_frames()
- print(f"Counted {num_frames} frames")
- except Exception as e:
- print(f"Error getting frame count: {e}")
- num_frames = None
-
- # If we couldn't determine the number of frames, read the last frame by iterating
- if num_frames is None or num_frames == float('inf'):
- print("Reading frames by iteration to find the last one")
- # Read frames until we reach the end
- frame_count = 0
- first_frame = None
- last_frame = None
- try:
- for frame in reader:
- if frame_count == 0:
- first_frame = frame
- last_frame = frame
- frame_count += 1
- # Print progress every 100 frames
- if frame_count % 100 == 0:
- print(f"Read {frame_count} frames...")
- print(f"Finished reading {frame_count} frames")
-
- # Save the first frame if available
- if first_frame is not None:
- print(f"Found first frame with shape: {first_frame.shape}")
- # DEBUG IMAGE SAVING REMOVED
- except Exception as e:
- print(f"Error reading frames: {e}")
-
- if last_frame is not None:
- print(f"Found last frame with shape: {last_frame.shape}")
-
- # DEBUG IMAGE SAVING REMOVED
- # Use the last frame for the thumbnail
- img = Image.fromarray(last_frame)
- img.thumbnail((100, 100))
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
- print("Successfully created thumbnail from last frame")
- else:
- print("No frames were read, using red thumbnail")
- # Fallback to red thumbnail if no frames were read - more visible for debugging
- img = Image.new('RGB', (100, 100), (255, 0, 0)) # Red for video
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
- else:
- # If we know the number of frames, try to get multiple frames for debugging
- try:
- # Try to get the first frame
- first_frame = None
- try:
- first_frame = reader.get_data(0)
- print(f"Got first frame with shape: {first_frame.shape}")
-
- # DEBUG IMAGE SAVING REMOVED
- except Exception as e:
- print(f"Error getting first frame: {e}")
-
- # Try to get a middle frame
- middle_frame = None
- try:
- middle_frame_idx = int(num_frames / 2)
- middle_frame = reader.get_data(middle_frame_idx)
- print(f"Got middle frame (frame {middle_frame_idx}) with shape: {middle_frame.shape}")
-
- # DEBUG IMAGE SAVING REMOVED
- except Exception as e:
- print(f"Error getting middle frame: {e}")
-
- # Try to get the last frame
- last_frame = None
- try:
- last_frame_idx = int(num_frames) - 1
- last_frame = reader.get_data(last_frame_idx)
- print(f"Got last frame (frame {last_frame_idx}) with shape: {last_frame.shape}")
-
- # DEBUG IMAGE SAVING REMOVED
- except Exception as e:
- print(f"Error getting last frame: {e}")
-
- # If we couldn't get the last frame directly, try to get it by iterating
- if last_frame is None:
- print("Trying to get last frame by iterating through all frames")
- try:
- for frame in reader:
- last_frame = frame
-
- if last_frame is not None:
- print(f"Got last frame by iteration with shape: {last_frame.shape}")
-
- # DEBUG IMAGE SAVING REMOVED
- except Exception as e:
- print(f"Error getting last frame by iteration: {e}")
-
- # Use the last frame for the thumbnail if available, otherwise use the middle or first frame
- frame_for_thumbnail = last_frame if last_frame is not None else (middle_frame if middle_frame is not None else first_frame)
-
- if frame_for_thumbnail is not None:
- # Convert to PIL Image and create a thumbnail
- img = Image.fromarray(frame_for_thumbnail)
- img.thumbnail((100, 100))
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
- print("Successfully created thumbnail from frame")
- else:
- print("No frames were extracted, using blue thumbnail")
- # Fallback to blue thumbnail if no frames were extracted
- img = Image.new('RGB', (100, 100), (0, 0, 255)) # Blue for video
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
- except Exception as e:
- # Fallback to blue thumbnail on error
- img = Image.new('RGB', (100, 100), (0, 0, 255)) # Blue for video
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
-
- # Close the reader
- try:
- reader.close()
- print("Successfully closed video reader")
- except Exception as e:
- print(f"Error closing reader: {e}")
-
- except Exception as e:
- print(f"Error extracting thumbnail from video: {e}")
- import traceback
- traceback.print_exc()
- # Fallback to bright green thumbnail on error to make it more visible
- img = Image.new('RGB', (100, 100), (0, 255, 0)) # Bright green for error
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
- print("Created bright green fallback thumbnail")
- else:
- # Handle other types
- self.thumbnail = None
- elif 'latent_type' in self.params:
- self.latent_type = self.params['latent_type']
- # Create a colored square based on latent type
- color_map = {
- "Black": (0, 0, 0),
- "White": (255, 255, 255),
- "Noise": (128, 128, 128),
- "Green Screen": (0, 177, 64)
- }
- color = color_map.get(self.latent_type, (0, 0, 0))
- img = Image.new('RGB', (100, 100), color)
- buffered = io.BytesIO()
- img.save(buffered, format="PNG")
- self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
-
-
-class VideoJobQueue:
- def __init__(self):
- self.queue = queue_module.Queue() # Using standard Queue instead of LifoQueue
- self.jobs = {}
- self.current_job = None
- self.lock = threading.Lock()
- self.worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
- self.worker_thread.start()
- self.worker_function = None # Will be set from outside
- self.is_processing = False # Flag to track if we're currently processing a job
-
- def set_worker_function(self, worker_function):
- """Set the worker function to use for processing jobs"""
- self.worker_function = worker_function
-
- def serialize_job(self, job):
- """Serialize a job to a JSON-compatible format"""
- try:
- # Create a simplified representation of the job
- serialized = {
- "id": job.id,
- "status": job.status.value,
- "created_at": job.created_at,
- "started_at": job.started_at,
- "completed_at": job.completed_at,
- "error": job.error,
- "result": job.result,
- "queue_position": job.queue_position,
- "generation_type": job.generation_type,
- }
-
- # Add simplified params (excluding complex objects)
- serialized_params = {}
- for k, v in job.params.items():
- if k not in ["input_image", "end_frame_image", "stream"]:
- # Try to include only JSON-serializable values
- try:
- # Test if value is JSON serializable
- json.dumps({k: v})
- serialized_params[k] = v
- except (TypeError, OverflowError):
- # Skip non-serializable values
- pass
-
- # Handle LoRA information specifically
- # Only include selected LoRAs for the generation
- if "selected_loras" in job.params and job.params["selected_loras"]:
- selected_loras = job.params["selected_loras"]
- # Ensure it's a list
- if not isinstance(selected_loras, list):
- selected_loras = [selected_loras] if selected_loras is not None else []
-
- # Get LoRA values if available
- lora_values = job.params.get("lora_values", [])
- if not isinstance(lora_values, list):
- lora_values = [lora_values] if lora_values is not None else []
-
- # Get loaded LoRA names
- lora_loaded_names = job.params.get("lora_loaded_names", [])
- if not isinstance(lora_loaded_names, list):
- lora_loaded_names = [lora_loaded_names] if lora_loaded_names is not None else []
-
- # Create LoRA data dictionary
- lora_data = {}
- for lora_name in selected_loras:
- try:
- # Find the index of the LoRA in loaded names
- idx = lora_loaded_names.index(lora_name) if lora_loaded_names else -1
- # Get the weight value
- weight = lora_values[idx] if lora_values and idx >= 0 and idx < len(lora_values) else 1.0
- # Handle weight as list
- if isinstance(weight, list):
- weight_value = weight[0] if weight and len(weight) > 0 else 1.0
- else:
- weight_value = weight
- # Store as float
- lora_data[lora_name] = float(weight_value)
- except (ValueError, IndexError):
- # Default weight if not found
- lora_data[lora_name] = 1.0
- except Exception as e:
- print(f"Error processing LoRA {lora_name}: {e}")
- lora_data[lora_name] = 1.0
-
- # Add to serialized params
- serialized_params["loras"] = lora_data
-
- serialized["params"] = serialized_params
-
- # Don't include the thumbnail as it can be very large and cause issues
- # if job.thumbnail:
- # serialized["thumbnail"] = job.thumbnail
-
- return serialized
- except Exception as e:
- print(f"Error serializing job {job.id}: {e}")
- # Return minimal information that should always be serializable
- return {
- "id": job.id,
- "status": job.status.value,
- "error": f"Error serializing: {str(e)}"
- }
-
- def save_queue_to_json(self):
- """Save the current queue to queue.json using the central metadata utility"""
- try:
- # Make a copy of job IDs to avoid holding the lock while serializing
- with self.lock:
- job_ids = list(self.jobs.keys())
-
- # Create a settings instance
- settings = Settings()
-
- # Create a directory to store queue images if it doesn't exist
- queue_images_dir = "queue_images"
- os.makedirs(queue_images_dir, exist_ok=True)
-
- # First, ensure all images are saved
- for job_id in job_ids:
- job = self.get_job(job_id)
- if job:
- # Save input image to disk if it exists and hasn't been saved yet
- if 'input_image' in job.params and isinstance(job.params['input_image'], np.ndarray) and not job.input_image_saved:
- input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
- try:
- Image.fromarray(job.params['input_image']).save(input_image_path)
- print(f"Saved input image for job {job_id} to {input_image_path}")
- # Mark the image as saved
- job.input_image_saved = True
- except Exception as e:
- print(f"Error saving input image for job {job_id}: {e}")
-
- # Save end frame image to disk if it exists and hasn't been saved yet
- if 'end_frame_image' in job.params and isinstance(job.params['end_frame_image'], np.ndarray) and not job.end_frame_image_saved:
- end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
- try:
- Image.fromarray(job.params['end_frame_image']).save(end_frame_image_path)
- print(f"Saved end frame image for job {job_id} to {end_frame_image_path}")
- # Mark the end frame image as saved
- job.end_frame_image_saved = True
- except Exception as e:
- print(f"Error saving end frame image for job {job_id}: {e}")
-
- # Now serialize jobs with the updated image saved flags
- serialized_jobs = {}
- for job_id in job_ids:
- job = self.get_job(job_id)
- if job:
- # Try to use metadata_utils.create_metadata if possible
- try:
- # Create metadata using the central utility
- metadata = create_metadata(job.params, job.id, settings.settings)
-
- # Add job status and other fields not included in metadata
- metadata.update({
- "id": job.id,
- "status": job.status.value,
- "created_at": job.created_at,
- "started_at": job.started_at,
- "completed_at": job.completed_at,
- "error": job.error,
- "result": job.result,
- "queue_position": job.queue_position,
- })
-
- # Add image paths to metadata if they've been saved
- if job.input_image_saved:
- input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
- if os.path.exists(input_image_path):
- metadata["saved_input_image_path"] = input_image_path
-
- if job.end_frame_image_saved:
- end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
- if os.path.exists(end_frame_image_path):
- metadata["saved_end_frame_image_path"] = end_frame_image_path
-
- serialized_jobs[job_id] = metadata
- except Exception as e:
- print(f"Error using metadata_utils for job {job_id}: {e}")
- # Fall back to the old serialization method
- serialized_jobs[job_id] = self.serialize_job(job)
-
- # Save to file
- with open("queue.json", "w") as f:
- json.dump(serialized_jobs, f, indent=2)
-
- # Clean up images for jobs that no longer exist
- self.cleanup_orphaned_images(job_ids)
- self.cleanup_orphaned_videos(job_ids)
-
- print(f"Saved {len(serialized_jobs)} jobs to queue.json")
- except Exception as e:
- print(f"Error saving queue to JSON: {e}")
-
- def cleanup_orphaned_videos(self, current_job_ids_uuids): # Renamed arg for clarity
- """
- Remove video files from input_files_dir for jobs that no longer exist
- or whose input_image_path does not point to them.
-
- Args:
- current_job_ids_uuids: List of job UUIDs currently in self.jobs
- """
- try:
- # Get the input_files_dir from settings to be robust
- settings = Settings()
- input_files_dir = settings.get("input_files_dir", "input_files")
- if not os.path.exists(input_files_dir):
- return
-
- # Normalize the managed input_files_dir path once
- norm_input_files_dir = os.path.normpath(input_files_dir)
- referenced_video_paths = set()
-
- with self.lock: # Access self.jobs safely
- for job_id_uuid in current_job_ids_uuids: # Iterate using the provided UUIDs
- job = self.jobs.get(job_id_uuid)
- if not (job and job.params):
- continue
-
- # Collect all potential video paths from the job parameters
- # Check for strings to avoid TypeError
- paths_to_consider = []
- p1 = job.params.get("input_image") # Primary path used by worker
- if isinstance(p1, str):
- paths_to_consider.append(p1)
-
- p2 = job.params.get("input_image_path") # Secondary/metadata path
- if isinstance(p2, str) and p2 not in paths_to_consider:
- paths_to_consider.append(p2)
-
- p3 = job.params.get("input_video") # Explicitly set during import
- if isinstance(p3, str) and p3 not in paths_to_consider:
- paths_to_consider.append(p3)
-
- for rel_or_abs_path in paths_to_consider:
- # Resolve to absolute path. If already absolute, abspath does nothing.
- # If relative, it's resolved against CWD (current working directory).
- abs_path = os.path.abspath(rel_or_abs_path)
- norm_abs_path = os.path.normpath(abs_path)
- # Check if this path is within the managed input_files_dir
- if norm_abs_path.startswith(norm_input_files_dir):
- referenced_video_paths.add(norm_abs_path)
-
- removed_count = 0
- for filename in os.listdir(input_files_dir):
- if filename.endswith(".mp4"): # Only process MP4 files
- file_path_to_check = os.path.normpath(os.path.join(input_files_dir, filename))
-
- if file_path_to_check not in referenced_video_paths:
- try:
- os.remove(file_path_to_check)
- removed_count += 1
- print(f"Removed orphaned video: {filename} (path: {file_path_to_check})")
- except Exception as e:
- print(f"Error removing orphaned video {filename}: {e}")
- if removed_count > 0:
- print(f"Cleaned up {removed_count} orphaned videos from {input_files_dir}")
- except Exception as e:
- print(f"Error cleaning up orphaned videos: {e}")
- import traceback
- traceback.print_exc()
-
- def cleanup_orphaned_images(self, current_job_ids):
- """
- Remove image files for jobs that no longer exist in the queue.
-
- Args:
- current_job_ids: List of job IDs currently in the queue
- """
- try:
- queue_images_dir = "queue_images"
- if not os.path.exists(queue_images_dir):
- return
-
- # Convert to set for faster lookups
- current_job_ids = set(current_job_ids)
-
- # Check all files in the queue_images directory
- removed_count = 0
- for filename in os.listdir(queue_images_dir):
- # Only process PNG files with our naming pattern
- if filename.endswith(".png") and ("_input.png" in filename or "_end_frame.png" in filename):
- # Extract job ID from filename
- parts = filename.split("_")
- if len(parts) >= 2:
- job_id = parts[0]
-
- # If job ID is not in current jobs, remove the file
- if job_id not in current_job_ids:
- file_path = os.path.join(queue_images_dir, filename)
- try:
- os.remove(file_path)
- removed_count += 1
- print(f"Removed orphaned image: {filename}")
- except Exception as e:
- print(f"Error removing orphaned image {filename}: {e}")
-
- if removed_count > 0:
- print(f"Cleaned up {removed_count} orphaned images")
- except Exception as e:
- print(f"Error cleaning up orphaned images: {e}")
-
-
- def synchronize_queue_images(self):
- """
- Synchronize the queue_images directory with the current jobs in the queue.
- This ensures all necessary images are saved and only images for removed jobs are deleted.
- """
- try:
- queue_images_dir = "queue_images"
- os.makedirs(queue_images_dir, exist_ok=True)
-
- # Get all current job IDs
- with self.lock:
- current_job_ids = set(self.jobs.keys())
-
- # Get all image files in the queue_images directory
- existing_image_files = set()
- if os.path.exists(queue_images_dir):
- for filename in os.listdir(queue_images_dir):
- if filename.endswith(".png") and ("_input.png" in filename or "_end_frame.png" in filename):
- existing_image_files.add(filename)
-
- # Extract job IDs from filenames
- file_job_ids = set()
- for filename in existing_image_files:
- # Extract job ID from filename (format: "{job_id}_input.png" or "{job_id}_end_frame.png")
- parts = filename.split("_")
- if len(parts) >= 2:
- job_id = parts[0]
- file_job_ids.add(job_id)
-
- # Find job IDs in files that are no longer in the queue
- removed_job_ids = file_job_ids - current_job_ids
-
- # Delete images for jobs that have been removed from the queue
- removed_count = 0
- for job_id in removed_job_ids:
- input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
- end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
-
- if os.path.exists(input_image_path):
- try:
- os.remove(input_image_path)
- removed_count += 1
- print(f"Removed image for deleted job: {input_image_path}")
- except Exception as e:
- print(f"Error removing image {input_image_path}: {e}")
-
- if os.path.exists(end_frame_image_path):
- try:
- os.remove(end_frame_image_path)
- removed_count += 1
- print(f"Removed image for deleted job: {end_frame_image_path}")
- except Exception as e:
- print(f"Error removing image {end_frame_image_path}: {e}")
-
- # Now ensure all current jobs have their images saved
- saved_count = 0
- with self.lock:
- for job_id, job in self.jobs.items():
- # Only save images for running or completed jobs
- if job.status in [JobStatus.RUNNING, JobStatus.COMPLETED]:
- # Save input image if it exists and hasn't been saved yet
- if 'input_image' in job.params and isinstance(job.params['input_image'], np.ndarray) and not job.input_image_saved:
- input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
- try:
- Image.fromarray(job.params['input_image']).save(input_image_path)
- job.input_image_saved = True
- saved_count += 1
- print(f"Saved input image for job {job_id}")
- except Exception as e:
- print(f"Error saving input image for job {job_id}: {e}")
-
- # Save end frame image if it exists and hasn't been saved yet
- if 'end_frame_image' in job.params and isinstance(job.params['end_frame_image'], np.ndarray) and not job.end_frame_image_saved:
- end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
- try:
- Image.fromarray(job.params['end_frame_image']).save(end_frame_image_path)
- job.end_frame_image_saved = True
- saved_count += 1
- print(f"Saved end frame image for job {job_id}")
- except Exception as e:
- print(f"Error saving end frame image for job {job_id}: {e}")
-
- # Save the queue to ensure the image paths are properly referenced
- self.save_queue_to_json()
-
- if removed_count > 0 or saved_count > 0:
- print(f"Queue image synchronization: removed {removed_count} images, saved {saved_count} images")
-
- except Exception as e:
- print(f"Error synchronizing queue images: {e}")
-
-
- def add_job(self, params, job_type=JobType.SINGLE, child_job_params_list=None, parent_job_id=None):
- """Add a job to the queue and return its ID"""
- job_id = str(uuid.uuid4())
-
- # For grid jobs, create child jobs first
- child_job_ids = []
- if job_type == JobType.GRID and child_job_params_list:
- with self.lock:
- for child_params in child_job_params_list:
- child_job_id = str(uuid.uuid4())
- child_job_ids.append(child_job_id)
- child_job = Job(
- id=child_job_id,
- params=child_params,
- status=JobStatus.PENDING,
- job_type=JobType.SINGLE, # Children are single jobs
- parent_job_id=job_id,
- created_at=time.time(),
- progress_data={},
- stream=AsyncStream(),
- input_image_saved=False,
- end_frame_image_saved=False
- )
- self.jobs[child_job_id] = child_job
- print(f" - Created child job {child_job_id} for grid job {job_id}")
-
- job = Job(
- id=job_id,
- params=params,
- status=JobStatus.PENDING,
- job_type=job_type,
- child_job_ids=child_job_ids,
- parent_job_id=parent_job_id,
- created_at=time.time(),
- progress_data={},
- stream=AsyncStream(),
- input_image_saved=False,
- end_frame_image_saved=False
- )
-
- with self.lock:
- print(f"Adding job {job_id} (type: {job_type.value}) to queue.")
- self.jobs[job_id] = job
- self.queue.put(job_id) # Only the parent (or single) job is added to the queue initially
-
- # Save the queue to JSON after adding a new job (outside the lock)
- try:
- self.save_queue_to_json()
- except Exception as e:
- print(f"Error saving queue to JSON after adding job: {e}")
-
- return job_id
-
- def get_job(self, job_id):
- """Get job by ID"""
- with self.lock:
- return self.jobs.get(job_id)
-
- def get_all_jobs(self):
- """Get all jobs"""
- with self.lock:
- return list(self.jobs.values())
-
- def cancel_job(self, job_id):
- """Cancel a pending job"""
- with self.lock:
- job = self.jobs.get(job_id)
- if not job:
- return False
-
- if job.status == JobStatus.PENDING:
- job.status = JobStatus.CANCELLED
- job.completed_at = time.time() # Mark completion time
- result = True
- elif job.status == JobStatus.RUNNING:
- # Send cancel signal to the job's stream
- if hasattr(job, 'stream') and job.stream:
- job.stream.input_queue.push('end')
-
- # Mark job as cancelled (this will be confirmed when the worker processes the end signal)
- job.status = JobStatus.CANCELLED
- job.completed_at = time.time() # Mark completion time
-
- # Let the worker loop handle the transition to the next job
- # This ensures the current job is fully processed before switching
- # DEBUG PRINT REMOVED
- result = True
- else:
- result = False
-
- # Save the queue to JSON after cancelling a job (outside the lock)
- if result:
- try:
- self.save_queue_to_json()
- except Exception as e:
- print(f"Error saving queue to JSON after cancelling job: {e}")
-
- return result
-
- def clear_queue(self):
- """Cancel all pending jobs in the queue"""
- cancelled_count = 0
- try:
- # First, make a copy of all pending job IDs to avoid modifying the dictionary during iteration
- with self.lock:
- # Get all pending job IDs
- pending_job_ids = [job_id for job_id, job in self.jobs.items()
- if job.status == JobStatus.PENDING]
-
- # Cancel each pending job individually
- for job_id in pending_job_ids:
- try:
- with self.lock:
- job = self.jobs.get(job_id)
- if job and job.status == JobStatus.PENDING:
- job.status = JobStatus.CANCELLED
- job.completed_at = time.time()
- cancelled_count += 1
- except Exception as e:
- print(f"Error cancelling job {job_id}: {e}")
-
- # Now clear the queue
- with self.lock:
- # Clear the queue (this doesn't affect running jobs)
- queue_items_cleared = 0
- try:
- while not self.queue.empty():
- try:
- self.queue.get_nowait()
- self.queue.task_done()
- queue_items_cleared += 1
- except queue_module.Empty:
- break
- except Exception as e:
- print(f"Error clearing queue: {e}")
-
- # Save the updated queue state
- try:
- self.save_queue_to_json()
- except Exception as e:
- print(f"Error saving queue state: {e}")
-
- # Synchronize queue images after clearing the queue
- if cancelled_count > 0:
- self.synchronize_queue_images()
-
- print(f"Cleared {cancelled_count} jobs from the queue")
- return cancelled_count
- except Exception as e:
- import traceback
- print(f"Error in clear_queue: {e}")
- traceback.print_exc()
- return 0
-
- def clear_completed_jobs(self):
- """Remove cancelled or completed jobs from the queue"""
- removed_count = 0
- try:
- # First, make a copy of all completed/cancelled job IDs to avoid modifying the dictionary during iteration
- with self.lock:
- # Get all completed or cancelled job IDs
- completed_job_ids = [job_id for job_id, job in self.jobs.items()
- if job.status in [JobStatus.COMPLETED, JobStatus.CANCELLED]]
-
- # Remove each completed/cancelled job individually
- for job_id in completed_job_ids:
- try:
- with self.lock:
- if job_id in self.jobs:
- del self.jobs[job_id]
- removed_count += 1
- except Exception as e:
- print(f"Error removing job {job_id}: {e}")
-
- # Save the updated queue state
- try:
- self.save_queue_to_json()
- except Exception as e:
- print(f"Error saving queue state: {e}")
-
- # Synchronize queue images after removing completed jobs
- if removed_count > 0:
- self.synchronize_queue_images()
-
- print(f"Removed {removed_count} completed/cancelled jobs from the queue")
- return removed_count
- except Exception as e:
- import traceback
- print(f"Error in clear_completed_jobs: {e}")
- traceback.print_exc()
- return 0
-
- def get_queue_position(self, job_id):
- """Get position in queue (0 = currently running)"""
- with self.lock:
- job = self.jobs.get(job_id)
- if not job:
- return None
-
- if job.status == JobStatus.RUNNING:
- return 0
-
- if job.status != JobStatus.PENDING:
- return None
-
- # Count pending jobs ahead in queue
- position = 1 # Start at 1 because 0 means running
- for j in self.jobs.values():
- if (j.status == JobStatus.PENDING and
- j.created_at < job.created_at):
- position += 1
- return position
-
- def update_job_progress(self, job_id, progress_data):
- """Update job progress data"""
- with self.lock:
- job = self.jobs.get(job_id)
- if job:
- job.progress_data = progress_data
-
- def export_queue_to_zip(self, output_path=None):
- """Export the current queue to a zip file containing queue.json and queue_images directory
-
- Args:
- output_path: Path to save the zip file. If None, uses 'queue_export.zip' in the configured output directory.
-
- Returns:
- str: Path to the created zip file
- """
- try:
- # Get the output directory from settings
- settings = Settings()
- output_dir = settings.get("output_dir", "outputs")
- os.makedirs(output_dir, exist_ok=True)
-
- # Use default path if none provided
- if output_path is None:
- output_path = os.path.join(output_dir, "queue_export.zip")
-
- # Make sure queue.json is up to date
- self.save_queue_to_json()
-
- # Create a zip file
- with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
- # Add queue.json to the zip file
- if os.path.exists("queue.json"):
- zipf.write("queue.json")
- print(f"Added queue.json to {output_path}")
- else:
- print("Warning: queue.json not found, creating an empty one")
- with open("queue.json", "w") as f:
- json.dump({}, f)
- zipf.write("queue.json")
-
- # Add queue_images directory to the zip file if it exists
- queue_images_dir = "queue_images"
- if os.path.exists(queue_images_dir) and os.path.isdir(queue_images_dir):
- for root, _, files in os.walk(queue_images_dir):
- for file in files:
- file_path = os.path.join(root, file)
- # Add file to zip with path relative to queue_images_dir
- arcname = os.path.join(os.path.basename(queue_images_dir), file)
- zipf.write(file_path, arcname)
- print(f"Added {file_path} to {output_path}")
- else:
- print(f"Warning: {queue_images_dir} directory not found or empty")
- # Create the directory if it doesn't exist
- os.makedirs(queue_images_dir, exist_ok=True)
-
- # Add input_files directory to the zip file if it exists
- input_files_dir = "input_files"
- if os.path.exists(input_files_dir) and os.path.isdir(input_files_dir):
- for root, _, files in os.walk(input_files_dir):
- for file in files:
- file_path = os.path.join(root, file)
- # Add file to zip with path relative to input_files_dir
- arcname = os.path.join(os.path.basename(input_files_dir), file)
- zipf.write(file_path, arcname)
- print(f"Added {file_path} to {output_path}")
- else:
- print(f"Warning: {input_files_dir} directory not found or empty")
- # Create the directory if it doesn't exist
- os.makedirs(input_files_dir, exist_ok=True)
-
- print(f"Queue exported to {output_path}")
- return output_path
-
- except Exception as e:
- import traceback
- print(f"Error exporting queue to zip: {e}")
- traceback.print_exc()
- return None
-
- def load_queue_from_json(self, file_path=None):
- """Load queue from a JSON file or zip file
-
- Args:
- file_path: Path to the JSON or ZIP file. If None, uses 'queue.json' in the current directory.
-
- Returns:
- int: Number of jobs loaded
- """
- try:
- # Import required modules
- import os
- import json
- from pathlib import PurePath
-
- # Use default path if none provided
- if file_path is None:
- file_path = "queue.json"
-
- # Check if file exists
- if not os.path.exists(file_path):
- print(f"Queue file not found: {file_path}")
- return 0
-
- # Check if it's a zip file
- if file_path.lower().endswith('.zip'):
- return self._load_queue_from_zip(file_path)
-
- # Load the JSON data
- with open(file_path, 'r') as f:
- serialized_jobs = json.load(f)
-
- # Count of jobs loaded
- loaded_count = 0
-
- # Process each job
- with self.lock:
- for job_id, job_data in serialized_jobs.items():
- # Skip if job already exists
- if job_id in self.jobs:
- print(f"Job {job_id} already exists, skipping")
- continue
-
- # Skip completed, failed, or cancelled jobs
- status = job_data.get('status')
- if status in ['completed', 'failed', 'cancelled']:
- print(f"Skipping job {job_id} with status {status}")
- continue
-
- # If the job was running when saved, we'll need to set it as the current job
- was_running = (status == 'running')
-
- # Extract relevant fields to construct params
- params = {
- # Basic parameters
- 'model_type': job_data.get('model_type', 'Original'),
- 'prompt_text': job_data.get('prompt', ''),
- 'n_prompt': job_data.get('negative_prompt', ''),
- 'seed': job_data.get('seed', 0),
- 'steps': job_data.get('steps', 25),
- 'cfg': job_data.get('cfg', 1.0),
- 'gs': job_data.get('gs', 10.0),
- 'rs': job_data.get('rs', 0.0),
- 'latent_type': job_data.get('latent_type', 'Black'),
- 'total_second_length': job_data.get('total_second_length', 6),
- 'blend_sections': job_data.get('blend_sections', 4),
- 'latent_window_size': job_data.get('latent_window_size', 9),
- 'resolutionW': job_data.get('resolutionW', 640),
- 'resolutionH': job_data.get('resolutionH', 640),
- 'use_magcache': job_data.get('use_magcache', False),
- 'magcache_threshold': job_data.get('magcache_threshold', 0.1),
- 'magcache_max_consecutive_skips': job_data.get('magcache_max_consecutive_skips', 2),
- 'magcache_retention_ratio': job_data.get('magcache_retention_ratio', 0.25),
-
- # Initialize image parameters
- 'input_image': None,
- 'end_frame_image': None,
- 'end_frame_strength': job_data.get('end_frame_strength', 1.0),
- 'use_teacache': job_data.get('use_teacache', True),
- 'teacache_num_steps': job_data.get('teacache_num_steps', 25),
- 'teacache_rel_l1_thresh': job_data.get('teacache_rel_l1_thresh', 0.15),
- 'has_input_image': job_data.get('has_input_image', True),
- 'combine_with_source': job_data.get('combine_with_source', False),
- }
-
- # Load input image from disk if saved path exists
- if "saved_input_image_path" in job_data and os.path.exists(job_data["saved_input_image_path"]):
- try:
- input_image_path = job_data["saved_input_image_path"]
- print(f"Loading input image from {input_image_path}")
- input_image = np.array(Image.open(input_image_path))
- params['input_image'] = input_image
- params['input_image_path'] = input_image_path # Store the path for reference
- params['has_input_image'] = True
- except Exception as e:
- print(f"Error loading input image for job {job_id}: {e}")
-
- # Load video from disk if saved path exists
- input_video_val = job_data.get("input_video") # Get value safely
- if isinstance(input_video_val, str): # Check if it's a string path
- if os.path.exists(input_video_val): # Now it's safe to call os.path.exists
- try:
- video_path = input_video_val # Use the validated string path
- print(f"Loading video from {video_path}")
- params['input_image'] = video_path
- params['input_image_path'] = video_path
- params['has_input_image'] = True
- except Exception as e:
- print(f"Error loading video for job {job_id}: {e}")
-
- # Load end frame image from disk if saved path exists
- if "saved_end_frame_image_path" in job_data and os.path.exists(job_data["saved_end_frame_image_path"]):
- try:
- end_frame_image_path = job_data["saved_end_frame_image_path"]
- print(f"Loading end frame image from {end_frame_image_path}")
- end_frame_image = np.array(Image.open(end_frame_image_path))
- params['end_frame_image'] = end_frame_image
- params['end_frame_image_path'] = end_frame_image_path # Store the path for reference
- # Make sure end_frame_strength is set if this is an endframe model
- if params['model_type'] == "Original with Endframe" or params['model_type'] == "F1 with Endframe":
- if 'end_frame_strength' not in params or params['end_frame_strength'] is None:
- params['end_frame_strength'] = job_data.get('end_frame_strength', 1.0)
- print(f"Set end_frame_strength to {params['end_frame_strength']} for job {job_id}")
- except Exception as e:
- print(f"Error loading end frame image for job {job_id}: {e}")
-
- # Add LoRA information if present
- if 'loras' in job_data:
- lora_data = job_data.get('loras', {})
- selected_loras = list(lora_data.keys())
- lora_values = list(lora_data.values())
- params['selected_loras'] = selected_loras
- params['lora_values'] = lora_values
-
- # Ensure the selected LoRAs are also in lora_loaded_names
- # This is critical for metadata_utils.create_metadata to find the LoRAs
- from modules.settings import Settings
- settings = Settings()
- lora_dir = settings.get("lora_dir", "loras")
-
- # Get the current lora_loaded_names from the system
- import os
- from pathlib import PurePath
- current_lora_names = []
- if os.path.isdir(lora_dir):
- for root, _, files in os.walk(lora_dir):
- for file in files:
- if file.endswith('.safetensors') or file.endswith('.pt'):
- lora_relative_path = os.path.relpath(os.path.join(root, file), lora_dir)
- lora_name = str(PurePath(lora_relative_path).with_suffix(''))
- current_lora_names.append(lora_name)
-
- # Combine the selected LoRAs with the current lora_loaded_names
- # This ensures that all selected LoRAs are in lora_loaded_names
- combined_lora_names = list(set(current_lora_names + selected_loras))
- params['lora_loaded_names'] = combined_lora_names
-
- print(f"Loaded LoRA data for job {job_id}: {lora_data}")
- print(f"Combined lora_loaded_names: {combined_lora_names}")
-
- # Get settings for output_dir and metadata_dir
- settings = Settings()
- output_dir = settings.get("output_dir")
- metadata_dir = settings.get("metadata_dir")
- input_files_dir = settings.get("input_files_dir")
-
- # Add these directories to the params
- params['output_dir'] = output_dir
- params['metadata_dir'] = metadata_dir
- params['input_files_dir'] = input_files_dir
-
- # Create a dummy preview image for the job
- dummy_preview = np.zeros((64, 64, 3), dtype=np.uint8)
-
- # Create progress data with the dummy preview
- from diffusers_helper.gradio.progress_bar import make_progress_bar_html
- initial_progress_data = {
- 'preview': dummy_preview,
- 'desc': 'Imported job...',
- 'html': make_progress_bar_html(0, 'Imported job...')
- }
-
- # Create a dummy preview image for the job
- dummy_preview = np.zeros((64, 64, 3), dtype=np.uint8)
-
- # Create progress data with the dummy preview
- from diffusers_helper.gradio.progress_bar import make_progress_bar_html
- initial_progress_data = {
- 'preview': dummy_preview,
- 'desc': 'Imported job...',
- 'html': make_progress_bar_html(0, 'Imported job...')
- }
-
- # Create a new job
- job = Job(
- id=job_id,
- params=params,
- status=JobStatus(job_data.get('status', 'pending')),
- created_at=job_data.get('created_at', time.time()),
- progress_data={},
- stream=AsyncStream(),
- # Mark images as saved if their paths exist in the job data
- input_image_saved="saved_input_image_path" in job_data and os.path.exists(job_data["saved_input_image_path"]),
- end_frame_image_saved="saved_end_frame_image_path" in job_data and os.path.exists(job_data["saved_end_frame_image_path"])
- )
-
- # Add job to the internal jobs dictionary
- self.jobs[job_id] = job
-
- # If a job was marked "running" in the JSON, reset it to "pending"
- # and add it to the processing queue.
- if was_running:
- print(f"Job {job_id} was 'running', resetting to 'pending' and adding to queue.")
- job.status = JobStatus.PENDING
- job.started_at = None # Clear started_at for re-queued job
- job.progress_data = {} # Reset progress
-
- # Add all non-completed/failed/cancelled jobs (now including reset 'running' ones) to the processing queue
- if job.status == JobStatus.PENDING:
- self.queue.put(job_id)
- loaded_count += 1
-
- # Synchronize queue images after loading the queue
- self.synchronize_queue_images()
-
- print(f"Loaded {loaded_count} pending jobs from {file_path}")
- return loaded_count
-
- except Exception as e:
- import traceback
- print(f"Error loading queue from JSON: {e}")
- traceback.print_exc()
- return 0
-
- def _load_queue_from_zip(self, zip_path):
- """Load queue from a zip file
-
- Args:
- zip_path: Path to the zip file
-
- Returns:
- int: Number of jobs loaded
- """
- try:
- # Create a temporary directory to extract the zip file
- temp_dir = "temp_queue_import"
- if os.path.exists(temp_dir):
- shutil.rmtree(temp_dir)
- os.makedirs(temp_dir, exist_ok=True)
-
- # Extract the zip file
- with zipfile.ZipFile(zip_path, 'r') as zipf:
- zipf.extractall(temp_dir)
-
- # Check if queue.json exists in the extracted files
- queue_json_path = os.path.join(temp_dir, "queue.json")
- if not os.path.exists(queue_json_path):
- print(f"queue.json not found in {zip_path}")
- shutil.rmtree(temp_dir)
- return 0
-
- # Define target_queue_images_dir and ensure it exists
- # This needs to be defined regardless of whether queue_images exists in the zip,
- # as it's used later for path updates.
- target_queue_images_dir = "queue_images"
- os.makedirs(target_queue_images_dir, exist_ok=True)
-
- # Check if queue_images directory exists in the extracted files
- queue_images_dir = os.path.join(temp_dir, "queue_images")
- if os.path.exists(queue_images_dir) and os.path.isdir(queue_images_dir):
- # Copy all files from the extracted queue_images directory to the target directory
- for file in os.listdir(queue_images_dir):
- src_path = os.path.join(queue_images_dir, file)
- dst_path = os.path.join(target_queue_images_dir, file)
- if os.path.isfile(src_path):
- shutil.copy2(src_path, dst_path)
- print(f"Copied {src_path} to {dst_path}")
-
- # Check if input_files directory exists in the extracted files
- input_files_dir = os.path.join(temp_dir, "input_files")
- print(f"DEBUG: Checking for input_files directory in zip: {input_files_dir}") # DEBUG
- if os.path.exists(input_files_dir) and os.path.isdir(input_files_dir):
- print(f"DEBUG: Found input_files directory in zip. Contents: {os.listdir(input_files_dir)}") # DEBUG
- # Copy the input_files directory to the current directory
- target_input_files_dir = "input_files"
- os.makedirs(target_input_files_dir, exist_ok=True)
-
- # Copy all files from the extracted input_files directory to the target directory
- for file in os.listdir(input_files_dir):
- print(f"DEBUG: Processing file from zip's input_files: {file}") # DEBUG
- src_path = os.path.join(input_files_dir, file)
- dst_path = os.path.join(target_input_files_dir, file)
- if os.path.isfile(src_path):
- print(f"DEBUG: Attempting to copy video file: {src_path} to {dst_path}") # DEBUG
- shutil.copy2(src_path, dst_path)
- print(f"Copied {src_path} to {dst_path}")
- else: # DEBUG
- print(f"DEBUG: Skipped copy, {src_path} is not a file.") # DEBUG
- else: # DEBUG
- print(f"DEBUG: Directory {input_files_dir} does not exist or is not a directory.") # DEBUG
-
- # Update paths in the queue.json file to reflect the new location of the images
- try:
- with open(queue_json_path, 'r') as f:
- queue_data = json.load(f)
-
- # Update paths for each job
- for job_id, job_data in queue_data.items():
- # Check for files with job_id in the name to identify input and end frame images
- input_image_filename = f"{job_id}_input.png"
- end_frame_image_filename = f"{job_id}_end_frame.png"
-
- # Check if these files exist in the target directory
- input_image_path = os.path.join(target_queue_images_dir, input_image_filename)
- end_frame_image_path = os.path.join(target_queue_images_dir, end_frame_image_filename)
-
- # Update paths in job_data
- if os.path.exists(input_image_path):
- job_data["saved_input_image_path"] = input_image_path
- print(f"Updated input image path for job {job_id}: {input_image_path}")
- elif "saved_input_image_path" in job_data:
- # Fallback to updating the existing path
- job_data["saved_input_image_path"] = os.path.join(target_queue_images_dir, os.path.basename(job_data["saved_input_image_path"]))
- print(f"Updated existing input image path for job {job_id}")
-
- if os.path.exists(end_frame_image_path):
- job_data["saved_end_frame_image_path"] = end_frame_image_path
- print(f"Updated end frame image path for job {job_id}: {end_frame_image_path}")
- elif "saved_end_frame_image_path" in job_data:
- # Fallback to updating the existing path
- job_data["saved_end_frame_image_path"] = os.path.join(target_queue_images_dir, os.path.basename(job_data["saved_end_frame_image_path"]))
- print(f"Updated existing end frame image path for job {job_id}")
-
- # Handle video path update for job_data["input_video"]
- current_input_video = job_data.get("input_video")
- current_input_image_path = job_data.get("input_image_path")
- model_type_for_job = job_data.get("model_type")
- video_extensions = ('.mp4', '.mov', '.avi', '.mkv', '.webm', '.flv', '.gif') # Add more if needed
-
- # Prioritize input_video if it's already a string path
- if isinstance(current_input_video, str):
- job_data["input_video"] = os.path.join("input_files", os.path.basename(current_input_video))
- print(f"Updated video path for job {job_id} from 'input_video': {job_data['input_video']}")
- # If input_video is None, but input_image_path is a video path (for Video/Video F1 models)
- elif current_input_video is None and \
- isinstance(current_input_image_path, str) and \
- model_type_for_job in ("Video", "Video F1") and \
- current_input_image_path.lower().endswith(video_extensions):
-
- video_basename = os.path.basename(current_input_image_path)
- job_data["input_video"] = os.path.join("input_files", video_basename)
- print(f"Updated video path for job {job_id} from 'input_image_path' ('{current_input_image_path}') to '{job_data['input_video']}'")
- elif current_input_video is None:
- # If input_video is None and input_image_path is not a usable video path, keep input_video as None
- print(f"Video path for job {job_id} is None and 'input_image_path' ('{current_input_image_path}') not used for 'input_video'. 'input_video' remains None.")
- # Write the updated queue.json back to the file
- with open(queue_json_path, 'w') as f:
- json.dump(queue_data, f, indent=2)
-
- print(f"Updated image paths in queue.json to reflect new location")
- except Exception as e:
- print(f"Error updating paths in queue.json: {e}")
-
- # Load the queue from the extracted queue.json
- loaded_count = self.load_queue_from_json(queue_json_path)
-
- # Clean up the temporary directory
- shutil.rmtree(temp_dir)
-
- return loaded_count
-
- except Exception as e:
- import traceback
- print(f"Error loading queue from zip: {e}")
- traceback.print_exc()
- # Clean up the temporary directory if it exists
- if os.path.exists(temp_dir):
- shutil.rmtree(temp_dir)
- return 0
-
- def _worker_loop(self):
- """Worker thread that processes jobs from the queue"""
- while True:
- try:
- # Get the next job ID from the queue
- try:
- job_id = self.queue.get(block=True, timeout=1.0)
- except queue_module.Empty:
- self._check_and_process_completed_grids()
- continue
-
- with self.lock:
- job = self.jobs.get(job_id)
- if not job:
- self.queue.task_done()
- continue
-
- # Skip cancelled jobs
- if job.status == JobStatus.CANCELLED:
- self.queue.task_done()
- continue
-
- # If it's a grid job, queue its children and mark it as running
- if job.job_type == JobType.GRID:
- print(f"Processing grid job {job.id}, adding {len(job.child_job_ids)} child jobs to queue.")
- job.status = JobStatus.RUNNING # Mark the grid job as running
- job.started_at = time.time()
- # Add child jobs to the front of the queue
- temp_queue = []
- while not self.queue.empty():
- temp_queue.append(self.queue.get())
- for child_id in reversed(job.child_job_ids): # Add in reverse to maintain order
- self.queue.put(child_id)
- for item in temp_queue:
- self.queue.put(item)
-
- self.queue.task_done()
- continue # Continue to the next iteration to process the first child job
-
- # If we're already processing a job, wait for it to complete
- if self.is_processing:
- # Check if this is the job that's already marked as running
- # This can happen if the job was marked as running but not yet processed
- if job.status == JobStatus.RUNNING and self.current_job and self.current_job.id == job_id:
- print(f"Job {job_id} is already marked as running, processing it now")
- # We'll process this job now
- pass
- else:
- # Put the job back in the queue
- self.queue.put(job_id)
- self.queue.task_done()
- time.sleep(0.1) # Small delay to prevent busy waiting
- continue
-
- # Check if there's a previously running job that was interrupted
- previously_running_job = None
- for j in self.jobs.values():
- if j.status == JobStatus.RUNNING and j.id != job_id:
- previously_running_job = j
- break
-
- # If there's a previously running job, process it first
- if previously_running_job:
- print(f"Found previously running job {previously_running_job.id}, processing it first")
- # Put the current job back in the queue
- self.queue.put(job_id)
- self.queue.task_done()
- # Process the previously running job
- job = previously_running_job
- job_id = previously_running_job.id
-
- # Create a new stream for the resumed job and initialize progress_data
- job.stream = AsyncStream()
- job.progress_data = {}
-
- # Push an initial progress update to the stream
- from diffusers_helper.gradio.progress_bar import make_progress_bar_html
- job.stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Resuming job...'))))
-
- print(f"Starting job {job_id}, current job was {self.current_job.id if self.current_job else 'None'}")
- job.status = JobStatus.RUNNING
- job.started_at = time.time()
- self.current_job = job
- self.is_processing = True
-
- job_completed = False
-
- try:
- if self.worker_function is None:
- raise ValueError("Worker function not set. Call set_worker_function() first.")
-
- # Start the worker function with the job parameters
- from diffusers_helper.thread_utils import async_run
- print(f"Starting worker function for job {job_id}")
-
- # Clean up params for the worker function
- worker_params = job.params.copy()
- if 'end_frame_image_original' in worker_params:
- del worker_params['end_frame_image_original']
- if 'end_frame_strength_original' in worker_params:
- del worker_params['end_frame_strength_original']
-
- async_run(
- self.worker_function,
- **worker_params,
- job_stream=job.stream
- )
- print(f"Worker function started for job {job_id}")
-
- # Process the results from the stream
- output_filename = None
-
- # Track activity time for logging purposes
- last_activity_time = time.time()
-
- while True:
- # Check if job has been cancelled before processing next output
- with self.lock:
- if job.status == JobStatus.CANCELLED:
- print(f"Job {job_id} was cancelled, breaking out of processing loop")
- job_completed = True
- break
-
- # Get current time for activity checks
- current_time = time.time()
-
- # Check for inactivity (no output for a while)
- if current_time - last_activity_time > 60: # 1 minute of inactivity
- print(f"Checking if job {job_id} is still active...")
- # Just a periodic check, don't break yet
-
- try:
- # Try to get data from the queue with a non-blocking approach
- flag, data = job.stream.output_queue.next()
-
- # Update activity time since we got some data
- last_activity_time = time.time()
-
- if flag == 'file':
- output_filename = data
- with self.lock:
- job.result = output_filename
-
- elif flag == 'progress':
- preview, desc, html = data
- with self.lock:
- job.progress_data = {
- 'preview': preview,
- 'desc': desc,
- 'html': html
- }
-
- elif flag == 'end':
- print(f"Received end signal for job {job_id}")
- job_completed = True
- break
-
- except IndexError:
- # Queue is empty, wait a bit and try again
- time.sleep(0.1)
- continue
- except Exception as e:
- print(f"Error processing job output: {e}")
- # Wait a bit before trying again
- time.sleep(0.1)
- continue
- except Exception as e:
- import traceback
- traceback.print_exc()
- print(f"Error processing job {job_id}: {e}")
- with self.lock:
- job.status = JobStatus.FAILED
- job.error = str(e)
- job.completed_at = time.time()
- job_completed = True
-
- finally:
- with self.lock:
- # Make sure we properly clean up the job state
- if job.status == JobStatus.RUNNING:
- if job_completed:
- job.status = JobStatus.COMPLETED
- else:
- # Something went wrong but we didn't mark it as completed
- job.status = JobStatus.FAILED
- job.error = "Job processing was interrupted"
-
- job.completed_at = time.time()
-
- print(f"Finishing job {job_id} with status {job.status}")
- self.is_processing = False
-
- # Check if there's another job in the queue before setting current_job to None
- # This helps prevent UI flashing when a job is cancelled
- next_job_id = None
- try:
- # Peek at the next job without removing it from the queue
- if not self.queue.empty():
- # We can't peek with the standard Queue, so we'll have to get creative
- # Store the queue items temporarily
- temp_queue = []
- while not self.queue.empty():
- item = self.queue.get()
- temp_queue.append(item)
- if next_job_id is None:
- next_job_id = item
-
- # Put everything back
- for item in temp_queue:
- self.queue.put(item)
- except Exception as e:
- print(f"Error checking for next job: {e}")
-
- # After a job completes or is cancelled, always set current_job to None
- self.current_job = None
-
- # The main loop's self.queue.get() will pick up the next available job.
- # No need to explicitly find and start the next job here.
-
- self.queue.task_done()
-
- # Save the queue to JSON after job completion (outside the lock)
- try:
- self.save_queue_to_json()
- except Exception as e:
- print(f"Error saving queue to JSON after job completion: {e}")
-
- except Exception as e:
- import traceback
- traceback.print_exc()
- print(f"Error in worker loop: {e}")
-
- # Make sure we reset processing state if there was an error
- with self.lock:
- self.is_processing = False
- if self.current_job:
- self.current_job.status = JobStatus.FAILED
- self.current_job.error = f"Worker loop error: {str(e)}"
- self.current_job.completed_at = time.time()
- self.current_job = None
-
- time.sleep(0.5) # Prevent tight loop on error
-
- def _check_and_process_completed_grids(self):
- """Check for completed grid jobs and process them."""
- with self.lock:
- # Find all running grid jobs
- running_grid_jobs = [job for job in self.jobs.values() if job.job_type == JobType.GRID and job.status == JobStatus.RUNNING]
-
- for grid_job in running_grid_jobs:
- # Check if all child jobs are completed
- child_jobs = [self.jobs.get(child_id) for child_id in grid_job.child_job_ids]
-
- if not all(child_jobs):
- print(f"Warning: Some child jobs for grid {grid_job.id} not found.")
- continue
-
- all_children_done = all(job.status in [JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED] for job in child_jobs)
-
- if all_children_done:
- print(f"All child jobs for grid {grid_job.id} are done. Assembling grid.")
- # Logic to assemble the grid
- # This is a placeholder for the actual grid assembly logic
- # For now, we'll just mark the grid job as completed.
-
- # Collect results from child jobs
- child_results = [child.result for child in child_jobs if child.status == JobStatus.COMPLETED and child.result]
-
- if not child_results:
- print(f"Grid job {grid_job.id} failed because no child jobs completed successfully.")
- grid_job.status = JobStatus.FAILED
- grid_job.error = "No child jobs completed successfully."
- grid_job.completed_at = time.time()
- continue
-
- # Placeholder for grid assembly.
- # In a real implementation, you would use a tool like FFmpeg or MoviePy to stitch the videos.
- # For this example, we'll just create a text file with the paths of the child videos.
- try:
- output_dir = grid_job.params.get("output_dir", "outputs")
- grid_filename = os.path.join(output_dir, f"grid_{grid_job.id}.txt")
- with open(grid_filename, "w") as f:
- f.write(f"Grid for job: {grid_job.id}\n")
- f.write("Child video paths:\n")
- for result_path in child_results:
- f.write(f"{result_path}\n")
-
- grid_job.result = grid_filename
- grid_job.status = JobStatus.COMPLETED
- print(f"Grid assembly for job {grid_job.id} complete. Result saved to {grid_filename}")
-
- except Exception as e:
- print(f"Error during grid assembly for job {grid_job.id}: {e}")
- grid_job.status = JobStatus.FAILED
- grid_job.error = f"Grid assembly failed: {e}"
-
- grid_job.completed_at = time.time()
- self.save_queue_to_json()
+import threading
+import time
+import uuid
+import json
+import os
+import zipfile
+import shutil
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, Any, Optional, List
+import queue as queue_module # Renamed to avoid conflicts
+import io
+import base64
+from PIL import Image
+import numpy as np
+
+from diffusers_helper.thread_utils import AsyncStream
+from modules.pipelines.metadata_utils import create_metadata
+from modules.settings import Settings
+from diffusers_helper.gradio.progress_bar import make_progress_bar_html
+
+
+# Simple LIFO queue implementation to avoid dependency on queue.LifoQueue
+class SimpleLifoQueue:
+ def __init__(self):
+ self._queue = []
+ self._mutex = threading.Lock()
+ self._not_empty = threading.Condition(self._mutex)
+
+ def put(self, item):
+ with self._mutex:
+ self._queue.append(item)
+ self._not_empty.notify()
+
+ def get(self):
+ with self._not_empty:
+ while not self._queue:
+ self._not_empty.wait()
+ return self._queue.pop()
+
+ def task_done(self):
+ pass # For compatibility with queue.Queue
+
+
+class JobStatus(Enum):
+ PENDING = "pending"
+ RUNNING = "running"
+ COMPLETED = "completed"
+ FAILED = "failed"
+ CANCELLED = "cancelled"
+
+
+class JobType(Enum):
+ SINGLE = "single"
+ GRID = "grid"
+
+
+@dataclass
+class Job:
+ id: str
+ params: Dict[str, Any]
+ status: JobStatus = JobStatus.PENDING
+ job_type: JobType = JobType.SINGLE
+ child_job_ids: List[str] = field(default_factory=list)
+ parent_job_id: Optional[str] = None
+ created_at: float = field(default_factory=time.time)
+ started_at: Optional[float] = None
+ completed_at: Optional[float] = None
+ error: Optional[str] = None
+ result: Optional[str] = None
+ progress_data: Optional[Dict] = None
+ queue_position: Optional[int] = None
+ stream: Optional[Any] = None
+ input_image: Optional[np.ndarray] = None
+ latent_type: Optional[str] = None
+ thumbnail: Optional[str] = None
+ generation_type: Optional[str] = None # Added generation_type
+ input_image_saved: bool = False # Flag to track if input image has been saved
+ end_frame_image_saved: bool = False # Flag to track if end frame image has been saved
+
+ def __post_init__(self):
+ # Store generation type
+ self.generation_type = self.params.get('model_type', 'Original') # Initialize generation_type
+
+ # Store input image or latent type
+ if 'input_image' in self.params and self.params['input_image'] is not None:
+ self.input_image = self.params['input_image']
+ # Create thumbnail
+ if isinstance(self.input_image, np.ndarray):
+ # Handle numpy array (image)
+ img = Image.fromarray(self.input_image)
+ img.thumbnail((100, 100))
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+ elif isinstance(self.input_image, str):
+ # Handle string (video path)
+ try:
+ print(f"Attempting to extract thumbnail from video: {self.input_image}")
+ # Try to extract frames from the video using imageio
+ import imageio
+
+ # Check if the file exists
+ if not os.path.exists(self.input_image):
+ print(f"Video file not found: {self.input_image}")
+ raise FileNotFoundError(f"Video file not found: {self.input_image}")
+
+ # Create outputs directory if it doesn't exist
+ os.makedirs("outputs", exist_ok=True)
+
+ # Try to open the video file
+ try:
+ reader = imageio.get_reader(self.input_image)
+ print(f"Successfully opened video file with imageio")
+ except Exception as e:
+ print(f"Failed to open video with imageio: {e}")
+ raise
+
+ # Get the total number of frames
+ num_frames = None
+ try:
+ # Try to get the number of frames from metadata
+ meta_data = reader.get_meta_data()
+ print(f"Video metadata: {meta_data}")
+ num_frames = meta_data.get('nframes')
+ if num_frames is None or num_frames == float('inf'):
+ print("Number of frames not available in metadata")
+ # If not available, try to count frames
+ if hasattr(reader, 'count_frames'):
+ print("Trying to count frames...")
+ num_frames = reader.count_frames()
+ print(f"Counted {num_frames} frames")
+ except Exception as e:
+ print(f"Error getting frame count: {e}")
+ num_frames = None
+
+ # If we couldn't determine the number of frames, read the last frame by iterating
+ if num_frames is None or num_frames == float('inf'):
+ print("Reading frames by iteration to find the last one")
+ # Read frames until we reach the end
+ frame_count = 0
+ first_frame = None
+ last_frame = None
+ try:
+ for frame in reader:
+ if frame_count == 0:
+ first_frame = frame
+ last_frame = frame
+ frame_count += 1
+ # Print progress every 100 frames
+ if frame_count % 100 == 0:
+ print(f"Read {frame_count} frames...")
+ print(f"Finished reading {frame_count} frames")
+
+ # Save the first frame if available
+ if first_frame is not None:
+ print(f"Found first frame with shape: {first_frame.shape}")
+ # DEBUG IMAGE SAVING REMOVED
+ except Exception as e:
+ print(f"Error reading frames: {e}")
+
+ if last_frame is not None:
+ print(f"Found last frame with shape: {last_frame.shape}")
+
+ # DEBUG IMAGE SAVING REMOVED
+ # Use the last frame for the thumbnail
+ img = Image.fromarray(last_frame)
+ img.thumbnail((100, 100))
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+ print("Successfully created thumbnail from last frame")
+ else:
+ print("No frames were read, using red thumbnail")
+ # Fallback to red thumbnail if no frames were read - more visible for debugging
+ img = Image.new('RGB', (100, 100), (255, 0, 0)) # Red for video
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+ else:
+ # If we know the number of frames, try to get multiple frames for debugging
+ try:
+ # Try to get the first frame
+ first_frame = None
+ try:
+ first_frame = reader.get_data(0)
+ print(f"Got first frame with shape: {first_frame.shape}")
+
+ # DEBUG IMAGE SAVING REMOVED
+ except Exception as e:
+ print(f"Error getting first frame: {e}")
+
+ # Try to get a middle frame
+ middle_frame = None
+ try:
+ middle_frame_idx = int(num_frames / 2)
+ middle_frame = reader.get_data(middle_frame_idx)
+ print(f"Got middle frame (frame {middle_frame_idx}) with shape: {middle_frame.shape}")
+
+ # DEBUG IMAGE SAVING REMOVED
+ except Exception as e:
+ print(f"Error getting middle frame: {e}")
+
+ # Try to get the last frame
+ last_frame = None
+ try:
+ last_frame_idx = int(num_frames) - 1
+ last_frame = reader.get_data(last_frame_idx)
+ print(f"Got last frame (frame {last_frame_idx}) with shape: {last_frame.shape}")
+
+ # DEBUG IMAGE SAVING REMOVED
+ except Exception as e:
+ print(f"Error getting last frame: {e}")
+
+ # If we couldn't get the last frame directly, try to get it by iterating
+ if last_frame is None:
+ print("Trying to get last frame by iterating through all frames")
+ try:
+ for frame in reader:
+ last_frame = frame
+
+ if last_frame is not None:
+ print(f"Got last frame by iteration with shape: {last_frame.shape}")
+
+ # DEBUG IMAGE SAVING REMOVED
+ except Exception as e:
+ print(f"Error getting last frame by iteration: {e}")
+
+ # Use the last frame for the thumbnail if available, otherwise use the middle or first frame
+ frame_for_thumbnail = last_frame if last_frame is not None else (middle_frame if middle_frame is not None else first_frame)
+
+ if frame_for_thumbnail is not None:
+ # Convert to PIL Image and create a thumbnail
+ img = Image.fromarray(frame_for_thumbnail)
+ img.thumbnail((100, 100))
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+ print("Successfully created thumbnail from frame")
+ else:
+ print("No frames were extracted, using blue thumbnail")
+ # Fallback to blue thumbnail if no frames were extracted
+ img = Image.new('RGB', (100, 100), (0, 0, 255)) # Blue for video
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+ except Exception as e:
+ # Fallback to blue thumbnail on error
+ img = Image.new('RGB', (100, 100), (0, 0, 255)) # Blue for video
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+
+ # Close the reader
+ try:
+ reader.close()
+ print("Successfully closed video reader")
+ except Exception as e:
+ print(f"Error closing reader: {e}")
+
+ except Exception as e:
+ print(f"Error extracting thumbnail from video: {e}")
+ import traceback
+ traceback.print_exc()
+ # Fallback to bright green thumbnail on error to make it more visible
+ img = Image.new('RGB', (100, 100), (0, 255, 0)) # Bright green for error
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+ print("Created bright green fallback thumbnail")
+ else:
+ # Handle other types
+ self.thumbnail = None
+ elif 'latent_type' in self.params:
+ self.latent_type = self.params['latent_type']
+ # Create a colored square based on latent type
+ color_map = {
+ "Black": (0, 0, 0),
+ "White": (255, 255, 255),
+ "Noise": (128, 128, 128),
+ "Green Screen": (0, 177, 64)
+ }
+ color = color_map.get(self.latent_type, (0, 0, 0))
+ img = Image.new('RGB', (100, 100), color)
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ self.thumbnail = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
+
+
+class VideoJobQueue:
+ def __init__(self):
+ self.queue = queue_module.Queue() # Using standard Queue instead of LifoQueue
+ self.jobs = {}
+ self.current_job = None
+ self.lock = threading.Lock()
+ self.worker_thread = threading.Thread(target=self._worker_loop, daemon=True)
+ self.worker_thread.start()
+ self.worker_function = None # Will be set from outside
+ self.is_processing = False # Flag to track if we're currently processing a job
+
+ def set_worker_function(self, worker_function):
+ """Set the worker function to use for processing jobs"""
+ self.worker_function = worker_function
+
+ def serialize_job(self, job):
+ """Serialize a job to a JSON-compatible format"""
+ try:
+ # Create a simplified representation of the job
+ serialized = {
+ "id": job.id,
+ "status": job.status.value,
+ "created_at": job.created_at,
+ "started_at": job.started_at,
+ "completed_at": job.completed_at,
+ "error": job.error,
+ "result": job.result,
+ "queue_position": job.queue_position,
+ "generation_type": job.generation_type,
+ }
+
+ # Add simplified params (excluding complex objects)
+ serialized_params = {}
+ for k, v in job.params.items():
+ if k not in ["input_image", "end_frame_image", "stream"]:
+ # Try to include only JSON-serializable values
+ try:
+ # Test if value is JSON serializable
+ json.dumps({k: v})
+ serialized_params[k] = v
+ except (TypeError, OverflowError):
+ # Skip non-serializable values
+ pass
+
+ # Handle LoRA information specifically
+ # Only include selected LoRAs for the generation
+ if "selected_loras" in job.params and job.params["selected_loras"]:
+ selected_loras = job.params["selected_loras"]
+ # Ensure it's a list
+ if not isinstance(selected_loras, list):
+ selected_loras = [selected_loras] if selected_loras is not None else []
+
+ # Get LoRA values if available
+ lora_values = job.params.get("lora_values", [])
+ if not isinstance(lora_values, list):
+ lora_values = [lora_values] if lora_values is not None else []
+
+ # Get loaded LoRA names
+ lora_loaded_names = job.params.get("lora_loaded_names", [])
+ if not isinstance(lora_loaded_names, list):
+ lora_loaded_names = [lora_loaded_names] if lora_loaded_names is not None else []
+
+ # Create LoRA data dictionary
+ lora_data = {}
+ for lora_name in selected_loras:
+ try:
+ # Find the index of the LoRA in loaded names
+ idx = lora_loaded_names.index(lora_name) if lora_loaded_names else -1
+ # Get the weight value
+ weight = lora_values[idx] if lora_values and idx >= 0 and idx < len(lora_values) else 1.0
+ # Handle weight as list
+ if isinstance(weight, list):
+ weight_value = weight[0] if weight and len(weight) > 0 else 1.0
+ else:
+ weight_value = weight
+ # Store as float
+ lora_data[lora_name] = float(weight_value)
+ except (ValueError, IndexError):
+ # Default weight if not found
+ lora_data[lora_name] = 1.0
+ except Exception as e:
+ print(f"Error processing LoRA {lora_name}: {e}")
+ lora_data[lora_name] = 1.0
+
+ # Add to serialized params
+ serialized_params["loras"] = lora_data
+
+ serialized["params"] = serialized_params
+
+ # Don't include the thumbnail as it can be very large and cause issues
+ # if job.thumbnail:
+ # serialized["thumbnail"] = job.thumbnail
+
+ return serialized
+ except Exception as e:
+ print(f"Error serializing job {job.id}: {e}")
+ # Return minimal information that should always be serializable
+ return {
+ "id": job.id,
+ "status": job.status.value,
+ "error": f"Error serializing: {str(e)}"
+ }
+
+ def save_queue_to_json(self):
+ """Save the current queue to queue.json using the central metadata utility"""
+ try:
+ # Make a copy of job IDs to avoid holding the lock while serializing
+ with self.lock:
+ job_ids = list(self.jobs.keys())
+
+ # Create a settings instance
+ settings = Settings()
+
+ # Create a directory to store queue images if it doesn't exist
+ queue_images_dir = "queue_images"
+ os.makedirs(queue_images_dir, exist_ok=True)
+
+ # First, ensure all images are saved
+ for job_id in job_ids:
+ job = self.get_job(job_id)
+ if job:
+ # Save input image to disk if it exists and hasn't been saved yet
+ if 'input_image' in job.params and isinstance(job.params['input_image'], np.ndarray) and not job.input_image_saved:
+ input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
+ try:
+ Image.fromarray(job.params['input_image']).save(input_image_path)
+ print(f"Saved input image for job {job_id} to {input_image_path}")
+ # Mark the image as saved
+ job.input_image_saved = True
+ except Exception as e:
+ print(f"Error saving input image for job {job_id}: {e}")
+
+ # Save end frame image to disk if it exists and hasn't been saved yet
+ if 'end_frame_image' in job.params and isinstance(job.params['end_frame_image'], np.ndarray) and not job.end_frame_image_saved:
+ end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
+ try:
+ Image.fromarray(job.params['end_frame_image']).save(end_frame_image_path)
+ print(f"Saved end frame image for job {job_id} to {end_frame_image_path}")
+ # Mark the end frame image as saved
+ job.end_frame_image_saved = True
+ except Exception as e:
+ print(f"Error saving end frame image for job {job_id}: {e}")
+
+ # Now serialize jobs with the updated image saved flags
+ serialized_jobs = {}
+ for job_id in job_ids:
+ job = self.get_job(job_id)
+ if job:
+ # Try to use metadata_utils.create_metadata if possible
+ try:
+ # Create metadata using the central utility
+ metadata = create_metadata(job.params, job.id, settings.settings)
+
+ # Add job status and other fields not included in metadata
+ metadata.update({
+ "id": job.id,
+ "status": job.status.value,
+ "created_at": job.created_at,
+ "started_at": job.started_at,
+ "completed_at": job.completed_at,
+ "error": job.error,
+ "result": job.result,
+ "queue_position": job.queue_position,
+ })
+
+ # Add image paths to metadata if they've been saved
+ if job.input_image_saved:
+ input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
+ if os.path.exists(input_image_path):
+ metadata["saved_input_image_path"] = input_image_path
+
+ if job.end_frame_image_saved:
+ end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
+ if os.path.exists(end_frame_image_path):
+ metadata["saved_end_frame_image_path"] = end_frame_image_path
+
+ serialized_jobs[job_id] = metadata
+ except Exception as e:
+ print(f"Error using metadata_utils for job {job_id}: {e}")
+ # Fall back to the old serialization method
+ serialized_jobs[job_id] = self.serialize_job(job)
+
+ # Save to file
+ with open("queue.json", "w") as f:
+ json.dump(serialized_jobs, f, indent=2)
+
+ # Clean up images for jobs that no longer exist
+ self.cleanup_orphaned_images(job_ids)
+ self.cleanup_orphaned_videos(job_ids)
+
+ print(f"Saved {len(serialized_jobs)} jobs to queue.json")
+ except Exception as e:
+ print(f"Error saving queue to JSON: {e}")
+
+ def cleanup_orphaned_videos(self, current_job_ids_uuids): # Renamed arg for clarity
+ """
+ Remove video files from input_files_dir for jobs that no longer exist
+ or whose input_image_path does not point to them.
+
+ Args:
+ current_job_ids_uuids: List of job UUIDs currently in self.jobs
+ """
+ try:
+ # Get the input_files_dir from settings to be robust
+ settings = Settings()
+ input_files_dir = settings.get("input_files_dir", "input_files")
+ if not os.path.exists(input_files_dir):
+ return
+
+ # Normalize the managed input_files_dir path once
+ norm_input_files_dir = os.path.normpath(input_files_dir)
+ referenced_video_paths = set()
+
+ with self.lock: # Access self.jobs safely
+ for job_id_uuid in current_job_ids_uuids: # Iterate using the provided UUIDs
+ job = self.jobs.get(job_id_uuid)
+ if not (job and job.params):
+ continue
+
+ # Collect all potential video paths from the job parameters
+ # Check for strings to avoid TypeError
+ paths_to_consider = []
+ p1 = job.params.get("input_image") # Primary path used by worker
+ if isinstance(p1, str):
+ paths_to_consider.append(p1)
+
+ p2 = job.params.get("input_image_path") # Secondary/metadata path
+ if isinstance(p2, str) and p2 not in paths_to_consider:
+ paths_to_consider.append(p2)
+
+ p3 = job.params.get("input_video") # Explicitly set during import
+ if isinstance(p3, str) and p3 not in paths_to_consider:
+ paths_to_consider.append(p3)
+
+ for rel_or_abs_path in paths_to_consider:
+ # Resolve to absolute path. If already absolute, abspath does nothing.
+ # If relative, it's resolved against CWD (current working directory).
+ abs_path = os.path.abspath(rel_or_abs_path)
+ norm_abs_path = os.path.normpath(abs_path)
+ # Check if this path is within the managed input_files_dir
+ if norm_abs_path.startswith(norm_input_files_dir):
+ referenced_video_paths.add(norm_abs_path)
+
+ removed_count = 0
+ for filename in os.listdir(input_files_dir):
+ if filename.endswith(".mp4"): # Only process MP4 files
+ file_path_to_check = os.path.normpath(os.path.join(input_files_dir, filename))
+
+ if file_path_to_check not in referenced_video_paths:
+ try:
+ os.remove(file_path_to_check)
+ removed_count += 1
+ print(f"Removed orphaned video: {filename} (path: {file_path_to_check})")
+ except Exception as e:
+ print(f"Error removing orphaned video {filename}: {e}")
+ if removed_count > 0:
+ print(f"Cleaned up {removed_count} orphaned videos from {input_files_dir}")
+ except Exception as e:
+ print(f"Error cleaning up orphaned videos: {e}")
+ import traceback
+ traceback.print_exc()
+
+ def cleanup_orphaned_images(self, current_job_ids):
+ """
+ Remove image files for jobs that no longer exist in the queue.
+
+ Args:
+ current_job_ids: List of job IDs currently in the queue
+ """
+ try:
+ queue_images_dir = "queue_images"
+ if not os.path.exists(queue_images_dir):
+ return
+
+ # Convert to set for faster lookups
+ current_job_ids = set(current_job_ids)
+
+ # Check all files in the queue_images directory
+ removed_count = 0
+ for filename in os.listdir(queue_images_dir):
+ # Only process PNG files with our naming pattern
+ if filename.endswith(".png") and ("_input.png" in filename or "_end_frame.png" in filename):
+ # Extract job ID from filename
+ parts = filename.split("_")
+ if len(parts) >= 2:
+ job_id = parts[0]
+
+ # If job ID is not in current jobs, remove the file
+ if job_id not in current_job_ids:
+ file_path = os.path.join(queue_images_dir, filename)
+ try:
+ os.remove(file_path)
+ removed_count += 1
+ print(f"Removed orphaned image: {filename}")
+ except Exception as e:
+ print(f"Error removing orphaned image {filename}: {e}")
+
+ if removed_count > 0:
+ print(f"Cleaned up {removed_count} orphaned images")
+ except Exception as e:
+ print(f"Error cleaning up orphaned images: {e}")
+
+
+ def synchronize_queue_images(self):
+ """
+ Synchronize the queue_images directory with the current jobs in the queue.
+ This ensures all necessary images are saved and only images for removed jobs are deleted.
+ """
+ try:
+ queue_images_dir = "queue_images"
+ os.makedirs(queue_images_dir, exist_ok=True)
+
+ # Get all current job IDs
+ with self.lock:
+ current_job_ids = set(self.jobs.keys())
+
+ # Get all image files in the queue_images directory
+ existing_image_files = set()
+ if os.path.exists(queue_images_dir):
+ for filename in os.listdir(queue_images_dir):
+ if filename.endswith(".png") and ("_input.png" in filename or "_end_frame.png" in filename):
+ existing_image_files.add(filename)
+
+ # Extract job IDs from filenames
+ file_job_ids = set()
+ for filename in existing_image_files:
+ # Extract job ID from filename (format: "{job_id}_input.png" or "{job_id}_end_frame.png")
+ parts = filename.split("_")
+ if len(parts) >= 2:
+ job_id = parts[0]
+ file_job_ids.add(job_id)
+
+ # Find job IDs in files that are no longer in the queue
+ removed_job_ids = file_job_ids - current_job_ids
+
+ # Delete images for jobs that have been removed from the queue
+ removed_count = 0
+ for job_id in removed_job_ids:
+ input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
+ end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
+
+ if os.path.exists(input_image_path):
+ try:
+ os.remove(input_image_path)
+ removed_count += 1
+ print(f"Removed image for deleted job: {input_image_path}")
+ except Exception as e:
+ print(f"Error removing image {input_image_path}: {e}")
+
+ if os.path.exists(end_frame_image_path):
+ try:
+ os.remove(end_frame_image_path)
+ removed_count += 1
+ print(f"Removed image for deleted job: {end_frame_image_path}")
+ except Exception as e:
+ print(f"Error removing image {end_frame_image_path}: {e}")
+
+ # Now ensure all current jobs have their images saved
+ saved_count = 0
+ with self.lock:
+ for job_id, job in self.jobs.items():
+ # Only save images for running or completed jobs
+ if job.status in [JobStatus.RUNNING, JobStatus.COMPLETED]:
+ # Save input image if it exists and hasn't been saved yet
+ if 'input_image' in job.params and isinstance(job.params['input_image'], np.ndarray) and not job.input_image_saved:
+ input_image_path = os.path.join(queue_images_dir, f"{job_id}_input.png")
+ try:
+ Image.fromarray(job.params['input_image']).save(input_image_path)
+ job.input_image_saved = True
+ saved_count += 1
+ print(f"Saved input image for job {job_id}")
+ except Exception as e:
+ print(f"Error saving input image for job {job_id}: {e}")
+
+ # Save end frame image if it exists and hasn't been saved yet
+ if 'end_frame_image' in job.params and isinstance(job.params['end_frame_image'], np.ndarray) and not job.end_frame_image_saved:
+ end_frame_image_path = os.path.join(queue_images_dir, f"{job_id}_end_frame.png")
+ try:
+ Image.fromarray(job.params['end_frame_image']).save(end_frame_image_path)
+ job.end_frame_image_saved = True
+ saved_count += 1
+ print(f"Saved end frame image for job {job_id}")
+ except Exception as e:
+ print(f"Error saving end frame image for job {job_id}: {e}")
+
+ # Save the queue to ensure the image paths are properly referenced
+ self.save_queue_to_json()
+
+ if removed_count > 0 or saved_count > 0:
+ print(f"Queue image synchronization: removed {removed_count} images, saved {saved_count} images")
+
+ except Exception as e:
+ print(f"Error synchronizing queue images: {e}")
+
+
+ def add_job(self, params, job_type=JobType.SINGLE, child_job_params_list=None, parent_job_id=None):
+ """Add a job to the queue and return its ID"""
+ job_id = str(uuid.uuid4())
+
+ # For grid jobs, create child jobs first
+ child_job_ids = []
+ if job_type == JobType.GRID and child_job_params_list:
+ with self.lock:
+ for child_params in child_job_params_list:
+ child_job_id = str(uuid.uuid4())
+ child_job_ids.append(child_job_id)
+ child_job = Job(
+ id=child_job_id,
+ params=child_params,
+ status=JobStatus.PENDING,
+ job_type=JobType.SINGLE, # Children are single jobs
+ parent_job_id=job_id,
+ created_at=time.time(),
+ progress_data={},
+ stream=AsyncStream(),
+ input_image_saved=False,
+ end_frame_image_saved=False
+ )
+ self.jobs[child_job_id] = child_job
+ print(f" - Created child job {child_job_id} for grid job {job_id}")
+
+ job = Job(
+ id=job_id,
+ params=params,
+ status=JobStatus.PENDING,
+ job_type=job_type,
+ child_job_ids=child_job_ids,
+ parent_job_id=parent_job_id,
+ created_at=time.time(),
+ progress_data={},
+ stream=AsyncStream(),
+ input_image_saved=False,
+ end_frame_image_saved=False
+ )
+
+ with self.lock:
+ print(f"Adding job {job_id} (type: {job_type.value}) to queue.")
+ self.jobs[job_id] = job
+ self.queue.put(job_id) # Only the parent (or single) job is added to the queue initially
+
+ # Save the queue to JSON after adding a new job (outside the lock)
+ try:
+ self.save_queue_to_json()
+ except Exception as e:
+ print(f"Error saving queue to JSON after adding job: {e}")
+
+ return job_id
+
+ def get_job(self, job_id):
+ """Get job by ID"""
+ with self.lock:
+ return self.jobs.get(job_id)
+
+ def get_all_jobs(self):
+ """Get all jobs"""
+ with self.lock:
+ return list(self.jobs.values())
+
+ def cancel_job(self, job_id):
+ """Cancel a pending job"""
+ with self.lock:
+ job = self.jobs.get(job_id)
+ if not job:
+ return False
+
+ if job.status == JobStatus.PENDING:
+ job.status = JobStatus.CANCELLED
+ job.completed_at = time.time() # Mark completion time
+ result = True
+ elif job.status == JobStatus.RUNNING:
+ # Send cancel signal to the job's stream
+ if hasattr(job, 'stream') and job.stream:
+ job.stream.input_queue.push('end')
+
+ # Mark job as cancelled (this will be confirmed when the worker processes the end signal)
+ job.status = JobStatus.CANCELLED
+ job.completed_at = time.time() # Mark completion time
+
+ # Let the worker loop handle the transition to the next job
+ # This ensures the current job is fully processed before switching
+ # DEBUG PRINT REMOVED
+ result = True
+ else:
+ result = False
+
+ # Save the queue to JSON after cancelling a job (outside the lock)
+ if result:
+ try:
+ self.save_queue_to_json()
+ except Exception as e:
+ print(f"Error saving queue to JSON after cancelling job: {e}")
+
+ return result
+
+ def clear_queue(self):
+ """Cancel all pending jobs in the queue"""
+ cancelled_count = 0
+ try:
+ # First, make a copy of all pending job IDs to avoid modifying the dictionary during iteration
+ with self.lock:
+ # Get all pending job IDs
+ pending_job_ids = [job_id for job_id, job in self.jobs.items()
+ if job.status == JobStatus.PENDING]
+
+ # Cancel each pending job individually
+ for job_id in pending_job_ids:
+ try:
+ with self.lock:
+ job = self.jobs.get(job_id)
+ if job and job.status == JobStatus.PENDING:
+ job.status = JobStatus.CANCELLED
+ job.completed_at = time.time()
+ cancelled_count += 1
+ except Exception as e:
+ print(f"Error cancelling job {job_id}: {e}")
+
+ # Now clear the queue
+ with self.lock:
+ # Clear the queue (this doesn't affect running jobs)
+ queue_items_cleared = 0
+ try:
+ while not self.queue.empty():
+ try:
+ self.queue.get_nowait()
+ self.queue.task_done()
+ queue_items_cleared += 1
+ except queue_module.Empty:
+ break
+ except Exception as e:
+ print(f"Error clearing queue: {e}")
+
+ # Save the updated queue state
+ try:
+ self.save_queue_to_json()
+ except Exception as e:
+ print(f"Error saving queue state: {e}")
+
+ # Synchronize queue images after clearing the queue
+ if cancelled_count > 0:
+ self.synchronize_queue_images()
+
+ print(f"Cleared {cancelled_count} jobs from the queue")
+ return cancelled_count
+ except Exception as e:
+ import traceback
+ print(f"Error in clear_queue: {e}")
+ traceback.print_exc()
+ return 0
+
+ def clear_completed_jobs(self):
+ """Remove cancelled or completed jobs from the queue"""
+ removed_count = 0
+ try:
+ # First, make a copy of all completed/cancelled job IDs to avoid modifying the dictionary during iteration
+ with self.lock:
+ # Get all completed or cancelled job IDs
+ completed_job_ids = [job_id for job_id, job in self.jobs.items()
+ if job.status in [JobStatus.COMPLETED, JobStatus.CANCELLED]]
+
+ # Remove each completed/cancelled job individually
+ for job_id in completed_job_ids:
+ try:
+ with self.lock:
+ if job_id in self.jobs:
+ del self.jobs[job_id]
+ removed_count += 1
+ except Exception as e:
+ print(f"Error removing job {job_id}: {e}")
+
+ # Save the updated queue state
+ try:
+ self.save_queue_to_json()
+ except Exception as e:
+ print(f"Error saving queue state: {e}")
+
+ # Synchronize queue images after removing completed jobs
+ if removed_count > 0:
+ self.synchronize_queue_images()
+
+ print(f"Removed {removed_count} completed/cancelled jobs from the queue")
+ return removed_count
+ except Exception as e:
+ import traceback
+ print(f"Error in clear_completed_jobs: {e}")
+ traceback.print_exc()
+ return 0
+
+ def get_queue_position(self, job_id):
+ """Get position in queue (0 = currently running)"""
+ with self.lock:
+ job = self.jobs.get(job_id)
+ if not job:
+ return None
+
+ if job.status == JobStatus.RUNNING:
+ return 0
+
+ if job.status != JobStatus.PENDING:
+ return None
+
+ # Count pending jobs ahead in queue
+ position = 1 # Start at 1 because 0 means running
+ for j in self.jobs.values():
+ if (j.status == JobStatus.PENDING and
+ j.created_at < job.created_at):
+ position += 1
+ return position
+
+ def update_job_progress(self, job_id, progress_data):
+ """Update job progress data"""
+ with self.lock:
+ job = self.jobs.get(job_id)
+ if job:
+ job.progress_data = progress_data
+
+ def export_queue_to_zip(self, output_path=None):
+ """Export the current queue to a zip file containing queue.json and queue_images directory
+
+ Args:
+ output_path: Path to save the zip file. If None, uses 'queue_export.zip' in the configured output directory.
+
+ Returns:
+ str: Path to the created zip file
+ """
+ try:
+ # Get the output directory from settings
+ settings = Settings()
+ output_dir = settings.get("output_dir", "outputs")
+ os.makedirs(output_dir, exist_ok=True)
+
+ # Use default path if none provided
+ if output_path is None:
+ output_path = os.path.join(output_dir, "queue_export.zip")
+
+ # Make sure queue.json is up to date
+ self.save_queue_to_json()
+
+ # Create a zip file
+ with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+ # Add queue.json to the zip file
+ if os.path.exists("queue.json"):
+ zipf.write("queue.json")
+ print(f"Added queue.json to {output_path}")
+ else:
+ print("Warning: queue.json not found, creating an empty one")
+ with open("queue.json", "w") as f:
+ json.dump({}, f)
+ zipf.write("queue.json")
+
+ # Add queue_images directory to the zip file if it exists
+ queue_images_dir = "queue_images"
+ if os.path.exists(queue_images_dir) and os.path.isdir(queue_images_dir):
+ for root, _, files in os.walk(queue_images_dir):
+ for file in files:
+ file_path = os.path.join(root, file)
+ # Add file to zip with path relative to queue_images_dir
+ arcname = os.path.join(os.path.basename(queue_images_dir), file)
+ zipf.write(file_path, arcname)
+ print(f"Added {file_path} to {output_path}")
+ else:
+ print(f"Warning: {queue_images_dir} directory not found or empty")
+ # Create the directory if it doesn't exist
+ os.makedirs(queue_images_dir, exist_ok=True)
+
+ # Add input_files directory to the zip file if it exists
+ input_files_dir = "input_files"
+ if os.path.exists(input_files_dir) and os.path.isdir(input_files_dir):
+ for root, _, files in os.walk(input_files_dir):
+ for file in files:
+ file_path = os.path.join(root, file)
+ # Add file to zip with path relative to input_files_dir
+ arcname = os.path.join(os.path.basename(input_files_dir), file)
+ zipf.write(file_path, arcname)
+ print(f"Added {file_path} to {output_path}")
+ else:
+ print(f"Warning: {input_files_dir} directory not found or empty")
+ # Create the directory if it doesn't exist
+ os.makedirs(input_files_dir, exist_ok=True)
+
+ print(f"Queue exported to {output_path}")
+ return output_path
+
+ except Exception as e:
+ import traceback
+ print(f"Error exporting queue to zip: {e}")
+ traceback.print_exc()
+ return None
+
+ def load_queue_from_json(self, file_path=None):
+ """Load queue from a JSON file or zip file
+
+ Args:
+ file_path: Path to the JSON or ZIP file. If None, uses 'queue.json' in the current directory.
+
+ Returns:
+ int: Number of jobs loaded
+ """
+ try:
+ # Import required modules
+ import os
+ import json
+ from pathlib import PurePath
+
+ # Use default path if none provided
+ if file_path is None:
+ file_path = "queue.json"
+
+ # Check if file exists
+ if not os.path.exists(file_path):
+ print(f"Queue file not found: {file_path}")
+ return 0
+
+ # Check if it's a zip file
+ if file_path.lower().endswith('.zip'):
+ return self._load_queue_from_zip(file_path)
+
+ # Load the JSON data
+ with open(file_path, 'r') as f:
+ serialized_jobs = json.load(f)
+
+ # Count of jobs loaded
+ loaded_count = 0
+
+ # Process each job
+ with self.lock:
+ for job_id, job_data in serialized_jobs.items():
+ # Skip if job already exists
+ if job_id in self.jobs:
+ print(f"Job {job_id} already exists, skipping")
+ continue
+
+ # Skip completed, failed, or cancelled jobs
+ status = job_data.get('status')
+ if status in ['completed', 'failed', 'cancelled']:
+ print(f"Skipping job {job_id} with status {status}")
+ continue
+
+ # If the job was running when saved, we'll need to set it as the current job
+ was_running = (status == 'running')
+
+ # Extract relevant fields to construct params
+ params = {
+ # Basic parameters
+ 'model_type': job_data.get('model_type', 'Original'),
+ 'prompt_text': job_data.get('prompt', ''),
+ 'n_prompt': job_data.get('negative_prompt', ''),
+ 'seed': job_data.get('seed', 0),
+ 'steps': job_data.get('steps', 25),
+ 'cfg': job_data.get('cfg', 1.0),
+ 'gs': job_data.get('gs', 10.0),
+ 'rs': job_data.get('rs', 0.0),
+ 'latent_type': job_data.get('latent_type', 'Black'),
+ 'total_second_length': job_data.get('total_second_length', 6),
+ 'blend_sections': job_data.get('blend_sections', 4),
+ 'latent_window_size': job_data.get('latent_window_size', 9),
+ 'resolutionW': job_data.get('resolutionW', 640),
+ 'resolutionH': job_data.get('resolutionH', 640),
+ 'use_magcache': job_data.get('use_magcache', False),
+ 'magcache_threshold': job_data.get('magcache_threshold', 0.1),
+ 'magcache_max_consecutive_skips': job_data.get('magcache_max_consecutive_skips', 2),
+ 'magcache_retention_ratio': job_data.get('magcache_retention_ratio', 0.25),
+
+ # Initialize image parameters
+ 'input_image': None,
+ 'end_frame_image': None,
+ 'end_frame_strength': job_data.get('end_frame_strength', 1.0),
+ 'use_teacache': job_data.get('use_teacache', True),
+ 'teacache_num_steps': job_data.get('teacache_num_steps', 25),
+ 'teacache_rel_l1_thresh': job_data.get('teacache_rel_l1_thresh', 0.15),
+ 'has_input_image': job_data.get('has_input_image', True),
+ 'combine_with_source': job_data.get('combine_with_source', False),
+ }
+
+ # Load input image from disk if saved path exists
+ if "saved_input_image_path" in job_data and os.path.exists(job_data["saved_input_image_path"]):
+ try:
+ input_image_path = job_data["saved_input_image_path"]
+ print(f"Loading input image from {input_image_path}")
+ input_image = np.array(Image.open(input_image_path))
+ params['input_image'] = input_image
+ params['input_image_path'] = input_image_path # Store the path for reference
+ params['has_input_image'] = True
+ except Exception as e:
+ print(f"Error loading input image for job {job_id}: {e}")
+
+ # Load video from disk if saved path exists
+ input_video_val = job_data.get("input_video") # Get value safely
+ if isinstance(input_video_val, str): # Check if it's a string path
+ if os.path.exists(input_video_val): # Now it's safe to call os.path.exists
+ try:
+ video_path = input_video_val # Use the validated string path
+ print(f"Loading video from {video_path}")
+ params['input_image'] = video_path
+ params['input_image_path'] = video_path
+ params['has_input_image'] = True
+ except Exception as e:
+ print(f"Error loading video for job {job_id}: {e}")
+
+ # Load end frame image from disk if saved path exists
+ if "saved_end_frame_image_path" in job_data and os.path.exists(job_data["saved_end_frame_image_path"]):
+ try:
+ end_frame_image_path = job_data["saved_end_frame_image_path"]
+ print(f"Loading end frame image from {end_frame_image_path}")
+ end_frame_image = np.array(Image.open(end_frame_image_path))
+ params['end_frame_image'] = end_frame_image
+ params['end_frame_image_path'] = end_frame_image_path # Store the path for reference
+ # Make sure end_frame_strength is set if this is an endframe model
+ if params['model_type'] == "Original with Endframe" or params['model_type'] == "F1 with Endframe":
+ if 'end_frame_strength' not in params or params['end_frame_strength'] is None:
+ params['end_frame_strength'] = job_data.get('end_frame_strength', 1.0)
+ print(f"Set end_frame_strength to {params['end_frame_strength']} for job {job_id}")
+ except Exception as e:
+ print(f"Error loading end frame image for job {job_id}: {e}")
+
+ # Add LoRA information if present
+ if 'loras' in job_data:
+ lora_data = job_data.get('loras', {})
+ selected_loras = list(lora_data.keys())
+ lora_values = list(lora_data.values())
+ params['selected_loras'] = selected_loras
+ params['lora_values'] = lora_values
+
+ # Ensure the selected LoRAs are also in lora_loaded_names
+ # This is critical for metadata_utils.create_metadata to find the LoRAs
+ from modules.settings import Settings
+ settings = Settings()
+ lora_dir = settings.get("lora_dir", "loras")
+
+ # Get the current lora_loaded_names from the system
+ import os
+ from pathlib import PurePath
+ current_lora_names = []
+ if os.path.isdir(lora_dir):
+ for root, _, files in os.walk(lora_dir):
+ for file in files:
+ if file.endswith('.safetensors') or file.endswith('.pt'):
+ lora_relative_path = os.path.relpath(os.path.join(root, file), lora_dir)
+ lora_name = str(PurePath(lora_relative_path).with_suffix(''))
+ current_lora_names.append(lora_name)
+
+ # Combine the selected LoRAs with the current lora_loaded_names
+ # This ensures that all selected LoRAs are in lora_loaded_names
+ combined_lora_names = list(set(current_lora_names + selected_loras))
+ params['lora_loaded_names'] = combined_lora_names
+
+ print(f"Loaded LoRA data for job {job_id}: {lora_data}")
+ print(f"Combined lora_loaded_names: {combined_lora_names}")
+
+ # Get settings for output_dir and metadata_dir
+ settings = Settings()
+ output_dir = settings.get("output_dir")
+ metadata_dir = settings.get("metadata_dir")
+ input_files_dir = settings.get("input_files_dir")
+
+ # Add these directories to the params
+ params['output_dir'] = output_dir
+ params['metadata_dir'] = metadata_dir
+ params['input_files_dir'] = input_files_dir
+
+ # Create a dummy preview image for the job
+ dummy_preview = np.zeros((64, 64, 3), dtype=np.uint8)
+
+ # Create progress data with the dummy preview
+ from diffusers_helper.gradio.progress_bar import make_progress_bar_html
+ initial_progress_data = {
+ 'preview': dummy_preview,
+ 'desc': 'Imported job...',
+ 'html': make_progress_bar_html(0, 'Imported job...')
+ }
+
+ # Create a dummy preview image for the job
+ dummy_preview = np.zeros((64, 64, 3), dtype=np.uint8)
+
+ # Create progress data with the dummy preview
+ from diffusers_helper.gradio.progress_bar import make_progress_bar_html
+ initial_progress_data = {
+ 'preview': dummy_preview,
+ 'desc': 'Imported job...',
+ 'html': make_progress_bar_html(0, 'Imported job...')
+ }
+
+ # Create a new job
+ job = Job(
+ id=job_id,
+ params=params,
+ status=JobStatus(job_data.get('status', 'pending')),
+ created_at=job_data.get('created_at', time.time()),
+ progress_data={},
+ stream=AsyncStream(),
+ # Mark images as saved if their paths exist in the job data
+ input_image_saved="saved_input_image_path" in job_data and os.path.exists(job_data["saved_input_image_path"]),
+ end_frame_image_saved="saved_end_frame_image_path" in job_data and os.path.exists(job_data["saved_end_frame_image_path"])
+ )
+
+ # Add job to the internal jobs dictionary
+ self.jobs[job_id] = job
+
+ # If a job was marked "running" in the JSON, reset it to "pending"
+ # and add it to the processing queue.
+ if was_running:
+ print(f"Job {job_id} was 'running', resetting to 'pending' and adding to queue.")
+ job.status = JobStatus.PENDING
+ job.started_at = None # Clear started_at for re-queued job
+ job.progress_data = {} # Reset progress
+
+ # Add all non-completed/failed/cancelled jobs (now including reset 'running' ones) to the processing queue
+ if job.status == JobStatus.PENDING:
+ self.queue.put(job_id)
+ loaded_count += 1
+
+ # Synchronize queue images after loading the queue
+ self.synchronize_queue_images()
+
+ print(f"Loaded {loaded_count} pending jobs from {file_path}")
+ return loaded_count
+
+ except Exception as e:
+ import traceback
+ print(f"Error loading queue from JSON: {e}")
+ traceback.print_exc()
+ return 0
+
+ def _load_queue_from_zip(self, zip_path):
+ """Load queue from a zip file
+
+ Args:
+ zip_path: Path to the zip file
+
+ Returns:
+ int: Number of jobs loaded
+ """
+ try:
+ # Create a temporary directory to extract the zip file
+ temp_dir = "temp_queue_import"
+ if os.path.exists(temp_dir):
+ shutil.rmtree(temp_dir)
+ os.makedirs(temp_dir, exist_ok=True)
+
+ # Extract the zip file
+ with zipfile.ZipFile(zip_path, 'r') as zipf:
+ zipf.extractall(temp_dir)
+
+ # Check if queue.json exists in the extracted files
+ queue_json_path = os.path.join(temp_dir, "queue.json")
+ if not os.path.exists(queue_json_path):
+ print(f"queue.json not found in {zip_path}")
+ shutil.rmtree(temp_dir)
+ return 0
+
+ # Define target_queue_images_dir and ensure it exists
+ # This needs to be defined regardless of whether queue_images exists in the zip,
+ # as it's used later for path updates.
+ target_queue_images_dir = "queue_images"
+ os.makedirs(target_queue_images_dir, exist_ok=True)
+
+ # Check if queue_images directory exists in the extracted files
+ queue_images_dir = os.path.join(temp_dir, "queue_images")
+ if os.path.exists(queue_images_dir) and os.path.isdir(queue_images_dir):
+ # Copy all files from the extracted queue_images directory to the target directory
+ for file in os.listdir(queue_images_dir):
+ src_path = os.path.join(queue_images_dir, file)
+ dst_path = os.path.join(target_queue_images_dir, file)
+ if os.path.isfile(src_path):
+ shutil.copy2(src_path, dst_path)
+ print(f"Copied {src_path} to {dst_path}")
+
+ # Check if input_files directory exists in the extracted files
+ input_files_dir = os.path.join(temp_dir, "input_files")
+ print(f"DEBUG: Checking for input_files directory in zip: {input_files_dir}") # DEBUG
+ if os.path.exists(input_files_dir) and os.path.isdir(input_files_dir):
+ print(f"DEBUG: Found input_files directory in zip. Contents: {os.listdir(input_files_dir)}") # DEBUG
+ # Copy the input_files directory to the current directory
+ target_input_files_dir = "input_files"
+ os.makedirs(target_input_files_dir, exist_ok=True)
+
+ # Copy all files from the extracted input_files directory to the target directory
+ for file in os.listdir(input_files_dir):
+ print(f"DEBUG: Processing file from zip's input_files: {file}") # DEBUG
+ src_path = os.path.join(input_files_dir, file)
+ dst_path = os.path.join(target_input_files_dir, file)
+ if os.path.isfile(src_path):
+ print(f"DEBUG: Attempting to copy video file: {src_path} to {dst_path}") # DEBUG
+ shutil.copy2(src_path, dst_path)
+ print(f"Copied {src_path} to {dst_path}")
+ else: # DEBUG
+ print(f"DEBUG: Skipped copy, {src_path} is not a file.") # DEBUG
+ else: # DEBUG
+ print(f"DEBUG: Directory {input_files_dir} does not exist or is not a directory.") # DEBUG
+
+ # Update paths in the queue.json file to reflect the new location of the images
+ try:
+ with open(queue_json_path, 'r') as f:
+ queue_data = json.load(f)
+
+ # Update paths for each job
+ for job_id, job_data in queue_data.items():
+ # Check for files with job_id in the name to identify input and end frame images
+ input_image_filename = f"{job_id}_input.png"
+ end_frame_image_filename = f"{job_id}_end_frame.png"
+
+ # Check if these files exist in the target directory
+ input_image_path = os.path.join(target_queue_images_dir, input_image_filename)
+ end_frame_image_path = os.path.join(target_queue_images_dir, end_frame_image_filename)
+
+ # Update paths in job_data
+ if os.path.exists(input_image_path):
+ job_data["saved_input_image_path"] = input_image_path
+ print(f"Updated input image path for job {job_id}: {input_image_path}")
+ elif "saved_input_image_path" in job_data:
+ # Fallback to updating the existing path
+ job_data["saved_input_image_path"] = os.path.join(target_queue_images_dir, os.path.basename(job_data["saved_input_image_path"]))
+ print(f"Updated existing input image path for job {job_id}")
+
+ if os.path.exists(end_frame_image_path):
+ job_data["saved_end_frame_image_path"] = end_frame_image_path
+ print(f"Updated end frame image path for job {job_id}: {end_frame_image_path}")
+ elif "saved_end_frame_image_path" in job_data:
+ # Fallback to updating the existing path
+ job_data["saved_end_frame_image_path"] = os.path.join(target_queue_images_dir, os.path.basename(job_data["saved_end_frame_image_path"]))
+ print(f"Updated existing end frame image path for job {job_id}")
+
+ # Handle video path update for job_data["input_video"]
+ current_input_video = job_data.get("input_video")
+ current_input_image_path = job_data.get("input_image_path")
+ model_type_for_job = job_data.get("model_type")
+ video_extensions = ('.mp4', '.mov', '.avi', '.mkv', '.webm', '.flv', '.gif') # Add more if needed
+
+ # Prioritize input_video if it's already a string path
+ if isinstance(current_input_video, str):
+ job_data["input_video"] = os.path.join("input_files", os.path.basename(current_input_video))
+ print(f"Updated video path for job {job_id} from 'input_video': {job_data['input_video']}")
+ # If input_video is None, but input_image_path is a video path (for Video/Video F1 models)
+ elif current_input_video is None and \
+ isinstance(current_input_image_path, str) and \
+ model_type_for_job in ("Video", "Video F1") and \
+ current_input_image_path.lower().endswith(video_extensions):
+
+ video_basename = os.path.basename(current_input_image_path)
+ job_data["input_video"] = os.path.join("input_files", video_basename)
+ print(f"Updated video path for job {job_id} from 'input_image_path' ('{current_input_image_path}') to '{job_data['input_video']}'")
+ elif current_input_video is None:
+ # If input_video is None and input_image_path is not a usable video path, keep input_video as None
+ print(f"Video path for job {job_id} is None and 'input_image_path' ('{current_input_image_path}') not used for 'input_video'. 'input_video' remains None.")
+ # Write the updated queue.json back to the file
+ with open(queue_json_path, 'w') as f:
+ json.dump(queue_data, f, indent=2)
+
+ print(f"Updated image paths in queue.json to reflect new location")
+ except Exception as e:
+ print(f"Error updating paths in queue.json: {e}")
+
+ # Load the queue from the extracted queue.json
+ loaded_count = self.load_queue_from_json(queue_json_path)
+
+ # Clean up the temporary directory
+ shutil.rmtree(temp_dir)
+
+ return loaded_count
+
+ except Exception as e:
+ import traceback
+ print(f"Error loading queue from zip: {e}")
+ traceback.print_exc()
+ # Clean up the temporary directory if it exists
+ if os.path.exists(temp_dir):
+ shutil.rmtree(temp_dir)
+ return 0
+
+ def _worker_loop(self):
+ """Worker thread that processes jobs from the queue"""
+ while True:
+ try:
+ # Get the next job ID from the queue
+ try:
+ job_id = self.queue.get(block=True, timeout=1.0)
+ except queue_module.Empty:
+ self._check_and_process_completed_grids()
+ continue
+
+ with self.lock:
+ job = self.jobs.get(job_id)
+ if not job:
+ self.queue.task_done()
+ continue
+
+ # Skip cancelled jobs
+ if job.status == JobStatus.CANCELLED:
+ self.queue.task_done()
+ continue
+
+ # If it's a grid job, queue its children and mark it as running
+ if job.job_type == JobType.GRID:
+ print(f"Processing grid job {job.id}, adding {len(job.child_job_ids)} child jobs to queue.")
+ job.status = JobStatus.RUNNING # Mark the grid job as running
+ job.started_at = time.time()
+ # Add child jobs to the front of the queue
+ temp_queue = []
+ while not self.queue.empty():
+ temp_queue.append(self.queue.get())
+ for child_id in reversed(job.child_job_ids): # Add in reverse to maintain order
+ self.queue.put(child_id)
+ for item in temp_queue:
+ self.queue.put(item)
+
+ self.queue.task_done()
+ continue # Continue to the next iteration to process the first child job
+
+ # If we're already processing a job, wait for it to complete
+ if self.is_processing:
+ # Check if this is the job that's already marked as running
+ # This can happen if the job was marked as running but not yet processed
+ if job.status == JobStatus.RUNNING and self.current_job and self.current_job.id == job_id:
+ print(f"Job {job_id} is already marked as running, processing it now")
+ # We'll process this job now
+ pass
+ else:
+ # Put the job back in the queue
+ self.queue.put(job_id)
+ self.queue.task_done()
+ time.sleep(0.1) # Small delay to prevent busy waiting
+ continue
+
+ # Check if there's a previously running job that was interrupted
+ previously_running_job = None
+ for j in self.jobs.values():
+ if j.status == JobStatus.RUNNING and j.id != job_id:
+ previously_running_job = j
+ break
+
+ # If there's a previously running job, process it first
+ if previously_running_job:
+ print(f"Found previously running job {previously_running_job.id}, processing it first")
+ # Put the current job back in the queue
+ self.queue.put(job_id)
+ self.queue.task_done()
+ # Process the previously running job
+ job = previously_running_job
+ job_id = previously_running_job.id
+
+ # Create a new stream for the resumed job and initialize progress_data
+ job.stream = AsyncStream()
+ job.progress_data = {}
+
+ # Push an initial progress update to the stream
+ from diffusers_helper.gradio.progress_bar import make_progress_bar_html
+ job.stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Resuming job...'))))
+
+ print(f"Starting job {job_id}, current job was {self.current_job.id if self.current_job else 'None'}")
+ job.status = JobStatus.RUNNING
+ job.started_at = time.time()
+ self.current_job = job
+ self.is_processing = True
+
+ job_completed = False
+
+ try:
+ if self.worker_function is None:
+ raise ValueError("Worker function not set. Call set_worker_function() first.")
+
+ # Start the worker function with the job parameters
+ from diffusers_helper.thread_utils import async_run
+ print(f"Starting worker function for job {job_id}")
+
+ # Clean up params for the worker function
+ worker_params = job.params.copy()
+ if 'end_frame_image_original' in worker_params:
+ del worker_params['end_frame_image_original']
+ if 'end_frame_strength_original' in worker_params:
+ del worker_params['end_frame_strength_original']
+
+ async_run(
+ self.worker_function,
+ **worker_params,
+ job_stream=job.stream
+ )
+ print(f"Worker function started for job {job_id}")
+
+ # Process the results from the stream
+ output_filename = None
+
+ # Track activity time for logging purposes
+ last_activity_time = time.time()
+
+ while True:
+ # Check if job has been cancelled before processing next output
+ with self.lock:
+ if job.status == JobStatus.CANCELLED:
+ print(f"Job {job_id} was cancelled, breaking out of processing loop")
+ job_completed = True
+ break
+
+ # Get current time for activity checks
+ current_time = time.time()
+
+ # Check for inactivity (no output for a while)
+ if current_time - last_activity_time > 60: # 1 minute of inactivity
+ print(f"Checking if job {job_id} is still active...")
+ # Just a periodic check, don't break yet
+
+ try:
+ # Try to get data from the queue with a non-blocking approach
+ flag, data = job.stream.output_queue.next()
+
+ # Update activity time since we got some data
+ last_activity_time = time.time()
+
+ if flag == 'file':
+ output_filename = data
+ with self.lock:
+ job.result = output_filename
+
+ elif flag == 'progress':
+ preview, desc, html = data
+ with self.lock:
+ job.progress_data = {
+ 'preview': preview,
+ 'desc': desc,
+ 'html': html
+ }
+
+ elif flag == 'end':
+ print(f"Received end signal for job {job_id}")
+ job_completed = True
+ break
+
+ except IndexError:
+ # Queue is empty, wait a bit and try again
+ time.sleep(0.1)
+ continue
+ except Exception as e:
+ print(f"Error processing job output: {e}")
+ # Wait a bit before trying again
+ time.sleep(0.1)
+ continue
+ except Exception as e:
+ import traceback
+ traceback.print_exc()
+ print(f"Error processing job {job_id}: {e}")
+ with self.lock:
+ job.status = JobStatus.FAILED
+ job.error = str(e)
+ job.completed_at = time.time()
+ job_completed = True
+
+ finally:
+ with self.lock:
+ # Make sure we properly clean up the job state
+ if job.status == JobStatus.RUNNING:
+ if job_completed:
+ job.status = JobStatus.COMPLETED
+ else:
+ # Something went wrong but we didn't mark it as completed
+ job.status = JobStatus.FAILED
+ job.error = "Job processing was interrupted"
+
+ job.completed_at = time.time()
+
+ print(f"Finishing job {job_id} with status {job.status}")
+ self.is_processing = False
+
+ # Check if there's another job in the queue before setting current_job to None
+ # This helps prevent UI flashing when a job is cancelled
+ next_job_id = None
+ try:
+ # Peek at the next job without removing it from the queue
+ if not self.queue.empty():
+ # We can't peek with the standard Queue, so we'll have to get creative
+ # Store the queue items temporarily
+ temp_queue = []
+ while not self.queue.empty():
+ item = self.queue.get()
+ temp_queue.append(item)
+ if next_job_id is None:
+ next_job_id = item
+
+ # Put everything back
+ for item in temp_queue:
+ self.queue.put(item)
+ except Exception as e:
+ print(f"Error checking for next job: {e}")
+
+ # After a job completes or is cancelled, always set current_job to None
+ self.current_job = None
+
+ # The main loop's self.queue.get() will pick up the next available job.
+ # No need to explicitly find and start the next job here.
+
+ self.queue.task_done()
+
+ # Save the queue to JSON after job completion (outside the lock)
+ try:
+ self.save_queue_to_json()
+ except Exception as e:
+ print(f"Error saving queue to JSON after job completion: {e}")
+
+ except Exception as e:
+ import traceback
+ traceback.print_exc()
+ print(f"Error in worker loop: {e}")
+
+ # Make sure we reset processing state if there was an error
+ with self.lock:
+ self.is_processing = False
+ if self.current_job:
+ self.current_job.status = JobStatus.FAILED
+ self.current_job.error = f"Worker loop error: {str(e)}"
+ self.current_job.completed_at = time.time()
+ self.current_job = None
+
+ time.sleep(0.5) # Prevent tight loop on error
+
+ def _check_and_process_completed_grids(self):
+ """Check for completed grid jobs and process them."""
+ with self.lock:
+ # Find all running grid jobs
+ running_grid_jobs = [job for job in self.jobs.values() if job.job_type == JobType.GRID and job.status == JobStatus.RUNNING]
+
+ for grid_job in running_grid_jobs:
+ # Check if all child jobs are completed
+ child_jobs = [self.jobs.get(child_id) for child_id in grid_job.child_job_ids]
+
+ if not all(child_jobs):
+ print(f"Warning: Some child jobs for grid {grid_job.id} not found.")
+ continue
+
+ all_children_done = all(job.status in [JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED] for job in child_jobs)
+
+ if all_children_done:
+ print(f"All child jobs for grid {grid_job.id} are done. Assembling grid.")
+ # Logic to assemble the grid
+ # This is a placeholder for the actual grid assembly logic
+ # For now, we'll just mark the grid job as completed.
+
+ # Collect results from child jobs
+ child_results = [child.result for child in child_jobs if child.status == JobStatus.COMPLETED and child.result]
+
+ if not child_results:
+ print(f"Grid job {grid_job.id} failed because no child jobs completed successfully.")
+ grid_job.status = JobStatus.FAILED
+ grid_job.error = "No child jobs completed successfully."
+ grid_job.completed_at = time.time()
+ continue
+
+ # Placeholder for grid assembly.
+ # In a real implementation, you would use a tool like FFmpeg or MoviePy to stitch the videos.
+ # For this example, we'll just create a text file with the paths of the child videos.
+ try:
+ output_dir = grid_job.params.get("output_dir", "outputs")
+ grid_filename = os.path.join(output_dir, f"grid_{grid_job.id}.txt")
+ with open(grid_filename, "w") as f:
+ f.write(f"Grid for job: {grid_job.id}\n")
+ f.write("Child video paths:\n")
+ for result_path in child_results:
+ f.write(f"{result_path}\n")
+
+ grid_job.result = grid_filename
+ grid_job.status = JobStatus.COMPLETED
+ print(f"Grid assembly for job {grid_job.id} complete. Result saved to {grid_filename}")
+
+ except Exception as e:
+ print(f"Error during grid assembly for job {grid_job.id}: {e}")
+ grid_job.status = JobStatus.FAILED
+ grid_job.error = f"Grid assembly failed: {e}"
+
+ grid_job.completed_at = time.time()
+ self.save_queue_to_json()
diff --git a/modules/xy_plot_ui.py b/modules/xy_plot_ui.py
index 068d59db..542520d9 100644
--- a/modules/xy_plot_ui.py
+++ b/modules/xy_plot_ui.py
@@ -262,11 +262,10 @@ def create_xy_plot_ui(lora_names, default_prompt, DUMMY_LORA_NAME):
xy_plot_end_frame_image_original = gr.Image(
sources='upload',
type="numpy",
- label="End Frame (Optional)",
- height=420,
+ label="End Frame (Optional)",
+ height=420,
elem_classes="contain-image",
image_mode="RGB",
- show_download_button=False,
show_label=True,
container=True
)
diff --git a/run.bat b/run.bat
index e4c2b5e8..2a5b0bff 100644
--- a/run.bat
+++ b/run.bat
@@ -1,22 +1,22 @@
-@echo off
-echo Starting FramePack-Studio...
-
-REM Check if Python is installed (basic check)
-where python >nul 2>&1
-if %errorlevel% neq 0 (
- echo Error: Python is not installed or not in your PATH. Cannot run studio.py.
- goto end
-)
-
-if exist "%cd%/venv/Scripts/python.exe" (
-
-"%cd%/venv/Scripts/python.exe" studio.py
-
-) else (
-
-echo Error: Virtual Environment for Python not found. Did you install correctly?
-goto end
-
-)
-
+@echo off
+echo Starting FramePack-Studio...
+
+REM Check if Python is installed (basic check)
+where python >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Error: Python is not installed or not in your PATH. Cannot run studio.py.
+ goto end
+)
+
+if exist "%cd%/venv/Scripts/python.exe" (
+
+"%cd%/venv/Scripts/python.exe" studio.py
+
+) else (
+
+echo Error: Virtual Environment for Python not found. Did you install correctly?
+goto end
+
+)
+
:end
\ No newline at end of file
diff --git a/start.bat b/start.bat
new file mode 100644
index 00000000..1ff4901b
--- /dev/null
+++ b/start.bat
@@ -0,0 +1,38 @@
+@echo off
+echo ============================================
+echo FramePack-Studio Launcher
+echo ============================================
+echo.
+
+REM Check if venv exists
+if not exist "%cd%\venv\Scripts\python.exe" (
+ echo Error: Virtual environment not found!
+ echo Please run install_40xx_50xx.bat first to set up the environment.
+ echo.
+ pause
+ exit /b 1
+)
+
+REM Check if studio.py exists
+if not exist "%cd%\studio.py" (
+ echo Error: studio.py not found in current directory.
+ echo Make sure you're running this script from the FramePack-Studio directory.
+ echo.
+ pause
+ exit /b 1
+)
+
+echo Starting FramePack-Studio...
+echo Using virtual environment at: %cd%\venv
+echo.
+
+REM Activate venv and run studio.py
+call "%cd%\venv\Scripts\activate.bat"
+python studio.py
+
+REM Deactivate when done (this line is reached when studio.py exits)
+deactivate
+
+echo.
+echo FramePack-Studio closed.
+pause
diff --git a/studio.py b/studio.py
index 0289afff..c27afc11 100644
--- a/studio.py
+++ b/studio.py
@@ -134,7 +134,9 @@ def verify_lora_state(transformer, label=""):
# Load models
text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
+text_encoder.config.output_hidden_states = True # Ensure hidden states are enabled
text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
+text_encoder_2.config.output_hidden_states = True # Ensure hidden states are enabled
tokenizer = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer')
tokenizer_2 = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='tokenizer_2')
vae = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='vae', torch_dtype=torch.float16).cpu()
@@ -473,7 +475,7 @@ def process(
queue_status = update_queue_status()
# Return immediately after adding to queue
# Return separate updates for start_button and end_button to prevent cross-contamination
- return None, job_id, None, '', f'Job added to queue. Job ID: {job_id}', gr.update(value="🚀 Add to Queue", interactive=True), gr.update(value="❌ Cancel Current Job", interactive=True)
+ return None, job_id, None, '', f'Job added to queue. Job ID: {job_id}', gr.update(value="[+] Add to Queue", interactive=True), gr.update(value="[X] Cancel Current Job", interactive=True)
@@ -565,7 +567,7 @@ def get_preview_updates(preview_value):
force_update = True
# Yield a temporary update to show we're switching jobs
right_preview, top_preview = get_preview_updates(None)
- yield last_video, right_preview, top_preview, '', 'Switching to current job...', gr.update(interactive=True), gr.update(value="❌ Cancel Current Job", visible=True)
+ yield last_video, right_preview, top_preview, '', 'Switching to current job...', gr.update(interactive=True), gr.update(value="[X] Cancel Current Job", visible=True)
continue
# Check if we're waiting for a job transition
@@ -583,7 +585,7 @@ def get_preview_updates(preview_value):
job_id = current_job.id
force_update = True
right_preview, top_preview = get_preview_updates(None)
- yield last_video, right_preview, top_preview, '', 'Switching to current job...', gr.update(interactive=True), gr.update(value="❌ Cancel Current Job", visible=True)
+ yield last_video, right_preview, top_preview, '', 'Switching to current job...', gr.update(interactive=True), gr.update(value="[X] Cancel Current Job", visible=True)
continue
else:
# If still waiting, sleep briefly and continue
@@ -594,7 +596,7 @@ def get_preview_updates(preview_value):
if not job:
# Correctly yield 7 items for the startup/no-job case
# This ensures the status text goes to the right component and the buttons are set correctly.
- yield None, None, None, 'No job ID provided', '', gr.update(value="🚀 Add to Queue", interactive=True, visible=True), gr.update(interactive=False, visible=False)
+ yield None, None, None, 'No job ID provided', '', gr.update(value="[+] Add to Queue", interactive=True, visible=True), gr.update(interactive=False, visible=False)
return
# If a new video file is available, yield it immediately
@@ -616,7 +618,7 @@ def get_preview_updates(preview_value):
if last_job_status != JobStatus.RUNNING:
# Check if the button text is already "Cancelling..." - if so, don't change it
# This prevents the button from changing back to "Cancel Current Job" during cancellation
- button_update = gr.update(interactive=True, value="❌ Cancel Current Job", visible=True)
+ button_update = gr.update(interactive=True, value="[X] Cancel Current Job", visible=True)
else:
# Keep current text and state - important to not override "Cancelling..." text
button_update = gr.update(interactive=True, visible=True)
@@ -658,19 +660,19 @@ def get_preview_updates(preview_value):
elif job.status == JobStatus.COMPLETED:
# Show the final video and reset the button text
right_preview, top_preview = get_preview_updates(last_preview)
- yield job.result, right_preview, top_preview, 'Completed', make_progress_bar_html(100, 'Completed'), gr.update(value="🚀 Add to Queue"), gr.update(interactive=True, value="❌ Cancel Current Job", visible=False)
+ yield job.result, right_preview, top_preview, 'Completed', make_progress_bar_html(100, 'Completed'), gr.update(value="[+] Add to Queue"), gr.update(interactive=True, value="[X] Cancel Current Job", visible=False)
break
elif job.status == JobStatus.FAILED:
# Show error and reset the button text
right_preview, top_preview = get_preview_updates(last_preview)
- yield job.result, right_preview, top_preview, f'Error: {job.error}', make_progress_bar_html(0, 'Failed'), gr.update(value="🚀 Add to Queue"), gr.update(interactive=True, value="❌ Cancel Current Job", visible=False)
+ yield job.result, right_preview, top_preview, f'Error: {job.error}', make_progress_bar_html(0, 'Failed'), gr.update(value="[+] Add to Queue"), gr.update(interactive=True, value="[X] Cancel Current Job", visible=False)
break
elif job.status == JobStatus.CANCELLED:
# Show cancelled message and reset the button text
right_preview, top_preview = get_preview_updates(last_preview)
- yield job.result, right_preview, top_preview, 'Job cancelled', make_progress_bar_html(0, 'Cancelled'), gr.update(interactive=True), gr.update(interactive=True, value="❌ Cancel Current Job", visible=False)
+ yield job.result, right_preview, top_preview, 'Job cancelled', make_progress_bar_html(0, 'Cancelled'), gr.update(interactive=True), gr.update(interactive=True, value="[X] Cancel Current Job", visible=False)
break
# Update last_job_status for the next iteration
diff --git a/update.bat b/update.bat
index 4c37cbf7..740c6efa 100644
--- a/update.bat
+++ b/update.bat
@@ -1,65 +1,65 @@
-@echo off
-echo FramePack-Studio Update Script
-
-REM Check if Git is installed (basic check)
-where git >nul 2>&1
-if %errorlevel% neq 0 (
- echo Error: Git is not installed or not in your PATH. Unable to update.
- goto end
-)
-
-REM Check if Python is installed (basic check)
-where python >nul 2>&1
-if %errorlevel% neq 0 (
- echo Error: Python is not installed or not in your PATH. Unable to update dependencies.
- REM Continue with Git pull, but warn about dependencies
- echo Warning: Python is not available, skipping dependency update.
- goto git_pull
-)
-
-
-:git_pull
-echo Pulling latest changes from Git...
-git pull
-
-REM Check if git pull was successful
-if %errorlevel% neq 0 (
- echo Error: Failed to pull latest changes from Git. Please resolve any conflicts manually.
- goto end
-)
-
-echo Git pull successful.
-
-REM Attempt to update dependencies if Virtual Environment is available
-if exist "%cd%/venv/Scripts/python.exe" (
-
-if %errorlevel% equ 0 (
- echo Updating dependencies using pip...
- REM This assumes there's a requirements.txt file in the root
- REM Using --upgrade to update existing packages
- "%cd%/venv/Scripts/python.exe" -m pip install --upgrade -r requirements.txt
-
- REM Check if pip update was successful
- if %errorlevel% neq 0 (
- echo Warning: Failed to update dependencies. You may need to update them manually.
- ) else (
- echo Dependency update successful.
- )
-) else (
- echo Skipping dependency update as Python is not available.
-)
-
-) else (
-
-echo Error: Virtual Environment for Python not found. Did you install correctly?
-goto end
-
-)
-
-
-
-echo Update complete.
-
-:end
-echo Exiting update script.
+@echo off
+echo FramePack-Studio Update Script
+
+REM Check if Git is installed (basic check)
+where git >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Error: Git is not installed or not in your PATH. Unable to update.
+ goto end
+)
+
+REM Check if Python is installed (basic check)
+where python >nul 2>&1
+if %errorlevel% neq 0 (
+ echo Error: Python is not installed or not in your PATH. Unable to update dependencies.
+ REM Continue with Git pull, but warn about dependencies
+ echo Warning: Python is not available, skipping dependency update.
+ goto git_pull
+)
+
+
+:git_pull
+echo Pulling latest changes from Git...
+git pull
+
+REM Check if git pull was successful
+if %errorlevel% neq 0 (
+ echo Error: Failed to pull latest changes from Git. Please resolve any conflicts manually.
+ goto end
+)
+
+echo Git pull successful.
+
+REM Attempt to update dependencies if Virtual Environment is available
+if exist "%cd%/venv/Scripts/python.exe" (
+
+if %errorlevel% equ 0 (
+ echo Updating dependencies using pip...
+ REM This assumes there's a requirements.txt file in the root
+ REM Using --upgrade to update existing packages
+ "%cd%/venv/Scripts/python.exe" -m pip install --upgrade -r requirements.txt
+
+ REM Check if pip update was successful
+ if %errorlevel% neq 0 (
+ echo Warning: Failed to update dependencies. You may need to update them manually.
+ ) else (
+ echo Dependency update successful.
+ )
+) else (
+ echo Skipping dependency update as Python is not available.
+)
+
+) else (
+
+echo Error: Virtual Environment for Python not found. Did you install correctly?
+goto end
+
+)
+
+
+
+echo Update complete.
+
+:end
+echo Exiting update script.
pause
\ No newline at end of file