Running a Stable Diffusion Cluster on GCP with tensorflow-serving (Part 2) | by Thushan Ganegedara | Mar, 2023

Creating the artifacts and deploying the mannequin on the cluster

In part 1, we realized the best way to use terraform to arrange and handle our infrastructure conveniently. On this half, we’ll proceed on our journey to deploy a operating Steady Diffusion mannequin on the provisioned cluster.

Be aware: You’ll be able to observe this tutorial end-to-end even for those who’re a free person (so long as you have got a few of free tier credit left).

All pictures, until in any other case famous, are by the creator

Github: https://github.com/thushv89/tf-serving-gke

Let’s check out what the ultimate outcome can be.

What’s Steady Diffusion anyway?

There are 5 most important elements that builds up the Stable Diffusion model:

• Tokenizer — Tokenizes a given string to an inventory of tokens (numerical IDs)
• Textual content encoder — Takes within the tokenized textual content and produces a textual content embedding
• Diffusion mannequin — Takes within the textual content embedding and a latent picture (initially noise) as an enter and incrementally refine the latent picture to encoder increasingly helpful info (visually pleasing)
• Decoder — Takes within the last latent picture and produces an precise picture
• Picture encoder (used for the in-painting function — we’ll be ignoring this for this train)

The principal ground-shattering thought behind secure diffusion (diffusion fashions) is,

For those who add a little bit of noise to a picture steadily for a lot of steps, you’ll find yourself with a picture containing noise. By reversing the order of the method, you possibly can have an enter (noise) and a goal (authentic picture). Then a mannequin is skilled to foretell the unique picture from noise.

All the above elements work cohesively to realize this concept.

Storing the Steady Diffusion mannequin

Code: https://github.com/thushv89/tf-serving-gke/blob/master/notebooks/savedmodel_stable_diffusion.ipynb

To be able to assemble a Stable Diffusion model, we’ll be utilizing the keras_cv library, which features a assortment of fashionable deep studying imaginative and prescient fashions for picture classification, segmentation, generative AI, and so on. You will discover a tutorial here, which explains the best way to use the StableDiffusion in keras_cv. You’ll be able to open up a pocket book and play with the mannequin to familiarize your self.

Our aim right here is to save lots of the StableDiffusionmannequin within the SavedModel format; the go-to commonplace for serializing TensorFlow fashions. One essential requirement to do that is ensuring all operations used are TensorFlow graph appropriate. Sadly, this isn’t the case.

• The present model of the mannequin makes use of a TensorFlow graph incompatible tokenizer, so it must be introduced out of the packaged mannequin and utilized in a separate step.
• The present model makes use of predict_on_batch so as to generate a picture, which isn’t supported by TensorFlow graph constructing.

Fixing the mannequin

To be able to patch up the keen mode StableDiffusion mannequin, we’ll create a brand new mannequin referred to as StableDiffusionNoTokenizer. By means of out this new mannequin, we’ll substitute all predict_on_batch() calls with graph appropriate __call__() calls. We’ll even be decoupling the tokenization course of from the mannequin because the identify suggests. Moreover, within the generate_image() perform, we’ll be changing,

