Spaces:
Sleeping
Sleeping
| from openai import OpenAI | |
| from pydantic import BaseModel | |
| import replicate | |
| from dotenv import load_dotenv | |
| import os | |
| from uuid import uuid4 | |
| import requests | |
| import boto3 | |
| import base64 | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from meta_data import meta_data_helper_function | |
| load_dotenv() | |
| gpt_client = OpenAI(api_key=os.getenv("GEMENI_KEY"), | |
| base_url="https://generativelanguage.googleapis.com/v1beta/openai/") | |
| class GeneratedPrompt(BaseModel): | |
| prompt: str | |
| def image_url_to_base64(url): | |
| response = requests.get(url, stream=True, timeout=60) | |
| response.raise_for_status() | |
| content_type = response.headers.get("Content-Type", "image/png") | |
| encoded_string = base64.b64encode(response.content).decode("utf-8") | |
| return f"data:{content_type};base64,{encoded_string}" | |
| def generate_prompt(url): | |
| system_prompt = ( | |
| "You are a top-tier performance digital marketer and creative strategist with 15+ years of expertise in affiliate marketing. " | |
| "Your objective is to analyze the provided winning ad image, deconstruct its concept, visual composition, and color scheme, and generate a fresh, conversion-focused ad visual tailored. " | |
| "The new design should convey a same as original image sentiment. Create a visually compelling ad optimized for Facebook Ads that is scroll-stopping, pattern-interrupting, and designed to drive high CTR and Conversion Rate. " | |
| "Utilize striking color combinations, dynamic contrast levels, and strategic layout compositions to command attention while aligning with the target audience avatar. " | |
| "Make sure the images should be realistic, not be stocky at all and raw which should look like they are shot from an iPhone. Make sure you clearly mention the text overlays and style of them how to put them on image." | |
| ) | |
| messages = [ | |
| {"role": "system", "content": [{"type": "text", "text": system_prompt}]}, | |
| {"role": "user", "content": [ | |
| {"type": "text", "text": "Generate the prompt for the given ad image."}, | |
| {"type": "image_url", "image_url": {"url": url}} | |
| ]} | |
| ] | |
| response = gpt_client.beta.chat.completions.parse( | |
| model="gemini-2.0-flash", | |
| messages=messages, | |
| response_format=GeneratedPrompt, | |
| ) | |
| final_prompt = response.choices[0].message.parsed.prompt | |
| return final_prompt | |
| def generate_images(prompt): | |
| replicate_client = replicate.Client(api_token=os.getenv("REPLICATE_API_TOKEN")) | |
| output = replicate_client.run( | |
| "google/imagen-4-ultra", | |
| input={"prompt": prompt, "aspect_ratio": "1:1"} | |
| ) | |
| urls = [] | |
| if isinstance(output, list) and output: | |
| first = output[0] | |
| url = getattr(first, "url", str(first)) | |
| urls = [url] | |
| elif isinstance(output, str): | |
| urls = [output] | |
| elif hasattr(output, "url"): | |
| urls = [getattr(output, "url")] | |
| return urls[0] | |
| def fetch_image_bytes(url): | |
| r = requests.get(url, timeout=60) | |
| r.raise_for_status() | |
| return r.content | |
| def init_s3(): | |
| return boto3.client( | |
| "s3", | |
| endpoint_url=os.getenv("R2_ENDPOINT"), | |
| aws_access_key_id=os.getenv("R2_ACCESS_KEY"), | |
| aws_secret_access_key=os.getenv("R2_SECRET_KEY"), | |
| region_name="auto", | |
| ) | |
| def upload_to_r2(image_bytes): | |
| meta_data_img = meta_data_helper_function(image_bytes) | |
| s3_client = init_s3() | |
| filename = f"{uuid4().hex}.png" | |
| file_key = f"infinityverse/{filename}" | |
| s3_client.put_object( | |
| Bucket=os.getenv("R2_BUCKET_NAME"), | |
| Key=file_key, | |
| Body=meta_data_img, | |
| ContentType="image/png", | |
| ) | |
| return f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}' | |
| def get_images(reference_images): | |
| if isinstance(reference_images, list): | |
| items = reference_images | |
| elif isinstance(reference_images, dict): | |
| items = reference_images.get("items", []) | |
| else: | |
| items = [] | |
| results = [] | |
| max_workers = min(32, (os.cpu_count() or 1) * 2) | |
| for image in items: | |
| num = int(image["num"]) | |
| ref_url = image["ref_url"] | |
| base = image_url_to_base64(ref_url) | |
| urls = [] | |
| with ThreadPoolExecutor(max_workers=min(max_workers, max(1, num))) as ex: | |
| futures = [ | |
| ex.submit( | |
| lambda b=base: upload_to_r2( | |
| fetch_image_bytes( | |
| generate_images( | |
| generate_prompt(b) | |
| ) | |
| ) | |
| ) | |
| ) | |
| for _ in range(num) | |
| ] | |
| for f in as_completed(futures): | |
| try: | |
| urls.append(f.result()) | |
| except Exception: | |
| pass | |
| results.append({"ref_url": ref_url, "urls": urls}) | |
| return results | |