|
from diffusers import DiffusionPipeline |
|
import os |
|
import sys |
|
from huggingface_hub import HfApi, hf_hub_download |
|
from .tools import build_dataset_json_from_list |
|
import torch |
|
|
|
class MOSDiffusionPipeline(DiffusionPipeline): |
|
|
|
def __init__(self, config_yaml, list_inference, reload_from_ckpt=None, base_folder=None): |
|
""" |
|
Initialize the MOS Diffusion pipeline and download the necessary files/folders. |
|
|
|
Args: |
|
config_yaml (str): Path to the YAML configuration file. |
|
list_inference (str): Path to the file containing inference prompts. |
|
reload_from_ckpt (str, optional): Checkpoint path to reload from. |
|
base_folder (str, optional): Base folder to store downloaded files. Defaults to the current working directory. |
|
""" |
|
super().__init__() |
|
|
|
|
|
self.base_folder = base_folder if base_folder else os.getcwd() |
|
self.repo_id = "jadechoghari/qa-mdt" |
|
self.config_yaml = config_yaml |
|
self.list_inference = list_inference |
|
self.reload_from_ckpt = reload_from_ckpt |
|
config_yaml_path = os.path.join(self.config_yaml) |
|
self.configs = self.load_yaml(config_yaml_path) |
|
if self.reload_from_ckpt is not None: |
|
self.configs["reload_from_ckpt"] = self.reload_from_ckpt |
|
|
|
self.dataset_key = build_dataset_json_from_list(self.list_inference) |
|
self.exp_name = os.path.basename(self.config_yaml.split(".")[0]) |
|
self.exp_group_name = os.path.basename(os.path.dirname(self.config_yaml)) |
|
|
|
def download_required_folders(self): |
|
""" |
|
Downloads the necessary folders from the Hugging Face Hub if they are not already available locally. |
|
""" |
|
api = HfApi() |
|
|
|
files = api.list_repo_files(repo_id=self.repo_id) |
|
|
|
required_folders = ["audioldm_train", "checkpoints", "infer", "log", "taming", "test_prompts"] |
|
|
|
files_to_download = [f for f in files if any(f.startswith(folder) for folder in required_folders)] |
|
|
|
for file in files_to_download: |
|
local_file_path = os.path.join(self.base_folder, file) |
|
if not os.path.exists(local_file_path): |
|
downloaded_file = hf_hub_download(repo_id=self.repo_id, filename=file) |
|
|
|
os.makedirs(os.path.dirname(local_file_path), exist_ok=True) |
|
|
|
os.rename(downloaded_file, local_file_path) |
|
|
|
sys.path.append(self.base_folder) |
|
|
|
def load_yaml(self, yaml_path): |
|
""" |
|
Helper method to load the YAML configuration. |
|
""" |
|
import yaml |
|
with open(yaml_path, "r") as f: |
|
return yaml.safe_load(f) |
|
|
|
|
|
@torch.no_grad() |
|
def __call__(self, *args, **kwargs): |
|
""" |
|
Run the MOS Diffusion Pipeline. This method calls the infer function from infer_mos5.py. |
|
""" |
|
from .infer.infer_mos5 import infer |
|
|
|
infer( |
|
dataset_key=self.dataset_key, |
|
configs=self.configs, |
|
config_yaml_path=self.config_yaml, |
|
exp_group_name=self.exp_group_name, |
|
exp_name=self.exp_name |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
pipeline = MOSDiffusionPipeline( |
|
config_yaml="audioldm_train/config/mos_as_token/qa_mdt.yaml", |
|
list_inference="test_prompts/good_prompts_1.lst", |
|
reload_from_ckpt="checkpoints/checkpoint_389999.ckpt", |
|
base_folder=None |
|
) |
|
|
|
|
|
pipeline() |
|
|
