gabe-init
fdab0ac836
- Google Image Search integration for ComfyUI - Secure credential management via config.json - Returns images as ComfyUI tensors - MIT License
113 lines
4.2 KiB
Python
113 lines
4.2 KiB
Python
from googleapiclient.discovery import build
|
|
import requests
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import torch
|
|
import numpy as np
|
|
import json
|
|
import os
|
|
|
|
class GoogleImageSearchNode:
|
|
def __init__(self):
|
|
self.config = self.load_config()
|
|
|
|
def load_config(self):
|
|
config_path = os.path.join(os.path.dirname(__file__), 'config.json')
|
|
if os.path.exists(config_path):
|
|
try:
|
|
with open(config_path, 'r') as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
print(f"Error loading config.json: {e}")
|
|
return {}
|
|
else:
|
|
print("config.json not found. Please copy config.json.example to config.json and add your credentials.")
|
|
return {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"search_query": ("STRING", {"default": "cat"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "search_image"
|
|
CATEGORY = "Custom Nodes/Google"
|
|
|
|
def search_image(self, search_query):
|
|
api_key = self.config.get('api_key', '')
|
|
search_engine_id = self.config.get('search_engine_id', '')
|
|
|
|
if not api_key or not search_engine_id:
|
|
print("Error: API key and Search Engine ID are required. Please configure config.json")
|
|
# Return a red error image
|
|
error_tensor = torch.zeros(1, 64, 64, 3)
|
|
error_tensor[0, :, :, 0] = 1
|
|
return (error_tensor,)
|
|
|
|
try:
|
|
print(f"Starting search for: {search_query}")
|
|
|
|
# Initialize the Custom Search API service
|
|
service = build("customsearch", "v1", developerKey=api_key)
|
|
|
|
# Perform the search with the provided Search Engine ID
|
|
result = service.cse().list(
|
|
q=search_query,
|
|
cx=search_engine_id,
|
|
searchType='image',
|
|
num=1,
|
|
safe='off' # Disable SafeSearch to get more results
|
|
).execute()
|
|
|
|
print(f"Search API response received")
|
|
|
|
# Get the first image URL
|
|
if 'items' in result and len(result['items']) > 0:
|
|
image_url = result['items'][0]['link']
|
|
print(f"Found image URL: {image_url}")
|
|
|
|
# Download the image
|
|
response = requests.get(image_url, timeout=10)
|
|
response.raise_for_status() # Raise an exception for bad status codes
|
|
|
|
print(f"Image downloaded successfully")
|
|
|
|
# Convert to PIL Image
|
|
img = Image.open(BytesIO(response.content))
|
|
|
|
# Convert to RGB if necessary
|
|
if img.mode != 'RGB':
|
|
img = img.convert('RGB')
|
|
|
|
print(f"Image converted to RGB format: {img.size}")
|
|
|
|
# Convert PIL image to numpy array
|
|
img_array = np.array(img).astype(np.float32) / 255.0
|
|
|
|
# Convert numpy array to tensor with correct shape [B, H, W, C]
|
|
img_tensor = torch.from_numpy(img_array)
|
|
img_tensor = img_tensor.unsqueeze(0) # Add batch dimension
|
|
|
|
print(f"Successfully created tensor with shape: {img_tensor.shape}")
|
|
return (img_tensor,)
|
|
else:
|
|
print(f"No images found in API response: {result}")
|
|
raise Exception("No images found in search results")
|
|
|
|
except Exception as e:
|
|
print(f"Detailed error in Google Image Search: {type(e).__name__}: {str(e)}")
|
|
# Return a small red tensor to indicate error
|
|
error_tensor = torch.zeros(1, 64, 64, 3) # [B, H, W, C] format
|
|
error_tensor[0, :, :, 0] = 1 # Red channel
|
|
return (error_tensor,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GoogleImageSearchNode": GoogleImageSearchNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GoogleImageSearchNode": "Google Image Search"
|
|
} |