Update run.py

run.py

@@ -9,7 +9,7 @@ import time
 processor = TrOCRProcessor.from_pretrained("arcma/decap")
 model = VisionEncoderDecoderModel.from_pretrained("arcma/decap")
 model.eval()
-# torch.compile(model)
+torch.compile(model)
 
 def check(x):
     if len(x) < 6:
@@ -18,10 +18,10 @@ def check(x):
         return False
     return True
 
-def process_image(image):
-    pixel_values = processor(image, return_tensors="pt").pixel_values
+@torch.jit.script
+def process_image(pixel_values):
     with torch.no_grad():
-        generated_ids = model.generate(pixel_values, num_beams=4, num_return_sequences=4)
+        generated_ids = model.generate(pixel_values, num_beams=1, num_return_sequences=1)
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
     generated_text = [x for x in generated_text if check(x)]
     return generated_text[0]
@@ -37,8 +37,8 @@ def process_html(html):
             )
         )
     )
-    
-    return process_image(orig_im)
+    pixel_values = processor(orig_im, return_tensors="pt").pixel_values
+    return process_image(pixel_values)