Create pipeline.py

pipeline.py

@@ -0,0 +1,68 @@
+from diffusers import DiffusionPipeline
+import torch
+import os
+import yaml
+from audioldm_train.utilities.tools import build_dataset_json_from_list
+from infer_mos5 import infer  # Importing the infer function
+
+class MOSDiffusionPipeline(DiffusionPipeline):
+
+    def __init__(self, config_yaml, list_inference, reload_from_ckpt=None):
+        """
+        Initialize the MOS Diffusion pipeline.
+
+        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.
+        """
+        super().__init__()
+
+        # we load and process the yaml config
+        self.config_yaml = config_yaml
+        self.list_inference = list_inference
+        self.reload_from_ckpt = reload_from_ckpt
+
+        # we load the yaml config
+        config_yaml_path = os.path.join(self.config_yaml)
+        self.configs = yaml.load(open(config_yaml_path, "r"), Loader=yaml.FullLoader)
+
+        # override checkpoint if provided--
+        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))
+
+    @torch.no_grad()
+    def __call__(self, *args, **kwargs):
+        """
+        Run the MOS Diffusion Pipeline. This method calls the infer function from infer_mos5.py.
+
+        Args:
+            *args: Additional arguments.
+            **kwargs: Keyword arguments that may contain overrides for configurations.
+
+        Returns:
+            None. Inference is performed and samples are generated.
+        """
+        # here call the infer function to perform the inference
+        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
+        )
+
+# # This is an example of how to use the pipeline
+# if __name__ == "__main__":
+#     pipeline = MOSDiffusionPipeline(
+#         config_yaml="audioldm_train/config/mos_as_token/qa_mdt.yaml",
+#         list_inference="/content/qa-mdt/test_prompts/good_prompts_1.lst",
+#         reload_from_ckpt="/content/qa-mdt/checkpoint_389999.ckpt"
+#     )
+
+#     # Run the pipeline
+#     pipeline()