import modules.scripts as scripts from modules.processing import StableDiffusionProcessingTxt2Img import os import re import json from dotenv import load_dotenv import requests from typing import Dict, List, Literal, TypedDict, Union, Optional, Any, Callable, Annotated import functools # Try to import LangGraph related libraries try: from langgraph.graph import StateGraph, END, START from langgraph.checkpoint.memory import MemorySaver LANGGRAPH_AVAILABLE = True except ImportError: LANGGRAPH_AVAILABLE = False print("Warning: LangGraph library not installed, using simplified implementation") print("Can be installed via 'pip install langgraph'") # Try to import Pydantic try: from pydantic import BaseModel, Field PYDANTIC_AVAILABLE = True except ImportError: PYDANTIC_AVAILABLE = False print("Warning: Pydantic library not installed, using simplified implementation") print("Can be installed via 'pip install pydantic'") # Load environment variables load_dotenv() # Get DeepSeek API key DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY") # Define state type class PromptState(TypedDict): original_prompt: str language: str # Changed: from is_chinese to language, can be "english", "chinese", "other", etc. translated_prompt: Optional[str] optimized_prompt: Optional[str] error: Optional[str] class PromptTemplate(BaseModel): """Prompt template Pydantic model""" name: str = Field(..., description="Template name") content: str = Field(..., description="Template content") def __str__(self) -> str: return self.content.strip() class PromptTemplates(BaseModel): """Collection of prompt templates""" txt2img_optimizer: PromptTemplate = Field( default=PromptTemplate( name="Stable Diffusion Prompt Optimizer", content=""" You are an expert prompt engineer for Stable Diffusion image generation with deep knowledge of how SD models interpret text. Your task is to transform standard prompts into highly optimized versions that produce exceptional quality images. Follow these guidelines: 1. Maintain the original subject and core concept 2. Enhance with precise descriptive adjectives and specific details 3. Add appropriate artistic style references (artists, movements, platforms) 4. Incorporate quality-boosting terms (masterpiece, best quality, highly detailed) 5. Apply technical enhancements through brackets for emphasis: - Use (term) for 1.1x emphasis - Use ((term)) for 1.2x emphasis - Use [term] for 0.9x emphasis - Use [[term]] for 0.8x emphasis - Use :1.x for specific weighting 6. Structure prompts effectively: - Main subject first with strongest emphasis - Scene details and environment - Style, quality, and technical terms last Return ONLY the optimized prompt without explanations or commentary. Preserve all special formatting like (), [], {}, :1.2, etc. from the original prompt. """ ), description="Stable Diffusion prompt optimization template" ) language_detector: PromptTemplate = Field( default=PromptTemplate( name="Language Detector", content=""" You are a language detection expert. Your task is to identify if the given text is in English or not. Analyze the provided text and determine if it's in English. Return ONLY 'yes' if the text is primarily in English, or 'no' if it's primarily in another language. If the text is primarily in English or contains mostly English words with a few non-English terms, return 'yes'. If the text is primarily in another language, return 'no'. Return ONLY 'yes' or 'no' without any explanations or additional text. """ ), description="Language detection template" ) universal_translator: PromptTemplate = Field( default=PromptTemplate( name="Universal Translator", content=""" You are a professional translator specializing in translating text to English for image generation. Your task is to accurately translate prompts from any language to English while preserving the original meaning and intent. Follow these guidelines: 1. Maintain the core subject and concept of the original prompt 2. Preserve any special formatting like (), [], {}, :1.2, etc. 3. Translate cultural-specific terms appropriately for an international audience 4. Keep artistic style references intact 5. Ensure the translation is natural and fluent in English Return ONLY the translated English prompt without explanations or commentary. """ ), description="Universal translation template" ) def get(self, template_name: str) -> PromptTemplate: """Get template by name""" if hasattr(self, template_name): return getattr(self, template_name) raise ValueError(f"Template not found: {template_name}") # Create template instance TEMPLATES = PromptTemplates() # Agent functions def router_agent(state: PromptState) -> Dict[str, Any]: """Determine the language of the prompt""" prompt = state["original_prompt"] if not prompt: return {"language": "unknown"} try: headers = { "Content-Type": "application/json", "Authorization": f"Bearer {DEEPSEEK_API_KEY}" } # Use predefined language detection template detector_template = TEMPLATES.get("language_detector") payload = { "model": "deepseek-chat", "messages": [ {"role": "system", "content": detector_template.content}, {"role": "user", "content": f"Is this text in English? {prompt}"} ], "temperature": 0.1, "max_tokens": 10 } response = requests.post( "https://api.deepseek.com/v1/chat/completions", headers=headers, json=payload ) if response.status_code == 200: result = response.json() is_english = result["choices"][0]["message"]["content"].strip().lower() == "yes" language = "english" if is_english else "other" print(f"RouterAgent: Prompt '{prompt}' detected as '{'English' if language == 'english' else 'Non-English'}'") return {"language": language} else: print(f"RouterAgent: Language detection failed - {response.status_code} - {response.text}") # Fallback to simple detection non_ascii_chars = 0 for char in prompt: if ord(char) > 127: non_ascii_chars += 1 language = "english" if (non_ascii_chars / len(prompt) < 0.3) else "other" print(f"RouterAgent: Prompt '{prompt}' detected as '{'English' if language == 'english' else 'Non-English'}' (simple detection)") return {"language": language} except Exception as e: print(f"RouterAgent: Language detection failed - {str(e)}") # Fallback to simple detection non_ascii_chars = 0 for char in prompt: if ord(char) > 127: non_ascii_chars += 1 language = "english" if (non_ascii_chars / len(prompt) < 0.3) else "other" print(f"RouterAgent: Prompt '{prompt}' detected as '{'English' if language == 'english' else 'Non-English'}' (simple detection)") return {"language": language} def translator_agent(state: PromptState) -> Dict[str, Any]: """Translate non-English prompts to English""" prompt = state["original_prompt"] language = state["language"] if language == "english": print(f"TranslatorAgent: Prompt is already in English, no translation needed") return {"translated_prompt": prompt} if not DEEPSEEK_API_KEY: print(f"TranslatorAgent: Warning - DEEPSEEK_API_KEY not set, using simplified translation") return {"error": "DEEPSEEK_API_KEY not set", "translated_prompt": prompt} try: headers = { "Content-Type": "application/json", "Authorization": f"Bearer {DEEPSEEK_API_KEY}" } # Use predefined universal translation template translator_template = TEMPLATES.get("universal_translator") payload = { "model": "deepseek-chat", "messages": [ {"role": "system", "content": translator_template.content}, {"role": "user", "content": f"Translate this prompt from {language} to English: {prompt}"} ], "temperature": 0.1, "max_tokens": 1000 } response = requests.post( "https://api.deepseek.com/v1/chat/completions", headers=headers, json=payload ) if response.status_code == 200: result = response.json() translated_text = result["choices"][0]["message"]["content"].strip() print(f"TranslatorAgent: Translation result - '{translated_text}'") return {"translated_prompt": translated_text} else: print(f"TranslatorAgent: Translation failed - {response.status_code} - {response.text}") return {"error": f"Translation API error: {response.status_code}", "translated_prompt": prompt} except Exception as e: print(f"TranslatorAgent: Translation failed - {str(e)}") return {"error": f"Translation error: {str(e)}", "translated_prompt": prompt} def optimizer_agent(state: PromptState) -> Dict[str, Any]: """Optimize English prompts""" # Determine the prompt to optimize prompt_to_optimize = state.get("translated_prompt") or state["original_prompt"] if not DEEPSEEK_API_KEY: print("OptimizerAgent: Warning - DEEPSEEK_API_KEY not set, using local optimization") optimized = local_optimize(prompt_to_optimize) return {"optimized_prompt": optimized} try: headers = { "Content-Type": "application/json", "Authorization": f"Bearer {DEEPSEEK_API_KEY}" } # Use predefined optimization template optimizer_template = TEMPLATES.get("txt2img_optimizer") payload = { "model": "deepseek-chat", "messages": [ {"role": "system", "content": optimizer_template.content}, {"role": "user", "content": f"Optimize this prompt: {prompt_to_optimize}"} ], "temperature": 0.3, "max_tokens": 1000 } response = requests.post( "https://api.deepseek.com/v1/chat/completions", headers=headers, json=payload ) if response.status_code == 200: result = response.json() enhanced_text = result["choices"][0]["message"]["content"].strip() print(f"OptimizerAgent: Optimization result - '{enhanced_text}'") return {"optimized_prompt": enhanced_text} else: print(f"OptimizerAgent: Optimization failed - {response.status_code} - {response.text}") optimized = local_optimize(prompt_to_optimize) return {"error": f"Optimization API error: {response.status_code}", "optimized_prompt": optimized} except Exception as e: print(f"OptimizerAgent: Optimization failed - {str(e)}") optimized = local_optimize(prompt_to_optimize) return {"error": f"Optimization error: {str(e)}", "optimized_prompt": optimized} def local_optimize(prompt: str) -> str: """Local prompt optimization method (used when API is unavailable)""" # Example optimization: add quality-boosting keywords quality_terms = ["high quality", "detailed", "sharp focus"] style_terms = ["masterpiece", "best quality"] # Check if prompt already contains these terms optimized = prompt # Add quality terms for term in quality_terms: if term.lower() not in optimized.lower(): if optimized.strip().endswith(('，', '。', ',', '.')): optimized = f"{optimized} {term}" else: optimized = f"{optimized}, {term}" # Add style terms (at the beginning) for term in reversed(style_terms): if term.lower() not in optimized.lower(): optimized = f"{term}, {optimized}" print(f"OptimizerAgent: Local optimization result - '{optimized}'") return optimized # Define routing logic def should_translate(state: PromptState) -> Literal["translator", "optimizer"]: """Determine if translation is needed""" if state.get("language", "") != "english": return "translator" else: return "optimizer" # Create LangGraph workflow def create_prompt_optimization_graph(): """Create prompt optimization workflow graph""" # If LangGraph is not available, return None if not LANGGRAPH_AVAILABLE: return None # Create state graph workflow = StateGraph(PromptState) # Add nodes workflow.add_node("router", router_agent) workflow.add_node("translator", translator_agent) workflow.add_node("optimizer", optimizer_agent) # Add edges # From start to router workflow.set_entry_point("router") # From router to translator or optimizer (based on language) workflow.add_conditional_edges( "router", should_translate, { "translator": "translator", "optimizer": "optimizer" } ) # From translator to optimizer workflow.add_edge("translator", "optimizer") # From optimizer to end workflow.add_edge("optimizer", END) # Compile workflow return workflow.compile() # Simplified workflow (used when LangGraph is not available) def simple_prompt_optimization_workflow(prompt: str) -> str: """Simplified prompt optimization workflow""" print(f"\n--- Simplified workflow started ---") print(f"Original prompt: '{prompt}'") # Initialize state state = PromptState( original_prompt=prompt, language="unknown", translated_prompt=None, optimized_prompt=None, error=None ) # Step 1: Router - determine language router_result = router_agent(state) state["language"] = router_result["language"] # Step 2: Translator - translate if not English if state["language"] != "english": translator_result = translator_agent(state) state["translated_prompt"] = translator_result.get("translated_prompt") if "error" in translator_result: state["error"] = translator_result["error"] # Step 3: Optimizer - optimize prompt optimizer_result = optimizer_agent(state) state["optimized_prompt"] = optimizer_result.get("optimized_prompt") if "error" in optimizer_result and not state["error"]: state["error"] = optimizer_result["error"] print(f"Final optimized prompt: '{state['optimized_prompt']}'") print(f"--- Simplified workflow finished ---\n") return state["optimized_prompt"] or prompt class PromptOptimizer(scripts.Script): # Class-level flag to track if initialization message has been shown _init_message_shown = False def __init__(self): super().__init__() # Show initialization message only once if not PromptOptimizer._init_message_shown: print("\n\n=== Txt2Img Prompt Optimizer (Multilingual) script loaded ===\n\n") PromptOptimizer._init_message_shown = True # Try to create LangGraph workflow self.graph = create_prompt_optimization_graph() # If LangGraph is not available, use simplified workflow if self.graph is None and not PromptOptimizer._init_message_shown: print("Using simplified prompt optimization workflow") # Track processed prompts to avoid duplicates self.processed_prompts = set() def title(self): return "Txt2Img Prompt Optimizer (Multilingual)" # Return AlwaysVisible to show script in UI def show(self, is_img2img): return scripts.AlwaysVisible # No UI elements needed def ui(self, is_img2img): return [] # Optimize prompt before processing def process(self, p): # Only optimize Txt2Img processing objects if not isinstance(p, StableDiffusionProcessingTxt2Img): return p # Record original prompt original_prompt = p.prompt print(f"\n=== Original prompt ===\n{original_prompt}\n") # Optimize main prompt (if not already processed) if p.prompt not in self.processed_prompts: optimized_prompt = self.optimize_prompt(p.prompt) p.prompt = optimized_prompt # Ensure all_prompts also uses optimized prompt if hasattr(p, 'all_prompts') and p.all_prompts: p.all_prompts = [optimized_prompt] * len(p.all_prompts) # Ensure main_prompt also uses optimized prompt if hasattr(p, 'main_prompt'): p.main_prompt = optimized_prompt self.processed_prompts.add(optimized_prompt) # Record optimization information (optional, for verification) if not hasattr(p, 'extra_generation_params'): p.extra_generation_params = {} p.extra_generation_params['Prompt optimized'] = True # Record final prompt sent to model print(f"\n=== Final prompt sent to model ===\n{p.prompt}\n") # Add post-processing hook to ensure prompt remains optimized original_setup_prompts = p.setup_prompts def patched_setup_prompts(): # Call original method original_setup_prompts() # Ensure prompt remains optimized if p.prompt in self.processed_prompts: p.all_prompts = [p.prompt] * len(p.all_prompts) p.main_prompt = p.prompt # Replace method p.setup_prompts = patched_setup_prompts return p def postprocess(self, p, processed): """Post-process after image generation""" # Nothing to do here return processed def optimize_prompt(self, prompt): """Optimize prompt to improve generation quality""" if not prompt: return prompt # Use LangGraph workflow or simplified workflow if self.graph is not None: # Use LangGraph workflow try: print(f"\n--- LangGraph started ---") print(f"Original prompt: '{prompt}'") # Create initial state initial_state = PromptState( original_prompt=prompt, language="unknown", translated_prompt=None, optimized_prompt=None, error=None ) # Execute workflow final_state = self.graph.invoke(initial_state) optimized = final_state.get("optimized_prompt") or prompt print(f"Final optimized prompt: '{optimized}'") print(f"--- LangGraph finished ---\n") return optimized except Exception as e: print(f"LangGraph failed: {str(e)}") # Fallback to simplified workflow return simple_prompt_optimization_workflow(prompt) else: # Use simplified workflow return simple_prompt_optimization_workflow(prompt) # Check if we're in a Stable Diffusion Webui environment try: from modules.processing import StableDiffusionProcessingTxt2Img except ImportError: # Create a mock class for testing outside of webui class StableDiffusionProcessingTxt2Img: """Mock class for testing outside of webui""" def __init__(self): self.prompt = "" self.all_prompts = [] self.main_prompt = "" self.extra_generation_params = {} def setup_prompts(self): """Mock setup_prompts method""" pass # For standalone testing if __name__ == "__main__": # Test the prompt optimization workflow test_prompts = [ "a cat", "beautiful landscape", "portrait of a woman", "科幻城市", # Chinese: "sci-fi city" "美丽的山水画", # Chinese: "beautiful landscape painting" ] print("Testing prompt optimization workflow...") # Initialize optimizer optimizer = PromptOptimizer() # Test each prompt for prompt in test_prompts: print(f"\nTesting prompt: '{prompt}'") optimized = optimizer.optimize_prompt(prompt) print(f"Optimized: '{optimized}'")