diff --git a/Dockerfile b/Dockerfile
index c321422..65dd34a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base
+FROM nvidia/cuda:12.6.0-runtime-ubuntu24.04 AS base
 
 # Install base packages
 RUN apt-get update && \
@@ -11,7 +11,7 @@ RUN apt-get update && \
     git \
     gnupg \
     jq \
-    libcurl4-nss-dev \
+    libcurl4-gnutls-dev \
     libvulkan1 \
     libvulkan-dev \
     mesa-vulkan-drivers \
@@ -46,8 +46,6 @@ RUN mkdir -p /usr/local/lib/llama-cpp-python
 RUN tar -xf /tmp/vulkan-libs.tar -C /usr/local/lib/llama-cpp-python
 RUN cp -f /usr/local/lib/llama-cpp-python/*.so* /usr/local/lib/
 RUN rm /tmp/vulkan-libs.tar
-# Set Vulkan environment to use CPU fallback (LavaPipe)
-ENV VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/lvp_icd.x86_64.json
 
 # Create directory structure
 RUN mkdir -p \
diff --git a/scripts/local.sh b/scripts/local.sh
index 0407e54..1b1a6f3 100755
--- a/scripts/local.sh
+++ b/scripts/local.sh
@@ -21,6 +21,13 @@ docker build \
     --tag sia \
     .
 
+GPU_FLAGS=""
+# Check if AMD GPU devices are available
+if [ -e /dev/dri ] && [ -e /dev/kfd ]; then
+    echo "AMD GPU detected, enabling GPU acceleration"
+    GPU_FLAGS="--device=/dev/dri --device=/dev/kfd --group-add video"
+fi
+
 docker run \
     --init \
     --rm \
@@ -28,6 +35,7 @@ docker run \
     -p 8080:8080 \
     -p 8000:8000 \
     --env-file .env \
+    $GPU_FLAGS \
     -v /$(pwd)/models/:/root/models/ \
     -v /$(pwd)/iterations/:/root/data/iterations/ \
     -v /$(pwd)/tasks/:/root/data/tasks/ \
diff --git a/sia/config.py b/sia/config.py
index 093c476..918e226 100644
--- a/sia/config.py
+++ b/sia/config.py
@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import Dict
 import argparse
 import os
-import tomli as tomllib
+import tomllib
 
 @dataclass
 class Config:
diff --git a/sia/web_agent.py b/sia/web_agent.py
index 87c325b..abbd13f 100644
--- a/sia/web_agent.py
+++ b/sia/web_agent.py
@@ -49,6 +49,7 @@ class WebAgent(BaseAgent):
 
         self._update_compiled_context()
         self._working_memory.add_change_handler(self._update_compiled_context)
+        self._parser._io_buffer.add_change_handler(lambda *_: self._update_compiled_context())
     
     @property
     def response_buffer(self) -> ResponseBuffer:
diff --git a/tools/gemma_infer/pyproject.toml b/tools/gemma_infer/pyproject.toml
index 3a1cea3..1bbf173 100644
--- a/tools/gemma_infer/pyproject.toml
+++ b/tools/gemma_infer/pyproject.toml
@@ -16,6 +16,7 @@ dependencies = [
     "sentencepiece>=0.2.0",
     "tiktoken>=0.9.0",
     "transformers>=4.0.0",
+    "vulkan",
     "xml_schema_validator @ file:///root/sia/lib/xml_schema_validator",
 ]
 
diff --git a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
index e688dec..008a209 100644
--- a/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
+++ b/tools/gemma_infer/src/gemma_infer/gemma_llm_engine.py
@@ -15,19 +15,20 @@ llama_cpp._lib.ggml_backend_load_all()
 
 class GemmaLlmEngine(LlmEngine):
     def __init__(
-            self,
-            model: str,
-            tokenizer: str,
-            token_limit: int,
-        ):
+        self,
+        model: str,
+        tokenizer: str,
+        token_limit: int,
+    ):
         self._model = model
         self._tokenizer = AutoTokenizer.from_pretrained(tokenizer)
         self._token_limit = token_limit
+
         self._llm = Llama(
             model_path=model,
-            n_gpu_layers=100,
+            n_gpu_layers=0,
             n_ctx=token_limit,
-            #verbose=False,  # Disable most logging
+            flash_attn=True,
         )
 
     def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
@@ -56,7 +57,9 @@ class GemmaLlmEngine(LlmEngine):
         yield from skip_prefix(content_generator(), prefix)
 
     def token_count(self, system: str, context: str) -> int:
-        return len(self._format_messages(system, context, None))
+        prompt = self._format_messages(system, context, None)
+        tokens = self._tokenizer.encode(prompt)
+        return len(tokens)
 
     def token_limit(self) -> int:
         return self._token_limit