enabling zerogpu

app.py

@@ -20,24 +20,31 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
+import spaces
+import torch
+
 import gradio as gr
 
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"current device is: {device}")
+
 fw_modelcard = "amurienne/gallek-m2m100"
 bw_modelcard = "amurienne/kellag-m2m100"
 
 fw_model = AutoModelForSeq2SeqLM.from_pretrained(fw_modelcard)
 fw_tokenizer = AutoTokenizer.from_pretrained(fw_modelcard)
 
-fw_translation_pipeline = pipeline("translation", model=fw_model, tokenizer=fw_tokenizer, src_lang='fr', tgt_lang='br', max_length=400, device="cpu")
+fw_translation_pipeline = pipeline("translation", model=fw_model, tokenizer=fw_tokenizer, src_lang='fr', tgt_lang='br', max_length=400, device=device)
 
 bw_model = AutoModelForSeq2SeqLM.from_pretrained(bw_modelcard)
 bw_tokenizer = AutoTokenizer.from_pretrained(bw_modelcard)
 
-bw_translation_pipeline = pipeline("translation", model=bw_model, tokenizer=bw_tokenizer, src_lang='br', tgt_lang='fr', max_length=400, device="cpu")
+bw_translation_pipeline = pipeline("translation", model=bw_model, tokenizer=bw_tokenizer, src_lang='br', tgt_lang='fr', max_length=400, device=device)
 
 # translation function
+@spaces.GPU
 def translate(text, direction):
     if direction == "fr_to_br":
         return fw_translation_pipeline("traduis de français en breton: " + text)[0]['translation_text']