From c46bfda0eefe0d6c852cc7dbfea38edb2f9bbb0b Mon Sep 17 00:00:00 2001
From: TashaSkyUp <tashaskyup@gmail.com>
Date: Wed, 27 Nov 2024 22:04:51 -0700
Subject: [PATCH] fix for vision in v2 litellm node

---
 audio.py |  9 ++++++++-
 image.py | 40 +++++++++++++++++++++++++++-------------
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/audio.py b/audio.py
index 25182a5..a449fb4 100644
--- a/audio.py
+++ b/audio.py
@@ -3,6 +3,8 @@ import shutil
 import copy
 import imageio_ffmpeg as ffmpeg
 import joblib
+
+from comfy_extras.nodes_audio import insert_or_replace_vorbis_comment
 from .config import config_settings
 
 temp_dir = config_settings["tmp_dir"]
@@ -558,15 +560,20 @@ class EtkSaveAudio:
     OUTPUT_NODE = True
 
     def save_audio(self, audio, filename=""):
+        import torchaudio
+        import io
+
+
         results = list()
         filename += self.prefix_append
 
+
         for (batch_number, waveform) in enumerate(audio["waveform"].cpu()):
             file = f"{filename}_{batch_number:05}.flac"
 
             buff = io.BytesIO()
             torchaudio.save(buff, waveform, audio["sample_rate"], format="FLAC")
-            buff = insert_or_replace_vorbis_comment(buff, metadata)
+            buff = insert_or_replace_vorbis_comment(buff, {})
 
             with open(file, 'wb') as f:
                 f.write(buff.getbuffer())
diff --git a/image.py b/image.py
index 7d1db29..4b7b4de 100644
--- a/image.py
+++ b/image.py
@@ -135,37 +135,50 @@ def torch_image_show(image):
 
 @ETK_image_base
 class LoadImageFromPath:
-    """receives a comfyui "IMAGE" tensor of BHWC and a path to a file converts the tensor to a PIL image and saves it to the path"""
 
     @classmethod
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "image": ("IMAGE",),
-                "path": ("STRING", {"default": "output/image.png"}),
+                "image_path": ("STRING", {"default": "output/image.png"}),
             },
         }
 
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("saved_path",)
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
     FUNCTION = "execute"
     OUTPUT_NODE = True
 
-    def execute(self, image: torch.Tensor, path: str):
+    def execute(self, image_path: str):
+        import torch
         from PIL import Image
         import numpy as np
         from pathlib import Path
 
-        path = Path(path)
-        path.parent.mkdir(parents=True, exist_ok=True)
+        path = Path(image_path)
+        loaded_image = Image.open(path)
+        loaded_image = np.array(loaded_image).astype(np.float32) / 255.0
+        loaded_image = torch.from_numpy(loaded_image)[None,]
+
+        # torch tensor mus be (B,W,H,C)
+        if len(loaded_image.shape) == 3:
+            loaded_image = loaded_image.unsqueeze(0)
+        return (loaded_image,)
+
+    @classmethod
+    def IS_CHANGED(s, image_path):
+        import hashlib
+        from pathlib import Path
+        # use the file path and modified date + file size (all as a string) to determine if the file has changed
+        path = Path(image_path)
+        m = hashlib.sha256()
+        m.update(str(path.stat().st_mtime).encode())
+        m.update(str(path.stat().st_size).encode())
+        return m.digest().hex()
+
 
-        # Convert torch tensor to PIL Image
-        pil_image = Image.fromarray((image.squeeze().cpu().numpy() * 255).astype(np.uint8))
 
-        # Save image
-        pil_image.save(path)
 
-        return (str(path),)
 
 
 @ETK_image_base
@@ -2372,6 +2385,7 @@ class StripAlphaChannel:
         >>> result = sac.strip_alpha_channel(image)
         >>> torch_image_show(result[0])
         """
+        from copy import deepcopy
 
         return (deepcopy(image[..., :3]),)