Update app.py

app.py

@@ -1,62 +1,119 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import spaces
 
-"""
-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)",
-        ),
-    ],
-)
+import os
+import warnings
+import shutil
+import time
+from threading import Thread
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, AutoProcessor
+from transformers import TextIteratorStreamer
+import torch
+from dc.model import *
+from dc.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from dc.conversation import conv_templates, SeparatorStyle
+from PIL import Image
+
+
+processor = AutoProcessor.from_pretrained('HuanjinYao/DenseConnector-v1.5-8B')
+tokenizer = AutoTokenizer.from_pretrained('HuanjinYao/DenseConnector-v1.5-8B', use_fast=False)
+model = LlavaLlamaForCausalLM.from_pretrained('HuanjinYao/DenseConnector-v1.5-8B', low_cpu_mem_usage=True, **kwargs)
+image_processor = model.get_vision_tower()
+if not vision_tower.is_loaded:
+    vision_tower.load_model()
+vision_tower.to(device=device, dtype=torch.float16)
+image_processor = vision_tower.image_processor
+
+model.to('cuda')
+
+# model.generation_config.eos_token_id = 128009
+tokenizer.unk_token = "<|reserved_special_token_0|>"
+tokenizer.pad_token = tokenizer.unk_token
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
+
+
+@spaces.GPU
+def bot_streaming(message, history):
+    print(message)
+    if message["files"]:
+        # message["files"][-1] is a Dict or just a string
+        if type(message["files"][-1]) == dict:
+            image = message["files"][-1]["path"]
+        else:
+            image = message["files"][-1]
+    else:
+        # if there's no image uploaded for this turn, look for images in the past turns
+        # kept inside tuples, take the last one
+        for hist in history:
+            if type(hist[0]) == tuple:
+                image = hist[0][0]
+    try:
+        if image is None:
+            # Handle the case where image is None
+            gr.Error("You need to upload an image for LLaVA to work.")
+    except NameError:
+        # Handle the case where 'image' is not defined at all
+        gr.Error("You need to upload an image for LLaVA to work.")
+
+    
+    conv = conv_templates['llama_3'].copy()
+    if len(history) == 0:
+        user = DEFAULT_IMAGE_TOKEN + '\n' + message['text']
+    else:
+        for idx, (user, assistant) in enumerate(history):
+            # conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+            if idx == 0:
+                user = DEFAULT_IMAGE_TOKEN + '\n' + user
+            conv.append_message(conv.roles[0], user)
+            conv.append_message(conv.roles[1], assistant)
+    conv.append_message(conv.roles[0], user)
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt()
+
+    image = Image.open(os.path.join(image, image_file)).convert('RGB')
+    image_tensor = image_processor([image], image_processor, self.model_config)[0]
+
+    inputs = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt')
+
+    
+    streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": False, "skip_prompt": True})
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False, eos_token_id = terminators)
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+
+    buffer = ""
+    # time.sleep(0.5)
+    for new_text in streamer:
+        if "<|eot_id|>" in new_text:
+            new_text = new_text.split("<|eot_id|>")[0]
+        buffer += new_text
+
+        generated_text_without_prompt = buffer
+        # time.sleep(0.06)
+        yield generated_text_without_prompt
+
+
+chatbot=gr.Chatbot(placeholder=PLACEHOLDER,scale=1)
+chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
+with gr.Blocks(fill_height=True, ) as demo:
+    gr.ChatInterface(
+    fn=bot_streaming,
+    title="LLaVA Llama-3-8B",
+    examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]},
+              {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
+    description="Try [LLaVA Llama-3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",
+    stop_btn="Stop Generation",
+    multimodal=True,
+    textbox=chat_input,
+    chatbot=chatbot,
+    )
 
 
 if __name__ == "__main__":