Update main.py

main.py

@@ -0,0 +1,162 @@
+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