Update app.py

app.py

@@ -1,34 +1,33 @@
 import streamlit as st
 import transformers
 import torch
+import requests
+from PIL import Image
+from transformers import MllamaForConditionalGeneration, AutoProcessor
 
 HF_TOKEN=st.secrets["hf_token"]
 # Load the model and pipeline
 model_id = "meta-llama/Llama-3.2-11B-Vision"
+# Streamlit user interface
+st.title("LLM Model Inference")
+st.write(f"**Using model:** {model_id}")
 
 # Set up the pipeline with the Hugging Face token
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16, "use_auth_token": HF_TOKEN}
+model = MllamaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
 )
+processor = AutoProcessor.from_pretrained(model_id)
+
+url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
+prompt = "<|image|><|begin_of_text|>If I had to write a haiku for this one"
+inputs = processor(image, prompt, return_tensors="pt").to(model.device)
+
+output = model.generate(**inputs, max_new_tokens=30)
+st.write(processor.decode(output[0]))
 
-# Streamlit user interface
-st.title("LLM Model Inference")
-st.write(f"**Using model:** {model_id}")
-input_text = st.text_input("Enter your prompt:")
-
-if st.button("Generate"):
-    if input_text:  # Check if the input is not empty
-        # Generate text using the pipeline
-        messages = [
-        {"role": "system", "content": "You are a question answering assistant."},
-        {"role": "user", "content": input_text}
-        ]
-        response = pipeline(messages, max_new_tokens=30)  
-        st.write("Generated Response:")
-        st.write(response[0]['generated_text'][-1]['content'])  
-    else:
-        st.error("Please enter a prompt to generate text.")