From b5ccd59ea02bc43c355a19bd810c6ec5e833c912 Mon Sep 17 00:00:00 2001 From: Katie Schilling <8411213+katieschilling@users.noreply.github.com> Date: Fri, 16 Jan 2026 14:41:21 -0700 Subject: [PATCH 1/5] feat(search): return 6 images instead of 1 in search results The template already supported multiple images, but the search was limited to 1 result. Now returns top 6 matches by similarity. Assisted-by: Claude Opus 4.5 via Claude Code Co-Authored-By: Claude Opus 4.5 --- main.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/main.py b/main.py index fbd8548..782957c 100644 --- a/main.py +++ b/main.py @@ -100,8 +100,7 @@ def perform_search(screenshots: pxt.Table, query: str) -> pxt.ResultSet: uuid=screenshots.uuid, url=screenshots.image.fileurl, ) - # .limit(6) - .limit(1) + .limit(6) ) return results.collect() From 6e67f848afa952076b919580d85c6357f358fef3 Mon Sep 17 00:00:00 2001 From: Katie Schilling <8411213+katieschilling@users.noreply.github.com> Date: Tue, 20 Jan 2026 09:11:49 -0700 Subject: [PATCH 2/5] feat: add image generation UI with Reve integration - Add UI to select search result image and enter generation prompt - Integrate Pixeltable's Reve edit() for image-to-image generation - Add /api/generate-image endpoint with computed column pipeline - Add generate_result.html and generate_error.html templates - Add CSS styles for generation panel, buttons, and textarea - Make OpenAI client optional (not needed for Reve) - Update pixeltable version requirement to >=0.5.1 Note: Requires REVE_API_KEY environment variable for generation. Assisted-by: Claude Opus 4.5 via Claude Code Co-Authored-By: Claude Opus 4.5 --- main.py | 83 ++++++++++++++-- pyproject.toml | 2 +- static/css/main.css | 102 ++++++++++++++++++++ templates/partials/api/generate_error.html | 11 +++ templates/partials/api/generate_result.html | 31 ++++++ templates/partials/api/search.html | 77 +++++++++++---- 6 files changed, 279 insertions(+), 27 deletions(-) create mode 100644 templates/partials/api/generate_error.html create mode 100644 templates/partials/api/generate_result.html diff --git a/main.py b/main.py index 782957c..ff1d05e 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,6 @@ import base64 import boto3 +import io import os import pixeltable as pxt import PIL.Image @@ -13,9 +14,23 @@ from urllib.parse import urlparse from uuid_extensions import uuid7 from pixeltable.functions.huggingface import clip +from pixeltable.functions import reve + + +def pil_to_base64(img: PIL.Image.Image) -> str: + """Convert a PIL Image to a base64 string.""" + if img is None: + return "" + buffered = io.BytesIO() + img.save(buffered, format="PNG") + return base64.b64encode(buffered.getvalue()).decode("utf-8") load_dotenv() -oai = OpenAI() + +# OpenAI client is optional - only initialize if API key is available +oai = None +if os.environ.get("OPENAI_API_KEY"): + oai = OpenAI() @pxt.udf @@ -66,11 +81,6 @@ def encode_image(file_path): return base64_image -# @pxt.udf -# def image_edit(prompt: str, input: PIL.Image.Image) -> PIL.Image.Image: -# pass - - generated_images = pxt.create_table( "generated_images", { @@ -87,9 +97,10 @@ def encode_image(file_path): input_image=get_image(generated_images.input_image_id), if_exists="ignore", ) -# generated_images.add_computed_column( -# gen_image=image_generations(generated_images.prompt, model="gpt-image-1") -# ) +generated_images.add_computed_column( + gen_image=reve.edit(generated_images.input_image, generated_images.prompt), + if_exists="ignore", +) def perform_search(screenshots: pxt.Table, query: str) -> pxt.ResultSet: @@ -107,6 +118,7 @@ def perform_search(screenshots: pxt.Table, query: str) -> pxt.ResultSet: app = Flask(__name__) htmx = HTMX(app) +app.jinja_env.filters["b64encode"] = pil_to_base64 tigris = boto3.client( "s3", endpoint_url="https://t3.storage.dev", @@ -181,3 +193,56 @@ def api_search(): query=query, results=results, ) + + +@app.route("/api/generate-image", methods=["POST"]) +def api_generate_image(): + if not htmx: + return redirect("/") + + image_id = request.form.get("image_id", "") + prompt = request.form.get("prompt", "") + + if not image_id or not prompt: + return render_template( + "partials/api/generate_error.html", + error="Please select an image and enter a prompt.", + ) + + # Insert into generated_images table - Pixeltable will compute the generated image + generated_images.insert([{"input_image_id": image_id, "prompt": prompt}]) + + # Get the most recently generated image + result = ( + generated_images.where( + (generated_images.input_image_id == image_id) + & (generated_images.prompt == prompt) + ) + .select( + uuid=generated_images.uuid, + input_image=generated_images.input_image, + gen_image=generated_images.gen_image, + prompt=generated_images.prompt, + ) + .order_by(generated_images.uuid, asc=False) + .limit(1) + .collect() + ) + + if not result: + return render_template( + "partials/api/generate_error.html", + error="Failed to generate image. Please try again.", + ) + + row = result[0] + return render_template( + "partials/api/generate_result.html", + generated_image=row["gen_image"], + input_image=row["input_image"], + prompt=row["prompt"], + ) + + +if __name__ == "__main__": + app.run(debug=True, host="0.0.0.0", port=5000) diff --git a/pyproject.toml b/pyproject.toml index 630aec0..65e23e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ "ipykernel>=7.1.0", "ipython>=9.7.0", "openai>=2.8.1", - "pixeltable==0.5.0", + "pixeltable>=0.5.1", "pixeltable-yolox>=0.4.2", "pydantic>=2.12.4", "s3fs>=2025.10.0", diff --git a/static/css/main.css b/static/css/main.css index 6dd4819..6d2c4e3 100644 --- a/static/css/main.css +++ b/static/css/main.css @@ -307,3 +307,105 @@ input.form-control.search::placeholder { h1 { margin-bottom: 0.5rem; } + +/* Form controls for generation panel */ +textarea.form-control { + background-color: #282828; + border: 2px solid #665c54; + color: #f9f5d7; + border-radius: 0.5rem; + padding: 0.75rem 1rem; + font-family: "Iosevka Aile Iaso", sans-serif; + font-size: 1rem; + transition: all 0.3s ease; +} + +textarea.form-control:focus { + outline: none; + border-color: #b16286; + box-shadow: 0 0 0 0.2rem rgba(177, 98, 134, 0.25); + background-color: #3c3836; +} + +textarea.form-control::placeholder { + color: #bdae93; + opacity: 0.7; +} + +/* Button styles */ +.btn { + background-color: #b16286; + color: #f9f5d7; + border: none; + border-radius: 0.5rem; + padding: 0.5rem 1.5rem; + font-family: "Iosevka Aile Iaso", sans-serif; + font-size: 1rem; + cursor: pointer; + transition: all 0.2s ease; +} + +.btn:hover { + background-color: #d3869b; + transform: translateY(-1px); +} + +.btn:active { + transform: translateY(0); +} + +.btn:disabled { + background-color: #665c54; + cursor: not-allowed; + transform: none; +} + +/* Generation panel styles */ +#generation-panel { + background-color: #282828; + border: 1px solid #3c3836; + border-radius: 0.5rem; + padding: 1.5rem; + margin-top: 2rem; +} + +#generation-panel h3 { + color: #f9f5d7; + margin-bottom: 1rem; +} + +/* Light mode styles */ +@media (prefers-color-scheme: light) { + textarea.form-control { + background-color: #fbf1c7; + border-color: #928374; + color: #1d2021; + } + + textarea.form-control:focus { + border-color: #b16286; + background-color: #f9f5d7; + } + + textarea.form-control::placeholder { + color: #665c54; + } + + .btn { + background-color: #b16286; + color: #f9f5d7; + } + + .btn:hover { + background-color: #d3869b; + } + + #generation-panel { + background-color: #fbf1c7; + border-color: #ebdbb2; + } + + #generation-panel h3 { + color: #1d2021; + } +} diff --git a/templates/partials/api/generate_error.html b/templates/partials/api/generate_error.html new file mode 100644 index 0000000..43bfe9a --- /dev/null +++ b/templates/partials/api/generate_error.html @@ -0,0 +1,11 @@ +
+

