From c0f2c852c41240d88cf6ad120bca45aaf82a03cb Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Fri, 1 Mar 2024 15:04:30 -0500
Subject: [PATCH] Transparent image infotext + saving support

---
 lib_layerdiffusion/models.py    | 21 ++++++++++++++++++++-
 scripts/forge_layerdiffusion.py | 12 +++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/lib_layerdiffusion/models.py b/lib_layerdiffusion/models.py
index d7cb040..d203ef4 100644
--- a/lib_layerdiffusion/models.py
+++ b/lib_layerdiffusion/models.py
@@ -3,6 +3,7 @@
 import cv2
 import numpy as np
 
+from PIL import Image
 from tqdm import tqdm
 from typing import Optional, Tuple
 from diffusers.configuration_utils import ConfigMixin, register_to_config
@@ -11,6 +12,9 @@
 import ldm_patched.modules.model_management as model_management
 from ldm_patched.modules.model_patcher import ModelPatcher
 
+from modules import images, processing
+from modules.shared import opts
+
 
 def zero_module(module):
     """
@@ -262,7 +266,22 @@ def wrapper(func, latent):
 
             pngs = torch.cat([fg, alpha], dim=3)
             pngs = (pngs * 255.0).detach().cpu().float().numpy().clip(0, 255).astype(np.uint8)
-            for png in pngs:
+            for i, png in enumerate(pngs):
+                png = Image.fromarray(png)
+                infotext = processing.Processed(p, []).infotext(p, i)
+                
+                if getattr(opts, 'layerdiffusion_save_transparent_images', False):
+                    images.save_image(
+                        image=png,
+                        path=p.outpath_samples,
+                        basename="",
+                        seed=p.seeds[i],
+                        prompt=p.prompts[i],
+                        extension=getattr(opts, 'samples_format', 'png'),
+                        info=infotext,
+                        p=p,
+                        suffix="-transparent"
+                    )
                 p.extra_result_images.append(png)
 
             return vis
diff --git a/scripts/forge_layerdiffusion.py b/scripts/forge_layerdiffusion.py
index 2ddf821..f87914f 100644
--- a/scripts/forge_layerdiffusion.py
+++ b/scripts/forge_layerdiffusion.py
@@ -5,7 +5,7 @@
 import numpy as np
 import copy
 
-from modules import scripts
+from modules import scripts, script_callbacks, shared
 from modules.processing import StableDiffusionProcessing
 from lib_layerdiffusion.enums import ResizeMode
 from lib_layerdiffusion.utils import rgba2rgbfp32, to255unit8, crop_and_resize_image
@@ -246,3 +246,13 @@ def conditioning_modifier(model, x, timestep, uncond, cond, cond_scale, model_op
         p.sd_model.forge_objects.unet = unet
         p.sd_model.forge_objects.vae = vae
         return
+
+def on_ui_settings():
+    section = ("Layer Diffusion", "Layer Diffusion")
+
+    shared.opts.add_option(
+        "layerdiffusion_save_transparent_images",
+        shared.OptionInfo(False, "Save images with transparency", section=section),
+    )
+
+script_callbacks.on_ui_settings(on_ui_settings)