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
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import os
|
|
import json
|
|
import configparser
|
|
from pathlib import Path
|
|
|
|
class GoogleAIPluginConfig:
|
|
def __init__(self):
|
|
self.config_dir = Path.home() / ".comfyui" / "google_ai_studio"
|
|
self.config_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
self.config_file = self.config_dir / "config.ini"
|
|
self.secrets_file = self.config_dir / "secrets.json"
|
|
|
|
self.config = configparser.ConfigParser()
|
|
self._load_config()
|
|
|
|
def _load_config(self):
|
|
"""Load configuration with defaults"""
|
|
defaults = {
|
|
'api': {
|
|
'timeout': '120',
|
|
'max_retries': '3',
|
|
'default_model': 'gemini-2.5-flash',
|
|
'rate_limit_buffer': '0.1'
|
|
},
|
|
'generation': {
|
|
'default_temperature': '0.7',
|
|
'default_aspect_ratio': '1:1',
|
|
'cache_results': 'true'
|
|
},
|
|
'optimization': {
|
|
'default_enhancement': 'general',
|
|
'max_iterations': '3',
|
|
'quality_threshold': '0.7'
|
|
}
|
|
}
|
|
|
|
self.config.read_dict(defaults)
|
|
|
|
if self.config_file.exists():
|
|
self.config.read(self.config_file)
|
|
|
|
def get_api_key(self, fallback_env_vars=None):
|
|
"""Get API key from multiple sources"""
|
|
if fallback_env_vars is None:
|
|
fallback_env_vars = ["GOOGLE_API_KEY", "GEMINI_API_KEY"]
|
|
|
|
# Try environment variables first
|
|
for env_var in fallback_env_vars:
|
|
api_key = os.getenv(env_var)
|
|
if api_key:
|
|
return api_key
|
|
|
|
# Try secrets file
|
|
if self.secrets_file.exists():
|
|
try:
|
|
with open(self.secrets_file, 'r') as f:
|
|
secrets = json.load(f)
|
|
return secrets.get("google_api_key")
|
|
except (json.JSONDecodeError, IOError):
|
|
pass
|
|
|
|
return None
|
|
|
|
def save_api_key(self, api_key):
|
|
"""Securely save API key"""
|
|
secrets = {}
|
|
if self.secrets_file.exists():
|
|
try:
|
|
with open(self.secrets_file, 'r') as f:
|
|
secrets = json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
secrets = {}
|
|
|
|
secrets["google_api_key"] = api_key
|
|
|
|
with open(self.secrets_file, 'w') as f:
|
|
json.dump(secrets, f)
|
|
|
|
# Set restrictive permissions on Unix-like systems
|
|
try:
|
|
os.chmod(self.secrets_file, 0o600)
|
|
except (OSError, AttributeError):
|
|
# Windows or permission issues
|
|
pass
|
|
|
|
# Global configuration instance
|
|
plugin_config = GoogleAIPluginConfig() |