From 82137692c24f1e4a95169b218b5fa0d420c00e55 Mon Sep 17 00:00:00 2001
From: pliny <133052465+elder-plinius@users.noreply.github.com>
Date: Thu, 30 Nov 2023 18:58:53 -0800
Subject: [PATCH] Update stable_diffusion.py

---
 swarms/models/stable_diffusion.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/swarms/models/stable_diffusion.py b/swarms/models/stable_diffusion.py
index 0ba9b463..78d417a9 100644
--- a/swarms/models/stable_diffusion.py
+++ b/swarms/models/stable_diffusion.py
@@ -2,6 +2,7 @@ import base64
 import os
 import requests
 import uuid
+import shutil
 from dotenv import load_dotenv
 from typing import List
 
@@ -135,3 +136,15 @@ class StableDiffusion:
             image_paths.append(image_path)
 
         return image_paths
+
+    def generate_and_move_image(self, prompt, iteration, folder_path):
+        # Generate the image
+        image_paths = self.run(prompt)
+        if not image_paths:
+            return None
+
+        # Move the image to the specified folder
+        src_image_path = image_paths[0]
+        dst_image_path = os.path.join(folder_path, f"image_{iteration}.jpg")
+        shutil.move(src_image_path, dst_image_path)
+        return dst_image_path