From 5c3e49c1f0ddde8986c9e1cb4587e8fd0c02cdfb Mon Sep 17 00:00:00 2001
From: Tasha Upchurch <TashaSkyUp@gmail.com>
Date: Sun, 20 Aug 2023 13:46:37 -0400
Subject: [PATCH] Enhanced image processing and memory management in image.py

This commit refactors the handling of the 'steps' attribute, making it configurable via arguments, and adds memory cleanup logic within the 'stackme' method. Additionally, conditional cloning of images ensures that operations are performed only if the images are not None. Commented-out code provides a placeholder for potential future modifications.
---
 image.py | 34 ++++++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/image.py b/image.py
index e406e58..f6a2645 100644
--- a/image.py
+++ b/image.py
@@ -236,6 +236,7 @@ class TinyTxtToImg:
         self.seed = kwargs.get("seed", random.randint(0, 2 ** 32 - 1))
         self.one_seed_per_batch = kwargs.get("one_seed", False)
         self.denoise = kwargs.get("denoise", 1.0)
+        self.steps = kwargs.get("steps", 10)
 
         if not hasattr(TinyTxtToImg, "share_mdl"):
             TinyTxtToImg.share_mdl = None
@@ -273,7 +274,7 @@ class TinyTxtToImg:
                                                 TinyTxtToImg.share_vae
                                                 )
 
-        self.steps = 10
+
         self.cfg = 8
         self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0]
         self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
@@ -320,8 +321,8 @@ class TinyTxtToImg:
                     print(e)
                     raise ValueError("advanced clip encoder failed")
 
-        if not self.latent_image:
-            self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0]
+        #if not self.latent_image == None:
+            #self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0]
 
         self.KSampler = KSampler()
         self.samples = self.KSampler.sample(
@@ -1277,11 +1278,36 @@ class StackImages:
     CATEGORY = "ETK/image"
 
     def stackme(self, image1, image2=None, number_of_images=2, mode="all"):
+        # import modules to help clear memory
+        import gc
+        import torch
+        import torch.cuda
+
+
+
         """
         stacks two images on top of each other
         images are (B, W, H, C) torch tensors
         """
-        image1 = image1.clone()
+        if image1 is not None:
+            image1 = image1.clone()
+
+        if image2 is not None:
+            image2 = image2.clone()
+
+        #gc.collect()
+        #torch.cuda.empty_cache()
+
+        del gc.garbage[:]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+
+
+
+
+
 
         if image2 is not None:
             image2 = image2.clone()