diff --git a/DESCRIPTION b/DESCRIPTION
index fdc308d..44a9082 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,6 +1,6 @@
 Package: pixeltrix
 Title: Simple Interactive Pixel Art
-Version: 0.2.0
+Version: 0.2.1
 Authors@R: 
     person("Matt", "Dray", , "mwdray@gmail.com", role = c("aut", "cre"))
 Description: A very simple 'pixel art' tool that lets you click squares
diff --git a/NEWS.md b/NEWS.md
index ca6a9ed..936bed0 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,9 @@
+# pixeltrix 0.2.1
+
+* Added additional input checks to check-utils.R.
+* `draw_pixels()` now finds `n_states` (#24).
+* Allow for an increased number of states in `edit_pixels()` (#25).
+
 # pixeltrix 0.2.0
 
 * A named character of colours is now provided as an extra attribute to matrices output from `click_pixels()` (#3, #17, thanks @TimTaylor).
@@ -6,7 +12,6 @@
 * Updated and expanded function documentation and README (#21).
 * Reused input checks have been generalised into 'R/utils-check.R'.
 * Expanded tests to cover argument input errors.
-* 
 
 # pixeltrix 0.1.3
 
diff --git a/R/animate.R b/R/animate.R
index d4d4eda..98d8553 100644
--- a/R/animate.R
+++ b/R/animate.R
@@ -50,9 +50,11 @@ frame_pixels <- function(
     grid     = TRUE
 ) {
 
-  .check_n_numeric(n_rows, n_cols, n_states)
+  .check_n_arg_numeric(n_rows)
+  .check_n_arg_numeric(n_cols)
+  .check_n_arg_numeric(n_states)
   .check_colours_char(colours)
-  .check_colours_len(colours, n_states)
+  .check_colours_len(n_states, colours)
   .check_grid(grid)
 
   m_list  <- list()
@@ -129,26 +131,14 @@ gif_pixels <- function(
     ...
 ) {
 
-  if (
-    !is.list(frames) |
-    !all(sapply(frames, function(x) identical(dim(x), dim(frames[[1]]))))
-  ) {
-    stop(
-      "Argument 'frames' must be a list of matrices of the same dimensions ",
-      "(preferably produced by the frame_pixels() function).",
-      call. = FALSE
-    )
-  }
+  .check_frames_dims(frames)
+  .check_file_gif(file)
 
-  if (
-    !inherits(file, "character") |
-    length(file) != 1 |
-    tools::file_ext(file) != "gif"
-  ) {
-    stop(
-      "Argument 'file' must be a character-string filepath ending '.gif'.",
-      call. = FALSE
-    )
+  # Retrieve n_states from attributes or matrix values
+  if (!is.null(attr(frames[[1]], "colours"))) {
+    n_states <- length(attr(frames[[1]], "colours"))
+  } else if (is.null(attr(frames, "colours"))) {
+    n_states <- max(unique(unlist(frames))) + 1L
   }
 
   # If the first frame has a 'colours' attribute, then use these
@@ -163,7 +153,7 @@ gif_pixels <- function(
   }
 
   .check_colours_char(colours)
-  .check_colours_unique(frames, colours)
+  .check_colours_states(frames, n_states, colours)
 
   # Write to
   gifski::save_gif(
diff --git a/R/click.R b/R/click.R
index 9cb32d7..0e6d75f 100644
--- a/R/click.R
+++ b/R/click.R
@@ -48,27 +48,33 @@ click_pixels <- function(
     grid     = TRUE
 ) {
 
-  .check_n_numeric(n_rows, n_cols, n_states)
+  # Check inputs
+  .check_n_arg_numeric(n_rows)
+  .check_n_arg_numeric(n_cols)
+  .check_n_arg_numeric(n_states)
   .check_colours_char(colours)
-  .check_colours_len(colours, n_states)
+  .check_colours_len(n_states, colours)
   .check_grid(grid)
 
+  # Convert to integer if required
   n_rows   <- .convert_to_int(n_rows)
   n_cols   <- .convert_to_int(n_cols)
   n_states <- .convert_to_int(n_states)
 
+  # Generate a palette of gradated greys if colours not provided by user
   if (is.null(colours)) {
     get_greys <- grDevices::colorRampPalette(c("white", "grey20"))
-    colours   <- get_greys(n_states)  # gradated colours from white to dark grey
+    colours   <- get_greys(n_states)
   }
 
+  # Initiate matrix, draw, let user interact
   m <- matrix(0L, n_rows, n_cols)
-
   .plot_canvas(m, n_states, colours)
   if (grid) .add_grid(m)
   m <- .repeat_loop(m, n_states, colours, grid)
 
-  attr(m, "colours")  <- stats::setNames(colours, seq(0, n_states - 1))
+  # Add colours as an attribute to returned matrix
+  attr(m, "colours") <- stats::setNames(colours, seq(0, n_states - 1))
 
   m
 
@@ -134,58 +140,38 @@ edit_pixels <- function(
     grid     = TRUE
 ) {
 
+  # Check inputs
   .check_matrix(m)
   .check_grid(grid)
+  .check_n_arg_numeric(n_states, null_allowed = TRUE)
+  .check_n_states_size(m, n_states)
 
-  if (!is.null(n_states)) {
-    if (!is.numeric(n_states)) {
-      stop(
-        "Argument 'n_states' must be a numeric value or NULL.",
-        call. = FALSE
-      )
-    }
-  }
-
-  if (!is.null(n_states) && n_states < max(m + 1L)) {
-    stop(
-      "The number of states, 'n_states', can't be less than ",
-      "the maximum value in the provided matrix, 'm'.",
-      call. = FALSE
-    )
-  }
-
-  # Coerce n_states to integer, if provided
-  if (!is.null(n_states)) {
+  # Handle n_states
+  if (!is.null(n_states)) {  # if provided, convert to integer
     n_states <- as.integer(n_states)
+  } else if (is.null(n_states) & !is.null(attr(m, "colours"))) {  # via attribute
+    n_states <- length(attr(m, "colours"))
+  } else if (is.null(n_states) & is.null(attr(m, "colours"))) {  # via matrix
+    n_states <- max(unique(as.vector(m)) + 1L)
   }
 
-  # Otherwise get n_state from attributes
-  if (is.null(n_states) & !is.null(attr(m, "colours"))) {
-    n_states <- length(attr(m, "colours"))  # n colours, so n states
-  }
-
-  # Otherwise take n_states from content of input matrix
-  if (is.null(n_states) & is.null(attr(m, "colours"))) {
-    n_states <- length(unique(as.vector(m)))
-  }
-
-  # Take colours from attributes of input matrix, if present
-  if (is.null(colours) & !is.null(attr(m, "colours"))) {
+  # Handle colours if not provided
+  if (is.null(colours) & !is.null(attr(m, "colours"))) {  # via attribute
     colours <- attr(m, "colours")
-  }
-
-  # If no 'colours' attribute and colours is NULL, then choose gradated greys
-  if (is.null(colours)) {
+  } else if (is.null(colours)) {  # otherwise a grey palette
     get_greys <- grDevices::colorRampPalette(c("white", "grey20"))
-    colours   <- get_greys(n_states)  # gradated colours from white to dark grey
+    colours <- get_greys(n_states)
   }
 
-  # .check_colours_unique(m, colours)
+  # Check n_states and colours values match
+  .check_colours_states(m, n_states, colours)
 
+  # Draw matrix, let user interact
   .plot_canvas(m, n_states, colours)
   if (grid) .add_grid(m)
   m <- .repeat_loop(m, n_states, colours, grid)
 
+  # Add colours as an attribute to returned matrix
   attr(m, "colours") <- stats::setNames(colours, seq(0, n_states - 1))
 
   m
diff --git a/R/draw.R b/R/draw.R
index 9431803..2db74f3 100644
--- a/R/draw.R
+++ b/R/draw.R
@@ -25,20 +25,28 @@ draw_pixels <- function(m, colours = NULL) {
 
   .check_matrix(m)
 
+  # Retrieve n_states from attributes or matrix values
+  if (!is.null(attr(m, "colours"))) {
+    n_states <- length(attr(m, "colours"))
+  } else if (is.null(attr(m, "colours"))) {
+    n_states <- max(unique(as.vector(m))) + 1L
+  }
+
   # Take colours from attributes of input matrix, if present
   if (is.null(colours) & !is.null(attr(m, "colours"))) {
     colours <- attr(m, "colours")
   }
 
-  # If matrix has no 'colour' attribute, create gradated grey palette
-  if (is.null(colours)) {
+  # If matrix has no 'colours' attribute, create gradated grey palette
+  if (is.null(colours) & is.null(attr(m, "colours"))) {
     get_greys <- grDevices::colorRampPalette(c("white", "grey20"))
     colours   <- get_greys(n_states)  # gradated colours from white to dark grey
   }
 
-  .check_colours_unique(m, colours)
+  # Check number of colours provided
+  .check_colours_states(m, n_states, colours)
 
-  par_start <- graphics::par(mar = rep(0, 4))
+  par_start <- graphics::par(mar = rep(0, 4))  # set margins, store previous par
 
   graphics::image(
     t(m[nrow(m):1, ]),  # reverse matrix rows and transpose
diff --git a/R/utils-check.R b/R/utils-check.R
index f3aa39c..003dd70 100644
--- a/R/utils-check.R
+++ b/R/utils-check.R
@@ -9,14 +9,41 @@
 
 }
 
-.check_n_numeric <- function(n_rows, n_cols, n_states) {
+.check_n_arg_numeric <- function(n_arg, null_allowed = FALSE) {
 
-  if (
-    is.logical(n_rows) | is.logical(n_states) | is.logical(n_cols) |
-    !is.numeric(c(n_rows, n_cols, n_states))
-  ) {
+  if (!null_allowed) {  # used in click_pixels, where defaults are provided
+
+    if (is.logical(n_arg) | !is.numeric(c(n_arg))) {
+      stop(
+        "Argument '", deparse(substitute(n_arg)), "' must be numeric.",
+        call. = FALSE
+      )
+    }
+
+  }
+
+  if (null_allowed) {  # used in edit_pixels, where default n_states is NULL
+
+    if (!is.null(n_arg) && !is.numeric(n_arg)) {
+      stop(
+        "Argument '", deparse(substitute(n_arg)), "' must be numeric or NULL.",
+        call. = FALSE
+      )
+    }
+
+  }
+
+
+}
+
+.check_n_states_size <- function(m, n_states) {
+
+  m_max <- max(m + 1L)
+
+  if (!is.null(n_states) && n_states < m_max) {
     stop(
-      "Arguments 'n_rows', 'n_cols' and 'n_states' must be numeric values.",
+      "Argument 'n_states' (", n_states, " detected) must be equal or greater ",
+      "than the maximum value in the matrix 'm' (", m_max, " detected).",
       call. = FALSE
     )
   }
@@ -36,33 +63,58 @@
 
 }
 
-.check_colours_len <- function(colours, n_states) {
+.check_colours_len <- function(n_states, colours) {
 
   if (!is.null(colours) && (length(colours) != n_states)) {
     stop(
-      "Argument 'colours' must be a character vector of length 'n_states'.",
+      "Argument 'colours' (", length(colours), " values provided) must be a ",
+      "character vector of length 'n_states' (", n_states, ").",
       call. = FALSE
     )
   }
 
 }
 
-.check_colours_unique <- function(object, colours) {
+.check_colours_states <- function(object, n_states, colours) {
 
-  if (is.list(object)) {
-    object <- unlist(object)
-  }
+  if (is.null(n_states)) {
+
+    if (is.matrix(object)) {  # edit_pixels() is a matrix
+
+      colours_attr <- attr(object, "colours")
+
+      if (!is.null(colours_attr)) {
+        n_states <- length(colours_attr)
+      }
+
+      if (is.null(colours_attr)) {
+        object <- as.vector(object)
+        n_states <- max(unique(object)) + 1L
+      }
 
-  if (is.matrix(object)) (
-    object <- as.vector(object)
-  )
+    }
+
+    if (is.list(object)) {  # frame_pixels() input is a list
+
+      colours_attr <- attr(object[[1]], "colours")
+
+      if (!is.null(colours_attr)) {
+        n_states <- length(colours_attr)
+      }
+
+      if (is.null(colours_attr)) {
+        object <- unlist(object)
+        n_states <- max(unique(object)) + 1L
+      }
+
+    }
 
-  states_len <- length(unique(object))
+  }
 
-  if (length(colours) != states_len) {
+  if (length(colours) != n_states) {
     stop(
-      "Length of argument 'colours' should match the number of unique ",
-      "pixel states (", states_len, ").",
+      "Number of colours (", length(colours), " detected) should match ",
+      "the number of pixel states (", n_states, " detected).",
       call. = FALSE
     )
   }
@@ -79,3 +131,33 @@
   }
 
 }
+
+.check_frames_dims <- function(frames) {
+
+  if (
+    !is.list(frames) |
+    !all(sapply(frames, function(frame) identical(dim(frame), dim(frames[[1]]))))
+  ) {
+    stop(
+      "Argument 'frames' must be a list of matrices of the same dimensions ",
+      "(preferably produced by the frame_pixels() function).",
+      call. = FALSE
+    )
+  }
+
+}
+
+.check_file_gif <- function(file) {
+
+  if (
+    !inherits(file, "character") |
+    length(file) != 1 |
+    tools::file_ext(file) != "gif"
+  ) {
+    stop(
+      "Argument 'file' must be a character-string filepath ending '.gif'.",
+      call. = FALSE
+    )
+  }
+
+}
diff --git a/R/utils-click.R b/R/utils-click.R
index 8e9fe7b..560dc10 100644
--- a/R/utils-click.R
+++ b/R/utils-click.R
@@ -14,7 +14,6 @@
     grDevices::dev.off()
 
     .plot_canvas(m, n_states, colours)
-
     if (grid) .add_grid(m)
 
   }
@@ -151,13 +150,22 @@
 
 }
 
-# Force numeric inputs to integer
-.convert_to_int <- function(n) {
+# Force numeric inputs to integer with warning
+.convert_to_int <- function(n_arg) {
+
+  if (!is.null(n_arg) && is.numeric(n_arg)) {
+
+    if (n_arg != as.integer(n_arg)) {
+      warning(
+        "Argument '", deparse(substitute(n_arg)), "' converted to integer.",
+        call. = FALSE
+      )
+    }
+
+    n_arg <- as.integer(n_arg)
 
-  if (!is.null(n) && is.numeric(n)) {
-    n <- as.integer(n)
   }
 
-  n
+  n_arg
 
 }
diff --git a/README.Rmd b/README.Rmd
index 1ef3bd3..07defa4 100644
--- a/README.Rmd
+++ b/README.Rmd
@@ -143,7 +143,7 @@ Some known limitations are that:
 
 * you can only click one pixel at a time
 * each click only increments the pixel state by 1
-* you can't change the number of states nor the colour palette on the fly
+* you can't change the number of pixel states on the fly, nor the colour palette
 * I've built and tested the package in RStudio on macOS, but it may not work perfectly for your platform and editor
 
 Issues or pull requests are welcome.
diff --git a/README.md b/README.md
index 27c4079..91f2291 100644
--- a/README.md
+++ b/README.md
@@ -170,8 +170,8 @@ Some known limitations are that:
 
 -   you can only click one pixel at a time
 -   each click only increments the pixel state by 1
--   you can’t change the number of states nor the colour palette on the
-    fly
+-   you can’t change the number of pixel states on the fly, nor the
+    colour palette
 -   I’ve built and tested the package in RStudio on macOS, but it may
     not work perfectly for your platform and editor
 
diff --git a/man/logo.Rd b/man/logo.Rd
index c6df706..eb4b5d1 100644
--- a/man/logo.Rd
+++ b/man/logo.Rd
@@ -3,7 +3,7 @@
 \docType{data}
 \name{logo}
 \alias{logo}
-\title{Example Sprite Matrix: 'pixeltrix' Package Logo}
+\title{Example Sprite Matrix: 'pixeltrix' Package Logo Text}
 \format{
 A matrix with 10 rows and 21 columns, taking integers from 0 to 2.
 }
@@ -11,8 +11,8 @@ A matrix with 10 rows and 21 columns, taking integers from 0 to 2.
 logo
 }
 \description{
-A matrix output from \code{\link{frame_pixels}} that represent the
-pixels of the pixeltrix package logo, which is the word 'pixeltrix' written
+A matrix output from \code{\link{click_pixels}} that represents the
+text of the pixeltrix package logo, which is the word 'pixeltrix' written
 in a font with letters that are 3 by 3 pixels.
 }
 \examples{
diff --git a/man/pixeltrix-package.Rd b/man/pixeltrix-package.Rd
index e597219..3a76c48 100644
--- a/man/pixeltrix-package.Rd
+++ b/man/pixeltrix-package.Rd
@@ -6,6 +6,8 @@
 \alias{pixeltrix-package}
 \title{pixeltrix: Simple Interactive Pixel Art}
 \description{
+\if{html}{\figure{logo.png}{options: align='right' alt='logo' width='120'}}
+
 A very simple 'pixel art' tool that lets you click squares interactively in a plot window and returns your final image as a matrix. Provide multiple frames to create an animated gif.
 }
 \seealso{