76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
import torch
|
|
from PIL import Image
|
|
import base64
|
|
import io
|
|
import importlib.util
|
|
import sys
|
|
import numpy as np # Added numpy import
|
|
|
|
def ensure_package(package_name, version=None):
|
|
try:
|
|
if version:
|
|
importlib.import_module(f"{package_name}=={version}")
|
|
else:
|
|
importlib.import_module(package_name)
|
|
except ImportError:
|
|
import subprocess
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", f"{package_name}{f'=={version}' if version else ''}"])
|
|
|
|
def tensor2pil(image: torch.Tensor) -> Image.Image:
|
|
# Assuming input is BHWC (Batch, Height, Width, Channels)
|
|
# Squeeze batch dim, move to CPU, convert to numpy
|
|
img_np = image.squeeze(0).cpu().numpy() # Shape (H, W, C)
|
|
# Convert float [0,1] to uint8 [0,255] if necessary
|
|
if img_np.dtype == np.float32 or img_np.dtype == np.float64:
|
|
img_np = np.clip(img_np * 255, 0, 255).astype(np.uint8)
|
|
elif img_np.dtype != np.uint8:
|
|
# Raise error for unexpected types instead of relying on .byte()
|
|
raise TypeError(f"tensor2pil cannot handle numpy dtype: {img_np.dtype}")
|
|
|
|
# Handle grayscale images (H, W, 1) -> (H, W) for PIL
|
|
if img_np.ndim == 3 and img_np.shape[2] == 1:
|
|
img_np = np.squeeze(img_np, axis=2)
|
|
|
|
return Image.fromarray(img_np)
|
|
|
|
def pil2base64(image: Image.Image) -> str:
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
|
|
# --- Utilities for Dynamic/Flexible Nodes ---
|
|
|
|
class AnyType(str):
|
|
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
|
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
class FlexibleOptionalInputType(dict):
|
|
"""A special class to make flexible nodes that pass data to our python handlers.
|
|
|
|
Enables both flexible/dynamic input types (like for Any Switch) or a dynamic number of inputs
|
|
(like for Any Switch, Context Switch, Context Merge, Power Lora Loader, etc).
|
|
|
|
Note, for ComfyUI, all that's needed is the `__contains__` override below, which tells ComfyUI
|
|
that our node will handle the input, regardless of what it is.
|
|
|
|
However, with https://github.com/comfyanonymous/ComfyUI/pull/2666 a large change would occur
|
|
requiring more details on the input itself. There, we need to return a list/tuple where the first
|
|
item is the type. This can be a real type, or use the AnyType for additional flexibility.
|
|
|
|
This should be forwards compatible unless more changes occur in the PR.
|
|
"""
|
|
def __init__(self, type):
|
|
self.type = type
|
|
|
|
def __getitem__(self, key):
|
|
return (self.type, )
|
|
|
|
def __contains__(self, key):
|
|
return True
|
|
|
|
# Global instance for convenience
|
|
any_type = AnyType("*")
|