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:
parent
6a6b6c4c13
commit
bce3b9b20a
8 changed files with 381 additions and 21 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue