Update main.py

main.py

@@ -18,74 +18,41 @@ class GeneratedPrompt(BaseModel):
     prompt: str
 
 def image_url_to_base64(url):
-    response = requests.get(url, stream=True)
+    response = requests.get(url, stream=True, timeout=60)
     response.raise_for_status()
-
-    content_type = response.headers['Content-Type']
-
-    encoded_string = base64.b64encode(response.content).decode('utf-8')
-
+    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."""
-
+    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
-                    }
-                }
-            ]
-        }
+        {"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
-
+    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"
-        }
+        input={"prompt": prompt, "aspect_ratio": "1:1"}
     )
-
-    print(output)
     urls = []
     if isinstance(output, list) and output:
         first = output[0]
@@ -95,7 +62,6 @@ def generate_images(prompt):
         urls = [output]
     elif hasattr(output, "url"):
         urls = [getattr(output, "url")]
-
     return urls[0]
 
 def fetch_image_bytes(url):
@@ -104,7 +70,6 @@ def fetch_image_bytes(url):
     return r.content
 
 def init_s3():
-
     return boto3.client(
         "s3",
         endpoint_url=os.getenv("R2_ENDPOINT"),
@@ -115,7 +80,6 @@ def init_s3():
 
 def upload_to_r2(image_bytes):
     s3_client = init_s3()
-
     filename = f"{uuid4().hex}.png"
     file_key = f"infinityverse/{filename}"
     s3_client.put_object(
@@ -124,11 +88,16 @@ def upload_to_r2(image_bytes):
         Body=image_bytes,
         ContentType="image/png",
     )
-    r2_url = f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}'
-    return r2_url
+    return f'{os.getenv("NEW_BASE").rstrip("/")}/{file_key}'
 
 def get_images(reference_images):
-    items = reference_images.get("items", [])
+    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)
 
@@ -137,7 +106,6 @@ def get_images(reference_images):
         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(
@@ -156,7 +124,5 @@ def get_images(reference_images):
                     urls.append(f.result())
                 except Exception:
                     pass
-
         results.append({"ref_url": ref_url, "urls": urls})
-
-    return results
+    return results