Upload UltravoxPipeline

config.json

@@ -1,6 +1,6 @@
 {
   "architectures": [
-    "ultravox_model.UltravoxModel"
+    "UltravoxModel"
   ],
   "audio_config": {
     "_name_or_path": "facebook/wav2vec2-base-960h",
@@ -14,10 +14,20 @@
     "model_type": "wav2vec2"
   },
   "audio_model_id": "facebook/wav2vec2-base-960h",
+  "audio_model_lora_config": {
+    "lora_alpha": 8,
+    "r": 0,
+    "target_modules": [
+      "k_proj",
+      "q_proj",
+      "linear_k",
+      "linear_q"
+    ]
+  },
   "audio_token_index": 32000,
   "auto_map": {
     "AutoConfig": "ultravox_config.UltravoxConfig",
-    "AutoModel": "ultravox_model.UltravoxModel"
+    "AutoModelForCausalLM": "ultravox_model.UltravoxModel"
   },
   "custom_pipelines": {
     "ultravox-pipeline": {
@@ -31,7 +41,7 @@
       },
       "impl": "ultravox_pipeline.UltravoxPipeline",
       "pt": [
-        "ultravox_model.UltravoxModel"
+        "UltravoxModel"
       ],
       "tf": [],
       "type": "multimodal"
@@ -58,6 +68,16 @@
     "torch_dtype": "bfloat16"
   },
   "text_model_id": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+  "text_model_lora_config": {
+    "lora_alpha": 8,
+    "r": 0,
+    "target_modules": [
+      "k_proj",
+      "q_proj",
+      "linear_k",
+      "linear_q"
+    ]
+  },
   "torch_dtype": "float32",
   "transformers_version": "4.41.2",
   "vocab_size": 32000

ultravox_model.py

@@ -397,6 +397,11 @@ class UltravoxProjector(nn.Sequential):
         return hidden_states
 
 
+UltravoxModel.register_for_auto_class()
 UltravoxModel.register_for_auto_class("AutoModelForCausalLM")
+transformers.AutoModelForCausalLM.register(
+    ultravox_config.UltravoxConfig, UltravoxModel
+)
+
 
 transformers.activations.ACT2FN["swiglu"] = SwiGLU