Update app.py

app.py

@@ -1,23 +1,166 @@
+import spaces
+import os
+import re
+import time
 import gradio as gr
-from pytube import YouTube
-
-def download_video(url):
-    try:
-        yt = YouTube(url)
-        stream = yt.streams.get_highest_resolution()
-        stream.download(output_path="downloads", filename=f"{yt.title}.mp4")
-        return f"Downloaded: {yt.title}.mp4"
-    except Exception as e:
-        return f"An error occurred: {e}"
-
-# Create a Gradio interface
-iface = gr.Interface(
-    fn=download_video,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter YouTube URL..."),
-    outputs="text",
-    title="YouTube Video Downloader",
-    description="Enter the YouTube video URL and click 'Submit' to download the video."
-)
-
-# Launch the interface
-iface.launch()
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import TextIteratorStreamer
+from threading import Thread
+
+model_name = 'AIDC-AI/Ovis1.6-Llama3.2-3B'
+
+# load model
+model = AutoModelForCausalLM.from_pretrained(model_name,
+                                             torch_dtype=torch.bfloat16,
+                                             multimodal_max_length=8192,
+                                             token=os.getenv('HUGGINGFACE_TOKEN'),
+                                             trust_remote_code=True).to(device='cuda')
+text_tokenizer = model.get_text_tokenizer()
+visual_tokenizer = model.get_visual_tokenizer()
+streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_tokens=True)
+image_placeholder = '<image>'
+cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+def submit_chat(chatbot, text_input):
+    response = ''
+    chatbot.append((text_input, response))
+    return chatbot ,''
+
+@spaces.GPU
+def ovis_chat(chatbot, image_input):
+    # preprocess inputs
+    conversations = []
+    response = ""
+    text_input = chatbot[-1][0]
+    for query, response in chatbot[:-1]:
+        conversations.append({
+            "from": "human",
+            "value": query
+        })
+        conversations.append({
+            "from": "gpt",
+            "value": response
+        })
+    text_input = text_input.replace(image_placeholder, '')
+    conversations.append({
+        "from": "human",
+        "value": text_input
+    })
+    if image_input is not None:
+        conversations[0]["value"] = image_placeholder + '\n' + conversations[0]["value"]
+    prompt, input_ids, pixel_values = model.preprocess_inputs(conversations, [image_input])
+    attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
+    input_ids = input_ids.unsqueeze(0).to(device=model.device)
+    attention_mask = attention_mask.unsqueeze(0).to(device=model.device)
+    if image_input is None:
+        pixel_values = [None]
+    else:
+        pixel_values = [pixel_values.to(dtype=visual_tokenizer.dtype, device=visual_tokenizer.device)]
+
+    with torch.inference_mode():
+        gen_kwargs = dict(
+            max_new_tokens=512,
+            do_sample=False,
+            top_p=None,
+            top_k=None,
+            temperature=None,
+            repetition_penalty=None,
+            eos_token_id=model.generation_config.eos_token_id,
+            pad_token_id=text_tokenizer.pad_token_id,
+            use_cache=True
+        )
+    response = ""
+    thread = Thread(target=model.generate, 
+                kwargs={"inputs": input_ids,
+                        "pixel_values": pixel_values,
+                        "attention_mask": attention_mask,
+                        "streamer": streamer,
+                        **gen_kwargs})
+    thread.start()
+    for new_text in streamer:
+        response += new_text
+        chatbot[-1][1] = response
+        yield chatbot
+    thread.join()
+    # debug
+    print('*'*60)
+    print('*'*60)
+    print('OVIS_CONV_START')
+    for i, (request, answer) in enumerate(chatbot[:-1], 1):
+        print(f'Q{i}:\n {request}')
+        print(f'A{i}:\n {answer}')
+    print('New_Q:\n', text_input)
+    print('New_A:\n', response)
+    print('OVIS_CONV_END')
+
+def clear_chat():
+    return [], None, ""
+
+with open(f"{cur_dir}/resource/logo.svg", "r", encoding="utf-8") as svg_file:
+    svg_content = svg_file.read()
+font_size = "2.5em"
+svg_content = re.sub(r'(<svg[^>]*)(>)', rf'\1 height="{font_size}" style="vertical-align: middle; display: inline-block;"\2', svg_content)
+html = f"""
+<p align="center" style="font-size: {font_size}; line-height: 1;">
+    <span style="display: inline-block; vertical-align: middle;">{svg_content}</span>
+    <span style="display: inline-block; vertical-align: middle;">{model_name.split('/')[-1]}</span>
+</p>
+<center><font size=3><b>Ovis</b> has been open-sourced on <a href='https://huggingface.co/{model_name}'>😊 Huggingface</a> and <a href='https://github.com/AIDC-AI/Ovis'>🌟 GitHub</a>. If you find Ovis useful, a like❤️ or a star🌟 would be appreciated.</font></center>
+"""
+
+latex_delimiters_set = [{
+        "left": "\\(",
+        "right": "\\)",
+        "display": False 
+    }, {
+        "left": "\\begin{equation}",
+        "right": "\\end{equation}",
+        "display": True 
+    }, {
+        "left": "\\begin{align}",
+        "right": "\\end{align}",
+        "display": True
+    }, {
+        "left": "\\begin{alignat}",
+        "right": "\\end{alignat}",
+        "display": True
+    }, {
+        "left": "\\begin{gather}",
+        "right": "\\end{gather}",
+        "display": True
+    }, {
+        "left": "\\begin{CD}",
+        "right": "\\end{CD}",
+        "display": True
+    }, {
+        "left": "\\[",
+        "right": "\\]",
+        "display": True
+    }]
+
+text_input = gr.Textbox(label="prompt", placeholder="Enter your text here...", lines=1, container=False)
+with gr.Blocks(title=model_name.split('/')[-1], theme=gr.themes.Ocean()) as demo:
+    gr.HTML(html)
+    with gr.Row():
+        with gr.Column(scale=3):
+            image_input = gr.Image(label="image", height=350, type="pil")
+            gr.Examples(
+                examples=[
+                    [f"{cur_dir}/examples/case1.png", "explain this model to me."],
+                    [f"{cur_dir}/examples/case2.png", "Which city is in the picture?"],
+                ],
+                inputs=[image_input, text_input]
+            )
+        with gr.Column(scale=7):
+            chatbot = gr.Chatbot(label="Ovis", layout="panel", height=600, show_copy_button=True, latex_delimiters=latex_delimiters_set)
+            text_input.render()
+            with gr.Row():
+                send_btn = gr.Button("Send", variant="primary")
+                clear_btn = gr.Button("Clear", variant="secondary")
+
+    send_click_event = send_btn.click(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
+    submit_event = text_input.submit(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
+    clear_btn.click(clear_chat, outputs=[chatbot, image_input, text_input])
+
+demo.launch()