Update app.py

app.py

@@ -15,18 +15,6 @@ from transformers.generation.streamers import TextIteratorStreamer
 import spaces
 import torch
 
-
-if torch.cuda.is_available():
-    device = "cuda"
-    print("Using GPU")
-else:
-    device = "cpu"
-    print("Using CPU")
-
-device = "cuda:0"
-tokenizer, model, image_processor, context_len = load_pretrained_model(model_path="./checkpoints/llava-phi3-3.8b-lora", model_name="llava-phi3-3.8b-lora", model_base="microsoft/Phi-3-mini-128k-instruct", load_8bit=False, load_4bit=False, device=device)
-model.to(device)
-
 def load_image(image_file):
     if image_file.startswith('http') or image_file.startswith('https'):
         response = requests.get(image_file)
@@ -37,7 +25,10 @@ def load_image(image_file):
 
 @spaces.GPU
 def bot_streaming(message, history):
-    print(message)
+    device = "cuda:0"
+    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path="./checkpoints/llava-phi3-3.8b-lora", model_name="llava-phi3-3.8b-lora", model_base="microsoft/Phi-3-mini-128k-instruct", load_8bit=False, load_4bit=False, device=device)
+    model.to(device)
+    
     if message["files"]:
         # message["files"][-1] is a Dict or just a string
         if type(message["files"][-1]) == dict: