diff --git a/src/pytti/update_func.py b/src/pytti/update_func.py
index 5ef85c4..1d3606e 100644
--- a/src/pytti/update_func.py
+++ b/src/pytti/update_func.py
@@ -16,6 +16,8 @@
     freeze_vram_usage,
 )
 
+from pytti.image_models.pixel import PixelImage
+
 from pytti.Transforms import (
     animate_2d,
     zoom_2d,