keras-hub-nightly 0.16.1.dev202410030339__py3-none-any.whl → 0.16.1.dev202410040340__py3-none-any.whl

keras_hub/src/version_utils.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.16.1.dev202410030339"
+__version__ = "0.16.1.dev202410040340"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.16.1.dev202410030339.dist-info → keras_hub_nightly-0.16.1.dev202410040340.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: keras-hub-nightly
-Version: 0.16.1.dev202410030339
+Version: 0.16.1.dev202410040340
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

{keras_hub_nightly-0.16.1.dev202410030339.dist-info → keras_hub_nightly-0.16.1.dev202410040340.dist-info}/RECORD

@@ -1,15 +1,15 @@
 keras_hub/__init__.py,sha256=QGdXyHgYt6cMUAP1ebxwc6oR86dE0dkMxNy2eOCQtFo,855
 keras_hub/api/__init__.py,sha256=spMxsgqzjpeuC8rY4WP-2kAZ2qwwKRSbFwddXgUjqQE,524
 keras_hub/api/bounding_box/__init__.py,sha256=T8R_X7BPm0et1xaZq8565uJmid7dylsSFSj4V-rGuFQ,1097
-keras_hub/api/layers/__init__.py,sha256=P1Zn4sjTx1OnmlRyX8-QRxSe-2gkvyQ-90BzCjqr3oU,2227
+keras_hub/api/layers/__init__.py,sha256=XImD0tHdnDR1a7q3u-Pw-VRMASi9sDtrV6hr2beVYTw,2331
 keras_hub/api/metrics/__init__.py,sha256=So8Ec-lOcTzn_UUMmAdzDm8RKkPu2dbRUm2px8gpUEI,381
-keras_hub/api/models/__init__.py,sha256=dyancDilnzbHByiTYQNhqfm6JFeZH_DKHl4PZuvWoA0,13994
+keras_hub/api/models/__init__.py,sha256=sMfVpa2N90cG7qjkwSEI_x3uCvZNwQqFbedn5wcUzbE,14311
 keras_hub/api/samplers/__init__.py,sha256=n-_SEXxr2LNUzK2FqVFN7alsrkx1P_HOVTeLZKeGCdE,730
 keras_hub/api/tokenizers/__init__.py,sha256=_f-r_cyUM2fjBB7iO84ThOdqqsAxHNIewJ2EBDlM0cA,2524
 keras_hub/api/utils/__init__.py,sha256=Gp1E6gG-RtKQS3PBEQEOz9PQvXkXaJ0ySGMqZ7myN7A,215
 keras_hub/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_hub/src/api_export.py,sha256=9pQZK27JObxWZ96QPLBp1OBsjWigh1iuV6RglPGMRk0,1499
-keras_hub/src/version_utils.py,sha256=Q7sWkBqN11QJLqnWmwU9B2XhXWRKLr1vv199Ud-cp4A,222
+keras_hub/src/version_utils.py,sha256=9pXCOZsdoqgw8IovZaLqnG-5LTgH73BlKBZXdOvDpYc,222
 keras_hub/src/bounding_box/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_hub/src/bounding_box/converters.py,sha256=a5po8DBm87oz2EXfi-0uEZHCMlCJPIb4-MaZIdYx3Dg,17865
 keras_hub/src/bounding_box/formats.py,sha256=YmskOz2BOSat7NaE__J9VfpSNGPJJR0znSzA4lp8MMI,3868
@@ -56,13 +56,13 @@ keras_hub/src/models/feature_pyramid_backbone.py,sha256=clEW-TTQSVJ_5qFNdDF0iABk
 keras_hub/src/models/image_classifier.py,sha256=yt6cjhPfqs8A_eWXBsXdXFzn-aRgH2rVHUq7Zu7CyK8,7804
 keras_hub/src/models/image_classifier_preprocessor.py,sha256=YdewYfMPVHI7gdhbBI-zVcy4NSfg0bhiOHTmGEKoOYI,2668
 keras_hub/src/models/image_segmenter.py,sha256=C1bzIO59pG58iist5GLn_qnlotDpcAVxPV_8a68BkAc,2876
