Update app.py

app.py

@@ -1,135 +1,101 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
-import spaces #0.32.0
-import torch
-import os
-import platform
-import requests
+from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
 from PIL import Image
+import requests
+import torch
+from threading import Thread
+import gradio as gr
+from gradio import FileData
+import time
+import spaces
+ckpt = "meta-llama/Llama-3.2-11B-Vision-Instruct"
+model = MllamaForConditionalGeneration.from_pretrained(ckpt,
+    torch_dtype=torch.bfloat16).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
 
 
-model = ""
-duration = None
-token = os.getenv('deepseekv2')
-provider = None #'fal-ai' #None #replicate # sambanova
-mode = "text-to-text"
-
-print(f"Is CUDA available: {torch.cuda.is_available()}")
-print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-print(f"CUDA version: {torch.version.cuda}")
-print(f"Python version: {platform.python_version()}")
-print(f"Pytorch version: {torch.__version__}")
-print(f"Gradio version: {gr. __version__}")
-# print(f"HFhub version: {huggingface_hub.__version__}")
-
-
-"""
-Packages ::::::::::
-Is CUDA available: True
-CUDA device: NVIDIA A100-SXM4-80GB MIG 3g.40gb
-CUDA version: 12.1
-Python version: 3.10.13
-Pytorch version: 2.4.0+cu121
-Gradio version: 5.0.1
-"""
-
-
-def choose_model(model_name):
-    if model_name == "DeepSeek-R1-Distill-Qwen-1.5B":
-        model = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-
-    elif model_name == "DeepSeek-R1-Distill-Qwen-32B":
-        model = "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"
-        
-    elif model_name == "Llama3-8b-Instruct":    
-        model = "meta-llama/Meta-Llama-3-8B-Instruct"
-
-    elif model_name == "Llama3.1-8b-Instruct":
-        model = "meta-llama/Llama-3.1-8B-Instruct"
-
-    elif model_name == "Llama2-13b-chat":
-        model = "meta-llama/Llama-2-13b-chat-hf"
-
-    elif model_name == "Llama-3.2-11B-Vision-Instruct":
-        model = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-        mode = "image-to-text"
-        return model
-
-    elif model_name == "Gemma-2-2b":
-        model = "google/gemma-2-2b-it"
-
-    elif model_name == "Gemma-7b":
-        model = "google/gemma-7b"
+@spaces.GPU
+def bot_streaming(message, history, max_new_tokens=250):
     
-    elif model_name == "Mixtral-8x7B-Instruct":
-        model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-
-    elif model_name == "Microsoft-phi-2":
-        model = "microsoft/phi-2"
-
-    elif model_name == "Qwen2.5-Coder-32B-Instruct":
-        model = "Qwen/Qwen2.5-Coder-32B-Instruct"
-
-    else:    # default to zephyr if no model chosen
-        model = "HuggingFaceH4/zephyr-7b-beta"
-
-    mode = "text-to-text"
-    return model
+    txt = message["text"]
+    ext_buffer = f"{txt}"
     
-
-@spaces.GPU(duration=duration)
-def respond(message, history: list[tuple[str, str]], image=None, model, system_message, max_tokens, temperature, top_p):
-
-    print(model)
-    model_name = choose_model(model)
-
-    client = InferenceClient(model_name, provider=provider, token=os.getenv('deepseekv2'))
+    messages= [] 
+    images = []
     
-    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
 
-    # else:
-    #     url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
-    #     image = Image.open(requests.get(url, stream=True).raw)
-
-
-demo = gr.ChatInterface(
-
-    respond,
-    input_components=[gr.Textbox(label="Message"), gr.Image(label="image", type="pil")],
-    multimodal=True,
-    stop_btn = "Stop generation",
-    # multimodal = True,
-    title="Ask me anything",
-    description="Hi there! I am your friendly AI chatbot. Choose from different language models under the Additional Inputs tab below.",
-    examples=[["Explain quantum computing"], ["Explain forex trading"], ["What is the capital of China?"], ["Make a poem about nature"]],
-    additional_inputs=[
-        gr.Dropdown(["DeepSeek-R1-Distill-Qwen-1.5B", "DeepSeek-R1-Distill-Qwen-32B", "Gemma-2-2b", "Gemma-7b", "Llama2-13b-chat", "Llama3-8b-Instruct", "Llama3.1-8b-Instruct", "Llama-3.2-11B-Vision-Instruct", "Microsoft-phi-2", "Mixtral-8x7B-Instruct", "Qwen2.5-Coder-32B-Instruct", "Zephyr-7b-beta"], label="Select Model"),
-        gr.Textbox(value="You are a friendly and helpful Chatbot, be concise and straight to the point, avoid excessive reasoning.", 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)")
+    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]}]})
+
+    # 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}]})
+
+
+    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
+
+    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)
+
+    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
+
+
+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)
 
-
-)
-
-
-if __name__ == "__main__":
-    demo.launch(share=True)