Update app.py

app.py

@@ -1,209 +1,319 @@
-import os
-import time
-import torch
-import requests
+# Copyright (c) 2025 Team OpthChat.
+#
+# This source code is based on by web_demo_mm.py, by Alibaba Cloud.
+# Licensed under Apache License 2.0
 
-from PIL import Image
-from collections.abc import Iterator
+import os
+import copy
+import re
+from argparse import ArgumentParser
 from threading import Thread
 
 import gradio as gr
-from gradio import FileData
-
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
+import torch
 from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration, TextIteratorStreamer
+
+DEFAULT_CKPT_PATH = 'farrell236/test_model'
+AUTH_TOKEN = os.environ.get("HF_spaces")
+
+def _get_args():
+    parser = ArgumentParser()
+    parser.add_argument('-c',
+                        '--checkpoint-path',
+                        type=str,
+                        default=DEFAULT_CKPT_PATH,
+                        help='Checkpoint name or path, default to %(default)r')
+    parser.add_argument('--cpu-only', action='store_true', help='Run demo with CPU only')
+    parser.add_argument('--flash-attn2',
+                        action='store_true',
+                        default=False,
+                        help='Enable flash_attention_2 when loading the model.')
+    parser.add_argument('--share',
+                        action='store_true',
+                        default=False,
+                        help='Create a publicly shareable link for the interface.')
+    parser.add_argument('--inbrowser',
+                        action='store_true',
+                        default=False,
+                        help='Automatically launch the interface in a new tab on the default browser.')
+    parser.add_argument('--server-port', type=int, default=7860, help='Demo server port.')
+    parser.add_argument('--server-name', type=str, default='0.0.0.0', help='Demo server name.')
+
+    args = parser.parse_args()
+    return args
 
