malepati/custom_template_working

1import asyncio2import os3import aiohttp4from google import genai5from google.genai.types import GenerateContentConfig6from openai import AsyncOpenAI7from models.image_prompt import ImagePrompt8from models.sql.image_asset import ImageAsset9from utils.download_helpers import download_file10from utils.get_env import get_pexels_api_key_env11from utils.get_env import get_pixabay_api_key_env12from utils.image_provider import (13    is_image_generation_disabled,14    is_pixels_selected,15    is_pixabay_selected,16    is_gemini_flash_selected,17    is_dalle3_selected,18)19import uuid20 21 22class ImageGenerationService:23    def __init__(self, output_directory: str):24        self.output_directory = output_directory25        self.is_image_generation_disabled = is_image_generation_disabled()26        self.image_gen_func = self.get_image_gen_func()27 28    def get_image_gen_func(self):29        if self.is_image_generation_disabled:30            return None31 32        if is_pixabay_selected():33            return self.get_image_from_pixabay34        elif is_pixels_selected():35            return self.get_image_from_pexels36        elif is_gemini_flash_selected():37            return self.generate_image_google38        # elif is_dalle3_selected():39        #     return self.generate_image_openai40        return None41 42    def is_stock_provider_selected(self):43        return is_pixels_selected() or is_pixabay_selected()44 45    async def generate_image(self, prompt: ImagePrompt) -> str | ImageAsset:46        """47        Generates an image based on the provided prompt.48        - If no image generation function is available, returns a placeholder image.49        - If the stock provider is selected, it uses the prompt directly,50        otherwise it uses the full image prompt with theme.51        - Output Directory is used for saving the generated image not the stock provider.52        """53        if self.is_image_generation_disabled:54            print("Image generation is disabled. Using placeholder image.")55            return "/static/images/placeholder.jpg"56 57        if not self.image_gen_func:58            print("No image generation function found. Using placeholder image.")59            return "/static/images/placeholder.jpg"60 61        image_prompt = prompt.get_image_prompt(62            with_theme=not self.is_stock_provider_selected()63        )64        print(f"Request - Generating Image for {image_prompt}")65 66        try:67            if self.is_stock_provider_selected():68                image_path = await self.image_gen_func(image_prompt)69            else:70                image_path = await self.image_gen_func(71                    image_prompt, self.output_directory72                )73            if image_path:74                if image_path == "/static/images/placeholder.jpg":75                    return image_path76                77                if image_path.startswith("http"):78                    return image_path79                elif os.path.exists(image_path):80                    # Convert absolute path to relative web path if it is in the app_data directory81                    if "app_data" in image_path:82                         try:83                             # Extract relative path from app_data (e.g. /app_data/images/foo.png)84                             # Assuming path ends with app_data/images/filename.png85                             # We want /app_data/images/...86                             # Find index of app_data87                             idx = image_path.lower().find("app_data")88                             if idx != -1:89                                 relative_path = image_path[idx:].replace("\\", "/")90                                 return ImageAsset(91                                    path=f"/{relative_path}", # Prepend / for web URL92                                    is_uploaded=False,93                                    extras={94                                        "prompt": prompt.prompt,95                                        "theme_prompt": prompt.theme_prompt,96                                    },97                                )98                         except Exception as e:99                             print(f"Error converting path: {e}")100 101                    return ImageAsset(102                        path=image_path,103                        is_uploaded=False,104                        extras={105                            "prompt": prompt.prompt,106                            "theme_prompt": prompt.theme_prompt,107                        },108                    )109            raise Exception(f"Image not found at {image_path}")110 111        except Exception as e:112            print(f"Error generating image: {e}")113            return "/static/images/placeholder.jpg"114 115    async def generate_image_openai(self, prompt: str, output_directory: str) -> str:116        client = AsyncOpenAI(timeout=30.0)117        try:118            result = await client.images.generate(119                model="dall-e-3",120                prompt=prompt,121                n=1,122                quality="standard",123                size="1024x1024",124            )125            image_url = result.data[0].url126            return await download_file(image_url, output_directory)127        except Exception as e:128            print(f"Error calling OpenAI DALL-E: {e}")129            return "/static/images/placeholder.jpg"130 131    async def generate_image_google(self, prompt: str, output_directory: str) -> str:132        client = genai.Client()133        response = await asyncio.to_thread(134            client.models.generate_content,135            model="gemini-2.5-flash-image-preview",136            contents=[prompt],137            config=GenerateContentConfig(response_modalities=["TEXT", "IMAGE"]),138        )139 140        for part in response.candidates[0].content.parts:141            if part.text is not None:142                print(part.text)143            elif part.inline_data is not None:144                image_path = os.path.join(output_directory, f"{uuid.uuid4()}.jpg")145                with open(image_path, "wb") as f:146                    f.write(part.inline_data.data)147 148        return image_path149 150    async def get_image_from_pexels(self, prompt: str) -> str:151        async with aiohttp.ClientSession(trust_env=True) as session:152            response = await session.get(153                f"https://api.pexels.com/v1/search?query={prompt}&per_page=1",154                headers={"Authorization": f"{get_pexels_api_key_env()}"},155            )156            data = await response.json()157            image_url = data["photos"][0]["src"]["large"]158            return image_url159 160    async def get_image_from_pixabay(self, prompt: str) -> str:161        async with aiohttp.ClientSession(trust_env=True) as session:162            response = await session.get(163                f"https://pixabay.com/api/?key={get_pixabay_api_key_env()}&q={prompt}&image_type=photo&per_page=3"164            )165            data = await response.json()166            image_url = data["hits"][0]["largeImageURL"]167            return image_url168