Make the OpenAI image URL fetchable, keep WebM audio, stream example imports

Four review items on the diffusion Studio work:

- response_format=url returned the bearer-gated gallery route, which a standard image
  client downloads with no Authorization header, so the default response format was
  unusable. Mint a short-lived HMAC link instead (the shape RAG already uses for pdf.js)
  served by a signed route, and leave the gallery route itself bearer-only.
- A manual gpu_layers=0 load carrying speculative_type="off" -- a value the UI persists
  and sends -- read as GPU-bearing, so it took the GPU arbiter and evicted a resident
  image/video pipeline even though the launcher hides the GPUs for it. Canonicalize the
  mode and exempt "off".
- The curated example import prepared the whole split before the loop stopped at the
  10-100 image cap; m1guelpf/nouns is 49,859 rows / 328 MB. Stream instead, with the
  prepared load kept as a fallback for a repo that cannot stream.
- WebM export dropped the audio track an LTX-2 clip carries, silently, on the format
  offered for web embeds. Mux it as Opus through a resampler + FIFO, and keep exporting
  the video alone on a build without libopus.
This commit is contained in:
Daniel Han 2026-07-26 23:39:49 +00:00
commit bce3b9b20a
8 changed files with 381 additions and 21 deletions

View file

@ -1431,8 +1431,12 @@ def zero_vram_chat_load(
if gpu_memory_mode != "manual" or gpu_layers != 0:
return False
# Any speculative mode may launch a GPU drafter, and only the request's own knobs are known here,
# so treat every non-empty selection as GPU-bearing.
if needs_mmproj or speculative_type:
# so treat every selection as GPU-bearing -- except "off", which the resolver never emits a
# drafter for. Canonicalize first: the UI persists and sends the literal "off", which a bare
# truthiness test read as "speculation requested" and so evicted a resident image/video pipeline
# for a CPU-only load.
spec_mode = _canonicalize_spec_mode(speculative_type)
if needs_mmproj or spec_mode not in (None, "off"):
return False
if LlamaCppBackend._is_vulkan_backend():
return False

View file

@ -116,17 +116,61 @@ def _transcode_webm(path: Path) -> bytes:
# Realtime settings: VP9's default "good" profile is slow; cpu-used 8 + row-mt is much faster at
# a small quality cost, right for a download button.
out_v.options = {"deadline": "realtime", "cpu-used": "8", "row-mt": "1"}
for frame in src.decode(in_v):
for packet in out_v.encode(frame.reformat(format = "yuv420p")):
dst.mux(packet)
# An LTX-2 clip carries a synchronized audio track, and WebM is offered as the web-embed
# format, so dropping the track would silently hand back half the generated result. Opus is
# WebM's audio codec: resample to its 48 kHz grid and hand the encoder whole frames through a
# FIFO (libopus takes a fixed 20 ms frame, 960 samples at 48 kHz).
in_a = src.streams.audio[0] if src.streams.audio else None
out_a = fifo = resampler = None
if in_a is not None:
try:
stereo = (getattr(in_a.codec_context.layout, "nb_channels", 1) or 1) > 1
layout = "stereo" if stereo else "mono"
out_a = dst.add_stream("libopus", rate = 48000, layout = layout)
resampler = av.audio.resampler.AudioResampler(
format = out_a.format.name, layout = layout, rate = 48000
)
fifo = av.audio.fifo.AudioFifo()
except Exception: # noqa: BLE001 -- a build without libopus still exports the video
out_a = fifo = resampler = None
def _drain_audio(flush: bool = False) -> None:
# frame_size is 0 until the container starts writing; 960 is libopus' own frame.
size = out_a.frame_size or 960
while True:
frame = fifo.read(size, partial = flush)
if frame is None:
break
for packet in out_a.encode(frame):
dst.mux(packet)
# Demux both streams together so the muxer sees them interleaved rather than buffering
# every video packet until the audio arrives.
for packet in src.demux(*([in_v] + ([in_a] if out_a is not None else []))):
if packet.dts is None: # flush packet from the demuxer
continue
if packet.stream is in_v:
for frame in packet.decode():
for out_packet in out_v.encode(frame.reformat(format = "yuv420p")):
dst.mux(out_packet)
continue
for frame in packet.decode():
for resampled in resampler.resample(frame):
# Let the FIFO time the output: the resampler's frames do not line up with
# Opus' fixed frame size.
resampled.pts = None
fifo.write(resampled)
_drain_audio()
for packet in out_v.encode():
dst.mux(packet)
if out_a is not None:
_drain_audio(flush = True)
for packet in out_a.encode():
dst.mux(packet)
except RuntimeError:
raise
except Exception as exc: # noqa: BLE001 -- surface as "encoder unavailable"
raise RuntimeError(f"WebM export failed (libvpx-vp9 unavailable?): {exc}") from exc
# Audio dropped: Opus muxing needs a 48 kHz resample chain and most clips are silent (the
# original MP4 keeps the audio).
return buf.getvalue()

View file

@ -10,7 +10,10 @@ import sys
import time
import uuid
from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request, status
import hashlib as _hashlib
import hmac as _hmac
import secrets as _secrets
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from fastapi.responses import StreamingResponse, JSONResponse, Response
from starlette.requests import ClientDisconnect
from typing import Any, Callable, List, Literal, Optional, Union
@ -16500,13 +16503,71 @@ def _parse_openai_image_size(size: str) -> tuple[int, int]:
return width, height
def _absolute_image_url(request: Request, relative: str) -> str:
"""Join a relative gallery path onto the request's own scheme+host, for the
response_format=url links. Like every Studio route, the target needs the
bearer token; b64_json avoids that for clients that can't carry it."""
# response_format=url links have to be fetchable by whoever received them: an OpenAI client hands
# data[].url back to the caller, who downloads it with a plain GET and no Authorization header, so a
# link to the bearer-gated gallery route answered 401 and the default response format was unusable.
# Mint a short-lived HMAC link instead -- the shape RAG already uses to feed pdf.js range requests --
# and leave the gallery route itself bearer-only. One hour matches OpenAI's own URL lifetime, and the
# per-process secret means a restart invalidates every outstanding link.
_IMAGE_LINK_TTL = 3600
_IMAGE_LINK_SECRET = _secrets.token_bytes(32)
def _sign_image_id(image_id: str) -> str:
exp = int(time.time()) + _IMAGE_LINK_TTL
payload = f"{image_id}.{exp}"
sig = _hmac.new(_IMAGE_LINK_SECRET, payload.encode(), _hashlib.sha256).hexdigest()
return f"{payload}.{sig}"
def _verify_image_link_token(token: str) -> Optional[str]:
"""The image id a valid, unexpired token names, else None. Gallery ids are
``[A-Za-z0-9_-]`` so the two dots always split id / expiry / signature."""
try:
image_id, exp_s, sig = token.rsplit(".", 2)
except ValueError:
return None
expected = _hmac.new(
_IMAGE_LINK_SECRET, f"{image_id}.{exp_s}".encode(), _hashlib.sha256
).hexdigest()
if not _hmac.compare_digest(sig, expected):
return None
try:
if int(exp_s) < int(time.time()):
return None
except ValueError:
return None
return image_id
def _absolute_image_url(request: Request, image_id: str) -> str:
"""The absolute, directly fetchable link for one gallery image, on the request's own
scheme+host. Signed rather than bearer-gated (see above), so a standard image client can
download it; b64_json still avoids the round trip entirely."""
relative = f"/api/inference/images/gallery/{image_id}/file-signed?token={_sign_image_id(image_id)}"
return str(request.base_url).rstrip("/") + relative
@studio_router.get("/images/gallery/{image_id}/file-signed")
async def get_gallery_image_file_signed(image_id: str, token: str = Query(...)):
"""Serve one gallery PNG gated by the HMAC token instead of the bearer, for the
response_format=url links a plain image client downloads. Same ownership gate as the
authenticated route, and the token names the single image it may serve."""
from core.inference import image_gallery
if _verify_image_link_token(token) != image_id:
raise HTTPException(status_code = 401, detail = "Invalid or expired image link.")
path = await asyncio.to_thread(image_gallery.owned_image_path, image_id)
if path is None:
raise HTTPException(status_code = 404, detail = "Image not found.")
data = await asyncio.to_thread(path.read_bytes)
return Response(
content = data,
media_type = "image/png",
headers = {"Cache-Control": "private, max-age=31536000, immutable"},
)
@router.post(
"/images/generations",
response_model = ImageGenerationResponse,
@ -16627,7 +16688,7 @@ async def openai_image_generations(
raise RuntimeError("generated image could not be read back for encoding")
items.append(ImageGenerationData(b64_json = encoded))
else:
items.append(ImageGenerationData(url = _absolute_image_url(request, record["url"])))
items.append(ImageGenerationData(url = _absolute_image_url(request, record["id"])))
return items
try:

View file

@ -2291,6 +2291,19 @@ def _detect_image_column(features) -> Optional[str]:
return None
def _detect_image_column_from_row(row: dict) -> Optional[str]:
"""Image column picked from one materialized row, for a streamed dataset that arrives with no
feature metadata to inspect."""
try:
from PIL.Image import Image as PILImage
except Exception: # noqa: BLE001 -- no Pillow -> the caller reports "no image column"
return None
for col, value in row.items():
if isinstance(value, PILImage):
return col
return None
def _detect_caption_column(entry: dict, columns: list[str]) -> Optional[str]:
"""Pick the caption column: the entry's declared one if present, else a common name."""
declared = entry.get("caption_column")
@ -2310,18 +2323,38 @@ def _materialize_hf_dataset(entry: dict, dest: Path, cap: int) -> int:
kwargs = {"split": "train"}
if entry.get("no_checks"):
kwargs["verification_mode"] = "no_checks"
ds = load_dataset(entry["repo"], **kwargs)
image_col = _detect_image_column(ds.features)
if image_col is None:
# Stream rather than prepare the whole split: the loop keeps at most `cap` rows (10-100) while
# these curated repos run to 49,859 rows / 328 MB (m1guelpf/nouns) and 1,000 rows / 237 MB
# (huggan/smithsonian_butterflies_subset), all of which a prepared load downloads and converts
# before the first row is read. A repo that cannot stream (loading script, no listed data files)
# falls back to the prepared load so the one-click import still works.
try:
ds = load_dataset(entry["repo"], streaming = True, **kwargs)
features = ds.features
except Exception: # noqa: BLE001 -- not streamable; the prepared load is the fallback
ds = load_dataset(entry["repo"], **kwargs)
features = ds.features
# Streaming can hand back a dataset whose features are only known once a row is read, so the
# columns are resolved from the first row in that case.
image_col = _detect_image_column(features) if features else None
if image_col is None and features:
raise HTTPException(
status_code = 502,
detail = f"'{entry['repo']}' has no image column to import.",
)
caption_col = _detect_caption_column(entry, list(ds.features.keys()))
caption_col = _detect_caption_column(entry, list(features.keys())) if features else None
written = 0
for row in ds:
if written >= cap:
break
if image_col is None:
image_col = _detect_image_column_from_row(row)
if image_col is None:
raise HTTPException(
status_code = 502,
detail = f"'{entry['repo']}' has no image column to import.",
)
caption_col = _detect_caption_column(entry, list(row.keys()))
img = row[image_col]
if img is None:
continue

View file

@ -305,17 +305,21 @@ class _FakeDS:
return iter(self._rows)
def _install_fake_load_dataset(monkeypatch, n_rows):
calls = {"count": 0}
def _install_fake_load_dataset(monkeypatch, n_rows, features = "default", streamable = True):
calls = {"count": 0, "streaming": [], "features": features}
rows = [
{"image": Image.new("RGB", (8, 8), (i * 30 % 255, 60, 90)), "prompt": f"caption {i}"}
for i in range(n_rows)
]
features = {"image": _FakeImageFeature(), "prompt": object()}
if features == "default":
features = {"image": _FakeImageFeature(), "prompt": object()}
def fake_load(repo, **kwargs):
calls["count"] += 1
calls["streaming"].append(bool(kwargs.get("streaming")))
assert kwargs.get("split") == "train"
if kwargs.get("streaming") and not streamable:
raise ValueError("Loading a dataset cached in a LocalFileSystem is not supported")
return _FakeDS(rows, features)
import datasets
@ -367,6 +371,48 @@ def test_import_example_respects_cap(client, ds_root, monkeypatch):
assert r.json()["image_count"] == 2
def test_import_example_streams_instead_of_preparing_the_whole_split(
client, ds_root, monkeypatch
):
# The cap keeps 10-100 rows, while the curated repos run to 49,859 rows / 328 MB
# (m1guelpf/nouns), all of which a prepared load downloads and converts before the first row is
# read. The import must ask for a streamed split.
calls = _install_fake_load_dataset(monkeypatch, n_rows = 3)
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"})
assert r.status_code == 200, r.text
assert r.json()["imported"] == 3
assert calls["streaming"] == [True]
def test_import_example_falls_back_when_the_repo_cannot_stream(client, ds_root, monkeypatch):
# A repo with a loading script or no listed data files cannot stream; the one-click import must
# still work through the prepared load rather than 502.
calls = _install_fake_load_dataset(monkeypatch, n_rows = 3, streamable = False)
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"})
assert r.status_code == 200, r.text
assert r.json()["imported"] == 3
assert calls["streaming"] == [True, False]
def test_import_example_resolves_columns_from_the_first_row_without_features(
client, ds_root, monkeypatch
):
# A streamed dataset can arrive with no feature metadata to inspect, so the image and caption
# columns come from the first row instead.
_install_fake_load_dataset(monkeypatch, n_rows = 2, features = None)
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"})
assert r.status_code == 200, r.text
assert r.json()["imported"] == 2
assert r.json()["caption_count"] == 2
def test_import_example_without_an_image_column_maps_to_502(client, ds_root, monkeypatch):
# No image anywhere in the row: still a clean 502, not a KeyError 500.
_install_fake_load_dataset(monkeypatch, n_rows = 2, features = {"prompt": object()})
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "tuxemon"})
assert r.status_code == 502
def test_import_example_unknown_id_404(client, ds_root):
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "does-not-exist"})
assert r.status_code == 404

