Update app.py

app.py

@@ -7,7 +7,8 @@ import torchvision.transforms as T
 from transformers import AutoTokenizer
 import gradio as gr
 from resnet50 import build_model
-from utils import generate_similiarity_map, post_process, load_tokenizer, build_transform_R50
+# from utils import generate_similiarity_map, post_process, load_tokenizer, build_transform_R50
+from utils import generate_similiarity_map, get_transform, post_process, load_tokenizer, build_transform_R50
 from utils import IMAGENET_MEAN, IMAGENET_STD
 from internvl.train.dataset import dynamic_preprocess
 from internvl.model.internvl_chat import InternVLChatModel
@@ -42,20 +43,16 @@ def load_model(check_type):
     elif 'TokenFD' in check_type:
         model_path = CHECKPOINTS[check_type]
         tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False, use_auth_token=HF_TOKEN)
-        model = InternVLChatModel.from_pretrained(model_path, torch_dtype=torch.bfloat16).eval()
-        transform = T.Compose([
-            T.Lambda(lambda img: img.convert('RGB')),
-            T.Resize((224, 224)),
-            T.ToTensor(),
-            T.Normalize(IMAGENET_MEAN, IMAGENET_STD)
-        ])
-    
+        # model = InternVLChatModel.from_pretrained(model_path, torch_dtype=torch.bfloat16).eval()
+        model = InternVLChatModel.from_pretrained(checkpoint_vit_english, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16 ,load_in_8bit=False, load_in_4bit=False).eval()
+        transform = get_transform(is_train=False, image_size=model.config.force_image_size)
+        
     return model.to(device), tokenizer, transform, device
 
 def process_image(model, tokenizer, transform, device, check_type, image, text):
     
     src_size = image.size
-    if 'TokenOCR' in check_type:
+    if 'TokenFD' in check_type:
         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,