change in model config

custom_model.py

@@ -23,10 +23,12 @@ class SpleeterModel(PreTrainedModel):
         Args:
             audio_path (str): Path to the input audio file.
         Returns:
-            dict: Separated stems.
+            path: Separated stems.
         """
         return self.separator.separate_to_file(audio_path, "separated_audio")
 
 
 AutoConfig.register("spleeter", SpleeterConfig)
 AutoModel.register(SpleeterConfig, SpleeterModel)
+SpleeterConfig.register_for_auto_class()
+SpleeterModel.register_for_auto_class("AutoModel")

test.py

@@ -1,22 +1,22 @@
-from transformers import AutoConfig, AutoModel
-from custom_model import SpleeterModel
+# from transformers import AutoConfig, AutoModel
+# from custom_model import SpleeterModel
 
-config = AutoConfig.from_pretrained("niral-env/youtube_spleeter")
-print("----"*30)
-print(config)
-print("----"*30)
+# config = AutoConfig.from_pretrained("niral-env/youtube_spleeter")
+# print("----"*30)
+# print(config)
+# print("----"*30)
 
-model = SpleeterModel(config)
+# model = SpleeterModel(config)
 
-print(model)
-result = model.forward("vocals.wav")
+# print(model)
+# result = model.forward("vocals.wav")
 
-print(result)
+# print(result)
 
 
-# from transformers import AutoModel
-# model = AutoModel.from_pretrained("niral-env/youtube_spleeter")
-# print(model)
-# result = model.forward("vocals.wav")
+from transformers import AutoModel
+model = AutoModel.from_pretrained("niral-env/youtube_spleeter", trust_remote_code=True)
+print(model)
+result = model.forward("vocals.wav")
 
-# print(result)
+print(result)