-DESCRIPTION = """\
-# Qwen2.5-VL-32B-Instruct
-"""
-
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-
-auth_token = os.environ.get("HF_spaces")
-
-model_id = 'Qwen/Qwen2.5-VL-3B-Instruct'
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    'farrell236/test_model',
-    use_auth_token=auth_token,
-    # torch_dtype=torch.bfloat16,
-    # attn_implementation="flash_attention_2",
-    device_map="auto"
-)
-processor = AutoProcessor.from_pretrained(model_id)
-
-
-import base64
-from PIL import Image
-import io
-
-# Function to encode the image (scaled down by half)
-def encode_image(image_path, scale=0.25):
-    with Image.open(image_path) as img:
-        # Resize image to half its size
-        new_size = (int(img.width * scale), int(img.height * scale))
-        img = img.resize(new_size)
-
-        # Save the resized image to a bytes buffer
-        buffer = io.BytesIO()
-        img.save(buffer, format="JPEG")  # Change format if needed (e.g., JPEG)
-        buffer.seek(0)
-
-        # Encode to base64
-        return base64.b64encode(buffer.read()).decode('utf-8')
-
-
-def generate(
-    message: str,
-    history: list[dict],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    num_beams: int = 1,
-    repetition_penalty: float = 1.2,
-) -> Iterator[str]:
-
-    txt = message["text"]
-    ext_buffer = f"{txt}"
-
-    messages= []
-    images = []
-
-
-    for i, msg in enumerate(history):
-        if isinstance(msg[0], tuple):
-            print('HIT2', msg[0])
-            messages.append({"role": "user", "content": [
-                {"type": "text", "text": history[i+1][0]},
-                {"type": "image", "image": f"data:image/jpeg;base64,{encode_image(msg[0][0])}"}
-            ]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
-        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
-            # messages are already handled
-            pass
-        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
-            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
-
-    # add current message
-    if len(message["files"]) == 1:
-
-        if isinstance(message["files"][0], str): # examples
-            base64_image = encode_image(message["files"][0])
-        else: # regular input
-            base64_image = encode_image(message["files"][0]["path"])
-        messages.append({"role": "user", "content": [
-            {"type": "text", "text": txt},
-            {"type": "image", "image": f"data:image/jpeg;base64,{base64_image}"}]})
+
+def _load_model_processor(args):
+    if args.cpu_only:
+        device_map = 'cpu'
+    else:
+        device_map = 'auto'
+
+    # Check if flash-attn2 flag is enabled and load model accordingly
+    if args.flash_attn2:
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            args.checkpoint_path,
+            use_auth_token=AUTH_TOKEN,
+            torch_dtype='auto',
+            attn_implementation='flash_attention_2',
+            device_map=device_map)
     else:
-        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
-
-
-    texts = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[texts],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    ) # .to("cuda")
-    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
-
-    generation_kwargs = dict(
-        inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=num_beams,
-        # repetition_penalty=repetition_penalty,
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.checkpoint_path, device_map=device_map)
+
+    processor = AutoProcessor.from_pretrained(args.checkpoint_path)
+    return model, processor
+
+
+def _parse_text(text):
+    lines = text.split('\n')
+    lines = [line for line in lines if line != '']
+    count = 0
+    for i, line in enumerate(lines):
+        if '```' in line:
+            count += 1
+            items = line.split('`')
+            if count % 2 == 1:
+                lines[i] = f'<pre><code class="language-{items[-1]}">'
+            else:
+                lines[i] = '<br></code></pre>'
+        else:
+            if i > 0:
+                if count % 2 == 1:
+                    line = line.replace('`', r'\`')
+                    line = line.replace('<', '&lt;')
+                    line = line.replace('>', '&gt;')
+                    line = line.replace(' ', '&nbsp;')
+                    line = line.replace('*', '&ast;')
+                    line = line.replace('_', '&lowbar;')
+                    line = line.replace('-', '&#45;')
+                    line = line.replace('.', '&#46;')
+                    line = line.replace('!', '&#33;')
+                    line = line.replace('(', '&#40;')
+                    line = line.replace(')', '&#41;')
+                    line = line.replace('$', '&#36;')
+                lines[i] = '<br>' + line
+    text = ''.join(lines)
+    return text
+
+
+def _remove_image_special(text):
+    text = text.replace('<ref>', '').replace('</ref>', '')
+    return re.sub(r'<box>.*?(</box>|$)', '', text)
+
+
+def _is_video_file(filename):
+    video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.mpeg']
+    return any(filename.lower().endswith(ext) for ext in video_extensions)
+
+
+def _gc():
+    import gc
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _transform_messages(original_messages):
+    transformed_messages = []
+    for message in original_messages:
+        new_content = []
+        for item in message['content']:
+            if 'image' in item:
+                new_item = {'type': 'image', 'image': item['image']}
+            elif 'text' in item:
+                new_item = {'type': 'text', 'text': item['text']}
+            elif 'video' in item:
+                new_item = {'type': 'video', 'video': item['video']}
+            else:
+                continue
+            new_content.append(new_item)
+
+        new_message = {'role': message['role'], 'content': new_content}
+        transformed_messages.append(new_message)
+
+    return transformed_messages
+
+
+def _launch_demo(args, model, processor):
+
+    def call_local_model(model, processor, messages,
+                         max_tokens=1024, temperature=0.6,
+                         top_p=0.9, top_k=50,
+                         repetition_penalty=1.2):
+
+        messages = _transform_messages(messages)
+
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors='pt')
+        inputs = inputs.to(model.device)
+
+        tokenizer = processor.tokenizer
+        streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+
+        gen_kwargs = {'max_new_tokens': max_tokens,
+                      'streamer': streamer,
+                      'temperature': temperature,
+                      'top_p': top_p,
+                      'top_k': top_k,
+                      'repetition_penalty': repetition_penalty,
+                      **inputs}
+
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+
+        generated_text = ''
+        for new_text in streamer:
+            generated_text += new_text
+            yield generated_text
+
+    def create_predict_fn():
+
+        def predict(_chatbot, task_history,
+                    max_tokens, temperature, top_p, top_k, repetition_penalty):
+            nonlocal model, processor
+            chat_query = _chatbot[-1][0]
+            query = task_history[-1][0]
+            if len(chat_query) == 0:
+                _chatbot.pop()
+                task_history.pop()
+                return _chatbot
+            print('User: ' + _parse_text(query))
+            history_cp = copy.deepcopy(task_history)
+            full_response = ''
+            messages = []
+            content = []
+            for q, a in history_cp:
+                if isinstance(q, (tuple, list)):
+                    if _is_video_file(q[0]):
+                        content.append({'video': f'file://{q[0]}'})
+                    else:
+                        content.append({'image': f'file://{q[0]}'})
+                else:
+                    content.append({'text': q})
+                    messages.append({'role': 'user', 'content': content})
+                    messages.append({'role': 'assistant', 'content': [{'text': a}]})
+                    content = []
+            messages.pop()
+
+            for response in call_local_model(model, processor, messages):
+                _chatbot[-1] = (_parse_text(chat_query), _remove_image_special(_parse_text(response)))
+
+                yield _chatbot
+                full_response = _parse_text(response)
+
+            task_history[-1] = (query, full_response)
+            print('Qwen-VL-Chat: ' + _parse_text(full_response))
+            yield _chatbot
+
+        return predict
+
+    def create_regenerate_fn():
+
+        def regenerate(_chatbot, task_history):
+            nonlocal model, processor
+            if not task_history:
+                return _chatbot
+            item = task_history[-1]
+            if item[1] is None:
+                return _chatbot
+            task_history[-1] = (item[0], None)
+            chatbot_item = _chatbot.pop(-1)
+            if chatbot_item[0] is None:
+                _chatbot[-1] = (_chatbot[-1][0], None)
+            else:
+                _chatbot.append((chatbot_item[0], None))
+            _chatbot_gen = predict(_chatbot, task_history)
+            for _chatbot in _chatbot_gen:
+                yield _chatbot
+
+        return regenerate
+
+    predict = create_predict_fn()
+    regenerate = create_regenerate_fn()
+
+    def add_text(history, task_history, text):
+        task_text = text
+        history = history if history is not None else []
+        task_history = task_history if task_history is not None else []
+        history = history + [(_parse_text(text), None)]
+        task_history = task_history + [(task_text, None)]
+        return history, task_history, ''
+
+    def add_file(history, task_history, file):
+        history = history if history is not None else []
+        task_history = task_history if task_history is not None else []
+        history = history + [((file.name,), None)]
+        task_history = task_history + [((file.name,), None)]
+        return history, task_history
+
+    def reset_user_input():
+        return gr.update(value='')
+
+    def reset_state(_chatbot, task_history):
+        task_history.clear()
+        _chatbot.clear()
+        _gc()
+        return []
+
+    with gr.Blocks() as demo:
+        gr.Markdown("""\
+<p align="center"><img src="https://home.mmc.edu/wp-content/uploads/2017/10/nih-logo-color.png" style="height: 80px"/><p>
+<center><font size=6>Qwen2.5-VL (model_a) for OpthChat</center>
+<center><font size=4></center>
+<center><font size=4></center>
+<center><font size=4></center>
+        """)
+
+        chatbot = gr.Chatbot(label='Qwen2.5-VL', elem_classes='control-height', height=500)
+
+        with gr.Accordion("Generation Parameters", open=False):
+            max_tokens = gr.Slider(64, 4096, value=512, step=64, label="Max Tokens")
+            temperature = gr.Slider(0.0, 2.0, value=0.6, step=0.1, label="Temperature")
+            top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)")
+            top_k = gr.Slider(0, 100, value=50, step=1, label="Top-k")
+            repetition_penalty = gr.Slider(0.5, 2.0, value=1.2, step=0.1, label="Repetition Penalty")
+
+        query = gr.Textbox(lines=2, label='Input')
+        task_history = gr.State([])
+
+        with gr.Row():
+            addfile_btn = gr.UploadButton('📁 Upload', file_types=['image', 'video'])
+            submit_btn = gr.Button('🚀 Submit')
+            regen_btn = gr.Button('♻️️ Regenerate')
+            empty_bin = gr.Button('🧹 Clear History')
+
+        submit_btn.click(add_text,
+                         [chatbot, task_history, query],
+                         [chatbot, task_history]).then(predict,
+                                                       [chatbot, task_history, max_tokens,
+                                                        temperature, top_p, top_k, repetition_penalty],
+                                                       [chatbot], show_progress=True)
+        submit_btn.click(reset_user_input, [], [query])
+        empty_bin.click(reset_state, [chatbot, task_history], [chatbot], show_progress=True)
+        regen_btn.click(regenerate, [chatbot, task_history], [chatbot], show_progress=True)
+        addfile_btn.upload(add_file, [chatbot, task_history, addfile_btn], [chatbot, task_history], show_progress=True)
+
+        gr.Markdown("""\
+<font size=2>Note: This demo is governed by the original license of Qwen2.5-VL,
+WebUI based on [Qwen2.5-VL](https://github.com/QwenLM/Qwen2.5-VL/blob/main/web_demo_mm.py). 
+Developed by Alibaba Cloud, modified by Team OpthChat
+        """)
+
+    demo.queue().launch(
+        share=args.share,
+        inbrowser=args.inbrowser,
+        server_port=args.server_port,
+        server_name=args.server_name,
     )
