diff --git a/DESCRIPTION b/DESCRIPTION
index 798d1d3..a979c3e 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,17 +1,17 @@
 Package: datamodelr
 Type: Package
 Title: Define and plot data model diagrams
-Version: 0.1.2.9002
-Date: 2016-04-12
+Version: 0.2.1.9001
+Date: 2017-01-12
 Author: Darko Bergant
 Maintainer: Darko Bergant <darko.bergant@gmail.com>
 URL: https://github.com/bergant/datamodelr
 Description: Data model definition with YAML file,
     extract from R data frames,
     reverse-engineering from PostgreSQL and SQL Server and
-    rendering with DiagrammeR graph objects.
+    rendering with DiagrammeR/graphviz.
 Depends:
-    R (>= 3.2)
+    R (>= 3.3)
 License: MIT + file LICENSE
 Suggests:
     yaml,
diff --git a/R/graph.R b/R/graph.R
index a1ab801..ad52823 100644
--- a/R/graph.R
+++ b/R/graph.R
@@ -189,19 +189,21 @@ dm_create_graph_list <- function(dm, view_type = "all",
     })
 
   nodes <-
-    list(
+    data.frame(
       nodes = names(tables),
       label = g_labels,
       shape = "plaintext",
       type = "upper",
-      segment = dm$tables[order(dm$tables$table), "segment"]
+      segment = dm$tables[order(dm$tables$table), "segment"],
+
+      stringsAsFactors = FALSE
     )
 
 
   if(!is.null(dm$references)) {
     edges <-
       with(dm$references[dm$references$ref_col_num == 1,],
-           DiagrammeR::create_edges(from = table, to = ref))
+           data.frame(from = table, to = ref, stringsAsFactors = FALSE))
   } else {
     edges <- NULL
   }
