Update app.py

app.py

@@ -7,6 +7,9 @@ import os
 from huggingface_hub import hf_hub_download
 from diffusers import StableDiffusionXLPipeline, AutoencoderKL, EulerAncestralDiscreteScheduler
 from compel import Compel, ReturnedEmbeddingsType
+from PIL import Image, PngImagePlugin
+import json
+import io
 
 # =====================================
 # Prompt weights
@@ -216,6 +219,25 @@ def add_comma_after_pattern_ti(text):
     pattern = re.compile(r'\b\w+_\d+\b')
     modified_text = pattern.sub(lambda x: x.group() + ',', text)
     return modified_text
+
+# Add metadata to the image
+def add_metadata_to_image(image, metadata):
+    metadata_str = json.dumps(metadata)
+    
+    # Convert PIL Image to PNG with metadata
+    img_with_metadata = image.copy()
+    
+    # Create a PngInfo object and add metadata
+    png_info = PngImagePlugin.PngInfo()
+    png_info.add_text("parameters", metadata_str)
+    
+    # Save to a byte buffer with metadata
+    buffer = io.BytesIO()
+    img_with_metadata.save(buffer, format="PNG", pnginfo=png_info)
+    
+    # Reopen from buffer to get the image with metadata
+    buffer.seek(0)
+    return Image.open(buffer)
     
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>你现在运行在CPU上 但是此项目只支持GPU.</p>"
@@ -289,7 +311,24 @@ def infer(
         generator=generator,
         use_resolution_binning=use_resolution_binning,
     ).images[0]
-    return image, seed
+    # Create metadata dictionary
+    metadata = {
+        "prompt": original_prompt,
+        "processed_prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "seed": seed,
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "model": "MiaoMiao Harem 1.6",
+        "use_resolution_binning": use_resolution_binning
+    }
+    
+    # Add metadata to the image
+    image_with_metadata = add_metadata_to_image(image, metadata)
+    
+    return image_with_metadata, seed
 
 examples = [
     "nahida (genshin impact)",