Update app.py

app.py

@@ -5,6 +5,7 @@ import json
 import time
 import asyncio
 from threading import Thread
+from typing import Iterable
 
 import gradio as gr
 import spaces
@@ -23,14 +24,105 @@ from transformers import (
     AutoTokenizer,
 )
 from transformers.image_utils import load_image
+from gradio.themes import Soft
+from gradio.themes.utils import colors, fonts, sizes
+
+colors.light_salmon = colors.Color(
+    name="light_salmon",
+    c50="#FFF9F2",
+    c100="#FFEC C6",
+    c200="#FFD9B3",
+    c300="#FFC6A0",
+    c400="#FFB38D",
+    c500="#FFA07A",
+    c600="#E6906E",
+    c700="#CC8062",
+    c800="#B37056",
+    c900="#99604A",
+    c950="#80503E",
+)
+
+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 LightSalmonTheme(Soft):
+    def __init__(
+        self,
+        *,
+        primary_hue: colors.Color | str = colors.gray,
+        secondary_hue: colors.Color | str = colors.light_salmon, # Use the new color
+        neutral_hue: colors.Color | str = colors.slate,
+        text_size: sizes.Size | str = sizes.text_lg,
+        font: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("Inconsolata"), "Arial", "sans-serif",
+        ),
+        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
+            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
+        ),
+    ):
+        super().__init__(
+            primary_hue=primary_hue,
+            secondary_hue=secondary_hue,
+            neutral_hue=neutral_hue,
+            text_size=text_size,
+            font=font,
+            font_mono=font_mono,
+        )
+        super().set(
+            background_fill_primary="*primary_50",
+            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_hover="white",
+            button_primary_background_fill="linear-gradient(90deg, *secondary_400, *secondary_400)",
+            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_600)",
+            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)",
+            button_cancel_background_fill=f"linear-gradient(90deg, {colors.red_gray.c400}, {colors.red_gray.c500})",
+            button_cancel_background_fill_dark=f"linear-gradient(90deg, {colors.red_gray.c700}, {colors.red_gray.c800})",
+            button_cancel_background_fill_hover=f"linear-gradient(90deg, {colors.red_gray.c500}, {colors.red_gray.c600})",
+            button_cancel_background_fill_hover_dark=f"linear-gradient(90deg, {colors.red_gray.c800}, {colors.red_gray.c900})",
+            button_cancel_text_color="white",
+            button_cancel_text_color_dark="white",
+            button_cancel_text_color_hover="white",
+            button_cancel_text_color_hover_dark="white",
+            slider_color="*secondary_300",
+            slider_color_dark="*secondary_600",
+            block_title_text_weight="600",
+            block_border_width="3px",
+            block_shadow="*shadow_drop_lg",
+            button_primary_shadow="*shadow_drop_lg",
+            button_large_padding="11px",
+            color_accent_soft="*primary_100",
+            block_label_background_fill="*primary_200",
+        )
+
+light_salmon_theme = LightSalmonTheme()
+
+css = """
+#main-title h1 {
+    font-size: 2.3em !important;
+}
+#output-title h2 {
+    font-size: 2.1em !important;
+}
+"""
 
-# Constants for text generation
 MAX_MAX_NEW_TOKENS = 4096
 DEFAULT_MAX_NEW_TOKENS = 2048
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-# Let the environment (e.g., Hugging Face Spaces) determine the device.
-# This avoids conflicts with the CUDA environment setup by the platform.
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
@@ -43,12 +135,7 @@ if torch.cuda.is_available():
     print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
 
 print("Using device:", device)
-# --- Model Loading ---
 
-# To address the warnings, we add `use_fast=False` to ensure we use the
-# processor version the model was originally saved with.
-
-# Load DREX-062225-exp
 MODEL_ID_X = "prithivMLmods/DREX-062225-exp"
 processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True, use_fast=False)
 model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained(
@@ -57,7 +144,6 @@ model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load typhoon-ocr-3b
 MODEL_ID_T = "scb10x/typhoon-ocr-3b"
 processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True, use_fast=False)
 model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
