Error `Cannot import name 'wrappers' from 'tensorflow.python.keras.layers'`?

The code is giving the following error message

• Cannot import name 'wrappers' from 'tensorflow.python.keras.layers' – and ImportError: graphviz or pydot are not available.

Even after installing the graphviz and pydot using the !apt-get -qq install -y graphviz && pip install pydot still not able to genrate model.png.

i m running the following code in colab

try:
  # %tensorflow_version only exists in Colab.
  %tensorflow_version 2.x
except Exception:
  pass

import tensorflow as tf
from tensorflow.python.keras.utils.vis_utils import plot_model
import pydot
from tensorflow.keras.models import Model

def build_model_with_sequential():
    
    # instantiate a Sequential class and linearly stack the layers of your model
    seq_model = tf.keras.models.Sequential([tf.keras.layers.Flatten(input_shape=(28, 28)),
                                            tf.keras.layers.Dense(128, activation=tf.nn.relu),
                                            tf.keras.layers.Dense(10, activation=tf.nn.softmax)])
    return seq_model

def build_model_with_functional():
    
    # instantiate the input Tensor
    input_layer = tf.keras.Input(shape=(28, 28))
    
    # stack the layers using the syntax: new_layer()(previous_layer)
    flatten_layer = tf.keras.layers.Flatten()(input_layer)
    first_dense = tf.keras.layers.Dense(128, activation=tf.nn.relu)(flatten_layer)
    output_layer = tf.keras.layers.Dense(10, activation=tf.nn.softmax)(first_dense)
    
    # declare inputs and outputs
    func_model = Model(inputs=input_layer, outputs=output_layer)
    
    return func_model

model = build_model_with_functional()
#model = build_model_with_sequential()

# Plot model graph
plot_model(model, show_shapes=True, show_layer_names=True, to_file='model.png')

ERROR

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
<ipython-input-5-aaa8f6f7cc6e> in <module>
      3 
      4 # Plot model graph
----> 5 plot_model(model, show_shapes=True, show_layer_names=True, to_file='model.png')

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/vis_utils.py in plot_model(model, to_file, show_shapes, show_dtype, show_layer_names, rankdir, expand_nested, dpi)
    327       rankdir=rankdir,
    328       expand_nested=expand_nested,
--> 329       dpi=dpi)
    330   to_file = path_to_string(to_file)
    331   if dot is None:

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/vis_utils.py in model_to_dot(model, show_shapes, show_dtype, show_layer_names, rankdir, expand_nested, dpi, subgraph)
     96     ImportError: if graphviz or pydot are not available.
     97   """
---> 98   from tensorflow.python.keras.layers import wrappers
     99   from tensorflow.python.keras.engine import sequential
    100   from tensorflow.python.keras.engine import functional

ImportError: cannot import name 'wrappers' from 'tensorflow.python.keras.layers' (/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/layers/__init__.py)

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
-------------------------------------------------------------------------

>Solution :

Seems to be some import conflicts. Rather use tf.keras for everything:

try:
  # %tensorflow_version only exists in Colab.
  %tensorflow_version 2.x
except Exception:
  pass

import tensorflow as tf
import pydot

def build_model_with_sequential():
    
    # instantiate a Sequential class and linearly stack the layers of your model
    seq_model = tf.keras.models.Sequential([tf.keras.layers.Flatten(input_shape=(28, 28)),
                                            tf.keras.layers.Dense(128, activation=tf.nn.relu),
                                            tf.keras.layers.Dense(10, activation=tf.nn.softmax)])
    return seq_model

def build_model_with_functional():
    
    # instantiate the input Tensor
    input_layer = tf.keras.Input(shape=(28, 28))
    
    # stack the layers using the syntax: new_layer()(previous_layer)
    flatten_layer = tf.keras.layers.Flatten()(input_layer)
    first_dense = tf.keras.layers.Dense(128, activation=tf.nn.relu)(flatten_layer)
    output_layer = tf.keras.layers.Dense(10, activation=tf.nn.softmax)(first_dense)
    
    # declare inputs and outputs
    func_model = tf.keras.Model(inputs=input_layer, outputs=output_layer)
    
    return func_model

model = build_model_with_functional()
tf.keras.utils.plot_model(model, show_shapes=True, show_layer_names=True, to_file='model.png')