File size: 3,368 Bytes
672864c 8f9661d 672864c 8f9661d 672864c 8f9661d f076978 8f9661d 672864c 8f9661d 672864c 8f9661d 672864c 8f9661d 672864c 8f9661d 672864c 8f9661d 672864c 8f9661d 672864c 8f9661d 672864c 8f9661d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import os
import streamlit as st
from huggingface_hub import login
from transformers import MllamaForConditionalGeneration, AutoProcessor
from PIL import Image
import torch
# Step 1: Log in to Hugging Face with your access token from secrets
huggingface_token = os.getenv("HUGGINGFACE_TOKEN") # Fetch the token from environment
if huggingface_token:
login(token=huggingface_token) # Authenticate using the token
else:
st.error("Hugging Face token not found. Please set it in the Secrets section.")
# Step 2: Load the model and processor
try:
model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(
model_name,
token=huggingface_token,
torch_dtype=torch.bfloat16,
device_map="auto",
)
processor = AutoProcessor.from_pretrained(
model_name,
use_auth_token=huggingface_token,
)
st.success("Model and processor loaded successfully!")
except Exception as e:
st.error(f"Error loading model or processor: {str(e)}")
# Step 3: Create a simple Streamlit app
def main():
st.title("Llama 3.2 11B Vision Model")
st.write("Upload an image and enter a prompt to generate output.")
# Upload image
image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
prompt = st.text_area("Enter your prompt here:")
if st.button("Generate Output"):
if image_file and prompt:
# Load image
image = Image.open(image_file).convert("RGB")
st.image(image, caption="Uploaded Image", use_column_width=True)
try:
# Prepare the messages in the format expected by the processor
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image"}
]
}
]
# Apply chat template
input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
# Prepare inputs for the model
inputs = processor(
text=input_text,
images=[image],
return_tensors="pt"
).to("cuda" if torch.cuda.is_available() else "cpu")
# Generate output
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=250,
)
# Decode the output
output_text = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
# Extract the generated response
# Remove the prompt part from the output_text
if input_text in output_text:
generated_output = output_text.replace(input_text, "").strip()
else:
generated_output = output_text.strip()
st.write("Generated Output:", generated_output)
except Exception as e:
st.error(f"Error during prediction: {str(e)}")
else:
st.warning("Please upload an image and enter a prompt.")
if __name__ == "__main__":
main()
|