syedfaisalabrar commited on
Commit
244c35c
·
verified ·
1 Parent(s): 359b1b4

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -3
app.py CHANGED
@@ -11,7 +11,14 @@ from frontPrompt import main as main_f
11
  model_path = "best.pt"
12
  model = YOLO(model_path)
13
 
14
-
 
 
 
 
 
 
 
15
 
16
  def preprocessing(image):
17
  """Apply three enhancement filters, including brightness reduction, and resize."""
@@ -84,9 +91,9 @@ def crop_image(image, bounding_boxes):
84
  def vision_ai_api(image, doc_type):
85
 
86
  if doc_type == "front":
87
- results = main_f(image)
88
  if doc_type == "back":
89
- results = main_b(image)
90
 
91
  return results
92
 
 
11
  model_path = "best.pt"
12
  model = YOLO(model_path)
13
 
14
+ model = AutoModel.from_pretrained(
15
+ path,
16
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
17
+ # load_in_8bit=True,
18
+ low_cpu_mem_usage=True,
19
+ use_flash_attn=True,
20
+ trust_remote_code=True).eval()
21
+ tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
22
 
23
  def preprocessing(image):
24
  """Apply three enhancement filters, including brightness reduction, and resize."""
 
91
  def vision_ai_api(image, doc_type):
92
 
93
  if doc_type == "front":
94
+ results = main_f(image,model,tokenizer)
95
  if doc_type == "back":
96
+ results = main_b(image,model,tokenizer)
97
 
98
  return results
99