update

app.py

@@ -24,11 +24,12 @@ from transformers import AutoProcessor
 repo_id = "microsoft/OmniParser-v2.0"  # HF repo
 local_dir = "weights"  # Target local directory
 
+dtype = torch.bfloat16
 som_generator = MarkHelper()
 magma_som_prompt = "<image>\nIn this view I need to click a button to \"{}\"? Provide the coordinates and the mark index of the containing bounding box if applicable."
 magma_qa_prompt = "<image>\n{} Answer the question briefly."
 magma_model_id = "microsoft/Magma-8B"
-magam_model = AutoModelForCausalLM.from_pretrained(magma_model_id, trust_remote_code=True)
+magam_model = AutoModelForCausalLM.from_pretrained(magma_model_id, trust_remote_code=True, torch_dtype=dtype)
 magma_processor = AutoProcessor.from_pretrained(magma_model_id, trust_remote_code=True)
 magam_model.to("cuda")
 
@@ -91,7 +92,7 @@ def get_som_response(instruction, image_som):
     inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
     inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
     inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    inputs = inputs.to("cuda")
+    inputs = inputs.to("cuda").to(dtype)
 
     magam_model.generation_config.pad_token_id = magma_processor.tokenizer.pad_token_id
     with torch.inference_mode():