-keras_hub/src/models/image_segmenter_preprocessor.py,sha256=vJoZc1OebQWlqUP_ygCS7P1Pyq1KmmUc-0V_-maDzX4,2658
+keras_hub/src/models/image_segmenter_preprocessor.py,sha256=IMmVJWBc0VZ1-5jLmFmmwQ3q_oQnhIfCE9A6nS1ss8Q,3743
 keras_hub/src/models/masked_lm.py,sha256=uXO_dE_hILlOC9jNr6oK6IHi9IGUqLyNGvr6nMt8Rk0,3576
 keras_hub/src/models/masked_lm_preprocessor.py,sha256=g8vrnyYwqdnSw5xppROM1Gzo_jmMWKYZoQCsKdfrFKk,5656
 keras_hub/src/models/preprocessor.py,sha256=pJodz7KRVncvsC3o4qoKDYWP2J0a8E9CD6oVGYgJzIM,7970
 keras_hub/src/models/seq_2_seq_lm.py,sha256=w0gX-5YZjatfvAJmFAgSHyqS_BLqc8FF8DPLGK8mrgI,1864
 keras_hub/src/models/seq_2_seq_lm_preprocessor.py,sha256=HUHRbWRG5SF1pPpotGzBhXlrMh4pLFxgAoFk05FIrB4,9687
-keras_hub/src/models/task.py,sha256=MfrzIoj3XFaRiNlUg-K6D8l-ylWfpzBjjmSy-guXtG8,13935
+keras_hub/src/models/task.py,sha256=2iapEFHvzyl0ASlH6yzQA2OHSr1jV1V-pLtagHdBncQ,14416
 keras_hub/src/models/text_classifier.py,sha256=VBDvQUHTpJPqKp7A4VAtm35FOmJ3yMo0DW6GdX67xG0,4159
 keras_hub/src/models/text_classifier_preprocessor.py,sha256=EoWp-GHnaLnAKTdAzDmC-soAV92ATF3QozdubdV2WXI,4722
 keras_hub/src/models/text_to_image.py,sha256=N42l1W8YEUBHOdGiT4BQNqzTpgjB2O5dtLU5FbKpMy0,10792
@@ -115,6 +115,13 @@ keras_hub/src/models/deberta_v3/deberta_v3_tokenizer.py,sha256=zEMCLy9eCiBEpA_xM
 keras_hub/src/models/deberta_v3/disentangled_attention_encoder.py,sha256=JbdQV1ZHFX4_PcJhemuHQ5YPkJC-XujgbTyjCrdL7nk,8556
 keras_hub/src/models/deberta_v3/disentangled_self_attention.py,sha256=3l7Hy7JfiZDDDFE6uTqSuFjg802kXD9acA7aHKRdzJk,13122
 keras_hub/src/models/deberta_v3/relative_embedding.py,sha256=3WIQ1nWcEhfWF0U9DcKyYz3AAhO3Pmg7ykpzrYe0Jgw,2886
+keras_hub/src/models/deeplab_v3/__init__.py,sha256=FHAUPM4a1DJj4EsNTbYEd1riNq__uHU4eB3t3Z1zgj0,288
+keras_hub/src/models/deeplab_v3/deeplab_v3_backbone.py,sha256=WyFhuLcjFPFVuNL09bvW_Jsja7WjEw3zazKCOwbFDTM,7709
+keras_hub/src/models/deeplab_v3/deeplab_v3_image_converter.py,sha256=mRkH3HdhpV0fCcQcVXEvIX7SNk-bAMb3SAHzgK-FD5c,371
+keras_hub/src/models/deeplab_v3/deeplab_v3_image_segmeter_preprocessor.py,sha256=hR9S6lNYamY0EBDBo3e1qTCiwtftmLXrN-UYuzfw5Io,581
+keras_hub/src/models/deeplab_v3/deeplab_v3_layers.py,sha256=qmEiolOOriLAojXB67xXW9IOo717kaCGeDVZJLaGY98,7834
+keras_hub/src/models/deeplab_v3/deeplab_v3_presets.py,sha256=JE-Uv_CXDVnwNkTcy5GtzrIQXvcXSP7fdVpfhmk8qhg,122
+keras_hub/src/models/deeplab_v3/deeplab_v3_segmenter.py,sha256=tiMDcCFANHMUx3IVtW3r1P_JTazgPPsbW4IktIytKEU,3650
 keras_hub/src/models/densenet/__init__.py,sha256=r7StyamnWeeZxOk9r4ZYNbS_YVhu9YGPyXhNxljvdPg,269
 keras_hub/src/models/densenet/densenet_backbone.py,sha256=dN9lUwKzO3E2HthNV2x54ozeBEQ0ilNs5uYHshFQpT0,6723
 keras_hub/src/models/densenet/densenet_image_classifier.py,sha256=ptuV6PwgoUpmrSPqX7-a85IpWsElwcCv_G5IVkP9E_Q,530