@@ -66,7 +152,6 @@ model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load olmOCR-7B-0225-preview
 MODEL_ID_O = "allenai/olmOCR-7B-0225-preview"
 processor_o = AutoProcessor.from_pretrained(MODEL_ID_O, trust_remote_code=True, use_fast=False)
 model_o = Qwen2VLForConditionalGeneration.from_pretrained(
@@ -75,7 +160,6 @@ model_o = Qwen2VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load Lumian-VLR-7B-Thinking
 MODEL_ID_J = "prithivMLmods/Lumian-VLR-7B-Thinking"
 SUBFOLDER = "think-preview"
 processor_j = AutoProcessor.from_pretrained(MODEL_ID_J, trust_remote_code=True, subfolder=SUBFOLDER, use_fast=False)
@@ -86,20 +170,14 @@ model_j = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load openbmb/MiniCPM-V-4
 MODEL_ID_V4 = 'openbmb/MiniCPM-V-4'
 model_v4 = AutoModel.from_pretrained(
     MODEL_ID_V4,
     trust_remote_code=True,
     torch_dtype=torch.bfloat16,
-    # Using 'sdpa' can sometimes cause issues in certain environments,
-    # letting transformers choose the default is safer.
-    # attn_implementation='sdpa'
 ).eval().to(device)
 tokenizer_v4 = AutoTokenizer.from_pretrained(MODEL_ID_V4, trust_remote_code=True, use_fast=False)
 
-# --- Refactored Model Dictionary ---
-# This simplifies model selection in the generation functions.
 MODELS = {
     "DREX-062225-7B-exp": (processor_x, model_x),
     "Typhoon-OCR-3B": (processor_t, model_t),
@@ -109,15 +187,10 @@ MODELS = {
 
 
 def downsample_video(video_path):
-    """
-    Downsamples the video to evenly spaced frames.
-    Each frame is returned as a PIL image along with its timestamp.
-    """
     vidcap = cv2.VideoCapture(video_path)
     total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
     fps = vidcap.get(cv2.CAP_PROP_FPS)
     frames = []
-    # Use a maximum of 10 frames to avoid excessive memory usage
     frame_indices = np.linspace(0, total_frames - 1, min(total_frames, 10), dtype=int)
     for i in frame_indices:
         vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
@@ -137,14 +210,10 @@ def generate_image(model_name: str, text: str, image: Image.Image,
                    top_p: float = 0.9,
                    top_k: int = 50,
                    repetition_penalty: float = 1.2):
-    """
-    Generates responses using the selected model for image input.
-    """
     if image is None:
         yield "Please upload an image.", "Please upload an image."
         return
 
-    # Handle MiniCPM-V-4 separately due to its different API
     if model_name == "openbmb/MiniCPM-V-4":
         msgs = [{'role': 'user', 'content': [image, text]}]
         try:
@@ -158,13 +227,12 @@ def generate_image(model_name: str, text: str, image: Image.Image,
             yield f"Error: {e}", f"Error: {e}"
         return
 
-    # Use the dictionary for other models
     if model_name not in MODELS:
         yield "Invalid model selected.", "Invalid model selected."
         return
     processor, model = MODELS[model_name]
 
-    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}]
+    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}]
     prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = processor(
         text=[prompt_full], images=[image], return_tensors="pt", padding=True,
@@ -187,9 +255,6 @@ def generate_video(model_name: str, text: str, video_path: str,
                    top_p: float = 0.9,
                    top_k: int = 50,
                    repetition_penalty: float = 1.2):
-    """
-    Generates responses using the selected model for video input.
-    """
     if video_path is None:
         yield "Please upload a video.", "Please upload a video."
         return
@@ -199,14 +264,11 @@ def generate_video(model_name: str, text: str, video_path: str,
         yield "Could not process video.", "Could not process video."
         return
 
-    # Handle MiniCPM-V-4 separately
     if model_name == "openbmb/MiniCPM-V-4":
         images = [frame for frame, ts in frames_with_ts]
-        # For video, the prompt includes the text and then all the image frames
         content = [text] + images
         msgs = [{'role': 'user', 'content': content}]
         try:
-            # The .chat API still takes a single image argument, typically the first frame
             answer = model_v4.chat(
                 image=images[0].convert('RGB'), msgs=msgs, tokenizer=tokenizer_v4,
                 max_new_tokens=max_new_tokens, temperature=temperature,
@@ -217,17 +279,15 @@ def generate_video(model_name: str, text: str, video_path: str,
             yield f"Error: {e}", f"Error: {e}"
         return
 
-    # Use the dictionary for other models
     if model_name not in MODELS:
         yield "Invalid model selected.", "Invalid model selected."
         return
     processor, model = MODELS[model_name]
 
-    # Prepare messages for Qwen-style models
     messages = [{"role": "user", "content": [{"type": "text", "text": text}]}]
     images_for_processor = []
     for frame, timestamp in frames_with_ts:
-        messages[0]["content"].append({"type": "image", "image": frame})
+        messages[0]["content"].insert(0, {"type": "image"})
         images_for_processor.append(frame)
 
     prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
@@ -250,43 +310,34 @@ def generate_video(model_name: str, text: str, video_path: str,
         time.sleep(0.01)
         yield buffer, buffer
 
-
-# Define examples for image and video inference
 image_examples = [
-    ["Describe the safety measures in the image. Conclude (Safe / Unsafe)..", "images/5.jpg"],
-    ["Convert this page to doc [markdown] precisely.", "images/3.png"],
-    ["Convert this page to doc [markdown] precisely.", "images/4.png"],
-    ["Explain the creativity in the image.", "images/6.jpg"],
-    ["Convert this page to doc [markdown] precisely.", "images/1.png"],
-    ["Convert chart to OTSL.", "images/2.png"]
+    ["Describe the safety measures in the image. Conclude (Safe / Unsafe)..", "examples/images/5.jpg"],
+    ["Convert this page to doc [markdown] precisely.", "examples/images/3.png"],
+    ["Convert this page to doc [markdown] precisely.", "examples/images/4.png"],
+    ["Explain the creativity in the image.", "examples/images/6.jpg"],
+    ["Convert this page to doc [markdown] precisely.", "examples/images/1.png"],
+    ["Convert chart to OTSL.", "examples/images/2.png"]
 ]
 
 video_examples = [
-    ["Explain the video in detail.", "videos/2.mp4"],
-    ["Explain the ad in detail.", "videos/1.mp4"]
+    ["Explain the video in detail.", "examples/videos/2.mp4"],
+    ["Explain the ad in detail.", "examples/videos/1.mp4"]
 ]
 
-css = """
-.submit-btn { background-color: #2980b9 !important; color: white !important; }
-.submit-btn:hover { background-color: #3498db !important; }
-.canvas-output { border: 2px solid #4682B4; border-radius: 10px; padding: 20px; }
-"""
-
-# Create the Gradio Interface
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    gr.Markdown("# **[Multimodal VLM Thinking](https://huggingface.co/collections/prithivMLmods/multimodal-implementations-67c9982ea04b39f0608badb0)**")
+with gr.Blocks(theme=light_salmon_theme, css=css) as demo:
+    gr.Markdown("# **Multimodal VLM Thinking**", elem_id="main-title")
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=2):
             with gr.Tabs():
                 with gr.TabItem("Image Inference"):
                     image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
                     image_upload = gr.Image(type="pil", label="Image", height=290)
-                    image_submit = gr.Button("Submit", elem_classes="submit-btn")
+                    image_submit = gr.Button("Submit", variant="primary")
                     gr.Examples(examples=image_examples, inputs=[image_query, image_upload])
                 with gr.TabItem("Video Inference"):
                     video_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
                     video_upload = gr.Video(label="Video", height=290)
-                    video_submit = gr.Button("Submit", elem_classes="submit-btn")
+                    video_submit = gr.Button("Submit", variant="primary")
                     gr.Examples(examples=video_examples, inputs=[video_query, video_upload])
 
             with gr.Accordion("Advanced options", open=False):
@@ -296,12 +347,11 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 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.2)
 
-        with gr.Column():
-            with gr.Column(elem_classes="canvas-output"):
-                gr.Markdown("## Output")
-                output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=5, show_copy_button=True)
-                with gr.Accordion("(Result.md)", open=False):
-                    markdown_output = gr.Markdown(label="(Result.Md)")
+        with gr.Column(scale=3):
+            gr.Markdown("## Output", elem_id="output-title")
+            output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=10, show_copy_button=True)
+            with gr.Accordion("(Result.md)", open=False):
+                markdown_output = gr.Markdown(label="(Result.Md)")
             model_choice = gr.Radio(
                 choices=["Lumian-VLR-7B-Thinking", "openbmb/MiniCPM-V-4", "Typhoon-OCR-3B", "DREX-062225-7B-exp", "olmOCR-7B-0225-preview"],
                 label="Select Model",