use model.ocr for 0.8b model

app.py

@@ -125,15 +125,28 @@ def inference(image_input,
     )
 
     # Call model.chat with history
-    response_text, new_state = model.chat(
-        tokenizer,
-        image_input,
-        user_message,
-        max_tiles = int(tile_num),
-        generation_config=generation_config,
-        history=state,
-        return_history=True
-    )
+    if '2b' in model_name.lower():
+        response_text, new_state = model.chat(
+            tokenizer,
+            image_input,
+            user_message,
+            max_tiles = int(tile_num),
+            generation_config=generation_config,
+            history=state,
+            return_history=True
+        )
+        
+
+    if '0.8b' in model_name.lower():
+        response_text, new_state = model.ocr(
+            tokenizer,
+            image_input,
+            user_message,
+            max_tiles = int(tile_num),
+            generation_config=generation_config,
+            history=state,
+            return_history=True
+        )
     
     # update the satet with new_state
     state = new_state
@@ -197,17 +210,28 @@ def regenerate_response(chatbot,
         top_p= float(top_p),
     )
     
-    
+
     # Regenerate the response
-    response_text, new_state = model.chat(
-        tokenizer,
-        image_input,
-        last_user_message,
-        max_tiles = int(tile_num),
-        generation_config=generation_config,
-        history=state,  # Exclude last assistant's response
-        return_history=True
-    )
+    if '2b' in model_name.lower():
+        response_text, new_state = model.chat(
+            tokenizer,
+            image_input,
+            last_user_message,
+            max_tiles = int(tile_num),
+            generation_config=generation_config,
+            history=state,  # Exclude last assistant's response
+            return_history=True
+        )
+    if '0.8b' in model_name.lower():
+        response_text, new_state = model.ocr(
+            tokenizer,
+            image_input,
+            last_user_message,
+            max_tiles = int(tile_num),
+            generation_config=generation_config,
+            history=state,  # Exclude last assistant's response
+            return_history=True
+        )
     
     # Update the state with new_state
     state = new_state
@@ -304,7 +328,7 @@ with gr.Blocks() as demo:
                 interactive=True, 
                 label="Top P")
             max_new_tokens_input = gr.Slider(
-                minimum=0, 
+                minimum=64, 
                 maximum=4096, 
                 step=64, 
                 value=1024,