Switched to optimum.nvidia

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import pipeline
+from optimum.nvidia.pipelines import pipeline
 import torch
 import os
 import sys
@@ -18,7 +18,7 @@ if not hf_token:
 def load_pipeline(model_name):
     with st.spinner(f'Loading {model_name}... This may take several minutes.'):
         try:
-            pipe = pipeline("text-generation", model=model_name,device=-1)
+            pipe = pipeline("text-generation", model=model_name,use_fp8=True)
         except Exception as e:
             st.error(f"An error occurred: {e}")
             st.stop()

requirements.txt

@@ -1,5 +1,4 @@
 streamlit
-transformers
 torch
-bitsandbytes==0.39.1
-accelerate
+--extra-index-url https://pypi.nvidia.com
+optimum-nvidia