app.py

# AI Image Creator: Enhanced UI and UX
# Part 3: Processing Logic, Prompt Enhancement, and Application Flow

# =============== PROMPT ENHANCEMENT LOGIC ===============

# Function to enhance prompt with Llama 4 with improved logical understanding
def enhance_prompt_with_llama(user_input, creation_type, art_style, mood):
    """
    Enhance user input with Llama 4 model to create detailed image generation prompts
    
    Args:
        user_input (str): User's original description
        creation_type (str): Selected creation type (e.g., "Digital Art")
        art_style (str): Selected art style (e.g., "Photorealistic")
        mood (str): Selected mood (e.g., "Peaceful")
        
    Returns:
        str: Enhanced prompt optimized for image generation
    """
    try:
        if not use_llama or llama_client is None:
            logger.warning("Llama enhancement not available, using fallback")
            return enhance_prompt_fallback(user_input, creation_type, art_style, mood)
        
        logger.info(f"Enhancing prompt with Llama 4 for creation type: {creation_type}, art style: {art_style}")
        
        # Enhanced Llama 4 system prompt with detailed instructions
        system_prompt = """You are a world-class prompt engineer who specializes in creating detailed, effective prompts for text-to-image AI models. 

Your task is to transform a user's simple description into a comprehensive, detailed image generation prompt that will create stunning visuals. Consider all the provided elements (description, creation type, art style, mood) and combine them into a cohesive, detailed prompt.

MOST IMPORTANTLY - ADD LOGICAL DETAILS:
- Analyze what the user wants and add logical details that would make the scene realistic or coherent
- If describing something fantastical (e.g., "flying cat"), add logical details about how this could work (e.g., "a cat with majestic feathered wings spread wide")
- Think about environment, lighting, perspective, time of day, weather, and other contextual elements
- Create a vivid, imaginable scene with spatial relationships clearly defined

PROMPT STRUCTURE GUIDELINES:
1. Start with the core subject and its primary characteristics
2. Add environment and setting details
3. Describe lighting, atmosphere, and mood
4. Include specific visual style and artistic technique references
5. Add technical quality terms (8K, detailed, masterful, etc.)

FORMAT YOUR RESPONSE AS A SINGLE PARAGRAPH with no additional comments, explanations, or bullet points. Use natural language without awkward comma separations. Aim for 75-150 words.

AVOID:
- Do not include quotation marks in your response
- Do not preface with "here's a prompt" or similar text
- Do not use placeholders
- Do not add negative prompts
- Do not write in list format or use bullet points

Respond only with the enhanced prompt and nothing else."""

        # Get creation type description
        creation_info = CREATION_TYPES.get(creation_type, {"description": "Create a detailed image", "icon": "🎨"})
        creation_description = creation_info["description"]
        
        # Get art style description
        style_info = ART_STYLES.get(art_style, {"description": "with detailed and professional quality", "icon": "🖌️"})
        style_description = style_info["description"]
        
        # Get mood description
        mood_info = MOODS.get(mood, {"description": "atmospheric", "icon": "✨"})
        mood_description = mood_info["description"]
        
        # Prepare the user prompt for Llama
        user_prompt = f"""Description: {user_input}
Creation Type: {creation_type} - {creation_description}
Art Style: {art_style} - {style_description}
Mood: {mood} - {mood_description}

Please create a comprehensive, detailed image generation prompt that combines all these elements."""

        try:
            # Request enhancement from Llama 4
            completion = llama_client.chat.completions.create(
                model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                max_tokens=500,
                temperature=0.7,  # Slight creativity while maintaining coherence
            )
            enhanced = completion.choices[0].message.content
            logger.info(f"Llama 4 enhanced prompt: {enhanced[:100]}...")
            return enhanced if enhanced else user_input
        except Exception as e:
            logger.error(f"Error during Llama enhancement: {str(e)}")
            return enhance_prompt_fallback(user_input, creation_type, art_style, mood)
    except Exception as e:
        logger.error(f"Error in Llama enhancement: {str(e)}")
        return enhance_prompt_fallback(user_input, creation_type, art_style, mood)

# Fallback prompt enhancement without Llama
def enhance_prompt_fallback(user_input, creation_type, art_style, mood):
    """
    Enhance user input without requiring Llama API using rule-based enhancement
    
    Args:
        user_input (str): User's original description
        creation_type (str): Selected creation type (e.g., "Digital Art")
        art_style (str): Selected art style (e.g., "Photorealistic")
        mood (str): Selected mood (e.g., "Peaceful")
        
    Returns:
        str: Enhanced prompt using predefined rules and templates
    """
    logger.info(f"Using fallback enhancement for: {user_input[:50]}...")
    
    # Quality terms by creation type
    quality_terms = {
        "Realistic Photo": [
            "photorealistic", "high resolution", "detailed", 
            "natural lighting", "sharp focus", "professional photography",
            "crisp details", "realistic textures", "DSLR photo"
        ],
        "Digital Art": [
            "vibrant colors", "clean lines", "digital illustration", 
            "polished", "professional digital art", "detailed rendering",
            "digital painting", "colorful", "vector-like precision"
        ],
        "Fantasy Illustration": [
            "magical atmosphere", "fantasy art", "detailed illustration", 
            "epic", "otherworldly", "imaginative scene",
            "fantasy environment", "magical lighting", "mythical qualities"
        ],
        "Concept Art": [
            "professional concept art", "detailed design", "conceptual illustration", 
            "industry standard", "visual development", "production artwork",
            "concept design", "detailed environment", "character design"
        ],
        "Anime/Manga": [
            "anime style", "manga illustration", "cel shaded", 
            "Japanese animation", "2D character art", "anime aesthetic",
            "clean linework", "anime proportions", "stylized features"
        ],
        "Oil Painting": [
            "oil on canvas", "textured brushwork", "rich colors", 
            "traditional painting", "artistic brushstrokes", "gallery quality",
            "glazed layers", "impasto technique", "classical painting style"
        ],
        "Watercolor": [
            "watercolor painting", "soft color bleeding", "delicate washes", 
            "transparent layers", "loose brushwork", "gentle transitions",
            "watercolor paper texture", "wet-on-wet technique", "fluid color blending"
        ],
        "Sketch": [
            "detailed sketch", "pencil drawing", "line art", 
            "hand-drawn", "fine details", "shading techniques",
            "graphite", "charcoal texture", "gestural lines"
        ],
        "3D Rendering": [
            "3D render", "volumetric lighting", "ray tracing", 
            "3D modeling", "realistic textures", "computer graphics",
            "physically based rendering", "global illumination", "ambient occlusion"
        ],
        "Pixel Art": [
            "pixel art", "8-bit style", "retro game aesthetic", 
            "limited color palette", "pixelated", "nostalgic game art",
            "16-bit look", "pixel perfect", "dithering effects"
        ]
    }
    
    # Style modifiers for different art styles - more detailed descriptions
    style_modifiers = {
        "Photorealistic": "highly detailed photorealistic style with perfect lighting, natural shadows, and lifelike textures",
        "Impressionist": "impressionist style with visible brushstrokes capturing light and atmosphere over precise details, reminiscent of Claude Monet",
        "Surrealist": "surrealist style with dreamlike and impossible elements, juxtaposed reality, inspired by Salvador Dali",
        "Pop Art": "pop art style with bold colors, sharp lines, halftone patterns and cultural references, like Andy Warhol",
        "Minimalist": "minimalist style with simplified forms, limited color palette, clean composition, and essential elements only",
        "Abstract": "abstract style using non-representational shapes, colors, and forms to express emotion rather than reality",
        "Cubist": "cubist style with geometric forms, multiple perspectives shown simultaneously, fractured surfaces, like Pablo Picasso",
        "Art Nouveau": "art nouveau style with ornate flowing lines inspired by natural forms, decorative elegance, and organic shapes",
        "Gothic": "gothic style with dark atmosphere, dramatic elements, pointed arches, and medieval-inspired architecture",
        "Cyberpunk": "cyberpunk style with neon colors, high tech low life aesthetic, futuristic technology, and urban decay",
        "Steampunk": "steampunk style with Victorian aesthetics, brass machinery, steam-powered technology, and retrofuturistic design",
        "Retro/Vintage": "retro style with nostalgic elements from past decades, aged texture, and period-appropriate colors and design",
        "Art Deco": "art deco style with geometric patterns, bold colors, symmetry, luxurious materials, and streamlined forms",
        "Baroque": "baroque style with dramatic lighting, rich details, contrast, dynamic composition, and ornate decorations",
        "Ukiyo-e": "ukiyo-e style Japanese woodblock print aesthetic with flat areas of color, strong outlines, and traditional subjects",
        "Comic Book": "comic book style with bold outlines, vibrant colors, dynamic action poses, and expressive characters",
        "Psychedelic": "psychedelic style with vibrant swirling colors, abstract patterns, distorted perspective, and 1960s-inspired visuals",
        "Vaporwave": "vaporwave aesthetic with glitch art, pastel colors, 80s/90s nostalgia, ancient statues, and digital elements",
        "Studio Ghibli": "Studio Ghibli anime style with whimsical detailed environments, soft colors, and charming character design",
        "Hyperrealism": "hyperrealistic style with extreme detail beyond photography, perfect textures, and meticulous precision"
    }
    
    # Mood modifiers for different moods - enhanced descriptions
    mood_modifiers = {
        "Happy": "bright cheerful atmosphere with warm golden lighting, vibrant colors, and uplifting elements",
        "Sad": "melancholic atmosphere with muted colors, soft shadows, and somber emotional tone",
        "Mysterious": "enigmatic atmosphere with shadows, fog, hidden elements, and dramatic lighting contrasts",
        "Peaceful": "serene calm atmosphere with gentle lighting, soft colors, and tranquil composition",
        "Tense": "suspenseful atmosphere with dramatic lighting, stark contrasts, and unsettling composition",
        "Whimsical": "playful whimsical atmosphere with imaginative elements, saturated colors, and fantastical details",
        "Dark": "dark gloomy atmosphere with deep shadows, limited lighting, and ominous elements",
        "Energetic": "dynamic vibrant atmosphere with strong colors, motion effects, and active composition",
        "Romantic": "soft romantic atmosphere with dreamy lighting, gentle colors, and intimate ambiance",
        "Epic": "grand epic atmosphere with dramatic scale, sweeping vistas, and majestic lighting"
    }
    
    # Get terms for the specific creation type, or use generic terms
    type_terms = quality_terms.get(creation_type, [
        "high quality", "detailed", "professional", "masterful", "high resolution", "sharp details"
    ])
    
    # Common quality terms enhanced with trending and technical terms
    common_terms = [
        "8K resolution", "highly detailed", "professional", "masterpiece",
        "trending on artstation", "award winning", "stunning", "intricate details",
        "perfect composition", "cinematic lighting"
    ]
    
    # Get style modifier
    style_modifier = style_modifiers.get(art_style, "detailed professional style")
    
    # Get mood modifier
    mood_modifier = mood_modifiers.get(mood, "atmospheric")
    
    # Basic prompt structure - core subject and style elements
    prompt_parts = [
        user_input,
        style_modifier,
        mood_modifier
    ]
    
    # Add randomly selected quality terms for variety
    selected_type_terms = random.sample(type_terms, min(3, len(type_terms)))
    selected_common_terms = random.sample(common_terms, min(3, len(common_terms)))
    
    # Combine terms
    quality_description = ", ".join(selected_type_terms + selected_common_terms)
    
    # Final enhanced prompt
    enhanced_prompt = f"{', '.join(prompt_parts)}, {quality_description}"
    
    logger.info(f"Fallback enhanced prompt: {enhanced_prompt[:100]}...")
    return enhanced_prompt

