Tonic commited on
Commit
cc155bb
Β·
unverified Β·
1 Parent(s): 17e5ee4

continues to remove cache parameter

Browse files
Files changed (1) hide show
  1. app.py +48 -155
app.py CHANGED
@@ -248,113 +248,25 @@ def image_to_base64(image):
248
  image.save(buffered, format="PNG")
249
  return base64.b64encode(buffered.getvalue()).decode()
250
 
251
- def safe_model_chat(model, tokenizer, image_path, **kwargs):
252
  """
253
- Safe wrapper for model.chat to handle DynamicCache and other compatibility issues
254
- Optimized for ZeroGPU environments
255
  """
256
- try:
257
- # First attempt: normal call
258
- return model.chat(tokenizer, image_path, **kwargs)
259
- except AttributeError as e:
260
- if "get_max_length" in str(e):
261
- # Try to fix the cache issue by clearing it
262
- try:
263
- # Clear any existing cache
264
- if hasattr(model, 'clear_cache'):
265
- model.clear_cache()
266
- elif hasattr(model, '_clear_cache'):
267
- model._clear_cache()
268
-
269
- # Try to clear cache from transformers
270
- try:
271
- from transformers import GenerationConfig
272
- if hasattr(GenerationConfig, 'clear_cache'):
273
- GenerationConfig.clear_cache()
274
- except:
275
- pass
276
-
277
- # Retry the call
278
- return model.chat(tokenizer, image_path, **kwargs)
279
- except:
280
- # If still failing, try with different parameters
281
- try:
282
- # Remove any cache-related parameters
283
- kwargs_copy = kwargs.copy()
284
- for key in list(kwargs_copy.keys()):
285
- if 'cache' in key.lower():
286
- del kwargs_copy[key]
287
- return model.chat(tokenizer, image_path, **kwargs_copy)
288
- except:
289
- raise Exception("Model compatibility issue: DynamicCache error. Please try again.")
290
- else:
291
- raise e
292
- except Exception as e:
293
- # Handle other potential issues
294
- if "attention_mask" in str(e).lower():
295
- # Try to handle attention mask issues
296
- try:
297
- return model.chat(tokenizer, image_path, **kwargs)
298
- except:
299
- raise Exception(f"Attention mask error: {str(e)}")
300
- else:
301
- raise e
302
-
303
- def safe_model_chat_crop(model, tokenizer, image_path, **kwargs):
304
- """
305
- Safe wrapper for model.chat_crop to handle DynamicCache and other compatibility issues
306
- Optimized for ZeroGPU environments
307
- """
308
- try:
309
- # First attempt: normal call
310
- return model.chat_crop(tokenizer, image_path, **kwargs)
311
- except AttributeError as e:
312
- if "get_max_length" in str(e):
313
- # Try to fix the cache issue by clearing it
314
- try:
315
- # Clear any existing cache
316
- if hasattr(model, 'clear_cache'):
317
- model.clear_cache()
318
- elif hasattr(model, '_clear_cache'):
319
- model._clear_cache()
320
-
321
- # Try to clear cache from transformers
322
- try:
323
- from transformers import GenerationConfig
324
- if hasattr(GenerationConfig, 'clear_cache'):
325
- GenerationConfig.clear_cache()
326
- except:
327
- pass
328
-
329
- # Retry the call
330
- return model.chat_crop(tokenizer, image_path, **kwargs)
331
- except:
332
- # If still failing, try with different parameters
333
- try:
334
- # Remove any cache-related parameters
335
- kwargs_copy = kwargs.copy()
336
- for key in list(kwargs_copy.keys()):
337
- if 'cache' in key.lower():
338
- del kwargs_copy[key]
339
- return model.chat_crop(tokenizer, image_path, **kwargs_copy)
340
- except:
341
- raise Exception("Model compatibility issue: DynamicCache error. Please try again.")
342
- else:
343
- raise e
344
- except Exception as e:
345
- # Handle other potential issues
346
- if "attention_mask" in str(e).lower():
347
- # Try to handle attention mask issues
348
- try:
349
- return model.chat_crop(tokenizer, image_path, **kwargs)
350
- except:
351
- raise Exception(f"Attention mask error: {str(e)}")
352
- else:
353
- raise e
354
-
355
 
356
  @spaces.GPU()
357
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
 
 
 
358
  if image is None:
359
  return "Error: No image provided", None, None
360
 
@@ -381,61 +293,22 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
381
  else:
382
  return "Error: Unsupported image format", None, None
383
 
384
- # Wrap model calls in try-except to handle DynamicCache errors
385
  try:
386
  if task == "Plain Text OCR":
387
- # Use cache manager for safer calls
388
- try:
389
- res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type='ocr')
390
- except:
391
- try:
392
- # Fallback to direct call
393
- res = cache_manager.direct_call('chat', tokenizer, image_path, ocr_type='ocr')
394
- except:
395
- # Final fallback to legacy call
396
- res = cache_manager.legacy_call('chat', tokenizer, image_path, ocr_type='ocr')
397
  return res, None, unique_id
398
  else:
399
  if task == "Format Text OCR":