-    generated_text = ""
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    buffer = ""
-
-    for new_text in streamer:
-        buffer += new_text
-        generated_text_without_prompt = buffer
-        time.sleep(0.01)
-        yield buffer
-
-demo = gr.ChatInterface(fn=generate, title="Multimodal Qwen", examples=[
-    [{"text": """\
-    You are a highly experienced ophthalmologist specializing in retinal diseases.
-    You will be shown a color fundus photograph of a patient's eye.
-    Your task is to identify key retinal features and return a structured response.
-    You must only respond in JSON format using the following fields:
-    - ADVAMD: 1 if advanced age-related macular degeneration is present, otherwise 0
-    - PIG: 1 if abnormal pigmentary is present, otherwise 0
-    - DRUS: 0 if no drusen or small drusen, 1 if intermediate or medium drusen, 2 if large drusen
-    - RPD: 1 if reticular pseudodrusen are present, otherwise 0
-    - NVAMD: 1 if neovascular AMD is present, otherwise 0
-    - GA: 1 if geographic atrophy is present, otherwise 0
-
-    Do not include any explanation, just return the JSON object.
-
-    Please assess this fundus image and return your findings in the specified JSON format.""",
-    "files":["./examples/ret-hem250-304.jpg"]},
-    1024],
-    ],
-        textbox=gr.MultimodalTextbox(),
-        additional_inputs = [
-            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="Beam Search",
-                minimum=1,
-                maximum=1,
-                step=1,
-                value=1,
-            ),
-            gr.Slider(
-                label="Repetition penalty",
-                minimum=1.0,
-                maximum=2.0,
-                step=0.05,
-                value=1.2,
-            ),
-        ],
-      cache_examples=False,
-      description=DESCRIPTION,
-      stop_btn="Stop Generation",
-      fill_height=True,
-    multimodal=True)
-
-if __name__ == "__main__":
-    demo.launch()
+
+
+def main():
+    args = _get_args()
+    model, processor = _load_model_processor(args)
+    _launch_demo(args, model, processor)
+
+
+if __name__ == '__main__':
+    main()