Updated main.py

main.py

@@ -1,21 +1,22 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoProcessor
+from PIL import Image
+import requests
 
-# Load the tokenizer and model from Hugging Face
-model_name = "impactframes/molmo-7B-D-bnb-4bit"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load the processor and model
+processor = AutoProcessor.from_pretrained('allenai/Molmo-7B-D-0924', trust_remote_code=True, device_map='auto')
+model = AutoModelForCausalLM.from_pretrained('allenai/Molmo-7B-D-0924', trust_remote_code=True, device_map='auto')
 
-# Example input prompt
-prompt = "What is the meaning of life?"
+# Download an image
+image_url = "https://picsum.photos/id/237/536/354"
+image = Image.open(requests.get(image_url, stream=True).raw)
 
-# Tokenize the input
-inputs = tokenizer(prompt, return_tensors="pt")
+# Process the image with some input text
+inputs = processor(images=[image], text="Describe this image.")
+inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
 
-# Generate output
-with torch.no_grad():
-    outputs = model.generate(inputs.input_ids, max_length=100)
+# Generate text based on the input
+output = model.generate_from_batch(inputs, max_new_tokens=200)
 
-# Decode the output
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-print(response)
+# Decode and print the generated text
+generated_text = processor.tokenizer.decode(output[0], skip_special_tokens=True)
+print(generated_text)