From 76eb3d72db783aace51413139c29a45ef278345d Mon Sep 17 00:00:00 2001 From: Shellishack <40737228+Shellishack@users.noreply.github.com> Date: Thu, 4 Sep 2025 16:53:15 +0800 Subject: [PATCH] Add saver nodes --- .env.example | 10 + pyproject.toml | 9 +- src/comfyui_remote_media_io/loader_nodes.py | 368 ++++++++++++++++++ src/comfyui_remote_media_io/nodes.py | 388 +------------------ src/comfyui_remote_media_io/saver_nodes.py | 389 ++++++++++++++++++++ 5 files changed, 786 insertions(+), 378 deletions(-) create mode 100644 .env.example create mode 100644 src/comfyui_remote_media_io/loader_nodes.py create mode 100644 src/comfyui_remote_media_io/saver_nodes.py diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..1bdf5f7 --- /dev/null +++ b/.env.example @@ -0,0 +1,10 @@ +# Set these environment variables for development. +# If you install via ComfyUI Package Manager, make sure +# these environment variables are set in the host machine. + +# Azure Blob Storage configuration +AZURE_TENANT_ID= +AZURE_CLIENT_ID= +AZURE_CLIENT_SECRET= +AZURE_STORAGE_ACCOUNT_NAME= +AZURE_STORAGE_CONTAINER= diff --git a/pyproject.toml b/pyproject.toml index 812dde6..383f5fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "comfyui_remote_media_io" -version = "0.0.2" +version = "0.0.4" description = "Load media (image/video/audio) from remote URL, and save results to remote storage." authors = [ {name = "Ellis Shang", email = "you@gmail.com"} @@ -13,9 +13,7 @@ readme = "README.md" license = {text = "MIT license"} requires-python = ">=3.10" classifiers = [] -dependencies = [ - -] +dependencies = [] @@ -27,7 +25,8 @@ dev = [ "pre-commit", # runs linting on commit "pytest", # testing "ruff", # linting - "Pillow", + "azure-identity", + "azure-storage-blob" ] [project.urls] diff --git a/src/comfyui_remote_media_io/loader_nodes.py b/src/comfyui_remote_media_io/loader_nodes.py new file mode 100644 index 0000000..dd3d2d1 --- /dev/null +++ b/src/comfyui_remote_media_io/loader_nodes.py @@ -0,0 +1,368 @@ +from fractions import Fraction +from inspect import cleandoc +import time +import requests +from PIL import Image, ImageSequence, ImageOps +from io import BytesIO +import torch +import numpy as np +import av + + +class LoadRemoteImage: + """ + A node to load image from URL. + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tulple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Secound value is a config for type "INT", "STRING" or "FLOAT". + """ + return { + "required": { + "image_url": ("STRING", {"multiline": True}), + }, + } + + RETURN_TYPES = ("IMAGE",) + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "load_image" + + # OUTPUT_NODE = False + # OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node + + CATEGORY = "image" + + def pil_to_tensor(self, pil_image): + """ + Convert a PIL image (or sequence of images) to torch tensors. + + Args: + pil_image (PIL.Image): Input image, possibly multi-frame (e.g. GIF). + + Returns: + tuple[torch.Tensor, torch.Tensor]: + - image_tensor: Float tensor in [0,1], shape (frames, H, W, 3) + - mask_tensor: Float tensor in [0,1], shape (frames, H, W) + """ + image_tensors = [] + mask_tensors = [] + + for frame in ImageSequence.Iterator(pil_image): + # Handle EXIF orientation + frame = ImageOps.exif_transpose(frame) + + # Normalize 32-bit integer images + if frame.mode == "I": + frame = frame.point(lambda p: p * (1 / 255)) + + # Convert to RGB float tensor + rgb_array = np.array(frame.convert("RGB"), dtype=np.float32) / 255.0 + rgb_tensor = torch.from_numpy(rgb_array).unsqueeze(0) # (1, H, W, 3) + + # Extract alpha channel as mask if available + if "A" in frame.getbands(): + alpha_array = np.array(frame.getchannel("A"), dtype=np.float32) / 255.0 + mask_tensor = 1.0 - torch.from_numpy(alpha_array) # invert alpha + else: + # Fallback: dummy mask (64x64, all zeros) + mask_tensor = torch.zeros((64, 64), dtype=torch.float32) + + image_tensors.append(rgb_tensor) + mask_tensors.append(mask_tensor.unsqueeze(0)) # add frame dimension + + # Concatenate frames if multiple + if len(image_tensors) > 1: + image_tensor = torch.cat(image_tensors, dim=0) + mask_tensor = torch.cat(mask_tensors, dim=0) + else: + image_tensor = image_tensors[0] + mask_tensor = mask_tensors[0] + + return image_tensor, mask_tensor + + def load_image(self, image_url): + # download image + response = requests.get(image_url) + img = Image.open(BytesIO(response.content)) + return self.pil_to_tensor(img) + + """ + The node will always be re executed if any of the inputs change but + this method can be used to force the node to execute again even when the inputs don't change. + You can make this node return a number or a string. This value will be compared to the one returned the last time the node was + executed, if it is different the node will be executed again. + This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash + changes between executions the LoadImage node is executed again. + """ + + @classmethod + def IS_CHANGED(cls, image_url): + return image_url + + +class LoadRemoteVideo: + """ + A node to load video from URL. + + Class methods + ------------- + INPUT_TYPES (dict): + Defines required input parameters for the node. + IS_CHANGED: + Optional method to control when the node is re-executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tuple. + FUNCTION (`str`): + The entry-point method to run. + CATEGORY (`str`): + The category the node should appear in the UI. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "video_url": ("STRING", {"multiline": True}), + "start_second": ("INT", {"multiline": False}), + "video_length_seconds": ("INT", {"multiline": False}), + }, + } + + RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT") + RETURN_NAMES = ("IMAGES", "AUDIO", "FPS") + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "load_video" + CATEGORY = "image/video" + + def load_video(self, video_url, start_second, video_length_seconds): + """ + Downloads a video from a URL and decodes only the requested segment into tensors + (both video and audio) in a single loop. + """ + + # download to buffer + download_start = time.time() + response = requests.get(video_url, stream=True) + response.raise_for_status() + buffer = BytesIO(response.content) + download_end = time.time() + print(f"Finished downloading in {download_end - download_start:.2f}s.") + + # open container + container = av.open(buffer) + + # pick streams + video_stream = next(s for s in container.streams if s.type == "video") + audio_stream = next((s for s in container.streams if s.type == "audio"), None) + frame_rate = video_stream.average_rate or Fraction(25, 1) + + # compute bounds + start_pts = int(start_second / av.time_base) # in microseconds + end_second = start_second + video_length_seconds + + # seek close to start + container.seek(start_pts, any_frame=False, backward=True, stream=video_stream) + + frames = [] + audio_frames = [] + + # decode interleaved + for packet in container.demux((video_stream, audio_stream) if audio_stream else (video_stream,)): + for frame in packet.decode(): + t = float(frame.pts * frame.time_base) + if t < start_second: + continue + if t >= end_second: + container.close() + images_tensor = torch.cat(frames, dim=0) if frames else None + audio_tensor = None + if audio_frames: + waveform = torch.cat(audio_frames, dim=1).unsqueeze(0) + audio_tensor = { + "waveform": waveform, + "sample_rate": audio_stream.rate, + } + return images_tensor, audio_tensor, frame_rate + + if packet.stream.type == "video": + img = frame.to_ndarray(format="rgb24") + img = torch.from_numpy(img).float() / 255.0 + frames.append(img.unsqueeze(0)) + + elif packet.stream.type == "audio": + arr = frame.to_ndarray() # (channels, samples) + audio_frames.append(torch.from_numpy(arr)) + + # finalize + images_tensor = torch.cat(frames, dim=0) if frames else None + audio_tensor = None + if audio_frames: + waveform = torch.cat(audio_frames, dim=1).unsqueeze(0) + audio_tensor = {"waveform": waveform, "sample_rate": audio_stream.rate} + + return images_tensor, audio_tensor, frame_rate + + @classmethod + def IS_CHANGED(cls, video_url, start_second, video_length_seconds): + return video_url + + +class LoadRemoteAudio: + """ + A example node + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tulple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Secound value is a config for type "INT", "STRING" or "FLOAT". + """ + return { + "required": { + "audio_url": ("STRING", {"multiline": True}), + }, + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("AUDIO",) + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "load_audio" + CATEGORY = "audio" + + def load_audio(self, audio_url): + """ + Downloads audio from a URL and decodes only the requested segment + into a tensor. + """ + + # download to buffer + download_start = time.time() + response = requests.get(audio_url, stream=True) + response.raise_for_status() + buffer = BytesIO(response.content) + download_end = time.time() + print(f"Finished downloading in {download_end - download_start:.2f}s.") + + # open container + container = av.open(buffer) + + # pick audio stream + audio_stream = next((s for s in container.streams if s.type == "audio"), None) + if not audio_stream: + raise ValueError("No audio stream found in the file.") + + audio_frames = [] + + # decode interleaved + for packet in container.demux(audio_stream): + for frame in packet.decode(): + arr = frame.to_ndarray() # (channels, samples) + audio_frames.append(torch.from_numpy(arr)) + + print(f"Decoded {len(audio_frames)} audio frames.") + + waveform = torch.cat(audio_frames, dim=1).unsqueeze(0) + audio_tensor = {"waveform": waveform, "sample_rate": audio_stream.rate} + + print(f"Final audio tensor shape: {audio_tensor['waveform'].shape}, sample rate: {audio_tensor['sample_rate']}") + + return (audio_tensor,) + + """ + The node will always be re executed if any of the inputs change but + this method can be used to force the node to execute again even when the inputs don't change. + You can make this node return a number or a string. This value will be compared to the one returned the last time the node was + executed, if it is different the node will be executed again. + This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash + changes between executions the LoadImage node is executed again. + """ + + @classmethod + def IS_CHANGED(s, audio_url): + return audio_url diff --git a/src/comfyui_remote_media_io/nodes.py b/src/comfyui_remote_media_io/nodes.py index ac0375b..bce4f49 100644 --- a/src/comfyui_remote_media_io/nodes.py +++ b/src/comfyui_remote_media_io/nodes.py @@ -1,384 +1,23 @@ -from fractions import Fraction -from inspect import cleandoc -import time -import requests -from PIL import Image, ImageSequence, ImageOps -from io import BytesIO -import torch -import numpy as np -import av - - -class LoadRemoteImage: - """ - A node to load image from URL. - - Class methods - ------------- - INPUT_TYPES (dict): - Tell the main program input parameters of nodes. - IS_CHANGED: - optional method to control when the node is re executed. - - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tulple. - RETURN_NAMES (`tuple`): - Optional: The name of each output in the output tulple. - FUNCTION (`str`): - The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() - OUTPUT_NODE ([`bool`]): - If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. - The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. - Assumed to be False if not present. - CATEGORY (`str`): - The category the node should appear in the UI. - execute(s) -> tuple || None: - The entry point method. The name of this method must be the same as the value of property `FUNCTION`. - For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. - """ - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - """ - Return a dictionary which contains config for all input fields. - Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". - Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. - The type can be a list for selection. - - Returns: `dict`: - - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - - Value input_fields (`dict`): Contains input fields config: - * Key field_name (`string`): Name of a entry-point method's argument - * Value field_config (`tuple`): - + First value is a string indicate the type of field or a list for selection. - + Secound value is a config for type "INT", "STRING" or "FLOAT". - """ - return { - "required": { - "image_url": ("STRING", {"multiline": True}), - }, - } - - RETURN_TYPES = ("IMAGE",) - DESCRIPTION = cleandoc(__doc__) - FUNCTION = "load_image" - - # OUTPUT_NODE = False - # OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node - - CATEGORY = "image" - - def pil_to_tensor(self, pil_image): - """ - Convert a PIL image (or sequence of images) to torch tensors. - - Args: - pil_image (PIL.Image): Input image, possibly multi-frame (e.g. GIF). - - Returns: - tuple[torch.Tensor, torch.Tensor]: - - image_tensor: Float tensor in [0,1], shape (frames, H, W, 3) - - mask_tensor: Float tensor in [0,1], shape (frames, H, W) - """ - image_tensors = [] - mask_tensors = [] - - for frame in ImageSequence.Iterator(pil_image): - # Handle EXIF orientation - frame = ImageOps.exif_transpose(frame) - - # Normalize 32-bit integer images - if frame.mode == "I": - frame = frame.point(lambda p: p * (1 / 255)) - - # Convert to RGB float tensor - rgb_array = np.array(frame.convert("RGB"), dtype=np.float32) / 255.0 - rgb_tensor = torch.from_numpy(rgb_array).unsqueeze(0) # (1, H, W, 3) - - # Extract alpha channel as mask if available - if "A" in frame.getbands(): - alpha_array = np.array(frame.getchannel("A"), dtype=np.float32) / 255.0 - mask_tensor = 1.0 - torch.from_numpy(alpha_array) # invert alpha - else: - # Fallback: dummy mask (64x64, all zeros) - mask_tensor = torch.zeros((64, 64), dtype=torch.float32) - - image_tensors.append(rgb_tensor) - mask_tensors.append(mask_tensor.unsqueeze(0)) # add frame dimension - - # Concatenate frames if multiple - if len(image_tensors) > 1: - image_tensor = torch.cat(image_tensors, dim=0) - mask_tensor = torch.cat(mask_tensors, dim=0) - else: - image_tensor = image_tensors[0] - mask_tensor = mask_tensors[0] - - return image_tensor, mask_tensor - - def load_image(self, image_url): - # download image - response = requests.get(image_url) - img = Image.open(BytesIO(response.content)) - return self.pil_to_tensor(img) - - """ - The node will always be re executed if any of the inputs change but - this method can be used to force the node to execute again even when the inputs don't change. - You can make this node return a number or a string. This value will be compared to the one returned the last time the node was - executed, if it is different the node will be executed again. - This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash - changes between executions the LoadImage node is executed again. - """ - - @classmethod - def IS_CHANGED(cls, image_url): - return image_url - - -class LoadRemoteVideo: - """ - A node to load video from URL. - - Class methods - ------------- - INPUT_TYPES (dict): - Defines required input parameters for the node. - IS_CHANGED: - Optional method to control when the node is re-executed. - - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tuple. - FUNCTION (`str`): - The entry-point method to run. - CATEGORY (`str`): - The category the node should appear in the UI. - """ - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "video_url": ("STRING", {"multiline": True}), - "start_second": ("INT", {"multiline": False}), - "video_length_seconds": ("INT", {"multiline": False}), - }, - } - - RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT") - RETURN_NAMES = ("IMAGES", "AUDIO", "FPS") - DESCRIPTION = cleandoc(__doc__) - FUNCTION = "load_video" - CATEGORY = "image/video" - - def load_video(self, video_url, start_second, video_length_seconds): - """ - Downloads a video from a URL and decodes only the requested segment into tensors - (both video and audio) in a single loop. - """ - - # download to buffer - download_start = time.time() - response = requests.get(video_url, stream=True) - response.raise_for_status() - buffer = BytesIO(response.content) - download_end = time.time() - print(f"Finished downloading in {download_end - download_start:.2f}s.") - - # open container - container = av.open(buffer) - - # pick streams - video_stream = next(s for s in container.streams if s.type == "video") - audio_stream = next((s for s in container.streams if s.type == "audio"), None) - frame_rate = video_stream.average_rate or Fraction(25, 1) - - # compute bounds - start_pts = int(start_second / av.time_base) # in microseconds - end_second = start_second + video_length_seconds - - # seek close to start - container.seek(start_pts, any_frame=False, backward=True, stream=video_stream) - - frames = [] - audio_frames = [] - - # decode interleaved - for packet in container.demux( - (video_stream, audio_stream) if audio_stream else (video_stream,) - ): - for frame in packet.decode(): - t = float(frame.pts * frame.time_base) - if t < start_second: - continue - if t >= end_second: - container.close() - images_tensor = torch.cat(frames, dim=0) if frames else None - audio_tensor = None - if audio_frames: - waveform = torch.cat(audio_frames, dim=1).unsqueeze(0) - audio_tensor = { - "waveform": waveform, - "sample_rate": audio_stream.rate, - } - return images_tensor, audio_tensor, frame_rate - - if packet.stream.type == "video": - img = frame.to_ndarray(format="rgb24") - img = torch.from_numpy(img).float() / 255.0 - frames.append(img.unsqueeze(0)) - - elif packet.stream.type == "audio": - arr = frame.to_ndarray() # (channels, samples) - audio_frames.append(torch.from_numpy(arr)) - - # finalize - images_tensor = torch.cat(frames, dim=0) if frames else None - audio_tensor = None - if audio_frames: - waveform = torch.cat(audio_frames, dim=1).unsqueeze(0) - audio_tensor = {"waveform": waveform, "sample_rate": audio_stream.rate} - - return images_tensor, audio_tensor, frame_rate - - @classmethod - def IS_CHANGED(cls, video_url, start_second, video_length_seconds): - return video_url - - -class LoadRemoteAudio: - """ - A example node - - Class methods - ------------- - INPUT_TYPES (dict): - Tell the main program input parameters of nodes. - IS_CHANGED: - optional method to control when the node is re executed. - - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tulple. - RETURN_NAMES (`tuple`): - Optional: The name of each output in the output tulple. - FUNCTION (`str`): - The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() - OUTPUT_NODE ([`bool`]): - If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. - The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. - Assumed to be False if not present. - CATEGORY (`str`): - The category the node should appear in the UI. - execute(s) -> tuple || None: - The entry point method. The name of this method must be the same as the value of property `FUNCTION`. - For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. - """ - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - """ - Return a dictionary which contains config for all input fields. - Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". - Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. - The type can be a list for selection. - - Returns: `dict`: - - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - - Value input_fields (`dict`): Contains input fields config: - * Key field_name (`string`): Name of a entry-point method's argument - * Value field_config (`tuple`): - + First value is a string indicate the type of field or a list for selection. - + Secound value is a config for type "INT", "STRING" or "FLOAT". - """ - return { - "required": { - "audio_url": ("STRING", {"multiline": True}), - }, - } - - RETURN_TYPES = ("AUDIO",) - RETURN_NAMES = ("AUDIO",) - DESCRIPTION = cleandoc(__doc__) - FUNCTION = "load_audio" - CATEGORY = "audio" - - def load_audio(self, audio_url): - """ - Downloads audio from a URL and decodes only the requested segment - into a tensor. - """ - - # download to buffer - download_start = time.time() - response = requests.get(audio_url, stream=True) - response.raise_for_status() - buffer = BytesIO(response.content) - download_end = time.time() - print(f"Finished downloading in {download_end - download_start:.2f}s.") - - # open container - container = av.open(buffer) - - # pick audio stream - audio_stream = next((s for s in container.streams if s.type == "audio"), None) - if not audio_stream: - raise ValueError("No audio stream found in the file.") - - audio_frames = [] - - # decode interleaved - for packet in container.demux(audio_stream): - for frame in packet.decode(): - - arr = frame.to_ndarray() # (channels, samples) - audio_frames.append(torch.from_numpy(arr)) - - print(f"Decoded {len(audio_frames)} audio frames.") - - waveform = torch.cat(audio_frames, dim=1).unsqueeze(0) - audio_tensor = {"waveform": waveform, "sample_rate": audio_stream.rate} - - print( - f"Final audio tensor shape: {audio_tensor['waveform'].shape}, sample rate: {audio_tensor['sample_rate']}" - ) - - return (audio_tensor,) - - """ - The node will always be re executed if any of the inputs change but - this method can be used to force the node to execute again even when the inputs don't change. - You can make this node return a number or a string. This value will be compared to the one returned the last time the node was - executed, if it is different the node will be executed again. - This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash - changes between executions the LoadImage node is executed again. - """ - - @classmethod - def IS_CHANGED(s, audio_url): - return audio_url +from .loader_nodes import ( + LoadRemoteAudio, + LoadRemoteImage, + LoadRemoteVideo, +) +from .saver_nodes import SaveAudioToRemote, SaveImageToRemote, SaveVideoToRemote # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique + + NODE_CLASS_MAPPINGS = { "LoadRemoteImage": LoadRemoteImage, "LoadRemoteVideo": LoadRemoteVideo, "LoadRemoteAudio": LoadRemoteAudio, + # This nodes require a input (bucket saving path & name) + "SaveImageToRemote": SaveImageToRemote, + "SaveVideoToRemote": SaveVideoToRemote, + "SaveAudioToRemote": SaveAudioToRemote, } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -386,4 +25,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LoadRemoteImage": "Load Remote Image", "LoadRemoteVideo": "Load Remote Video", "LoadRemoteAudio": "Load Remote Audio", + "SaveImageToRemote": "Save Image To Remote", + "SaveVideoToRemote": "Save Video To Remote", + "SaveAudioToRemote": "Save Audio To Remote", } diff --git a/src/comfyui_remote_media_io/saver_nodes.py b/src/comfyui_remote_media_io/saver_nodes.py new file mode 100644 index 0000000..01a6a87 --- /dev/null +++ b/src/comfyui_remote_media_io/saver_nodes.py @@ -0,0 +1,389 @@ +from inspect import cleandoc +import io +import os +from PIL import Image +from io import BytesIO +import numpy as np +import av +from azure.identity import DefaultAzureCredential +from azure.storage.blob import BlobServiceClient +import torchaudio + +from dotenv import load_dotenv, find_dotenv + +load_dotenv(find_dotenv()) + +azure_credential = DefaultAzureCredential() +blob_service_client = BlobServiceClient( + account_url=f"https://{os.getenv('AZURE_STORAGE_ACCOUNT_NAME')}.blob.core.windows.net", + credential=azure_credential, +) + + +def get_blob_client(blob_name: str): + print("Container name:", os.getenv("AZURE_STORAGE_CONTAINER")) + blob_client = blob_service_client.get_blob_client(container=os.getenv("AZURE_STORAGE_CONTAINER"), blob=blob_name) + + return blob_client + + +class SaveImageToRemote: + """ + A node to load image from URL. + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tulple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Secound value is a config for type "INT", "STRING" or "FLOAT". + """ + return { + "required": { + "image": ("IMAGE", {"tooltip": "The image to save."}), + "filename": ( + "STRING", + {"tooltip": "The full path to save the file."}, + ), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("URL",) + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "save_image" + + OUTPUT_NODE = True + # OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node + + CATEGORY = "image" + + def tensor_to_pil(self, image_tensor): + """ + Convert a torch tensor back to a list of PIL images. + + Args: + image_tensor (torch.Tensor): Float tensor in [0,1], + shape (frames, H, W, 3) or (H, W, 3) + + Returns: + list[PIL.Image.Image]: One or more PIL images (frames). + """ + # Ensure 4D shape: (frames, H, W, 3) + if image_tensor.ndim == 3: + image_tensor = image_tensor.unsqueeze(0) + + frames = [] + for i in range(image_tensor.shape[0]): + rgb = image_tensor[i].clamp(0, 1).numpy() # (H, W, 3) + rgb_uint8 = (rgb * 255).astype(np.uint8) + frame = Image.fromarray(rgb_uint8, mode="RGB") + frames.append(frame) + + return frames + + def save_image(self, image, filename): + format = filename.split(".")[-1].lower() + + if format not in ["png"]: + raise ValueError("SaveImageToRemote only supports PNG format.") + + frames = self.tensor_to_pil(image) + + file = BytesIO() + frames[0].save(file, format=format) + file.seek(0) + + # upload image to Azure storage + try: + blob_client = get_blob_client(filename) + blob_client.upload_blob(file, overwrite=True) + except Exception as e: + print("Error uploading to Azure Blob Storage:", e) + raise e + + return ( + f"https://{os.getenv('AZURE_STORAGE_ACCOUNT_NAME')}.blob.core.windows.net/{os.getenv('AZURE_STORAGE_CONTAINER')}/{filename}", + ) + + +class SaveVideoToRemote: + """ + A node to load video from URL. + + Class methods + ------------- + INPUT_TYPES (dict): + Defines required input parameters for the node. + IS_CHANGED: + Optional method to control when the node is re-executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tuple. + FUNCTION (`str`): + The entry-point method to run. + CATEGORY (`str`): + The category the node should appear in the UI. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "video": ("VIDEO", {"tooltip": "The video to save."}), + "filename": ( + "STRING", + {"tooltip": "The full path to save the file."}, + ), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("URL",) + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "save_video" + CATEGORY = "image/video" + + OUTPUT_NODE = True + + def save_video(self, video, filename): + format = filename.split(".")[-1].lower() + + # Save to os temp folder + output_path = f"/tmp/{filename}" + + # Use VideoFromFile.save_to from ComfyUI python backend. + video.save_to( + output_path, + format=format, + codec="h264", + ) + + file = open(output_path, "rb") + + # upload video to Azure storage + try: + blob_client = get_blob_client(filename) + blob_client.upload_blob(file, overwrite=True) + except Exception as e: + print("Error uploading to Azure Blob Storage:", e) + raise e + finally: + file.close() + os.remove(output_path) + + return ( + f"https://{os.getenv('AZURE_STORAGE_ACCOUNT_NAME')}.blob.core.windows.net/{os.getenv('AZURE_STORAGE_CONTAINER')}/{filename}", + ) + + +class SaveAudioToRemote: + """ + A example node + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tulple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Secound value is a config for type "INT", "STRING" or "FLOAT". + """ + return { + "required": { + "audio": ("AUDIO", {"tooltip": "The audio to save."}), + "filename": ("STRING", {"tooltip": "The full path to save the file."}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("URL",) + DESCRIPTION = cleandoc(__doc__) + FUNCTION = "save_audio" + CATEGORY = "audio" + + OUTPUT_NODE = True + + def save_audio( + self, + audio, + filename, + quality="128k", + ): + format = filename.split(".")[-1].lower() + + # Opus supported sample rates + OPUS_RATES = [8000, 12000, 16000, 24000, 48000] + + # Only process 1 batch at a time for now + batch_number = audio["waveform"].shape[0] + if batch_number > 1: + raise ValueError("SaveAudioToRemote only supports batch size of 1.") + + for batch_number, waveform in enumerate(audio["waveform"].cpu()): + output_path = f"/tmp/{filename}" + + # Use original sample rate initially + sample_rate = audio["sample_rate"] + + # Handle Opus sample rate requirements + if format == "opus": + if sample_rate > 48000: + sample_rate = 48000 + elif sample_rate not in OPUS_RATES: + # Find the next highest supported rate + for rate in sorted(OPUS_RATES): + if rate > sample_rate: + sample_rate = rate + break + if sample_rate not in OPUS_RATES: # Fallback if still not supported + sample_rate = 48000 + + # Resample if necessary + if sample_rate != audio["sample_rate"]: + waveform = torchaudio.functional.resample(waveform, audio["sample_rate"], sample_rate) + + # Create output with specified format + output_buffer = io.BytesIO() + output_container = av.open(output_buffer, mode="w", format=format) + + # Set up the output stream with appropriate properties + if format == "opus": + out_stream = output_container.add_stream("libopus", rate=sample_rate) + if quality == "64k": + out_stream.bit_rate = 64000 + elif quality == "96k": + out_stream.bit_rate = 96000 + elif quality == "128k": + out_stream.bit_rate = 128000 + elif quality == "192k": + out_stream.bit_rate = 192000 + elif quality == "320k": + out_stream.bit_rate = 320000 + elif format == "mp3": + out_stream = output_container.add_stream("libmp3lame", rate=sample_rate) + if quality == "V0": + # TODO i would really love to support V3 and V5 but there doesn't seem to be a way to set the qscale level, the property below is a bool + out_stream.codec_context.qscale = 1 + elif quality == "128k": + out_stream.bit_rate = 128000 + elif quality == "320k": + out_stream.bit_rate = 320000 + else: # format == "flac": + out_stream = output_container.add_stream("flac", rate=sample_rate) + + frame = av.AudioFrame.from_ndarray( + waveform.movedim(0, 1).reshape(1, -1).float().numpy(), + format="flt", + layout="mono" if waveform.shape[0] == 1 else "stereo", + ) + frame.sample_rate = sample_rate + frame.pts = 0 + output_container.mux(out_stream.encode(frame)) + + # Flush encoder + output_container.mux(out_stream.encode(None)) + + # Close containers + output_container.close() + + # Write the output to file + output_buffer.seek(0) + with open(output_path, "wb") as f: + f.write(output_buffer.getbuffer()) + + file = open(output_path, "rb") + # upload audio to Azure storage + try: + blob_client = get_blob_client(filename) + blob_client.upload_blob(file, overwrite=True) + except Exception as e: + print("Error uploading to Azure Blob Storage:", e) + raise e + finally: + file.close() + os.remove(output_path) + + return ( + f"https://{os.getenv('AZURE_STORAGE_ACCOUNT_NAME')}.blob.core.windows.net/{os.getenv('AZURE_STORAGE_CONTAINER')}/{filename}", + )