import os import torch import torchvision.transforms as transforms from modelscope import snapshot_download from PIL import Image MODEL_DIR = snapshot_download( f"ccmusic-database/bel_canto", cache_dir=f"{os.getcwd()}/__pycache__", ) TEMP_DIR = f"{os.getcwd()}/flagged" def toCUDA(x): if hasattr(x, "cuda"): if torch.cuda.is_available(): return x.cuda() return x def find_wav_files(folder_path=f"{MODEL_DIR}/examples"): wav_files = [] for root, _, files in os.walk(folder_path): for file in files: if file.endswith(".wav"): file_path = os.path.join(root, file) wav_files.append(file_path) return wav_files def get_modelist(model_dir=MODEL_DIR): try: entries = os.listdir(model_dir) except OSError as e: print(f"无法访问 {model_dir}: {e}") return # 遍历所有条目 output = [] for entry in entries: # 获取完整路径 full_path = os.path.join(model_dir, entry) # 跳过'.git'文件夹 if entry == ".git" or entry == "examples": print(f"跳过 .git 或 examples 文件夹: {full_path}") continue # 检查条目是文件还是目录 if os.path.isdir(full_path): # 打印目录路径 output.append(os.path.basename(full_path)) return output def embed_img(img_path: str, input_size=224): transform = transforms.Compose( [ transforms.Resize([input_size, input_size]), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] ) img = Image.open(img_path).convert("RGB") return transform(img).unsqueeze(0)