Spaces:
Running
Running
attempts to fix mcp
Browse files- __pycache__/app.cpython-313.pyc +0 -0
- __pycache__/globe.cpython-313.pyc +0 -0
- app.py +182 -27
__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 |
-
|
|
|
|
|
|
|
|
|
22 |
|
23 |
-
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
24 |
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
73 |
-
|
74 |
-
|
75 |
-
|
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:
|