diff --git a/README.md b/README.md index 5c67c4b..177bb74 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ Need to deploy your Worker to Cloudflare? Python Workers are in open beta and ha - [**`dynamic-py-py/`**](dynamic-py-py) — shows how to load and run a Python Worker dynamically at runtime using a [Worker Loader](https://developers.cloudflare.com/workers/runtime-apis/bindings/worker-loader/) binding. - [**`django/`**](django) — runs a naive Django WSGI application directly on Python Workers. - [**`django-todo-d1/`**](django-todo-d1) — uses Django with D1 for a basic TODO application. - +- [**`image-redraw/`**](image-redraw) — an example that combines [FastAPI](https://fastapi.tiangolo.com/), [R2](https://developers.cloudflare.com/r2/), [Queues](https://developers.cloudflare.com/queues/), [Workflows](https://developers.cloudflare.com/workflows/) and [Workers AI](https://developers.cloudflare.com/workers-ai/) to redraw uploaded images. ## Open Beta and Limits diff --git a/image-redraw/README.md b/image-redraw/README.md new file mode 100644 index 0000000..333a223 --- /dev/null +++ b/image-redraw/README.md @@ -0,0 +1,74 @@ +# Image Redraw — FastAPI + R2 + Queues + Workflows + Workers AI + +[![Deploy to Cloudflare](https://deploy.workers.cloudflare.com/button)](https://deploy.workers.cloudflare.com/?url=https://github.com/cloudflare/python-workers-examples/tree/main/20-image-redraw) + +Upload a picture and get it back redrawn as a wobbly MS Paint doodle. + +This example demonstrates how to leverage multiple Cloudflare services to +create a full-featured image processing pipeline using Python Workers. + +## What it uses + +- [Python Workers](https://developers.cloudflare.com/workers/languages/python/) — the runtime +- [FastAPI](https://fastapi.tiangolo.com/) — the HTTP API, served over ASGI +- [R2](https://developers.cloudflare.com/r2/) — stores uploaded originals and redrawn outputs +- [Queues](https://developers.cloudflare.com/queues/) — hands work off the request path +- [Workflows](https://developers.cloudflare.com/workflows/) — durable, retrying background execution +- [Workers AI](https://developers.cloudflare.com/workers-ai/) — image generation guided by the uploaded reference +- [Pillow](https://pillow.readthedocs.io/en/stable/) — image normalization inside the Workflow +- [Static Assets](https://developers.cloudflare.com/workers/static-assets/) — the plain HTML/CSS/JS frontend + +```mermaid +flowchart LR + Browser["Browser"] + Worker["Python Worker
FastAPI"] + Queue[["Queue"]] + Workflow["Workflow"] + R2[("R2")] + AI["Workers AI"] + + Browser --> Worker + Worker --> R2 + Worker --> Queue + Queue --> Workflow + Workflow --> R2 + Workflow --> AI +``` + +## The flow + +1. The browser POSTs the image bytes to the Worker. +2. FastAPI validates the image and stores the original in R2, and + enqueues the job. +3. The queue consumer turns each batch of IDs into Workflow instances. +4. The Workflow reads the original from R2, normalizes it with Pillow, calls + Workers AI, and stores the result back in R2. + + +## Setup + +First ensure that `uv` is installed: +https://docs.astral.sh/uv/getting-started/installation/#standalone-installer + +**Workers AI is a remote binding, even during local development.** `wrangler.jsonc` +declares `"ai": { "binding": "AI", "remote": true }`, so inference always runs on +Cloudflare's network and bills against your account. Log in before starting the +dev server: + +```sh +uv run pwrangler login +``` + +## How to Run + +```sh +uv run pywrangler dev +``` + +Then open http://localhost:8787/ in your browser. + +## How to deploy + +```sh +uv run pywrangler deploy +``` diff --git a/image-redraw/package.json b/image-redraw/package.json new file mode 100644 index 0000000..2a460ff --- /dev/null +++ b/image-redraw/package.json @@ -0,0 +1,13 @@ +{ + "name": "image-redraw-worker", + "version": "0.0.0", + "private": true, + "scripts": { + "deploy": "uv run pywrangler deploy", + "dev": "uv run pywrangler dev", + "start": "uv run pywrangler dev" + }, + "devDependencies": { + "wrangler": "^4.114.0" + } +} diff --git a/image-redraw/public/app.js b/image-redraw/public/app.js new file mode 100644 index 0000000..931bdaa --- /dev/null +++ b/image-redraw/public/app.js @@ -0,0 +1,279 @@ +const REFERENCE_SIZE = 511; +const MAX_UPLOAD_BYTES = 5_000_000; +const POLL_INTERVAL_MS = 2000; +const MAX_POLLS = 150; +const MAX_POLL_FAILURES = 3; +const TYPE_LABELS = { + "image/png": "PNG", + "image/jpeg": "JPEG", + "image/webp": "WebP", +}; +const ACCEPTED_TYPES = Object.keys(TYPE_LABELS); + +const byId = (id) => document.getElementById(id); +const uploadForm = byId("upload-form"); +const fileInput = byId("file-input"); +const submitButton = byId("submit-button"); +const uploadStatus = byId("upload-status"); +const preview = byId("preview"); +const previewImage = byId("preview-image"); +const previewCaption = byId("preview-caption"); +const compare = byId("compare"); +const compareEmpty = byId("compare-empty"); +const compareMeta = byId("compare-meta"); +const originalImage = byId("original-image"); +const outputImage = byId("output-image"); +const gallery = byId("gallery"); +const galleryStatus = byId("gallery-status"); +const refreshButton = byId("refresh-button"); + +let rawUpload = null; +let previewUrl = null; +let selectedJobId = null; + +const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); +const shortId = (jobId) => jobId.slice(0, 8); + +function setStatus(element, message, tone = "info") { + element.textContent = message; + element.dataset.tone = tone; +} + +function formatTime(isoString) { + const date = new Date(isoString); + if (Number.isNaN(date.getTime())) return "unknown time"; + return date.toLocaleString([], { + month: "short", + day: "numeric", + hour: "2-digit", + minute: "2-digit", + }); +} + +function formatBytes(bytes) { + const kilobytes = bytes / 1024; + if (kilobytes < 1000) return `${kilobytes.toFixed(1)} KB`; + return `${(kilobytes / 1024).toFixed(2)} MB`; +} + +async function requestJson(path, options) { + let response; + try { + response = await fetch(path, options); + } catch { + throw new Error("Network error. Is the Worker still running?"); + } + + const data = await response.json().catch(() => null); + if (!response.ok) { + const detail = typeof data?.detail === "string" ? data.detail : null; + throw new Error(detail ?? `Request failed with status ${response.status}`); + } + return data; +} + +// Decodes only to prove the bytes are an image and to read the source size. +// The bitmap is discarded; the original File is what gets uploaded. +async function readSourceSize(file) { + let bitmap; + try { + bitmap = await createImageBitmap(file); + } catch { + throw new Error("That file could not be decoded as an image."); + } + const size = { width: bitmap.width, height: bitmap.height }; + bitmap.close?.(); + return size; +} + +async function inspectFile(file) { + if (!ACCEPTED_TYPES.includes(file.type)) { + throw new Error("Only PNG, JPEG and WebP are supported."); + } + if (file.size > MAX_UPLOAD_BYTES) { + throw new Error( + `That file is ${file.size.toLocaleString()} bytes, over the ${MAX_UPLOAD_BYTES.toLocaleString()} byte limit. Try a smaller picture.`, + ); + } + + const { width, height } = await readSourceSize(file); + const label = TYPE_LABELS[file.type]; + return { + file, + caption: + `${width}x${height} ${label} (${file.type}), ${formatBytes(file.size)}`, + }; +} + +async function showPreview() { + rawUpload = null; + preview.hidden = true; + if (previewUrl) URL.revokeObjectURL(previewUrl); + previewUrl = null; + + const file = fileInput.files?.[0]; + if (!file) { + setStatus(uploadStatus, ""); + return; + } + + setStatus(uploadStatus, "Checking your picture..."); + try { + rawUpload = await inspectFile(file); + previewUrl = URL.createObjectURL(rawUpload.file); + previewImage.src = previewUrl; + previewImage.alt = + "The original picture you chose, shown at full size before it is uploaded."; + previewCaption.textContent = rawUpload.caption; + preview.hidden = false; + setStatus(uploadStatus, 'Ready. Press "Redraw it!".'); + } catch (error) { + setStatus(uploadStatus, error.message, "error"); + } +} + +// Resolves with the terminal job object so the caller can read job.reason. +async function pollJob(jobId) { + let failures = 0; + + for (let attempt = 1; attempt <= MAX_POLLS; attempt += 1) { + await sleep(POLL_INTERVAL_MS); + + let job; + try { + job = await requestJson(`/api/jobs/${jobId}`); + failures = 0; + } catch (error) { + failures += 1; + if (failures >= MAX_POLL_FAILURES) { + throw new Error(`Lost contact with the Worker. ${error.message}`); + } + continue; + } + + if (job.status === "complete" || job.status === "failed") return job; + setStatus( + uploadStatus, + `Job ${shortId(jobId)} is ${job.status}... (checked ${attempt} times)`, + ); + } + + throw new Error("This redraw is taking too long. Try Refresh later."); +} + +function selectJob(job) { + selectedJobId = job.jobId; + originalImage.src = job.originalUrl; + originalImage.alt = `The picture you uploaded for job ${shortId(job.jobId)}.`; + outputImage.src = job.outputUrl; + outputImage.alt = `Workers AI redraw of job ${shortId(job.jobId)}, in a clumsy MS Paint style.`; + compareMeta.textContent = `Job ${job.jobId} - finished ${formatTime(job.completedAt)}`; + compare.hidden = false; + compareEmpty.hidden = true; + + for (const card of gallery.querySelectorAll(".card")) { + card.setAttribute("aria-pressed", String(card.dataset.jobId === job.jobId)); + } +} + +function createCard(job) { + const thumb = document.createElement("img"); + thumb.src = job.outputUrl; + thumb.loading = "lazy"; + thumb.alt = `Redrawn picture from job ${shortId(job.jobId)}`; + + const time = document.createElement("span"); + time.textContent = formatTime(job.completedAt); + + const card = document.createElement("button"); + card.type = "button"; + card.className = "card"; + card.dataset.jobId = job.jobId; + card.setAttribute("aria-pressed", String(job.jobId === selectedJobId)); + card.append(thumb, time); + card.addEventListener("click", () => selectJob(job)); + + const item = document.createElement("li"); + item.append(card); + return item; +} + +function renderGallery(jobs) { + if (jobs.length === 0) { + const empty = document.createElement("li"); + empty.className = "empty"; + empty.textContent = "No pictures yet. Upload one to start the gallery."; + gallery.replaceChildren(empty); + return; + } + gallery.replaceChildren(...jobs.map(createCard)); +} + +async function loadGallery(jobIdToSelect) { + refreshButton.disabled = true; + setStatus(galleryStatus, "Loading gallery..."); + + try { + const jobs = (await requestJson("/api/jobs"))?.jobs ?? []; + renderGallery(jobs); + setStatus( + galleryStatus, + jobs.length === 1 ? "1 picture saved." : `${jobs.length} pictures saved.`, + ); + + const target = jobs.find((job) => job.jobId === jobIdToSelect); + if (target) selectJob(target); + } catch (error) { + setStatus(galleryStatus, error.message, "error"); + } finally { + refreshButton.disabled = false; + } +} + +uploadForm.addEventListener("submit", async (event) => { + event.preventDefault(); + if (!rawUpload) await showPreview(); + if (!rawUpload) { + if (!fileInput.files?.length) { + setStatus(uploadStatus, "Choose a picture first.", "error"); + } + fileInput.focus(); + return; + } + + const upload = rawUpload.file; + fileInput.disabled = true; + submitButton.disabled = true; + try { + setStatus(uploadStatus, "Uploading to the Worker..."); + const created = await requestJson("/api/jobs", { + method: "POST", + headers: { "Content-Type": upload.type }, + body: upload, + }); + + setStatus(uploadStatus, `Job ${shortId(created.jobId)} is queued...`); + const job = await pollJob(created.jobId); + if (job.status === "complete") { + setStatus(uploadStatus, "Finished! Behold the artwork.", "done"); + await loadGallery(job.jobId); + } else { + // The backend only ever sends a reason it is happy to show a visitor. + const reason = + typeof job.reason === "string" && job.reason + ? job.reason + : "The Workflow gave up on that one."; + setStatus(uploadStatus, reason, "error"); + } + } catch (error) { + setStatus(uploadStatus, error.message, "error"); + } finally { + fileInput.disabled = false; + submitButton.disabled = false; + } +}); + +fileInput.addEventListener("change", () => void showPreview()); +refreshButton.addEventListener("click", () => void loadGallery()); + +void loadGallery(); diff --git a/image-redraw/public/index.html b/image-redraw/public/index.html new file mode 100644 index 0000000..356c94e --- /dev/null +++ b/image-redraw/public/index.html @@ -0,0 +1,92 @@ + + + + + + Image Redraw - Python Workers + Workers AI + + + + +
+
+

Image Redraw

+

+ Upload a photo. A Python Worker and Workers AI redraw it + with clunky confidence. +

+
+ +
+
+
+

untitled - Paint

+ +
+
+
+ + +

+ PNG, JPEG or WebP. +

+ +
+ +

+ + +
+
+ +
+
+

before-and-after.bmp

+ +
+
+

+ Nothing open yet. Redraw a picture, or pick one from My Pictures. +

+ +

+
+
+
+ +
+
+ + +
+
+
+

Finished redraws, newest first.

+ +
+ + +
+
+
+ + + + diff --git a/image-redraw/public/style.css b/image-redraw/public/style.css new file mode 100644 index 0000000..acf1100 --- /dev/null +++ b/image-redraw/public/style.css @@ -0,0 +1,275 @@ +/* A deliberately clumsy MS Paint / Windows 95 look, in plain CSS. */ +:root { + --face: #c0c0c0; + --face-light: #dfdfdf; + --face-shadow: #808080; + --white: #ffffff; + --ink: #000000; + --ink-muted: #3a3a3a; + --desktop: #008080; + --title: #000080; + --yellow: #ffff00; + --red: #d40000; + --green: #007700; + --space-1: 0.25rem; + --space-2: 0.5rem; + --space-3: 1rem; + --space-4: 1.5rem; + --font-display: "Comic Sans MS", "Chalkboard SE", cursive; + --font-ui: Tahoma, Verdana, sans-serif; + --font-mono: "Courier New", Courier, monospace; + --bevel-out: inset -2px -2px 0 var(--face-shadow), + inset 2px 2px 0 var(--face-light), inset -3px -3px 0 var(--ink); + --bevel-in: inset 2px 2px 0 var(--face-shadow), + inset -2px -2px 0 var(--face-light), inset 3px 3px 0 var(--ink); +} + +* { + box-sizing: border-box; +} +body { + margin: 0; + background: var(--desktop); + color: var(--ink); + font: 0.9375rem/1.5 var(--font-ui); +} +img { + display: block; + max-width: 100%; +} +[hidden] { + display: none !important; +} +:focus-visible { + outline: 3px dotted var(--ink); + outline-offset: 2px; +} + +/* Desktop and header */ +.desktop { + max-width: 58rem; + margin: 0 auto; + padding: var(--space-4) var(--space-3); +} +.page-header { + text-align: center; + margin-bottom: var(--space-4); +} +.page-header h1 { + margin: 0; + font: 2.25rem var(--font-display); + color: var(--white); + text-shadow: 3px 3px 0 var(--ink); + transform: rotate(-1deg); +} +.page-header p { + max-width: 34rem; + margin: var(--space-3) auto 0; + padding: var(--space-2); + background: var(--white); + border: 2px solid var(--ink); + font-size: 0.8125rem; + transform: rotate(0.8deg); +} + +/* Windows */ +.columns { + display: grid; + grid-template-columns: repeat(auto-fit, minmax(19rem, 1fr)); + gap: var(--space-4); +} +.window { + background: var(--face); + border: 2px solid var(--ink); + box-shadow: 6px 6px 0 rgba(0, 0, 0, 0.4); + padding: 3px; + margin-bottom: var(--space-4); + transform: rotate(-0.8deg); +} +.titlebar { + display: flex; + align-items: center; + justify-content: space-between; + padding: var(--space-1) var(--space-2); + background: var(--title); +} +.titlebar h2 { + margin: 0; + font-size: 0.8125rem; + color: var(--white); +} +.titlebar-buttons { + padding: 0 var(--space-1); + background: var(--face); + box-shadow: var(--bevel-out); + font-size: 0.6875rem; + letter-spacing: 0.3em; +} +.window-body { + padding: var(--space-3); +} + +/* Form controls */ +form { + display: grid; + gap: var(--space-2); +} +.field-label { + font: 1.125rem var(--font-display); +} +.hint { + margin: 0; + font-size: 0.8125rem; + color: var(--ink-muted); +} +input[type="file"] { + padding: var(--space-2); + background: var(--white); + border: 0; + box-shadow: var(--bevel-in); + font: 0.8125rem var(--font-mono); +} +.btn { + justify-self: start; + padding: var(--space-2) var(--space-3); + border: 0; + background: var(--face); + box-shadow: var(--bevel-out); + font: 0.9375rem var(--font-ui); + color: var(--ink); + cursor: pointer; +} +.btn:hover:not(:disabled), +.card:hover { + background: var(--face-light); +} +.btn:active:not(:disabled) { + box-shadow: var(--bevel-in); +} +.btn:disabled { + color: var(--face-shadow); + cursor: not-allowed; +} +.btn-go { + background: var(--yellow); + font: 1.5rem var(--font-display); + transform: rotate(0.9deg); +} + +/* Status, preview and comparison */ +.status { + min-height: 1.5rem; + margin: var(--space-2) 0 0; + font: 0.8125rem var(--font-mono); +} +.status[data-tone="error"], +.status[data-tone="done"] { + padding: var(--space-2); + background: var(--white); + border: 2px solid currentcolor; + font-weight: bold; +} +.status[data-tone="error"] { + color: var(--red); +} +.status[data-tone="done"] { + color: var(--green); +} +figure { + margin: var(--space-3) 0 0; +} +figure img { + width: 100%; + aspect-ratio: 1; + object-fit: contain; + background: var(--white); + box-shadow: var(--bevel-in); +} +figcaption, +.meta, +.card { + font: 0.6875rem var(--font-mono); + color: var(--ink-muted); +} +.compare { + display: grid; + grid-template-columns: 1fr 1fr; + gap: var(--space-3); +} +.meta { + margin: var(--space-3) 0 0; + word-break: break-all; +} +.empty { + margin: 0; + padding: var(--space-4); + background: var(--white); + border: 3px dashed var(--face-shadow); + font-family: var(--font-display); + text-align: center; + color: var(--ink-muted); +} + +/* Gallery */ +.gallery-header { + display: flex; + align-items: center; + justify-content: space-between; + gap: var(--space-3); + flex-wrap: wrap; +} +#gallery { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(8.5rem, 1fr)); + gap: var(--space-3); + list-style: none; + margin: var(--space-2) 0 0; + padding: 0; +} +#gallery .empty { + grid-column: 1 / -1; +} +.card { + display: grid; + gap: var(--space-1); + padding: var(--space-2); + border: 0; + background: var(--face); + box-shadow: var(--bevel-out); + cursor: pointer; +} +.card[aria-pressed="true"] { + background: var(--yellow); + box-shadow: var(--bevel-in); + color: var(--ink); + font-weight: bold; + transform: rotate(-1.1deg); +} +.card img { + aspect-ratio: 1; + object-fit: contain; + background: var(--white); + border: 2px solid var(--ink); + image-rendering: pixelated; /* crunchy on purpose */ +} + +/* Small screens: straighten up so nothing runs off the edge */ +@media (max-width: 40rem) { + .window, + .btn-go, + .page-header h1, + .page-header p, + .card[aria-pressed="true"] { + transform: none; + } + .compare { + grid-template-columns: 1fr; + } +} +@media (prefers-reduced-motion: reduce) { + * { + animation-duration: 0.01ms !important; + transition-duration: 0.01ms !important; + scroll-behavior: auto !important; + } +} diff --git a/image-redraw/pyproject.toml b/image-redraw/pyproject.toml new file mode 100644 index 0000000..080ef12 --- /dev/null +++ b/image-redraw/pyproject.toml @@ -0,0 +1,16 @@ +[project] +name = "image-redraw-worker" +version = "0.1.0" +description = "Image redraw example using R2, Queues, Workflows and Workers AI" +readme = "README.md" +requires-python = ">=3.13" +dependencies = [ + "fastapi", + "pillow", +] + +[dependency-groups] +dev = [ + "workers-py", + "workers-runtime-sdk", +] diff --git a/image-redraw/src/entry.py b/image-redraw/src/entry.py new file mode 100644 index 0000000..86710d6 --- /dev/null +++ b/image-redraw/src/entry.py @@ -0,0 +1,27 @@ +from image_redraw.api import app +from image_redraw.constants import is_job_id +from image_redraw.workflow import RedrawWorkflow +from workers import WorkerEntrypoint, asgi + +# Re-export RedrawWorkflow so workerd can find workflow classes +__all__ = ["Default", "RedrawWorkflow"] + + +class Default(WorkerEntrypoint): + async def fetch(self, request): + return await asgi.fetch(app, request, self.env) + + async def queue(self, batch, env, ctx): + pending = [] + for message in batch.messages: + body = message.body + job_id = body.get("jobId") if isinstance(body, dict) else None + if not is_job_id(job_id): + print(f"Skipping malformed queue message {message.id}") + message.ack() + continue + pending.append((message, {"id": job_id, "params": {"jobId": job_id}})) + if pending: + await self.env.REDRAW_WORKFLOW.create_batch([spec for _, spec in pending]) + for message, _ in pending: + message.ack() diff --git a/image-redraw/src/image_redraw/__init__.py b/image-redraw/src/image_redraw/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/image-redraw/src/image_redraw/api.py b/image-redraw/src/image_redraw/api.py new file mode 100644 index 0000000..7faa49a --- /dev/null +++ b/image-redraw/src/image_redraw/api.py @@ -0,0 +1,153 @@ +import uuid +from datetime import UTC, datetime +from typing import Any + +from fastapi import FastAPI, HTTPException, Request, Response + +from .constants import ( + ALLOWED_CONTENT_TYPES, + GALLERY_SIZE, + LIST_PAGE_SIZE, + OUTPUT_PREFIX, + STATUS_MAP, + failure_key, + is_job_id, + job_id_from_output_key, + original_key, + output_key, +) + +app = FastAPI() + +IMAGE_HEADERS = {"X-Content-Type-Options": "nosniff"} + + +def declared_length(request: Request) -> int | None: + try: + return int(request.headers["content-length"]) + except (KeyError, ValueError): + # A missing or malformed header falls through to the post-read check. + return None + + +@app.post("/api/jobs", status_code=202) +async def create_job(request: Request) -> dict[str, str]: + content_type = request.headers.get("content-type", "") + if not content_type.startswith(ALLOWED_CONTENT_TYPES): + raise HTTPException(415, "Send a raw image/png, image/jpeg or image/webp body.") + + image = await request.body() + if not image: + raise HTTPException(400, "The request body is empty.") + + env = request.scope["env"] + job_id = uuid.uuid4().hex + created_at = datetime.now(UTC).isoformat() + + await env.REDRAW_BUCKET.put( + original_key(job_id), + image, + httpMetadata={"contentType": content_type}, + customMetadata={"createdAt": created_at}, + ) + + try: + await env.REDRAW_QUEUE.send({"jobId": job_id}) + except Exception: + await env.REDRAW_BUCKET.delete(original_key(job_id)) + raise + return {"jobId": job_id, "status": "queued", "createdAt": created_at} + + +@app.get("/api/jobs") +async def list_jobs(request: Request) -> dict[str, list[dict[str, str]]]: + env = request.scope["env"] + bucket = env.REDRAW_BUCKET + jobs = [] + cursor = None + + while True: + options: dict[str, Any] = {"prefix": OUTPUT_PREFIX, "limit": LIST_PAGE_SIZE} + if cursor: + options["cursor"] = cursor + + listed = await bucket.list(**options) + for obj in listed["objects"]: + job_id = job_id_from_output_key(obj.key) + if not is_job_id(job_id): + continue + jobs.append( + { + "jobId": job_id, + "completedAt": obj.uploaded.toISOString(), + "originalUrl": f"/api/images/original/{job_id}", + "outputUrl": f"/api/images/output/{job_id}", + } + ) + + # R2 only returns a cursor while more pages remain. + cursor = listed["cursor"] if listed["truncated"] else None + if not cursor: + break + + jobs.sort(key=lambda job: job["completedAt"], reverse=True) + return {"jobs": jobs[:GALLERY_SIZE]} + + +@app.get("/api/jobs/{job_id}") +async def get_job(job_id: str, request: Request) -> dict[str, str]: + if not is_job_id(job_id): + raise HTTPException(404, "Job not found.") + + env = request.scope["env"] + bucket = env.REDRAW_BUCKET + + if await bucket.head(output_key(job_id)) is not None: + return { + "jobId": job_id, + "status": "complete", + "outputUrl": f"/api/images/output/{job_id}", + } + + failure = await bucket.get(failure_key(job_id)) + if failure is not None: + blob = await failure.blob() + return {"jobId": job_id, "status": "failed", "reason": await blob.text()} + + if await bucket.head(original_key(job_id)) is None: + raise HTTPException(404, "Job not found.") + + try: + instance = await env.REDRAW_WORKFLOW.get(job_id) + status = await instance.status() + except Exception: + # The original exists, so the consumer just has not created the instance yet. + return {"jobId": job_id, "status": "queued"} + return {"jobId": job_id, "status": STATUS_MAP.get(status["status"], "running")} + + +@app.get("/api/images/{kind}/{job_id}") +async def get_image(kind: str, job_id: str, request: Request) -> Response: + if not is_job_id(job_id): + raise HTTPException(404, "Image not found.") + + if kind == "original": + key = original_key(job_id) + elif kind == "output": + key = output_key(job_id) + else: + raise HTTPException(404, "Image kind must be 'original' or 'output'.") + + env = request.scope["env"] + obj = await env.REDRAW_BUCKET.get(key) + if obj is None: + raise HTTPException(404, "Image not found.") + + http_metadata = obj.httpMetadata + media_type = http_metadata.contentType if http_metadata is not None else None + blob = await obj.blob() + return Response( + content=await blob.bytes(), + media_type=media_type or "application/octet-stream", + headers=IMAGE_HEADERS, + ) diff --git a/image-redraw/src/image_redraw/constants.py b/image-redraw/src/image_redraw/constants.py new file mode 100644 index 0000000..cd6fda0 --- /dev/null +++ b/image-redraw/src/image_redraw/constants.py @@ -0,0 +1,83 @@ +import re + +MODEL = "@cf/black-forest-labs/flux-2-klein-4b" +# FLUX takes multipart form fields, so every option is sent as a string. +AI_OPTIONS = { + "prompt": ( + "redraw the scene in the reference image as a clumsy MS Paint drawing, " + "keeping the same subject and composition but with wobbly mouse-drawn " + "outlines, flat bucket-fill colors, jagged pixelated edges, childlike " + "and amateur" + ), + "guidance": "2.5", + "width": "512", + "height": "512", +} +AI_RETRIES = {"retries": {"limit": 5, "delay": "5 seconds", "backoff": "exponential"}} + +AI_ERROR_CODE_PATTERN = re.compile(r"(\d+)\s*:") +# 3030 means the safety filter refused the generated image. This app +# deliberately treats it as final rather than retrying, to avoid repeated +# billable inference on a reference the filter is likely to refuse again. +AI_SAFETY_ERROR_CODE = 3030 + +ORIGINAL_PREFIX = "originals/" +OUTPUT_PREFIX = "outputs/" +FAILURE_PREFIX = "failures/" + +ALLOWED_CONTENT_TYPES = ("image/png", "image/jpeg", "image/webp") + +# A deliberately small, fixed canvas for this example: it keeps inference cheap +# and every reference picture uniform. It is a choice made here, not a limit +# documented by the model. +TARGET_SIZE = (511, 511) +CANVAS_COLOR = (255, 255, 255) +JPEG_QUALITY = 82 +MAX_SOURCE_PIXELS = {"JPEG": 40_000_000, "PNG": 8_000_000, "WEBP": 4_000_000} + +# The gallery shows the newest finished redraws; R2 lists keys in lexicographic +# order, so every page has to be read before the newest ones can be picked. +GALLERY_SIZE = 20 +LIST_PAGE_SIZE = 100 + +JOB_ID_PATTERN = re.compile(r"[0-9a-f]{32}") +STATUS_MAP = { + "queued": "queued", + "complete": "complete", + "errored": "failed", + "terminated": "failed", +} +SAFETY_REJECTED_REASON = ( + "Workers AI refused to redraw that picture. Try a different one." +) +INVALID_IMAGE_REASON = "That upload could not be decoded as a usable image." +MISSING_ORIGINAL_REASON = "The uploaded picture is no longer available." +INVALID_OUTPUT_REASON = "Workers AI returned something that was not a picture." + + +def is_job_id(value: object) -> bool: + # Job IDs are uuid4().hex, and a malformed one would fail the whole create_batch. + return isinstance(value, str) and JOB_ID_PATTERN.fullmatch(value) is not None + + +def original_key(job_id: str) -> str: + return f"{ORIGINAL_PREFIX}{job_id}" + + +def output_key(job_id: str) -> str: + return f"{OUTPUT_PREFIX}{job_id}" + + +def failure_key(job_id: str) -> str: + return f"{FAILURE_PREFIX}{job_id}" + + +def job_id_from_output_key(key: str) -> str: + return key.removeprefix(OUTPUT_PREFIX) + + +def ai_error_code(message: str | None) -> int | None: + if not message: + return None + match = AI_ERROR_CODE_PATTERN.match(message.strip()) + return int(match.group(1)) if match else None diff --git a/image-redraw/src/image_redraw/workflow.py b/image-redraw/src/image_redraw/workflow.py new file mode 100644 index 0000000..71d9626 --- /dev/null +++ b/image-redraw/src/image_redraw/workflow.py @@ -0,0 +1,168 @@ +import base64 +import io +import json + +from PIL import Image, ImageOps +from workers import Blob, FormData, Response, WorkflowEntrypoint +from workers.workflows import NonRetryableError + +from .constants import ( + AI_OPTIONS, + AI_RETRIES, + AI_SAFETY_ERROR_CODE, + CANVAS_COLOR, + INVALID_IMAGE_REASON, + INVALID_OUTPUT_REASON, + JPEG_QUALITY, + MAX_SOURCE_PIXELS, + MISSING_ORIGINAL_REASON, + MODEL, + SAFETY_REJECTED_REASON, + TARGET_SIZE, + ai_error_code, + failure_key, + original_key, + output_key, +) + +# Magic bytes, because the model tells us nothing about the format it picked. +IMAGE_SIGNATURES = ( + (b"\x89PNG", "image/png"), + (b"\xff\xd8\xff", "image/jpeg"), + (b"RIFF", "image/webp"), +) + + +class UnusableImageError(Exception): + pass + + +def sniff_content_type(image_bytes: bytes) -> str | None: + for signature, content_type in IMAGE_SIGNATURES: + if image_bytes.startswith(signature): + return content_type + return None + + +def resize(image_bytes: bytes) -> bytes: + try: + with Image.open(io.BytesIO(image_bytes)) as image: + source_format = image.format or "" + pixel_limit = MAX_SOURCE_PIXELS.get(source_format) + if pixel_limit is None: + raise UnusableImageError(f"Unsupported image format {source_format!r}.") + + width, height = image.size + if width * height > pixel_limit: + raise UnusableImageError( + f"{source_format} input is {width}x{height}, " + f"over the {pixel_limit} pixel budget." + ) + + if source_format == "JPEG": + image.draft("RGB", TARGET_SIZE) + ImageOps.exif_transpose(image, in_place=True) + + image.thumbnail(TARGET_SIZE, Image.Resampling.LANCZOS, reducing_gap=2.0) + return _encode_centered_jpeg(image) + except UnusableImageError: + raise + except (Image.DecompressionBombError, OSError, ValueError) as exc: + raise UnusableImageError(f"Pillow could not decode the image: {exc}") from exc + + +def _encode_centered_jpeg(image: Image.Image) -> bytes: + if image.mode in ("RGBA", "LA") or "transparency" in image.info: + # Dropping alpha without a mask would render transparency as black. + image = image.convert("RGBA") + mask = image.getchannel("A") + else: + image = image.convert("RGB") + mask = None + + canvas = Image.new("RGB", TARGET_SIZE, CANVAS_COLOR) + left = (TARGET_SIZE[0] - image.width) // 2 + top = (TARGET_SIZE[1] - image.height) // 2 + canvas.paste(image, (left, top), mask) + + buffer = io.BytesIO() + canvas.save(buffer, format="JPEG", quality=JPEG_QUALITY) + return buffer.getvalue() + + +class RedrawWorkflow(WorkflowEntrypoint): + async def run(self, event, step): + job_id = event["payload"]["jobId"] + bucket = self.env.REDRAW_BUCKET + source_key = original_key(job_id) + target_key = output_key(job_id) + + @step.do() + async def verify_original(): + if await bucket.head(source_key) is None: + raise NonRetryableError(f"No original stored for job {job_id}.") + return source_key + + @step.do(config=AI_RETRIES) + async def redraw(verify_original): + source = await bucket.get(verify_original) + if source is None: + return json.dumps({"ok": False, "reason": MISSING_ORIGINAL_REASON}) + original = await source.blob() + + try: + reference = resize(await original.bytes()) + except UnusableImageError as exc: + print(f"Job {job_id} has an unusable original: {exc}") + return json.dumps({"ok": False, "reason": INVALID_IMAGE_REASON}) + + form = FormData() + for field, value in AI_OPTIONS.items(): + form[field] = value + form.append("input_image_0", Blob(reference, "image/jpeg"), "input.jpg") + + serialized = Response(form) + multipart_type = serialized.headers["content-type"] + + try: + generated = await self.env.AI.run( + MODEL, + { + "multipart": { + "body": serialized.body, + "contentType": multipart_type, + } + }, + ) + except Exception as exc: + if ai_error_code(getattr(exc, "message", None)) != AI_SAFETY_ERROR_CODE: + raise + # Returning rather than raising checkpoints the step, which is + # what guarantees no further inference happens. + return json.dumps({"ok": False, "reason": SAFETY_REJECTED_REASON}) + + # FLUX replies with JSON holding a base64 image, so decode before storing. + image = base64.b64decode(generated["image"]) + content_type = sniff_content_type(image) + if content_type is None: + return json.dumps({"ok": False, "reason": INVALID_OUTPUT_REASON}) + + await bucket.put( + target_key, + image, + httpMetadata={"contentType": content_type}, + customMetadata={"jobId": job_id}, + ) + return json.dumps({"ok": True, "key": target_key}) + + result = json.loads(await redraw()) + if not result["ok"]: + await bucket.put( + failure_key(job_id), + result["reason"], + httpMetadata={"contentType": "text/plain"}, + customMetadata={"jobId": job_id}, + ) + # Raised outside the retrying step so the failure is final. + raise NonRetryableError(result["reason"]) + return None diff --git a/image-redraw/wrangler.jsonc b/image-redraw/wrangler.jsonc new file mode 100644 index 0000000..3cd4337 --- /dev/null +++ b/image-redraw/wrangler.jsonc @@ -0,0 +1,52 @@ +{ + "$schema": "node_modules/wrangler/config-schema.json", + "name": "image-redraw-worker", + "main": "src/entry.py", + "compatibility_date": "2026-08-21", + "compatibility_flags": [ + "python_workers", + "python_workflows_implicit_dependencies" + ], + "assets": { + "directory": "./public", + "run_worker_first": [ + "/api/*" + ] + }, + "r2_buckets": [ + { + "binding": "REDRAW_BUCKET", + "bucket_name": "image-redraw" + } + ], + "queues": { + "producers": [ + { + "binding": "REDRAW_QUEUE", + "queue": "image-redraw-jobs" + } + ], + "consumers": [ + { + "queue": "image-redraw-jobs", + "max_batch_size": 10, + "max_batch_timeout": 2, + "max_retries": 3 + } + ] + }, + "workflows": [ + { + "name": "image-redraw", + "binding": "REDRAW_WORKFLOW", + "class_name": "RedrawWorkflow" + } + ], + "ai": { + "binding": "AI", + "remote": true + }, + "observability": { + "enabled": true + } +}