@@ -14,12 +14,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+from google.cloud.aiplatform.helpers import _decorators
 from google.cloud.aiplatform.v1beta1.schema import predict
 from google.cloud.aiplatform.v1beta1.schema import trainingjob
-
+from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1 import (
+types as instance,
+)
+from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1 import (
+types as params,
+)
+from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1 import (
+types as prediction,
+)
+from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1 import (
+types as definition,
+)
 
 __all__ = (
 "predict",
 "trainingjob",
 )
+
+enhanced_types_packages = [
+instance,
+params,
+prediction,
+definition,
+]
+
+for pkg in enhanced_types_packages:
+_decorators._add_methods_to_classes_in_package(pkg)

@@ -14,9 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1 import types as pkg
-
 from google.cloud.aiplatform.v1beta1.schema.predict.instance_v1beta1.types.image_classification import (
 ImageClassificationPredictionInstance,
 )
@@ -56,4 +53,3 @@
 "VideoClassificationPredictionInstance",
 "VideoObjectTrackingPredictionInstance",
 )
-_decorators._add_methods_to_classes_in_package(pkg)

@@ -14,9 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1 import types as pkg
-
 from google.cloud.aiplatform.v1beta1.schema.predict.params_v1beta1.types.image_classification import (
 ImageClassificationPredictionParams,
 )
@@ -44,4 +41,3 @@
 "VideoClassificationPredictionParams",
 "VideoObjectTrackingPredictionParams",
 )
-_decorators._add_methods_to_classes_in_package(pkg)

@@ -14,11 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1 import (
-types as pkg,
-)
-
 from google.cloud.aiplatform.v1beta1.schema.predict.prediction_v1beta1.types.classification import (
 ClassificationPredictionResult,
 )
@@ -66,4 +61,3 @@
 "VideoClassificationPredictionResult",
 "VideoObjectTrackingPredictionResult",
 )
-_decorators._add_methods_to_classes_in_package(pkg)

@@ -14,11 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from google.cloud.aiplatform.helpers import _decorators
-from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1 import (
-types as pkg,
-)
-
 from google.cloud.aiplatform.v1beta1.schema.trainingjob.definition_v1beta1.types.automl_forecasting import (
 AutoMlForecasting,
 )
@@ -134,4 +129,3 @@
 "AutoMlVideoObjectTrackingInputs",
 "ExportEvaluatedDataItemsConfig",
 )
-_decorators._add_methods_to_classes_in_package(pkg)