diff --git a/agent_jarvis/dockerfile b/agent_jarvis/dockerfile
index f647592..6975245 100644
--- a/agent_jarvis/dockerfile
+++ b/agent_jarvis/dockerfile
@@ -45,4 +45,7 @@ COPY ./agent_jarvis /root/jarvis
 COPY ./example_modules /root/example_modules
 COPY ./example_services /root/example_services
 
+# Download the u2net model
+RUN cd /root/example_modules/demo_modules/ && python rembg.module.py
+
 CMD ["bash", "/root/jarvis/supervisor.d/entrypoint.sh"]
diff --git a/example_modules/demo_modules/rembg.module.py b/example_modules/demo_modules/rembg.module.py
index 95bf95c..dab055c 100644
--- a/example_modules/demo_modules/rembg.module.py
+++ b/example_modules/demo_modules/rembg.module.py
@@ -2,10 +2,6 @@ import io
 import base64
 from typing import List
 
-from jarvis.functional_modules.caller_context import CallerContext
-from jarvis.functional_modules.functional_module import functional_module
-from jarvis.logger import logger
-
 
 def reg_or_not():
     try:
@@ -16,7 +12,13 @@ def reg_or_not():
         return
 
     session = new_session(model_name="u2net")
+    if __name__ == "__main__":
+        # If we run this script directly, download the model and exit
+        return
 
+    from jarvis.functional_modules.caller_context import CallerContext
+    from jarvis.functional_modules.functional_module import functional_module
+    from jarvis.logger import logger
     @functional_module(
         name="remove_bg",
         description="Remove the background of last image",
@@ -38,4 +40,5 @@ def reg_or_not():
         await context.reply_image_base64(result)
         return "Success"
 
+
 reg_or_not()