400
- try:
401
- res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
402
- except:
403
- try:
404
- res = cache_manager.direct_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
405
- except:
406
- res = cache_manager.legacy_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
407
  elif task == "Fine-grained OCR (Box)":
408
- try:
409
- res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
410
- except:
411
- try:
412
- res = cache_manager.direct_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
413
- except:
414
- res = cache_manager.legacy_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
415
  elif task == "Fine-grained OCR (Color)":
416
- try:
417
- res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
418
- except:
419
- try:
420
- res = cache_manager.direct_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
421
- except:
422
- res = cache_manager.legacy_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
423
  elif task == "Multi-crop OCR":
424
- try:
425
- res = cache_manager.safe_call('chat_crop', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
426
- except:
427
- try:
428
- res = cache_manager.direct_call('chat_crop', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
429
- except:
430
- res = cache_manager.legacy_call('chat_crop', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
431
  elif task == "Render Formatted OCR":
432
- try:
433
- res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
434
- except:
435
- try:
436
- res = cache_manager.direct_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
437
- except:
438
- res = cache_manager.legacy_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
439
 
440
  if os.path.exists(result_path):
441
  with open(result_path, 'r') as f:
@@ -443,19 +316,39 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
443
  return res, html_content, unique_id
444
  else:
445
  return res, None, unique_id
446
- except AttributeError as e:
447
- if "get_max_length" in str(e):
448
- # Handle DynamicCache compatibility issue
449
- return "Error: Model compatibility issue detected. Please try again or contact support.", None, None
450
- else:
451
- raise e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
452
 
453
  except Exception as e:
454
  return f"Error: {str(e)}", None, None
455
  finally:
456
  if os.path.exists(image_path):
457
  os.remove(image_path)
458
-
459
  def update_image_input(task):
460
  if task == "Fine-grained OCR (Color)":
461
  return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
 
248
  image.save(buffered, format="PNG")
249
  return base64.b64encode(buffered.getvalue()).decode()
250
 
251
+ def direct_model_call(model, method_name, *args, **kwargs):
252
  """
253
+ Direct model call without any cache-related parameters
 
254
  """
255
+ # Create a clean kwargs dict without any cache-related parameters
256
+ clean_kwargs = {}
257
+ for key, value in kwargs.items():
258
+ if 'cache' not in key.lower():
259
+ clean_kwargs[key] = value
260
+
261
+ # Get the method and call it directly
262
+ method = getattr(model, method_name)
263
+ return method(*args, **clean_kwargs)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
264
 
265
  @spaces.GPU()
266
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
267
+ """
268
+ Process image with OCR using ZeroGPU-compatible approach
269
+ """
270
  if image is None:
271
  return "Error: No image provided", None, None
272
 
 
293
  else:
294
  return "Error: Unsupported image format", None, None
295
 
296
+ # Use direct model calls without any cache management
297
  try:
298
  if task == "Plain Text OCR":
299
+ res = direct_model_call(model, 'chat', tokenizer, image_path, ocr_type='ocr')
 
 
 
 
 
 
 
 
 
300
  return res, None, unique_id
301
  else:
302
  if task == "Format Text OCR":
303
+ res = direct_model_call(model, 'chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
 
 
 
 
 
 
304
  elif task == "Fine-grained OCR (Box)":
305
+ res = direct_model_call(model, 'chat', tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
 
 
 
 
 
 
306
  elif task == "Fine-grained OCR (Color)":
307
+ res = direct_model_call(model, 'chat', tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
 
 
 
 
 
 
308
  elif task == "Multi-crop OCR":
309
+ res = direct_model_call(model, 'chat_crop', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
 
 
 
 
 
 
310
  elif task == "Render Formatted OCR":
311
+ res = direct_model_call(model, 'chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
 
 
 
 
 
 
312
 
313
  if os.path.exists(result_path):
314
  with open(result_path, 'r') as f:
 
316
  return res, html_content, unique_id
317
  else:
318
  return res, None, unique_id
319
+ except Exception as e:
320
+ # If direct call fails, try with cache manager as fallback
321
+ try:
322
+ if task == "Plain Text OCR":
323
+ res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type='ocr')
324
+ return res, None, unique_id
325
+ else:
326
+ if task == "Format Text OCR":
327
+ res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
328
+ elif task == "Fine-grained OCR (Box)":
329
+ res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
330
+ elif task == "Fine-grained OCR (Color)":
331
+ res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
332
+ elif task == "Multi-crop OCR":
333
+ res = cache_manager.safe_call('chat_crop', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
334
+ elif task == "Render Formatted OCR":
335
+ res = cache_manager.safe_call('chat', tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
336
+
337
+ if os.path.exists(result_path):
338
+ with open(result_path, 'r') as f:
339
+ html_content = f.read()
340
+ return res, html_content, unique_id
341
+ else:
342
+ return res, None, unique_id
343
+ except Exception as fallback_error:
344
+ return f"Error: {str(fallback_error)}", None, None
345
 
346
  except Exception as e:
347
  return f"Error: {str(e)}", None, None
348
  finally:
349
  if os.path.exists(image_path):
350
  os.remove(image_path)
351
+
352
  def update_image_input(task):
353
  if task == "Fine-grained OCR (Color)":
354
  return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)