From b223d55cc29a9f882644212df789a884e04c43f6 Mon Sep 17 00:00:00 2001
From: Jeffreytsai1004 <jeffreytsai1004@gmail.com>
Date: Sun, 16 Feb 2025 17:09:55 +0800
Subject: [PATCH] Update

---
 download_resource.bat | 20 ++++++++++----------
 install.bat           | 11 ++++++-----
 run.bat               | 18 ++++++++++++++++++
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/download_resource.bat b/download_resource.bat
index 8d6bab0..764e3a6 100644
--- a/download_resource.bat
+++ b/download_resource.bat
@@ -97,12 +97,10 @@ echo.
 REM 创建resource目录结构
 mkdir resource 2>nul
 cd resource
-mkdir NeuralHaircut\pretrained_models\diffusion_prior 2>nul
-mkdir NeuralHaircut\pretrained_models\strand_prior 2>nul
-mkdir NeuralHaircut\PIXIE 2>nul
+mkdir NeuralHaircut 2>nul
 mkdir Matte-Anything 2>nul
-mkdir openpose\models\pose\coco 2>nul
-mkdir hyperIQA\pretrained 2>nul
+mkdir openpose 2>nul
+mkdir hyperIQA 2>nul
 
 REM 验证已有文件
 echo 正在检查已有资源文件...
@@ -112,7 +110,7 @@ SET /A EXISTING_FILES=0
 SET MISSING_LIST=
 
 REM Neural Haircut Diffusion Prior 检查
-IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
+IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\diffusion_prior\dif_ckpt.pt" (
     echo [√] Neural Haircut Diffusion Prior已存在
     SET /A EXISTING_FILES+=1
 ) ELSE (
@@ -122,7 +120,7 @@ IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckp
 )
 
 REM Neural Haircut Strand Prior 检查
-IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\strand_ckpt.pt" (
+IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\strand_prior\strand_ckpt.pt" (
     echo [√] Neural Haircut Strand Prior已存在
     SET /A EXISTING_FILES+=1
 ) ELSE (
@@ -233,7 +231,9 @@ tar -xvzf models.tar.gz
 del models.tar.gz
 
 REM 下载 hyperIQA 模型
-cd ..\hyperIQA\pretrained
+cd ..\hyperIQA
+mkdir pretrained 2>nul
+cd pretrained
 python -m gdown "%HYPERIQA_MODEL_URL%"
 
 cd ..\..\..
@@ -256,14 +256,14 @@ echo 正在验证下载文件...
 cd %PROJECT_DIR%
 SET VERIFY_FAILED=0
 
-IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
+IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\diffusion_prior\dif_ckpt.pt" (
     echo [×] 缺少Neural Haircut Diffusion Prior
     SET /A VERIFY_FAILED+=1
 ) ELSE (
     echo [√] Neural Haircut Diffusion Prior验证成功
 )
 
-IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\strand_ckpt.pt" (
+IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\strand_prior\strand_ckpt.pt" (
     echo [×] 缺少Neural Haircut Strand Prior
     SET /A VERIFY_FAILED+=1
 ) ELSE (
diff --git a/install.bat b/install.bat
index 704fddd..814e1b6 100644
--- a/install.bat
+++ b/install.bat
@@ -164,11 +164,10 @@ IF NOT EXIST "%PROJECT_DIR%\resource" (
     echo.
     echo resource/
     echo ├── NeuralHaircut/
-    echo │   ├── pretrained_models/
-    echo │   │   ├── diffusion_prior/
-    echo │   │   │   └── dif_ckpt.pt
-    echo │   │   └── strand_prior/
-    echo │   │       └── strand_ckpt.pt
+    echo │   ├── diffusion_prior/
+    echo │   │   └── dif_ckpt.pt
+    echo │   ├── strand_prior/
+    echo │   │   └── strand_ckpt.pt
     echo │   └── PIXIE/
     echo │       └── pixie_data/
     echo ├── Matte-Anything/
@@ -191,6 +190,8 @@ echo 正在复制模型文件...
 
 REM 创建必要的目录
 mkdir "%PROJECT_DIR%\ext\NeuralHaircut" 2>nul
+mkdir "%PROJECT_DIR%\ext\NeuralHaircut\diffusion_prior" 2>nul
+mkdir "%PROJECT_DIR%\ext\NeuralHaircut\strand_prior" 2>nul
 mkdir "%PROJECT_DIR%\ext\PIXIE" 2>nul
 mkdir "%PROJECT_DIR%\ext\Matte-Anything\pretrained" 2>nul
 mkdir "%PROJECT_DIR%\ext\openpose\models\pose\coco" 2>nul
diff --git a/run.bat b/run.bat
index 81aa267..31a8799 100644
--- a/run.bat
+++ b/run.bat
@@ -340,3 +340,21 @@ IF %ERRORLEVEL% NEQ 0 (
     echo ERROR: Failed to run command
     exit /b 1
 )
+
+REM 检查模型文件
+IF NOT EXIST "%PROJECT_DIR%\ext\NeuralHaircut\diffusion_prior\dif_ckpt.pt" (
+    echo 错误：未找到Neural Haircut模型，请先运行install.bat
+    exit /b 1
+)
+
+REM 检查输入数据
+IF NOT EXIST "%DATA_PATH%" (
+    echo 错误：未找到数据目录 %DATA_PATH%
+    exit /b 1
+)
+
+REM 检查CUDA设备
+nvidia-smi >nul 2>nul || (
+    echo 错误：未找到NVIDIA GPU或驱动未正确安装
+    exit /b 1
+)