View file

@ -1192,6 +1192,19 @@ def test_zero_vram_chat_load_refuses_every_gpu_companion(not_vulkan):
assert zero("manual", 0, ["--model-draft", "/tmp/d.gguf", "--spec-draft-ngl", "0"]) is True
def test_zero_vram_chat_load_exempts_disabled_speculation(not_vulkan):
# "off" is a canonical mode the UI persists and sends, and the resolver emits no drafter for it,
# so a CPU-only load carrying it holds no VRAM either. Legacy spellings canonicalize the same way,
# while every mode that MAY resolve to a drafter stays GPU-bearing.
zero = llama_cpp_module.zero_vram_chat_load
assert zero("manual", 0, [], False, "off") is True
assert zero("manual", 0, [], False, " OFF ") is True
assert zero("manual", 0, [], False, "") is True
assert zero("manual", 0, [], False, "auto") is False
assert zero("manual", 0, [], False, "mtp") is False
assert zero("manual", 0, [], False, "default") is False
def test_zero_vram_chat_load_is_skipped_on_vulkan(monkeypatch):
# Vulkan builds are exempt from the CPU-only mask at launch, so the arbiter gate must match.
monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda *a, **k: True))

View file

@ -207,7 +207,9 @@ def test_url_response_shape(client):
assert len(body["data"]) == 1
item = body["data"][0]
assert "url" in item and "b64_json" not in item # exclude_none drops the unused key
assert item["url"].endswith("/file")
# Signed link, not the bearer-gated /file route: an OpenAI client downloads this URL with a plain
# GET and no Authorization header.
assert "/images/gallery/img0/file-signed?token=" in item["url"]
# Z-Image-Turbo defaults (9 steps, 0 guidance) flow into the backend call.
assert client.backend.calls[0] == dict(
prompt = "a sloth", width = 256, height = 256, steps = 9, guidance = 0.0, batch_size = 1
@ -381,3 +383,69 @@ def test_auth_required():
# No dependency override: the real auth dependency runs and rejects.
resp = TestClient(app).post("/v1/images/generations", json = {"prompt": "p"})
assert resp.status_code in (401, 403)
def _signed_link_app(monkeypatch, backend, png: "object"):
"""An app carrying BOTH surfaces: /v1 for the compat POST and /api/inference for the gallery,
which is where the returned link points. get_current_subject is overridden for the POST only in
the sense that the signed route never depends on it."""
from routes.inference import studio_router
app = FastAPI()
install_api_error_handlers(app)
app.include_router(router, prefix = "/v1")
app.include_router(studio_router, prefix = "/api/inference")
app.dependency_overrides[get_current_subject] = lambda: "test-user"
monkeypatch.setattr(gallery_module, "owned_image_path", lambda i: png if i == "img0" else None)
return TestClient(app)
def test_url_response_link_is_fetchable_without_the_bearer(monkeypatch, tmp_path):
# The point of response_format=url: an image client hands data[].url to a plain downloader that
# sends no Authorization header. Fetch the returned link with the auth header stripped and assert
# the PNG comes back.
png = tmp_path / "img0.png"
png.write_bytes(b"\x89PNG\r\n\x1a\nfake")
backend = _FakeBackend()
monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: backend)
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
store = {}
def _save(image, meta):
image_id = f"img{len(store)}"
record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"}
store[image_id] = record
return record
monkeypatch.setattr(gallery_module, "save", _save)
cli = _signed_link_app(monkeypatch, backend, png)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
assert resp.status_code == 200
url = resp.json()["data"][0]["url"]
fetched = cli.get(url.replace("http://testserver", ""), headers = {})
assert fetched.status_code == 200
assert fetched.headers["content-type"] == "image/png"
assert fetched.content == b"\x89PNG\r\n\x1a\nfake"
def test_signed_image_link_rejects_tampering_and_expiry(monkeypatch, tmp_path):
# The token names one image and carries its own expiry, so a swapped id, a forged signature and a
# stale link all 401 rather than serving the gallery.
import routes.inference as inference_routes
png = tmp_path / "img0.png"
png.write_bytes(b"\x89PNG\r\n\x1a\nfake")
backend = _FakeBackend()
cli = _signed_link_app(monkeypatch, backend, png)
token = inference_routes._sign_image_id("img0")
base = "/api/inference/images/gallery"
assert cli.get(f"{base}/img0/file-signed?token={token}").status_code == 200
# Same signature, different image.
assert cli.get(f"{base}/img1/file-signed?token={token}").status_code == 401
assert cli.get(f"{base}/img0/file-signed?token=img0.9999999999.dead").status_code == 401
assert cli.get(f"{base}/img0/file-signed?token=nonsense").status_code == 401
monkeypatch.setattr(inference_routes, "_IMAGE_LINK_TTL", -1)
expired = inference_routes._sign_image_id("img0")
assert cli.get(f"{base}/img0/file-signed?token={expired}").status_code == 401

