add image generation for placeholder images

backend/image_generation.py

@@ -0,0 +1,83 @@
+import asyncio
+import os
+import re
+from openai import AsyncOpenAI
+from bs4 import BeautifulSoup
+
+
+async def process_tasks(prompts):
+    tasks = [generate_image(prompt) for prompt in prompts]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    processed_results = []
+    for result in results:
+        if isinstance(result, Exception):
+            print(f"An exception occurred: {result}")
+            processed_results.append(None)
+        else:
+            processed_results.append(result)
+
+    return processed_results
+
+
+async def generate_image(prompt):
+    client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+    image_params = {
+        "model": "dall-e-3",
+        "quality": "standard",
+        "style": "natural",
+        "n": 1,
+        "size": "1024x1024",
+        "prompt": prompt,
+    }
+    res = await client.images.generate(**image_params)
+    return res.data[0].url
+
+
+def extract_dimensions(url):
+    # Regular expression to match numbers in the format '300x200'
+    matches = re.findall(r"(\d+)x(\d+)", url)
+
+    if matches:
+        width, height = matches[0]  # Extract the first match
+        width = int(width)
+        height = int(height)
+        return (width, height)
+    else:
+        return (100, 100)
+
+
+async def generate_images(code):
+    # Find all images and extract their alt texts
+    soup = BeautifulSoup(code, "html.parser")
+    images = soup.find_all("img")
+    alts = [img.get("alt", None) for img in images]
+
+    # Exclude images with no alt text
+    alts = [alt for alt in alts if alt is not None]
+
+    # Remove duplicates
+    prompts = list(set(alts))
+
+    # Generate images
+    results = await process_tasks(prompts)
+
+    # Create a dict mapping alt text to image URL
+    mapped_image_urls = dict(zip(prompts, results))
+
+    # Replace alt text with image URLs
+    for img in images:
+        new_url = mapped_image_urls[img.get("alt")]
+
+        if new_url:
+            # Set width and height attributes
+            width, height = extract_dimensions(img["src"])
+            img["width"] = width
+            img["height"] = height
+            # Replace img['src'] with the mapped image URL
+            img["src"] = new_url
+        else:
+            print("Image generation failed for alt text:" + img.get("alt"))
+
+    # Return the modified HTML
+    return str(soup)

backend/main.py

@@ -1,19 +1,26 @@
 # Load environment variables first
-import json
 from dotenv import load_dotenv
-import os
-from datetime import datetime
-
-from prompts import assemble_prompt
 
 load_dotenv()
 
 
+import json
+import os
+import traceback
+from datetime import datetime
 from fastapi import FastAPI, WebSocket
+
 from llm import stream_openai_response
+from mock import MOCK_HTML, mock_completion
+from image_generation import generate_images
+from prompts import assemble_prompt
 
 app = FastAPI()
 
+# Useful for debugging purposes when you don't want to waste GPT4-Vision credits
+# Setting to True will stream a mock response instead of calling the OpenAI API
+SHOULD_MOCK_AI_RESPONSE = False
+
 
 def write_logs(prompt_messages, completion):
     # Create run_logs directory if it doesn't exist
@@ -41,14 +48,31 @@ async def stream_code_test(websocket: WebSocket):
 
     prompt_messages = assemble_prompt(params["image"])
 
-    completion = await stream_openai_response(
+    if SHOULD_MOCK_AI_RESPONSE:
-        prompt_messages,
+        completion = await mock_completion(process_chunk)
-        lambda x: process_chunk(x),
+    else:
-    )
+        completion = await stream_openai_response(
+            prompt_messages,
+            lambda x: process_chunk(x),
+        )
 
     # Write the messages dict into a log so that we can debug later
     write_logs(prompt_messages, completion)
 
-    await websocket.send_json({"type": "status", "value": "Code generation complete."})
+    # Generate images
+    await websocket.send_json({"type": "status", "value": "Generating images..."})
 
-    await websocket.close()
+    try:
+        updated_html = await generate_images(completion)
+        await websocket.send_json({"type": "setCode", "value": updated_html})
+        await websocket.send_json(
+            {"type": "status", "value": "Code generation complete."}
+        )
+    except Exception as e:
+        traceback.print_exc()
+        print("Image generation failed", e)
+        await websocket.send_json(
+            {"type": "status", "value": "Image generation failed but code is complete."}
+        )
+    finally:
+        await websocket.close()

backend/mock.py

