microsoft/Magma-Gaming
14
1import os2# add a command for installing flash-attn3os.system('pip install flash-attn --no-build-isolation')4os.system("pip install gradio==4.44.1")5 6import pygame7import numpy as np8import gradio as gr9import time10import torch11from PIL import Image12from transformers import AutoModelForCausalLM, AutoProcessor13import re14import random15 16pygame.mixer.quit() # Disable sound17 18# Constants19WIDTH, HEIGHT = 800, 80020GRID_SIZE = 8021WHITE = (255, 255, 255)22GREEN = (34, 139, 34) # Forest green - more like an apple23RED = (200, 50, 50)24BLACK = (0, 0, 0)25GRAY = (128, 128, 128)26YELLOW = (218, 165, 32) # Golden yellow color27 28# Directions29UP = (0, -1)30DOWN = (0, 1)31LEFT = (-1, 0)32RIGHT = (1, 0)33STATIC = (0, 0)34 35ACTIONS = ["up", "down", "left", "right", "static"]36 37# Load AI Model38dtype = torch.bfloat1639magma_model_id = "microsoft/Magma-8B"40magam_model = AutoModelForCausalLM.from_pretrained(magma_model_id, trust_remote_code=True, torch_dtype=dtype)41magma_processor = AutoProcessor.from_pretrained(magma_model_id, trust_remote_code=True)42magam_model.to("cuda")43 44magma_img = pygame.image.load("./assets/images/magma_game_thin.png")45magma_img = pygame.transform.scale(magma_img, (GRID_SIZE, GRID_SIZE))46 47class MagmaFindGPU:48 def __init__(self):49 self.reset()50 self.step_count = 051 52 def reset(self):53 self.snake = [(5, 5)]54 self.direction = RIGHT55 self.score = 056 self.game_over = False57 self.step_count = 058 self.place_target()59 60 def place_target(self):61 while True:62 target_x = np.random.randint(1, WIDTH // GRID_SIZE - 1)63 target_y = np.random.randint(1, HEIGHT // GRID_SIZE - 1)64 if (target_x, target_y) not in self.snake:65 self.target = (target_x, target_y)66 break67 68 def step(self, action):69 if action == "up":70 self.direction = UP71 elif action == "down":72 self.direction = DOWN73 elif action == "left":74 self.direction = LEFT75 elif action == "right":76 self.direction = RIGHT77 elif action == "static":78 self.direction = STATIC79 80 if self.game_over:81 self.reset()82 return self.render(), self.score83 84 new_head = (self.snake[0][0] + self.direction[0], self.snake[0][1] + self.direction[1])85 86 if new_head[0] < 0 or new_head[1] < 0 or new_head[0] >= WIDTH // GRID_SIZE or new_head[1] >= HEIGHT // GRID_SIZE:87 self.game_over = True88 return self.render(), self.score89 90 self.snake = [new_head] # Keep only the head (single block snake)91 self.step_count += 192 93 # Check if the target is covered by four surrounding squares94 head_x, head_y = self.snake[0]95 neighbors = set([(head_x, head_y - 1), (head_x, head_y + 1), (head_x - 1, head_y), (head_x + 1, head_y)])96 97 if neighbors.issuperset(set([self.target])):98 self.score += 199 self.place_target()100 101 return self.render(), self.score102 103 def render(self):104 pygame.init()105 surface = pygame.Surface((WIDTH, HEIGHT))106 surface.fill(BLACK)107 108 head_x, head_y = self.snake[0]109 surface.blit(magma_img, (head_x * GRID_SIZE, head_y * GRID_SIZE)) 110 111 # pygame.draw.rect(surface, RED, (self.snake[0][0] * GRID_SIZE, self.snake[0][1] * GRID_SIZE, GRID_SIZE, GRID_SIZE))112 pygame.draw.rect(surface, GREEN, (self.target[0] * GRID_SIZE, self.target[1] * GRID_SIZE, GRID_SIZE, GRID_SIZE))113 114 # Draw four surrounding squares with labels115 head_x, head_y = self.snake[0]116 neighbors = [(head_x, head_y - 1), (head_x, head_y + 1), (head_x - 1, head_y), (head_x + 1, head_y)]117 labels = ["1", "2", "3", "4"]118 font = pygame.font.Font(None, 48)119 120 # clone surface121 surface_nomark = surface.copy()122 for i, (nx, ny) in enumerate(neighbors):123 if 0 <= nx < WIDTH // GRID_SIZE and 0 <= ny < HEIGHT // GRID_SIZE:124 pygame.draw.rect(surface, RED, (nx * GRID_SIZE, ny * GRID_SIZE, GRID_SIZE, GRID_SIZE), GRID_SIZE)125 # pygame.draw.rect(surface_nomark, RED, (nx * GRID_SIZE, ny * GRID_SIZE, GRID_SIZE, GRID_SIZE), GRID_SIZE)126 127 text = font.render(labels[i], True, WHITE)128 text_rect = text.get_rect(center=(nx * GRID_SIZE + GRID_SIZE // 2, ny * GRID_SIZE + GRID_SIZE // 2))129 surface.blit(text, text_rect)130 131 return np.array(pygame.surfarray.array3d(surface_nomark)).swapaxes(0, 1), np.array(pygame.surfarray.array3d(surface)).swapaxes(0, 1)132 133 def get_state(self):134 return self.render()135 136game = MagmaFindGPU()137 138def play_game():139 state, state_som = game.get_state()140 pil_img = Image.fromarray(state_som)141 convs = [142 {"role": "system", "content": "You are an agent that can see, talk, and act. Avoid hitting the wall."}, 143 {"role": "user", "content": "<image_start><image><image_end>\nWhich mark is closer to green block? Answer with a single number."},144 ]145 prompt = magma_processor.tokenizer.apply_chat_template(convs, tokenize=False, add_generation_prompt=True)146 inputs = magma_processor(images=[pil_img], texts=prompt, return_tensors="pt")147 inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)148 inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0) 149 inputs = inputs.to("cuda").to(dtype)150 generation_args = { 151 "max_new_tokens": 10, 152 "temperature": 0.3, 153 "do_sample": True, 154 "use_cache": True,155 "num_beams": 1,156 }157 with torch.inference_mode():158 generate_ids = magam_model.generate(**inputs, **generation_args)159 generate_ids = generate_ids[:, inputs["input_ids"].shape[-1] :]160 action = magma_processor.decode(generate_ids[0], skip_special_tokens=True).strip()161 # extract mark id fro action use re162 match = re.search(r'\d+', action)163 if match:164 action = match.group(0)165 if action.isdigit() and 1 <= int(action) <= 4:166 action = ACTIONS[int(action) - 1]167 else:168 # random choose one from the pool169 action = random.choice(ACTIONS[:-1])170 else:171 action = random.choice(ACTIONS[:-1])172 173 img, score = game.step(action)174 img = img[0]175 return img, f"Score: {score}"176 177def reset_game():178 game.reset()179 return game.render()[0], "Score: 0"180 181MARKDOWN = """182<div align="center">183<h2>Magma: A Foundation Model for Multimodal AI Agents</h2>184 185\[[arXiv Paper](https://www.arxiv.org/pdf/2502.13130)\] \[[Project Page](https://microsoft.github.io/Magma/)\] \[[Github Repo](https://github.com/microsoft/Magma)\] \[[Hugging Face Model](https://huggingface.co/microsoft/Magma-8B)\] 186 187This demo is powered by [Gradio](https://gradio.app/).188 189<b>Goal: Collects the green blocks by automatically moving up, down, left and right.</b>190 191</div>192"""193 194with gr.Blocks() as interface:195 gr.Markdown(MARKDOWN)196 with gr.Row():197 image_output = gr.Image(label="Game Screen")198 with gr.Column():199 score_output = gr.Text(label="Score", elem_classes="large-text")200 gr.HTML("""201 <style>202 .large-text textarea {203 font-size: 24px !important;204 }205 </style>206 """)207 start_btn = gr.Button("Start/Reset Game")208 209 interface.load(fn=play_game, every=1, inputs=[], outputs=[image_output, score_output])210 start_btn.click(fn=reset_game, inputs=[], outputs=[image_output, score_output])211 212interface.launch()213