From 7ec0ceecc9d79560bc474c0c046a3aa56d0db232 Mon Sep 17 00:00:00 2001
From: Limitex <76650151+Limitex@users.noreply.github.com>
Date: Fri, 5 Jan 2024 16:27:19 +0900
Subject: [PATCH] add img2img

---
 nodes.py | 23 +++++++++++++++++++----
 utils.py | 11 +++++++++--
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/nodes.py b/nodes.py
index 84d3edf..36e4346 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2,7 +2,7 @@ import copy
 import os
 import torch
 from safetensors.torch import load_file
-from .utils import SCHEDULERS, convert_images_to_tensors, token_auto_concat_embeds, vae_pt_to_vae_diffuser
+from .utils import SCHEDULERS, token_auto_concat_embeds, vae_pt_to_vae_diffuser, convert_images_to_tensors, convert_tensors_to_images, resize_images
 from comfy.model_management import get_torch_device
 import folder_paths
 from streamdiffusion import StreamDiffusion
@@ -294,6 +294,9 @@ class StreamDiffusionSampler:
                 "num": ("INT", {"default": 1, "min": 1, "max": 10000}),
                 "warmup": ("INT", {"default": 1, "min": 0, "max": 10000}),
             },
+            "optional" : {
+                "image" : ("IMAGE", )
+            }
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -302,7 +305,7 @@ class StreamDiffusionSampler:
 
     CATEGORY = "Diffusers/StreamDiffusion"
 
-    def sample(self, stream: StreamDiffusion, positive, negative, steps, cfg, delta, seed, num, warmup):
+    def sample(self, stream: StreamDiffusion, positive, negative, steps, cfg, delta, seed, num, warmup, image = None):
         stream.prepare(
             prompt = positive,
             negative_prompt = negative,
@@ -312,13 +315,25 @@ class StreamDiffusionSampler:
             seed = seed
         )
         
+        if image != None:
+            image = convert_tensors_to_images(image)
+            image = resize_images(image, (stream.width, stream.height))
+            
         for _ in range(warmup):
             stream()
 
         result = []
         for _ in range(num):
-            x_output = stream.txt2img()
-            result.append(postprocess_image(x_output, output_type="pil")[0])
+            x_outputs = []
+            if image is None:
+                x_outputs.append(stream.txt2img())
+            else:
+                stream(image[0])
+                for i in image[1:] + image[-1:]:
+                    x_outputs.append(stream(i))
+            for x_output in x_outputs:
+                result.append(postprocess_image(x_output, output_type="pil")[0])
+        
         return (convert_images_to_tensors(result),)
 
 class StreamDiffusionWarmup:
diff --git a/utils.py b/utils.py
index 59f5c9c..866b013 100644
--- a/utils.py
+++ b/utils.py
@@ -2,6 +2,7 @@ import io
 import torch
 import requests
 import numpy as np
+from PIL import Image
 from omegaconf import OmegaConf
 from torchvision.transforms import ToTensor
 from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
@@ -202,5 +203,11 @@ def vae_pt_to_vae_diffuser(
     vae.save_pretrained(output_path)
 
 
-def convert_images_to_tensors(images):
-    return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images])
\ No newline at end of file
+def convert_images_to_tensors(images: list[Image.Image]):
+    return torch.stack([np.transpose(ToTensor()(image), (1, 2, 0)) for image in images])
+
+def convert_tensors_to_images(images: torch.tensor):
+    return [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images]
+
+def resize_images(images: list[Image.Image], size: tuple[int, int]):
+    return [image.resize(size) for image in images]
\ No newline at end of file