Tonic commited on
Commit
a33adcb
Β·
unverified Β·
1 Parent(s): 79ffa77

attempts to fix mcp

Browse files
__pycache__/app.cpython-313.pyc ADDED
Binary file (24.1 kB). View file
 
__pycache__/globe.cpython-313.pyc ADDED
Binary file (2.84 kB). View file
 
app.py CHANGED
@@ -16,17 +16,85 @@ import time
16
  import shutil
17
  import cv2
18
  import re
 
19
 
 
 
 
 
 
 
 
 
 
 
20
 
21
- model_name = 'ucaslcl/GOT-OCR2_0'
 
 
 
 
22
 
23
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
24
 
25
- tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
26
- config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
27
- model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map=device, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
28
- model = model.eval().to(device)
29
- model.config.pad_token_id = tokenizer.eos_token_id
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
30
 
31
  UPLOAD_FOLDER = "./uploads"
32
  RESULTS_FOLDER = "./results"
@@ -40,6 +108,84 @@ def image_to_base64(image):
40
  image.save(buffered, format="PNG")
41
  return base64.b64encode(buffered.getvalue()).decode()
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
 
44
  @spaces.GPU()
45
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
@@ -69,27 +215,36 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
69
  else:
70
  return "Error: Unsupported image format", None, None
71
 
72
- if task == "Plain Text OCR":
73
- res = model.chat(tokenizer, image_path, ocr_type='ocr')
74
- return res, None, unique_id
75
- else:
76
- if task == "Format Text OCR":
77
- res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
78
- elif task == "Fine-grained OCR (Box)":
79
- res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
80
- elif task == "Fine-grained OCR (Color)":
81
- res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
82
- elif task == "Multi-crop OCR":
83
- res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
84
- elif task == "Render Formatted OCR":
85
- res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
86
-
87
- if os.path.exists(result_path):
88
- with open(result_path, 'r') as f:
89
- html_content = f.read()
90
- return res, html_content, unique_id
91
- else:
92
  return res, None, unique_id
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93
  except Exception as e:
94
  return f"Error: {str(e)}", None, None
95
  finally:
 
16
  import shutil
17
  import cv2
18
  import re
19
+ import warnings
20
 
21
+ # Try to import spaces module for ZeroGPU compatibility
22
+ try:
23
+ import spaces
24
+ SPACES_AVAILABLE = True
25
+ except ImportError:
26
+ SPACES_AVAILABLE = False
27
+ # Create a dummy decorator for local development
28
+ def dummy_gpu_decorator(func):
29
+ return func
30
+ spaces = type('spaces', (), {'GPU': dummy_gpu_decorator})()
31
 
32
+ # Suppress specific warnings that are known issues with GOT-OCR
33
+ warnings.filterwarnings("ignore", message="The attention mask and the pad token id were not set")
34
+ warnings.filterwarnings("ignore", message="Setting `pad_token_id` to `eos_token_id`")
35
+ warnings.filterwarnings("ignore", message="The attention mask is not set and cannot be inferred")
36
+ warnings.filterwarnings("ignore", message="The `seen_tokens` attribute is deprecated")
37
 
 
38
 
39
+ def initialize_model_safely():
40
+ """
41
+ Safely initialize the GOT-OCR model with proper error handling for ZeroGPU
42
+ """
43
+ model_name = 'ucaslcl/GOT-OCR2_0'
44
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
45
+
46
+ try:
47
+ # Initialize tokenizer with proper settings
48
+ tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
49
+
50
+ # Set pad token properly
51
+ if tokenizer.pad_token is None:
52
+ tokenizer.pad_token = tokenizer.eos_token
53
+
54
+ config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
55
+
56
+ model = AutoModel.from_pretrained(
57
+ 'ucaslcl/GOT-OCR2_0',
58
+ trust_remote_code=True,
59
+ low_cpu_mem_usage=True,
60
+ device_map=device,
61
+ use_safetensors=True,
62
+ pad_token_id=tokenizer.eos_token_id,
63
+ use_cache=True,
64
+ torch_dtype=torch.float16 if device == 'cuda' else torch.float32
65
+ )
66
+
67
+ model = model.eval().to(device)
68
+ model.config.pad_token_id = tokenizer.eos_token_id
69
+
70
+ # Ensure the model has proper tokenizer settings
71
+ if hasattr(model, 'config'):
72
+ model.config.pad_token_id = tokenizer.eos_token_id
73
+ model.config.eos_token_id = tokenizer.eos_token_id
74
+
75
+ return model, tokenizer
76
+
77
+ except Exception as e:
78
+ print(f"Error initializing model: {str(e)}")
79
+ # Fallback initialization
80
+ try:
81
+ tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
82
+ if tokenizer.pad_token is None:
83
+ tokenizer.pad_token = tokenizer.eos_token
84
+
85
+ model = AutoModel.from_pretrained(
86
+ 'ucaslcl/GOT-OCR2_0',
87
+ trust_remote_code=True,
88
+ low_cpu_mem_usage=True,
89
+ device_map=device,
90
+ use_safetensors=True
91
+ )
92
+ model = model.eval().to(device)
93
+ return model, tokenizer
94
+ except Exception as fallback_error:
95
+ raise Exception(f"Failed to initialize model: {str(e)}. Fallback also failed: {str(fallback_error)}")
96
+
97
+ model, tokenizer = initialize_model_safely()
98
 
