Use keras.Model for VAEImageDecoder and follows the coding style …

…in `VAEAttention`

keras_nlp/src/models/stable_diffusion_v3/vae_attention.py

@@ -77,9 +77,6 @@ def build(self, input_shape):
         self.value_conv2d.build(input_shape)
         self.output_conv2d.build(input_shape)
 
-    def compute_output_shape(self, input_shape):
-        return input_shape
-
     def call(self, inputs, training=None):
         x = self.group_norm(inputs)
         query = self.query_conv2d(x)
@@ -112,10 +109,18 @@ def call(self, inputs, training=None):
         x = self.output_conv2d(x)
         if self.data_format == "channels_first":
             x = ops.transpose(x, (0, 3, 1, 2))
-        x = layers.Add()([x, inputs])
+        x = ops.add(x, inputs)
         return x
 
     def get_config(self):
         config = super().get_config()
-        config.update({"filters": self.filters, "groups": self.groups})
+        config.update(
+            {
+                "filters": self.filters,
+                "groups": self.groups,
+            }
+        )
         return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

keras_nlp/src/models/stable_diffusion_v3/vae_image_decoder.py

@@ -11,18 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import keras
 from keras import layers
 
-from keras_nlp.src.models.backbone import Backbone
 from keras_nlp.src.models.stable_diffusion_v3.vae_attention import VAEAttention
 from keras_nlp.src.utils.keras_utils import standardize_data_format
 
 
-class VAEImageDecoder(Backbone):
+class VAEImageDecoder(keras.Model):
     def __init__(
         self,
-        stackwise_num_filters=[512, 512, 256, 128],
-        stackwise_num_blocks=[3, 3, 3, 3],
+        stackwise_num_filters,
+        stackwise_num_blocks,
         output_channels=3,
         latent_shape=(None, None, 16),
         data_format=None,