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
238 lines
9.7 KiB
Python
238 lines
9.7 KiB
Python
import json
|
|
import asyncio
|
|
from typing import List, Dict, Optional
|
|
|
|
from .google_ai_client import GoogleAIStudioClient
|
|
|
|
class PromptOptimizer:
|
|
"""Advanced prompt optimization using multiple techniques"""
|
|
|
|
def __init__(self, api_key: Optional[str] = None):
|
|
self.api_key = api_key
|
|
self.enhancement_templates = {
|
|
"photography": "Professional photography style: {prompt}, shot with high-end camera, perfect lighting, composition following rule of thirds",
|
|
"artistic": "Artistic masterpiece: {prompt}, in the style of renowned artists, museum quality, detailed brushwork",
|
|
"cinematic": "Cinematic scene: {prompt}, dramatic lighting, film photography, depth of field, professional color grading",
|
|
"technical": "Technical illustration: {prompt}, precise details, technical accuracy, clean lines, professional diagram style",
|
|
"fantasy": "Fantasy artwork: {prompt}, magical atmosphere, mystical lighting, detailed fantasy elements",
|
|
"sci-fi": "Science fiction: {prompt}, futuristic design, advanced technology, dramatic sci-fi atmosphere"
|
|
}
|
|
|
|
async def optimize_prompt(self, original_prompt: str,
|
|
target_model: str = "gemini-2.5-flash",
|
|
enhancement_type: str = "general",
|
|
iterations: int = 1) -> Dict[str, any]:
|
|
"""Optimize prompt using multiple enhancement techniques"""
|
|
|
|
optimization_results = {
|
|
"original_prompt": original_prompt,
|
|
"optimized_prompts": [],
|
|
"enhancement_scores": [],
|
|
"final_prompt": original_prompt,
|
|
"enhancement_details": []
|
|
}
|
|
|
|
current_prompt = original_prompt
|
|
|
|
for i in range(iterations):
|
|
# Apply template-based enhancement
|
|
if enhancement_type in self.enhancement_templates:
|
|
current_prompt = self.enhancement_templates[enhancement_type].format(prompt=current_prompt)
|
|
|
|
# Use LLM for advanced optimization
|
|
enhanced_prompt = await self.llm_enhance_prompt(current_prompt, target_model, enhancement_type)
|
|
|
|
# Score the enhancement
|
|
score = await self.score_prompt_quality(enhanced_prompt, target_model)
|
|
|
|
optimization_results["optimized_prompts"].append(enhanced_prompt)
|
|
optimization_results["enhancement_scores"].append(score)
|
|
optimization_results["enhancement_details"].append({
|
|
"iteration": i + 1,
|
|
"enhancement_type": enhancement_type,
|
|
"score": score,
|
|
"prompt_length": len(enhanced_prompt)
|
|
})
|
|
|
|
current_prompt = enhanced_prompt
|
|
|
|
# Select best prompt based on scores
|
|
if optimization_results["enhancement_scores"]:
|
|
best_idx = max(range(len(optimization_results["enhancement_scores"])),
|
|
key=lambda i: optimization_results["enhancement_scores"][i])
|
|
optimization_results["final_prompt"] = optimization_results["optimized_prompts"][best_idx]
|
|
optimization_results["best_iteration"] = best_idx + 1
|
|
|
|
return optimization_results
|
|
|
|
async def llm_enhance_prompt(self, prompt: str, target_model: str, enhancement_type: str = "general") -> str:
|
|
"""Use LLM to enhance prompt for specific model"""
|
|
|
|
enhancement_instructions = {
|
|
"general": "Improve this prompt for AI image generation by adding specific visual details and technical terms",
|
|
"photography": "Enhance this prompt with professional photography terminology and technical specifications",
|
|
"artistic": "Refine this prompt with artistic terminology and style references",
|
|
"cinematic": "Improve this prompt with cinematic and film-making terminology",
|
|
"technical": "Enhance this prompt with precise technical language and specifications"
|
|
}
|
|
|
|
enhancement_prompt = f"""
|
|
You are an expert prompt engineer for AI content generation.
|
|
Task: {enhancement_instructions.get(enhancement_type, enhancement_instructions["general"])}
|
|
|
|
Original prompt: "{prompt}"
|
|
Target model: {target_model}
|
|
Enhancement style: {enhancement_type}
|
|
|
|
Requirements:
|
|
1. Maintain the core concept and intent
|
|
2. Add specific visual details that improve generation quality
|
|
3. Include relevant technical parameters when appropriate
|
|
4. Remove redundant or conflicting elements
|
|
5. Optimize for the specific capabilities of {target_model}
|
|
6. Keep the enhanced prompt concise but descriptive
|
|
|
|
Return only the optimized prompt, no explanations or additional text.
|
|
"""
|
|
|
|
async with GoogleAIStudioClient(self.api_key) as client:
|
|
try:
|
|
response = await client.generate_content(
|
|
prompt=enhancement_prompt,
|
|
model="gemini-2.5-pro",
|
|
temperature=0.3,
|
|
max_tokens=500
|
|
)
|
|
|
|
# Extract enhanced prompt from response
|
|
candidates = response.get("candidates", [])
|
|
if candidates:
|
|
content = candidates[0].get("content", {})
|
|
parts = content.get("parts", [])
|
|
if parts:
|
|
enhanced = parts[0].get("text", prompt).strip()
|
|
# Clean up any markdown formatting
|
|
enhanced = enhanced.replace("**", "").replace("*", "")
|
|
return enhanced
|
|
|
|
except Exception as e:
|
|
print(f"Enhancement failed: {e}")
|
|
|
|
return prompt
|
|
|
|
async def score_prompt_quality(self, prompt: str, target_model: str) -> float:
|
|
"""Score prompt quality using multiple metrics"""
|
|
|
|
scoring_prompt = f"""
|
|
Rate this AI content generation prompt on a scale of 0.0 to 1.0 based on:
|
|
|
|
1. Clarity and specificity (30% weight) - How clear and specific are the instructions?
|
|
2. Appropriate detail level (30% weight) - Is the detail level appropriate for generation?
|
|
3. Technical compatibility (20% weight) - How well does it work with {target_model}?
|
|
4. Creative potential (20% weight) - How likely is it to produce interesting results?
|
|
|
|
Prompt to evaluate: "{prompt}"
|
|
|
|
Consider:
|
|
- Specificity of visual elements
|
|
- Technical terminology usage
|
|
- Prompt length and structure
|
|
- Potential for quality output
|
|
|
|
Return only a single decimal number between 0.0 and 1.0 (e.g., 0.75).
|
|
No explanations, just the score.
|
|
"""
|
|
|
|
try:
|
|
async with GoogleAIStudioClient(self.api_key) as client:
|
|
response = await client.generate_content(
|
|
prompt=scoring_prompt,
|
|
model="gemini-2.5-flash",
|
|
temperature=0.1,
|
|
max_tokens=10
|
|
)
|
|
|
|
# Extract score
|
|
candidates = response.get("candidates", [])
|
|
if candidates:
|
|
content = candidates[0].get("content", {})
|
|
parts = content.get("parts", [])
|
|
if parts:
|
|
score_text = parts[0].get("text", "0.5").strip()
|
|
try:
|
|
# Extract number from text
|
|
import re
|
|
numbers = re.findall(r'\d+\.?\d*', score_text)
|
|
if numbers:
|
|
score = float(numbers[0])
|
|
return min(max(score, 0.0), 1.0) # Clamp between 0 and 1
|
|
except ValueError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# Default fallback score
|
|
return 0.5
|
|
|
|
async def analyze_prompt_structure(self, prompt: str) -> Dict[str, any]:
|
|
"""Analyze prompt structure and provide suggestions"""
|
|
|
|
analysis_prompt = f"""
|
|
Analyze this AI generation prompt and provide structured feedback:
|
|
|
|
Prompt: "{prompt}"
|
|
|
|
Analyze:
|
|
1. Word count and length appropriateness
|
|
2. Specificity level (vague vs specific)
|
|
3. Technical terminology usage
|
|
4. Style consistency
|
|
5. Potential improvements
|
|
|
|
Provide analysis in JSON format:
|
|
{{
|
|
"word_count": number,
|
|
"specificity_level": "low/medium/high",
|
|
"technical_terms": ["list", "of", "terms"],
|
|
"style_consistency": "consistent/inconsistent",
|
|
"suggested_improvements": ["suggestion1", "suggestion2"],
|
|
"overall_assessment": "brief summary"
|
|
}}
|
|
"""
|
|
|
|
try:
|
|
async with GoogleAIStudioClient(self.api_key) as client:
|
|
response = await client.generate_content(
|
|
prompt=analysis_prompt,
|
|
model="gemini-2.5-pro",
|
|
temperature=0.2,
|
|
max_tokens=1000
|
|
)
|
|
|
|
candidates = response.get("candidates", [])
|
|
if candidates:
|
|
content = candidates[0].get("content", {})
|
|
parts = content.get("parts", [])
|
|
if parts:
|
|
analysis_text = parts[0].get("text", "{}").strip()
|
|
try:
|
|
# Extract JSON from response
|
|
import re
|
|
json_match = re.search(r'\{.*\}', analysis_text, re.DOTALL)
|
|
if json_match:
|
|
return json.loads(json_match.group())
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
except Exception:
|
|
pass
|
|
|
|
# Default analysis
|
|
return {
|
|
"word_count": len(prompt.split()),
|
|
"specificity_level": "medium",
|
|
"technical_terms": [],
|
|
"style_consistency": "unknown",
|
|
"suggested_improvements": ["Add more specific details"],
|
|
"overall_assessment": "Analysis unavailable"
|
|
} |