Upload 2 files

dc.py

@@ -323,6 +323,7 @@ class GuiSD:
         images, seed, image_list, metadata = model(**pipe_params)
         progress(1, desc="Inference completed.")
         if not isinstance(images, list): images = [images]
+        images = save_images(images, metadata)
         img = []
         for image in images:
             img.append((image, None))
@@ -678,17 +679,11 @@ class GuiSD:
         }
 
         # Maybe fix lora issue: 'Cannot copy out of meta tensor; no data!''
-        self.model.pipe.to("cuda:0" if torch.cuda.is_available() else "cpu")
+        #self.model.pipe.to("cuda:0" if torch.cuda.is_available() else "cpu")
 
         progress(1, desc="Inference preparation completed. Starting inference...")
 
-        info_state = f"PROCESSING "
-        info_state += ">"
-        info_state = f"COMPLETED. Seeds: {str(seed)}"
-        if vae_msg:
-            info_state = info_state + "<br>" + vae_msg
-        if msg_lora:
-            info_state = info_state + "<br>" + "<br>".join(msg_lora)
+        info_state = "" # for yield version
         return self.infer_short(self.model, pipe_params, progress), info_state
 ## END MOD
 
@@ -1310,3 +1305,18 @@ def process_style_prompt(prompt: str, neg_prompt: str, styles_key: str = "None",
     return gr.update(value=prompt), gr.update(value=neg_prompt)
 
 
+from PIL import Image
+def save_images(images: list[Image.Image], metadatas: list[str]):
+    from PIL import PngImagePlugin
+    try:
+        output_images = []
+        for image, metadata in zip(images, metadatas):
+            info = PngImagePlugin.PngInfo()
+            info.add_text("metadata", metadata)
+            savefile = "image.png"
+            image.save(savefile, "PNG", pnginfo=info)
+            output_images.append(str(Path(savefile).resolve()))
+        return output_images
+    except Exception as e:
+        print(f"Failed to save image file: {e}")
+        raise Exception(f"Failed to save image file:") from e