Update app.py

app.py

@@ -115,22 +115,22 @@ def load_model(check_type):
 def process_image(model, tokenizer, transform, device, check_type, image, text):
     global current_vis, current_bpe, current_index
     src_size = image.size
-    # Ensure all processing is done on the correct device
-    image = image.to(device)
 
+    # Convert PIL Image to Tensor and move to the appropriate device
     if 'TokenOCR' in check_type:
+        # If dynamic preprocessing is required, handle differently
         images, target_ratio = dynamic_preprocess(image, min_num=1, max_num=12, 
                                                   image_size=model.config.force_image_size,
                                                   use_thumbnail=model.config.use_thumbnail,
                                                   return_ratio=True)
-        pixel_values = torch.stack([transform(img) for img in images]).to(device)
+        pixel_values = torch.stack([transform(img).to(device) for img in images])
     else:
-        pixel_values = torch.stack([transform(image)]).to(device)
+        # Standard image processing for a single image
+        pixel_values = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
         target_ratio = (1, 1)
 
     text += ' '
-    input_ids = tokenizer(text)['input_ids'][1:]
-    input_ids = torch.tensor(input_ids, device=device)
+    input_ids = tokenizer(text, return_tensors='pt').input_ids.to(device)  # Ensure tokens are on the same device
 
     with torch.no_grad():
         if 'R50' in check_type:
@@ -147,14 +147,14 @@ def process_image(model, tokenizer, transform, device, check_type, image, text):
         resized_size = size1 if size1 is not None else size2
 
     attn_map = similarity.reshape(len(text_embeds), resized_size[0], resized_size[1])
-    all_bpe_strings = [tokenizer.decode(input_id) for input_id in input_ids]
-    current_vis = generate_similiarity_map([image.cpu()], attn_map.cpu(), 
-                                           [tokenizer.decode([i]) for i in input_ids], 
+    current_vis = generate_similiarity_map([image], attn_map, 
+                                           [tokenizer.decode([i]) for i in input_ids.squeeze()], 
                                            [], target_ratio, src_size)
 
-    current_bpe = [tokenizer.decode([i]) for i in input_ids]
+    current_bpe = [tokenizer.decode([i]) for i in input_ids.squeeze()]
     current_bpe[-1] = text
-    return image.cpu(), current_vis[0], current_bpe[0]
+    return image, current_vis[0], current_bpe[0]
+
 
 # 事件处理函数
 def update_index(change):