Secure Text-To-Image Generation in Snowflake with Stable Diffusion and Snowpark Container Services

Published in

Snowflake Builders Blog: Data Engineers, App Developers, AI/ML, & Data Science

14 min readNov 24, 2023

Author: Rachel Blum, Co-Author: Brian Hess

With the recent advances in LLM and Generative AI solutions, many Snowflake customers are asking how they can securely leverage these tools and capabilities to drive revenue or customer satisfaction and loyalty while protecting data and customer privacy. One of the key capabilities that GenAI drives is the ability to scale creative processes that have been historically manual and time consuming (and, hence, quite cost consuming). For example, using Generative AI, organizations will be able to automatically generate marketing and advertising messages, emails or even banner images personalized for their loyalty customers based on behavioral, psychographic and/or demographic data housed in their CDP, and to do so in a privacy compliant manner.

In this article we will be introducing the basic deployment of text-to-image generative AI models in Snowflake using Snowpark Container Services (SPCS). We will be using two pre-trained diffusion models from Hugging Face to securely create images from prompts generated from your Snowflake data all within the secure confines of your Snowflake instance. Specifically, we will be using Stability AI Stability Diffusion Base XL 1.0 (SDBXL1.0) model as well as a smaller model from Segmind Small SD (Small-SD) for testing.

While Large Language Models (LLMs) like LLAMA-2 are primarily trained for text-to-text tasks, Large Multimodal Models (LMMs) involve training generative models to convert between two or more different kinds of data, like text to image, image to text, and image plus text to image, and even text or image to video. While the proprietary multimodal models like OpenAI’s DALLE-2 and Google’s Imagen currently produce the highest quality image results, in order to leverage them for this use case, potentially sensitive or personally identifying data would have to be exposed outside of your data platform to the model APIs. SDBXL1.0 is currently considered the best open source model for LMM processing, but as of today, it still does require a significant amount of skilled prompt engineering or configuration (i.e., using ComfyUI) to get closer to the capabilities of the leading proprietary models and In a future article we will plan to cover deploying ComfyUI as part of this text-to-image processing architecture. It is worth noting that at the rate of acceleration of open source models and the ability to easily swap out models with this architecture, an introduction to deploying diffusion models securely in Snowflake is a great way to get ready for what is coming.

In this article we will:

Create a Snowpark Container Service (currently in preview) to host a Python Flask app and pre-trained diffusion models from Hugging Face for processing text-to-image tasks
Create a service function that calls the Flask app with a prompt generated from customer data
Create a Streamlit in Snowflake UI to call the function and visualize the model’s response as an image.

Pre-reqs for this application include:

A Snowflake account that has been enabled for the Snowpark Container Services preview
ACCOUNTADMIN access in this Snowflake account
A Hugging Face token

Features

Snowflake Snowpark Container Services — Snowpark Container Services is a fully managed solution designed for the seamless deployment, management, and scaling of containerized services, jobs, and functions within the Snowflake ecosystem. By leveraging containerization, Snowpark Container Services frees developers to utilize the bespoke environments of their choosing to support their custom functionality. This offering eliminates the need to transfer data out of Snowflake and comes with built-in security, configuration, and operational best practices as part of its fully managed services. Snowpark Container Services unlocks a range of functionalities, including the ability to run long-lasting services, deploy machine learning models, incorporate custom libraries, use native code, and work with various languages such as C++. This flexibility extends to deploying these functionalities on a diverse set of CPU and GPU compute pools.

Flask — Flask is a lightweight and flexible web framework for Python that is commonly used to build web applications. Flask apps integrate easily when creating and deploying SPCS Service Functions.

Service Function — A service function is a UDF interface used to seamlessly communicate with the containerized Flask application or service from Snowflake SQL.

Streamlit in Snowflake — Streamlit is an open-source Python library that makes it easy to create and share custom web apps for machine learning and data science. By using Streamlit you can quickly build and deploy powerful data applications. For more information about the open-source library, see the Streamlit Library documentation.