# =============== IMAGE GENERATION FUNCTIONS ===============

# Generate image function with loading state handling and retry mechanism
def generate_image(description, creation_type, art_style, mood, model_name, retries=1):
    """
    Generate image based on user inputs by enhancing prompt and calling image model API
    
    Args:
        description (str): User's original description
        creation_type (str): Selected creation type
        art_style (str): Selected art style
        mood (str): Selected mood
        model_name (str): Model identifier
        retries (int): Number of retries if generation fails
        
    Returns:
        tuple: (image, status_message, enhanced_prompt)
    """
    try:
        # Validate input
        if not description.strip():
            return None, "Please enter a description for your image", ""
            
        logger.info(f"Generating image with model: {model_name}")
        
        # Enhance prompt with Llama or fallback
        enhanced_prompt = enhance_prompt_with_llama(description, creation_type, art_style, mood)
        
        # Validate client availability
        if hf_client is None:
            logger.error("Hugging Face client not available")
            return None, "Error: Unable to connect to image generation service. Please try again later.", enhanced_prompt
        
        # Add negative prompt to avoid common issues
        negative_prompt = "low quality, blurry, distorted, deformed, disfigured, bad anatomy, watermark, signature, text, poorly drawn, amateur, ugly"
        
        try:
            # Generate image with progress tracking
            logger.info(f"Sending request to model {model_name} with prompt: {enhanced_prompt[:100]}...")
            
            # Log start time for performance tracking
            start_time = time.time()
            
            # Generate the image
            image = hf_client.text_to_image(
                prompt=enhanced_prompt,
                model=model_name,
                negative_prompt=negative_prompt
            )
            
            # Calculate generation time
            generation_time = time.time() - start_time
            logger.info(f"Image generated successfully in {generation_time:.2f} seconds")
            
            # Success message with generation details
            if use_llama:
                enhancement_method = "Llama 4 AI"
            else:
                enhancement_method = "rule-based enhancement"
                
            success_message = f"Image created successfully in {generation_time:.1f}s using {model_name.split('/')[-1]} model and {enhancement_method}"
            
            return image, success_message, enhanced_prompt
            
        except Exception as e:
            error_message = str(e)
            logger.error(f"Error during image generation: {error_message}")
            
            # Retry logic for transient errors
            if retries > 0:
                logger.info(f"Retrying image generation, {retries} attempts remaining")
                time.sleep(1)  # Small delay before retry
                return generate_image(description, creation_type, art_style, mood, model_name, retries - 1)
            
            # Format user-friendly error message
            if "429" in error_message:
                friendly_error = "Server is currently busy. Please try again in a few moments."
            elif "401" in error_message or "403" in error_message:
                friendly_error = "Authentication error with the image service. Please check API settings."
            elif "timeout" in error_message.lower():
                friendly_error = "Request timed out. The server might be under heavy load."
            else:
                friendly_error = f"Error generating image: {error_message}"
                
            return None, friendly_error, enhanced_prompt
    except Exception as e:
        logger.error(f"Unexpected error in generate_image: {str(e)}")
        return None, f"Unexpected error: {str(e)}", ""

