iimmortall commited on
Commit
fa186dc
·
verified ·
1 Parent(s): 413a033

Update app_docker.py

Browse files
Files changed (1) hide show
  1. app_docker.py +15 -8
app_docker.py CHANGED
@@ -11,13 +11,20 @@ from torchvision.transforms import ToTensor, ToPILImage
11
 
12
 
13
  # -------------------------- HuggingFace -------------------------------
14
- from huggingface_hub import hf_hub_download, snapshot_download
15
- model_name = "iimmortall/UltraFusion"
16
- auth_token = os.getenv("HF_AUTH_TOKEN")
17
- model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app", force_download=True)
18
- # import sys
19
- # sys.path.append("/data")
20
- # sys.path.append("/data/")
 
 
 
 
 
 
 
21
 
22
  from ultrafusion_utils import load_model, run_ultrafusion, check_input
23
  PYCUDA_FLAG = True
@@ -31,7 +38,7 @@ RUN_TIMES = 0
31
 
32
  to_tensor = ToTensor()
33
  to_pil = ToPILImage()
34
- ultrafusion_pipe, flow_model = load_model()
35
 
36
  device = "cuda" if torch.cuda.is_available() else "cpu"
37
  if torch.cuda.is_available():
 
11
 
12
 
13
  # -------------------------- HuggingFace -------------------------------
14
+ # 1. Download the model online
15
+ # from huggingface_hub import hf_hub_download, snapshot_download
16
+ # model_name = "iimmortall/UltraFusion"
17
+ # auth_token = os.getenv("HF_AUTH_TOKEN")
18
+ # model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app", force_download=True)
19
+ # model_folder = ""
20
+
21
+ # 2. using pre-download model
22
+ # from huggingface_hub import hf_hub_download, snapshot_download
23
+ # model_name = "iimmortall/UltraFusion"
24
+ # auth_token = os.getenv("HF_AUTH_TOKEN")
25
+ # model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/data", force_download=True)
26
+ model_folder = "/data"
27
+ sys.path.append(f"{model_folder}")
28
 
29
  from ultrafusion_utils import load_model, run_ultrafusion, check_input
30
  PYCUDA_FLAG = True
 
38
 
39
  to_tensor = ToTensor()
40
  to_pil = ToPILImage()
41
+ ultrafusion_pipe, flow_model = load_model(model_folder)
42
 
43
  device = "cuda" if torch.cuda.is_available() else "cpu"
44
  if torch.cuda.is_available():