70 lines
2.6 KiB
Python
70 lines
2.6 KiB
Python
import base64
|
|
import io
|
|
from PIL import Image
|
|
from typing import Optional, Union
|
|
import torch
|
|
|
|
class ImagePreprocessor:
|
|
"""
|
|
Custom node for preprocessing images before passing them to LLMs.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {"default": None}),
|
|
},
|
|
"optional": {
|
|
"format": ("STRING", {"default": "PNG"}),
|
|
"quality": (["High", "Medium", "Low"], {"default": "High"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("processed_image",)
|
|
FUNCTION = "preprocess"
|
|
CATEGORY = "DeepSeek_Toolkit"
|
|
|
|
def preprocess(self, image: Optional[Union[str, Image.Image]], format: str = "PNG", quality: str = "High"):
|
|
quality_map = {"High": 95, "Medium": 75, "Low": 50}
|
|
quality_str = quality # Ensure quality is treated as a string
|
|
quality = quality_map.get(quality_str, 95) # Default to High if invalid value
|
|
print(f"Selected quality: {quality}") # Debugging line
|
|
|
|
if image is None:
|
|
raise ValueError("Image input cannot be None")
|
|
|
|
# Convert image to PIL object if it's a tensor
|
|
if isinstance(image, torch.Tensor):
|
|
image = image.squeeze(0).cpu().numpy() # [H, W, C]
|
|
if image.shape[2] == 1: # Grayscale image
|
|
image = image.squeeze(-1)
|
|
image = (image * 255).astype('uint8')
|
|
image = Image.fromarray(image)
|
|
|
|
# Ensure the image is a PIL object
|
|
elif not isinstance(image, Image.Image):
|
|
raise ValueError("Unsupported image type. Expected torch.Tensor or PIL.Image.")
|
|
|
|
# Resize image based on quality
|
|
size_map = {"High": 1024, "Medium": 768, "Low": 512}
|
|
max_size = size_map.get(quality_str, 1024)
|
|
print(f"Resizing image to max dimension: {max_size}")
|
|
image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
|
|
|
|
# Convert PIL image to base64 string
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format=format, quality=quality)
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
image_url = f"data:image/{format.lower()};base64,{img_str}"
|
|
print(f"Image size with quality {quality}: {buffered.tell()} bytes")
|
|
print(f"Generated image URL: {image_url[:50]}...") # 打印生成的图像 URL 前 50 个字符
|
|
|
|
return (image_url,)
|
|
|
|
|
|
# Register the node with ComfyUI
|
|
NODE_CLASS_MAPPINGS = {"ImagePreprocessor": ImagePreprocessor}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"ImagePreprocessor": "Image Preprocessor"}
|