@@ -0,0 +1,140 @@
+import asyncio
+
+
+async def mock_completion(process_chunk):
+    code_to_return = MOCK_HTML_2
+
+    for i in range(0, len(code_to_return), 10):
+        await process_chunk(code_to_return[i : i + 10])
+        await asyncio.sleep(0.01)
+
+    return code_to_return
+
+
+MOCK_HTML = """<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Product Showcase</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <link href="https://fonts.googleapis.com/css2?family=Roboto:wght@400;500;700&display=swap" rel="stylesheet">
+    <style>
+        body {
+            font-family: 'Roboto', sans-serif;
+        }
+    </style>
+</head>
+<body class="bg-black text-white">
+    <nav class="py-6">
+        <div class="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 flex justify-between items-center">
+            <div class="flex items-center">
+                <img src="https://placehold.co/24x24" alt="Company Logo" class="mr-8">
+                <a href="#" class="text-white text-sm font-medium mr-4">Store</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">Mac</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">iPad</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">iPhone</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">Watch</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">Vision</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">AirPods</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">TV & Home</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">Entertainment</a>
+                <a href="#" class="text-white text-sm font-medium mr-4">Accessories</a>
+                <a href="#" class="text-white text-sm font-medium">Support</a>
+            </div>
+            <div class="flex items-center">
+                <a href="#" class="text-white text-sm font-medium mr-4"><i class="fas fa-search"></i></a>
+                <a href="#" class="text-white text-sm font-medium"><i class="fas fa-shopping-bag"></i></a>
+            </div>
+        </div>
+    </nav>
+
+    <main class="mt-8">
+        <div class="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8">
+            <div class="text-center">
+                <img src="https://placehold.co/100x100" alt="Brand Logo" class="mx-auto mb-4">
+                <h1 class="text-5xl font-bold mb-4">WATCH SERIES 9</h1>
+                <p class="text-2xl font-medium mb-8">Smarter. Brighter. Mightier.</p>
+                <div class="flex justify-center space-x-4">
+                    <a href="#" class="text-blue-600 text-sm font-medium">Learn more ></a>
+                    <a href="#" class="text-blue-600 text-sm font-medium">Buy ></a>
+                </div>
+            </div>
+            <div class="flex justify-center mt-12">
+                <img src="https://placehold.co/500x300" alt="Product image of a smartwatch with a pink band and a circular interface displaying various health metrics." class="mr-8">
+                <img src="https://placehold.co/500x300" alt="Product image of a smartwatch with a blue band and a square interface showing a classic analog clock face." class="ml-8">
+            </div>
+        </div>
+    </main>
+</body>
+</html>"""
+
+MOCK_HTML_2 = """
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>The New York Times - News</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <link href="https://fonts.googleapis.com/css2?family=Libre+Franklin:wght@300;400;700&display=swap" rel="stylesheet">
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css">
+    <style>
+        body {
+            font-family: 'Libre Franklin', sans-serif;
+        }
+    </style>
+</head>
+<body class="bg-gray-100">
+    <div class="container mx-auto px-4">
+        <header class="border-b border-gray-300 py-4">
+            <div class="flex justify-between items-center">
+                <div class="flex items-center space-x-4">
+                    <button class="text-gray-700"><i class="fas fa-bars"></i></button>
+                    <button class="text-gray-700"><i class="fas fa-search"></i></button>
+                    <div class="text-xs uppercase tracking-widest">Tuesday, November 14, 2023<br>Today's Paper</div>
+                </div>
+                <div>
+                    <img src="https://placehold.co/200x50?text=The+New+York+Times+Logo" alt="The New York Times Logo" class="h-8">
+                </div>
+                <div class="flex items-center space-x-4">
+                    <button class="bg-black text-white px-4 py-1 text-xs uppercase tracking-widest">Give the times</button>
+                    <div class="text-xs">Account</div>
+                </div>
+            </div>
+            <nav class="flex justify-between items-center py-4">
+                <div class="flex space-x-4">
+                    <a href="#" class="text-xs uppercase tracking-widest text-gray-700">U.S.</a>
+                    <!-- Add other navigation links as needed -->
+                </div>
+                <div class="flex space-x-4">
+                    <a href="#" class="text-xs uppercase tracking-widest text-gray-700">Cooking</a>
+                    <!-- Add other navigation links as needed -->
+                </div>
+            </nav>
+        </header>
+        <main>
+            <section class="py-6">
+                <div class="grid grid-cols-3 gap-4">
+                    <div class="col-span-2">
+                        <article class="mb-4">
+                            <h2 class="text-xl font-bold mb-2">Israeli Military Raids Gaza’s Largest Hospital</h2>
+                            <p class="text-gray-700 mb-2">Israeli troops have entered the Al-Shifa Hospital complex, where conditions have grown dire and Israel says Hamas fighters are embedded.</p>
+                            <a href="#" class="text-blue-600 text-sm">See more updates <i class="fas fa-external-link-alt"></i></a>
+                        </article>
+                        <!-- Repeat for each news item -->
+                    </div>
+                    <div class="col-span-1">
+                        <article class="mb-4">
+                            <img src="https://placehold.co/300x200?text=News+Image" alt="Flares and plumes of smoke over the northern Gaza skyline on Tuesday." class="mb-2">
+                            <h2 class="text-xl font-bold mb-2">From Elvis to Elopements, the Evolution of the Las Vegas Wedding</h2>
+                            <p class="text-gray-700 mb-2">The glittering city that attracts thousands of couples seeking unconventional nuptials has grown beyond the drive-through wedding.</p>
+                            <a href="#" class="text-blue-600 text-sm">8 MIN READ</a>
+                        </article>
+                        <!-- Repeat for each news item -->
+                    </div>
+                </div>
+            </section>
+        </main>
+    </div>
+</body>
+</html>
+"""

