update app

app.py

@@ -11,6 +11,7 @@ import spaces
 import torch
 import numpy as np
 from PIL import Image
+import cv2
 
 from transformers import (
     Qwen2VLForConditionalGeneration,
@@ -25,35 +26,28 @@ from gradio.themes.utils import colors, fonts, sizes
 
 # --- Theme and CSS Definition ---
 
-# Define the Thistle color palette
-colors.thistle = colors.Color(
-    name="thistle",
-    c50="#F9F5F9",
-    c100="#F0E8F1",
-    c200="#E7DBE8",
-    c300="#DECEE0",
-    c400="#D2BFD8",
-    c500="#D8BFD8",  # Thistle base color
-    c600="#B59CB7",
-    c700="#927996",
-    c800="#6F5675",
-    c900="#4C3454",
-    c950="#291233",
+# Define the SteelBlue color palette
+colors.steel_blue = colors.Color(
+    name="steel_blue",
+    c50="#EBF3F8",
+    c100="#D3E5F0",
+    c200="#A8CCE1",
+    c300="#7DB3D2",
+    c400="#529AC3",
+    c500="#4682B4",  # SteelBlue base color
+    c600="#3E72A0",
+    c700="#36638C",
+    c800="#2E5378",
+    c900="#264364",
+    c950="#1E3450",
 )
 
-colors.red_gray = colors.Color(
-    name="red_gray",
-    c50="#f7eded", c100="#f5dcdc", c200="#efb4b4", c300="#e78f8f",
-    c400="#d96a6a", c500="#c65353", c600="#b24444", c700="#8f3434",
-    c800="#732d2d", c900="#5f2626", c950="#4d2020",
-)
-
-class ThistleTheme(Soft):
+class SteelBlueTheme(Soft):
     def __init__(
         self,
         *,
         primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.thistle, # Use the new color
+        secondary_hue: colors.Color | str = colors.steel_blue,
         neutral_hue: colors.Color | str = colors.slate,
         text_size: sizes.Size | str = sizes.text_lg,
         font: fonts.Font | str | Iterable[fonts.Font | str] = (
@@ -76,19 +70,19 @@ class ThistleTheme(Soft):
             background_fill_primary_dark="*primary_900",
             body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
             body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
-            button_primary_text_color="black",
+            button_primary_text_color="white",
             button_primary_text_color_hover="white",
-            button_primary_background_fill="linear-gradient(90deg, *secondary_400, *secondary_500)",
-            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
+            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_800)",
+            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",
             button_secondary_text_color="black",
             button_secondary_text_color_hover="white",
             button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
             button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
             button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
             button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
-            slider_color="*secondary_400",
+            slider_color="*secondary_500",
             slider_color_dark="*secondary_600",
             block_title_text_weight="600",
             block_border_width="3px",
@@ -100,7 +94,7 @@ class ThistleTheme(Soft):
         )
 
 # Instantiate the new theme
-thistle_theme = ThistleTheme()
+steel_blue_theme = SteelBlueTheme()
 
 css = """
 #main-title h1 {
@@ -109,56 +103,12 @@ css = """
 #output-title h2 {
     font-size: 2.1em !important;
 }
-:root {
-    --color-grey-50: #f9fafb;
-    --banner-background: var(--secondary-400);
-    --banner-text-color: var(--primary-100);
-    --banner-background-dark: var(--secondary-800);
-    --banner-text-color-dark: var(--primary-100);
-    --banner-chrome-height: calc(16px + 43px);
-    --chat-chrome-height-wide-no-banner: 320px;
-    --chat-chrome-height-narrow-no-banner: 450px;
-    --chat-chrome-height-wide: calc(var(--chat-chrome-height-wide-no-banner) + var(--banner-chrome-height));
-    --chat-chrome-height-narrow: calc(var(--chat-chrome-height-narrow-no-banner) + var(--banner-chrome-height));
-}
-.banner-message { background-color: var(--banner-background); padding: 5px; margin: 0; border-radius: 5px; border: none; }
-.banner-message-text { font-size: 13px; font-weight: bolder; color: var(--banner-text-color) !important; }
-body.dark .banner-message { background-color: var(--banner-background-dark) !important; }
-body.dark .gradio-container .contain .banner-message .banner-message-text { color: var(--banner-text-color-dark) !important; }
-.toast-body { background-color: var(--color-grey-50); }
-.html-container:has(.css-styles) { padding: 0; margin: 0; }
-.css-styles { height: 0; }
-.model-message { text-align: end; }
-.model-dropdown-container { display: flex; align-items: center; gap: 10px; padding: 0; }
-.user-input-container .multimodal-textbox{ border: none !important; }
-.control-button { height: 51px; }
-button.cancel { border: var(--button-border-width) solid var(--button-cancel-border-color); background: var(--button-cancel-background-fill); color: var(--button-cancel-text-color); box-shadow: var(--button-cancel-shadow); }
-button.cancel:hover, .cancel[disabled] { background: var(--button-cancel-background-fill-hover); color: var(--button-cancel-text-color-hover); }
-.opt-out-message { top: 8px; }
-.opt-out-message .html-container, .opt-out-checkbox label { font-size: 14px !important; padding: 0 !important; margin: 0 !important; color: var(--neutral-400) !important; }
-div.block.chatbot { height: calc(100svh - var(--chat-chrome-height-wide)) !important; max-height: 900px !important; }
-div.no-padding { padding: 0 !important; }
-@media (max-width: 1280px) { div.block.chatbot { height: calc(100svh - var(--chat-chrome-height-wide)) !important; } }
-@media (max-width: 1024px) {
-    .responsive-row { flex-direction: column; }
-    .model-message { text-align: start; font-size: 10px !important; }
-    .model-dropdown-container { flex-direction: column; align-items: flex-start; }
-    div.block.chatbot { height: calc(100svh - var(--chat-chrome-height-narrow)) !important; }
-}
-@media (max-width: 400px) {
-    .responsive-row { flex-direction: column; }
-    .model-message { text-align: start; font-size: 10px !important; }
-    .model-dropdown-container { flex-direction: column; align-items: flex-start; }
-    div.block.chatbot { max-height: 360px !important; }
-}
-@media (max-height: 932px) { .chatbot { max-height: 500px !important; } }
-@media (max-height: 1280px) { div.block.chatbot { max-height: 800px !important; } }
 """
 
 # Constants for text generation
+MAX_MAX_NEW_TOKENS = 4096
 DEFAULT_MAX_NEW_TOKENS = 1024
-# Increased max_length to accommodate more complex inputs, especially with multiple images
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -221,7 +171,9 @@ model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 
 
 @spaces.GPU
-def generate_image(model_name: str, text: str, image: Image.Image):
+def generate_image(model_name: str, text: str, image: Image.Image,
+                   max_new_tokens: int, temperature: float, top_p: float,
+                   top_k: int, repetition_penalty: float):
     """
     Generates responses using the selected model for image input.
     Yields raw text and Markdown-formatted text.
@@ -266,7 +218,16 @@ def generate_image(model_name: str, text: str, image: Image.Image):
     ).to(device)
 
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": DEFAULT_MAX_NEW_TOKENS}
+    generation_kwargs = {
+        **inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "do_sample": True,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "repetition_penalty": repetition_penalty,
+    }
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     buffer = ""
@@ -286,12 +247,20 @@ image_examples = [
 
 
 # Create the Gradio Interface
-with gr.Blocks(css=css, theme=thistle_theme) as demo:
+with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     gr.Markdown("# **Multimodal OCR**", elem_id="main-title")
     with gr.Row():
         with gr.Column(scale=2):
             image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
             image_upload = gr.Image(type="pil", label="Upload Image", height=290)
+            
+            with gr.Accordion("Advanced options", open=False):
+                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
+                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.7)
+                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
+                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
+                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.1)
+
             image_submit = gr.Button("Submit", variant="primary")
             gr.Examples(
                 examples=image_examples,
@@ -312,7 +281,7 @@ with gr.Blocks(css=css, theme=thistle_theme) as demo:
 
     image_submit.click(
         fn=generate_image,
-        inputs=[model_choice, image_query, image_upload],
+        inputs=[model_choice, image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
         outputs=[output, markdown_output]
     )