add image generation for placeholder images
This commit is contained in:
Abi Raja
parent
3f071e7a68
commit
51c7334c0e
83
backend/image_generation.py
Normal file
83
backend/image_generation.py
Normal file
@ -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)
|
||||
@ -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,6 +48,9 @@ async def stream_code_test(websocket: WebSocket):
|
||||
|
||||
prompt_messages = assemble_prompt(params["image"])
|
||||
|
||||
if SHOULD_MOCK_AI_RESPONSE:
|
||||
completion = await mock_completion(process_chunk)
|
||||
else:
|
||||
completion = await stream_openai_response(
|
||||
prompt_messages,
|
||||
lambda x: process_chunk(x),
|
||||
@ -49,6 +59,20 @@ async def stream_code_test(websocket: WebSocket):
|
||||
# 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..."})
|
||||
|
||||
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()
|
||||
|
||||
140
backend/mock.py
Normal file
140
backend/mock.py
Normal file
@ -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>
|
||||
"""
|
||||
33
backend/poetry.lock
generated
33
backend/poetry.lock
generated
@ -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"
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -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} />
|
||||
</>
|
||||
|
||||
@ -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);
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user