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
222 lines
8.6 KiB
Python
222 lines
8.6 KiB
Python
"""
|
|
Google AI Studio Integration for ComfyUI
|
|
|
|
A comprehensive plugin that integrates Google AI Studio's powerful AI models
|
|
including Gemini 2.5 Flash and Pro models for content generation, image creation,
|
|
video planning, and advanced prompt optimization.
|
|
|
|
Features:
|
|
- Image generation with AI-powered descriptions
|
|
- Video storyboard and planning tools
|
|
- Advanced prompt optimization and analysis
|
|
- Multiple model support (Gemini 2.5 Flash/Pro)
|
|
- Async processing for better performance
|
|
- Rate limiting and error handling
|
|
"""
|
|
|
|
import logging
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# Set up logging with proper format
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='[%(name)s] %(levelname)s: %(message)s'
|
|
)
|
|
logger = logging.getLogger("ComfyUI-GoogleAIStudio")
|
|
|
|
# Add the plugin directory to Python path
|
|
plugin_dir = Path(__file__).parent
|
|
sys.path.insert(0, str(plugin_dir))
|
|
|
|
def check_environment():
|
|
"""Check plugin environment and dependencies"""
|
|
issues = []
|
|
|
|
# Check API key
|
|
api_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
|
|
if not api_key:
|
|
issues.append("No Google AI Studio API key found in environment")
|
|
logger.warning("Set GOOGLE_API_KEY environment variable for full functionality")
|
|
|
|
# Check dependencies
|
|
required_deps = ['torch', 'PIL', 'aiohttp', 'numpy']
|
|
for dep in required_deps:
|
|
try:
|
|
__import__(dep)
|
|
except ImportError:
|
|
issues.append(f"Missing dependency: {dep}")
|
|
|
|
return len(issues) == 0, issues
|
|
|
|
try:
|
|
# Check environment first
|
|
env_ok, env_issues = check_environment()
|
|
if env_issues:
|
|
for issue in env_issues:
|
|
logger.warning(issue)
|
|
|
|
# Import all node classes (using the approach that works)
|
|
# Try relative imports first, fallback to absolute if needed
|
|
try:
|
|
from .nodes.image_generation import GoogleAIImageGenerator, GoogleAITextToImage
|
|
from .nodes.video_generation import GoogleAIVideoGenerator, GoogleAIVideoStoryboard
|
|
from .nodes.prompt_optimizer import PromptOptimizerNode, PromptAnalyzerNode
|
|
from .nodes.config_node import GoogleAIConfigNode, GoogleAIKeyManagerNode, GoogleAIConnectionTesterNode
|
|
logger.info("Using relative imports (recommended)")
|
|
except ImportError:
|
|
# Fallback to absolute imports if relative fails
|
|
from nodes.image_generation import GoogleAIImageGenerator, GoogleAITextToImage
|
|
from nodes.video_generation import GoogleAIVideoGenerator, GoogleAIVideoStoryboard
|
|
from nodes.prompt_optimizer import PromptOptimizerNode, PromptAnalyzerNode
|
|
from nodes.config_node import GoogleAIConfigNode, GoogleAIKeyManagerNode, GoogleAIConnectionTesterNode
|
|
logger.info("Using absolute imports (fallback)")
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Node class mappings for ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
# Configuration nodes (most important - put first)
|
|
"GoogleAIConfig": GoogleAIConfigNode,
|
|
"GoogleAIKeyManager": GoogleAIKeyManagerNode,
|
|
"GoogleAIConnectionTester": GoogleAIConnectionTesterNode,
|
|
|
|
# Generation nodes
|
|
"GoogleAIImageGenerator": GoogleAIImageGenerator,
|
|
"GoogleAITextToImage": GoogleAITextToImage,
|
|
"GoogleAIVideoGenerator": GoogleAIVideoGenerator,
|
|
"GoogleAIVideoStoryboard": GoogleAIVideoStoryboard,
|
|
|
|
# Utility nodes
|
|
"PromptOptimizer": PromptOptimizerNode,
|
|
"PromptAnalyzer": PromptAnalyzerNode,
|
|
}
|
|
|
|
# Display names for the ComfyUI interface
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
# Configuration nodes (显眼的配置节点)
|
|
"GoogleAIConfig": "[CONFIG] Google AI Settings",
|
|
"GoogleAIKeyManager": "[KEY] API Key Manager",
|
|
"GoogleAIConnectionTester": "[TEST] Connection Tester",
|
|
|
|
# Generation nodes
|
|
"GoogleAIImageGenerator": "[AI] Google AI Image Generator",
|
|
"GoogleAITextToImage": "[ART] Google AI Text to Image",
|
|
"GoogleAIVideoGenerator": "[VIDEO] Google AI Video Generator",
|
|
"GoogleAIVideoStoryboard": "[STORY] Google AI Video Storyboard",
|
|
|
|
# Utility nodes
|
|
"PromptOptimizer": "[OPT] Prompt Optimizer",
|
|
"PromptAnalyzer": "[ANALYZE] Prompt Analyzer",
|
|
}
|
|
|
|
# Web directory for custom JavaScript/CSS
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
logger.info("Google AI Studio plugin loaded successfully!")
|
|
logger.info(f"Loaded {len(NODE_CLASS_MAPPINGS)} node types")
|
|
|
|
# Validate all registered nodes
|
|
validation_errors = []
|
|
for node_id, node_class in NODE_CLASS_MAPPINGS.items():
|
|
try:
|
|
# Check required methods
|
|
if not hasattr(node_class, 'INPUT_TYPES'):
|
|
validation_errors.append(f"{node_id}: Missing INPUT_TYPES method")
|
|
if not hasattr(node_class, 'RETURN_TYPES'):
|
|
validation_errors.append(f"{node_id}: Missing RETURN_TYPES attribute")
|
|
if not hasattr(node_class, 'FUNCTION'):
|
|
validation_errors.append(f"{node_id}: Missing FUNCTION attribute")
|
|
|
|
# Check if function exists
|
|
func_name = getattr(node_class, 'FUNCTION', None)
|
|
if func_name and not hasattr(node_class, func_name):
|
|
validation_errors.append(f"{node_id}: Function '{func_name}' not found")
|
|
|
|
except Exception as e:
|
|
validation_errors.append(f"{node_id}: Validation error - {e}")
|
|
|
|
if validation_errors:
|
|
logger.error("Node validation errors found:")
|
|
for error in validation_errors:
|
|
logger.error(f" - {error}")
|
|
else:
|
|
logger.info("All nodes passed validation")
|
|
|
|
# Print available nodes for debugging
|
|
logger.info("Registered nodes:")
|
|
for node_class, display_name in NODE_DISPLAY_NAME_MAPPINGS.items():
|
|
logger.info(f" - {node_class}: {display_name}")
|
|
|
|
# Environment status summary
|
|
if env_ok:
|
|
logger.info("Environment: Ready ✓")
|
|
else:
|
|
logger.warning(f"Environment: {len(env_issues)} issues found")
|
|
|
|
except ImportError as e:
|
|
logger = logging.getLogger(__name__)
|
|
logger.error(f"Failed to import Google AI Studio nodes: {e}")
|
|
logger.error("This might be due to missing dependencies. Please run: pip install -r requirements.txt")
|
|
|
|
# Provide empty mappings to prevent ComfyUI errors
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
except Exception as e:
|
|
logger = logging.getLogger(__name__)
|
|
logger.error(f"Unexpected error loading Google AI Studio plugin: {e}")
|
|
|
|
# Provide empty mappings to prevent ComfyUI errors
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
# Export for ComfyUI (必须导出这些变量)
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
|
|
|
|
# Plugin metadata
|
|
__version__ = "1.0.0"
|
|
__author__ = "FloatFloat-chan (浮浮酱)"
|
|
__description__ = "Google AI Studio integration for ComfyUI with advanced prompt optimization"
|
|
__license__ = "MIT"
|
|
|
|
# Print startup banner (Windows compatible)
|
|
def print_startup_banner():
|
|
try:
|
|
banner = f"""
|
|
╔══════════════════════════════════════════════════╗
|
|
║ Google AI Studio Plugin for ComfyUI ║
|
|
║ Version {__version__} ║
|
|
║ Created by FloatFloat-chan ║
|
|
║ ║
|
|
║ [AI] Image Generation [VIDEO] Video Planning ║
|
|
║ [OPT] Prompt Optimizer [CONFIG] Smart Setup ║
|
|
║ ║
|
|
║ Successfully Loaded & Ready! ║
|
|
╚══════════════════════════════════════════════════╝
|
|
"""
|
|
print(banner)
|
|
except UnicodeEncodeError:
|
|
# Fallback for systems with encoding issues
|
|
print("=" * 50)
|
|
print(f" Google AI Studio Plugin v{__version__} - LOADED ")
|
|
print("=" * 50)
|
|
|
|
# Show startup banner if environment is ready
|
|
if 'NODE_CLASS_MAPPINGS' in globals() and len(NODE_CLASS_MAPPINGS) > 0:
|
|
print_startup_banner()
|
|
|
|
# Additional startup info
|
|
logger.info(f"Plugin Info:")
|
|
logger.info(f" Version: {__version__}")
|
|
logger.info(f" Author: {__author__}")
|
|
logger.info(f" Description: {__description__}")
|
|
logger.info(f" Available at: https://github.com/comfyui-extensions/google-ai-studio")
|
|
logger.info(f" Documentation: See README.md and QUICK_START.md")
|
|
else:
|
|
logger.error("Plugin failed to load properly - no nodes registered") |