"""Image/tensor format helpers. The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv). Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License. """ import base64 import mimetypes from io import BytesIO import av import numpy as np import torch from comfy_api.input import AudioInput, VideoInput from comfy_extras import nodes_audio from PIL import Image from .logger import get_logger logger = get_logger(__name__) def base64_to_image_tensor(base64_str: str, mode: str = "RGB") -> torch.Tensor: """ Convert base64-encoded image to ComfyUI image tensor. Args: base64_str: Base64-encoded image string mode: PIL image mode (default RGB for transparency support) Returns: torch.Tensor with shape (1, H, W, C) in float32 [0, 1] range Raises: ValueError: If base64 string is invalid or image cannot be decoded """ if base64_str.startswith("data:image"): _, base64_str = base64_str.split(",", 1) try: # Decode base64 to bytes image_bytes = base64.b64decode(base64_str) except Exception as e: raise ValueError(f"Invalid base64 string: {e}") # Create BytesIO object for PIL image_bytesio = BytesIO(image_bytes) # Open with PIL and convert to desired mode try: pil_image = Image.open(image_bytesio) pil_image = pil_image.convert(mode) except Exception as e: raise RuntimeError(f"Failed to open image: {e}") image_array = np.asarray(pil_image).astype(np.float32) / 255.0 image_tensor = torch.from_numpy(image_array) return image_tensor def image_tensor_to_png_bytes(tensor: torch.Tensor, filename: str = "image.png") -> BytesIO: """ Convert ComfyUI image tensor to PNG BytesIO for multipart upload. This function converts a ComfyUI IMAGE tensor to a PNG-encoded BytesIO object suitable for multipart/form-data upload. The BytesIO object has its .name attribute set, which is required by aiohttp for file uploads. Args: tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1] filename: Name attribute to set on BytesIO (default: "image.png") Returns: BytesIO object containing PNG-encoded image with .name attribute set Raises: ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.) """ if tensor.ndim != 4: raise ValueError(f"Expected 4D tensor with shape (B, H, W, C), got {tensor.ndim}D tensor") image_tensor = tensor[0] # Shape: (H, W, C) image_np = (image_tensor.cpu().numpy() * 255.0).astype(np.uint8) pil_image = Image.fromarray(image_np) # Save to BytesIO as image file img_bytes = BytesIO() # Set name attribute (required for multipart upload and mimetype detection) img_bytes.name = filename try: pil_image.save(img_bytes) except Exception as e: raise RuntimeError(f"Failed to save image as file: {e}") # Reset position to beginning img_bytes.seek(0) return img_bytes def image_tensor_to_base64(tensor: torch.Tensor, filename: str = "image.png") -> str: """ Convert ComfyUI image tensor to base64-encoded image string. Args: tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1] filename: Name attribute to set on BytesIO (default: "image.png") format: File format of the output image file buffer (default: "PNG") Returns: Base64-encoded image string Raises: ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.) """ img_bytes = image_tensor_to_png_bytes(tensor, filename) img_bytes.seek(0) byte_data = img_bytes.read() base64_str = base64.b64encode(byte_data).decode("utf-8") mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" return f"data:{mime_type};base64,{base64_str}" def video_to_bytes(video: VideoInput, filename: str = "video.mp4") -> BytesIO: output_buffer = BytesIO() output_buffer.name = filename video.save_to(output_buffer) output_buffer.seek(0) return output_buffer def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str: video_buffer = video_to_bytes(video, filename) video_buffer.seek(0) byte_data = video_buffer.read() base64_str = base64.b64encode(byte_data).decode("utf-8") mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" return f"data:{mime_type};base64,{base64_str}" def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO: waveform = audio["waveform"][0] # Shape: (C, T) sample_rate = audio["sample_rate"] format = filename.rsplit(".", maxsplit=1)[1] layout = "mono" if waveform.shape[0] == 1 else "stereo" output_buffer = BytesIO() output_buffer.name = filename output_container = av.open(output_buffer, mode="w", format=format) if format == "opus": out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout) if quality == "64k": out_stream.bit_rate = 64000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif quality == "96k": out_stream.bit_rate = 96000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif quality == "128k": out_stream.bit_rate = 128000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif quality == "192k": out_stream.bit_rate = 192000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif quality == "320k": out_stream.bit_rate = 320000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif format == "mp3": out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout) if quality == "V0": out_stream.codec_context.qscale = 1 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif quality == "128k": out_stream.bit_rate = 128000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py elif quality == "320k": out_stream.bit_rate = 320000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py else: # format == "flac": out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout) frame = av.AudioFrame.from_ndarray( waveform.movedim(0, 1).reshape(1, -1).float().numpy(), format="flt", layout=layout, ) frame.sample_rate = sample_rate frame.pts = 0 output_container.mux(out_stream.encode(frame)) # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py # Flush encoder output_container.mux(out_stream.encode(None)) # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py output_container.close() output_buffer.seek(0) return output_buffer def audio_to_base64(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> str: audio_buffer = audio_to_bytes(audio, filename, quality) audio_buffer.seek(0) byte_data = audio_buffer.read() base64_str = base64.b64encode(byte_data).decode("utf-8") mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" return f"data:{mime_type};base64,{base64_str}" def bytes_to_audio(audio_bytes: bytes) -> AudioInput: """ Convert audio bytes to ComfyUI audio tensor. Args: audio_bytes: Audio file bytes Returns: torch.Tensor with shape (B, C, T) in float32 range [-1, 1] """ audio_buffer = BytesIO(audio_bytes) waveform, sample_rate = nodes_audio.load(audio_buffer) # type: ignore # Although expect string argument, it calls av.open underneath, which supports BytesIO (file-like) return {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate} def base64_to_audio(base64_str: str) -> AudioInput: """ Convert base64-encoded audio to ComfyUI audio tensor. Args: base64_str: Base64-encoded audio string Returns: torch.Tensor with shape (B, C, T) in float32 range [-1, 1] """ if base64_str.startswith("data:audio"): _, base64_str = base64_str.split(",", 1) try: # Decode base64 to bytes audio_bytes = base64.b64decode(base64_str) except Exception as e: raise ValueError(f"Invalid base64 string: {e}") return bytes_to_audio(audio_bytes)