Image Generation with Streamlit and Google Vertex AI: A Guide for Developers
Introduction
Image generation is an emerging field with numerous applications, from marketing to e-commerce. This blog post focuses on how developers can quickly get started with this technology using Streamlit for the frontend and Google Vertex AI for the backend.
The Importance of Image Generation
Before diving into the technical aspects, let’s consider why image generation is a compelling tool:
- Customization: Generate images that cater to specific user needs or preferences.
- Speed: Create images programmatically without the need for manual design efforts.
- Versatility: Applicable to various sectors like advertising, content creation, and more.
Getting Started with Streamlit
Streamlit provides an easy-to-use framework for developing web applications. For this example, we’ll use it to create a simple user interface.
# Streamlit code to upload a JSON file
st.chat_message("user").write("Please upload a file")
menufile = st.file_uploader("File", type=["json"])Here, Streamlit sets up a UI prompt and a file uploader for JSON files, making it accessible even for developers new to web development.
Google Vertex AI for Backend
Google Vertex AI offers pre-trained machine learning models, one of which is for image generation. The API is straightforward, making it convenient for developers to integrate into their applications.
# Initialize and use Vertex AI for image generation
vertexai.init(project=PROJECT_ID, location=LOCATION, credentials=credentials)
model = ImageGenerationModel.from_pretrained("imagegeneration@002")
images1 = model.generate_images(prompt=text, number_of_images=4, seed=1)In these lines of code, Vertex AI is initialized, a pre-trained model for image generation is loaded, and then images are generated based on a given text prompt.
Putting It All Together
Code Flow
- Initialize Google Cloud Credentials: To interact with Google Vertex AI.
credentials = service_account.Credentials.from_service_account_info(st.secrets["gcs_connections"])- Credentials will need to be stored in .streamlit/secrets.toml (or in Streamlit Cloud Secrets if you’re running hosted). You will need to create a Google Cloud service account, assign it the vertex ai service agent role, and then generate a Service Account Key in JSON format. Take the fields from the JSON and put them into the TOML file. It should look something like this:
[gcs_connections]
type = "service_account"
project_id = "skeenan"
private_key_id = "some_private_key"
private_key = "-----BEGIN PRIVATE KEY-----\nYour Key Here\n-----END PRIVATE KEY-----\n"
client_email = "my_app@skeenan.iam.gserviceaccount.com"
client_id = "1234567890"
auth_uri = "https://accounts.google.com/o/oauth2/auth"
token_uri = "https://oauth2.googleapis.com/token"
auth_provider_x509_cert_url = "https://www.googleapis.com/oauth2/v1/certs"
client_x509_cert_url = "https://www.googleapis.com/robot/v1/metadata/x509/my_app%40skeenan.iam.gserviceaccount.com"- Streamlit UI for JSON Upload: Allow the user to upload a JSON file. We use this file to determine what images to render for this use case
menufile = st.file_uploader("File", type=["json"])- Text Extraction and Prompt Formation: Formulate the text prompt for image generation based on the uploaded JSON data.
data = json.loads(bytes_data) text = "Generate an image based on: {content}".format(content=data['content'][0])- Image Generation with Vertex AI: Generate images.
images = model.generate_images(prompt=text, number_of_images=4, seed=1)- Display Images using Streamlit: Show the generated images on the UI.
st.image("./{}1.png".format(fileid))Conclusion
For developers interested in exploring image generation, Streamlit and Google Vertex AI offer a simple yet powerful way to start. This combination enables rapid prototyping, allowing you to focus on the application logic rather than getting bogged down by the intricacies of machine learning or web development. With such tools at your disposal, diving into the world of custom image generation has never been easier.
