Create huggingface_integration.py

huggingface_integration.py

@@ -0,0 +1,16 @@
+from transformers import PreTrainedModel, PretrainedConfig
+import torch
+from audio_craft_model import AudioCraftGenerator
+
+class AudioCraftForHuggingFace(PreTrainedModel):
+    def __init__(self, config: PretrainedConfig, model_name='facebook/audiogen-medium', duration=5):
+        super(AudioCraftForHuggingFace, self).__init__(config)
+        self.audio_craft_generator = AudioCraftGenerator(model_name, duration)
+
+    def forward(self, descriptions):
+        with torch.no_grad():
+            wav = self.audio_craft_generator(descriptions)
+        return wav
+
+    def save_wav(self, wav, idx):
+        return self.audio_craft_generator.save_wav(wav, idx)