Administrator
86386baa1a
Features: - 🤖 AI image generation with Gemini models - 🎬 Video planning and storyboard creation - ✨ Advanced prompt optimization and analysis - 🔧 Smart API key configuration and management - 📊 Multiple model support (Gemini 2.5 Flash/Pro) - 🛡️ Robust error handling and rate limiting - 📖 Comprehensive documentation and quick start guide Created by FloatFloat-chan with ❤️ for the ComfyUI community
110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
import torch
|
|
import numpy as np
|
|
import base64
|
|
import io
|
|
from PIL import Image
|
|
from typing import Union, Tuple
|
|
|
|
def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
|
"""Convert ComfyUI tensor to PIL Image"""
|
|
# ComfyUI images are [B, H, W, C] or [H, W, C] tensors
|
|
if len(tensor.shape) == 4:
|
|
# Take first image from batch
|
|
np_image = tensor[0].cpu().numpy()
|
|
else:
|
|
np_image = tensor.cpu().numpy()
|
|
|
|
# Convert from float [0,1] to uint8 [0,255]
|
|
if np_image.dtype == np.float32 or np_image.dtype == np.float64:
|
|
np_image = (np_image * 255).astype(np.uint8)
|
|
|
|
# Ensure we have 3 channels (RGB)
|
|
if np_image.shape[-1] == 1:
|
|
np_image = np.repeat(np_image, 3, axis=-1)
|
|
elif np_image.shape[-1] == 4:
|
|
# Convert RGBA to RGB
|
|
np_image = np_image[:, :, :3]
|
|
|
|
return Image.fromarray(np_image)
|
|
|
|
def pil_to_tensor(pil_image: Image.Image) -> torch.Tensor:
|
|
"""Convert PIL Image to ComfyUI tensor format"""
|
|
# Ensure RGB format
|
|
if pil_image.mode != "RGB":
|
|
pil_image = pil_image.convert("RGB")
|
|
|
|
# Convert to numpy array
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
|
|
# Convert to tensor with batch dimension: [1, H, W, C]
|
|
tensor = torch.from_numpy(np_image).unsqueeze(0)
|
|
|
|
return tensor
|
|
|
|
def image_to_base64(image: Union[torch.Tensor, Image.Image], format: str = "PNG") -> str:
|
|
"""Convert image to base64 string"""
|
|
if isinstance(image, torch.Tensor):
|
|
pil_image = tensor_to_pil(image)
|
|
else:
|
|
pil_image = image
|
|
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format=format)
|
|
image_bytes = buffer.getvalue()
|
|
|
|
return base64.b64encode(image_bytes).decode('utf-8')
|
|
|
|
def base64_to_image(b64_string: str) -> torch.Tensor:
|
|
"""Convert base64 string to ComfyUI tensor"""
|
|
try:
|
|
image_bytes = base64.b64decode(b64_string)
|
|
pil_image = Image.open(io.BytesIO(image_bytes))
|
|
return pil_to_tensor(pil_image)
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to decode base64 image: {e}")
|
|
|
|
def resize_image_aspect_ratio(image: Image.Image, target_ratio: str, max_size: int = 1024) -> Image.Image:
|
|
"""Resize image to target aspect ratio while maintaining quality"""
|
|
# Parse aspect ratio
|
|
if target_ratio == "1:1":
|
|
target_width, target_height = 1, 1
|
|
elif target_ratio == "16:9":
|
|
target_width, target_height = 16, 9
|
|
elif target_ratio == "9:16":
|
|
target_width, target_height = 9, 16
|
|
elif target_ratio == "4:3":
|
|
target_width, target_height = 4, 3
|
|
elif target_ratio == "3:4":
|
|
target_width, target_height = 3, 4
|
|
else:
|
|
return image # Return original if unknown ratio
|
|
|
|
# Calculate target dimensions
|
|
target_aspect = target_width / target_height
|
|
current_width, current_height = image.size
|
|
current_aspect = current_width / current_height
|
|
|
|
if current_aspect > target_aspect:
|
|
# Image is wider than target, crop width
|
|
new_width = int(current_height * target_aspect)
|
|
left = (current_width - new_width) // 2
|
|
image = image.crop((left, 0, left + new_width, current_height))
|
|
elif current_aspect < target_aspect:
|
|
# Image is taller than target, crop height
|
|
new_height = int(current_width / target_aspect)
|
|
top = (current_height - new_height) // 2
|
|
image = image.crop((0, top, current_width, top + new_height))
|
|
|
|
# Resize to max_size while maintaining aspect ratio
|
|
current_width, current_height = image.size
|
|
if max(current_width, current_height) > max_size:
|
|
if current_width > current_height:
|
|
new_width = max_size
|
|
new_height = int(max_size / target_aspect)
|
|
else:
|
|
new_height = max_size
|
|
new_width = int(max_size * target_aspect)
|
|
|
|
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
|
|
return image |