@@ -236,10 +238,6 @@ dm_create_graph <- function(dm, rankdir = "BT", graph_name = "Data Model",
 
   if(!is.data_model(dm)) stop("Input must be a data model object.")
 
-  if( !requireNamespace("DiagrammeR", quietly = TRUE)) {
-    stop("DiagrammeR package needed for this function to work. Please install it.",
-         call. = FALSE)
-  }
 
   if(!all(col_attr %in% names(dm$columns) )) {
     stop("Not all col_attr in data model column attributes.")
@@ -251,18 +249,19 @@ dm_create_graph <- function(dm, rankdir = "BT", graph_name = "Data Model",
     warning("The number of tables to render is 0.")
   }
   graph <-
-    DiagrammeR::create_graph(
+    list(
       graph_attrs = sprintf('rankdir=%s tooltip="%s" %s', rankdir, graph_name, graph_attrs),
       node_attrs = sprintf('margin=0 fontcolor = "#444444" %s', node_attrs),
-      nodes_df = do.call(DiagrammeR::create_nodes, g_list$nodes),
-      edges_df = if(!is.null(g_list$edges)) do.call(DiagrammeR::create_edges, g_list$edges) else NULL,
+      nodes_df = g_list$nodes,
+      edges_df = g_list$edges,
       edge_attrs = c('color = "#555555"',"arrowsize = 1", edge_attrs)
     )
+  class(graph) <- c("datamodelr_graph", class(graph))
 
   # re-create dot code for data model
   # (DiagrammeR does not support yet the HTML labels and clusters (v.0.8))
   graph$dot_code <- dot_graph(graph)
-  sessionInfo()
+
   graph
 
 }
@@ -270,34 +269,32 @@ dm_create_graph <- function(dm, rankdir = "BT", graph_name = "Data Model",
 
 #' Render graph
 #'
-#' A wrapper around DiagrammeR::render_graph
+#' Using DiagrammeR to render datamodelr graph object
 #'
-#' @param graph	a \pkg{DiagrammeR} dgr_graph object
-#' @param  output	string specifying the output type; graph (the default) renders
-#'   the graph using the grViz function,
-#'   vivagraph renders the graph using the vivagraph function,
-#'   visNetwork renders the graph using the visnetwork function,
-#'   DOT outputs DOT code for the graph, and SVG provides SVG code for the rendered graph.
-#' @param layout	a string specifying a layout type for a vivagraph rendering of the graph,
-#'   either forceDirected or constant.
+#' @param graph	a graph object
 #' @param width	an optional parameter for specifying the width of the resulting
 #' graphic in pixels.
 #' @param height an optional parameter for specifying the height of the resulting graphic in pixels.
 #' @export
-dm_render_graph <- function (graph, output = "graph", layout = NULL, width = NULL,
-                       height = NULL) {
+dm_render_graph <- function (graph, width = NULL, height = NULL) {
+
+  if( !requireNamespace("DiagrammeR", quietly = TRUE)) {
+    stop("DiagrammeR package needed for this function to work. Please install it.",
+         call. = FALSE)
+  }
+
 
-  if(substring(graph$dot_code, 1, 11) != "#data_model") {
+  if(is.null(graph$dot_code)) {
     graph$dot_code <- dot_graph(graph)
   }
 
-  DiagrammeR::render_graph(graph, output, layout, width, height)
+  DiagrammeR::grViz(graph$dot_code, allow_subst = FALSE, width, height)
 }
 
 
 dot_graph <- function(graph) {
 
-  graph_type <- ifelse(graph$directed, "digraph", "graph")
+  graph_type <- "digraph"
 
   dot_attr <- paste0(
     sprintf("graph [%s]\n\n", paste(graph$graph_attrs, collapse = ", ")),
@@ -306,7 +303,7 @@ dot_graph <- function(graph) {
   )
   segments <- unique(graph$nodes_df$segment)
   segments <- segments[!is.na(segments)]
-  segments <- setNames(1:(length(segments)), segments)
+  segments <- stats::setNames(1:(length(segments)), segments)
 
   dot_nodes <- sapply(seq_len(nrow(graph$nodes_df)), function(n) {
     node <- graph$nodes_df[n,]
diff --git a/man/dm_render_graph.Rd b/man/dm_render_graph.Rd
index b326b21..fc1b5e6 100644
--- a/man/dm_render_graph.Rd
+++ b/man/dm_render_graph.Rd
@@ -4,20 +4,10 @@
 \alias{dm_render_graph}
 \title{Render graph}
 \usage{
-dm_render_graph(graph, output = "graph", layout = NULL, width = NULL,
-  height = NULL)
+dm_render_graph(graph, width = NULL, height = NULL)
 }
 \arguments{
-\item{graph}{a \pkg{DiagrammeR} dgr_graph object}
-
-\item{output}{string specifying the output type; graph (the default) renders
-the graph using the grViz function,
-vivagraph renders the graph using the vivagraph function,
-visNetwork renders the graph using the visnetwork function,
-DOT outputs DOT code for the graph, and SVG provides SVG code for the rendered graph.}
-
-\item{layout}{a string specifying a layout type for a vivagraph rendering of the graph,
-either forceDirected or constant.}
+\item{graph}{a graph object}
 
 \item{width}{an optional parameter for specifying the width of the resulting
 graphic in pixels.}
@@ -25,6 +15,6 @@ graphic in pixels.}
 \item{height}{an optional parameter for specifying the height of the resulting graphic in pixels.}
 }
 \description{
-A wrapper around DiagrammeR::render_graph
+Using DiagrammeR to render datamodelr graph object
 }
 
diff --git a/tests/testthat/test_graph.R b/tests/testthat/test_graph.R
index c7f1b21..e71c431 100644
--- a/tests/testthat/test_graph.R
+++ b/tests/testthat/test_graph.R
@@ -37,7 +37,7 @@ test_that("Example - simple", {
     ")
 
   graph <- dm_create_graph(dm)
-  expect_is(graph, "dgr_graph")
+  expect_is(graph, "datamodelr_graph")
 })
 
 test_that("Example - graph - example.yml", {
@@ -45,5 +45,5 @@ test_that("Example - graph - example.yml", {
   file_path <- system.file("samples/example.yml", package = "datamodelr")
   dm <- dm_read_yaml(file_path)
   graph <- dm_create_graph(dm)
-  expect_is(graph, "dgr_graph")
+  expect_is(graph, "datamodelr_graph")
 })