Guess who or what app using Hugging Face Deep Learning container model on Vertex AI

Source notebook

Repo path: open-models/use-cases/guess_app.ipynb · Open on GitHub · intermediate

Builds a Gradio riddle game using Gemini and a Hugging Face FLUX model deployed on Vertex AI.

Summary

This notebook shows how to register and deploy the Hugging Face FLUX.1-dev text-to-image model on Vertex AI using a Deep Learning container. It then uses Gemini to solve a riddle, generate an image prompt for the answer, call the deployed FLUX endpoint, and display the result in a Gradio app. The workflow also covers authentication, IAM/API requirements, endpoint cleanup, and model cleanup.

Key code patterns

Initialize Vertex AI

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
aiplatform.init(project=PROJECT_ID, location=LOCATION)
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the project and region used by both Vertex AI model deployment and Gemini calls.

Register Hugging Face model

flux_model = aiplatform.Model.upload(
    display_name="flux--generate",
    serving_container_image_uri="us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-2.ubuntu2204.py310",
    serving_container_environment_variables={
        "HF_MODEL_ID": "black-forest-labs/FLUX.1-dev",
        "HF_TASK": "text-to-image",
        "HF_TOKEN": get_token(),
    },
)
flux_model.wait()

Uses a Hugging Face Deep Learning container and environment variables to register FLUX in Vertex AI Model Registry.

Deploy to endpoint

endpoint = aiplatform.Endpoint.create(display_name="flux--generate-endpoint")
deployed_flux_model = flux_model.deploy(
    endpoint=endpoint,
    machine_type="g2-standard-48",
    accelerator_type="NVIDIA_L4",
    accelerator_count=4,
    sync=False,
)

Creates a Vertex AI endpoint backed by GPU resources for managed FLUX inference.

Predict and decode image

response = ENDPOINT.predict(
    instances=[image_gen_prompt],
    parameters={"width": 512, "height": 512, "num_inference_steps": 8, "guidance_scale": 3.5},
)
image = Image.open(io.BytesIO(base64.b64decode(response.predictions[0])))

Calls the deployed endpoint and converts the base64 prediction payload into a PIL image.

Gemini riddle flow

def guess_game(riddle: str) -> tuple[Image.Image, str, str]:
    answer = generate_subject(riddle)
    prompt = generate_prompt(answer)
    image = generate_image(prompt)
    return image, answer, prompt

Chains Gemini subject extraction, Gemini prompt generation, and FLUX image generation into one app workflow.

Gradio event wiring

submit_btn.click(
    guess_game,
    inputs=[prompt_input],
    outputs=[image_output, answer_output, image_prompt],
)
app.launch()

Exposes the riddle-to-image pipeline through an interactive Gradio interface.

Models & APIs used

  • Models: black-forest-labs/FLUX.1-dev, gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI Model Registry, Vertex AI Prediction, Artifact Registry
  • SDKs / libraries: google-cloud-aiplatform, huggingface_hub, gradio, vertexai, Pillow

When to use this

Use this pattern when building a small app that combines Gemini text reasoning with an open text-to-image model deployed on Vertex AI.

Gotchas & caveats

  • Colab may require a runtime restart after installing google-cloud-aiplatform, huggingface_hub, and gradio.
  • Local JupyterLab may require gcloud auth login, while Colab uses google.colab.auth.authenticate_user.
  • Hugging Face authentication is required with interpreter_login so get_token can supply HF_TOKEN.
  • The Google Cloud project must have Vertex AI API and Artifact Registry API enabled.
  • The runtime identity needs Artifact Registry Reader and Vertex AI User roles.
  • The deployment requests g2-standard-48 with four NVIDIA_L4 accelerators, which may require regional quota and availability.
  • Gemini safety thresholds are explicitly set to OFF in the app code.
  • The notebook includes cleanup for the Gradio app, deployed model, and uploaded model to avoid leaving resources running.

Best practices

  • Initialize aiplatform and vertexai with explicit project and location values.
  • Retrieve the Hugging Face token with get_token instead of hardcoding it in the notebook.
  • Register the Hugging Face model in Vertex AI Model Registry before deploying it to an endpoint.
  • Use separate helper functions for Gemini content generation, subject extraction, prompt generation, and image generation.
  • Use temperature 0 and candidate_count 1 for Gemini riddle solving and prompt generation.
  • Decode the endpoint response from base64 before displaying it as an image.
  • Close the Gradio app and delete the deployed model and uploaded model during cleanup.