From 7a3314a92792b976eeb15ddbe0b0122d909a1d8a Mon Sep 17 00:00:00 2001
From: ManInDark <61268856+ManInDark@users.noreply.github.com>
Date: Fri, 4 Aug 2023 23:16:43 +0200
Subject: [PATCH] Added mapping for images and automatically adding them to
 database

---
 ui/easydiffusion/easydb/mappings.py  | 30 ++++++++++++++++++++
 ui/easydiffusion/utils/save_utils.py | 42 ++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 ui/easydiffusion/easydb/mappings.py

diff --git a/ui/easydiffusion/easydb/mappings.py b/ui/easydiffusion/easydb/mappings.py
new file mode 100644
index 00000000..745f33b2
--- /dev/null
+++ b/ui/easydiffusion/easydb/mappings.py
@@ -0,0 +1,30 @@
+from sqlalchemy import Column, Integer, String, Float, Boolean
+from sqlalchemy.ext.declarative import declarative_base
+
+Base = declarative_base()
+
+class Image(Base):
+    __tablename__ = 'images'
+
+    path = Column(String, primary_key=True)
+    seed = Column(Integer)
+    use_stable_diffusion_model = Column(String)
+    clip_skip = Column(Boolean)
+    use_vae_model = Column(String)
+    sampler_name = Column(String)
+    width = Column(Integer)
+    height = Column(Integer)
+    num_inference_steps = Column(Integer)
+    guidance_scale = Column(Float)
+    lora = Column(String)
+    use_hypernetwork_model = Column(String)
+    tiling = Column(String)
+    use_face_correction = Column(String)
+    use_upscale = 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)
+
+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 49743554..0082d7f2 100644
--- a/ui/easydiffusion/utils/save_utils.py
+++ b/ui/easydiffusion/utils/save_utils.py
@@ -140,6 +140,48 @@ def save_images_to_disk(
             output_quality=output_format.output_quality,
             output_lossless=output_format.output_lossless,
         )
+
+        for i in range(len(filtered_images)):
+            path_i = f"{os.path.join(save_dir_path, make_filename(i))}.{output_format.output_format.lower()}"
+
+            def createLoraString(metadata_entries, i):
+                if metadata_entries[i]["use_lora_model"] is None:
+                    return "None"
+                elif isinstance(metadata_entries[i]["use_lora_model"], list):
+                    loraString = ""
+                    for j in range(len(metadata_entries[i]["use_lora_model"])):
+                        loraString += metadata_entries[i]["use_lora_model"][j] + ":" + str(metadata_entries[i]["lora_alpha"][j]) + " "
+                    return loraString.trim()
+                else:
+                    return metadata_entries[i]["use_lora_model"] + ":" + str(metadata_entries[i]["lora_alpha"])
+
+            from easydiffusion.easydb.mappings import Image
+            from easydiffusion.easydb.database import SessionLocal
+
+            session = SessionLocal()
+            
+            img = Image(
+                path = path_i,
+                seed = metadata_entries[i]["seed"],
+                use_stable_diffusion_model = metadata_entries[i]["use_stable_diffusion_model"],
+                clip_skip = metadata_entries[i]["clip_skip"],
+                use_vae_model = metadata_entries[i]["use_vae_model"],
+                sampler_name = metadata_entries[i]["sampler_name"],
+                width = metadata_entries[i]["width"],
+                height = metadata_entries[i]["height"],
+                num_inference_steps = metadata_entries[i]["num_inference_steps"],
+                guidance_scale = metadata_entries[i]["guidance_scale"],
+                lora = createLoraString(metadata_entries, i),
+                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"]
+            )
+
+            session.add(img)
+            session.commit()
+            session.close()
+        
         if task_data.metadata_output_format:
             for metadata_output_format in task_data.metadata_output_format.split(","):
                 if metadata_output_format.lower() in ["json", "txt", "embed"]: