_LAUNCH_STAGE check in _model_garden breaks code

[khaerensml6]

Environment details

• OS type and version: ubuntu
• Python version: 3.10
• pip version: 22.0.2
• google-cloud-aiplatform version: 1.34.0

Steps to reproduce

1. Try to use "generate_images" to generate images from a pretrained model:
2. read error

Code example

 model = ImageGenerationModel.from_pretrained("imagegeneration@002")

    images = model.generate_images(
        prompt=prompt,
        number_of_images=4,
    )

Stack trace

Exception has occurred: ValueError
The model you are trying to instantiate does not support the launch stage: PRIVATE_PREVIEW
  File "/home/karel/Desktop/work/store-n8/recipe_processing/recipe_enrichment_mlpipelines/mlpipelines/components/image_generation/src/gcp_helpers.py", line 74, in generate_images
    model = ImageGenerationModel.from_pretrained("imagegeneration@002")
  File "/home/karel/Desktop/work/store-n8/recipe_processing/recipe_enrichment_mlpipelines/mlpipelines/components/image_generation/src/gcp_helpers.py", line 92, in <module>
    images = generate_images(prompt)
ValueError: The model you are trying to instantiate does not support the launch stage: PRIVATE_PREVIEW

After some debugging, I found the error is thrown because of a check in _model_garden : here

If I change the _LAUNCH_STAGE in _model_garden to _SDK_PRIVATE_PREVIEW_LAUNCH_STAGE, everything works again.

Don't know what the point of this check is, it's fairly easy to circumvent, but it's a bit annoying.

Hacky Solution

import vertexai._model_garden._model_garden_models as _model_garden_models
_model_garden_models._ModelGardenModel._LAUNCH_STAGE = _model_garden_models._SDK_PRIVATE_PREVIEW_LAUNCH_STAGE

[sasha-gitg]

This should be fixed. Please reopen if needed.