diff --git a/__init__.py b/__init__.py index a5fca5a..cec28bc 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,17 @@ -from .nodes import LoadImageBase64, LoadAudioBase64 +# from .nodes import LoadImageBase64, LoadAudioBase64, LastFrameToBase64 +from .nodes import ImageToBase64, Base64ToImage + NODE_CLASS_MAPPINGS = { - "LoadImageBase64": LoadImageBase64, - "LoadAudioBase64": LoadAudioBase64 + "ImageToBase64": ImageToBase64, + "Base64ToImage": Base64ToImage + # "LoadAudioBase64": LoadAudioBase64, + # "LastFrameBase64": LastFrameToBase64 } NODE_DISPLAY_NAME_MAPPINGS = { - "LoadImageBase64": "Load Image (Base64)", - "LoadAudioBase64": "Load Audio (Base64)" + "ImageToBase64": "Encode an image to Base64 string.", + "Base64ToImage": "Decode Base64 string to an image." + # "LoadAudioBase64": "Load Audio (Base64)", + # "LastFrameToBase64": "Load Last Frame (Base64)" } \ No newline at end of file diff --git a/nodes.py b/nodes.py index d31f223..9fb0d90 100644 --- a/nodes.py +++ b/nodes.py @@ -1,8 +1,9 @@ from base64 import b64decode -from .utils import get_image_mask, get_audio +# from .utils import convert_base64_to_image, convert_base_to_audio, convert_tensor_to_base +from .utils import convert_base64_to_image, convert_image_to_base64 -class LoadImageBase64: +class Base64ToImage: @classmethod def INPUT_TYPES(s): return { @@ -11,39 +12,80 @@ class LoadImageBase64: } } - RETURN_TYPES = ("IMAGE", "MASK") - RETURN_NAMES = ("image", "mask") + RETURN_TYPES = ("IMAGE","MASK",) + RETURN_NAMES = ("image","mask",) CATEGORY = "Base64-Nodes" - FUNCTION = "load_base64_image" + FUNCTION = "base64_to_image" - def load_base64_image(self, base64): - image_data = b64decode(base64) - image, mask = get_image_mask(image_data) - return (image, mask) + def base64_to_image(self, base64): + image, mask = convert_base64_to_image(base64) + return (image, mask,) -class LoadAudioBase64: +class ImageToBase64: @classmethod def INPUT_TYPES(s): return { "required": { - "base64": ("STRING", {"multiline": True}) - }, - "optional" : { - "start_time": ("FLOAT", {"default": 0, "min": 0,}), - "duration": ("FLOAT" , {"default": 0, "min": 0, "max": 10000000, "step": 0.01}), + "image": ("IMAGE",) } } - RETURN_TYPES = ("AUDIO",) - RETURN_NAMES = ("audio",) + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("base64",) CATEGORY = "Base64-Nodes" - FUNCTION = "load_base64_audio" + FUNCTION = "image_to_base64" + + def image_to_base64(self, image): + base64 = convert_image_to_base64(image) + return (base64,) - def load_base64_audio(self, base64, start_time=0, duration=0): - audio_data = b64decode(base64) - audio = get_audio(audio_data, start_time, duration) - return (audio,) \ No newline at end of file +# class LoadAudioBase64: +# @classmethod +# def INPUT_TYPES(s): +# return { +# "required": { +# "base64": ("STRING", {"multiline": True}) +# }, +# "optional" : { +# "start_time": ("FLOAT", {"default": 0, "min": 0,}), +# "duration": ("FLOAT" , {"default": 0, "min": 0, "max": 10000000, "step": 0.01}), +# } +# } + +# RETURN_TYPES = ("AUDIO",) +# RETURN_NAMES = ("audio",) + +# CATEGORY = "Base64-Nodes" + +# FUNCTION = "load_base64_audio" + +# def load_base64_audio(self, base64, start_time=0, duration=0): +# audio_data = b64decode(base64) +# audio = convert_base_to_audio(audio_data, start_time, duration) +# return (audio,) + +# class LastFrameToBase64: +# @classmethod +# def INPUT_TYPES(s): +# return { +# "required": { +# "images": ("IMAGE",) +# } +# } + +# RETURN_TYPES = ("STRING",) +# RETURN_NAMES = ("base64",) + +# CATEGORY = "Base64-Nodes" + +# FUNCTION = "load_images" + +# def load_images(self, images): +# last_image = images[-1] +# b64_str = convert_tensor_to_base(last_image) +# print(images[-1].shape) +# return (b64_str,) \ No newline at end of file diff --git a/utils.py b/utils.py index bac9634..7675d69 100644 --- a/utils.py +++ b/utils.py @@ -1,33 +1,72 @@ -from PIL import Image -import numpy as np +from base64 import b64decode, b64encode from io import BytesIO import torch -from imageio_ffmpeg import get_ffmpeg_exe +from PIL import Image +import numpy as np +import torchvision.transforms as transforms import subprocess import re +from imageio_ffmpeg import get_ffmpeg_exe FFMPEG_PATH = get_ffmpeg_exe() ENCODE_ARGS = ("utf-8", 'backslashreplace') -def get_image_mask(image_data): +def convert_base64_to_image(base64_str: str) -> tuple[torch.Tensor, torch.Tensor]: + """ + Decodes a Base64-encoded image into a PyTorch image tensor and a mask. + """ try: - image = Image.open(BytesIO(image_data)) + image = Image.open(BytesIO(b64decode(base64_str))) + W, H = image.size if "A" in image.getbands(): - mask = np.array(image.getchannel("A")).astype(np.float32) / 255.0 - mask = 1.0 - torch.from_numpy(mask) + alpha = np.array(image.getchannel("A")).astype(np.float32) / 255.0 + mask = torch.from_numpy(1.0 - alpha) else: - mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + mask = torch.zeros((H, W), dtype=torch.float32) image = image.convert("RGB") image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] + image = torch.from_numpy(image).contiguous()[None,] return image, mask except Exception as e: - raise RuntimeError(f"Failed to load Base64 image: {e}") + raise RuntimeError(f"Failed to decode image from Base64 string: {e}") + +def convert_image_to_base64(image_tensor: torch.Tensor) -> str: + """ + Encodes a PyTorch image tensor into a Base64-encoded image. + """ + try: + # Handle batch dimension + if image_tensor.ndim == 4: + image_tensor = image_tensor[0] -def get_audio(audio_data, start_time=0, duration=0): + print(image_tensor.shape) + + # Ensure (C, H, W) format + if image_tensor.ndim == 3 and image_tensor.shape[2] in [1, 3, 4]: + image_tensor = image_tensor.permute(2, 0, 1) + + print(image_tensor.shape) + + # Normalize pixel values (0.0, 1.0) + image_tensor = torch.clamp(image_tensor, 0.0, 1.0) + + image = transforms.ToPILImage()(image_tensor) + + buffer = BytesIO() + image.save(buffer, format="PNG") + + return b64encode(buffer.getvalue()).decode("utf-8") + + except Exception as e: + raise RuntimeError(f"Failed to encode image into Base64 string: {e}") + +def convert_base64_to_audio(base64_str: str, start_time=0, duration=0): + """ + Decodes a Base64-encoded audio file into a PyTorch waveform using FFmpeg. + """ try: args = [FFMPEG_PATH, '-i', 'pipe:0', '-vn', '-acodec', 'pcm_f32le', '-f', 'f32le', 'pipe:1'] @@ -39,7 +78,7 @@ def get_audio(audio_data, start_time=0, duration=0): args.insert(2, str(duration)) try: - res = subprocess.run(args, input=audio_data, capture_output=True, check=True) + res = subprocess.run(args, input=base64_str, capture_output=True, check=True) audio = torch.frombuffer(res.stdout, dtype=torch.float32) @@ -54,8 +93,7 @@ def get_audio(audio_data, start_time=0, duration=0): audio = audio.reshape((-1, channels)).transpose(0, 1).unsqueeze(0) return {'waveform': audio, 'sample_rate': sample_rate} - except subprocess.CalledProcessError as e: raise RuntimeError(f"FFmpeg failed:\n{e.stderr.decode(*ENCODE_ARGS)}") except Exception as e: - raise RuntimeError(f"Failed to load Base64 audio: {e}") \ No newline at end of file + raise RuntimeError(f"Failed to decode audio: {e}") \ No newline at end of file