"""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 fractions import Fraction from io import BytesIO import av import numpy as np import torch from av.audio.frame import AudioFrame from av.audio.resampler import AudioResampler from av.video.frame import VideoFrame from comfy_api.input import AudioInput, VideoInput from comfy_api.latest import InputImpl, Types 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 bytes_to_video(video_bytes: bytes) -> VideoInput: video_buffer = BytesIO(video_bytes) try: with av.open(video_buffer, mode="r") as container: video_stream = next((s for s in container.streams if s.type == "video"), None) if video_stream is None: raise ValueError("No video stream found in decoded payload.") frames = [] for frame in container.decode(video_stream): if not isinstance(frame, VideoFrame): continue image = frame.to_ndarray(format="rgb24") frames.append(torch.from_numpy(image).float() / 255.0) if len(frames) == 0: raise ValueError("No video frames found in decoded payload.") images = torch.stack(frames, dim=0) frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1) audio: AudioInput | None = None if len(container.streams.audio): audio_stream = container.streams.audio[-1] audio_frames = [] resampler = AudioResampler(format="fltp") for frame in container.decode(audio_stream): if not isinstance(frame, AudioFrame): continue resampled = resampler.resample(frame) if not isinstance(resampled, list): resampled = [resampled] for audio_frame in resampled: if audio_frame is not None: audio_frames.append(audio_frame.to_ndarray()) if len(audio_frames) > 0: audio_data = np.concatenate(audio_frames, axis=1) sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1 audio = { "waveform": torch.from_numpy(audio_data).unsqueeze(0), "sample_rate": sample_rate, } components = Types.VideoComponents( images=images, frame_rate=frame_rate, audio=audio, metadata=container.metadata if container.metadata else None, ) except Exception as e: raise RuntimeError(f"Failed to decode video: {e}") return InputImpl.VideoFromComponents(components) def base64_to_video(base64_str: str) -> VideoInput: if base64_str.startswith("data:video"): _, base64_str = base64_str.split(",", 1) try: video_bytes = base64.b64decode(base64_str) except Exception as e: raise ValueError(f"Invalid base64 string: {e}") return bytes_to_video(video_bytes) 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)