Update app.py

app.py

@@ -1,36 +1,100 @@
+from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
+from PIL import Image
 import requests
 import torch
-from PIL import Image
-from transformers import pipeline
+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)
+
+
+@spaces.GPU
+def bot_streaming(message, history, max_new_tokens=250):
+    
+    txt = message["text"]
+    ext_buffer = f"{txt}"
+    
+    messages= [] 
+    images = []
+    
 
-pipe = pipeline("text-generation", model="meta-llama/Llama-3.2-1B")
+    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]}]})
 
-# Load model directly
-from transformers import AutoTokenizer, AutoModelForCausalLM
+    # 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}]})
 
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
 
-# Load the processor
-processor = AutoProcessor.from_pretrained(model_id)
+    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)
 
-# Define the function to generate text based on input prompt
-def generate_text(prompt):
-    if llm_pipeline is None:
-        return "Error: Model not loaded."
-    result = llm_pipeline(prompt, max_length=100, num_return_sequences=1)
-    return result[0]['generated_text']
+    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
 
-# Create the Gradio interface
-interface = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(lines=7, label="Input Prompt"),
-    outputs="text",
-    title="Large Language Model Text Generation",
-    description="Enter a prompt to generate text using a large language model."
-)
 
-print("Launching the Gradio interface...")
-# Launch the interface
-interface.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)