Update app.py

app.py

@@ -1,34 +1,28 @@
-import gradio as gr
-from transformers import pipeline
+import requests
+import torch
+from PIL import Image
+from transformers import MllamaForConditionalGeneration, AutoProcessor
 
-# Load the large language model (LLM)
-try:
-# Load model directly
-from transformers import AutoProcessor, AutoModelForPreTraining
+model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
 
-processor = AutoProcessor.from_pretrained("meta-llama/Llama-3.2-11B-Vision-Instruct")
-model = AutoModelForPreTraining.from_pretrained("meta-llama/Llama-3.2-11B-Vision-Instruct") # You can use a different model here
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    llm_pipeline = None
+model = MllamaForConditionalGeneration.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+processor = AutoProcessor.from_pretrained(model_id)
 
-# 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']
+url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
 
-# 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."
-)
+messages = [
+    {"role": "user", "content": [
+        {"type": "image"},
+        {"type": "text", "text": "If I had to write a haiku for this one, it would be: "}
+    ]}
+]
+input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
+inputs = processor(image, input_text, return_tensors="pt").to(model.device)
 
-print("Launching the Gradio interface...")
-# Launch the interface
-interface.launch()
+output = model.generate(**inputs, max_new_tokens=30)
+print(processor.decode(output[0]))