Spaces:
Sleeping
Sleeping
File size: 2,356 Bytes
1061354 | 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 | from PIL import Image
from io import BytesIO
from pathlib import Path
from typing import Optional
from loguru import logger
from src.llm.gemini_client import GeminiImageGeneration
from src.prompts.imagen_prompt import IMAGEN_PROMPT_TEMPLATE
from src.config import settings
class ImageGenerator:
"""
Class to handle image generation using Gemini Image Generation API.
This class provides methods to generate image prompts based on product details
and save the generated images to a specified directory.
"""
def __init__(self) -> None:
self.imagen = GeminiImageGeneration()
self.prompt_template = IMAGEN_PROMPT_TEMPLATE
self.save_dir = Path(settings.UPLOAD_DIR)
self.save_dir.mkdir(parents=True, exist_ok=True)
def generate_image_prompt(
self,
product_name: str,
brand_name: str,
product_description: str
) -> Optional[str]:
"""
Generate image prompt for the given product details.
Args:
product_name: Name of the product
brand_name: Brand name of the product
product_description: Description of the product
Returns:
Generated image prompt string
"""
try:
file_name = f"{product_name.replace(' ', '_')}_{brand_name.replace(' ', '_')}.png"
file_path = self.save_dir / file_name
prompt = self.prompt_template.format(
product_name=product_name,
brand_name=brand_name,
product_description=product_description
)
response = self.imagen.generate_image(prompt=prompt)
for part in response.candidates[0].content.parts:
if part.text is not None:
logger.info(f"Generated text: {part.text}")
elif part.inline_data is not None:
image_data = part.inline_data.data
image = Image.open(BytesIO(image_data))
image.save(file_path)
logger.info(f"Image saved to: {file_path}")
return str(file_path)
except Exception as e:
logger.error(f"Error generating image prompt: {e}")
raise RuntimeError(f"Failed to generate image prompt: {e}")
|