Update src/pipelines/embeddings.py

src/pipelines/embeddings.py

@@ -1,4 +1,5 @@
 import random
+import spaces
 import numpy as np
 import gradio as gr
 import matplotlib.pyplot as plt
@@ -16,17 +17,21 @@ age = get_axis_embeddings(young, old)
 gender = get_axis_embeddings(masculine, feminine)
 royalty = get_axis_embeddings(common, elite)
 
-images = []
-for example in examples:
-    image = pipe(
-        prompt=example,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-    ).images[0]
-    buffer = BytesIO()
-    image.save(buffer, format="JPEG")
-    encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
-    images.append("data:image/jpeg;base64, " + encoded_image)
+@spaces.GPU(enable_queue=True)
+def generate_examples(examples=examples, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale):
+    images = []
+    for example in examples:
+        image = pipe(
+            prompt=example,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        ).images[0]
+        buffer = BytesIO()
+        image.save(buffer, format="JPEG")
+        encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        images.append("data:image/jpeg;base64, " + encoded_image)
+
+generate_examples()
 
 axis = np.vstack([gender, royalty, age])
 axis[1] = calculate_residual(axis, axis_names)