fix: load model revision based on input

app.py

@@ -6,29 +6,32 @@ import torch
 import html
 from transformers import DonutProcessor, VisionEncoderDecoderModel
 
-pretrained_repo_name = 'ivelin/donut-refexp-click'
-pretrained_revision = 'main'
-# revision can be git commit hash, branch or tag
-# use 'main' for latest revision
-print(f"Loading model checkpoint: {pretrained_repo_name}")
-
-processor = DonutProcessor.from_pretrained(
-    pretrained_repo_name, revision=pretrained_revision, use_auth_token="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")
-processor.image_processor.do_align_long_axis = False
-# do not manipulate image size and position
-processor.image_processor.do_resize = False
-processor.image_processor.do_thumbnail = False
-processor.image_processor.do_pad = False
-# processor.image_processor.do_rescale = False
-processor.image_processor.do_normalize = True
-print(f'processor image size: {processor.image_processor.size}')
-
-model = VisionEncoderDecoderModel.from_pretrained(
-    pretrained_repo_name, use_auth_token="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", revision=pretrained_revision)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
 
+global model, processor
+
+
+def load_model(pretrained_revision: str = 'main'):
+    global model, processor
+    pretrained_repo_name = 'ivelin/donut-refexp-click'
+    # revision can be git commit hash, branch or tag
+    # use 'main' for latest revision
+    print(f"Loading model checkpoint from repo: {pretrained_repo_name}, revision: {pretrained_revision}")    
+    processor = DonutProcessor.from_pretrained(
+        pretrained_repo_name, revision=pretrained_revision, use_auth_token="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")
+    processor.image_processor.do_align_long_axis = False
+    # do not manipulate image size and position
+    processor.image_processor.do_resize = False
+    processor.image_processor.do_thumbnail = False
+    processor.image_processor.do_pad = False
+    # processor.image_processor.do_rescale = False
+    processor.image_processor.do_normalize = True
+    print(f'processor image size: {processor.image_processor.size}')
+    model = VisionEncoderDecoderModel.from_pretrained(
+        pretrained_repo_name, use_auth_token="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", revision=pretrained_revision)
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+  
 
 def prepare_image_for_encoder(image=None, output_image_size=None):
     """
@@ -86,10 +89,12 @@ def translate_point_coords_from_out_to_in(point=None, input_image_size=None, out
         f"translated point={point}, resized_image_size: {resized_width, resized_height}")
 
 
-def process_refexp(image: Image, prompt: str, revision: str = 'main'):
+def process_refexp(image: Image, prompt: str, model_revision: str = 'main'):
 
     print(f"(image, prompt): {image}, {prompt}")
-    print(f"model checkpoint revision: {revision}")
+    print(f"model checkpoint revision: {model_revision}")
+
+    load_model(model_revision)
 
     # trim prompt to 80 characters and normalize to lowercase
     prompt = prompt[:80].lower()