pe.py

"""
Prompt Enhancement Module for Z-Image Turbo
Enhances user prompts for better image generation quality
"""

def enhance_prompt(user_prompt: str, language: str = "auto") -> str:
    """
    Enhances the user's prompt with additional quality tags and descriptors
    
    Args:
        user_prompt: The original user prompt
        language: Language of the prompt (auto-detect or specify)
    
    Returns:
        Enhanced prompt string
    """
    
    # Quality enhancement tags
    quality_tags = [
        "high quality",
        "detailed",
        "professional",
        "sharp focus",
        "8k resolution",
    ]
    
    # Detect if prompt is in Chinese
    is_chinese = any('\u4e00' <= char <= '\u9fff' for char in user_prompt)
    
    if is_chinese:
        # Chinese quality enhancement
        quality_suffix = "高质量,精细细节,专业摄影,清晰对焦,8K分辨率"
    else:
        # English quality enhancement
        quality_suffix = ", ".join(quality_tags)
    
    # Don't add quality tags if they're already present
    prompt_lower = user_prompt.lower()
    if any(tag in prompt_lower for tag in ["high quality", "detailed", "8k", "高质量", "精细"]):
        return user_prompt
    
    # Combine original prompt with quality enhancement
    if is_chinese:
        enhanced = f"{user_prompt},{quality_suffix}"
    else:
        enhanced = f"{user_prompt}, {quality_suffix}"
    
    return enhanced


def get_negative_prompt(style: str = "default") -> str:
    """
    Returns appropriate negative prompt based on style
    
    Args:
        style: Style preset for negative prompts
    
    Returns:
        Negative prompt string
    """
    
    negative_prompts = {
        "default": "low quality, blurry, distorted, deformed, ugly, bad anatomy, bad proportions, watermark, signature, text, error, cropped, worst quality, low quality jpeg artifacts",
        "portrait": "low quality, blurry, distorted face, deformed face, bad eyes, bad hands, bad anatomy, ugly, worst quality, watermark, signature",
        "landscape": "low quality, blurry, distorted, ugly, bad composition, watermark, signature, text, people, humans",
        "artistic": "low quality, blurry, photograph, realistic, watermark, signature, text",
    }
    
    return negative_prompts.get(style, negative_prompts["default"])


def parse_resolution(resolution_str: str) -> tuple:
    """
    Parses resolution string into width and height
    
    Args:
        resolution_str: Resolution string like "1024x1024 ( 1:1 )"
    
    Returns:
        Tuple of (width, height)
    """
    try:
        # Extract the resolution part before the ratio
        res_part = resolution_str.split("(")[0].strip()
        width, height = map(int, res_part.split("x"))
        return (width, height)
    except Exception as e:
        print(f"Error parsing resolution: {e}")
        return (1024, 1024)  # Default fallback


def validate_prompt(prompt: str, max_length: int = 500) -> tuple:
    """
    Validates user prompt
    
    Args:
        prompt: User's prompt
        max_length: Maximum allowed prompt length
    
    Returns:
        Tuple of (is_valid, error_message)
    """
    if not prompt or not prompt.strip():
        return (False, "Prompt cannot be empty")
    
    if len(prompt) > max_length:
        return (False, f"Prompt too long. Maximum {max_length} characters allowed")
    
    # Check for potentially harmful content
    banned_words = ["nude", "nsfw", "explicit", "porn"]
    prompt_lower = prompt.lower()
    
    for word in banned_words:
        if word in prompt_lower:
            return (False, "Prompt contains inappropriate content")
    
    return (True, "")


def create_prompt_template(
    base_prompt: str,
    style: str = None,
    mood: str = None,
    lighting: str = None,
    camera: str = None
) -> str:
    """
    Creates a structured prompt from components
    
    Args:
        base_prompt: The main subject/scene description
        style: Art style (e.g., "photorealistic", "anime", "oil painting")
        mood: Mood/atmosphere (e.g., "dramatic", "peaceful", "mysterious")
        lighting: Lighting description (e.g., "golden hour", "studio lighting")
        camera: Camera/lens info (e.g., "50mm lens", "wide angle")
    
    Returns:
        Formatted prompt string
    """
    components = [base_prompt]
    
    if style:
        components.append(f"{style} style")
    
    if mood:
        components.append(f"{mood} atmosphere")
    
    if lighting:
        components.append(f"{lighting}")
    
    if camera:
        components.append(f"shot with {camera}")
    
    return ", ".join(components)


# Preset prompt templates
PROMPT_TEMPLATES = {
    "portrait": "portrait of {subject}, professional studio lighting, bokeh background, sharp focus, detailed face, 85mm lens, f/1.8",
    "landscape": "{scene}, golden hour lighting, dramatic sky, wide angle, highly detailed, professional landscape photography",
    "product": "{product}, clean white background, studio lighting, commercial photography, high detail, sharp focus",
    "artistic": "{subject}, artistic interpretation, creative composition, vibrant colors, unique perspective",
    "cinematic": "{scene}, cinematic lighting, dramatic atmosphere, film grain, anamorphic lens, color graded",
}


def apply_template(template_name: str, **kwargs) -> str:
    """
    Applies a preset template with user variables
    
    Args:
        template_name: Name of the template to use
        **kwargs: Variables to fill in the template
    
    Returns:
        Filled template string
    """
    template = PROMPT_TEMPLATES.get(template_name)
    if not template:
        return kwargs.get("subject", "") or kwargs.get("scene", "") or kwargs.get("product", "")
    
    try:
        return template.format(**kwargs)
    except KeyError as e:
        print(f"Missing template variable: {e}")
        return template


# Chinese-specific prompt enhancement
def enhance_chinese_prompt(prompt: str) -> str:
    """
    Specifically enhances Chinese prompts with culturally relevant tags
    
    Args:
        prompt: Chinese language prompt
    
    Returns:
        Enhanced Chinese prompt
    """
    
    # Common Chinese photography/art terms
    enhancement_tags = {
        "人物": "精致五官,自然光线,专业人像摄影",
        "风景": "壮丽景色,黄金时刻,广角镜头,高清细节",
        "古风": "中国传统美学,古典韵味,精美细节,水墨画风格",
        "现代": "现代简约,专业摄影,高级感,时尚",
    }
    
    # Detect category and add appropriate enhancement
    for category, enhancement in enhancement_tags.items():
        if category in prompt:
            if enhancement not in prompt:
                prompt = f"{prompt},{enhancement}"
            break
    
    return prompt


def smart_enhance(prompt: str, auto_detect: bool = True) -> str:
    """
    Intelligently enhances prompt based on content and language
    
    Args:
        prompt: User's original prompt
        auto_detect: Whether to auto-detect language and style
    
    Returns:
        Smartly enhanced prompt
    """
    
    # Validate first
    is_valid, error = validate_prompt(prompt)
    if not is_valid:
        return prompt  # Return original if invalid
    
    # Detect language
    is_chinese = any('\u4e00' <= char <= '\u9fff' for char in prompt)
    
    if is_chinese and auto_detect:
        prompt = enhance_chinese_prompt(prompt)
    
    # General enhancement
    prompt = enhance_prompt(prompt)
    
    return prompt