add sliders for variables , add cpu support

app.py

@@ -20,10 +20,12 @@ import re
 
 model_name = 'ucaslcl/GOT-OCR2_0'
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
 config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-model = model.eval().cuda()
+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)
+model = model.eval().to(device)
 model.config.pad_token_id = tokenizer.eos_token_id
 
 UPLOAD_FOLDER = "./uploads"
@@ -40,7 +42,7 @@ def image_to_base64(image):
 
 
 @spaces.GPU()
-def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
+def process_image(image, task, max_new_tokens, no_repeat_ngram_size, ocr_type=None, ocr_box=None, ocr_color=None):
     if image is None:
         return "Error: No image provided", None, None
     
@@ -49,7 +51,7 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
     result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
     
     try:
-        if isinstance(image, dict):  # If image is from ImageEditor
+        if isinstance(image, dict):
             composite_image = image.get("composite")
             if composite_image is not None:
                 if isinstance(composite_image, np.ndarray):
@@ -68,19 +70,19 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
             return "Error: Unsupported image format", None, None
 
         if task == "Plain Text OCR":
-            res = model.chat(tokenizer, image_path, ocr_type='ocr')
+            res = model.chat(tokenizer, image_path, ocr_type='ocr', max_new_tokens=max_new_tokens, no_repeat_ngram_size=no_repeat_ngram_size)
             return res, None, unique_id
         else:
             if task == "Format Text OCR":
-                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
+                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path, max_new_tokens=max_new_tokens, no_repeat_ngram_size=no_repeat_ngram_size)
             elif task == "Fine-grained OCR (Box)":
-                res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
+                res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path, max_new_tokens=max_new_tokens, no_repeat_ngram_size=no_repeat_ngram_size)
             elif task == "Fine-grained OCR (Color)":
-                res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
+                res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path, max_new_tokens=max_new_tokens, no_repeat_ngram_size=no_repeat_ngram_size)
             elif task == "Multi-crop OCR":
-                res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
+                res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path, max_new_tokens=max_new_tokens, no_repeat_ngram_size=no_repeat_ngram_size)
             elif task == "Render Formatted OCR":
-                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
+                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path, max_new_tokens=max_new_tokens, no_repeat_ngram_size=no_repeat_ngram_size)
             
             if os.path.exists(result_path):
                 with open(result_path, 'r') as f:
@@ -249,6 +251,10 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
                     label="OCR Color",
                     visible=False
                 )
+                with gr.Row():
+                    max_new_tokens_slider = gr.Slider(50, 500, step=10, value=150, label="Max New Tokens")
+                    no_repeat_ngram_size_slider = gr.Slider(1, 10, step=1, value=2, label="No Repeat N-gram Size")
+
                 submit_button = gr.Button("Process")
                 editor_submit_button = gr.Button("Process Edited Image", visible=False)