View file

@ -354,6 +354,97 @@ def _real_mp4_bytes(
return buf.getvalue()
def _real_mp4_with_audio(seconds: int = 1, size: int = 32, rate: int = 8) -> bytes:
# An LTX-2-shaped clip: video plus a synchronized audio track (a 440 Hz tone), so the WebM
# export can be checked for the track rather than assumed silent.
av = pytest.importorskip("av")
np = pytest.importorskip("numpy")
import io
import math
from fractions import Fraction
arate = 44100
buf = io.BytesIO()
with av.open(buf, "w", format = "mp4") as out:
video = out.add_stream("mpeg4", rate = rate)
video.width = video.height = size
video.pix_fmt = "yuv420p"
audio = out.add_stream("aac", rate = arate)
audio.layout = "stereo"
for i in range(seconds * rate):
frame = av.VideoFrame.from_ndarray(
np.full((size, size, 3), (i * 30) % 256, dtype = np.uint8), format = "rgb24"
)
for packet in video.encode(frame):
out.mux(packet)
written = 0
while written < seconds * arate:
count = min(1024, seconds * arate - written)
tone = np.array(
[
int(20000 * math.sin(2 * math.pi * 440 * (written + k) / arate))
for k in range(count)
],
dtype = np.int16,
)
# Packed s16 is one interleaved row.
frame = av.AudioFrame.from_ndarray(
np.repeat(tone, 2).reshape(1, count * 2), format = "s16", layout = "stereo"
)
frame.sample_rate = arate
frame.pts = written
frame.time_base = Fraction(1, arate)
for packet in audio.encode(frame):
out.mux(packet)
written += count
for packet in video.encode():
out.mux(packet)
for packet in audio.encode():
out.mux(packet)
return buf.getvalue()
def test_webm_export_keeps_the_audio_track():
"""WebM is offered as the web-embed format, so a clip generated with synchronized audio (LTX-2)
must not come back mute: the exporter has to mux the track as Opus, WebM's audio codec."""
av = pytest.importorskip("av")
import io
record = gallery.save(_real_mp4_with_audio(), _meta())
webm = gallery.transcode(record["id"], "webm")
assert webm is not None and webm[:4] == b"\x1a\x45\xdf\xa3"
with av.open(io.BytesIO(webm)) as container:
kinds = {(s.type, s.codec_context.name) for s in container.streams}
assert ("video", "vp9") in kinds
assert ("audio", "opus") in kinds
samples = 0
with av.open(io.BytesIO(webm)) as container:
for frame in container.decode(audio = 0):
samples += frame.samples
# A full second of 48 kHz audio survived (Opus pads its last 20 ms frame).
assert samples >= 48000, samples
def test_webm_export_still_works_without_an_audio_encoder(monkeypatch):
# A PyAV build with no libopus must keep exporting the video rather than failing the download.
av = pytest.importorskip("av")
import io
real_add_stream = av.container.OutputContainer.add_stream
def _no_opus(self, codec_name = None, *args, **kwargs):
if codec_name == "libopus":
raise ValueError("unknown encoder 'libopus'")
return real_add_stream(self, codec_name, *args, **kwargs)
monkeypatch.setattr(av.container.OutputContainer, "add_stream", _no_opus)
record = gallery.save(_real_mp4_with_audio(), _meta())
webm = gallery.transcode(record["id"], "webm")
assert webm is not None and webm[:4] == b"\x1a\x45\xdf\xa3"
with av.open(io.BytesIO(webm)) as container:
assert [s.type for s in container.streams] == ["video"]
def test_transcode_gif_and_webm_produce_real_containers():
record = gallery.save(_real_mp4_bytes(), _meta())
gif = gallery.transcode(record["id"], "gif")