26 lines
850 B
Python
26 lines
850 B
Python
"""
|
|
Image and tensor conversion utilities for ComfyUI-Distributed.
|
|
"""
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
def tensor_to_pil(img_tensor, batch_index=0):
|
|
"""Takes a batch of images in tensor form [B, H, W, C] and returns an RGB PIL Image."""
|
|
return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
|
|
|
|
def pil_to_tensor(image):
|
|
"""Takes a PIL image and returns a tensor of shape [1, H, W, C]."""
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image).unsqueeze(0)
|
|
if len(image.shape) == 3: # If grayscale, add channel dimension
|
|
image = image.unsqueeze(-1)
|
|
return image
|
|
|
|
def ensure_contiguous(tensor):
|
|
"""Ensure tensor is contiguous in memory."""
|
|
if not tensor.is_contiguous():
|
|
return tensor.contiguous()
|
|
return tensor
|
|
|