update

app.py

@@ -8,6 +8,9 @@ from transformers import AutoModelForCausalLM, AutoProcessor
 import re
 import random
 
+# add a command for installing flash-attn
+os.system('pip install flash-attn --no-build-isolation')
+
 pygame.mixer.quit()  # Disable sound
 
 # Constants
@@ -30,8 +33,9 @@ STATIC = (0, 0)
 ACTIONS = ["up", "down", "left", "right", "static"]
 
 # Load AI Model
+dtype = torch.bfloat16
 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")
 
@@ -137,7 +141,7 @@ def play_game():
     inputs = magma_processor(images=[pil_img], 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)
     generation_args = { 
         "max_new_tokens": 10, 
         "temperature": 0,