Streamlit in Snowflake helps developers securely build, deploy, and share Streamlit apps on Snowflake’s data cloud. Using Streamlit in Snowflake, you can build applications that process and use data in Snowflake without moving data or application code to an external system, and easily share your application with other Snowflake users.

Creating your Text-to-Image Application

To create your application you will need to:

Setup your database
Create, build, and push your container to Snowpark Container Services
Create your container service and service function
Setup your Streamlit in Snowflake application

Step 1: Setup your database

Run this script in your Snowflake Account via your Snowflake Snowsight UI. Make sure you have ACCOUNTADMIN access as it is required to create the initial security integration. The ACCOUNTADMIN will create a role here named SERVICEADMIN that will be used for the rest of the deployment.

The script then does the following:

Creates a database and schema that we will use for the rest of the example
Creates a table of dogs that we will use in this example
Creates two stages, one for the models and one for various specification files we will use
Creates a Snowflake virtual warehouse
Creates a image repository for our Docker images
Creates a compute pool with instances that have GPUs
Creates a helper stored procedure that can be used to write to stage locations

//Database & Role Setup

use role accountadmin;

CREATE SECURITY INTEGRATION IF NOT EXISTS snowservices_ingress_oauth
 TYPE=oauth
 OAUTH_CLIENT=snowservices_ingress
 ENABLED=true;

create or replace role serviceadmin;
grant  role serviceadmin to user {yourusername};
grant role accountadmin to role serviceadmin;


// Setup Sample Data

use role serviceadmin;

CREATE OR REPLACE DATABASE GENAI_SPCS_DB;
CREATE OR REPLACE SCHEMA GENAI_SPCS_SCHEMA;

CREATE OR REPLACE TABLE dogs (name varchar, breed varchar, cabinet varchar, paint varchar);
INSERT INTO dogs (name, breed, cabinet, paint)
VALUES
  ('Maverick', 'Black Labrador', 'Cherry', 'White'),
  ('Max', 'Golden Retriever', 'Oak', 'Green'),
  ('Luna', 'French Bulldog', 'Walnut', 'Cream'),
  ('Charlie', 'Dachshund', 'Espresso', 'Yellow'),
  ('Daisy', 'Shih Tzu', 'Cherry', 'Lavender'),
  ('Rocky', 'Labrador Retriever', 'Oak', 'Aqua'),
  ('Zoe', 'Pomeranian', 'Walnut', 'Gray'),
  ('Milo', 'Beagle', 'Espresso', 'Peach'),
  ('Chloe', 'Siberian Husky', 'White', 'Teal'),
  ('Teddy', 'Poodle', 'Blue', 'Cream');


// Setup Stages

create stage if not exists models encryption = (type = 'SNOWFLAKE_SSE');
create or replace stage  specs encryption = (type = 'SNOWFLAKE_SSE');

CREATE OR REPLACE WAREHOUSE SERVICE_WH
 WAREHOUSE_SIZE = XSMALL
 AUTO_SUSPEND = 120
 AUTO_RESUME = TRUE;

use warehouse service_wh;


//  Setup SPCS

CREATE OR REPLACE IMAGE REPOSITORY images;

// Run this command to find your repository URL and note for later

SHOW IMAGE REPOSITORIES;

CREATE OR REPLACE PROCEDURE PUT_TO_STAGE(STAGE VARCHAR,FILENAME VARCHAR, CONTENT VARCHAR)
RETURNS STRING
LANGUAGE PYTHON
RUNTIME_VERSION=3.8
PACKAGES=('snowflake-snowpark-python')
HANDLER='put_to_stage'
AS $$
import io
import os


def put_to_stage(session, stage, filename, content):
   local_path = '/tmp'
   local_file = os.path.join(local_path, filename)
   f = open(local_file, "w")
   f.write(content)
   f.close()
   session.file.put(local_file, '@'+stage, auto_compress=False, overwrite=True)
   return "saved file "+filename+" in stage "+stage
$$;

Step 2: Build your container

In order to prepare the container, we will need to create the 5 files as shown here:

- /genai_project

- /app

- generate_app.py

- Dockerfile

- download_model.py

- entrypoint.sh

- requirements.txt