@@ -263,14 +270,13 @@ keras_hub/src/models/sam/sam_presets.py,sha256=AfGUKNOkz0G11OMYqVebXKgEBar1qpIkA
 keras_hub/src/models/sam/sam_prompt_encoder.py,sha256=2foB7900QbzQfZjBo335XYsdjmhOnVT8fKD1CubJNVE,11801
 keras_hub/src/models/sam/sam_transformer.py,sha256=L2bdxdc2RUF1juRZ0F0Z6r0gTva1sUwEdjItJmKKf6w,5730
 keras_hub/src/models/stable_diffusion_3/__init__.py,sha256=ZKYQuaRObyhKq8GVAHmoRvlXp6FpU8ChvutVCHyXKuc,343
-keras_hub/src/models/stable_diffusion_3/flow_match_euler_discrete_scheduler.py,sha256=9SLbOpAv50q8yv8I6H4DHbsIgwNo8TJmwZfAH8Ew6Zw,2827
+keras_hub/src/models/stable_diffusion_3/flow_match_euler_discrete_scheduler.py,sha256=vtVhieAv277mAiZj7Kvvqg_Ba7klfQxZVk4PPxNNQ0s,3062
 keras_hub/src/models/stable_diffusion_3/mmdit.py,sha256=ntmxjDJtZbHDGVPPAnasVZyoOTp5bbMPhxM30SYmpoQ,25711
-keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py,sha256=9rWSG0C23_pwN1pymZbial3GX_UM4tmDLXtB4kTQ04w,22599
-keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_presets.py,sha256=gfF5ZOhJx03IQTPnb2Nf65i3pNz-fQlhdAJ3DjKHHZ8,658
-keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py,sha256=XH4osHG9EE1sJpfj7rf0bCqrIHpeXaswFoEojWnE0pw,4419
+keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py,sha256=D-U5T6UYKzraHLAgMa-LLcd40ZmX_5rmlybawT4ooHY,21398
+keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_presets.py,sha256=sBGVRFd-bYxcqydydOB70XpOtpTt6AVrTR3LV-LBFXY,662
+keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py,sha256=8nP3ejDOd1hqjYXJzbri62PgtclxGydw-8bw-qHIPdc,4414
 keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image_preprocessor.py,sha256=TB0KESt5dnFYiS292PbzB0LdiH23AD6aTSTGmQEuzGM,2742
 keras_hub/src/models/stable_diffusion_3/t5_encoder.py,sha256=oV7P1uwCKdGiD93zXq7kmqX0elMZQU4UvBa8wg6P1hs,5113
-keras_hub/src/models/stable_diffusion_3/vae_image_decoder.py,sha256=j4nrvfYhW-4njhhk4PFf-bWQF-EHzplvaT15Q7s5Pb4,10056
 keras_hub/src/models/t5/__init__.py,sha256=OWyoUeDY3v4DnO8Ry02DWV1bNSVGcC89PF9oCftyi1s,233
 keras_hub/src/models/t5/t5_backbone.py,sha256=AtE2VudEUkm7hE3p6JP_CfEAjt4pwgSKOBQ0B0BggQc,10258
 keras_hub/src/models/t5/t5_layer_norm.py,sha256=R8KPHFOq9N3SD013WjtloLWRzaEMNEyY0fbViNEFVXQ,630
@@ -279,6 +285,9 @@ keras_hub/src/models/t5/t5_preprocessor.py,sha256=UVOnCHUJF_MBcOyfR9G9oeRUEoN3Xo
 keras_hub/src/models/t5/t5_presets.py,sha256=95zU4cTNEZMH2yiCLptA9zhu2D4mE1Cay18K91nt7jM,3005
 keras_hub/src/models/t5/t5_tokenizer.py,sha256=pLTu15JeYSpVmy-2600vBc-Mxn_uHyTKts4PI2MxxBM,2517
 keras_hub/src/models/t5/t5_transformer_layer.py,sha256=uDeP84F1x7xJxki5iKe12Zn6eWD_4yVjoFXMuod-a3A,5347
