diff --git a/lib/modules/graph/index.js b/lib/modules/graph/index.js
index 75f1fb156..21d6ce32d 100644
--- a/lib/modules/graph/index.js
+++ b/lib/modules/graph/index.js
@@ -6,6 +6,7 @@ class GraphGenerator {
   constructor(embark, options) {
     const self = this;
     this.events = embark.events;
+    this.contracts = [];
 
     this.events.setCommandHandler("graph:create", function(options, cb) {
       self.generate(options);
@@ -20,13 +21,12 @@ class GraphGenerator {
     let relationshipString = "";
     let idMapping = {};
     let contractInheritance = {};
-    let contracts = {};
     let contractsDependencies = {};
 
     async.waterfall([
       function getContractList(next) {
-        self.events.request('contracts:list', (err, _contracts) => {
-          contracts = _contracts;
+        self.events.request('contracts:list', (err, contracts) => {
+          self.contracts = contracts;
           next();
         });
       },
@@ -37,7 +37,7 @@ class GraphGenerator {
         });
       },
       function (next) {
-        for (let contract of contracts) {
+        for (let contract of self.contracts) {
           if (options.skipUndeployed && !contract.deploy) continue;
 
           id++;
@@ -49,7 +49,7 @@ class GraphGenerator {
           contractLabel += `${contract.className}`;
           let tooltip = contract.className;
 
-          if (contract.instanceOf !== undefined && contracts[contract.instanceOf] !== undefined) {
+          if (contract.instanceOf !== undefined && self.getContract(contract.instanceOf) !== undefined) {
             contractInheritance[contract.className] = contract.instanceOf;
             contractLabel += ": " + contract.instanceOf;
             tooltip += " instance of " + contract.instanceOf;
@@ -106,7 +106,7 @@ class GraphGenerator {
           let contractDependencies = Array.from(new Set(contractsDependencies[c]));
           contractDependencies.forEach((d) => {
             if (idMapping[c] !== undefined && idMapping[d] !== undefined) {
-              if ((options.skipUndeployed && contracts[c].deploy && contracts[d].deploy) || !options.skipUndeployed) {
+              if ((options.skipUndeployed && self.getContract(c).deploy && self.getContract(d).deploy) || !options.skipUndeployed) {
                 relationshipString += `${idMapping[d]}->${idMapping[c]}[constraint=true, arrowtail=diamond, tooltip="${c} uses ${d}"]\n`;
               }
             }
@@ -116,13 +116,10 @@ class GraphGenerator {
       },
       function (next) {
         for (let c in contractInheritance){
-          if(options.skipUndeployed && !contracts[contractInheritance[c]].deploy) continue;
+          if(options.skipUndeployed && !self.getContract(contractInheritance[c]).deploy) continue;
 
           relationshipString += `${idMapping[contractInheritance[c]]}->${idMapping[c]}[tooltip="${c} instance of ${contractInheritance[c]}"]\n`;
         }
-        console.log("-----");
-        console.log(relationshipString);
-        console.log("-----");
         next();
       },
       function (next) {
@@ -134,7 +131,6 @@ class GraphGenerator {
             ${relationshipString}
         }`;
 
-        console.log(dot);
         let svg = Viz(dot);
 
         let filename = "diagram.svg";
@@ -147,6 +143,10 @@ class GraphGenerator {
     ], function(_err, _result) {
     });
   }
+
+  getContract(contractName) {
+    return this.contracts.find((contract) => { return contract.className === contractName });
+  }
 }
 
 module.exports = GraphGenerator;