Create Flask application

Create a file called generate_app.py with the code shown here and store it in a folder named app in your project folder. This Flask app exposes a single API endpoint at /generate that will receive a prompt and output an image generated by the pre-trained diffusion model (in this case output as a base64 encoded image).

# Save this file as generate_app.py in the app subfolder of your project

from flask import Flask, request, jsonify
from diffusers import DiffusionPipeline
import logging
import base64
import os
import PIL.Image
import torch
import sys
from functools import lru_cache


app = Flask(__name__)
repo = './app/models/'+os.getenv("HF_MODEL")


@lru_cache
def get_pipe():
   pipe = DiffusionPipeline.from_pretrained(repo, local_files_only=True)
   device = "cuda" if torch.cuda.is_available() else "cpu"
   pipe.to(device)
   return pipe

@app.route('/generate', methods=['POST','GET'])
def generate():
   try:
       request_data: dict = request.get_json(force=True)  # type: ignore  
       return_data = []
       for index, prompt in request_data["data"]:
           pipe = get_pipe()
           images = pipe(prompt).images
           image = images[0]
           image.save("output.png")
           with open("output.png", "rb") as f:
               content_bytes = f.read()
           content_b64encoded = base64.b64encode(content_bytes).decode()
           return_data.append([index, content_b64encoded])
       return jsonify({"data": return_data})
   except Exception as e:
       app.logger.exception(e)
       return jsonify(str(e)), 500


if __name__ == '__main__':
   app.run(host='0.0.0.0', port=5000)

This application creates a cacheable function called get_pipe() which initializes the diffusion pipeline for your selected pre-trained model. This will save a bit of time when repeatedly generating the images from function generate(). The get_pipe() function also identifies if GPU processing is available to the service, otherwise it will leverage CPUs. This can be helpful when testing locally on a machine that doesn’t have GPUs. The generate() function uses the pipeline for the selected pre-trained model, passes in a prompt, and outputs a base64 encoded image which can be decoded and displayed or saved to a table in your Snowflake database. It is decorated with Flask decorators to create the /generate API endpoint.

Create shell script

Create and save this script as entrypoint.sh to the main folder of your project. This is the script that will be used as the entrypoint in our Docker image. This script will:

Test to ensure that there is a Hugging Face model set (we will set this in the diffusion.yaml configuration specification).
Check the models stage to see if the model has already been downloaded so that it won’t need to redownload the model from Hugging Face every time you restart the service. If the model is not found it will download the image.
Run the Flask app.

#!/bin/bash


# Check if HUGGING_FACE_MODEL is set
if [ -z "${HUGGING_FACE_MODEL}" ]; then
   echo "Error: HUGGING_FACE_MODEL not set"
   exit 1
fi


# Extract org and repo from HUGGING_FACE_MODEL
export HF_ORG=$(echo $HUGGING_FACE_MODEL | cut -d'/' -f1)
export HF_MODEL=$(echo $HUGGING_FACE_MODEL | cut -d'/' -f2)
export TARGET_DIR="/demo/app/models/$HF_MODEL"
export TARGET_CONFIG="$TARGET_DIR/model_index.json"


if [ -f "$TARGET_CONFIG" ]; then
   echo "Model appears to exist in stage. Skipping download..."
else
   echo ""
   echo ""
   echo "The provided model does not exist in the stage."
   echo "This startup script will download it for you and save to stage. This can take a few minutes."
   echo "This will not need to download on future startups."
   echo ""
   echo ""
   python /download_model.py
fi


python app/generate_app.py

Create app to download model

Create and save this as download_model.py in the main folder of your project. If the model is not found in stage in the entrypoint.sh script (if it hasn’t been downloaded previously), this is the script that will perform the download from Hugging Face. It will require a Hugging Face token which we will provide in the diffusion.yaml configuration specification).

import os
from huggingface_hub import snapshot_download

hf_token = os.getenv("HF_TOKEN")
full_model = os.getenv("HUGGING_FACE_MODEL")
target_dir = os.getenv("TARGET_DIR")

if not os.path.exists(target_dir):
    os.makedirs(target_dir)


