Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import spaces
+import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 langs = {"German": "de", "Romanian": "ro", "English": "en", "French": "fr", "Spanish": "es", "Italian": "it"}
@@ -15,18 +16,18 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         try:
             model_name_full = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
             tokenizer = AutoTokenizer.from_pretrained(model_name_full)
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full, device_map="auto")
         except EnvironmentError:
             try :   
                 model_name_full = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
                 tokenizer = AutoTokenizer.from_pretrained(model_name_full)
-                model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full, device_map="auto")
             except EnvironmentError as error:
                 return f"Error finding model: {model_name_full}! Try other available language combination.", error
     elif model_name.startswith('facebook/nllb'):
         from languagecodes import nllb_language_codes
         tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=nllb_language_codes[sselected_language])
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
         translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=nllb_language_codes[sselected_language], tgt_lang=nllb_language_codes[tselected_language])
         translated_text = translator(input_text, max_length=360)
         return translated_text[0]['translation_text'], message_text
@@ -34,6 +35,10 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         tokenizer = T5Tokenizer.from_pretrained(model_name)
         model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
 
+    # Move the model to GPU if available
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+        print("CUDA is available! Using GPU.")
     if model_name.startswith("Helsinki-NLP"):
         prompt = input_text
     else: