diff --git a/Dockerfile b/Dockerfile
index a43dd0c..a1377f7 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -46,4 +46,4 @@ COPY --chown=appuser . .
 STOPSIGNAL SIGINT
 ENV LD_PRELOAD=libtcmalloc.so
 ENV PATH="$PATH:/home/appuser/.local/bin"
-CMD python3 "./kohya_gui.py" ${CLI_ARGS} --listen 0.0.0.0 --server_port 7860 --headless
+CMD python3 "./kohya_gui.py" ${CLI_ARGS} --listen 0.0.0.0 --server_port 7860
diff --git a/docker-compose.yaml b/docker-compose.yaml
index afaa1bb..c24588f 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -12,6 +12,7 @@ services:
     environment:
       CLI_ARGS: ""
       SAFETENSORS_FAST_GPU: 1
+      DISPLAY: $DISPLAY
     tmpfs:
       - /tmp      
     volumes:
@@ -21,6 +22,7 @@ services:
       - ./.cache/config:/app/appuser/.config
       - ./.cache/nv:/home/appuser/.nv 
       - ./.cache/keras:/home/appuser/.keras      
+      - /tmp/.X11-unix:/tmp/.X11-unix
     deploy:
       resources:
         reservations: