main.py

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

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)
    response.raise_for_status()

    content_type = response.headers['Content-Type']

    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."""

    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
    final_prompt = final_prompt.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"
        }
    )

    print(output)
    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):
    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=image_bytes,
        ContentType="image/png",
    )
    r2_url = f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}'
    return r2_url

def get_images(reference_images):
    items = reference_images.get("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