+keras_hub/src/models/vae/__init__.py,sha256=i3UaSW4IJf76O7lSPE1dyxOVjuHx8iAYKivqvUbDHOw,62
+keras_hub/src/models/vae/vae_backbone.py,sha256=aYf1sGteFJ7FyR3X8Ek6QBjAT5GjRtQTK2jXhYVJeM4,6671
+keras_hub/src/models/vae/vae_layers.py,sha256=N83CYM1zgbl1EIjAOs3cFCkJEdxvbXkgM9ghKyljFAg,27752
 keras_hub/src/models/vgg/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_hub/src/models/vgg/vgg_backbone.py,sha256=QnEDKn5n9bA9p3nvt5fBHnAssvnLxR0qv-oB372Ts0U,3702
 keras_hub/src/models/vgg/vgg_image_classifier.py,sha256=Dtq_HIJP6fHe8m7ZVLVn8IbHEsVMFWLvWMmn8TU1ntw,6600
@@ -350,7 +359,7 @@ keras_hub/src/utils/transformers/convert_mistral.py,sha256=kVhN9h1ZFVhwkNW8p3wnS
 keras_hub/src/utils/transformers/convert_pali_gemma.py,sha256=B1leeDw96Yvu81hYumf66hIid07k5NLqoeWAJgPnaLs,10649
 keras_hub/src/utils/transformers/preset_loader.py,sha256=GS44hZUuGQCtzsyn8z44ZpHdftd3DFemwV2hx2bQa-U,2738
 keras_hub/src/utils/transformers/safetensor_utils.py,sha256=rPK-Uw1CG0DX0d_UAD-r2cG9fw8GI8bvAlrcXfQ9g4c,3323
-keras_hub_nightly-0.16.1.dev202410030339.dist-info/METADATA,sha256=tLxESmpHL96pjwqK1gteBF1IdJ_CKtgBOvGEIG9gfyU,7458
-keras_hub_nightly-0.16.1.dev202410030339.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-keras_hub_nightly-0.16.1.dev202410030339.dist-info/top_level.txt,sha256=N4J6piIWBKa38A4uV-CnIopnOEf8mHAbkNXafXm_CuA,10
-keras_hub_nightly-0.16.1.dev202410030339.dist-info/RECORD,,
+keras_hub_nightly-0.16.1.dev202410040340.dist-info/METADATA,sha256=BuGQEiANSPxHuXnmW4K_EEYaGBE2ENMMSMARmCthK70,7458
+keras_hub_nightly-0.16.1.dev202410040340.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+keras_hub_nightly-0.16.1.dev202410040340.dist-info/top_level.txt,sha256=N4J6piIWBKa38A4uV-CnIopnOEf8mHAbkNXafXm_CuA,10
+keras_hub_nightly-0.16.1.dev202410040340.dist-info/RECORD,,

keras_hub/src/models/stable_diffusion_3/vae_image_decoder.py

