diff --git a/setup_diagrams.sh b/setup_diagrams.sh
index 8d1d530..ce4daec 100644
--- a/setup_diagrams.sh
+++ b/setup_diagrams.sh
@@ -4,20 +4,22 @@
 sh /setup_programs.sh
 
 # Install accompanying soft
-apt-get install -y python3 git python3-pip curl openjdk-8-jdk graphviz
+apt install -y python3 git python3-pip curl openjdk-8-jdk graphviz
 
 # Get diagrams
-cd /var
+cd /tmp
 git clone https://github.com/BITPlan/diagrams
+mv /tmp/diagrams/* /var/diagrams
+rm -r /tmp/diagrams/
 cd /var/diagrams
 
 # Replace sudo commands in script
-sed -i 's/sudo //g' ./install
-sed -i 's/apt-get install/apt-get install -y/g' ./install
+sed -i 's/sudo //g' ./scripts/install
+sed -i 's/apt-get install/apt-get install -y/g' ./scripts/install
 sed -i 's/sudo //g' ./run
 
 # Install diagrams
-./install
+./scripts/install
 
 # Setup environment
 export GRAPHVIZ_DOT=/usr/bin/dot
@@ -27,8 +29,8 @@ ln -s /usr/bin/mscgen /var/diagrams/bin/mscgen
 ln -s /usr/bin/dot /var/diagrams/bin/dot
 
 # Replace with edited files
-rm /var/diagrams/dgs/diagrams.py
-rm /var/diagrams/dgs/webserver.py
+mv /var/diagrams/dgs/diagrams.py /var/diagrams/dgs/orig_diagrams.py
+mv /var/diagrams/dgs/webserver.py /var/diagrams/dgs/orig_webserver.py
 wget https://git.dfosd.synology.me/docker/ubuntu-diagrams/raw/branch/master/diagrams.py -P /var/diagrams/dgs
 wget https://git.dfosd.synology.me/docker/ubuntu-diagrams/raw/branch/master/webserver.py -P /var/diagrams/dgs