Spaces:
Sleeping
Sleeping
| import os | |
| import configparser | |
| import PIL | |
| from PIL import Image | |
| from typing import Any, List, Optional, Literal | |
| from pydantic import BaseModel, Field | |
| import google.generativeai as genai | |
| import vertexai | |
| from vertexai.preview.vision_models import ImageGenerationModel, GeneratedImage | |
| from vertexai.preview.vision_models import Image as VertexImage | |
| # Load configuration | |
| config = configparser.ConfigParser() | |
| config.read('config.ini') | |
| api_key = config['credentials']['gemini_api_key'] | |
| proj_id = config['credentials']['gcp_project_id'] | |
| info = {"PROJECT_ID": proj_id, "LOCATION": "us-central1", "API_KEY": api_key} | |
| # Create directories for temporary and output images if they do not exist | |
| os.makedirs('images/temp', exist_ok=True) | |
| os.makedirs('images/output', exist_ok=True) | |
| class BannerGenerator(BaseModel): | |
| """Class responsible for generating banners using generative models.""" | |
| CONFIGS: dict | |
| topic: str | |
| images: List[str] | None = Field(default=None) | |
| aspect_ratio: Optional[Literal["1:1", "9:16", "16:9", "4:3", "3:4"]] = None # Added aspect_ratio attribute | |
| text_model: str = "gemini-1.5-flash" | |
| image_model: str = "imagen-3.0-generate-001" | |
| edit_model: str = "imagegeneration@006" | |
| pm: Any = None | |
| im: Any = None | |
| em: Any = None | |
| text_v0: str = None | |
| text_v1: str = None | |
| text_v2: str = None | |
| text_v3: str = None | |
| img_response_v1: Any = None | |
| img_response_v2: Any = None | |
| img_response_v3: Any = None | |
| launch_state: bool = False | |
| def __launch(self): | |
| """Launches the generative models and sets up environment.""" | |
| if not self.launch_state: | |
| vertexai.init(project=self.CONFIGS['PROJECT_ID'], location=self.CONFIGS['LOCATION']) | |
| genai.configure(api_key=self.CONFIGS['API_KEY']) | |
| self.pm = genai.GenerativeModel(self.text_model) | |
| self.im = ImageGenerationModel.from_pretrained(self.image_model) | |
| self.em = ImageGenerationModel.from_pretrained(self.edit_model) | |
| self.launch_state = True | |
| print("Model Launch successful!") | |
| def load_images(self) -> List[PIL.Image.Image]: | |
| """Loads images from file paths provided in the `images` attribute.""" | |
| self.__launch() | |
| loaded_images = [] | |
| for image_path in self.images: | |
| img = PIL.Image.open(image_path) | |
| if img.mode == 'RGBA': | |
| img = img.convert('RGB') # Convert to RGB if needed | |
| loaded_images.append(img) | |
| return loaded_images | |
| def extract_image_information(self) -> str: | |
| """Extracts information from images using the generative text model.""" | |
| images = self.load_images() | |
| extraction_prompt = '''Examine the set of images to extract information about the product (name, logo) in less than 80 words.''' | |
| model_input = [extraction_prompt] + images | |
| response = self.pm.generate_content(model_input) | |
| print("Attached images examined!") | |
| return response.text | |
| def extract_information(self) -> None: | |
| """Extracts information from the given topic and images using a detailed analysis.""" | |
| self.__launch() | |
| out_text = f"""Deep analyze text from retail advertising, marketing psychology, and thoroughly researched marketing studies perspective: {self.topic} | |
| Extract the following information: | |
| 0. Product: Product name, brand and supplier, logo, tagline, size, packaging if available | |
| 1. Objective: 1 word for primary goal of the banner. Example - Awareness, Engagement, Conversion, Branding | |
| 2. Festival: Event or occasion it may be tied to. Example - Christmas, Diwali, Black Friday, Summer Sale, New Year, Generic | |
| 3. Headline: Suggest a main text that captures attention. Example - Discover [product] for [festival], Shop now and save!, Limited time offer, Innovate your life with [product] | |
| 4. Subheadline: Optional additional supporting information to clarify the offer. Example - Get 50% off until [date], Exclusive deal for festive season, Hurry offer ends soon | |
| 5. CTA: Add a call to action. Example - Buy now, Shop the collection, Discover More, Sign up today | |
| 6. Color Scheme: Use color palette based on audience, occasion, or product tone. Example - Red & Gold (Festive, Urgency), Blue & White (Trust, Calm), Green & Brown (Eco-friendly, Natural), Black & White (Elegant, Minimal) | |
| 7. Promotional offer: Suggest 1 best promotional offer. Example - MAX ₹99 OFF, UP TO 60% OFF, UNDER ₹999, MIN ₹10 OFF, MIN 20% OFF, STARTS @₹99, FLAT ₹100 OFF, FLAT 20% OFF, ₹499 STORE, BUY 2 GET 1 FREE | |
| 8. Background color gradient: Dynamic color generation to match overall look and feel | |
| 9. Background theme: Festival oriented or generic if no festival | |
| """ | |
| self.text_v0 = self.pm.generate_content(out_text).text | |
| # Information consolidation | |
| out_text = f"Respond concisely and summarize in Python dictionary format only this: {self.text_v0}" | |
| if self.images: | |
| image_info = self.extract_image_information() | |
| out_text += ' Product insights: ' + image_info | |
| print(f"Product insights: {image_info}") | |
| self.text_v1 = self.pm.generate_content(out_text).text[9:-5] | |
| # Scrapper to ensure data integrity and consistency | |
| out_text = f"Respond concisely by scrapping all unavailable information in Python dictionary format only this: {self.text_v1}" | |
| self.text_v2 = self.pm.generate_content(out_text).text | |
| print("Information collection complete!") | |
| def create_text_prompt(self) -> None: | |
| """Creates a text prompt based on the extracted information.""" | |
| out_text = f"""Task: Fill in the values in this json: {self.text_v2} | |
| Guidelines: | |
| 1. It will be used to generate an ads banner. | |
| 2. Ensure it has all details pair-wise meticulously captured. | |
| 3. All unknown/missing/unprovided variables are replaced with the attributes of the most probable shopper for that product. | |
| 4. Recheck and identify all ambiguity or any text that leads to uncertainty. | |
| 5. Replace all uncertainty with targeted values that make the most sense for the given product. | |
| 6. Quantify everything possible, like high, medium, and lows to percentage values based on marketing and psychometric research studies. | |
| 7. All KPIs and qualitative measures are to be used subcontextually only. Remove any details about statistical testing or names of any performance KPIs. | |
| 8. Avoid sentences and use only necessary keywords. | |
| 9. Remove all redundant key-value pairs. | |
| """ | |
| self.text_v3 = self.pm.generate_content(out_text).text | |
| print("Information processed!") | |
| def generate_image(self) -> str: | |
| """Generates an image based on the given prompt and saves it in images/temp/.""" | |
| prompt = f"""Realistic, subcontextually implied qualitative attributes inspired, excellent image quality ad capturing every detail in json:{self.text_v3}""" | |
| # Adding the aspect ratio parameter to the image generation | |
| if self.aspect_ratio: | |
| self.img_response_v1 = self.im.generate_images(prompt=prompt, aspect_ratio=self.aspect_ratio) | |
| else: | |
| self.img_response_v1 = self.im.generate_images(prompt=prompt) | |
| # Save the generated image to images/temp/ | |
| temp_img_path = 'images/temp/temp.jpg' | |
| if os.path.exists(temp_img_path): | |
| os.remove(temp_img_path) | |
| self.img_response_v1.images[0].save(temp_img_path) | |
| print("Image v1 generated!") | |
| return temp_img_path | |
| def identify_lags(self) -> str: | |
| """Identifies quality issues in the generated image and provides suggestions for improvement.""" | |
| prompt = f"""Be direct. Quality check the banner out of 10 on: | |
| 1. Promotional offer present as per instructions below | |
| 2. Ensure ALL texts pass grammatical checks | |
| 3. Color pallette as per instructions below | |
| 4. Occassion or festival theme is present as per instructions below | |
| ONLY USE INFORMATION FROM {self.text_v3}. Don't DO NOT make up colors, promo or occasion. Make sure the promo and color pallete is followed as per above instructions. | |
| Precisely point out errors and corresponding actions to fix the image where score is below 8. | |
| Do not output anything about elements that need no change. | |
| """ | |
| temp_img_path = 'images/temp/temp.jpg' | |
| response = self.pm.generate_content([prompt, PIL.Image.open(temp_img_path)]) | |
| print(f'Lags identified: {response.text}') | |
| return response.text | |
| def fix_image(self, retest: bool = False) -> str: | |
| """Attempts to fix the identified lags in the generated image and saves it.""" | |
| prompt = f'Realistic, subcontextually implied qualitative attributes inspired, excellent image quality ad by: {self.identify_lags()}' | |
| temp_img_path = 'images/temp/temp.jpg' | |
| base_image = VertexImage.load_from_file(location=temp_img_path) | |
| self.img_response_v2 = self.em.edit_image( | |
| base_image=base_image, | |
| prompt=prompt, | |
| edit_mode="inpainting-insert", | |
| mask_mode="background" | |
| ) | |
| self.img_response_v2.images[0].save(temp_img_path) | |
| print("Image v2 generated!") | |
| if retest: | |
| prompt = f'Realistic, subcontextually implied qualitative attributes inspired, excellent image quality ad edit by: {self.identify_lags()}' | |
| self.img_response_v3 = self.em.edit_image( | |
| base_image=VertexImage.load_from_file(location=temp_img_path), | |
| prompt=prompt, | |
| edit_mode="inpainting-insert", | |
| mask_mode="background" | |
| ) | |
| self.img_response_v3.images[0].save(temp_img_path) | |
| print("Image v3 generated!") | |
| # Save final image to output directory with a name corresponding to the topic | |
| output_filename = ''.join(e for e in self.topic if e.isalnum() or e == ' ') | |
| output_filename = output_filename[:50].strip().replace(' ', '_') + ".jpg" | |
| output_img_path = os.path.join('images/output', output_filename) | |
| self.img_response_v2.images[0].save(output_img_path) | |
| print(f"Final image saved as: {output_img_path}") | |
| return output_img_path | |
| def execute(self, QC=False) -> str: | |
| """Executes the entire workflow to generate and refine the banner.""" | |
| self.extract_information() | |
| self.create_text_prompt() | |
| self.generate_image() | |
| final_image_path = self.fix_image(retest=QC) | |
| return final_image_path |