Update app.py

app.py

@@ -7,6 +7,8 @@ from PIL import Image
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import re
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # --------------------------- Configuration & Session State ---------------------------
 # Define maximum dimensions for the fortune image (in pixels)
 MAX_SIZE = (400, 400)
@@ -56,7 +58,7 @@ def load_finetuned_classifier_model(question):
 # Function to generate a detailed answer by combining the user's question and the fortune detail
 def generate_answer(question, fortune):
     tokenizer = AutoTokenizer.from_pretrained("tonyhui2234/finetuned_model_text_gen")
-    model = AutoModelForSeq2SeqLM.from_pretrained("tonyhui2234/finetuned_model_text_gen")
+    model = AutoModelForSeq2SeqLM.from_pretrained("tonyhui2234/finetuned_model_text_gen", device_map="auto")
     input_text = "Question: " + question + " Fortune: " + fortune
     inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
     outputs = model.generate(