99
  UPLOAD_FOLDER = "./uploads"
100
  RESULTS_FOLDER = "./results"
 
108
  image.save(buffered, format="PNG")
109
  return base64.b64encode(buffered.getvalue()).decode()
110
 
111
+ def safe_model_chat(model, tokenizer, image_path, **kwargs):
112
+ """
113
+ Safe wrapper for model.chat to handle DynamicCache and other compatibility issues
114
+ Optimized for ZeroGPU environments
115
+ """
116
+ try:
117
+ # First attempt: normal call
118
+ return model.chat(tokenizer, image_path, **kwargs)
119
+ except AttributeError as e:
120
+ if "get_max_length" in str(e):
121
+ # Try to fix the cache issue by clearing it
122
+ try:
123
+ if hasattr(model, 'clear_cache'):
124
+ model.clear_cache()
125
+ # Retry the call
126
+ return model.chat(tokenizer, image_path, **kwargs)
127
+ except:
128
+ # If still failing, try with different parameters
129
+ try:
130
+ # Remove any cache-related parameters
131
+ kwargs_copy = kwargs.copy()
132
+ if 'use_cache' in kwargs_copy:
133
+ del kwargs_copy['use_cache']
134
+ return model.chat(tokenizer, image_path, **kwargs_copy)
135
+ except:
136
+ raise Exception("Model compatibility issue: DynamicCache error. Please try again.")
137
+ else:
138
+ raise e
139
+ except Exception as e:
140
+ # Handle other potential issues
141
+ if "attention_mask" in str(e).lower():
142
+ # Try to handle attention mask issues
143
+ try:
144
+ return model.chat(tokenizer, image_path, **kwargs)
145
+ except:
146
+ raise Exception(f"Attention mask error: {str(e)}")
147
+ else:
148
+ raise e
149
+
150
+ def safe_model_chat_crop(model, tokenizer, image_path, **kwargs):
151
+ """
152
+ Safe wrapper for model.chat_crop to handle DynamicCache and other compatibility issues
153
+ Optimized for ZeroGPU environments
154
+ """
155
+ try:
156
+ # First attempt: normal call
157
+ return model.chat_crop(tokenizer, image_path, **kwargs)
158
+ except AttributeError as e:
159
+ if "get_max_length" in str(e):
160
+ # Try to fix the cache issue by clearing it
161
+ try:
162
+ if hasattr(model, 'clear_cache'):
163
+ model.clear_cache()
164
+ # Retry the call
165
+ return model.chat_crop(tokenizer, image_path, **kwargs)
166
+ except:
167
+ # If still failing, try with different parameters
168
+ try:
169
+ # Remove any cache-related parameters
170
+ kwargs_copy = kwargs.copy()
171
+ if 'use_cache' in kwargs_copy:
172
+ del kwargs_copy['use_cache']
173
+ return model.chat_crop(tokenizer, image_path, **kwargs_copy)
174
+ except:
175
+ raise Exception("Model compatibility issue: DynamicCache error. Please try again.")
176
+ else:
177
+ raise e
178
+ except Exception as e:
179
+ # Handle other potential issues
180
+ if "attention_mask" in str(e).lower():
181
+ # Try to handle attention mask issues
182
+ try:
183
+ return model.chat_crop(tokenizer, image_path, **kwargs)
184
+ except:
185
+ raise Exception(f"Attention mask error: {str(e)}")
186
+ else:
187
+ raise e
188
+
189
 
190
  @spaces.GPU()
191
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
 
215
  else:
216
  return "Error: Unsupported image format", None, None
217
 
218
+ # Wrap model calls in try-except to handle DynamicCache errors
219
+ try:
220
+ if task == "Plain Text OCR":
221
+ res = safe_model_chat(model, tokenizer, image_path, ocr_type='ocr')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
222
  return res, None, unique_id
223
+ else:
224
+ if task == "Format Text OCR":
225
+ res = safe_model_chat(model, tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
226
+ elif task == "Fine-grained OCR (Box)":
227
+ res = safe_model_chat(model, tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
228
+ elif task == "Fine-grained OCR (Color)":
229
+ res = safe_model_chat(model, tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
230
+ elif task == "Multi-crop OCR":
231
+ res = safe_model_chat_crop(model, tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
232
+ elif task == "Render Formatted OCR":
233
+ res = safe_model_chat(model, tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
234
+
235
+ if os.path.exists(result_path):
236
+ with open(result_path, 'r') as f:
237
+ html_content = f.read()
238
+ return res, html_content, unique_id
239
+ else:
240
+ return res, None, unique_id
241
+ except AttributeError as e:
242
+ if "get_max_length" in str(e):
243
+ # Handle DynamicCache compatibility issue
244
+ return "Error: Model compatibility issue detected. Please try again or contact support.", None, None
245
+ else:
246
+ raise e
247
+
248
  except Exception as e:
249
  return f"Error: {str(e)}", None, None
250
  finally: