Update pipeline.py

pipeline.py

@@ -2,7 +2,7 @@ from diffusers import DiffusionPipeline
 import os
 import sys
 from huggingface_hub import HfApi, hf_hub_download
-from .tools import build_dataset_json_from_list
+# from .tools import build_dataset_json_from_list
 import torch
 
 class MOSDiffusionPipeline(DiffusionPipeline):
@@ -67,28 +67,32 @@ class MOSDiffusionPipeline(DiffusionPipeline):
 
 
     @torch.no_grad()
-    def __call__(self, *args, **kwargs):
+    def __call__(self, prompt: str):
         """
         Run the MOS Diffusion Pipeline. This method calls the infer function from infer_mos5.py.
         """
         from .infer.infer_mos5 import infer
+        dataset_key = self.build_dataset_json_from_prompt(prompt)
 
+        # we run inference with the prompt - configs - and other settings
         infer(
-            dataset_key=self.dataset_key,
+            dataset_key=dataset_key,
             configs=self.configs,
             config_yaml_path=self.config_yaml,
-            exp_group_name=self.exp_group_name,
-            exp_name=self.exp_name
+            exp_group_name="qa_mdt",
+            exp_name="mos_as_token"
         )
 
+    def build_dataset_json_from_prompt(self, prompt: str):
+        """
+        Build dataset_key dynamically from the provided prompt.
+        """
+        # for simplicity let's just return the prompt as the dataset_key
+        return {"prompt": prompt}
+        
+
 # 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="test_prompts/good_prompts_1.lst",
-        reload_from_ckpt="checkpoints/checkpoint_389999.ckpt",
-        base_folder=None
-    )
-
-    # Run the pipeline
-    pipeline()
+    pipe = MOSDiffusionPipeline()
+    result = pipe("Generate a description of a sunny day.")
+    print(result)