Yw22 commited on
Commit
f739617
·
1 Parent(s): dff466e
Files changed (1) hide show
  1. examples/blobctrl/blobctrl_app.py +7 -9
examples/blobctrl/blobctrl_app.py CHANGED
@@ -64,16 +64,14 @@ if not os.path.exists(Dino_path):
64
  print(f"Dino checkpoints downloaded to {Dino_path}")
65
 
66
  # download SAM model
67
- SAM_path = "examples/blobctrl/models/sam/sam_vit_h_4b8939.pth"
68
  if not os.path.exists(SAM_path):
69
- os.makedirs(os.path.dirname(SAM_path), exist_ok=True)
70
- import urllib.request
71
- print(f"Downloading SAM model...")
72
- urllib.request.urlretrieve(
73
- "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
74
- SAM_path
75
  )
76
- print(f"SAM model downloaded to {SAM_path}")
77
 
78
 
79
  ## load models and pipeline
@@ -81,7 +79,7 @@ blobnet_path = "./examples/blobctrl/models/blobnet"
81
  unet_lora_path = "./examples/blobctrl/models/unet_lora"
82
  stabel_diffusion_model_path = "./examples/blobctrl/models/stable-diffusion-v1-5"
83
  dinov2_path = "./examples/blobctrl/models/dinov2-large"
84
- sam_path = "./examples/blobctrl/models/sam/sam_vit_h_4b8939.pth"
85
 
86
  ## unet
87
  print(f"Loading UNet...")
 
64
  print(f"Dino checkpoints downloaded to {Dino_path}")
65
 
66
  # download SAM model
67
+ SAM_path = "examples/blobctrl/models/sam"
68
  if not os.path.exists(SAM_path):
69
+ SAM_path = snapshot_download(
70
+ repo_id="kunkaran/sam_vit_h_4b8939.pth",
71
+ local_dir=SAM_path,
72
+ token=os.getenv("HF_TOKEN"),
 
 
73
  )
74
+ print(f"SAM checkpoints downloaded to {SAM_path}")
75
 
76
 
77
  ## load models and pipeline
 
79
  unet_lora_path = "./examples/blobctrl/models/unet_lora"
80
  stabel_diffusion_model_path = "./examples/blobctrl/models/stable-diffusion-v1-5"
81
  dinov2_path = "./examples/blobctrl/models/dinov2-large"
82
+ sam_path = "./examples/blobctrl/models/sam/sam_vt_h.pth"
83
 
84
  ## unet
85
  print(f"Loading UNet...")