timesteps = tf.vary(1, 1000, 1000 // num_steps)
alphas, alphas_prev = self._get_initial_alphas(timesteps)
progbar = keras.utils.Progbar(len(timesteps))
iteration = 0
for index, timestep in checklist(enumerate(timesteps))[::-1]:
latent_prev = latent  # Put aside the earlier latent vector
t_emb = self._get_timestep_embedding(timestep, batch_size)
unconditional_latent = self.diffusion_model.predict_on_batch(
[latent, t_emb, unconditional_context]
)
latent = self.diffusion_model.predict_on_batch(
[latent, t_emb, context]
)
latent = unconditional_latent + unconditional_guidance_scale * (
latent - unconditional_latent
)
a_t, a_prev = alphas[index], alphas_prev[index]
pred_x0 = (latent_prev - math.sqrt(1 - a_t) * latent) / math.sqrt(
a_t
)
latent = (
latent * math.sqrt(1.0 - a_prev) + math.sqrt(a_prev) * pred_x0
)
iteration += 1
progbar.replace(iteration)

with,

latent = self.diffusion_reverse_loop(
latent,
context=context, 
unconditional_context=unconditional_context, 
batch_size=batch_size, 
unconditional_guidance_scale=unconditional_guidance_scale,
num_steps=num_steps,
)

the place,

@tf.perform
def diffusion_reverse_loop(self, latent, context, unconditional_context, batch_size, unconditional_guidance_scale, num_steps):
index = num_steps -1
cond = tf.math.higher(index, -1)
timesteps = tf.vary(1, 1000, 1000 // num_steps)
alphas, alphas_prev = self._get_initial_alphas(timesteps)
iter_partial_fn = functools.partial(
self._diffusion_reverse_iter, 
timesteps=timesteps, 
alphas=alphas, 
alphas_prev=alphas_prev, 
context=context, 
unconditional_context=unconditional_context, 
batch_size=batch_size, 
unconditional_guidance_scale=unconditional_guidance_scale, 
num_steps=num_steps
)
latent, index = tf.while_loop(cond=lambda _, i: tf.math.higher(i, -1), physique=iter_partial_fn, loop_vars=[latent, index])
return latent 
@tf.perform
def _diffusion_reverse_iter(self, latent_prev, index, timesteps,  alphas, alphas_prev, context, unconditional_context, batch_size, unconditional_guidance_scale, num_steps):
t_emb = self._get_timestep_embedding(timesteps[index], batch_size)
combined_latent = self.diffusion_model(
[
tf.concat([latent_prev, latent_prev],axis=0), 
tf.concat([t_emb, t_emb], axis=0), 
tf.concat([context, unconditional_context], axis=0)
], coaching=False
)
latent, unconditional_latent = tf.break up(combined_latent, 2, axis=0)
latent = unconditional_latent + unconditional_guidance_scale * (
latent - unconditional_latent
)
a_t, a_prev = alphas[index], alphas_prev[index]
pred_x0 = (latent_prev - tf.math.sqrt(1 - a_t) * latent) / tf.math.sqrt(a_t)
latent = latent * tf.math.sqrt(1.0 - a_prev) + tf.math.sqrt(a_prev) * pred_x0
index -= 1
return latent, index

Two most important adjustments I’ve performed are:

• As a substitute of a Python for loop, I’ve used the tf.while_loop which is extra performant in TensorFlow.
• Mixed the 2 separate calls to the diffusion_model to a single name and later break up the outputs.

There are different adjustments equivalent to changing varied operations with TensorFlow equal (e.g. np.clip() -> tf.clip_by_value()), you possibly can evaluate and distinction the original model, with this version to check and distinction.

When working with TensorFlow within the graph execution mode, you need to use tf.print() statements so as to make sure the validity of the code throughout execution. Please consult with the Appendix for extra details about tf.print().

As soon as the underlying mannequin is mounted, we are able to create the next mannequin, which could be executed in graph mode with no hiccup.

class StableDiffusionTFModel(tf.keras.fashions.Mannequin):
def __init__(self):
tremendous().__init__()
self.image_width = self.image_height = 384
self.mannequin = StableDiffusionNoTokenizer(img_width=self.image_width, img_height=self.image_height, encoded_text_length=None, jit_compile=True)
# This forces the mannequin obtain its elements
# self.image_encoder is just required for in-painting - we'll ignore this performance on this excercise
self.text_encoder = self.mannequin.text_encoder
self.diffusion_model = self.mannequin.diffusion_model
self.decoder = self.mannequin.decoder
self.default_num_steps = tf.fixed(40) 
self.default_batch_size = tf.fixed(2)
# These damaging immediate tokens are borrowed from the unique secure diffusion mannequin
self.default_negative_prompt_tokens = tf.fixed(
[
49406, 8159, 267, 83, 3299, 267, 21101, 8893, 3500, 267, 21101, 
8893, 4804, 267, 21101, 8893, 1710, 267, 620, 539, 6481, 267, 
38626, 267, 12598, 943, 267, 4231, 34886, 267, 4231, 7072, 267, 
4231, 5706, 267, 1518, 15630, 267, 561, 6528, 267, 3417, 268, 
3272, 267, 1774, 620, 539, 6481, 267, 21977, 267, 2103, 794, 
267, 2103, 15376, 267, 38013, 267, 4160, 267, 2505, 2110, 267, 
782, 23257, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407
], dtype=tf.int32
)
def name(self, inputs):
encoded_text = self.text_encoder([inputs["tokens"], self.mannequin._get_pos_ids()], coaching=False)
pictures = self.mannequin.generate_image(
encoded_text, 
negative_prompt_tokens=inputs.get("negative_prompt_tokens", self.default_negative_prompt_tokens),
num_steps=inputs.get("num_steps", self.default_num_steps), 
batch_size=inputs.get("batch_size", self.default_batch_size)
)
return pictures
mannequin = StableDiffusionTFModel()

This mannequin takes within the following inputs:

• input_tokens: Tokenized illustration of the enter string
• negative_prompt_tokens: Tokenized illustration of the damaging immediate (extra about damaging prompting: right here)
• num_steps: Variety of steps to run the diffusion course of for
• batch_size: Variety of pictures to generate per picture

Right here’s an instance utilization of this mannequin:

# Tokenizing the prompts
tokenizer = SimpleTokenizer()
def generate_tokens(tokenizer, immediate, MAX_PROMPT_LENGTH):
inputs = tokenizer.encode(immediate)
if len(inputs) > MAX_PROMPT_LENGTH:
increase ValueError(
f"Immediate is simply too lengthy (ought to be <= {MAX_PROMPT_LENGTH} tokens)"
)
phrase = tf.concat([inputs, ([49407] * (MAX_PROMPT_LENGTH - len(inputs)))], axis=0)
return phrase
tokens = generate_tokens(tokenizer, "a ferrari automobile with wings", MAX_PROMPT_LENGTH)
# Invoking the mannequin
all_images = []
num_steps = 30
tokens = generate_tokens(tokenizer, "a fortress in Norway overlooking a glacier, panorama, surrounded by fairies combating trolls, sundown, prime quality", MAX_PROMPT_LENGTH)
neg_tokens = generate_tokens(tokenizer, "ugly, tiling, poorly drawn fingers, poorly drawn toes, poorly drawn face, out of body, mutation, mutated, additional limbs, additional legs, additional arms, disfigured, deformed, cross-eye, physique out of body, blurry, dangerous artwork, dangerous anatomy, blurred, textual content, watermark, grainy", MAX_PROMPT_LENGTH)
pictures = mannequin({
"tokens": tokens, 
"negative_prompt_tokens": neg_tokens,
"num_steps": tf.fixed(num_steps), 
"batch_size": tf.fixed(1)
})

Keep in mind that, I’m (i.e. Free person tier) closely restricted by quotas on this challenge.

• No GPU quota in any respect
• Max 8 N2 CPUs (For those who select N1 CPUs, you possibly can go as much as 12)

Due to this fact, I can not use any GPU cases or greater than 2 n2-standard-4cases. Steady Diffusion fashions are fairly sluggish subsequently we’ll be challenged by latency utilizing CPU cases.

Listed below are some particulars about how lengthy it takes underneath completely different parameters. The assessments had been don on a n2-standard-8 machine on Vertex AI workbench.

• Picture dimension (num_steps = 40)
  — 512×512 picture: 474s
  — 384×384 picture: 233s
• batch_size and num_steps
  — batch dimension = 1: 21.6s (num_steps=5), 67.7s (num_steps=20) and 99.5s (num_steps=30)
  — batch dimension = 2, 55.6s (num_steps=5), 121.1s (num_steps=20) and 180.2s (num_steps=30)
  — batch dimension=4, 21.6s (num_steps=5), 67.7s (num_steps=20) and 99.5s (num_steps=30)

As you possibly can see, growing the image_size , batch_size , num_steps result in elevated time consumption. Due to this fact balancing the computational price with the picture high quality, we selected the next parameters for our deployed mannequin.

• image_size: 384x384
• num_steps: 30
• batch_size: 1

As soon as the mannequin is created, add the mannequin to the created GCS bucket.

!gsutil -m cp -r  ./stable_diffusion_model gs://<challenge>-bucket/

This would be the knowledge supply we’ll be utilizing so as to deploy our mannequin as a prediction service.

Let’s once more take a second to understand some pictures generated by the mannequin, earlier than persevering with on to the following part.

Code: https://github.com/thushv89/tf-serving-gke/tree/master/infrastrcture

To deploy our mannequin and setup a prediction service, we’d like 3 configurations:

• configmap.yaml — This defines varied variables which might be required in the course of the deployment. For instance this is able to embody the placement of the SavedModel on GCS (i.e. accessible by the atmosphere variable MODEL_PATH)
• deployment.yaml — Deployment defines the pod specs (e.g. CPU) and containers it ought to be operating. On this case, we’ll be operating a single container operating tensorflow-serving serving the mannequin situated at MODEL_PATH .
• service.yaml — Service is the mechanism with which we expose our tensorflow-serving app operating in our pods. For instance we are able to inform it to reveal our pod(s) by a load balancer.

The deployment

Let’s first take a look at the spec of the deployment :

spec:
replicas: 1
selector:
matchLabels:
app: stable-diffusion
template:
metadata:
labels:
app: stable-diffusion
spec:
containers:
- identify: tf-serving
picture: "tensorflow/serving:2.11.0"
args:
- "--model_name=$(MODEL_NAME)"
- "--model_base_path=$(MODEL_PATH)"
- "--rest_api_timeout_in_ms=720000"
envFrom:
- configMapRef:
identify: tfserving-configs
imagePullPolicy: IfNotPresent
readinessProbe:
httpGet:
path: "/v1/fashions/stable-diffusion"
port: 8501
scheme: HTTP
initialDelaySeconds: 30
periodSeconds: 15
failureThreshold: 10
ports:
- identify: http
containerPort: 8501
protocol: TCP
- identify: grpc
containerPort: 8500
protocol: TCP
sources:
requests:
cpu: "3"
reminiscence: "12Gi"

We are able to make few attention-grabbing observations:

• We’re solely declaring a single reproduction within the script, scaling will likely be setup individually in and will likely be managed by an autoscaling coverage
• We offer a selector, which the service will search for in a deployment to make sure it’s serving on the proper deployment
• We expose two ports; 8501 (HTTP visitors) and 8500 (GRPC visitors)
• We’ll be requesting 3 “CPU time” and 12Gi per container.

Be aware 1:A node will likely be sometimes operating different pods necessitated by Kubernetes (i.e. DNS, monitoring, and so on.). Due to this fact such components should be taken under consideration when stipulating compute sources for the pod. You’ll be able to see that, although we’ve got a 4 CPUs in a node, we solely request 3 (you might request fractional CPU sources as properly — e.g. 3.5). You’ll be able to see the allocatable CPU/Reminiscence of every node in GCP (GCP Console → Clusters → Nodes → Click on on a node) or utilizing kubectl describe nodes.
In case your node is unable to meet the compute sources you specify Kubernetes will be unable to run the pods and throw an error (e.g. PodUnschedulable).

Be aware 2: One of the vital essential arguments you want to watch out about is --rest_api_timeout_in_ms=720000 . It takes about 250s to serve up a single request, so right here, we’re setting roughly three instances that point because the timeout, to account for any enqueued requests, after we ship parallel requests. For those who set this to a price that’s too small, your requests will timeout earlier than they’re full.

Defining the service

Right here we’re defining a LoadBalancer kind service, the place we’ll expose the stable-diffusion app by the GCP load balancer. On this strategy, you may be supplied with the load balancer’s IP deal with, the place the load balancer will route the visitors to the pods approaching to it. Customers will likely be making requests towards the IP deal with of the load balancer.

metadata:
identify: stable-diffusion
namespace: default
labels:
app: stable-diffusion
spec:
kind: LoadBalancer
ports:
- port: 8500
protocol: TCP
identify: tf-serving-grpc
- port: 8501
protocol: TCP
identify: tf-serving-http
selector:
app: stable-diffusion

Autoscaling

There’s an crucial matter that we’ve been pushing aside; scaling up our service. In actual world, you might have to serve 1000’s, tens of millions and even billions of shoppers. So as to take action, your service wants to have the ability to scale up/down the variety of nodes/pods within the cluster, based mostly on the demand. Thankfully GCP offers quite a lot of choices, from absolutely managed autoscaling to semi/absolutely person managed autoscaling. You’ll be able to be taught extra about these in this video.

Right here we’ll be utilizing a horizontal pod autoscaler (HPA). The horizontal pod autoscaler will scale up the variety of pods, based mostly on some threshold you present (e.g. CPU or reminiscence utilization). Right here’s an instance.

kubectl autoscale deployment stable-diffusion --cpu-percent=60 --min=1 --max=2

Right here we’re giving the HPA a minimal of 1, most of two pods, and asking it so as to add extra pods if the common CPU throughout present set of pods go above 60%.

Making use of the adjustments

We now obtained all of the constructing blocks prepared to begin our service. Merely run the next instructions.

gcloud container clusters get-credentials sd-cluster --zone us-central1-c && 
kubectl apply -f tf-serving/configmap.yaml && 
kubectl apply -f tf-serving/deployment.yaml && 
kubectl autoscale deployment stable-diffusion --cpu-percent=60 --min=1 --max=2 && 
kubectl apply -f tf-serving/service.yaml

To be able to predict you merely have to make a POST request to the proper URL, with a payload containing the enter to the mannequin.

Sequential predictions

As the primary instance, we present how one can make a collection of requests one after the opposite.

def predict_rest(json_data, url):
json_response = requests.publish(url, knowledge=json_data)
response = json.hundreds(json_response.textual content)
if "predictions" not in response:
print(response)
rest_outputs = np.array(response["predictions"])
return rest_outputs
url = f"http://{stable_diffusion_service_ip}:8501/v1/fashions/stable-diffusion:predict"
tokens_list = [
generate_tokens(tokenizer, "A wine glass made from lego bricks, rainbow colored liquid being poured into it, hyper realistic, high detail", MAX_PROMPT_LENGTH).numpy().tolist(),
generate_tokens(tokenizer, "A staircase made from color pencils, hyper realistic, high detail", MAX_PROMPT_LENGTH).numpy().tolist(),
generate_tokens(tokenizer, "A ferrari car in the space astronaut driving it, futuristic, hyper realistic, high detail", MAX_PROMPT_LENGTH).numpy().tolist(),
generate_tokens(tokenizer, "a dragon covered with weapons fighting an army, fire, explosions, hyper realistic, high detail", MAX_PROMPT_LENGTH).numpy().tolist(),
generate_tokens(tokenizer, "A sawing girl in a boat, hyper realistic, high detail", MAX_PROMPT_LENGTH).numpy().tolist(),
]
negative_tokens = generate_tokens(tokenizer, "ugly, tiling, poorly drawn fingers, poorly drawn toes, poorly drawn face, out of body, mutation, mutated, additional limbs, additional legs, additional arms, disfigured, deformed, cross-eye, physique out of body, blurry, dangerous artwork, dangerous anatomy, blurred, textual content, watermark, grainy", MAX_PROMPT_LENGTH).numpy().tolist()
all_images = []
all_data = []
for tokens, negative_tokens in zip(tokens_list, [negative_tokens for _ in range(5)]):
all_data.append(generate_json_data(tokens, negative_tokens))
all_images = [predict_rest(data, url) for data in all_data]

This took over 1600s once I ran the experiment. As you may think, this setup is sort of inefficient, and is unable to leverage the cluster’s potential to scale up.

Parallel predictions

You should utilize Python’s multiprocessing library to make parallel requests, which is extra remnant of real-world person requests.

def predict_rest(input_data, url):
json_data, sleep_time = input_data["data"], input_data["sleep_time"]
# We add a delay to simulate actual world person requests
time.sleep(sleep_time)
print("Making a request")
t1 = time.perf_counter()
json_response = requests.publish(url, knowledge=json_data)
response = json.hundreds(json_response.textual content)
outcome = np.array([])
attempt: 
outcome = np.array(response["predictions"])
besides KeyError:
print(f"Could not full the request {response}")
lastly:
t2 = time.perf_counter() 
print(f"It took {t2-t1}s to finish a single request")
return outcome
t1 = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
all_images_gen = executor.map(
functools.partial(predict_rest, url=url), 
[{"data": data, "sleep_time": min(i*20, 60)} for i, data in enumerate(all_data)]
)
all_images = [img for img in all_images_gen]
t2 = time.perf_counter()    
print(f"It took {t2-t1}s to finish {n_requests} requests")

This ran in 900s. Due to this fact, we’ve got a ~180% pace up, by scaling up the cluster to a most of two pods.

Be aware about setting timeouts
Watch out when establishing the parallel requests. For those who ship the parallel requests abruptly (since that is solely 6 requests), they’ll probably timeout. It’s because it takes time to create a brand new node and initialize a brand new pod. So if all requests are made immediately, the load balancer won’t even have time to see the 2nd node and find yourself making an attempt to serve all requests to the only node.
Your timeout outlined above is counted from the time the request is obtained (i.e. enter the tensorflow-serving queue), not from the time it begins serving the request. So if the request waits too lengthy within the queue that additionally counts for the timeout.

You’ll be able to monitor the compute metrics such because the CPU utilization and reminiscence consumption on GCP (GCP → Kubernetes Engine → Companies & Ingress → Choose your service).

On this 2 half tutorial, we,

• Setup the infrastructure utilizing terraform (an IaaS instrument), which primarily consisted of a cluster and a node-pool (Half 1)
• Deployed a mannequin and created a prediction service to serve person requests utilizing a Steady Diffusion mannequin (Half 2)

We setup this tutorial in a manner that it may be run by even a free tier person. We setup a cluster with 2 nodes and created 1 pod per node. Then we made each sequential and parallel predictions and noticed that parallel predictions result in ~180% features in greater throughput.

Subsequent steps:

• Model warmup — tensorflow-serving provides a simple solution to heat up a mannequin. You’ll be able to parse instance requests and they are going to be loaded and despatched to the mannequin, previous to serving precise person requests. This can scale back the latency for the preliminary person requests.
• Dynamic batching of requests — You’ll be able to select to dynamically batch the incoming requests. This can permit the mannequin to make predictions on a batch of inputs than predicting on every enter. Given sufficient reminiscence, this can probably to offer throughput features, permitting you to serve lot of requests inside affordable time bounds.

Debugging inside pods

After I’m making an attempt to get this up and operating, a painstaking situation I confronted was operating into the next wall of brick.

and once I go into one of many pods within the deployment, I get a extra wise (nonetheless inconspicuous) error. However it was nonetheless insufficient to place a finger on what precisely was unsuitable.

So I needed to discover a solution to microscopically examine the basis trigger. For that I first logged into the pod’s container in query,

kubectl exec --stdin --tty <container identify> -- /bin/bash

As soon as I’m in, all I obtained to do is capitalize on the “everything is a file” paradigm Linux thrives on. In different phrases, you possibly can merely faucet right into a file to see the output/error stream of a course of. For instance, in my case, tensorflow-serving course of had the PID 7, subsequently, /proc/7/fd/2 offers the error stream of that course of.

tail -n 10  /proc/7/fd/2

In right here, I used to be capable of see precisely why this wasn’t kick-starting. It was as a result of the container didn’t have the mandatory permission to entry the GCS bucket laid out in MODEL_PATH .

Utilizing tf.print for debugging

As you understand, TensorFlow provides two kinds of execution; crucial and declarative. Since we use the __call__() to name the mannequin (i.e. self.mannequin(<inputs>), these calls are executed as graph operations. It’s possible you’ll know already that graph execution is notoriously troublesome to debug, as a consequence of obscurities attributable to the inner graph. One resolution TensorFlow provides is the utilization of tf.print statements.

You’ll be able to place tf.print statements in your mannequin calls, and people print statements are added as operations to the graph, so you possibly can see values of executed tensors, and so on. permitting you to debug the code than throwing darts and hoping for the very best.

Be certain that your tf.print assertion is printing an enter that seems instantly earlier than the time you need it to be printed. For those who add impartial/dummy tf.print statements, they won’t get embedded within the graph within the appropriate place. This may increasingly provide the misleading feeling that some computation is occurring in a short time, because of the incorrect placement of the graph.

Be aware about machine varieties

There are two most important types of machines you need to use for this train; n1 and n2. N2 instances use 3rd generation Xeon processors which might be geared up particular instruction units (AVX-512) to hurry up operations equivalent to matrix multiplication. Due to this fact, CPU sure TensorFlow code runs quicker on n2 machines than on n1.

I’d prefer to acknowledge the ML Developer Programs and the group for the GCP credit offered to make this tutorial successful.