Update app.py

app.py

@@ -1,64 +1,103 @@
+from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
+from PIL import Image
+import requests
+import torch
+from threading import Thread
 import gradio as gr
-from huggingface_hub import InferenceClient
+from gradio import FileData
+import time
+import spaces
+ckpt = ""
+model = MllamaForConditionalGeneration.from_pretrained(ckpt,
+    torch_dtype=torch.bfloat16).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
 
-"""
-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("letianWoowoof/woowoofv1")
 
+processor = AutoProcessor.from_pretrained("letianWoowoof/woowoofv1")
+model = AutoModelForPreTraining.from_pretrained("letianWoowoof/woowoofv1")
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+@spaces.GPU
+def bot_streaming(message, history, max_new_tokens=250):
+    
+    txt = message["text"]
+    ext_buffer = f"{txt}"
+    
+    messages= [] 
+    images = []
+    
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    for i, msg in enumerate(history): 
+        if isinstance(msg[0], tuple):
+            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
+            images.append(Image.open(msg[0][0]).convert("RGB"))
+        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]}]})
 
-    messages.append({"role": "user", "content": message})
+    # add current message
+    if len(message["files"]) == 1:
+        
+        if isinstance(message["files"][0], str): # examples
+            image = Image.open(message["files"][0]).convert("RGB")
+        else: # regular input
+            image = Image.open(message["files"][0]["path"]).convert("RGB")
+        images.append(image)
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
+    else:
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
 
-    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
+    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
 
-        response += token
-        yield response
+    if images == []:
+        inputs = processor(text=texts, return_tensors="pt").to("cuda")
+    else:
+        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
+    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
 
-
-"""
-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)",
-        ),
-    ],
-)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    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
 
 
-if __name__ == "__main__":
-    demo.launch()
+demo = gr.ChatInterface(fn=bot_streaming, title="Multimodal Llama", examples=[
+    [{"text": "Which era does this piece belong to? Give details about the era.", "files":["./examples/rococo.jpg"]},
+    200],
+    [{"text": "Where do the droughts happen according to this diagram?", "files":["./examples/weather_events.png"]},
+    250],
+    [{"text": "What happens when you take out white cat from this chain?", "files":["./examples/ai2d_test.jpg"]},
+    250],
+    [{"text": "How long does it take from invoice date to due date? Be short and concise.", "files":["./examples/invoice.png"]},
+    250],
+    [{"text": "Where to find this monument? Can you give me other recommendations around the area?", "files":["./examples/wat_arun.jpg"]}, 
+    250],
+    ],
+      textbox=gr.MultimodalTextbox(), 
+      additional_inputs = [gr.Slider(
+              minimum=10,
+              maximum=500,
+              value=250,
+              step=10,
+              label="Maximum number of new tokens to generate",
+          )
+        ],
+      cache_examples=False,
+      description="Try Multimodal Llama by Meta with transformers in this demo. Upload an image, and start chatting about it, or simply try one of the examples below. To learn more about Llama Vision, visit [our blog post](https://huggingface.co/blog/llama32). ",
+      stop_btn="Stop Generation", 
+      fill_height=True,
+    multimodal=True)
+    
+demo.launch(debug=True)