ai_bot_app/module_enhanced_image.py at main · yahayuta/ai_bot_app · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
import os
from typing import Dict
import module_stability
import module_openai
from module_prompt_builder import prompt_builder

class EnhancedImageGenerator:
    """Enhanced image generator with advanced prompt handling and multiple AI providers."""

    def __init__(self):
        self.stability_key = os.environ.get('STABILITY_KEY', '')
        self.openai_key = os.environ.get('OPENAI_TOKEN', '')

    def generate_with_enhanced_prompt(self,
                                    user_input: str,
                                    provider: str = 'auto',
                                    image_path: str = None) -> Dict:
        """
        Generate image with enhanced prompt processing.

        Args:
            user_input: User's prompt or structured input
            provider: 'stability', 'openai', or 'auto'
            image_path: Path to save the generated image

        Returns:
            Dict with generation results and metadata
        """
        # Parse user input for structured components
        if '|' in user_input:
            components = prompt_builder.parse_user_input(user_input)
            enhanced_prompt = prompt_builder.build_enhanced_prompt(
                base_prompt=components['base_prompt'],
                style=components['style'],
                location=components['location'],
                art_style=components['art_style']
            )
        else:
            # Simple prompt - apply basic enhancement
            enhanced_prompt = prompt_builder.build_enhanced_prompt(
                base_prompt=user_input
            )

        # Choose provider
        if provider == 'auto':
            provider = self._select_best_provider(enhanced_prompt)

        # Generate image
        try:
            if provider == 'stability':
                result = self._generate_stability(enhanced_prompt, image_path)
            else:  # openai
                result = self._generate_openai(enhanced_prompt, image_path)

            result['enhanced_prompt'] = enhanced_prompt
            result['provider'] = provider
            return result

        except Exception as e:
            return {
                'success': False,
                'error': str(e),
                'enhanced_prompt': enhanced_prompt,
                'provider': provider
            }

    def _select_best_provider(self, prompt: str) -> str:
        """Select the best provider based on prompt characteristics."""
        # Stability AI is better for artistic styles and complex compositions
        artistic_keywords = ['oil painting', 'watercolor', 'anime', 'cartoon', 'artistic']
        if any(keyword in prompt.lower() for keyword in artistic_keywords):
            return 'stability'

        # OpenAI DALL-E is better for realistic and simple subjects
        return 'openai'

    def _generate_stability(self, prompt: str, image_path: str) -> Dict:
        """Generate image using Stability AI."""
        if not self.stability_key:
            raise Exception("Stability AI key not configured")

        module_stability.generate(prompt, image_path)

        return {
            'success': True,
            'provider': 'stability',
            'prompt': prompt
        }

    def _generate_openai(self, prompt: str, image_path: str) -> Dict:
        """Generate image using OpenAI DALL-E."""
        if not self.openai_key:
            raise Exception("OpenAI key not configured")

        module_openai.openai_create_image(prompt, image_path)

        return {
            'success': True,
            'provider': 'openai',
            'prompt': prompt
        }

    def generate_random_artwork(self, subject: str = None, image_path: str = None) -> Dict:
        """Generate artwork with random artistic style."""
        if not subject:
            subject = prompt_builder.get_random_topic()

        random_style = prompt_builder.get_random_style()
        random_location = prompt_builder.get_random_location()

        enhanced_prompt = prompt_builder.build_enhanced_prompt(
            base_prompt=subject,
            art_style=random_style,
            location=random_location
        )

        return self.generate_with_enhanced_prompt(enhanced_prompt, 'auto', image_path)

    def generate_anime_character(self, character_description: str, anime_style: str = None, image_path: str = None) -> Dict:
        """Generate anime-style character."""
        enhanced_prompt = prompt_builder.build_anime_prompt(character_description, anime_style)
        return self.generate_with_enhanced_prompt(enhanced_prompt, 'stability', image_path)

    def generate_photorealistic(self, subject: str, location: str = None, image_path: str = None) -> Dict:
        """Generate photorealistic image."""
        enhanced_prompt = prompt_builder.build_photorealistic_prompt(subject, location)
        return self.generate_with_enhanced_prompt(enhanced_prompt, 'openai', image_path)

# Global instance
enhanced_generator = EnhancedImageGenerator()