Update app_docker.py

app_docker.py

@@ -11,13 +11,20 @@ from torchvision.transforms import ToTensor, ToPILImage
 
 
 # -------------------------- HuggingFace -------------------------------
-from huggingface_hub import hf_hub_download, snapshot_download
-model_name = "iimmortall/UltraFusion"
-auth_token = os.getenv("HF_AUTH_TOKEN")
-model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app", force_download=True)
-# import sys
-# sys.path.append("/data")
-# sys.path.append("/data/")
+# 1. Download the model online
+# from huggingface_hub import hf_hub_download, snapshot_download
+# model_name = "iimmortall/UltraFusion"
+# auth_token = os.getenv("HF_AUTH_TOKEN")
+# model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app", force_download=True)
+# model_folder = ""
+
+# 2. using pre-download model
+# from huggingface_hub import hf_hub_download, snapshot_download
+# model_name = "iimmortall/UltraFusion"
+# auth_token = os.getenv("HF_AUTH_TOKEN")
+# model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/data", force_download=True)
+model_folder = "/data"
+sys.path.append(f"{model_folder}")
 
 from ultrafusion_utils import load_model, run_ultrafusion, check_input
 PYCUDA_FLAG = True
@@ -31,7 +38,7 @@ RUN_TIMES = 0
 
 to_tensor = ToTensor()
 to_pil = ToPILImage()
-ultrafusion_pipe, flow_model = load_model()
+ultrafusion_pipe, flow_model = load_model(model_folder)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 if torch.cuda.is_available():