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
@@ -212,6 +215,25 @@ def get_embed_new(prompt, pipeline, compel, only_convert_string=False, compel_pr
 
     return merge_embeds([prompt_attention_to_invoke_prompt(i) for i in global_prompt_chanks], compel)
 
+# 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)
+
 def add_comma_after_pattern_ti(text):
     pattern = re.compile(r'\b\w+_\d+\b')
     modified_text = pattern.sub(lambda x: x.group() + ',', text)
@@ -272,6 +294,8 @@ def infer(
     # 在 infer 函数中调用 get_embed_new
     if not use_negative_prompt:
         negative_prompt = ""
+    
+    original_prompt = prompt  # Store original prompt for metadata
     prompt = get_embed_new(prompt, pipe, compel, only_convert_string=True)
     negative_prompt = get_embed_new(negative_prompt, pipe, compel, only_convert_string=True)
     conditioning, pooled = compel([prompt, negative_prompt]) # 必须同时处理来保证长度相等
@@ -289,7 +313,25 @@ 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": "MiaoMiaoPixel_V1.0",
+        "use_resolution_binning": use_resolution_binning,
+        "PreUrl": "https://huggingface.co/spaces/Menyu/MiaoPixel"
+    }
+    # Add metadata to the image
+    image_with_metadata = add_metadata_to_image(image, metadata)
+    
+    return image_with_metadata, seed
 
 examples = [
     "nahida (genshin impact)",