Update app.py

app.py

@@ -1,64 +1,233 @@
+# Copyright (c) Kuaishou.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import os
+import base64
+import numpy as np
+os.system('pip install huggingface_hub gradio openai keye_vl_utils -U')
+
+from argparse import ArgumentParser
+from pathlib import Path
+
+import copy
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+import os
+import re
+import tempfile
+from openai import OpenAI
+from PIL import Image
+from io import BytesIO
+from keye_vl_utils import fetch_video, process_vision_info
+
+
+def _get_args():
+    parser = ArgumentParser()
+    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="127.0.0.1",
+                        help="Demo server name.")
+    args = parser.parse_args()
+    return args
+
+
+def _parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line != ""]
+    count = 0
+    for i, line in enumerate(lines):
+        line = line.replace("<think>", "**思考过程开始**：\n")
+        line = line.replace("</think>", "\n**思考过程结束**\n")
+        line = line.replace("<answer>", "**回答开始**：\n")
+        line = line.replace("</answer>", "\n**回答结束**\n")
+        line = line.replace("<analysis>", "**分析开始**：\n")
+        line = line.replace("</analysis>", "\n**分析结束**\n")
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                lines[i] = f'<pre><code class="language-{items[-1]}">'
+            else:
+                lines[i] = f"<br></code></pre>"
+        else:
+            if True or 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 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 video_processor(video_path):
+    video_inputs, video_sample_fps = fetch_video({"video": video_path, "max_frames": 8, "max_pixels": 256*28*28}, return_video_sample_fps=True)
+    video_input = video_inputs.permute(0, 2, 3, 1).numpy().astype(np.uint8)
+
+    # encode image with base64
+    base64_frames = []
+    for frame in video_input:
+        img = Image.fromarray(frame)
+        output_buffer = BytesIO()
+        img.save(output_buffer, format="jpeg")
+        byte_data = output_buffer.getvalue()
+        base64_str = base64.b64encode(byte_data).decode("utf-8")
+        base64_frames.append(base64_str)
+
+    video_info = {
+        "type": "video_url",
+        "video_url": {"url": f"data:video/jpeg;base64,{','.join(base64_frames)}"}
+    }
+    return video_info
+
+
+def _launch_demo(args):
+    EP_URL = os.getenv("ENDPOINT_URL")
+    HF_TOKEN = os.getenv("HF_TOKEN")
+    client = OpenAI(
+		base_url=f"{EP_URL}/v1/",
+		api_key=HF_TOKEN
+	)
+
+    def predict(_chatbot, task_history):
+        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]):
+                    video_info = video_processor(q[0])
+                    content.append(video_info)
+                else:
+                    # convert image to base64
+                    with open(q[0], 'rb') as img_file:
+                        img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
+                    content.append({'type': 'image_url', 'image_url': {"url": f"data:image/jpeg;base64,{img_base64}"}})
+            else:
+                content.append({'type': 'text', 'text': q})
+                messages.append({'role': 'user', 'content': content})
+                messages.append({'role': 'assistant', 'content': [{'type': 'text', 'text': a}]})
+                content = []
+        messages.pop()
+        
+        responses = client.chat.completions.create(
+            model="Kwai-Keye/Keye-VL-8B-Preview",
+            messages=messages,
+            top_p=0.95,
+            temperature=0.6,
+            stream=True,
+            timeout=360
+        )
+        response_text = ""
+        for response in responses:
+            response = response.choices[0].delta.content
+            response_text += response
+            _chatbot[-1] = (chat_query, _parse_text(response_text))
+            yield _chatbot
+        
+        response = response_text
+        _chatbot[-1] = (chat_query, _parse_text(response))
+        full_response = _parse_text(response)
+
+        task_history[-1] = (query, full_response)
+        print("Kwai-Keye-VL-Chat: " + full_response)
+        yield _chatbot
+
+
+    def regenerate(_chatbot, task_history):
+        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
+
+    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(task_history):
+        task_history.clear()
+        return []
+
+    with gr.Blocks() as demo:
+        gr.Markdown("""<center><font size=3> Kwai-Keye-VL Demo </center>""")
+
+        chatbot = gr.Chatbot(label='Kwai-Keye-VL', elem_classes="control-height", height=500)
+        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], [chatbot], show_progress=True
+        )
+        submit_btn.click(reset_user_input, [], [query])
+        empty_bin.click(reset_state, [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)
+
+    demo.queue(default_concurrency_limit=40).launch(
+        share=args.share,
+    )
+
+
+def main():
+    args = _get_args()
+    _launch_demo(args)
+
+
+if __name__ == '__main__':
+    main()