Assume that $AB=BA^2$.
Hint: We can rewrite the equation as $$ A^2 = B^{-1}AB. $$ So, $A^2$ and $A$ are similar, which means that they must have the same eigenvalues. So, the eigenvalues of $A$ satisfy either $\lambda_1^2 = \lambda_1, \lambda_2^2 = \lambda_2$ or $\lambda_1^2 = \lambda_2, \lambda_2^2 = \lambda_1$.
In either case: if $A$ is diagonalizable, we should find that $A^3 = I$.
If $A$ is not diagonalizable with repeated eigenvalue $1$, then it must be similar to the Jordan block $$ J = \pmatrix{1&1\\0&1}. $$
(Mildly incomplete)
Taking the determinant of both sides, $\det A \det B = \det B (\det A)^2$. Thus, $\det A = 1$.
Next, note that $B^{-1}AB =A^2$. Taking trace of both sides, and noting that $\mbox{tr}(P^{-1}QP) = Q$ for any pair of matrices $P,Q$ , we get $\mbox{tr}{A} = \mbox{tr}{A^2}$.
We know that the eigenvalues of $A^2$ are exactly the squares of the eigenvalues of $A$. Let $\lambda_1 , \lambda_2$ be the (complex) eigenvalues of $A$, then $\lambda_1\lambda_2 = 1$, and $\lambda_1^2 + \lambda_2^2 = \lambda_1 + \lambda_2$.
This has three solutions : $(\lambda_1, \lambda_2) = (1,1) , (\omega , \omega^2),(\omega^2 , \omega)$ where $\omega^3 = 1$ is a complex cube root of unity.
Consider the latter two cases. Here, $A$ is diagonalizable since it has distinct eigenvalues. Furthermore, by the Cayley-Hamilton theorem, we get $A^3 = I$, since each eigenvalue satisfies this. This corresponds to case $1$.
In the first case, $A$ has the same eigenvalues $1,1$. Thus, we have $(A-I)^2 = 0$. Now, we either have $A - I = 0$ or $A- I \neq 0$. In the first case, $A$ is diagonalizable, and $A^3 = I$ and we are back to case $1$.
Elsewise, if $A \neq I$, we get that $A$ can't be diagonalizable(why? because if $A$ were diagonalizable then $I$ is the only candidate diagonal matrix) , and $A^2 =2A-I$. Thus, we get $BA^2 = 2BA - B$. From here, I am unable to see why $B$ must be diagonalizable, but I hope this should also go through.
Here is an argument to show that $B$ is diagonalizable if $A^3\ne I$. Then by @Omnomnomnom's answer, $A$ is similar to $\pmatrix{1 & 1 \\ 0 & 1}$.
Using the Jordan normal form of $A$, we need to consider the case $A=\pmatrix{1 & 1 \\ 0 & 1}$ only. Then $B$ satisfies $AB=BA^2$, i.e., $$ \pmatrix{1 & 1 \\ 0 & 1} B = B \pmatrix{1 &2 \\ 0 & 1}. $$ Writing $B=\pmatrix{ a & b \\ c & d}$, this yields the conditions $c=0$, $d=2a$, and $B$ is of the form $$ B=\pmatrix{ a & b \\ 0 & 2a}, $$ which is clearly diagonalizable: the eigenvalues are $a,2a$. Since $B$ is assumed to be invertible, $a\ne0$, and $B$ has two distinct eigenvalues..