Update app.py

app.py

@@ -29,11 +29,20 @@ from transformers.image_utils import load_image
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
-colors.teal_gray = colors.Color(
-    name="teal_gray",
-    c50="#e8f1f4", c100="#cddde3", c200="#a8c3cf", c300="#7da6b8",
-    c400="#588aa2", c500="#3d6e87", c600="#335b70", c700="#2b495a",
-    c800="#2c5364", c900="#233f4b", c950="#1b323c",
+# Define a new "Thistle" color palette
+colors.thistle = colors.Color(
+    name="thistle",
+    c50="#F9F5F9",
+    c100="#F3ECF4",
+    c200="#E8D9EA",
+    c300="#DCC5E0",
+    c400="#D0B2D6",
+    c500="#D8BFD8",  # Thistle
+    c600="#B8A2B9",
+    c700="#98869A",
+    c800="#796A7C",
+    c900="#5C505D",
+    c950="#423A44",
 )
 
 colors.red_gray = colors.Color(
@@ -43,12 +52,12 @@ colors.red_gray = colors.Color(
     c800="#732d2d", c900="#5f2626", c950="#4d2020",
 )
 
-class Teals(Soft):
+class ThistleTheme(Soft):
     def __init__(
         self,
         *,
         primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.teal_gray,
+        secondary_hue: colors.Color | str = colors.thistle,
         neutral_hue: colors.Color | str = colors.slate,
         text_size: sizes.Size | str = sizes.text_md,
         font: fonts.Font | str | Iterable[fonts.Font | str] = (
@@ -69,12 +78,12 @@ class Teals(Soft):
         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="linear-gradient(135deg, *secondary_200, *secondary_100)",
             body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
-            button_primary_text_color="white",
-            button_primary_text_color_hover="black",
+            button_primary_text_color="*neutral_900",
+            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_300, *secondary_300)",
+            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",
@@ -88,10 +97,7 @@ class Teals(Soft):
             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="*secondary_400",
             slider_color_dark="*secondary_600",
             block_title_text_weight="600",
             block_border_width="3px",
@@ -99,57 +105,16 @@ class Teals(Soft):
             button_primary_shadow="*shadow_drop_lg",
             button_large_padding="11px",
             color_accent_soft="*primary_100",
-            block_label_background_fill="*primary_200",
+            block_label_background_fill="*secondary_200",
         )
 
-teals = Teals()
+thistle_theme = ThistleTheme()
 
 # --- Custom CSS ---
 css = """
 :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; } }
 """
 
 # --- App Constants & Setup ---
@@ -187,8 +152,7 @@ def downsample_video(video_path):
         success, image = vidcap.read()
         if success:
             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            pil_image = Image.fromarray(image)
-            frames.append(pil_image)
+            frames.append(Image.fromarray(image))
     vidcap.release()
     return frames
 
@@ -218,49 +182,58 @@ def load_and_preview_pdf(file_path: Optional[str]) -> Tuple[Optional[Image.Image
         pages = convert_pdf_to_images(file_path)
         if not pages:
             return None, state, '<div style="text-align:center;">Could not load file</div>'
-        state["pages"] = pages
-        state["total_pages"] = len(pages)
-        page_info_html = f'<div style="text-align:center;">Page 1 / {state["total_pages"]}</div>'
-        return pages[0], state, page_info_html
+        state["pages"], state["total_pages"] = pages, len(pages)
+        return pages[0], state, f'<div style="text-align:center;">Page 1 / {state["total_pages"]}</div>'
     except Exception as e:
         return None, state, f'<div style="text-align:center;">Failed to load preview: {e}</div>'
 
 def navigate_pdf_page(direction: str, state: Dict[str, Any]):
     if not state or not state["pages"]:
         return None, state, '<div style="text-align:center;">No file loaded</div>'
-    current_index = state["current_page_index"]
-    total_pages = state["total_pages"]
-    if direction == "prev":
-        new_index = max(0, current_index - 1)
-    elif direction == "next":
-        new_index = min(total_pages - 1, current_index + 1)
-    else:
-        new_index = current_index
-    state["current_page_index"] = new_index
-    image_preview = state["pages"][new_index]
-    page_info_html = f'<div style="text-align:center;">Page {new_index + 1} / {total_pages}</div>'
-    return image_preview, state, page_info_html
+    idx, total = state["current_page_index"], state["total_pages"]
+    new_idx = max(0, idx - 1) if direction == "prev" else min(total - 1, idx + 1)
+    state["current_page_index"] = new_idx
+    return state["pages"][new_idx], state, f'<div style="text-align:center;">Page {new_idx + 1} / {total}</div>'
 
 @spaces.GPU
-def generate_image(text: str, image: Image.Image, max_new_tokens: int = 1024, temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2):
-    if image is None:
-        yield "Please upload an image.", "Please upload an image."
-        return
-    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}]
+def model_stream_response(prompt_text: str, images: list, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float):
+    messages = [{"role": "user", "content": []}]
+    for img in images:
+        messages[0]["content"].append({"type": "image"})
+    messages[0]["content"].append({"type": "text", "text": prompt_text})
+    
     prompt_full = processor_q3vl.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor_q3vl(text=[prompt_full], images=[image], return_tensors="pt", padding=True).to(device)
+    inputs = processor_q3vl(text=[prompt_full], images=images, return_tensors="pt", padding=True).to(device)
+    
     streamer = TextIteratorStreamer(processor_q3vl, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": 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_q3vl.generate, kwargs=generation_kwargs)
     thread.start()
+    
     buffer = ""
     for new_text in streamer:
         buffer += new_text
+        yield buffer.replace("<|im_end|>", ""), buffer.replace("<|im_end|>", "")
         time.sleep(0.01)
-        yield buffer, buffer
 
-@spaces.GPU
-def generate_video(text: str, video_path: str, max_new_tokens: int = 1024, temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2):
+def generate_image(text: str, image: Image.Image, *args):
+    if image is None:
+        yield "Please upload an image.", "Please upload an image."
+        return
+    yield from model_stream_response(text, [image], *args)
+
+def generate_video(text: str, video_path: str, *args):
     if video_path is None:
         yield "Please upload a video.", "Please upload a video."
         return
@@ -268,69 +241,70 @@ def generate_video(text: str, video_path: str, max_new_tokens: int = 1024, tempe
     if not frames:
         yield "Could not process video.", "Could not process video."
         return
-    messages = [{"role": "user", "content": [{"type": "text", "text": text}]}]
-    for frame in frames:
-        messages[0]["content"].insert(0, {"type": "image"})
-    prompt_full = processor_q3vl.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = processor_q3vl(text=[prompt_full], images=frames, return_tensors="pt", padding=True).to(device)
-    streamer = TextIteratorStreamer(processor_q3vl, skip_prompt=True, skip_special_tokens=True)
-    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_q3vl.generate, kwargs=generation_kwargs)
-    thread.start()
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        buffer = buffer.replace("<|im_end|>", "")
-        time.sleep(0.01)
-        yield buffer, buffer
+    yield from model_stream_response(text, frames, *args)
 
-@spaces.GPU
-def generate_pdf(text: str, state: Dict[str, Any], max_new_tokens: int = 2048, temperature: float = 0.6, top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1.2):
+def generate_pdf(text: str, state: Dict[str, Any], *args):
     if not state or not state["pages"]:
         yield "Please upload a PDF file first.", "Please upload a PDF file first."
         return
-    page_images = state["pages"]
+    
     full_response = ""
-    for i, image in enumerate(page_images):
-        page_header = f"--- Page {i+1}/{len(page_images)} ---\n"
+    for i, image in enumerate(state["pages"]):
+        page_header = f"--- Page {i+1}/{len(state['pages'])} ---\n"
         yield full_response + page_header, full_response + page_header
-        messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": text}]}]
-        prompt_full = processor_q3vl.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor_q3vl(text=[prompt_full], images=[image], return_tensors="pt", padding=True).to(device)
-        streamer = TextIteratorStreamer(processor_q3vl, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-        thread = Thread(target=model_q3vl.generate, kwargs=generation_kwargs)
-        thread.start()
-        page_buffer = ""
-        for new_text in streamer:
-            page_buffer += new_text
-            yield full_response + page_header + page_buffer, full_response + page_header + page_buffer
-            time.sleep(0.01)
-        full_response += page_header + page_buffer + "\n\n"
+        
+        # This is a simplified approach. For true streaming of the whole PDF, a more complex logic would be needed.
+        # Here we just get the full response for the page and then append it.
+        final_page_text = ""
+        for page_text, _ in model_stream_response(text, [image], *args):
+             yield full_response + page_header + page_text, full_response + page_header + page_text
+             final_page_text = page_text
+        
+        full_response += page_header + final_page_text + "\n\n"
+
+def generate_caption(image: Image.Image, *args):
+    if image is None:
+        yield "Please upload an image for captioning.", "Please upload an image for captioning."
+        return
+    
+    system_prompt = (
+        "You are an AI assistant that rigorously follows this response protocol: For every input image, "
+        "your primary task is to write a precise caption that captures the essence of the image in clear, "
+        "concise, and contextually accurate language. Along with the caption, provide a structured set of "
+        "attributes describing the visual elements, including details such as objects, people, actions, "
+        "colors, environment, mood, and other notable characteristics. Ensure captions are precise, neutral, "
+        "and descriptive, avoiding unnecessary elaboration or subjective interpretation unless explicitly required. "
+        "Do not reference the rules or instructions in the output; only return the formatted caption, attributes, and class_name."
+    )
+    yield from model_stream_response(system_prompt, [image], *args)
 
 # --- Gradio Interface ---
 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"]]
 video_examples = [["Explain the video in detail.", "videos/2.mp4"]]
-pdf_examples = [["examples/sample-doc.pdf"]]
+caption_examples = [["images/3.png"], ["images/5.jpg"]]
 
-with gr.Blocks(theme=teals, css=css) as demo:
+with gr.Blocks(theme=thistle_theme, css=css) as demo:
     pdf_state = gr.State(value=get_initial_pdf_state())
     gr.Markdown("# **Qwen3-VL-Demo**")
+    
     with gr.Row():
         with gr.Column(scale=2):
             with gr.Tabs():
+                # Image Tab
                 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", variant="primary")
                     gr.Examples(examples=image_examples, inputs=[image_query, image_upload])
 
+                # Video Tab
                 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", variant="primary")
                     gr.Examples(examples=video_examples, inputs=[video_query, video_upload])
 
+                # PDF Tab
                 with gr.TabItem("PDF Inference"):
                     with gr.Row():
                         with gr.Column(scale=1):
@@ -338,19 +312,29 @@ with gr.Blocks(theme=teals, css=css) as demo:
                             pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"])
                             pdf_submit = gr.Button("Submit", variant="primary")
                         with gr.Column(scale=1):
-                            pdf_preview_img = gr.Image(label="PDF Preview", height=290)
+                            pdf_preview_img = gr.Image(label="PDF Preview", height=290, interactive=False)
                             with gr.Row():
                                 prev_page_btn = gr.Button("◀ Previous")
                                 page_info = gr.HTML('<div style="text-align:center;">No file loaded</div>')
                                 next_page_btn = gr.Button("Next ▶")
+                
+                # Caption Tab
+                with gr.TabItem("Caption"):
+                    caption_image_upload = gr.Image(type="pil", label="Image to Caption", height=290)
+                    caption_submit = gr.Button("Generate Caption", variant="primary")
+                    gr.Examples(examples=caption_examples, inputs=[caption_image_upload])
 
+            # Advanced Options
             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.6)
-                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.2)
+                adv_opts = [
+                    gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
+                    gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
+                    gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
+                    gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
+                    gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
+                ]
 
+        # Output Column
         with gr.Column(scale=3):
             gr.Markdown("## Output")
             output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=14, show_copy_button=True)
@@ -388,5 +372,14 @@ with gr.Blocks(theme=teals, css=css) as demo:
         inputs=[pdf_state],
         outputs=[pdf_preview_img, pdf_state, page_info]
     )
+    caption_submit.click(
+        fn=generate_caption,
+        inputs=[caption_image_upload] + adv_opts, 
+        outputs=[output, markdown_output])
+    
+    pdf_upload.change(fn=load_and_preview_pdf, inputs=[pdf_upload], outputs=[pdf_preview_img, pdf_state, page_info])
+    prev_page_btn.click(fn=lambda s: navigate_pdf_page("prev", s), inputs=[pdf_state], outputs=[pdf_preview_img, pdf_state, page_info])
+    next_page_btn.click(fn=lambda s: navigate_pdf_page("next", s), inputs=[pdf_state], outputs=[pdf_preview_img, pdf_state, page_info])
+
 if __name__ == "__main__":
     demo.queue(max_size=50).launch(mcp_server=True, ssr_mode=False, show_error=True)