Update app.py

app.py

@@ -24,6 +24,7 @@ IMAGE_SAFETY_PATCHES = {
     "default": "safety_patch.pt"
 }
 
+DEVICE = "cpu"
 
 def rtp_read(text_file):
     dataset = []
@@ -40,7 +41,7 @@ model = loaded_model_name = tokenizer = image_processor = context_len = my_gener
 def load_model_async(model_path, model_name):
     global tokenizer, model, image_processor, context_len, loaded_model_name, my_generator
     print(f"Loading {model_name} model ... ")
-    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, load_4bit=False)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, device_map=DEVICE, device=DEVICE)
     if "llava" in model_name.lower():
         loaded_model_name = "LLaVA"
     else:
@@ -90,13 +91,13 @@ def generate_answer(image, user_message: str, requested_model_name: str,
             image = load_image(image)
 
         # transform the image using the visual encoder (CLIP) of LLaVA 1.5; the processed image size would be PyTorch tensor whose shape is (336,336).
-        image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].cuda()
+        image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].to(DEVICE)
 
         if image_safety_patch != None:
             # make the image pixel values between (0,1)
             image = normalize(image)
             # load the safety patch tensor whose values are (0,1)
-            safety_patch = torch.load(image_safety_patch).cuda()
+            safety_patch = torch.load(image_safety_patch).to(DEVICE)
             # apply the safety patch to the input image, clamp it between (0,1) and denormalize it to the original pixel values
             safe_image = denormalize((image + safety_patch).clamp(0, 1))
             # make sure the image value is between (0,1)