first commit

predict.py

+import numpy as np
+import tempfile
+import shutil
+import os
+from PIL import Image
+import subprocess
+from cog import BasePredictor, Input, Path
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self,
+        image: Path = Input(
+            description="Input Image.",
+        ),
+    ) -> Path:
+        input_dir = "input_dir"
+        output_path = Path(tempfile.mkdtemp()) / "output.png"
+
+        try:
+            for d in [input_dir, "results"]:
+                if os.path.exists(input_dir):
+                    shutil.rmtree(input_dir)
+            os.makedirs(input_dir, exist_ok=False)
+
+            input_path = os.path.join(input_dir, os.path.basename(image))
+            shutil.copy(str(image), input_path)
+            subprocess.call(
+                [
+                    "python",
+                    "hat/test.py",
+                    "-opt",
+                    "options/test/HAT_SRx4_ImageNet-LR.yml",
+                ]
+            )
+            res_dir = os.path.join(
+                "results", "HAT_SRx4_ImageNet-LR", "visualization", "custom"
+            )
+            assert (
+                len(os.listdir(res_dir)) == 1
+            ), "Should contain only one result for Single prediction."
+            res = Image.open(os.path.join(res_dir, os.listdir(res_dir)[0]))
+            res.save(str(output_path))
+
+        finally:
+            pass
+            shutil.rmtree(input_dir)
+            shutil.rmtree("results")
+
+        return output_path

requirements.txt

+einops
+torch>=1.7
+basicsr==1.3.4.9

run.sh

+ulimit -u 200000
+echo "START TIME: $(date)"
+hostfile=./hostfile
+
+np=$(cat $hostfile|sort|uniq |wc -l)
+np=$(($np*8))
+
+echo $np
+
+nodename=$(cat $hostfile |sed -n "1p")
+dist_url=`echo $nodename | awk '{print $1}'`
+which mpirun
+
+# 添加pythonlib环境, 用户需修改为自己的环境变量地址
+export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/path/of/conda/envs/{env_name}/lib
+export PYTHON=python3
+export NPROC_PER_NODE=4
+
+# 使用mpirun 进行多卡训练，分别对单张卡进行单线程启动
+# -np 显卡数量
+# -x 将变量传递到single_process.sh脚本中
+mpirun -np $np --allow-run-as-root --hostfile hostfile --bind-to none -x dist_url -x PYTHON -x NPROC_PER_NODE `pwd`/single_process.sh
+echo "END TIME: $(date)"