snapshot_download(repo_id=full_model, token=hf_token, local_dir_use_symlinks=False,local_dir=target_dir)

Create Dockerfile

Create and save this as Dockerfile in the main folder of your project. To create your container, we recommend a base image of nvcr.io/nvidia/pytorch:23.06-py3 which provides optimized GPU-accelerated versions of popular deep learning frameworks like PyTorch. This Docker file will copy the required files to your image, install required libraries, expose port 5000 for your Flask app, and set the entrypoint of the image to execute entrypoint.sh.

FROM nvcr.io/nvidia/pytorch:23.06-py3

WORKDIR /demo
COPY /app ./app
COPY requirements.txt .
COPY entrypoint.sh /entrypoint.sh
COPY download_model.py /download_model.py
RUN chmod +x /entrypoint.sh

RUN apt update && apt install python3-pip -y
RUN pip3 install accelerate safetensors huggingface_hub fschat diffusers
RUN pip install --no-cache-dir -r requirements.txt

EXPOSE 5000

ENTRYPOINT ["/entrypoint.sh"]

Create requirements file

Save this file as requirements.txt in your main project directory. These are the Python dependencies needed for the Python scripts.

transformers==4.33.3
accelerate
safetensors
huggingface_hub
fschat
pyarrow==10.0.1
diffusers
flask
python-dotenv
pybase64
Pillow

Build and push your docker image to SPCS

In a terminal, use the following commands from your main project folder to build and push your image to Snowflake. We explicitly set the platform when we build the Docker image so that we can build from any machine, i.e., from a Mac with an ARM chip. This may not be required for your environment, but it doesn’t hurt to include it. Make sure to substitute your organization, account name and user name where applicable.

You can find your image repository URL in Step 1 where you ran SHOW IMAGE REPOSITORIES. Replace that where you see {image_repository_url} and replace your {your_username} with your Snowflake username. When you run docker login you will be prompted to enter your Snowflake password.

docker build --platform=linux/amd64 -t {your_username}/diffusion .   


docker login  {image_repository_url} -u {your_username}


docker tag  {yourusername}/diffusion:latest  {image_repository_url}/diffusion:latest


docker push {image_repository_url}/diffusion:latest

Once this has been completed you can check that your image has been uploaded to the repository by issuing this SQL (from Snowsight or SnowSQL or any SQL client). If you created the repository in a different database (genai_spcs_db) and/or schema (genai_spcs_schema), replace those values with your own.

select system$registry_list_images('/genai_spcs_db/genai_spcs_schema/images');

Step 3: Create your service and function in Snowflake

For this step you will need your free Hugging Face token and the name of your Hugging Face text-to-image diffusion model. We used Segmind’s Small-SD model for testing, but for better image generation quality the current gold standard of open-source diffusion models is Stability AI’s Stable Diffusion XL Base 1.0. Small-SD will run on GPU3 while our tests of Stable XL required GPU7 for the RAM. Fortunately, with this application you can easily test both or any models without having to rebuild or deploy your image to SPCS.

The service is defined by a YAML file, in our case the diffusion.yaml file, that defines:

The containers included. In our case, we just need one container, diffusion. For each container, we specify:
The image from the image repository
Any environment variables for the container. In our case we specify the Hugging Face token here, as well as the Hugging Face model to use. This allows us to change the token or model without having to rebuild the image.
Any stage-based volumes that should be mounted inside the container, and their mount location.
The endpoints that the service should expose from the containers. These are ports that have been exposed by the Docker image, and the exposed port number is the same for the service as is for the Docker image. We can also specify if the exposed endpoint should be publicly available or only internally available.
The stages that we would like to mount into the containers.
The network policy for ingress or egress to/from the containers.

In your Snowflake Snowsight UI run the following scripts:

Setup Segmind’s Small_SD Model for Testing

Note that when you test the service what you should get back is a very long string — that should be the base64 encoded image you generated with your diffusion model.

/// SETUP FOR SMALL MODEL

use role serviceadmin;

//Setup compute pool