backend/poetry.lock

@@ -22,6 +22,25 @@ doc = ["Sphinx", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-
 test = ["anyio[trio]", "coverage[toml] (>=4.5)", "hypothesis (>=4.0)", "mock (>=4)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
 trio = ["trio (<0.22)"]
 
+[[package]]
+name = "beautifulsoup4"
+version = "4.12.2"
+description = "Screen-scraping library"
+category = "main"
+optional = false
+python-versions = ">=3.6.0"
+files = [
+    {file = "beautifulsoup4-4.12.2-py3-none-any.whl", hash = "sha256:bd2520ca0d9d7d12694a53d44ac482d181b4ec1888909b035a3dbf40d0f57d4a"},
+    {file = "beautifulsoup4-4.12.2.tar.gz", hash = "sha256:492bbc69dca35d12daac71c4db1bfff0c876c00ef4a2ffacce226d4638eb72da"},
+]
+
+[package.dependencies]
+soupsieve = ">1.2"
+
+[package.extras]
+html5lib = ["html5lib"]
+lxml = ["lxml"]
+
 [[package]]
 name = "certifi"
 version = "2023.7.22"
@@ -284,6 +303,18 @@ files = [
     {file = "sniffio-1.3.0.tar.gz", hash = "sha256:e60305c5e5d314f5389259b7f22aaa33d8f7dee49763119234af3755c55b9101"},
 ]
 
+[[package]]
+name = "soupsieve"
+version = "2.5"
+description = "A modern CSS selector implementation for Beautiful Soup."
+category = "main"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "soupsieve-2.5-py3-none-any.whl", hash = "sha256:eaa337ff55a1579b6549dc679565eac1e3d000563bcb1c8ab0d0fefbc0c2cdc7"},
+    {file = "soupsieve-2.5.tar.gz", hash = "sha256:5663d5a7b3bfaeee0bc4372e7fc48f9cff4940b3eec54a6451cc5299f1097690"},
+]
+
 [[package]]
 name = "starlette"
 version = "0.27.0"
@@ -440,4 +471,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "37bf71ae4f77aaeda11cbb524e3999464fceb20706135680a4c77add8712a847"
+content-hash = "5e4aa03dda279f66a9b3d30f7327109bcfd395795470d95f8c563897ce1bff84"

backend/pyproject.toml

@@ -13,6 +13,7 @@ uvicorn = "^0.24.0.post1"
 websockets = "^12.0"
 openai = "^1.2.4"
 python-dotenv = "^1.0.0"
+beautifulsoup4 = "^4.12.2"
 
 [build-system]
 requires = ["poetry-core"]

frontend/src/App.tsx

@@ -12,7 +12,7 @@ function App() {
   );
   const [generatedCode, setGeneratedCode] = useState<string>("");
   const [referenceImages, setReferenceImages] = useState<string[]>([]);
-  const [console, setConsole] = useState<string[]>([]);
+  const [executionConsole, setExecutionConsole] = useState<string[]>([]);
   const [blobUrl, setBlobUrl] = useState("");
 
   const createBlobUrl = () => {
@@ -29,8 +29,11 @@ function App() {
       function (token) {
         setGeneratedCode((prev) => prev + token);
       },
+      function (code) {
+        setGeneratedCode(code);
+      },
       function (line) {
-        setConsole((prev) => [...prev, line]);
+        setExecutionConsole((prev) => [...prev, line]);
       },
       function () {
         setAppState("CODE_READY");
@@ -67,7 +70,7 @@ function App() {
             </div>
             <div className="bg-gray-400 px-4 py-2 rounded text-sm hidden">
               <h2 className="text-lg mb-4 border-b border-gray-800">Console</h2>
-              {console.map((line, index) => (
+              {executionConsole.map((line, index) => (
                 <div
                   key={index}
                   className="border-b border-gray-400 mb-2 text-gray-600 font-mono"
@@ -82,7 +85,7 @@ function App() {
             <>
               <div className="flex items-center gap-x-1">
                 <Spinner />
-                Generating...
+                {executionConsole.slice(-1)[0]}
               </div>
               <CodePreview code={generatedCode} />
             </>

frontend/src/generateCode.ts

@@ -7,6 +7,7 @@ const ERROR_MESSAGE =
 export function generateCode(
   imageUrl: string,
   onChange: (chunk: string) => void,
+  onSetCode: (code: string) => void,
   onStatusUpdate: (status: string) => void,
   onComplete: () => void
 ) {
@@ -29,6 +30,8 @@ export function generateCode(
       onChange(response.value);
     } else if (response.type === "status") {
       onStatusUpdate(response.value);
+    } else if (response.type === "setCode") {
+      onSetCode(response.value);
     }
   });