mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-02 17:00:12 +08:00
26 lines
788 B
Python
26 lines
788 B
Python
from io import BytesIO
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
|
|
def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor:
|
|
"""Converts image data from BytesIO to a torch.Tensor.
|
|
|
|
Args:
|
|
image_bytesio: BytesIO object containing the image data.
|
|
mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA").
|
|
|
|
Returns:
|
|
A torch.Tensor representing the image (1, H, W, C).
|
|
|
|
Raises:
|
|
PIL.UnidentifiedImageError: If the image data cannot be identified.
|
|
ValueError: If the specified mode is invalid.
|
|
"""
|
|
image = Image.open(image_bytesio)
|
|
image = image.convert(mode)
|
|
image_array = np.array(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(image_array).unsqueeze(0)
|