Merge pull request #24 from GooeyAI/seamless_v2

Upgrade SeamlessM4T to v2 and add text2speech + translation support

celeryconfig.py

@@ -3,6 +3,7 @@
 import typing
 
 from celery import Celery
+from celery.exceptions import WorkerShutdown
 from celery.signals import worker_init
 from kombu import Queue
 
@@ -34,12 +35,15 @@ def setup_queues(
     queue_prefix: str = os.environ.get("QUEUE_PREFIX", "gooey-gpu"),
 ):
     def init(**kwargs):
-        for model_id in model_ids:
-            try:
+        model_id = None
+        try:
+            for model_id in model_ids:
                 load_fn(model_id)
-            except Exception as e:
-                traceback.print_exc()
-                raise
+        except:
+            # for some reason, celery seems to swallow exceptions in init
+            print(f"Error loading {model_id}:")
+            traceback.print_exc()
+            raise WorkerShutdown()
 
     init_fns.append(init)
 

chart/model-values.yaml

@@ -202,16 +202,16 @@ deployments:
         bark
 
   - name: "common-seamless"
-    image: *commonImg
+    image: "crgooeyprodwestus1.azurecr.io/gooey-gpu-common:6"
     limits:
-      memory: "20Gi"
+      gpu: "10Gi"
+      memory: "28Gi"  # (220 / 80) * 10
       cpu: "1"
     env:
       IMPORTS: |-
         common.seamless_asr
       SEAMLESS_MODEL_IDS: |-
-        facebook/hf-seamless-m4t-large
-        facebook/hf-seamless-m4t-medium
+        facebook/seamless-m4t-v2-large
 
   - name: "common-diffusion-instruct-pix2pix"
     image: *commonImg

chart/templates/deployment.yaml

@@ -46,6 +46,10 @@ spec:
             value: "{{ $value }}"
           {{- end }}
           {{- end }}
+          {{- range $name, $value := .limits }}
+          - name: "RESOURCE_LIMITS_{{ $name | upper }}"
+            value: "{{ $value }}"
+          {{- end }}
         livenessProbe:
           exec:
             command: [ "bash", "-c", "celery inspect ping -d celery@$HOSTNAME" ]

common/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.2.0-devel-ubuntu20.04
+FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
 
 ARG DEBIAN_FRONTEND=noninteractive