mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #50 from fakerybakery/fix-downloads

Browse Source
This commit is contained in:
Gong Junmin
2025-05-10 14:05:54 +08:00
committed by GitHub
parent 3ec51ebc83 25b4c4bbcd
commit 0294766a74
1 changed files with 10 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+10 -141
View File
@@ -17,7 +17,7 @@ from loguru import logger
from tqdm import tqdm
import json
import math
from huggingface_hub import hf_hub_download
from huggingface_hub import hf_hub_download, snapshot_download
# from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from acestep.schedulers.scheduling_flow_match_euler_discrete import (
@@ -129,146 +129,15 @@ class ACEStepPipeline:
def load_checkpoint(self, checkpoint_dir=None):
device = self.device
dcae_model_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
vocoder_model_path = os.path.join(checkpoint_dir, "music_vocoder")
ace_step_model_path = os.path.join(checkpoint_dir, "ace_step_transformer")
text_encoder_model_path = os.path.join(checkpoint_dir, "umt5-base")
files_exist = (
os.path.exists(os.path.join(dcae_model_path, "config.json"))
and os.path.exists(
os.path.join(dcae_model_path, "diffusion_pytorch_model.safetensors")
)
and os.path.exists(os.path.join(vocoder_model_path, "config.json"))
and os.path.exists(
os.path.join(vocoder_model_path, "diffusion_pytorch_model.safetensors")
)
and os.path.exists(os.path.join(ace_step_model_path, "config.json"))
and os.path.exists(
os.path.join(ace_step_model_path, "diffusion_pytorch_model.safetensors")
)
and os.path.exists(os.path.join(text_encoder_model_path, "config.json"))
and os.path.exists(
os.path.join(text_encoder_model_path, "model.safetensors")
)
and os.path.exists(
os.path.join(text_encoder_model_path, "special_tokens_map.json")
)
)
if not files_exist:
logger.info(
f"Checkpoint directory {checkpoint_dir} is not complete, downloading from Hugging Face Hub"
)
# download music dcae model
hf_hub_download(
repo_id=REPO_ID,
subfolder="music_dcae_f8c8",
filename="config.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="music_dcae_f8c8",
filename="diffusion_pytorch_model.safetensors",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
# download vocoder model
hf_hub_download(
repo_id=REPO_ID,
subfolder="music_vocoder",
filename="config.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="music_vocoder",
filename="diffusion_pytorch_model.safetensors",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
# download ace_step transformer model
hf_hub_download(
repo_id=REPO_ID,
subfolder="ace_step_transformer",
filename="config.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="ace_step_transformer",
filename="diffusion_pytorch_model.safetensors",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
# download text encoder model
# os.makedirs(text_encoder_model_path, exist_ok=True) # hf_hub_download should create subdirectories
hf_hub_download(
repo_id=REPO_ID,
subfolder="umt5-base",
filename="config.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="umt5-base",
filename="model.safetensors",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="umt5-base",
filename="special_tokens_map.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="umt5-base",
filename="tokenizer_config.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
hf_hub_download(
repo_id=REPO_ID,
subfolder="umt5-base",
filename="tokenizer.json",
local_dir=checkpoint_dir,
local_dir_use_symlinks=False,
)
# Verify files were downloaded correctly
if not all([
os.path.exists(os.path.join(dcae_model_path, "config.json")),
os.path.exists(os.path.join(dcae_model_path, "diffusion_pytorch_model.safetensors")),
os.path.exists(os.path.join(vocoder_model_path, "config.json")),
os.path.exists(os.path.join(vocoder_model_path, "diffusion_pytorch_model.safetensors")),
os.path.exists(os.path.join(ace_step_model_path, "config.json")),
os.path.exists(os.path.join(ace_step_model_path, "diffusion_pytorch_model.safetensors")),
os.path.exists(os.path.join(text_encoder_model_path, "config.json")),
os.path.exists(os.path.join(text_encoder_model_path, "model.safetensors")),
os.path.exists(os.path.join(text_encoder_model_path, "special_tokens_map.json")),
]):
logger.error("Failed to download all required model files. Please check your internet connection and try again.")
logger.info(f"DCAE model path: {dcae_model_path}, files exist: {os.path.exists(os.path.join(dcae_model_path, 'config.json'))}")
logger.info(f"Vocoder model path: {vocoder_model_path}, files exist: {os.path.exists(os.path.join(vocoder_model_path, 'config.json'))}")
logger.info(f"ACE-Step model path: {ace_step_model_path}, files exist: {os.path.exists(os.path.join(ace_step_model_path, 'config.json'))}")
logger.info(f"Text encoder model path: {text_encoder_model_path}, files exist: {os.path.exists(os.path.join(text_encoder_model_path, 'config.json'))}")
raise RuntimeError("Model download failed. See logs for details.")
logger.info("Models downloaded successfully")
if checkpoint_dir is None:
checkpoint_dir_models = snapshot_download(REPO_ID)
else:
checkpoint_dir_models = snapshot_download(REPO_ID, cache_dir=checkpoint_dir)
dcae_model_path = os.path.join(checkpoint_dir_models, "music_dcae_f8c8")
vocoder_model_path = os.path.join(checkpoint_dir_models, "music_vocoder")
ace_step_model_path = os.path.join(checkpoint_dir_models, "ace_step_transformer")
text_encoder_model_path = os.path.join(checkpoint_dir_models, "umt5-base")
dcae_checkpoint_path = dcae_model_path
vocoder_checkpoint_path = vocoder_model_path
ace_step_checkpoint_path = ace_step_model_path