CREATE COMPUTE POOL IF NOT EXISTS gpu3_pool
MIN_NODES = 1
MAX_NODES = 1
INSTANCE_FAMILY = gpu_3;

show compute pools;


/// Create the Specification YAML
call put_to_stage('specs','diffusion.yaml',$$
spec:
 containers:
   - name: diffusion
     image: sfsenorthamerica-rbdemo02.registry.snowflakecomputing.com/genai_spcs_db/genai_spcs_schema/images/diffusion
     env:
         HF_TOKEN: '{your_huggingface_token}'
         HUGGING_FACE_MODEL: 'segmind/small-sd'
     volumeMounts:
       - name: models
         mountPath: demo/app/models
 endpoints:
   - name: pipe
     port: 5000
     public: true
 volumes:
   - name: models
     source: "@genai_spcs_db.genai_spcs_schema.models"
 networkPolicyConfig:
   allowInternetEgress: true
$$);

/// Create the Service
create service diffusion
in compute pool gpu3_pool
from @specs
spec='diffusion.yaml';

/// Check the status of the service
CALL SYSTEM$GET_SERVICE_STATUS('GENAI_SPCS_DB.GENAI_SPCS_SCHEMA.diffusion');

/// Retrieve the logs from the diffusion container in the service
call system$get_service_logs('diffusion', '0', 'diffusion', '100');

/// Create the service function
create or replace function diffusion(message varchar)
returns varchar
service=diffusion
endpoint=pipe
as '/generate';


// Test function

select diffusion('bear in a tree');

Setup for Stable Diffusion XL

To use the same image for Stability AI’s XL model, I simply switch out the HUGGING_FACE_MODEL environment variable in a new yaml, create a gpu7 compute pool, create a new service function called diffusion_xl using that new compute pool, and create a new function named diffusion_xl based on that service.

/// SETUP FOR XL Model

use role serviceadmin;


/// Create a GPU7 compute pool
CREATE COMPUTE POOL IF NOT EXISTS gpu7_pool
MIN_NODES = 1
MAX_NODES = 1
INSTANCE_FAMILY = gpu_7;


/// Create the diffusion_xl.yaml specification file
call put_to_stage('specs','diffusion_xl.yaml',$$
spec:
 containers:
   - name: diffusion
     image: sfsenorthamerica-rbdemo02.registry.snowflakecomputing.com/genai_spcs_db/genai_spcs_schema/images/diffusion
     env:
         HF_TOKEN: '{your_huggingface_token}'
         HUGGING_FACE_MODEL: 'stabilityai/stable-diffusion-xl-base-1.0'
     volumeMounts:
       - name: models
         mountPath: demo/app/models
 endpoints:
   - name: pipe
     port: 5000
     public: true
 volumes:
   - name: models
     source: "@genai_spcs_db.genai_spcs_schema.models"
 networkPolicyConfig:
   allowInternetEgress: true
$$);


/// Create the Service
create service diffusion_xl
in compute pool gpu7_pool
from @specs
spec='diffusion_xl.yaml';


/// Check the status of the service
CALL SYSTEM$GET_SERVICE_STATUS('GENAI_SPCS_DB.GENAI_SPCS_SCHEMA.diffusion_xl');

/// Retrieve the logs from the diffusion container in the service
call system$get_service_logs('diffusion_xl', '0', 'diffusion', '100');


/// Create the service function
create or replace function diffusion_xl(message varchar)
returns varchar
service=diffusion_xl
endpoint=pipe
as '/generate'
;


// Test function

select diffusion_xl('bear in a tree');

Step 4: Create Streamlit in Snowflake UI

Once you have tested that your function works (i.e., it returns a base64 encoded image string), you can set up this Streamlit in Snowflake app to enter manual prompts or generate prompts from your sample data to reveal your diffusion model generated image.

The simplest way to create this Streamlit in Snowflake is to open Snowsight and navigate to the Streamlit tab. From there, create a new Streamlit, specifying the warehouse we created in Step 1 and genai_spcs_db and genai_spcs_schema as the database and schema, respectively. Then, replace the boilerplate code with the following code:

