register pipeline

.gitignore

@@ -1 +1,2 @@
 .vscode
+__pycache__

README.md

@@ -55,3 +55,33 @@ Here are some results:
   <source src="https://huggingface.co/cubbk/orpheus-swedish/resolve/main/audios/6.wav" type="audio/wav">
   Your browser does not support the audio element.
 </audio>
+
+## How to use
+
+```
+%pip install transformers snac soundfile torch torchaudio
+```
+
+```py
+from transformers import pipeline
+
+# Use your custom task name and model repo id
+pipe = pipeline(task="orpheus-swedish", model="cubbk/orpheus-swedish", trust_remote_code=True)
+```
+
+```py
+from IPython.display import Audio, display
+
+prompt = [
+        "Enligt brittiska medier kommer ledarna att presentera en techpakt som ska stärka ländernas samarbete kring AI, kvantfysik och kärnkraft."
+    ]
+outputs = pipe(prompt)
+
+for i in range(len(outputs)):
+    print(prompt[i])
+    samples = outputs[i][0]
+    display(Audio(samples.detach().squeeze().to("cpu").numpy(), rate=24000))
+```
+
+Good to know: performs poorly on short text
+Cuts out at 14 sec.(didn't figure out why)

pipeline/pipeline.py

@@ -6,7 +6,7 @@ from huggingface_hub import snapshot_download
 
 
 class MyPipeline(Pipeline):
-    def __init__(self):
+    def __init__(self, model=None, tokenizer=None, feature_extractor=None, **kwargs):
         self.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
         self.snac_model = self.snac_model.to("cpu")
         print(
@@ -41,7 +41,11 @@ class MyPipeline(Pipeline):
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name, torch_dtype=torch.bfloat16
         )
-        self.model.cuda()
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+        else:
+            print("CUDA not available, running on CPU.")
+
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
 
         super().__init__(model=self.model, tokenizer=self.tokenizer)

pipeline/register_pipeline.py

@@ -0,0 +1,14 @@
+from .pipeline import MyPipeline
+from transformers.pipelines import PIPELINE_REGISTRY
+from transformers import AutoModelForCausalLM, pipeline as hf_pipeline
+
+# Register your custom pipeline under a custom task name
+PIPELINE_REGISTRY.register_pipeline(
+    "orpheus-swedish",  # custom task name
+    pipeline_class=MyPipeline,
+    pt_model=AutoModelForCausalLM,
+)
+
+# Instantiate your pipeline using the custom task name and your model repo id
+my_pipe = hf_pipeline(task="orpheus-swedish", model="cubbk/orpheus-swedish")
+my_pipe.push_to_hub("cubbk/orpheus-swedish")