Error

+

{{ error }}

+ +
diff --git a/templates/partials/api/generate_result.html b/templates/partials/api/generate_result.html new file mode 100644 index 0000000..9715513 --- /dev/null +++ b/templates/partials/api/generate_result.html @@ -0,0 +1,31 @@ +
+

Generated Image

+

Prompt: "{{ prompt }}"

+ +
+
+

Original

+ Original image +
+
+

Generated

+ Generated image +
+
+ + +
diff --git a/templates/partials/api/search.html b/templates/partials/api/search.html index 96c1a0a..a0c87ad 100644 --- a/templates/partials/api/search.html +++ b/templates/partials/api/search.html @@ -1,21 +1,64 @@
{% for result in results %} -
-
- - - Search result -
+
+ Search result
{% endfor %} -
\ No newline at end of file + + + + + \ No newline at end of file From 2411d35659a1651d93182194f271be45b0ab42db Mon Sep 17 00:00:00 2001 From: Katie Schilling <8411213+katieschilling@users.noreply.github.com> Date: Tue, 20 Jan 2026 09:28:18 -0700 Subject: [PATCH 3/5] fix(main): correct S3 endpoint and image type handling - Add Tigris S3 endpoint URL (t3.storage.dev) to storage_options - Remove invalid if_exists argument from load_dataset call - Store input_image directly as pxt.Image instead of using computed query - Update /api/generate-image to fetch image before insert Assisted-by: Claude Opus 4.5 via Claude Code Co-Authored-By: Claude Opus 4.5 --- main.py | 44 +++++++++++++++++++++++--------------------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/main.py b/main.py index ff1d05e..27c12d2 100644 --- a/main.py +++ b/main.py @@ -39,7 +39,10 @@ def gen_uuid() -> str: def import_screenshots(): data_files = "s3://xe-zohar-copy/ds/screenshots_sharded/*.parquet" - storage_options={"profile": "tigris-dev"} + storage_options = { + "profile": "tigris-dev", + "endpoint_url": "https://t3.storage.dev", + } dataset = load_dataset( "parquet", @@ -47,7 +50,6 @@ def import_screenshots(): data_files=data_files, streaming=False, storage_options=storage_options, - if_exists="ignore", ) screenshots = pxt.create_table("screenshots", source=dataset, if_exists="ignore") screenshots.add_embedding_index( @@ -66,15 +68,6 @@ def import_screenshots(): screenshots = cast(pxt.Table, screenshots) -@pxt.query -def get_image(image_id: str) -> PIL.Image.Image: - return ( - screenshots.where(screenshots.uuid == image_id) - .select(screenshots.image) - .limit(1) - ) - - def encode_image(file_path): with open(file_path, "rb") as f: base64_image = base64.b64encode(f.read()).decode("utf-8") @@ -84,7 +77,7 @@ def encode_image(file_path): generated_images = pxt.create_table( "generated_images", { - "input_image_id": pxt.String, + "input_image": pxt.Image, "prompt": pxt.String, }, if_exists="ignore", @@ -93,10 +86,6 @@ def encode_image(file_path): "creating the table did not result in creating the table" ) generated_images.add_computed_column(uuid=gen_uuid(), if_exists="ignore") -generated_images.add_computed_column( - input_image=get_image(generated_images.input_image_id), - if_exists="ignore", -) generated_images.add_computed_column( gen_image=reve.edit(generated_images.input_image, generated_images.prompt), if_exists="ignore", @@ -209,15 +198,28 @@ def api_generate_image(): error="Please select an image and enter a prompt.", ) + # Get the input image from screenshots table + input_image_result = ( + screenshots.where(screenshots.uuid == image_id) + .select(screenshots.image) + .limit(1) + .collect() + ) + + if not input_image_result: + return render_template( + "partials/api/generate_error.html", + error="Could not find the selected image.", + ) + + input_image = input_image_result[0]["image"] + # Insert into generated_images table - Pixeltable will compute the generated image - generated_images.insert([{"input_image_id": image_id, "prompt": prompt}]) + generated_images.insert([{"input_image": input_image, "prompt": prompt}]) # Get the most recently generated image result = ( - generated_images.where( - (generated_images.input_image_id == image_id) - & (generated_images.prompt == prompt) - ) + generated_images.where(generated_images.prompt == prompt) .select( uuid=generated_images.uuid, input_image=generated_images.input_image, From 7363bea747d6a6505a02e6406cdacf72c56db8c0 Mon Sep 17 00:00:00 2001 From: Katie Schilling <8411213+katieschilling@users.noreply.github.com> Date: Tue, 20 Jan 2026 09:31:26 -0700 Subject: [PATCH 4/5] fix(search): serve images via Flask route instead of S3 presigned URLs Images are stored locally by Pixeltable, not on S3. Added /api/image/ endpoint to serve images directly from the screenshots table. Assisted-by: Claude Opus 4.5 via Claude Code Co-Authored-By: Claude Opus 4.5 --- main.py | 42 +++++++++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/main.py b/main.py index 27c12d2..68f690f 100644 --- a/main.py +++ b/main.py @@ -7,7 +7,7 @@ from botocore.client import Config from datasets import load_dataset from dotenv import load_dotenv -from flask import Flask, request, render_template, send_from_directory, redirect +from flask import Flask, request, render_template, send_from_directory, redirect, send_file from flask_htmx import HTMX from openai import OpenAI from typing import cast @@ -132,6 +132,29 @@ def healthz(): return "OK" +@app.route("/api/image/") +def api_image(uuid): + """Serve an image from the screenshots table by UUID.""" + result = ( + screenshots.where(screenshots.uuid == uuid) + .select(screenshots.image) + .limit(1) + .collect() + ) + if not result: + return "Image not found", 404 + + img = result[0]["image"] + if img is None: + return "Image not found", 404 + + # Convert PIL Image to bytes and serve + buffered = io.BytesIO() + img.save(buffered, format="PNG") + buffered.seek(0) + return send_file(buffered, mimetype="image/png") + + @app.route("/") def index(): if htmx: @@ -154,26 +177,11 @@ def api_search(): results = [] for hit in hits: - print(hit) uuid = hit["uuid"] - url = hit["url"] - - # Parse S3 URL to extract bucket and key - parsed_url = urlparse(url) - bucket = parsed_url.netloc - key = parsed_url.path[1:] - print(bucket, key) - - presigned_url = tigris.generate_presigned_url( - "get_object", - Params={"Bucket": bucket, "Key": key}, - ExpiresIn=3600, - ) - results.append( { "id": uuid, - "url": presigned_url, + "url": f"/api/image/{uuid}", } ) From c6e8fe771a2db91f16fd940e0c75a091a829b623 Mon Sep 17 00:00:00 2001 From: Katie Schilling <8411213+katieschilling@users.noreply.github.com> Date: Thu, 22 Jan 2026 10:35:37 -0700 Subject: [PATCH 5/5] feat: fix image loading and add model selection for generation - Fix threading issue by returning base64 images directly in search - Add visible progress indicator during image generation - Add model selection dropdown (Reve, Fal AI/Replicate, HuggingFace) - Update backend to support multiple generation models - Add CSS styling for select dropdown Assisted-by: Claude Opus 4.5 via Claude Code Co-Authored-By: Claude Opus 4.5 --- main.py | 94 +++++++++++++++++++++--------- static/css/main.css | 22 ++++++- templates/partials/api/search.html | 53 ++++++++++------- 3 files changed, 119 insertions(+), 50 deletions(-) diff --git a/main.py b/main.py index 68f690f..a54a3b8 100644 --- a/main.py +++ b/main.py @@ -13,8 +13,8 @@ from typing import cast from urllib.parse import urlparse from uuid_extensions import uuid7 -from pixeltable.functions.huggingface import clip -from pixeltable.functions import reve +from pixeltable.functions.huggingface import clip, image_to_image +from pixeltable.functions import reve, replicate def pil_to_base64(img: PIL.Image.Image) -> str: @@ -98,7 +98,7 @@ def perform_search(screenshots: pxt.Table, query: str) -> pxt.ResultSet: screenshots.order_by(sim, asc=False) .select( uuid=screenshots.uuid, - url=screenshots.image.fileurl, + image=screenshots.image, ) .limit(6) ) @@ -178,10 +178,13 @@ def api_search(): for hit in hits: uuid = hit["uuid"] + image = hit["image"] + # Convert image to base64 data URL to avoid concurrent request issues + b64 = pil_to_base64(image) results.append( { "id": uuid, - "url": f"/api/image/{uuid}", + "url": f"data:image/png;base64,{b64}", } ) @@ -192,6 +195,32 @@ def api_search(): ) +def generate_with_reve(image: PIL.Image.Image, prompt: str) -> PIL.Image.Image: + """Generate image using Reve edit.""" + return reve.edit(image, prompt) + + +def generate_with_huggingface(image: PIL.Image.Image, prompt: str) -> PIL.Image.Image: + """Generate image using HuggingFace image-to-image.""" + return image_to_image( + image, + prompt, + model_id="timbrooks/instruct-pix2pix", + ) + + +def generate_with_replicate(image: PIL.Image.Image, prompt: str) -> PIL.Image.Image: + """Generate image using Replicate FLUX.""" + return replicate.run( + "black-forest-labs/flux-1.1-pro", + { + "prompt": prompt, + "image": image, + "prompt_upsampling": True, + }, + ) + + @app.route("/api/generate-image", methods=["POST"]) def api_generate_image(): if not htmx: @@ -199,6 +228,7 @@ def api_generate_image(): image_id = request.form.get("image_id", "") prompt = request.form.get("prompt", "") + model = request.form.get("model", "reve") if not image_id or not prompt: return render_template( @@ -222,35 +252,45 @@ def api_generate_image(): input_image = input_image_result[0]["image"] - # Insert into generated_images table - Pixeltable will compute the generated image - generated_images.insert([{"input_image": input_image, "prompt": prompt}]) - - # Get the most recently generated image - result = ( - generated_images.where(generated_images.prompt == prompt) - .select( - uuid=generated_images.uuid, - input_image=generated_images.input_image, - gen_image=generated_images.gen_image, - prompt=generated_images.prompt, - ) - .order_by(generated_images.uuid, asc=False) - .limit(1) - .collect() - ) - - if not result: + # Generate image using selected model + try: + if model == "reve": + # Use the computed column approach for Reve + generated_images.insert([{"input_image": input_image, "prompt": prompt}]) + result = ( + generated_images.where(generated_images.prompt == prompt) + .select( + input_image=generated_images.input_image, + gen_image=generated_images.gen_image, + prompt=generated_images.prompt, + ) + .order_by(generated_images.uuid, asc=False) + .limit(1) + .collect() + ) + if not result: + raise Exception("Failed to generate image") + generated_image = result[0]["gen_image"] + elif model == "huggingface": + generated_image = generate_with_huggingface(input_image, prompt) + elif model == "fal": + generated_image = generate_with_replicate(input_image, prompt) + else: + return render_template( + "partials/api/generate_error.html", + error=f"Unknown model: {model}", + ) + except Exception as e: return render_template( "partials/api/generate_error.html", - error="Failed to generate image. Please try again.", + error=f"Generation failed: {str(e)}", ) - row = result[0] return render_template( "partials/api/generate_result.html", - generated_image=row["gen_image"], - input_image=row["input_image"], - prompt=row["prompt"], + generated_image=generated_image, + input_image=input_image, + prompt=prompt, ) diff --git a/static/css/main.css b/static/css/main.css index 6d2c4e3..b672b42 100644 --- a/static/css/main.css +++ b/static/css/main.css @@ -309,7 +309,8 @@ h1 { } /* Form controls for generation panel */ -textarea.form-control { +textarea.form-control, +select.form-control { background-color: #282828; border: 2px solid #665c54; color: #f9f5d7; @@ -320,6 +321,15 @@ textarea.form-control { transition: all 0.3s ease; } +select.form-control { + cursor: pointer; + appearance: none; + background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='12' height='12' viewBox='0 0 12 12'%3E%3Cpath fill='%23bdae93' d='M6 8L1 3h10z'/%3E%3C/svg%3E"); + background-repeat: no-repeat; + background-position: right 1rem center; + padding-right: 2.5rem; +} + textarea.form-control:focus { outline: none; border-color: #b16286; @@ -376,13 +386,19 @@ textarea.form-control::placeholder { /* Light mode styles */ @media (prefers-color-scheme: light) { - textarea.form-control { + textarea.form-control, + select.form-control { background-color: #fbf1c7; border-color: #928374; color: #1d2021; } - textarea.form-control:focus { + select.form-control { + background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='12' height='12' viewBox='0 0 12 12'%3E%3Cpath fill='%23665c54' d='M6 8L1 3h10z'/%3E%3C/svg%3E"); + } + + textarea.form-control:focus, + select.form-control:focus { border-color: #b16286; background-color: #f9f5d7; } diff --git a/templates/partials/api/search.html b/templates/partials/api/search.html index a0c87ad..669f869 100644 --- a/templates/partials/api/search.html +++ b/templates/partials/api/search.html @@ -5,7 +5,7 @@ class="search-result-img" style="max-width: 14rem; cursor: pointer; border: 3px solid transparent; border-radius: 0.5rem; transition: border-color 0.2s;" src="{{result.url}}" - alt="Search result" + alt="" data-result-id="{{result.id}}" data-result-url="{{result.url}}" onclick="selectImageForGeneration(this, '{{result.id}}', '{{result.url}}')" @@ -14,33 +14,47 @@ {% endfor %} -