# Wrapper function for generate_image with status updates
def generate_with_status(description, creation_type_val, art_style_val, mood_val, model_name):
    """
    Wrapper for generate_image that handles UI status updates and parameter formatting
    
    Args:
        description (str): User's original description
        creation_type_val (str): Formatted creation type with icon
        art_style_val (str): Formatted art style with icon
        mood_val (str): Formatted mood with icon
        model_name (str): Formatted model name with icon
        
    Returns:
        tuple: (image, status_html, enhanced_prompt, parameters_html)
    """
    # Check if description is empty
    if not description.strip():
        return None, update_status("Please enter a description", is_error=True), "", ""
    
    # Extract keys from formatted values
    creation_key = extract_key(creation_type_val)
    art_key = extract_key(art_style_val)
    mood_key = extract_key(mood_val)
    
    # Get model key from formatted name
    model_key = None
    for key, info in IMAGE_MODELS.items():
        if f"{info['icon']} {info['display_name']}" == model_name:
            model_key = key
            break
    
    if not model_key:
        return None, update_status("Invalid model selection", is_error=True), "", ""
    
    try:
        # Generate the image
        image, message, enhanced_prompt = generate_image(
            description, creation_key, art_key, mood_key, model_key
        )
        
        if image is None:
            return None, update_status(message, is_error=True), "", ""
        
        # Format parameters display
        params_html = format_parameters(creation_type_val, art_style_val, mood_val, model_name)
        
        # Success message
        success_message = update_status(message)
        return image, success_message, enhanced_prompt, params_html
        
    except Exception as e:
        error_message = str(e)
        logger.error(f"Error in generate_with_status: {error_message}")
        return None, update_status(f"Error: {error_message}", is_error=True), "", ""