@@ -1,320 +0,0 @@
-import math
-
-from keras import layers
-from keras import ops
-
-from keras_hub.src.models.backbone import Backbone
-from keras_hub.src.utils.keras_utils import standardize_data_format
-
-
-class VAEAttention(layers.Layer):
-    def __init__(self, filters, groups=32, data_format=None, **kwargs):
-        super().__init__(**kwargs)
-        self.filters = filters
-        self.data_format = standardize_data_format(data_format)
-        gn_axis = -1 if self.data_format == "channels_last" else 1
-
-        self.group_norm = layers.GroupNormalization(
-            groups=groups,
-            axis=gn_axis,
-            epsilon=1e-6,
-            dtype="float32",
-            name="group_norm",
-        )
-        self.query_conv2d = layers.Conv2D(
-            filters,
-            1,
-            1,
-            data_format=self.data_format,
-            dtype=self.dtype_policy,
-            name="query_conv2d",
-        )
-        self.key_conv2d = layers.Conv2D(
-            filters,
-            1,
-            1,
-            data_format=self.data_format,
-            dtype=self.dtype_policy,
-            name="key_conv2d",
-        )
-        self.value_conv2d = layers.Conv2D(
-            filters,
-            1,
-            1,
-            data_format=self.data_format,
-            dtype=self.dtype_policy,
-            name="value_conv2d",
-        )
-        self.softmax = layers.Softmax(dtype="float32")
-        self.output_conv2d = layers.Conv2D(
-            filters,
-            1,
-            1,
-            data_format=self.data_format,
-            dtype=self.dtype_policy,
-            name="output_conv2d",
-        )
-
-        self.groups = groups
-        self._inverse_sqrt_filters = 1.0 / math.sqrt(float(filters))
-
-    def build(self, input_shape):
-        self.group_norm.build(input_shape)
-        self.query_conv2d.build(input_shape)
-        self.key_conv2d.build(input_shape)
-        self.value_conv2d.build(input_shape)
-        self.output_conv2d.build(input_shape)
-
-    def call(self, inputs, training=None):
-        x = self.group_norm(inputs)
-        query = self.query_conv2d(x)
-        key = self.key_conv2d(x)
-        value = self.value_conv2d(x)
-
-        if self.data_format == "channels_first":
-            query = ops.transpose(query, (0, 2, 3, 1))
-            key = ops.transpose(key, (0, 2, 3, 1))
-            value = ops.transpose(value, (0, 2, 3, 1))
-        shape = ops.shape(inputs)
-        b = shape[0]
-        query = ops.reshape(query, (b, -1, self.filters))
-        key = ops.reshape(key, (b, -1, self.filters))
-        value = ops.reshape(value, (b, -1, self.filters))
-
-        # Compute attention.
-        query = ops.multiply(
-            query, ops.cast(self._inverse_sqrt_filters, query.dtype)
-        )
-        # [B, H0 * W0, C], [B, H1 * W1, C] -> [B, H0 * W0, H1 * W1]
-        attention_scores = ops.einsum("abc,adc->abd", query, key)
-        attention_scores = ops.cast(
-            self.softmax(attention_scores), self.compute_dtype
-        )
-        # [B, H2 * W2, C], [B, H0 * W0, H1 * W1] -> [B, H1 * W1 ,C]
-        attention_output = ops.einsum("abc,adb->adc", value, attention_scores)
-        x = ops.reshape(attention_output, shape)
-
-        x = self.output_conv2d(x)
-        if self.data_format == "channels_first":
-            x = ops.transpose(x, (0, 3, 1, 2))
-        x = ops.add(x, inputs)
-        return x
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "filters": self.filters,
-                "groups": self.groups,
-            }
-        )
-        return config
-
-    def compute_output_shape(self, input_shape):
-        return input_shape
-
-
-def apply_resnet_block(x, filters, data_format=None, dtype=None, name=None):
-    data_format = standardize_data_format(data_format)
-    gn_axis = -1 if data_format == "channels_last" else 1
-    input_filters = x.shape[gn_axis]
-
-    residual = x
-    x = layers.GroupNormalization(
-        groups=32,
-        axis=gn_axis,
-        epsilon=1e-6,
-        dtype="float32",
-        name=f"{name}_norm1",
-    )(x)
-    x = layers.Activation("swish", dtype=dtype)(x)
-    x = layers.Conv2D(
-        filters,
-        3,
-        1,
-        padding="same",
-        data_format=data_format,
-        dtype=dtype,
-        name=f"{name}_conv1",
-    )(x)
-    x = layers.GroupNormalization(
-        groups=32,
-        axis=gn_axis,
-        epsilon=1e-6,
-        dtype="float32",
-        name=f"{name}_norm2",
-    )(x)
-    x = layers.Activation("swish", dtype=dtype)(x)
-    x = layers.Conv2D(
-        filters,
-        3,
-        1,
-        padding="same",
-        data_format=data_format,
-        dtype=dtype,
-        name=f"{name}_conv2",
-    )(x)
-    if input_filters != filters:
-        residual = layers.Conv2D(
-            filters,
-            1,
-            1,
-            data_format=data_format,
-            dtype=dtype,
-            name=f"{name}_residual_projection",
-        )(residual)
-    x = layers.Add(dtype=dtype)([residual, x])
-    return x
-
-
-class VAEImageDecoder(Backbone):
-    """Decoder for the VAE model used in Stable Diffusion 3.
-
-    Args:
-        stackwise_num_filters: list of ints. The number of filters for each
-            stack.
-        stackwise_num_blocks: list of ints. The number of blocks for each stack.
-        output_channels: int. The number of channels in the output.
-        latent_shape: tuple. The shape of the latent image.
-        data_format: `None` or str. If specified, either `"channels_last"` or
-            `"channels_first"`. The ordering of the dimensions in the
-            inputs. `"channels_last"` corresponds to inputs with shape
-            `(batch_size, height, width, channels)`
-            while `"channels_first"` corresponds to inputs with shape
-            `(batch_size, channels, height, width)`. It defaults to the
-            `image_data_format` value found in your Keras config file at
-            `~/.keras/keras.json`. If you never set it, then it will be
-            `"channels_last"`.
-        dtype: `None` or str or `keras.mixed_precision.DTypePolicy`. The dtype
-            to use for the model's computations and weights.
-    """
-
-    def __init__(
-        self,
-        stackwise_num_filters,
-        stackwise_num_blocks,
-        output_channels=3,
-        latent_shape=(None, None, 16),
-        data_format=None,
-        dtype=None,
-        **kwargs,
-    ):
-        data_format = standardize_data_format(data_format)
-        gn_axis = -1 if data_format == "channels_last" else 1
-
-        # === Functional Model ===
-        latent_inputs = layers.Input(shape=latent_shape)
-
-        x = layers.Conv2D(
-            stackwise_num_filters[0],
-            3,
-            1,
-            padding="same",
-            data_format=data_format,
-            dtype=dtype,
-            name="input_projection",
-        )(latent_inputs)
-        x = apply_resnet_block(
-            x,
-            stackwise_num_filters[0],
-            data_format=data_format,
-            dtype=dtype,
-            name="input_block0",
-        )
-        x = VAEAttention(
-            stackwise_num_filters[0],
-            data_format=data_format,
-            dtype=dtype,
-            name="input_attention",
-        )(x)
-        x = apply_resnet_block(
-            x,
-            stackwise_num_filters[0],
-            data_format=data_format,
-            dtype=dtype,
-            name="input_block1",
-        )
-
-        # Stacks.
-        for i, filters in enumerate(stackwise_num_filters):
-            for j in range(stackwise_num_blocks[i]):
-                x = apply_resnet_block(
-                    x,
-                    filters,
-                    data_format=data_format,
-                    dtype=dtype,
-                    name=f"block{i}_{j}",
-                )
-            if i != len(stackwise_num_filters) - 1:
-                # No upsamling in the last blcok.
-                x = layers.UpSampling2D(
-                    2,
-                    data_format=data_format,
-                    dtype=dtype,
-                    name=f"upsample_{i}",
-                )(x)
-                x = layers.Conv2D(
-                    filters,
-                    3,
-                    1,
-                    padding="same",
-                    data_format=data_format,
-                    dtype=dtype,
-                    name=f"upsample_{i}_conv",
-                )(x)
-
-        # Ouput block.
-        x = layers.GroupNormalization(
-            groups=32,
-            axis=gn_axis,
-            epsilon=1e-6,
-            dtype="float32",
-            name="output_norm",
-        )(x)
-        x = layers.Activation("swish", dtype=dtype, name="output_activation")(x)
-        image_outputs = layers.Conv2D(
-            output_channels,
-            3,
-            1,
-            padding="same",
-            data_format=data_format,
-            dtype=dtype,
-            name="output_projection",
-        )(x)
-        super().__init__(inputs=latent_inputs, outputs=image_outputs, **kwargs)
-
-        # === Config ===
-        self.stackwise_num_filters = stackwise_num_filters
-        self.stackwise_num_blocks = stackwise_num_blocks
-        self.output_channels = output_channels
-        self.latent_shape = latent_shape
-
-    @property
-    def scaling_factor(self):
-        """The scaling factor for the latent space.
-
-        This is used to scale the latent space to have unit variance when
-        training the diffusion model.
-        """
-        return 1.5305
-
-    @property
-    def shift_factor(self):
-        """The shift factor for the latent space.
-
-        This is used to shift the latent space to have zero mean when
-        training the diffusion model.
-        """
-        return 0.0609
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "stackwise_num_filters": self.stackwise_num_filters,
-                "stackwise_num_blocks": self.stackwise_num_blocks,
-                "output_channels": self.output_channels,
-                "image_shape": self.latent_shape,
-            }
-        )
-        return config