preload models for fast run

app.py

@@ -4,7 +4,6 @@ import os
 import torch
 import spaces
 
-
 # Define model paths
 MODEL_PATHS = {
     "Terjman-Nano-v2": "BounharAbdelaziz/Terjman-Nano-v2.0",
@@ -16,12 +15,50 @@ MODEL_PATHS = {
 # Load environment token
 TOKEN = os.environ['TOKEN']
 
+# Preload models and tokenizers
+def preload_models():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f"[INFO] Using device: {device}")
+
+    # Load Nano and Large models
+    nano_large_models = {}
+    for model_name in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
+        print(f"[INFO] Loading {model_name}...")
+        translator = pipeline(
+            "translation",
+            model=MODEL_PATHS[model_name],
+            token=TOKEN,
+            device=device if device.startswith("cuda") else -1
+        )
+        nano_large_models[model_name] = translator
+
+    # Load Ultra and Supreme models
+    ultra_supreme_models = {}
+    for model_name in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
+        print(f"[INFO] Loading {model_name}...")
+        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATHS[model_name], token=TOKEN).to(device)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATHS[model_name], token=TOKEN)
+        translator = pipeline(
+            "translation",
+            model=model,
+            tokenizer=tokenizer,
+            device=device if device.startswith("cuda") else -1,
+            src_lang="eng_Latn",
+            tgt_lang="ary_Arab"
+        )
+        ultra_supreme_models[model_name] = translator
+
+    return nano_large_models, ultra_supreme_models
+
+# Preload all models
+nano_large_models, ultra_supreme_models = preload_models()
+
 # Translation function for Nano and Large models
 @spaces.GPU
-def translate_nano_large(text, model_path):
-    translator = pipeline("translation", model=model_path, token=TOKEN)
+def translate_nano_large(text, model_name):
+    translator = nano_large_models[model_name]
     translated = translator(
-        text, 
+        text,
         max_length=512,
         num_beams=4,
         no_repeat_ngram_size=3,
@@ -35,30 +72,17 @@ def translate_nano_large(text, model_path):
 
 # Translation function for Ultra and Supreme models
 @spaces.GPU
-def translate_ultra_supreme(text, model_path):
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    print(f'[INFO] Using device: {device}')
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_path, token=TOKEN)
-    tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang="eng_Latn", tgt_lang="ary_Arab", token=TOKEN)
-    translator = pipeline(
-        "translation",
-        model=model,
-        tokenizer=tokenizer,
-        max_length=512,
-        src_lang="eng_Latn",  # Keep src_lang and tgt_lang in the pipeline
-        tgt_lang="ary_Arab",
-        device=device,
-    )
+def translate_ultra_supreme(text, model_name):
+    translator = ultra_supreme_models[model_name]
     translation = translator(text)[0]['translation_text']
     return translation
 
 # Main translation function
 def translate_text(text, model_choice):
-    model_path = MODEL_PATHS[model_choice]
     if model_choice in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
-        return translate_nano_large(text, model_path)
+        return translate_nano_large(text, model_choice)
     elif model_choice in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
-        return translate_ultra_supreme(text, model_path)
+        return translate_ultra_supreme(text, model_choice)
     else:
         return "Invalid model selection."