# =============== MAIN APPLICATION FLOW ===============

def main():
    """
    Main application entry point - creates UI and sets up event handlers
    """
    # Create the UI components
    interface, description_input, creation_type, art_style, mood_dropdown, model_selector, generate_button, image_output, generation_status, prompt_output, parameters_display, char_counter, creation_info, art_info, model_info = create_ui()
    
    # Set up event handlers
    
    # Character counter update
    description_input.change(
        fn=update_char_count,
        inputs=description_input,
        outputs=char_counter
    )
    
    # Creation type info update
    creation_type.change(
        fn=update_creation_info,
        inputs=creation_type,
        outputs=creation_info
    )
    
    # Art style info update
    art_style.change(
        fn=update_art_style_info,
        inputs=art_style,
        outputs=art_info
    )
    
    # Model info update
    model_selector.change(
        fn=update_model_info,
        inputs=model_selector,
        outputs=model_info
    )
    
    # Generate button click handler
    generate_button.click(
        fn=generate_with_status,
        inputs=[
            description_input, 
            creation_type, 
            art_style, 
            mood_dropdown, 
            model_selector
        ],
        outputs=[
            image_output, 
            generation_status, 
            prompt_output,
            parameters_display
        ]
    )
    
    # Load default values on page load
    def load_defaults():
        creation_val = f"{CREATION_TYPES['Digital Art']['icon']} Digital Art"
        art_val = f"{ART_STYLES['Photorealistic']['icon']} Photorealistic"
        model_val = f"{IMAGE_MODELS['stabilityai/stable-diffusion-xl-base-1.0']['icon']} {IMAGE_MODELS['stabilityai/stable-diffusion-xl-base-1.0']['display_name']}"
        
        creation_info = update_creation_info(creation_val)
        art_info = update_art_style_info(art_val)
        model_info = update_model_info(model_val)
        
        return creation_info, art_info, model_info
    
    # Load defaults when the interface loads
    interface.load(
        fn=load_defaults,
        outputs=[creation_info, art_info, model_info]
    )
    
    # Launch the interface with theme and analytics disabled for privacy
    interface.launch(
        share=False,  # Set to True to create a public link
        debug=False,  # Set to True for development 
        enable_queue=True,  # Enable request queue for better handling under load
    )

# =============== APP EXECUTION ===============

if __name__ == "__main__":
    # Install dependencies if needed
    try:
        import pkg_resources
        required_packages = ['huggingface_hub', 'gradio', 'pillow']
        for package in required_packages:
            try:
                pkg_resources.get_distribution(package)
            except pkg_resources.DistributionNotFound:
                logger.info(f"Installing required package: {package}")
                import subprocess
                subprocess.check_call(['pip', 'install', package])
                logger.info(f"Successfully installed {package}")
    except Exception as e:
        logger.warning(f"Error checking or installing dependencies: {str(e)}")
    
    # Start the application
    main()