2 files changed, 13 insertions, 7 deletions

diff --git a/js/controllers/elements/element-controller.js b/js/controllers/elements/element-controller.js
index b35251ad..efb33292 100755
--- a/js/controllers/elements/element-controller.js
+++ b/js/controllers/elements/element-controller.js
@@ -5,22 +5,27 @@ No rights, expressed or implied, whatsoever to this software are provided by Mot
 </copyright> */
 
 var Montage = require("montage/core/core").Montage,
-    NJComponent =   require("js/lib/nj-base").NJComponent;
+    Component = require("montage/ui/component").Component;
 
-var ElementController = exports.ElementController = Montage.create(NJComponent, {
+exports.ElementController = Montage.create(Component, {
 
     addElement: {
         value: function(el, styles) {
             this.application.ninja.currentDocument.documentRoot.appendChild(el);
-            // Nested elements -
-            // TODO make sure the CSS is correct before nesting elements
+            // Nested elements - TODO make sure the CSS is correct before nesting elements
             // this.application.ninja.currentSelectedContainer.appendChild(el);
-            this.application.ninja.stylesController.setElementStyles(el, styles);
+            if(styles) {
+                this.application.ninja.stylesController.setElementStyles(el, styles);
+            }
         }
     },
 
+    // Remove the element from the DOM and clear the GLWord.
     removeElement: {
         value: function(el) {
+            if(el.elementModel && el.elementModel.shapeModel && el.elementModel.shapeModel.GLWorld) {
+                el.elementModel.shapeModel.GLWorld.clearTree();
+            }
             el.parentNode.removeChild(el);
         }
     },
diff --git a/js/controllers/elements/shapes-controller.js b/js/controllers/elements/shapes-controller.js
index ef453bea..c9164e3e 100755
--- a/js/controllers/elements/shapes-controller.js
+++ b/js/controllers/elements/shapes-controller.js
@@ -72,8 +72,9 @@ exports.ShapesController = Montage.create(CanvasController, {
                     canvas.setAttribute("data-RDGE-id", njModule.NJUtils.generateRandom());
                     canvas.width = el.width;
                     canvas.height = el.height;
-                    this.application.ninja.elementMediator.replaceElement(el, canvas);
-                    NJevent("elementDeleted", el);
+                    canvas.elementModel = el.elementModel;
+                    this.application.ninja.currentDocument.documentRoot.replaceChild(canvas, el);
+                    NJevent("elementsRemoved", el);
                     el = canvas;
                     this.toggleWebGlMode(el, value);
                     el.elementModel.shapeModel.GLWorld.render();