diff --git a/README.md b/README.md index e4dd67ea..3b18b1e2 100644 --- a/README.md +++ b/README.md @@ -246,12 +246,13 @@ You can also pass `flow_project_id` per project on `POST /api/projects`. ### What does not work on the new API yet -Three capabilities have no captured payload, so they fail with -`UNSUPPORTED_ON_BATCH_API` rather than quietly producing the wrong thing: +1080p export is ported: FlowKit mirrors the Flow UI's `p0UkFb` high-resolution +Download request and polls the resulting media through `as29s`. The remaining +capabilities below still fail loudly rather than quietly producing the wrong thing: | Capability | Status | Workaround | |---|---|---| -| 4K/1080p upscale (`/fk-pipeline` last step) | unported | none — keep the 1080p render | +| 4K export | plan-gated / not live-verified | use 1080p; Google Flow exposes Full HD as the standard high-resolution export | | Reference-to-video (r2v) | unported | `FLOW_ALLOW_DEGRADED=1` → i2v off the first reference | | Start+end-frame chaining (`/fk-gen-chain-videos`) | unported | `FLOW_ALLOW_DEGRADED=1` → i2v off the start frame | | Omni Flash (`model_family=omni_flash`) | unported | use `model_family=veo` | diff --git a/agent/api/flow.py b/agent/api/flow.py index 99394096..eaf63fdd 100644 --- a/agent/api/flow.py +++ b/agent/api/flow.py @@ -64,6 +64,7 @@ class UpscaleVideoRequest(BaseModel): scene_id: str aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT" resolution: str = "VIDEO_RESOLUTION_4K" + project_id: Optional[str] = None class UploadImageRequest(BaseModel): diff --git a/agent/api/upscale_status.py b/agent/api/upscale_status.py new file mode 100644 index 00000000..6f9c0f9f --- /dev/null +++ b/agent/api/upscale_status.py @@ -0,0 +1,87 @@ +"""Explicit Full HD / 4K export endpoints for Google Flow videos.""" + +from typing import Literal + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel + +from agent.services.flow_client import get_flow_client +from agent.services.upscale_polling import annotate_upscale_polling, check_upscale_status + +router = APIRouter(prefix="/flow", tags=["flow"]) + + +class ExportVideoRequest(BaseModel): + media_id: str + scene_id: str = "export" + quality: Literal["1080p", "4k"] = "1080p" + aspect_ratio: str = "VIDEO_ASPECT_RATIO_LANDSCAPE" + project_id: str | None = None + + +class CheckExportStatusRequest(BaseModel): + workflows: list[dict] + + +@router.post("/export-video") +async def export_video(body: ExportVideoRequest): + """Start Google's native Full HD/4K export. + + This is the same Flow upsample operation exposed by the UI, presented as an + export/download-quality choice. 1080p is the default because Omni Flash's + generated file is normally 720p and Full HD is the expected downloadable + master. + """ + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + + resolution = ( + "VIDEO_RESOLUTION_1080P" + if body.quality == "1080p" + else "VIDEO_RESOLUTION_4K" + ) + result = await client.upscale_video( + media_id=body.media_id, + scene_id=body.scene_id, + aspect_ratio=body.aspect_ratio, + resolution=resolution, + project_id=body.project_id, + ) + if result.get("error") or ( + isinstance(result.get("status"), int) and result["status"] >= 400 + ): + raise HTTPException( + result.get("status", 502), + result.get("error", result.get("data")), + ) + + annotated = annotate_upscale_polling(result) + data = annotated.get("data", annotated) + if isinstance(data, dict): + data["export"] = { + "quality": body.quality, + "resolution": resolution, + "native_flow_export": True, + "next": "/api/flow/check-export-status", + } + return data + + +@router.post("/check-export-status") +@router.post("/check-upscale-status", include_in_schema=False) +async def check_export_status(body: CheckExportStatusRequest): + """Return a signed downloadable URL when the native Flow export is ready.""" + client = get_flow_client() + if not client.connected: + raise HTTPException(503, "Extension not connected") + try: + result = await check_upscale_status(body.workflows) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + + if result.get("status") == "COMPLETED": + result["download_ready"] = True + else: + result["download_ready"] = False + return result diff --git a/agent/main.py b/agent/main.py index 9649c807..14b30ce4 100644 --- a/agent/main.py +++ b/agent/main.py @@ -17,6 +17,7 @@ from agent.api.scenes import router as scenes_router from agent.api.requests import router as requests_router from agent.api.flow import router as flow_router +from agent.api.upscale_status import router as upscale_status_router from agent.api.reviews import router as reviews_router from agent.api.tts import router as tts_router from agent.api.materials import router as materials_router @@ -128,6 +129,7 @@ async def lifespan(app: FastAPI): app.include_router(scenes_router, prefix="/api") app.include_router(requests_router, prefix="/api") app.include_router(flow_router, prefix="/api") +app.include_router(upscale_status_router, prefix="/api") app.include_router(reviews_router, prefix="/api") app.include_router(tts_router, prefix="/api") app.include_router(materials_router, prefix="/api") diff --git a/agent/services/flow_batch.py b/agent/services/flow_batch.py index f0111632..33b1b10b 100644 --- a/agent/services/flow_batch.py +++ b/agent/services/flow_batch.py @@ -40,6 +40,7 @@ RPC_PROJECT_MEDIA = "Zzl0ze" RPC_MEDIA = "as29s" RPC_UPLOAD_IMAGE = "maseQ" +RPC_UPSCALE = "p0UkFb" CAPTCHA_IMAGE = "IMAGE_GENERATION" CAPTCHA_VIDEO = "VIDEO_GENERATION" @@ -357,6 +358,23 @@ def video_request(prompt: str, project_id: str, source_media_id: str, return build_envelope(RPC_GEN_VIDEO, inner) +def upscale_request(media_id: str, project_id: str, + aspect: Any = VIDEO_ASPECT_LANDSCAPE, + model: str = "veo_3_1_upsampler_1080p") -> str: + """Build Flow's migrated high-resolution download request (RPC p0UkFb).""" + item = [None] * 32 + item[0] = [None, media_id] + item[2] = 1 + item[4] = [None, str(uuid.uuid4()), None, None, _client_uuid()] + item[6] = resolve_video_aspect(aspect) + item[31] = model + return build_envelope(RPC_UPSCALE, [ + [item], + _context(project_id), + [_client_uuid()], + ]) + + def upload_request(image_b64: str, project_id: str, mime_type: str = "image/jpeg", file_name: str = "upload.jpg") -> str: """Put a local image into the project so it can be used as a reference. @@ -429,6 +447,14 @@ def read_uploaded_media_id(payload: Any) -> str: return media_id +def read_upscaled_media_id(payload: Any) -> str: + """Return the media id created by the p0UkFb upscale submit.""" + for text in _walk_strings(payload): + if isinstance(text, str) and text.endswith("_upsampled"): + return text + raise FlowBatchError("upscale response carried no upsampled media id") + + def read_operation(payload: Any) -> Operation: """`[null, 50, [[opId, projectId, sceneId, status, …]]]`. diff --git a/agent/services/flow_client.py b/agent/services/flow_client.py index 667e2d51..6b20148b 100644 --- a/agent/services/flow_client.py +++ b/agent/services/flow_client.py @@ -694,14 +694,33 @@ async def generate_video_from_references(self, reference_media_ids: list[str], async def upscale_video(self, media_id: str, scene_id: str, aspect_ratio: str = "VIDEO_ASPECT_RATIO_PORTRAIT", - resolution: str = "VIDEO_RESOLUTION_4K") -> dict: - """Upscale a video.""" + resolution: str = "VIDEO_RESOLUTION_4K", + project_id: str | None = None) -> dict: + """Upscale/export a video using Flow's migrated p0UkFb RPC.""" if not USE_BATCH_RPC: return await self._legacy_upscale_video(media_id, scene_id, aspect_ratio, resolution) - return {"error": _unsupported( - "video upscale", - "no upsampler rpc appears in the new frontend's captures", - )} + + model = UPSCALE_MODELS.get(resolution) + if not model: + return {"status": 400, "error": f"Unsupported upscale resolution: {resolution}"} + try: + pid = self._batch_project_id(project_id or "") + freq = fb.upscale_request(media_id, pid, aspect=aspect_ratio, model=model) + payload = await self._batch_payload( + fb.RPC_UPSCALE, freq, fb.CAPTCHA_VIDEO, timeout=120) + upscaled_media_id = fb.read_upscaled_media_id(payload) + except Exception as e: + return _batch_error(e) + + workflow = { + "name": upscaled_media_id, + "primary_media_id": upscaled_media_id, + "project_id": pid, + } + return {"status": 200, "data": { + "media": [{"name": upscaled_media_id}], + "workflows": [workflow], + }} async def check_video_status(self, operations: list[dict]) -> dict: """One poll round for each submitted operation. diff --git a/agent/services/upscale_polling.py b/agent/services/upscale_polling.py new file mode 100644 index 00000000..29d705e5 --- /dev/null +++ b/agent/services/upscale_polling.py @@ -0,0 +1,185 @@ +"""Headless polling for Flow video upscales. + +Google Flow's upsampler returns workflow descriptors whose logical +``primaryMediaId`` may not appear in ``flow.projectInitialData``. The browser UI +can still resolve completed media through ``media.getMediaUrlRedirect``. + +This module exposes a small active poller that treats a successful authenticated +media redirect as the completion signal, avoiding the legacy +``batchCheckAsyncVideoGenerationStatus`` and ``/v1/media/{id}`` paths. +""" + +from __future__ import annotations + +from urllib.parse import quote + +from agent.config import USE_BATCH_RPC +from agent.services.flow_client import get_flow_client + +_ALLOWED_MEDIA_URL_PREFIX = "https://flow-content.google/" + + +def _normalize_workflow(workflow: dict) -> dict | None: + if not isinstance(workflow, dict): + return None + name = workflow.get("name") + primary_media_id = workflow.get("primary_media_id") + if not primary_media_id: + metadata = workflow.get("metadata") + if isinstance(metadata, dict): + primary_media_id = metadata.get("primaryMediaId") + if not isinstance(name, str) or not name: + return None + if not isinstance(primary_media_id, str) or not primary_media_id: + return None + return {"name": name, "primary_media_id": primary_media_id} + + +def extract_upscale_workflows(result: dict) -> list[dict]: + if not isinstance(result, dict): + return [] + data = result.get("data") if isinstance(result.get("data"), dict) else result + workflows = data.get("workflows", []) if isinstance(data, dict) else [] + normalized = [] + for workflow in workflows: + item = _normalize_workflow(workflow) + if item: + normalized.append(item) + return normalized + + +def annotate_upscale_polling(result: dict) -> dict: + workflows = extract_upscale_workflows(result) + if not workflows: + return result + data = result.get("data") if isinstance(result.get("data"), dict) else result + if isinstance(data, dict): + data["flowkitPolling"] = { + "mode": "media_redirect", + "workflows": workflows, + } + return result + + +async def _fetch_media_url(client, media_id: str) -> dict: + if USE_BATCH_RPC: + result = await client.get_media(media_id) + data = result.get("data") if isinstance(result.get("data"), dict) else {} + video = data.get("video") if isinstance(data, dict) else None + candidate = video.get("fifeUrl") if isinstance(video, dict) else None + return { + "status": result.get("status", 200), + "data": { + "url": candidate, + "contentType": "video/mp4" if candidate else None, + }, + "error": result.get("error"), + } + + url = ( + "https://labs.google/fx/api/trpc/media.getMediaUrlRedirect" + f"?name={quote(media_id, safe='')}" + ) + return await client._send( + "trpc_request", + { + "url": url, + "method": "GET", + "headers": {"content-type": "application/json"}, + "responseMode": "url", + }, + timeout=15, + ) + + +def _parse_media_redirect(response: dict) -> tuple[str | None, str | None, str | None]: + if not isinstance(response, dict): + return None, None, "Flow media redirect returned an invalid response" + status = response.get("status") + data = response.get("data") if isinstance(response.get("data"), dict) else {} + candidate = data.get("url") + content_type = data.get("contentType") + if ( + isinstance(status, int) + and status < 400 + and isinstance(candidate, str) + and candidate.startswith(_ALLOWED_MEDIA_URL_PREFIX) + ): + return candidate, content_type if isinstance(content_type, str) else None, None + error = response.get("error") + if not error and isinstance(status, int) and status >= 400: + error = f"API_{status}" + if not error: + error = "media redirect not ready" + return None, content_type if isinstance(content_type, str) else None, str(error) + + +async def check_upscale_status( + workflows: list[dict], + include_encoded_video: bool = False, +) -> dict: + """Poll native Flow Full HD/4K export workflows without buffering the MP4.""" + normalized = [] + for workflow in workflows or []: + item = _normalize_workflow(workflow) + if item: + normalized.append(item) + if not normalized: + raise ValueError( + "Export polling requires workflow descriptors with name and " + "primary_media_id (or raw Flow metadata.primaryMediaId)" + ) + + client = get_flow_client() + items = [] + for workflow in normalized: + media_id = workflow["primary_media_id"] + response = await _fetch_media_url(client, media_id) + url, content_type, diagnostic = _parse_media_redirect(response) + if url: + media = { + "media_id": media_id, + "url": url, + "encoded_video_available": False, + "resolved_via": "as29s" if USE_BATCH_RPC else "media.getMediaUrlRedirect", + } + if content_type: + media["content_type"] = content_type + if include_encoded_video: + media["encoded_video"] = None + items.append({ + "name": workflow["name"], + "primary_media_id": media_id, + "done": True, + "status": "MEDIA_GENERATION_STATUS_SUCCESSFUL", + "error": None, + "media": media, + }) + continue + + probe = {} + if isinstance(response, dict): + if isinstance(response.get("status"), int): + probe["http_status"] = response["status"] + data = response.get("data") + if isinstance(data, dict) and isinstance(data.get("url"), str): + probe["resolved_url"] = data["url"] + if diagnostic: + probe["diagnostic"] = diagnostic + item = { + "name": workflow["name"], + "primary_media_id": media_id, + "done": False, + "status": "PENDING", + "error": None, + } + if probe: + item["probe"] = probe + items.append(item) + + all_done = bool(items) and all(item["done"] for item in items) + return { + "done": all_done, + "status": "COMPLETED" if all_done else "PENDING", + "workflows": items, + } diff --git a/docs/VIDEO_EXPORTS.md b/docs/VIDEO_EXPORTS.md new file mode 100644 index 00000000..de056793 --- /dev/null +++ b/docs/VIDEO_EXPORTS.md @@ -0,0 +1,50 @@ +# Video exports + +FlowKit treats output quality as an explicit export/download choice rather than +an implementation detail. + +## Omni Flash + +Gemini Omni Flash generation normally produces a 720p source video. Google Flow +provides a native Full HD export for that result. FlowKit exposes it directly as +**Export 1080p**; internally Google names the operation an upsample, but callers +do not need to reason about that implementation detail. + +## Export Full HD (recommended) + +```bash +curl -sS -X POST http://127.0.0.1:8100/api/flow/export-video \ + -H 'Content-Type: application/json' \ + -d '{ + "media_id": "", + "scene_id": "job-1", + "quality": "1080p", + "aspect_ratio": "VIDEO_ASPECT_RATIO_LANDSCAPE" + }' +``` + +`quality` defaults to `1080p`. `4k` remains an explicit API option, but it is plan-gated by Google Flow; the currently verified account exposes 4K as disabled while 1080p is available. + +The response contains `flowkitPolling.workflows`. Poll those descriptors: + +```bash +curl -sS -X POST http://127.0.0.1:8100/api/flow/check-export-status \ + -H 'Content-Type: application/json' \ + -d '{"workflows": }' +``` + +When `download_ready` becomes `true`, use +`workflows[].media.url` immediately. It is a short-lived signed +`flow-content.google` URL. + +## Compatibility + +The older endpoints remain supported: + +- `POST /api/flow/upscale-video` +- `POST /api/flow/check-upscale-status` + +They are aliases/low-level surfaces for the same Google Flow capability. New +integrations should use `export-video` and `check-export-status`, because those +names describe the user-visible operation: selecting the downloadable output +quality. diff --git a/tests/unit/test_flow_batch.py b/tests/unit/test_flow_batch.py index 63b8761d..e5e5c902 100644 --- a/tests/unit/test_flow_batch.py +++ b/tests/unit/test_flow_batch.py @@ -121,6 +121,21 @@ def test_a_hand_reframed_crop_overrides_the_default(self): crop = [None, 0.1, 1, 0.9] assert inner(fb.video_request("go", self.PID, "mid", crop=crop))[0][0][4][5] == crop + def test_upscale_matches_the_captured_1080p_slots(self): + payload = inner(fb.upscale_request( + "media-1", self.PID, + aspect="VIDEO_ASPECT_RATIO_LANDSCAPE", + model="veo_3_1_upsampler_1080p", + )) + item = payload[0][0] + assert len(item) == 32 + assert item[0] == [None, "media-1"] + assert item[2] == 1 + assert item[6] == fb.VIDEO_ASPECT_LANDSCAPE + assert item[31] == "veo_3_1_upsampler_1080p" + assert payload[1][5] == self.PID + assert fb.CAPTCHA_SLOT in json.dumps(payload) + class TestReaders: OP = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" @@ -135,6 +150,9 @@ def test_a_repeated_url_is_not_a_second_variant(self): url = f"https://{fb.MEDIA_HOST}/image/{self.MID}?sig=x" assert len(fb.read_images([url, url])) == 1 + def test_upscale_submit_reads_the_new_media_id(self): + assert fb.read_upscaled_media_id([[self.MID + "_upsampled"]]) == self.MID + "_upsampled" + def test_operation_reads_the_id_and_status(self): op = fb.read_operation([None, 50, [[self.OP, "proj", "scene", "CAE"]]]) assert (op.operation_id, op.status, op.done) == (self.OP, "CAE", True) diff --git a/tests/unit/test_flow_client_batch.py b/tests/unit/test_flow_client_batch.py index 733f5244..0b1b5b95 100644 --- a/tests/unit/test_flow_client_batch.py +++ b/tests/unit/test_flow_client_batch.py @@ -158,11 +158,26 @@ async def test_degraded_r2v_uses_the_first_reference_as_the_start_frame(self, cl payload = json.loads(json.loads(client.calls[0]["freq"])[0][0][1]) assert payload[0][0][4][1] == "ref-a" - async def test_upscale_is_unported_and_has_no_fallback(self, client, monkeypatch): - import agent.services.flow_client as module - monkeypatch.setattr(module, "FLOW_ALLOW_DEGRADED", True) - result = await client.upscale_video(MEDIA, "scene-1") - assert "UNSUPPORTED_ON_BATCH_API" in result["error"] + + +class TestUpscaleVideo: + async def test_1080p_submit_returns_a_pollable_workflow(self, client): + upscaled = MEDIA + "_upsampled" + client.responses[fb.RPC_UPSCALE] = { + "data": envelope(fb.RPC_UPSCALE, [[[[upscaled], "", None, None, 1]]]) + } + + result = await client.upscale_video( + MEDIA, + "scene-1", + aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE", + resolution="VIDEO_RESOLUTION_1080P", + ) + workflow = result["data"]["workflows"][0] + assert workflow["primary_media_id"] == upscaled + assert workflow["project_id"] == PROJECT + assert client.calls[0]["rpcid"] == fb.RPC_UPSCALE + assert client.calls[0]["captcha"] == fb.CAPTCHA_VIDEO class TestCheckVideoStatus: diff --git a/tests/unit/test_upscale_polling.py b/tests/unit/test_upscale_polling.py new file mode 100644 index 00000000..14826628 --- /dev/null +++ b/tests/unit/test_upscale_polling.py @@ -0,0 +1,67 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import agent.services.upscale_polling as polling +from agent.services.upscale_polling import annotate_upscale_polling, check_upscale_status + + +MEDIA = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + + +def test_annotate_upscale_polling_uses_workflow_primary_media_id(): + result = { + "status": 200, + "data": { + "workflows": [ + {"name": "wf-1", "metadata": {"primaryMediaId": MEDIA + "_upsampled"}} + ] + }, + } + annotated = annotate_upscale_polling(result) + assert annotated["data"]["flowkitPolling"] == { + "mode": "media_redirect", + "workflows": [ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ], + } + + +@pytest.mark.asyncio +async def test_batch_poll_resolves_completed_1080_media_through_as29s(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + client = MagicMock() + client.get_media = AsyncMock(return_value={ + "status": 200, + "data": { + "video": { + "fifeUrl": "https://flow-content.google/video/out?Signature=test" + } + }, + }) + with patch("agent.services.upscale_polling.get_flow_client", return_value=client): + result = await check_upscale_status([ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ]) + + assert result["done"] is True + assert result["status"] == "COMPLETED" + media = result["workflows"][0]["media"] + assert media["resolved_via"] == "as29s" + assert media["url"].startswith("https://flow-content.google/video/") + client.get_media.assert_awaited_once_with(MEDIA + "_upsampled") + + +@pytest.mark.asyncio +async def test_batch_poll_stays_pending_until_video_url_exists(monkeypatch): + monkeypatch.setattr(polling, "USE_BATCH_RPC", True) + client = MagicMock() + client.get_media = AsyncMock(return_value={"status": 200, "data": {"video": {}}}) + with patch("agent.services.upscale_polling.get_flow_client", return_value=client): + result = await check_upscale_status([ + {"name": "wf-1", "primary_media_id": MEDIA + "_upsampled"} + ]) + + assert result["done"] is False + assert result["status"] == "PENDING" + assert result["workflows"][0]["status"] == "PENDING"