import streamlit as st
from snowflake.snowpark.context import get_active_session
from PIL import Image
import base64
from io import BytesIO


function = 'diffusion_xl'
# diffusion or diffusion_xl


# Title
st.title("Snowflake Secure Text-to-Image Generation")
st.write(
    "")

# Get the current credentials
session = get_active_session()

# Function to decode base64 to image
def decode_base64_image(base64_string):
    try:
        # Decode the base64 string to binary data
        image_data = base64.b64decode(base64_string)
        # Create a BytesIO buffer and open the image with Pillow
        image = Image.open(BytesIO(image_data))
        return image
    except Exception as e:
        st.error(f"Error decoding base64 image: {e}")
        return None

option = st.radio("Enter your own promps or generate prompts from data",('Generate a sample prompt from data','Enter your own prompts') )


if option=='Generate a sample prompt from data':
    dog_df = session.sql("select name from GENAI_SPCS_DB.GENAI_SPCS_SCHEMA.dogs").to_pandas()
    dog_pick = st.selectbox('Pick a dog:',dog_df)  
    attributes = session.sql(f"select breed, cabinet, paint from GENAI_SPCS_DB.GENAI_SPCS_SCHEMA.dogs where name = '{dog_pick}' limit 1").to_pandas()
    gen_prompt = f"PROMPT: Realistic photograph of a happy {attributes['BREED'].iloc[0]} in a kitchen with {attributes['CABINET'].iloc[0]} wood cabinets and {attributes['PAINT'].iloc[0]} painted walls" 
    st.write(gen_prompt)
    run = st.button('Generate Image',key='sample')
    if run:
        prompt = '{ "steps": 40,"width": 1024,"height": 1024,"seed": 0,"cfg_scale": 5,"samples": 1,"style_preset": "photographic","text_prompts": [{"text": "'+ gen_prompt+ '","weight": 1},{"text": "blurry","weight": -1}],}'
        image_query =session.sql(f"select GENAI_SPCS_DB.GENAI_SPCS_SCHEMA.{function}('{gen_prompt}') as IMG_BASE64").to_pandas()
        base64_string = (image_query["IMG_BASE64"].head(1).item())
        image = decode_base64_image(base64_string)
        st.image(image)

if option=='Enter your own prompts':
    # Enter your own prompt
    prompt_enter = st.text_input('Please enter your text prompt here:')
    negprompt_enter = st.text_input ('Please enter your negative prompt here:', value='blurry')
    submit = st.button('Generate Image',key='manual')
    if submit:
        prompt = '{ "steps": 40,"width": 1024,"height": 1024,"seed": 0,"cfg_scale": 5,"samples": 1,"style_preset": "photographic","text_prompts": [{"text": "'+ prompt_enter+ '","weight": 1},{"text": "'+ negprompt_enter+ '","weight": -1}],}'
        image_query =session.sql(f"select GENAI_SPCS_DB.GENAI_SPCS_SCHEMA.{function}('{prompt}') as IMG_BASE64").to_pandas()
        base64_string = (image_query["IMG_BASE64"].head(1).item())
        image = decode_base64_image(base64_string)
        st.image(image)

To generate your image, select a dog’s name from the drop down list and click the Generate Image button.

You can also select the radio button for ‘Enter your own prompts’ to test your own prompt engineering skills.

Summary

Deploying open-source diffusion models like Stability AI’s Stable Diffusion Base XL 1.0 in Snowpark Container Services will allow companies to significantly accelerate their targeted marketing and advertising processes by securely and automatically analyzing customer data and generating personalized creative with near or close to near real time capabilities. It is likely that there will soon be significant advances including improved text-to-image, image-to-image and text- and image-to-video capabilities for these open-source models. For now, one of the best ways to better engineer the output of SDXLB1.0 is via ComfyUI as we will demonstrate in a future article. We will also be documenting how publishers can deploy this text-to-image processing architecture as part of their Snowflake Data Clean Room Native Application, joining data securely across publisher and advertiser accounts. Getting started with this basic architecture for deploying diffusion models in Snowflake is a great first step in leveraging Generative AI to drive greater business value while preserving data privacy.