From 265a45fcd7d1b1aee8c4cf2f8c4fa2b5267b8b16 Mon Sep 17 00:00:00 2001
From: ManInDark <61268856+ManInDark@users.noreply.github.com>
Date: Sun, 6 Aug 2023 18:27:09 +0200
Subject: [PATCH] Added prompt and negative_prompt to DB

---
 ui/easydiffusion/bucket_manager.py   | 2 +-
 ui/easydiffusion/easydb/mappings.py  | 6 ++++--
 ui/easydiffusion/utils/save_utils.py | 4 +++-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/ui/easydiffusion/bucket_manager.py b/ui/easydiffusion/bucket_manager.py
index 89537cb8..e4013dc5 100644
--- a/ui/easydiffusion/bucket_manager.py
+++ b/ui/easydiffusion/bucket_manager.py
@@ -106,7 +106,7 @@ def init():
         images = db.query(Image).all()
         sum_string = ""
         for img in images:
-            options = f"Path: {img.path}\nPrompt: \nNegative Prompt: \nSeed: {img.seed}\nModel: {img.use_stable_diffusion_model}\nSize: {img.height}x{img.width}\nSampler: {img.sampler_name}\nSteps: {img.num_inference_steps}\nGuidance Scale: {img.guidance_scale}\nLoRA: {img.lora}\nUpscaling: {img.use_upscale}\nFace Correction: {img.use_face_correction}\n"
+            options = f"Path: {img.path}\nPrompt: {img.prompt}\nNegative Prompt: {img.negative_prompt}\nSeed: {img.seed}\nModel: {img.use_stable_diffusion_model}\nSize: {img.height}x{img.width}\nSampler: {img.sampler_name}\nSteps: {img.num_inference_steps}\nGuidance Scale: {img.guidance_scale}\nLoRA: {img.lora}\nUpscaling: {img.use_upscale}\nFace Correction: {img.use_face_correction}\n"
             sum_string += f"<img src='/image/{img.path}' title='{options}'>"
         return Response(content=sum_string, media_type="text/html")
 
diff --git a/ui/easydiffusion/easydb/mappings.py b/ui/easydiffusion/easydb/mappings.py
index 745f33b2..ad68ecab 100644
--- a/ui/easydiffusion/easydb/mappings.py
+++ b/ui/easydiffusion/easydb/mappings.py
@@ -21,10 +21,12 @@ class Image(Base):
     tiling = Column(String)
     use_face_correction = Column(String)
     use_upscale = Column(String)
+    prompt = Column(String)
+    negative_prompt = Column(String)
 
     def __repr__(self):
-        return "<Image(path='%s', seed='%s', use_stable_diffusion_model='%s', clip_skip='%s', use_vae_model='%s', sampler_name='%s', width='%s', height='%s', num_inference_steps='%s', guidance_scale='%s', lora='%s', use_hypernetwork_model='%s', tiling='%s', use_face_correction='%s', use_upscale='%s')>" % (
-            self.path, self.seed, self.use_stable_diffusion_model, self.clip_skip, self.use_vae_model, self.sampler_name, self.width, self.height, self.num_inference_steps, self.guidance_scale, self.lora, self.use_hypernetwork_model, self.tiling, self.use_face_correction, self.use_upscale)
+        return "<Image(path='%s', seed='%s', use_stable_diffusion_model='%s', clip_skip='%s', use_vae_model='%s', sampler_name='%s', width='%s', height='%s', num_inference_steps='%s', guidance_scale='%s', lora='%s', use_hypernetwork_model='%s', tiling='%s', use_face_correction='%s', use_upscale='%s', prompt='%s', negative_prompt='%s')>" % (
+            self.path, self.seed, self.use_stable_diffusion_model, self.clip_skip, self.use_vae_model, self.sampler_name, self.width, self.height, self.num_inference_steps, self.guidance_scale, self.lora, self.use_hypernetwork_model, self.tiling, self.use_face_correction, self.use_upscale, self.prompt, self.negative_prompt)
 
 from easydiffusion.easydb.database import engine
 Image.metadata.create_all(engine)
\ No newline at end of file
diff --git a/ui/easydiffusion/utils/save_utils.py b/ui/easydiffusion/utils/save_utils.py
index 63d00494..bb508229 100644
--- a/ui/easydiffusion/utils/save_utils.py
+++ b/ui/easydiffusion/utils/save_utils.py
@@ -174,7 +174,9 @@ def save_images_to_disk(
                 use_hypernetwork_model = metadata_entries[i]["use_hypernetwork_model"],
                 tiling = metadata_entries[i]["tiling"],
                 use_face_correction = metadata_entries[i]["use_face_correction"],
-                use_upscale = metadata_entries[i]["use_upscale"]
+                use_upscale = metadata_entries[i]["use_upscale"],
+                prompt = metadata_entries[i]["prompt"],
+                negative_prompt = metadata_entries[i]["negative_prompt"]
             ))
             session.commit()
             session.close()