mbarnig/ACE-Step

1import json2from pathlib import Path3import random4 5 6DEFAULT_ROOT_DIR = "examples/default/input_params"7ZH_RAP_LORA_ROOT_DIR = "examples/zh_rap_lora/input_params"8 9class DataSampler:10    def __init__(self, root_dir=DEFAULT_ROOT_DIR):11        self.root_dir = root_dir12        self.input_params_files = list(Path(self.root_dir).glob("*.json"))13        self.zh_rap_lora_input_params_files = list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json"))14        self.zh_rap_lora_input_params_files += list(Path(ZH_RAP_LORA_ROOT_DIR).glob("*.json"))15 16    def load_json(self, file_path):17        with open(file_path, "r", encoding="utf-8") as f:18            return json.load(f)19 20    def sample(self, lora_name_or_path=None):21        if lora_name_or_path is None or lora_name_or_path == "none":22            json_path = random.choice(self.input_params_files)23            json_data = self.load_json(json_path)24        else:25            json_path = random.choice(self.zh_rap_lora_input_params_files)26            json_data = self.load_json(json_path)27            # Update the lora_name in the json_data28            json_data["lora_name_or_path"] = lora_name_or_path29 30        return json_data31