Update sd.py

sd.py

@@ -1,4 +1,5 @@
 import requests
+from PIL import Image, PngImagePlugin
 import base64
 import os
 from datetime import datetime
@@ -54,7 +55,7 @@ def upload_to_imgbb(image_data, file_name):
 def generate_sd(prompt):
     """
     Generate an image using the specified prompt and save it to the cache directory.
-
+    Embeds the YAML payload into the image metadata.
     Returns:
         image_data (bytes): Raw image data as bytes.
         image_path (str): Absolute path to the saved image file using forward slashes.
@@ -85,21 +86,28 @@ def generate_sd(prompt):
         if 'images' in r and r['images']:
             image_data = base64.b64decode(r['images'][0])
             timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-            image_filename = f"image_{timestamp}.jpg"
-            prompt_filename = f"image_{timestamp}.yaml"
+            image_filename = f"image_{timestamp}.png"
             image_path = os.path.abspath(os.path.join(CACHE_DIR, image_filename))
-            prompt_path = os.path.abspath(os.path.join(CACHE_DIR, prompt_filename))
-
-            # Save the prompt to a YAML file
-            with open(prompt_path, 'w') as f:
-                yaml.dump(payload, f)
 
-            # Save the image as a JPEG file
+            # Save the image as a PNG file
             with open(image_path, 'wb') as f:
                 f.write(image_data)
 
+            # Open the image with Pillow
+            img = Image.open(image_path)
+
+            # Convert payload to YAML string
+            yaml_data = yaml.dump(payload)
+
+            # Embed YAML data into image metadata
+            metadata = PngImagePlugin.PngInfo()
+            metadata.add_text("YAML", yaml_data)
+
+            # Save the image with updated metadata
+            img.save(image_path, "PNG", pnginfo=metadata)
+
             # Upload image to ImgBB
-            image_url = upload_to_imgbb(image_data, image_filename)
+            image_url = upload_to_imgbb(image_path, image_filename)
 
             # Convert the path to use forward slashes
             image_path = image_path.replace('\\', '/')
@@ -107,10 +115,10 @@ def generate_sd(prompt):
             return image_data, image_path
         else:
             print("No images found in the response.")
-            return None, None, None
+            return None, None
     except requests.RequestException as e:
         print(f"HTTP request error: {e}")
-        return None, None, None
+        return None, None
     except Exception as e:
         print(f"Error generating image: {e}")
-        return None, None, None
+        return None, None