Update ai images

ai_tasks/headlines_ai_images.py

@@ -7,11 +7,11 @@ from utils.io import print_system
 
 
 PROMPT = """
-Your goal is to create an ad for each of the following dimensions: {dimensions}.
-You will rely on an AI image generator to create the ads for you.
+Your goal is to create an ad with the following dimensions: {dimensions}.
+You will rely on an AI image generator to create an image for you.
 
-For each of the ad dimensions, create a prompt to be used by the AI image generator.
-Pick a dimension to map to from this list: 256x256, 512x512, 1024x1024.
+Create a prompt to be used by the AI image generator.
+Pick a dimension to map to from this list: 256x256, 512x512, 1024x1024. The AI image will have these dimensions.
 Create a headline.
 Consider that the headlines will be edited on top of the generated images.
 
@@ -20,20 +20,18 @@ Here is the summary of the website:
 
 Generate a JSON in the following format:
 ```
-[
-    {{
-        "ad_dimension":
-        "dimension_to_map":
-        "headline":
-        "prompt":
-    }}
-]
+{{
+    "ad_dimension":
+    "dimension_to_map":
+    "headline":
+    "prompt":
+}}
 ```
 """
 
 
-def generate_headlines(summary: str, dimensions: List[str]) -> List[Dict[str, str]]:
-    print_system("Generating headlines for website...")
+def generate_headline_and_prompt(summary: str, dimensions: str) -> Dict[str, str]:
+    print_system("Generating headline for website...")
     instructions = PROMPT.format(
         summary=summary,
         dimensions=dimensions,
@@ -42,8 +40,8 @@ def generate_headlines(summary: str, dimensions: List[str]) -> List[Dict[str, st
     return _parse_output(llm.next(messages, temperature=0))
 
 
-def _parse_output(assistant_message: str) -> List[Dict[str, str]]:
+def _parse_output(assistant_message: str) -> Dict[str, str]:
     # Might throw
-    match = re.search(r"(\[\s*{.*}\s*\])", assistant_message, re.DOTALL)
+    match = re.search("({.*})", assistant_message, re.DOTALL)
     json_request = match.group(0)  # type: ignore
     return json.loads(json_request)

ai_tasks/headlines_for_images.py

@@ -29,7 +29,7 @@ Original dimensions:
 """
 
 
-def get_headlines_for_images(
+def get_headline_for_image(
     summary: str, dimensions: str, image_labels: List[Dict]
 ) -> str:
     print_system("Generating ad from images...")

control_flow/main.py

@@ -2,8 +2,8 @@ from typing import Dict
 import json
 
 from ai import image
-from ai_tasks.headlines_for_images import get_headlines_for_images
-from ai_tasks.headlines_ai_images import generate_headlines
+from ai_tasks.headlines_for_images import get_headline_for_image
+from ai_tasks.headlines_ai_images import generate_headline_and_prompt
 from ai_tasks.text_summary import summarize_text
 from code_tasks.custom import get_image_info, run_parallel_jobs
 from code_tasks.images_in_url import get_images_from_url
@@ -34,17 +34,14 @@ def run():
     # AI tasks
     summary = summarize_text(text)
     print_assistant(summary)
-    headlines = get_headlines_for_images(summary, dimensions, image_info)
-    print_assistant(headlines)
-
-    # headlines_prompts = generate_headlines(summary, DIMENSIONS)
-    # print_system("Generating AI images...")
-    # run_parallel_jobs(gen_image, headlines_prompts)
-
-
-def gen_image(input: Dict[str, str]) -> None:
-    ai_image = image.urls(input["prompt"], size=input["dimension_to_map"])[0]
-    print_system(f"Prompt: {input['prompt']}")
-    print_assistant(input["ad_dimension"])
-    print_assistant(input["headline"])
+    # Pick an image and generate a headline
+    # headlines = get_headline_for_image(summary, dimensions, image_info)
+    # print_assistant(headlines)
+    # Generate a headline and an image
+    headline_prompt = generate_headline_and_prompt(summary, dimensions)
+    print_system("Generating AI images...")
+    ai_image = image.urls(headline_prompt["prompt"], size=headline_prompt["dimension_to_map"])[0]
+    print_system(f"Prompt: {headline_prompt['prompt']}")
+    print_assistant(headline_prompt["ad_dimension"])
+    print_assistant(headline_prompt["headline"])
     print_assistant(ai_image)