Use huggingface_hub to download models

app.py

@@ -9,6 +9,7 @@ import torch
 import torchvision
 from diffusers import AutoencoderKL, DDIMScheduler
 from einops import rearrange
+from huggingface_hub import hf_hub_download
 from omegaconf import OmegaConf
 from PIL import Image
 from torchvision import transforms
@@ -87,72 +88,19 @@ If you have any questions, please feel free to reach me out at <b>ywl@stu.pku.ed
 
 # """
 
-os.makedirs("models/personalized")
-os.makedirs("models/sd1-5")
-
-if not os.path.exists("models/flow_controlnet.ckpt"):
-    os.system(
-        f"wget -q https://huggingface.co/TencentARC/ImageConductor/resolve/main/flow_controlnet.ckpt?download=true -P models/"
-    )
-    os.system(f"mv models/flow_controlnet.ckpt?download=true models/flow_controlnet.ckpt")
-    print(
-        "flow_controlnet Download!",
-    )
-
-if not os.path.exists("models/image_controlnet.ckpt"):
-    os.system(
-        f"wget -q https://huggingface.co/TencentARC/ImageConductor/resolve/main/image_controlnet.ckpt?download=true -P models/"
-    )
-    os.system(f"mv models/image_controlnet.ckpt?download=true models/image_controlnet.ckpt")
-    print(
-        "image_controlnet Download!",
-    )
-
-if not os.path.exists("models/unet.ckpt"):
-    os.system(
-        f"wget -q https://huggingface.co/TencentARC/ImageConductor/resolve/main/unet.ckpt?download=true -P models/"
-    )
-    os.system(f"mv models/unet.ckpt?download=true models/unet.ckpt")
-    print(
-        "unet Download!",
-    )
+flow_controlnet_path = hf_hub_download("TencentARC/ImageConductor", "flow_controlnet.ckpt")
+image_controlnet_path = hf_hub_download("TencentARC/ImageConductor", "image_controlnet.ckpt")
+unet_path = hf_hub_download("TencentARC/ImageConductor", "unet.ckpt")
 
+helloobjects_path = hf_hub_download("TencentARC/ImageConductor", "helloobjects_V12c.safetensors")
+tusun_path = hf_hub_download("TencentARC/ImageConductor", "TUSUN.safetensors")
 
+os.makedirs("models/sd1-5", exist_ok=True)
+sd15_config_path = hf_hub_download("runwayml/stable-diffusion-v1-5", "config.json", subfolder="unet")
 if not os.path.exists("models/sd1-5/config.json"):
-    os.system(
-        f"wget -q https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/unet/config.json?download=true -P models/sd1-5/"
-    )
-    os.system(f"mv models/sd1-5/config.json?download=true  models/sd1-5/config.json")
-    print(
-        "config Download!",
-    )
-
-
+    os.symlink(sd15_config_path, "models/sd1-5/config.json")
 if not os.path.exists("models/sd1-5/unet.ckpt"):
-    os.system(f"cp -r models/unet.ckpt  models/sd1-5/unet.ckpt")
-
-# os.system(f'wget https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/unet/diffusion_pytorch_model.bin?download=true -P models/sd1-5/')
-
-if not os.path.exists("models/personalized/helloobjects_V12c.safetensors"):
-    os.system(
-        f"wget -q https://huggingface.co/TencentARC/ImageConductor/resolve/main/helloobjects_V12c.safetensors?download=true -P models/personalized"
-    )
-    os.system(
-        f"mv models/personalized/helloobjects_V12c.safetensors?download=true models/personalized/helloobjects_V12c.safetensors"
-    )
-    print(
-        "helloobjects_V12c Download!",
-    )
-
-
-if not os.path.exists("models/personalized/TUSUN.safetensors"):
-    os.system(
-        f"wget -q https://huggingface.co/TencentARC/ImageConductor/resolve/main/TUSUN.safetensors?download=true -P models/personalized"
-    )
-    os.system(f"mv models/personalized/TUSUN.safetensors?download=true models/personalized/TUSUN.safetensors")
-    print(
-        "TUSUN Download!",
-    )
+    os.symlink(unet_path, "models/sd1-5/unet.ckpt")
 
 # mv1 = os.system(f'mv /usr/local/lib/python3.10/site-packages/gradio/helpers.py /usr/local/lib/python3.10/site-packages/gradio/helpers_bkp.py')
 # mv2 = os.system(f'mv helpers.py /usr/local/lib/python3.10/site-packages/gradio/helpers.py')
@@ -245,11 +193,11 @@ IMAGE_PATH = {
 
 
 DREAM_BOOTH = {
-    "HelloObject": "models/personalized/helloobjects_V12c.safetensors",
+    "HelloObject": helloobjects_path,
 }
 
 LORA = {
-    "TUSUN": "models/personalized/TUSUN.safetensors",
+    "TUSUN": tusun_path,
 }
 
 LORA_ALPHA = {
@@ -632,6 +580,17 @@ def delete_last_step(tracking_points, first_frame_path, drag_mode):
     return {tracking_points_var: tracking_points, input_image: trajectory_map}
 
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+ImageConductor_net = ImageConductor(
+    device=device,
+    unet_path=unet_path,
+    image_controlnet_path=image_controlnet_path,
+    flow_controlnet_path=flow_controlnet_path,
+    height=256,
+    width=384,
+    model_length=16,
+)
+
 block = gr.Blocks(theme=gr.themes.Soft(radius_size=gr.themes.sizes.radius_none, text_size=gr.themes.sizes.text_md))
 with block:
     with gr.Row():
@@ -644,20 +603,6 @@ with block:
         with gr.Row(equal_height=True):
             gr.Markdown(instructions)
 
-    # device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    device = torch.device("cuda")
-    unet_path = "models/unet.ckpt"
-    image_controlnet_path = "models/image_controlnet.ckpt"
-    flow_controlnet_path = "models/flow_controlnet.ckpt"
-    ImageConductor_net = ImageConductor(
-        device=device,
-        unet_path=unet_path,
-        image_controlnet_path=image_controlnet_path,
-        flow_controlnet_path=flow_controlnet_path,
-        height=256,
-        width=384,
-        model_length=16,
-    )
     first_frame_path_var = gr.State(value=None)
     tracking_points_var = gr.State([])