From 883c9e9b46aef712ac71eee98c8795349404f194 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Tue, 26 Mar 2024 18:36:32 -0700
Subject: [PATCH] Add duration metrics to Stable Diffusion XL library.

PiperOrigin-RevId: 619375261
---
 test/gpu/stablediffusion/stablediffusion.go | 37 +++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/test/gpu/stablediffusion/stablediffusion.go b/test/gpu/stablediffusion/stablediffusion.go
index 7d1c2c90f..9b6cba9ab 100644
--- a/test/gpu/stablediffusion/stablediffusion.go
+++ b/test/gpu/stablediffusion/stablediffusion.go
@@ -149,6 +149,43 @@ func (i *XLImage) Image() (image.Image, error) {
 	return png.Decode(base64.NewDecoder(base64.StdEncoding, bytes.NewBufferString(strings.Join(i.data.ImagePNGBase64, ""))))
 }
 
+// TotalDuration returns the total time taken to generate the image.
+func (i *XLImage) TotalDuration() time.Duration {
+	return i.data.Done.Sub(i.data.Start)
+}
+
+// ColdBaseDuration returns time taken to run the base image generation model
+// the first time the image was generated (i.e. the model was cold).
+func (i *XLImage) ColdBaseDuration() time.Duration {
+	return i.data.ColdBaseDone.Sub(i.data.ColdStartImage)
+}
+
+// ColdRefinerDuration returns time taken to run the refiner model
+// the first time the image was generated (i.e. the model was cold).
+// Returns -1 if the refiner was not run.
+func (i *XLImage) ColdRefinerDuration() time.Duration {
+	if !i.Prompt.UseRefiner {
+		return -1
+	}
+	return i.data.ColdRefinerDone.Sub(i.data.ColdBaseDone)
+}
+
+// WarmBaseDuration returns time taken to run the base image generation model
+// the second time the image was generated (i.e. the model was warm).
+func (i *XLImage) WarmBaseDuration() time.Duration {
+	return i.data.WarmBaseDone.Sub(i.data.WarmStartImage)
+}
+
+// WarmRefinerDuration returns time taken to run the refiner model
+// the second time the image was generated (i.e. the model was warm).
+// Returns -1 if the refiner was not run.
+func (i *XLImage) WarmRefinerDuration() time.Duration {
+	if !i.Prompt.UseRefiner {
+		return -1
+	}
+	return i.data.WarmRefinerDone.Sub(i.data.WarmBaseDone)
+}
+
 // Generate generates an image with Stable Diffusion XL.
 func (xl *XL) Generate(ctx context.Context, prompt *XLPrompt) (*XLImage, error) {
 	argv := []string{