support downloading resource

This commit is contained in:
Artiprocher
2025-11-18 10:29:07 +08:00
parent 8332ecebb7
commit 1313f4dd63
2 changed files with 20 additions and 8 deletions

View File

@@ -4,6 +4,7 @@ from PIL import Image
from einops import repeat, reduce from einops import repeat, reduce
from typing import Optional, Union from typing import Optional, Union
from dataclasses import dataclass from dataclasses import dataclass
from huggingface_hub import snapshot_download as hf_snapshot_download
from modelscope import snapshot_download from modelscope import snapshot_download
import numpy as np import numpy as np
from PIL import Image from PIL import Image
@@ -196,6 +197,7 @@ class ModelConfig:
self.local_model_path = "./models" self.local_model_path = "./models"
if not skip_download: if not skip_download:
downloaded_files = glob.glob(self.origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id)) downloaded_files = glob.glob(self.origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id))
if self.download_resource.lower() == "modelscope":
snapshot_download( snapshot_download(
self.model_id, self.model_id,
local_dir=os.path.join(self.local_model_path, self.model_id), local_dir=os.path.join(self.local_model_path, self.model_id),
@@ -203,6 +205,16 @@ class ModelConfig:
ignore_file_pattern=downloaded_files, ignore_file_pattern=downloaded_files,
local_files_only=False local_files_only=False
) )
elif self.download_resource.lower() == "huggingface":
hf_snapshot_download(
self.model_id,
local_dir=os.path.join(self.local_model_path, self.model_id),
allow_patterns=allow_file_pattern,
ignore_patterns=downloaded_files,
local_files_only=False
)
else:
raise ValueError("`download_resource` should be `modelscope` or `huggingface`.")
# Let rank 1, 2, ... wait for rank 0 # Let rank 1, 2, ... wait for rank 0
if use_usp: if use_usp:

View File

@@ -14,7 +14,7 @@ else:
setup( setup(
name="diffsynth", name="diffsynth",
version="1.1.8", version="1.1.9",
description="Enjoy the magic of Diffusion models!", description="Enjoy the magic of Diffusion models!",
author="Artiprocher", author="Artiprocher",
packages=find_packages(), packages=find_packages(),