Update runpod handler for Seamless

server.py

@@ -21,65 +21,56 @@ from lang_list import (
     TEXT_SOURCE_LANGUAGE_NAMES,
 )
 
-# CHECKPOINTS_PATH = pathlib.Path(os.getenv("CHECKPOINTS_PATH", "/home/user/app/models"))
-# if not CHECKPOINTS_PATH.exists():
-#     snapshot_download(repo_id="facebook/seamless-m4t-v2-large", repo_type="model", local_dir=CHECKPOINTS_PATH)
-# asset_store.env_resolvers.clear()
-# asset_store.env_resolvers.append(lambda: "demo")
-# demo_metadata = [
-#     {
-#         "name": "seamlessM4T_v2_large@demo",
-#         "checkpoint": f"file://{CHECKPOINTS_PATH}/seamlessM4T_v2_large.pt",
-#         "char_tokenizer": f"file://{CHECKPOINTS_PATH}/spm_char_lang38_tc.model",
-#     },
-#     {
-#         "name": "vocoder_v2@demo",
-#         "checkpoint": f"file://{CHECKPOINTS_PATH}/vocoder_v2.pt",
-#     },
-# ]
-# asset_store.metadata_providers.append(InProcAssetMetadataProvider(demo_metadata))
+CHECKPOINTS_PATH = pathlib.Path(os.getenv("CHECKPOINTS_PATH", "/home/user/app/models"))
+if not CHECKPOINTS_PATH.exists():
+    snapshot_download(repo_id="facebook/seamless-m4t-v2-large", repo_type="model", local_dir=CHECKPOINTS_PATH)
+asset_store.env_resolvers.clear()
+asset_store.env_resolvers.append(lambda: "demo")
+demo_metadata = [
+    {
+        "name": "seamlessM4T_v2_large@demo",
+        "checkpoint": f"file://{CHECKPOINTS_PATH}/seamlessM4T_v2_large.pt",
+        "char_tokenizer": f"file://{CHECKPOINTS_PATH}/spm_char_lang38_tc.model",
+    },
+    {
+        "name": "vocoder_v2@demo",
+        "checkpoint": f"file://{CHECKPOINTS_PATH}/vocoder_v2.pt",
+    },
+]
+asset_store.metadata_providers.append(InProcAssetMetadataProvider(demo_metadata))
 
-# if torch.cuda.is_available():
-#     device = torch.device("cuda:0")
-#     dtype = torch.float16
-# else:
-#     device = torch.device("cpu")
-#     dtype = torch.float32
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+    dtype = torch.float16
+else:
+    device = torch.device("cpu")
+    dtype = torch.float32
 
-# translator = Translator(
-#     model_name_or_card="seamlessM4T_v2_large",
-#     vocoder_name_or_card="vocoder_v2",
-#     device=device,
-#     dtype=dtype,
-#     apply_mintox=True,
-# )
+translator = Translator(
+    model_name_or_card="seamlessM4T_v2_large",
+    vocoder_name_or_card="vocoder_v2",
+    device=device,
+    dtype=dtype,
+    apply_mintox=True,
+)
 
-# def run_t2tt(input_text: str, source_language: str, target_language: str) -> str:
-#     source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
-#     target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
-#     out_texts, _ = translator.predict(
-#         input=input_text,
-#         task_str="T2TT",
-#         src_lang=source_language_code,
-#         tgt_lang=target_language_code,
-#     )
-#     return str(out_texts[0])
+def run_t2tt(input_text: str, source_language: str, target_language: str) -> str:
+    source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
+    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
+    out_texts, _ = translator.predict(
+        input=input_text,
+        task_str="T2TT",
+        src_lang=source_language_code,
+        tgt_lang=target_language_code,
+    )
+    return str(out_texts[0])
 
 import runpod
 
-def is_even(job):   
-
-    job_input = job["input"]
-    the_number = job_input["number"]
-
-    if not isinstance(the_number, int):
-        return {"error": "Silly human, you need to pass an integer."}
-
-    if the_number % 2 == 0:
-        return True
-    
-    return False
-
-    # output_text = run_t2tt(input_text, source_language, target_language)
+def runpod_handler(job):
+    input_text = job["input_text"]
+    source_language = job["source_language"]
+    target_language = job["target_language"]
+    return run_t2tt(input_text, source_language, target_language)
 
-runpod.serverless.start({"handler": is_even})
+runpod.serverless.start({"handler": runpod_handler})

test_input.json

@@ -0,0 +1,7 @@
+{
+    "input": {
+        "input_text": "How are you doing today?",
+        "source_language": "English",
+        "target_language": "Mandarin Chinese"
+    }
+}