99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
from base64 import b64decode, b64encode
|
|
from io import BytesIO
|
|
import torch
|
|
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 convert_base64_to_image(base64_str: str) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Decodes a Base64-encoded image into a ComfyUI image tensor and an alpha mask.
|
|
"""
|
|
try:
|
|
image = Image.open(BytesIO(b64decode(base64_str)))
|
|
W, H = image.size
|
|
|
|
if "A" in image.getbands():
|
|
alpha = np.array(image.getchannel("A")).astype(np.float32) / 255.0
|
|
mask = torch.from_numpy(1.0 - alpha)
|
|
else:
|
|
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).contiguous()[None,]
|
|
|
|
return image, mask
|
|
except Exception as 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]
|
|
|
|
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']
|
|
|
|
if start_time > 0:
|
|
args.insert(1, '-ss')
|
|
args.insert(2, str(start_time))
|
|
if duration > 0:
|
|
args.insert(1, '-t')
|
|
args.insert(2, str(duration))
|
|
|
|
try:
|
|
audio_data = b64decode(base64_str)
|
|
res = subprocess.run(args, input=audio_data, capture_output=True, check=True)
|
|
|
|
audio = torch.frombuffer(res.stdout, dtype=torch.float32)
|
|
|
|
match = re.search(r', (\d+) Hz, (\w+),', res.stderr.decode(*ENCODE_ARGS))
|
|
if match:
|
|
sample_rate = int(match.group(1))
|
|
channels = {"mono": 1, "stereo": 2}[match.group(2)]
|
|
else:
|
|
sample_rate = 44100
|
|
channels = 2
|
|
|
|
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 decode audio: {e}") |