334 lines
16 KiB
Python
334 lines
16 KiB
Python
"""
|
|
Core Prompt Orchestrator
|
|
Lightweight orchestrator that coordinates all modifier types
|
|
"""
|
|
|
|
from typing import List, Dict, Optional, Callable
|
|
from .modifiers import ModifierCategory, LightingModifier, SeasonalModifier, LocationModifier
|
|
|
|
|
|
class PromptOrchestrator:
|
|
"""
|
|
Main orchestration class for image prompt variation automation.
|
|
Handles text input manipulation and modifier application.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the orchestrator - load all modifiers eagerly"""
|
|
# Import all modifiers at initialization (not lazy)
|
|
from modifiers.lighting_modifiers import get_lighting_modifiers
|
|
from modifiers.seasonal_modifiers import get_seasonal_modifiers
|
|
from modifiers.location_modifiers import get_location_modifiers
|
|
from modifiers.weather_modifiers import get_weather_modifiers
|
|
from modifiers.mood_modifiers import get_mood_modifiers
|
|
from modifiers.camera_modifiers import get_camera_modifiers
|
|
from modifiers.color_grading_modifiers import get_color_grading_modifiers
|
|
from modifiers.action_modifiers import get_action_modifiers
|
|
from modifiers.era_modifiers import get_era_modifiers
|
|
from modifiers.detail_modifiers import get_detail_modifiers
|
|
from modifiers.composition_modifiers import get_composition_modifiers
|
|
|
|
# Store as instance variables (not properties)
|
|
self.modifiers = {m.id: m for m in get_lighting_modifiers()}
|
|
self.seasonal_modifiers = {m.id: m for m in get_seasonal_modifiers()}
|
|
self.location_modifiers = {m.id: m for m in get_location_modifiers()}
|
|
self.weather_modifiers = {m.id: m for m in get_weather_modifiers()}
|
|
self.mood_modifiers = {m.id: m for m in get_mood_modifiers()}
|
|
self.camera_modifiers = {m.id: m for m in get_camera_modifiers()}
|
|
self.color_grading_modifiers = {m.id: m for m in get_color_grading_modifiers()}
|
|
self.action_modifiers = {m.id: m for m in get_action_modifiers()}
|
|
self.era_modifiers = {m.id: m for m in get_era_modifiers()}
|
|
self.detail_modifiers = {m.id: m for m in get_detail_modifiers()}
|
|
self.composition_modifiers = {m.id: m for m in get_composition_modifiers()}
|
|
self.custom_style_templates = self._load_style_templates()
|
|
|
|
def _load_style_templates(self) -> Dict[str, str]:
|
|
"""Load predefined style transformation templates"""
|
|
return {
|
|
"pop_art_1960s": "1960s pop art style with bright colors, bold graphics, and commercial aesthetics",
|
|
"bauhaus": "Bauhaus style with geometric forms, primary colors, and modernist principles",
|
|
"art_deco": "Art Deco style with geometric patterns, luxurious details, and streamlined forms",
|
|
"cyberpunk": "cyberpunk aesthetic with neon lights, high-tech elements, and dystopian atmosphere",
|
|
"minimalist": "minimalist style with clean lines, negative space, and essential elements only",
|
|
"impressionist": "impressionist painting style with visible brushstrokes and emphasis on light",
|
|
"noir": "film noir style with dramatic shadows, high contrast, and moody atmosphere",
|
|
"futuristic": "futuristic design with sleek forms, advanced materials, and sci-fi elements",
|
|
"vintage": "vintage aesthetic with retro colors, grain, and nostalgic atmosphere",
|
|
"brutalist": "brutalist architecture style with raw concrete, geometric forms, and bold structure",
|
|
}
|
|
|
|
def apply_lighting_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a lighting modifier to a base prompt"""
|
|
if modifier_id not in self.modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
cleaned_prompt = self._strip_lighting_keywords(base_prompt)
|
|
return f"{cleaned_prompt}, {modifier.description}"
|
|
|
|
def apply_seasonal_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a seasonal modifier to a base prompt"""
|
|
if modifier_id not in self.seasonal_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.seasonal_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
cleaned_prompt = self._strip_seasonal_keywords(base_prompt)
|
|
return f"{cleaned_prompt}, {modifier.description}"
|
|
|
|
def apply_location_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a location modifier to a base prompt"""
|
|
if modifier_id not in self.location_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.location_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
cleaned_prompt = self._strip_location_keywords(base_prompt)
|
|
return f"{cleaned_prompt}, {modifier.description}"
|
|
|
|
def apply_weather_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a weather modifier to a base prompt"""
|
|
if modifier_id not in self.weather_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.weather_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
cleaned_prompt = self._strip_weather_keywords(base_prompt)
|
|
return f"{cleaned_prompt}, {modifier.description}"
|
|
|
|
def apply_mood_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a mood modifier to a base prompt"""
|
|
if modifier_id not in self.mood_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.mood_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for mood
|
|
|
|
def apply_camera_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a camera modifier to a base prompt"""
|
|
if modifier_id not in self.camera_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.camera_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for camera
|
|
|
|
def apply_color_grading_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a color grading modifier to a base prompt"""
|
|
if modifier_id not in self.color_grading_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.color_grading_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for color grading
|
|
|
|
def apply_action_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply an action modifier to a base prompt"""
|
|
if modifier_id not in self.action_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.action_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for action
|
|
|
|
def apply_era_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply an era modifier to a base prompt"""
|
|
if modifier_id not in self.era_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.era_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for era
|
|
|
|
def apply_detail_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a detail modifier to a base prompt"""
|
|
if modifier_id not in self.detail_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.detail_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for detail
|
|
|
|
def apply_composition_modifier(self, base_prompt: str, modifier_id: int,
|
|
preserve_base: bool = True) -> str:
|
|
"""Apply a composition modifier to a base prompt"""
|
|
if modifier_id not in self.composition_modifiers:
|
|
raise ValueError(f"Invalid modifier_id: {modifier_id}. Must be 1-20.")
|
|
|
|
modifier = self.composition_modifiers[modifier_id]
|
|
|
|
if preserve_base:
|
|
return f"{base_prompt}, {modifier.description}"
|
|
else:
|
|
return f"{base_prompt}, {modifier.description}" # No keyword stripping for composition
|
|
|
|
def _strip_lighting_keywords(self, prompt: str) -> str:
|
|
"""Remove common lighting/time-of-day keywords from prompt"""
|
|
keywords = [
|
|
"sunrise", "sunset", "morning", "afternoon", "evening", "night", "twilight",
|
|
"blue hour", "golden hour", "dawn", "dusk", "midday", "noon",
|
|
"clear sky", "cloudy", "overcast", "back-light", "side-light", "rim-light",
|
|
"wet road", "wet street", "dry", "moonlit", "sunlit"
|
|
]
|
|
return self._strip_keywords(prompt, keywords)
|
|
|
|
def _strip_seasonal_keywords(self, prompt: str) -> str:
|
|
"""Remove common seasonal keywords from prompt"""
|
|
keywords = [
|
|
"spring", "summer", "autumn", "fall", "winter",
|
|
"blooming", "flowers", "foliage", "leaves", "snow", "frost", "ice",
|
|
"heat haze", "cold", "warm days", "crisp air"
|
|
]
|
|
return self._strip_keywords(prompt, keywords)
|
|
|
|
def _strip_location_keywords(self, prompt: str) -> str:
|
|
"""Remove common location keywords from prompt"""
|
|
keywords = [
|
|
"urban", "downtown", "city", "industrial", "suburban", "commercial",
|
|
"countryside", "rural", "farmland", "country road", "highway",
|
|
"interior", "showroom", "garage", "studio", "museum",
|
|
"coastal", "desert", "forest", "mountain", "bridge"
|
|
]
|
|
return self._strip_keywords(prompt, keywords)
|
|
|
|
def _strip_weather_keywords(self, prompt: str) -> str:
|
|
"""Remove common weather keywords from prompt"""
|
|
keywords = [
|
|
"clear", "cloudy", "overcast", "rain", "storm", "snow", "fog", "mist",
|
|
"drizzle", "downpour", "precipitation", "wind", "breeze", "calm"
|
|
]
|
|
return self._strip_keywords(prompt, keywords)
|
|
|
|
def _strip_keywords(self, prompt: str, keywords: List[str]) -> str:
|
|
"""Generic keyword stripping"""
|
|
cleaned = prompt
|
|
for keyword in keywords:
|
|
cleaned = cleaned.replace(keyword, "")
|
|
cleaned = cleaned.replace(keyword.title(), "")
|
|
cleaned = cleaned.replace(keyword.upper(), "")
|
|
|
|
cleaned = " ".join(cleaned.split())
|
|
cleaned = cleaned.replace(" ,", ",").replace(",,", ",")
|
|
cleaned = cleaned.strip(", ")
|
|
return cleaned
|
|
|
|
def apply_style_transformation(self, base_prompt: str, style: str,
|
|
maintain_composition: bool = True) -> str:
|
|
"""Apply a style transformation to a prompt"""
|
|
if style.lower().replace(" ", "_") in self.custom_style_templates:
|
|
style_desc = self.custom_style_templates[style.lower().replace(" ", "_")]
|
|
else:
|
|
style_desc = style
|
|
|
|
if maintain_composition:
|
|
return f"Transform to {style_desc} while maintaining the original composition. Base scene: {base_prompt}"
|
|
else:
|
|
return f"{base_prompt}, {style_desc}"
|
|
|
|
def get_modifiers_by_category(self, category: ModifierCategory) -> List[LightingModifier]:
|
|
"""Get all lighting modifiers of a specific category (DRY/WET)"""
|
|
return [m for m in self.modifiers.values() if m.category == category]
|
|
|
|
def generate_variation_batch(self, base_prompt: str,
|
|
modifier_ids: Optional[List[int]] = None,
|
|
include_styles: Optional[List[str]] = None) -> List[Dict[str, str]]:
|
|
"""Generate a batch of prompt variations"""
|
|
variations = []
|
|
|
|
if modifier_ids is None:
|
|
modifier_ids = list(range(1, 21))
|
|
|
|
for mod_id in modifier_ids:
|
|
modifier = self.modifiers[mod_id]
|
|
prompt = self.apply_lighting_modifier(base_prompt, mod_id)
|
|
variations.append({
|
|
"id": f"lighting_{mod_id}",
|
|
"type": "lighting",
|
|
"category": modifier.category.value,
|
|
"description": modifier.description,
|
|
"prompt": prompt
|
|
})
|
|
|
|
if include_styles:
|
|
for idx, style in enumerate(include_styles, 1):
|
|
prompt = self.apply_style_transformation(base_prompt, style)
|
|
variations.append({
|
|
"id": f"style_{idx}",
|
|
"type": "style",
|
|
"category": "transformation",
|
|
"description": style,
|
|
"prompt": prompt
|
|
})
|
|
|
|
return variations
|
|
|
|
def smart_variation(self, base_prompt: str, intent: str) -> str:
|
|
"""Apply intelligent variation based on natural language intent"""
|
|
intent_lower = intent.lower()
|
|
|
|
time_mappings = {
|
|
"morning": [1, 2, 3],
|
|
"sunrise": [1],
|
|
"afternoon": [5, 6, 14],
|
|
"golden hour": [7],
|
|
"sunset": [8],
|
|
"evening": [15],
|
|
"night": [11, 12],
|
|
"twilight": [9],
|
|
"blue hour": [10, 17, 18],
|
|
"wet": [17, 18, 19, 20]
|
|
}
|
|
|
|
for keyword, mod_ids in time_mappings.items():
|
|
if keyword in intent_lower:
|
|
return self.apply_lighting_modifier(base_prompt, mod_ids[0])
|
|
|
|
if "style" in intent_lower or "aesthetic" in intent_lower:
|
|
for style_name in self.custom_style_templates.keys():
|
|
if style_name.replace("_", " ") in intent_lower:
|
|
return self.apply_style_transformation(base_prompt, style_name)
|
|
|
|
return self.apply_style_transformation(base_prompt, intent)
|