From 089bcd2879301291e7d3d57194691800a3c6a101 Mon Sep 17 00:00:00 2001 From: pinokio Date: Fri, 4 Sep 2026 03:05:53 +0200 Subject: [PATCH 1/5] feat(tools): upscale still images through shared flow --- app/_launch_runtime.py | 730 +++++++++++++++--- app/docs/API.md | 18 + app/services/generation_provenance.py | 2 + .../tools/background_removal_request.py | 72 +- app/wgp.py | 11 +- tests/test_tools_upscale_contract.py | 197 +++++ ui/e2e/helpers/apiRoutes.ts | 140 +++- ui/e2e/specs/tools-background-removal.spec.ts | 72 +- ui/src/api/generation.ts | 25 +- ui/src/components/Sidebar/ToolsPanel.tsx | 31 +- .../components/Sidebar/ToolsSourcePanel.tsx | 127 ++- ui/src/i18n/locales/en/studio.json | 11 +- ui/src/i18n/locales/es/studio.json | 11 +- ui/src/stores/useStore.ts | 13 +- ui/tests/toolsPanel.test.tsx | 42 +- 15 files changed, 1319 insertions(+), 183 deletions(-) create mode 100644 tests/test_tools_upscale_contract.py diff --git a/app/_launch_runtime.py b/app/_launch_runtime.py index 19850d944..aebd8796f 100644 --- a/app/_launch_runtime.py +++ b/app/_launch_runtime.py @@ -822,6 +822,11 @@ def _public_generation_details(params: dict | None) -> dict: "generation_mode": params.get("generation_mode"), "execution_mode": params.get("_execution_mode"), "simulated": params.get("_execution_mode") == "simulate", + "capability": params.get("capability"), + "source_kind": params.get("source_kind"), + "source_asset_id": params.get("source_asset_id"), + "source_filename": params.get("source_filename"), + "method": params.get("method"), "resolution": params.get("resolution"), "seed": params.get("seed"), "steps": params.get("num_inference_steps"), @@ -22721,6 +22726,188 @@ def _apply_spatial_upsampling_to_file(video_path: str, method: str, job: dict = # See memory/project_tools_postprocessing.md. # ============================================================================ + +_TOOL_UPSCALE_METHODS = frozenset({ + "flashvsr2", "flashvsr3", "flashvsr4", "flashvsr2pass2", + "flashvsr2pass4", "lanczos1.5", "lanczos2", +}) +_TOOL_SOURCE_EXTENSIONS = { + "image": frozenset({ + ".bmp", ".gif", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp", + }), + "video": frozenset({ + ".avi", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm", ".wmv", + }), +} + + +def _tool_asset_roots() -> list[dict[str, str]]: + """Return every explicit root that the Tools asset catalog can own.""" + roots: list[dict[str, str]] = [] + seen: set[str] = set() + for item in _list_workspaces(): + if not isinstance(item, dict): + continue + workspace = str(item.get("name") or "").strip() + if not workspace or workspace in seen: + continue + try: + directory = _workspace_dir(workspace) + except Exception: + continue + roots.append({"workspace_id": workspace, "path": directory}) + seen.add(workspace) + uploads = os.path.realpath(os.path.abspath(os.path.join(os.getcwd(), "uploads"))) + if uploads not in {os.path.realpath(os.path.abspath(item["path"])) for item in roots}: + roots.append({"workspace_id": "__uploads__", "path": uploads}) + return roots + + +def _tool_source_kind_from_name(value: str) -> str | None: + from urllib.parse import unquote, urlsplit + + raw = str(value or "").strip() + if not raw: + return None + parsed = urlsplit(raw) + candidate = unquote( + parsed.path + if parsed.scheme or raw.startswith("/api/v1/") + else raw.split("?", 1)[0].split("#", 1)[0] + ) + extension = os.path.splitext(os.path.basename(candidate))[1].casefold() + for kind, extensions in _TOOL_SOURCE_EXTENSIONS.items(): + if extension in extensions: + return kind + return None + + +def _resolve_tool_source( + body: dict, + *, + expected_kinds: tuple[str, ...] = ("image", "video"), +) -> tuple[str, str, str, str, str, str, str]: + """Resolve one exact Tools source and its destination context. + + Source resolution is shared with the background-removal boundary so + canonical asset IDs, workspace-qualified file URLs, uploads, symlinks and + traversal attempts receive the same confinement rules. The return tuple + is ``path, filename, source_workspace, kind, asset_id, destination, + output_dir``. + """ + from shared.tools.background_removal_request import ( + RemoveBackgroundRequest, + UPSCALE_IMAGE_EXTENSIONS, + UPSCALE_VIDEO_EXTENSIONS, + destination_context, + resolve_source, + ) + from services.asset_catalog import find_asset + + if not isinstance(body, dict): + raise HTTPException(status_code=400, detail="Request body must be an object") + + for field, limit in ( + ("source", 1200), ("source_path", 1200), ("video_path", 1200), + ("source_workspace", 160), ("workspace", 160), + ): + value = body.get(field) + if value is not None and ( + not isinstance(value, str) + or len(value) > limit + or "\x00" in value + ): + raise HTTPException(status_code=400, detail=f"Invalid {field}") + + source_values = [ + body.get(key) for key in ("source", "source_path", "video_path") + if body.get(key) not in (None, "") + ] + if len(source_values) > 1 and any(str(value) != str(source_values[0]) for value in source_values[1:]): + raise HTTPException(status_code=409, detail="Conflicting source paths") + source_value = str(source_values[0]).strip() if source_values else None + + asset_id = body.get("asset_id") + if asset_id is not None: + if ( + not isinstance(asset_id, str) + or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:-]{0,179}", asset_id) + ): + raise HTTPException(status_code=400, detail="Invalid asset ID") + + requested_kind = body.get("source_kind") + if requested_kind is not None: + if not isinstance(requested_kind, str) or requested_kind.casefold() not in _TOOL_SOURCE_EXTENSIONS: + raise HTTPException(status_code=400, detail="Source kind must be image or video") + requested_kind = requested_kind.casefold() + + asset = None + if asset_id: + asset = find_asset(_tool_asset_roots(), asset_id) + if asset is None: + raise HTTPException(status_code=404, detail="Source asset not found") + asset_kind = str(asset.get("kind") or "").casefold() + if asset_kind not in _TOOL_SOURCE_EXTENSIONS: + raise HTTPException(status_code=400, detail="Source asset must be an image or video") + if requested_kind and requested_kind != asset_kind: + raise HTTPException(status_code=400, detail="Source kind does not match asset") + source_kind = requested_kind or asset_kind + else: + source_kind = requested_kind or _tool_source_kind_from_name(source_value or "") + if source_kind is None: + raise HTTPException(status_code=400, detail="Source format is not supported") + + if source_kind not in expected_kinds: + expected = " or ".join(expected_kinds) + raise HTTPException(status_code=400, detail=f"Source must be {expected}") + allowed_extensions = ( + UPSCALE_IMAGE_EXTENSIONS + if source_kind == "image" else UPSCALE_VIDEO_EXTENSIONS + ) + if source_value and _tool_source_kind_from_name(source_value) != source_kind: + raise HTTPException(status_code=400, detail="Source kind does not match file format") + + payload = RemoveBackgroundRequest( + asset_id=asset_id, + source=source_value, + source_workspace=body.get("source_workspace"), + workspace=body.get("workspace"), + provenance=body.get("provenance") if isinstance(body.get("provenance"), dict) else {}, + ) + destination_workspace, output_dir = destination_context( + payload, + get_active_workspace=_get_active_workspace, + list_workspaces=_list_workspaces, + workspace_dir=_workspace_dir, + ) + source_path, source_filename, source_workspace = resolve_source( + payload, + destination_workspace=destination_workspace, + workspace_dir=_workspace_dir, + uploads_dir=lambda: os.path.join(os.getcwd(), "uploads"), + asset_finder=lambda value: find_asset(_tool_asset_roots(), value), + expected_kind=source_kind, + allowed_extensions=allowed_extensions, + source_label=source_kind, + ) + source_asset_id = asset_id or ( + "asset_unmanaged_" + + uuid.uuid5( + uuid.NAMESPACE_URL, + f"hocuspocus:unmanaged:{source_workspace}:{source_filename}", + ).hex + ) + return ( + source_path, + source_filename, + source_workspace, + source_kind, + source_asset_id, + destination_workspace, + output_dir, + ) + + def _resolve_tool_clip_path(raw_path, workspace=None): """Resolve a Tools input path. Accepts an absolute path, a filename in the active workspace output dir, or a name/relative path under uploads/. @@ -22741,6 +22928,61 @@ def _resolve_tool_clip_path(raw_path, workspace=None): return raw_path if os.path.isfile(raw_path) else None +def _upscale_tool_image( + source_path: str, + output_path: str, + method: str, + *, + seed: int = -1, + abort_callback=None, + progress_callback=None, +) -> tuple[int, int]: + """Upscale one still image through the existing spatial adapter. + + The FlashVSR bridge accepts the same ``[C, F, H, W]`` tensor used by the + video path. ``still_image=True`` tells that adapter to use its image + inference path; Lanczos remains the stateless implementation in WGP. No + video decoder, audio extractor, or video writer is involved here. + """ + from PIL import Image + from shared.utils.utils import convert_image_to_tensor, convert_tensor_to_image + + if callable(abort_callback) and abort_callback(): + raise InterruptedError("Image upscale was cancelled") + temporary_path = f"{output_path}.tmp-{uuid.uuid4().hex}" + try: + with Image.open(source_path) as opened: + image = wgp.convert_image(opened).copy() + sample = convert_image_to_tensor(image).unsqueeze(1) + if callable(progress_callback): + progress_callback("Upscaling image", 5, 0, 1) + sample = wgp.perform_spatial_upsampling( + sample, + method, + seed=seed, + abort_callback=abort_callback, + progress_callback=progress_callback, + still_image=True, + ) + if callable(abort_callback) and abort_callback(): + raise InterruptedError("Image upscale was cancelled") + if sample is None: + if callable(abort_callback) and abort_callback(): + raise InterruptedError("Image upscale was cancelled") + raise RuntimeError("Image upsampler returned no result") + result = convert_tensor_to_image(sample, 0).convert("RGB") + result.save(temporary_path, format="PNG") + os.replace(temporary_path, output_path) + return result.size + except Exception: + try: + if os.path.isfile(temporary_path): + os.remove(temporary_path) + except OSError: + pass + raise + + def _write_tool_sidecar( out_dir, filename, @@ -22753,39 +22995,78 @@ def _write_tool_sidecar( task_id=None, root_task_id=None, workspace=None, + generation_mode="video", + source_asset_id=None, + source_kind=None, + provenance=None, + inputs=None, + parents=None, + transformations=None, + technical=None, ): """Write a .meta.json sidecar so a Tools output shows up in the gallery with the right mode + edit_sub_mode tag (mirrors _run_sfx_generation).""" + public_params = { + key: value + for key, value in (params or {}).items() + if not str(key).startswith("_") + and key not in {"source_path", "video_path", "voice_ref_paths"} + } + if source_asset_id: + public_params.setdefault("source_asset_id", source_asset_id) + if source_kind: + public_params.setdefault("source_kind", source_kind) + public_params.setdefault("source_filename", source_name) + provenance = provenance if isinstance(provenance, dict) else {} + capability = provenance.get("capability") or tool sidecar = { - "params": {**params, "edit_sub_mode": tool}, - "generation_mode": "video", + "params": {**public_params, "edit_sub_mode": tool}, + "generation_mode": generation_mode, "tool": tool, + "capability": capability, "tool_source": source_name, + "source_asset_id": source_asset_id, + "source_kind": source_kind, "job_id": job_id, "task_id": task_id, "root_task_id": root_task_id or task_id, "generation_time": round(elapsed), "created_at": time.time(), } + if inputs: + sidecar["inputs"] = list(inputs) + if parents: + sidecar["parents"] = list(parents) + if transformations: + sidecar["transformations"] = list(transformations) + if technical: + sidecar["technical"] = dict(technical) try: publish_generation_sidecar( os.path.join(out_dir, filename), sidecar, - workspace_id=workspace, - tool=tool, + workspace_id=provenance.get("workspace_id") or workspace, + output_folder=workspace, + tool=provenance.get("tool") or tool, + actor=provenance.get("actor"), + capability=capability, ) except Exception: pass def _run_tool_upscale(job_id: str): - """Background worker: upscale an existing clip with the configured spatial - upsampler (FlashVSR / Lanczos), preserving the original audio. Thin extract - of edit_video's postprocessing path — no model generation/Gradio state.""" + """Background worker for the shared image/video upscale action. + + Video sources retain the existing audio-preserving path. Still images are + dispatched before any video metadata/decoder/audio calls and use the same + spatial adapter with its ``still_image`` mode. + """ job = _jobs[job_id] start_time = None abort_state = {"abort": False} audio_tracks = [] + final_path = None with _coordinated_generation_slot( job, description="HocusPocus Lab GPU tool · upscale", ) as acquired: @@ -22809,29 +23090,28 @@ def _run_tool_upscale(job_id: str): wgp.save_path = out_dir method = params.get("method") or "flashvsr2" - video_source = _resolve_tool_clip_path(params.get("video_path"), workspace) - if not video_source: + if method not in _TOOL_UPSCALE_METHODS: + raise ValueError("Unsupported upscale method") + source_kind = str(params.get("source_kind") or "video").casefold() + source_value = ( + params.get("source_path") + if source_kind == "image" else params.get("video_path") + ) or params.get("_source_path") or params.get("source") + source_path = _resolve_tool_clip_path(source_value, workspace) + if not source_path: finish_job( - job, "failed", error="Input clip not found", - message="Error: input clip not found", + job, "failed", error="Input source not found", + message="Error: input source not found", ) return False + if source_kind not in _TOOL_SOURCE_EXTENSIONS: + raise ValueError("Unsupported source kind") + source_extension = os.path.splitext(source_path)[1].casefold() + if source_extension not in _TOOL_SOURCE_EXTENSIONS[source_kind]: + raise ValueError("Source kind does not match file format") before = set(os.listdir(out_dir)) if os.path.isdir(out_dir) else set() - from shared.utils.utils import get_video_info - fps, _width, _height, _frames = get_video_info(video_source) - - # Preserve original audio — re-muxed onto the upscaled video. - audio_tracks, audio_metadata = wgp.extract_audio_tracks(video_source) - has_audio = len(audio_tracks) > 0 - - if not update_job( - job, message="Upscaling...", phase="Upscaling", progress=5, - ): - wgp.cleanup_temp_audio_files(audio_tracks) - return False - def _abort(): return bool(abort_state.get("abort")) or is_cancel_requested(job) @@ -22856,77 +23136,173 @@ def _progress(phase, current_step=None, total_steps=None): if changes: update_job(job, **changes) - container = wgp.server_config.get("video_container", "mp4") - codec = wgp.server_config.get("video_output_codec", None) - final_path = wgp.get_available_filename(out_dir, os.path.basename(video_source), "_upscaled", force_extension=f".{container}") - - if wgp.flashvsr.is_upsampling(method): - # Chunked engine (shared with the post-generation pass) — - # bounds RAM on long clips. The previous unchunked path let - # FlashVSR allocate its float32 output buffer for the WHOLE - # video: a 4-minute 2x upscale tried 280+ GB and died in - # DefaultCPUAllocator. - tmp_path = _chunked_flashvsr_upscale(video_source, method, job=job, abort_check=_abort, progress_callback=_progress) - if tmp_path is None or _abort(): - if tmp_path and os.path.isfile(tmp_path): + source_filename = str( + params.get("source_filename") or os.path.basename(source_path) + ) + final_path = None + image_size = None + if source_kind == "image": + if not update_job( + job, message="Upscaling image...", phase="Upscaling", progress=5, + ): + return False + final_path = wgp.get_available_filename( + out_dir, source_filename, "_upscaled", force_extension=".png", + ) + image_size = _upscale_tool_image( + source_path, + final_path, + method, + seed=int(params.get("seed", -1)), + abort_callback=_abort, + progress_callback=_progress, + ) + else: + from shared.utils.utils import get_video_info + fps, _width, _height, _frames = get_video_info(source_path) + + # Preserve original audio — re-muxed onto the upscaled video. + audio_tracks, audio_metadata = wgp.extract_audio_tracks(source_path) + has_audio = len(audio_tracks) > 0 + + if not update_job( + job, message="Upscaling...", phase="Upscaling", progress=5, + ): + wgp.cleanup_temp_audio_files(audio_tracks) + return False + + container = wgp.server_config.get("video_container", "mp4") + codec = wgp.server_config.get("video_output_codec", None) + final_path = wgp.get_available_filename( + out_dir, source_filename, "_upscaled", + force_extension=f".{container}", + ) + + if wgp.flashvsr.is_upsampling(method): + # Chunked engine (shared with the post-generation pass) — + # bounds RAM on long clips. The previous unchunked path let + # FlashVSR allocate its float32 output buffer for the WHOLE + # video: a 4-minute 2x upscale tried 280+ GB and died in + # DefaultCPUAllocator. + tmp_path = _chunked_flashvsr_upscale( + source_path, method, job=job, abort_check=_abort, + progress_callback=_progress, + ) + if tmp_path is None or _abort(): + if tmp_path and os.path.isfile(tmp_path): + try: + os.remove(tmp_path) + except OSError: + pass + wgp.cleanup_temp_audio_files(audio_tracks) + return False + if has_audio: + wgp.combine_video_with_audio_tracks( + tmp_path, audio_tracks, final_path, + audio_metadata=audio_metadata, + ) try: os.remove(tmp_path) except OSError: pass - wgp.cleanup_temp_audio_files(audio_tracks) - return False - if has_audio: - wgp.combine_video_with_audio_tracks(tmp_path, audio_tracks, final_path, audio_metadata=audio_metadata) - try: - os.remove(tmp_path) - except OSError: - pass - wgp.cleanup_temp_audio_files(audio_tracks) + wgp.cleanup_temp_audio_files(audio_tracks) + else: + os.replace(tmp_path, final_path) else: - os.replace(tmp_path, final_path) - else: - # Lanczos & friends — cheap stateless resize, legacy inline path. - sample = wgp.get_resampled_video(video_source, 0, wgp.max_source_video_frames, fps) - sample = sample.permute(-1, 0, 1, 2) # [F,H,W,C] -> [C,F,H,W] - sample = wgp.perform_spatial_upsampling( - sample, method, seed=int(params.get("seed", -1)), - abort_callback=_abort, progress_callback=_progress, - ) + # Lanczos & friends — cheap stateless resize, legacy inline path. + sample = wgp.get_resampled_video( + source_path, 0, wgp.max_source_video_frames, fps, + ) + sample = sample.permute(-1, 0, 1, 2) # [F,H,W,C] -> [C,F,H,W] + sample = wgp.perform_spatial_upsampling( + sample, method, seed=int(params.get("seed", -1)), + abort_callback=_abort, progress_callback=_progress, + ) - if _abort(): - return False + if _abort(): + return False - output_fps = round(fps) - if has_audio: - tmp_path = wgp.get_available_filename(out_dir, os.path.basename(video_source), "_uptmp", force_extension=f".{container}") - wgp.save_video(tensor=sample[None], save_file=tmp_path, fps=output_fps, nrow=1, normalize=True, value_range=(-1, 1), codec_type=codec, container=container) - wgp.combine_video_with_audio_tracks(tmp_path, audio_tracks, final_path, audio_metadata=audio_metadata) - try: - os.remove(tmp_path) - except OSError: - pass - wgp.cleanup_temp_audio_files(audio_tracks) - else: - wgp.save_video(tensor=sample[None], save_file=final_path, fps=output_fps, nrow=1, normalize=True, value_range=(-1, 1), codec_type=codec, container=container) + output_fps = round(fps) + if has_audio: + tmp_path = wgp.get_available_filename( + out_dir, source_filename, "_uptmp", + force_extension=f".{container}", + ) + wgp.save_video( + tensor=sample[None], save_file=tmp_path, + fps=output_fps, nrow=1, normalize=True, + value_range=(-1, 1), codec_type=codec, container=container, + ) + wgp.combine_video_with_audio_tracks( + tmp_path, audio_tracks, final_path, + audio_metadata=audio_metadata, + ) + try: + os.remove(tmp_path) + except OSError: + pass + wgp.cleanup_temp_audio_files(audio_tracks) + else: + wgp.save_video( + tensor=sample[None], save_file=final_path, fps=output_fps, + nrow=1, normalize=True, value_range=(-1, 1), + codec_type=codec, container=container, + ) - sample = None + sample = None after = set(os.listdir(out_dir)) if os.path.isdir(out_dir) else set() new_files = sorted(f for f in (after - before) if not f.endswith(".meta.json") and "_uptmp" not in f) - record_job_outputs(job, new_files) if is_cancel_requested(job): + # A cancellation can win just after the adapter commits its + # bytes. Do not expose that late result in Activity or leave a + # derived artifact behind when the source was never published. + for fname in new_files: + try: + os.remove(os.path.join(out_dir, fname)) + except OSError: + pass return False + record_job_outputs(job, new_files) + source_asset_id = params.get("source_asset_id") + source_ref = { + "id": source_asset_id, + "kind": source_kind, + "uri": source_filename, + "role": "source", + } for fname in new_files: _write_tool_sidecar( out_dir, fname, - source_name=os.path.basename(video_source), + source_name=source_filename, tool="upscale", - params={"method": method, "model_type": "post_processing"}, + params={ + "method": method, + "model_type": "post_processing", + "source_asset_id": source_asset_id, + "source_kind": source_kind, + "source_filename": source_filename, + }, elapsed=time.time() - start_time, job_id=job_id, task_id=job.get("task_id"), root_task_id=job.get("root_task_id"), workspace=job.get("workspace"), + generation_mode=source_kind, + source_asset_id=source_asset_id, + source_kind=source_kind, + provenance=job.get("provenance"), + inputs=[source_ref] if source_asset_id else [], + parents=[source_ref] if source_asset_id else [], + transformations=[{ + "type": "upscale", + "backend": "flashvsr" if method.startswith("flashvsr") else "lanczos", + "method": method, + }], + technical=( + {"width": image_size[0], "height": image_size[1], "output": "png"} + if image_size else {"output": "video"} + ), ) completed = finish_job( @@ -22936,8 +23312,16 @@ def _progress(phase, current_step=None, total_steps=None): phase="", message="Done", ) - print(f"[Tools/upscale] {os.path.basename(video_source)} -> {new_files} ({wgp.format_time(time.time() - start_time)})") + print(f"[Tools/upscale] {source_filename} -> {new_files} ({wgp.format_time(time.time() - start_time)})") return completed + except InterruptedError: + if final_path and os.path.isfile(final_path): + try: + os.remove(final_path) + except OSError: + pass + acknowledge_cancel(job) + return False except Exception as e: traceback.print_exc() finish_job(job, "failed", error=str(e), message=f"Error: {e}") @@ -22991,6 +23375,13 @@ def _run_tool_revoice(job_id: str): message="Error: input clip not found", ) return False + if str(params.get("source_kind") or "video").casefold() != "video": + raise ValueError("Revoice only supports video sources") + if os.path.splitext(video_source)[1].casefold() not in _TOOL_SOURCE_EXTENSIONS["video"]: + raise ValueError("Revoice source format is not supported") + source_filename = str( + params.get("source_filename") or os.path.basename(video_source) + ) mode = params.get("mode", "single") voice_refs = [] @@ -23066,14 +23457,41 @@ def _run_tool_revoice(job_id: str): _write_tool_sidecar( out_dir, fname, - source_name=os.path.basename(video_source), + source_name=source_filename, tool="revoice", - params={"mode": mode, "model_type": "post_processing"}, + params={ + "mode": mode, + "model_type": "post_processing", + "source_asset_id": params.get("source_asset_id"), + "source_kind": "video", + "source_filename": source_filename, + }, elapsed=time.time() - start_time, job_id=job_id, task_id=job.get("task_id"), root_task_id=job.get("root_task_id"), workspace=job.get("workspace"), + generation_mode="video", + source_asset_id=params.get("source_asset_id"), + source_kind="video", + provenance=job.get("provenance"), + inputs=[{ + "id": params.get("source_asset_id"), + "kind": "video", + "uri": source_filename, + "role": "source", + }], + parents=[{ + "id": params.get("source_asset_id"), + "kind": "video", + "uri": source_filename, + "role": "source", + }], + transformations=[{ + "type": "revoice", + "backend": "seedvc", + "mode": mode, + }], ) completed = finish_job( @@ -23095,34 +23513,82 @@ def _run_tool_revoice(job_id: str): @api.post("/api/v1/tools/upscale") async def tools_upscale(request: Request): - """Upscale an existing clip (a gallery output or an uploaded file) with the - configured spatial upsampler. Returns a job_id; poll /api/v1/status/{job_id}. + """Upscale one image or video through the shared Tools action. - Body: { video_path: str, method?: str (default "flashvsr2"), - seed?: int, workspace?: str } + Images use ``source``/``source_kind=image`` and produce a PNG. Videos may + continue sending the legacy ``video_path`` field and retain the existing + audio-preserving pipeline. """ body = await request.json() - video_path = body.get("video_path") - if not video_path: - raise HTTPException(status_code=400, detail="video_path is required") - workspace = body.get("workspace") or _get_active_workspace() - resolved = _resolve_tool_clip_path(video_path, workspace) - if not resolved: - raise HTTPException(status_code=400, detail=f"Clip not found: {video_path}") + if not isinstance(body, dict): + raise HTTPException(status_code=400, detail="Request body must be an object") + method = body.get("method") or "flashvsr2" + if not isinstance(method, str) or method not in _TOOL_UPSCALE_METHODS: + raise HTTPException(status_code=400, detail="Unsupported upscale method") + ( + resolved, + source_filename, + source_workspace, + source_kind, + source_asset_id, + workspace, + output_dir, + ) = _resolve_tool_source(body, expected_kinds=("image", "video")) + from services.generation_provenance import normalize_submission_provenance + + provenance = normalize_submission_provenance({ + **(body.get("provenance") if isinstance(body.get("provenance"), dict) else {}), + "capability": "upscale", + # Workspace identity is runtime-owned, not browser-authored. + "workspace_id": workspace, + }) + source_ref = { + "id": source_asset_id, + "kind": source_kind, + "uri": source_filename, + "role": "source", + } + transformation = { + "type": "upscale", + "backend": "flashvsr" if str(method).startswith("flashvsr") else "lanczos", + "method": method, + } + try: + raw_seed = body.get("seed", -1) + if isinstance(raw_seed, bool): + raise ValueError + seed = int(raw_seed) + except (TypeError, ValueError): + raise HTTPException(status_code=400, detail="Seed must be an integer") from None job_id = uuid.uuid4().hex[:8] job = { "id": job_id, "status": "queued", "progress": 0, "step": 0, "total_steps": 0, "phase": "", "message": "Queued (upscale)", "created_at": time.time(), "params": { - "video_path": resolved, - "method": body.get("method") or "flashvsr2", - "seed": body.get("seed", -1), + # Keep the physical source private to the in-memory worker. The + # sidecar writer deliberately strips it before publishing. + "source_path": resolved if source_kind == "image" else None, + "video_path": resolved if source_kind == "video" else None, + "source": source_filename, + "source_filename": source_filename, + "source_workspace": source_workspace, + "source_asset_id": source_asset_id, + "source_kind": source_kind, + "method": method, + "seed": seed, "model_type": "post_processing", - "generation_mode": "video", + "generation_mode": source_kind, + "provider": "local", + "capability": "upscale", + "inputs": [source_ref], + "parents": [source_ref], + "transformations": [transformation], + "_non_durable_tool": "upscale", }, "output_files": [], "error": None, - "workspace": workspace, "out_dir": _workspace_dir(workspace), + "workspace": workspace, "out_dir": output_dir, + "provenance": provenance, } _register_manual_generation_job(job) worker = _run_generation if execution_mode.policy().simulated else _run_tool_upscale @@ -23143,19 +23609,49 @@ async def tools_revoice(request: Request): workspace?: str } """ body = await request.json() - video_path = body.get("video_path") - if not video_path: - raise HTTPException(status_code=400, detail="video_path is required") - workspace = body.get("workspace") or _get_active_workspace() - resolved = _resolve_tool_clip_path(video_path, workspace) - if not resolved: - raise HTTPException(status_code=400, detail=f"Clip not found: {video_path}") + ( + resolved, + source_filename, + source_workspace, + _source_kind, + source_asset_id, + workspace, + output_dir, + ) = _resolve_tool_source(body, expected_kinds=("video",)) + from services.generation_provenance import normalize_submission_provenance + + provenance = normalize_submission_provenance({ + **(body.get("provenance") if isinstance(body.get("provenance"), dict) else {}), + "capability": "revoice", + "workspace_id": workspace, + }) voice_refs = body.get("voice_ref_paths") if not voice_refs and body.get("voice_ref_path"): voice_refs = [body.get("voice_ref_path")] if not voice_refs: raise HTTPException(status_code=400, detail="At least one voice_ref_path is required") + if ( + not isinstance(voice_refs, list) + or len(voice_refs) > 2 + or any( + not isinstance(value, str) + or not value.strip() + or len(value) > 1200 + or "\x00" in value + for value in voice_refs + ) + ): + raise HTTPException(status_code=400, detail="Invalid voice reference paths") + resolved_voice_refs = [] + for value in voice_refs: + resolved_voice_refs.append( + _resolve_request_media_path( + value, + workspace=workspace, + kinds=("audio", "video"), + ) + ) mode = body.get("mode", "single") if mode not in ("single", "two"): @@ -23167,15 +23663,23 @@ async def tools_revoice(request: Request): "phase": "", "message": "Queued (revoice)", "created_at": time.time(), "params": { "video_path": resolved, - "voice_ref_paths": voice_refs, + "source": source_filename, + "source_filename": source_filename, + "source_workspace": source_workspace, + "source_asset_id": source_asset_id, + "source_kind": "video", + "voice_ref_paths": resolved_voice_refs, "mode": mode, "diffusion_steps": body.get("diffusion_steps", 25), "cfg_rate": body.get("cfg_rate", 0.5), "model_type": "post_processing", "generation_mode": "video", + "capability": "revoice", + "_non_durable_tool": "revoice", }, "output_files": [], "error": None, - "workspace": workspace, "out_dir": _workspace_dir(workspace), + "workspace": workspace, "out_dir": output_dir, + "provenance": provenance, } _register_manual_generation_job(job) worker = _run_generation if execution_mode.policy().simulated else _run_tool_revoice @@ -23215,9 +23719,23 @@ def publish_progress(message: str, value: int, step: int, total: int) -> None: return False output_name = os.path.basename(generated_path) record_job_outputs(job, [output_name]) + tool_job = str(params.get("_non_durable_tool") or "") in {"upscale", "revoice"} + sidecar_params = params + if tool_job: + # Simulation still exercises the real queue/task/manifest path, but + # host filesystem paths and uploaded voice references must never enter + # the durable sidecar. + sidecar_params = { + key: value + for key, value in params.items() + if not str(key).startswith("_") + and key not in {"source_path", "video_path", "voice_ref_paths"} + } sidecar = { - "params": params, + "params": sidecar_params, "generation_mode": params.get("generation_mode"), + "tool": "tools" if tool_job else None, + "capability": params.get("capability") if tool_job else None, "job_id": job.get("id"), "task_id": job.get("task_id"), "root_task_id": job.get("root_task_id") or job.get("task_id"), @@ -23227,6 +23745,10 @@ def publish_progress(message: str, value: int, step: int, total: int) -> None: "simulated": True, "execution_mode": "simulate", } + if tool_job: + sidecar["inputs"] = params.get("inputs") or [] + sidecar["parents"] = params.get("parents") or [] + sidecar["transformations"] = params.get("transformations") or [] _publish_generation_sidecar_for_studio_job(job, generated_path, sidecar) if not finalize: return update_job( @@ -35774,6 +36296,10 @@ def _publish_generation_task(job: dict) -> dict: }.get(mode, "Generation job") if str(provenance.get("capability") or "") == "remove_background": task_title = "Tools · Remove background" + elif str(provenance.get("capability") or "") == "upscale": + task_title = "Tools · Upscale" + elif str(provenance.get("capability") or "") == "revoice": + task_title = "Tools · Revoice" return _upsert_canonical_task( workspace, task_id, diff --git a/app/docs/API.md b/app/docs/API.md index 90d95703e..e3d2a20e0 100644 --- a/app/docs/API.md +++ b/app/docs/API.md @@ -726,6 +726,24 @@ running, completed, failed or cancelled state; the derived asset exposes the source asset ID, tool/capability, instruction, model/backend, timings and transparent-PNG technical metadata. +## Tools upscale + +`POST /api/v1/tools/upscale` is one shared post-processing action for either a +still image or a video. Send `{ "source": "image.png", "source_kind": +"image", "asset_id": "...", "source_workspace": "...", "method": +"flashvsr2", "workspace": "default" }` for an image, or keep the legacy +`video_path` field with `source_kind: "video"` for a clip. Supported image +formats are `.bmp`, `.gif`, `.jpeg`, `.jpg`, `.png`, `.tif`, `.tiff`, and +`.webp`; supported video formats are `.avi`, `.m4v`, `.mkv`, `.mov`, `.mp4`, +`.mpeg`, `.mpg`, `.webm`, and `.wmv`. The source must be an exact asset, upload, or +file inside the selected workspace roots; path traversal and mismatched asset +IDs/kinds are rejected. Images use the existing spatial upsampler in still +mode and produce a new PNG beside the source. Videos retain the existing +audio-preserving pipeline and produce a new video. Neither path overwrites its +source. Poll the returned job with `GET /api/v1/status/{job_id}` and cancel it +with `POST /api/v1/cancel/{job_id}`. Activity and the canonical asset manifest +retain the source lineage, method, workspace, provenance, and execution mode. + ## Gallery mix kinds `GET /api/v1/outputs` accepts `result_kind=music_video|trailer|series_episode` (plus the existing `media_type`, `multiclip_only`, `favorites_only`, `search`, `workspace`, `limit`, `offset`). Classification lives in `services.output_result_kind` and applies only to **assembled** filenames (`multiclip`, `_mv.mp4`, `_movie.mp4`, `_rejoin_multiclip.mp4`, `_series_assembly`). Requesting `series_episode` also matches `chapter`. When `result_kind` is set, pagination is bypassed and every match is returned. diff --git a/app/services/generation_provenance.py b/app/services/generation_provenance.py index 667f66914..7bfa9e5dd 100644 --- a/app/services/generation_provenance.py +++ b/app/services/generation_provenance.py @@ -24,6 +24,8 @@ _TRUSTED_TOOL_BY_CAPABILITY = { "generate_story_song": "story_lab", "start_director_production": "director", + "upscale": "tools", + "revoice": "tools", "remove_background": "tools", } _TASK_ENTITY_FIELDS = ( diff --git a/app/shared/tools/background_removal_request.py b/app/shared/tools/background_removal_request.py index 0559d15e7..e9177d17e 100644 --- a/app/shared/tools/background_removal_request.py +++ b/app/shared/tools/background_removal_request.py @@ -18,6 +18,17 @@ from services.character_kit_face_cleanup import IMAGE_EXTENSIONS, _contained +# Keep the accepted formats at the source boundary. ``IMAGE_EXTENSIONS`` is +# intentionally the narrower set used by background removal; Tools -> Upscale +# can also consume the image formats Pillow/asset-catalog already understand. +UPSCALE_IMAGE_EXTENSIONS = frozenset( + IMAGE_EXTENSIONS | {".bmp", ".gif", ".tif", ".tiff"} +) +UPSCALE_VIDEO_EXTENSIONS = frozenset({ + ".avi", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm", ".wmv", +}) + + class RemoveBackgroundRequest(BaseModel): """Canonical request accepted by both Tools and the Wizard adapter.""" @@ -71,10 +82,15 @@ def child_workspace_name(path: str, default_root: str) -> str | None: return first if not first.startswith((".", "_")) and _valid_workspace_name(first) and first != "default" else None -def _safe_image_in_root(filename: str, root: str) -> str | None: +def _safe_image_in_root( + filename: str, + root: str, + allowed_extensions: Iterable[str] = IMAGE_EXTENSIONS, +) -> str | None: if not isinstance(filename, str) or not filename or os.path.basename(filename) != filename: return None - if os.path.splitext(filename)[1].casefold() not in IMAGE_EXTENSIONS: + extensions = {str(value).casefold() for value in allowed_extensions} + if os.path.splitext(filename)[1].casefold() not in extensions: return None root_real = os.path.realpath(os.path.abspath(root)) candidate = os.path.realpath(os.path.abspath(os.path.join(root_real, filename))) @@ -111,6 +127,7 @@ def _asset_location( destination_workspace: str, workspace_dir: Callable[[str], str], uploads_dir: Callable[[], str], + allowed_extensions: Iterable[str] = IMAGE_EXTENSIONS, ) -> tuple[str, str] | None: for location in _asset_locations(asset, source_workspace, destination_workspace): scope = str(location.get("workspace_id") or "").strip() @@ -120,7 +137,7 @@ def _asset_location( root = _location_root(scope, workspace_dir, uploads_dir) if root is None: continue - path = _safe_image_in_root(filename, root) + path = _safe_image_in_root(filename, root, allowed_extensions) if path: return path, scope return None @@ -133,12 +150,18 @@ def _source_from_asset( workspace_dir: Callable[[str], str], uploads_dir: Callable[[], str], asset_finder: Callable[[str], Mapping[str, Any] | None], + expected_kind: str = "image", + allowed_extensions: Iterable[str] = IMAGE_EXTENSIONS, + source_label: str = "image", ) -> tuple[str, str, str]: asset = asset_finder(payload.asset_id or "") if not asset: raise HTTPException(status_code=404, detail="Source asset not found") - if str(asset.get("kind") or "") != "image": - raise HTTPException(status_code=400, detail="Source asset must be an image") + if expected_kind and str(asset.get("kind") or "") != expected_kind: + raise HTTPException( + status_code=400, + detail=f"Source asset must be a {source_label}", + ) scope = _explicit_source_workspace(payload) location = _asset_location( asset, @@ -146,6 +169,7 @@ def _source_from_asset( destination_workspace=destination_workspace, workspace_dir=workspace_dir, uploads_dir=uploads_dir, + allowed_extensions=allowed_extensions, ) if not location: raise HTTPException(status_code=404, detail="Source asset location is unavailable") @@ -168,7 +192,7 @@ def _parse_source(payload: RemoveBackgroundRequest) -> tuple[str, str, str, str (is_api_path and not is_virtual) or (not is_api_path and source_without_query != source_name and not os.path.isabs(source_without_query)) ): - raise HTTPException(status_code=400, detail="Source image path is not allowed") + raise HTTPException(status_code=400, detail="Source path is not allowed") scope = _explicit_source_workspace(payload) return raw, source_without_query, source_name, scope, is_virtual @@ -181,6 +205,8 @@ def _absolute_source( destination_workspace: str, workspace_dir: Callable[[str], str], uploads_dir: Callable[[], str], + allowed_extensions: Iterable[str] = IMAGE_EXTENSIONS, + source_label: str = "image", ) -> tuple[str, str, str]: absolute = os.path.realpath(os.path.abspath(source)) uploads_root = os.path.realpath(os.path.abspath(uploads_dir())) @@ -201,18 +227,19 @@ def _absolute_source( try: root = os.path.realpath(os.path.abspath(workspace_dir(child))) except Exception as exc: - raise HTTPException(status_code=400, detail="Source image path is not allowed") from exc + raise HTTPException(status_code=400, detail="Source path is not allowed") from exc scope = child else: scope, root = destination_workspace, destination_root else: - raise HTTPException(status_code=400, detail="Source image path is not allowed") + raise HTTPException(status_code=400, detail="Source path is not allowed") if not _contained(absolute, root) or absolute == root: - raise HTTPException(status_code=400, detail="Source image path is not allowed") - if os.path.splitext(source_name)[1].casefold() not in IMAGE_EXTENSIONS: - raise HTTPException(status_code=400, detail="Source must be an image") + raise HTTPException(status_code=400, detail="Source path is not allowed") + extensions = {str(value).casefold() for value in allowed_extensions} + if os.path.splitext(source_name)[1].casefold() not in extensions: + raise HTTPException(status_code=400, detail=f"Source must be a {source_label}") if not os.path.isfile(absolute): - raise HTTPException(status_code=404, detail="Source image not found") + raise HTTPException(status_code=404, detail=f"Source {source_label} not found") return absolute, source_name, scope or destination_workspace @@ -249,14 +276,16 @@ def _named_source( destination_workspace: str, workspace_dir: Callable[[str], str], uploads_dir: Callable[[], str], + allowed_extensions: Iterable[str] = IMAGE_EXTENSIONS, + source_label: str = "image", ) -> tuple[str, str, str]: _validate_virtual_scope(source_without_query if is_virtual else "", scope) scope = _canonical_source_scope(source_without_query, scope, destination_workspace) _validate_source_scope(scope) root = uploads_dir() if scope == "__uploads__" else workspace_dir(scope) - path = _safe_image_in_root(source_name, root) + path = _safe_image_in_root(source_name, root, allowed_extensions) if not path: - raise HTTPException(status_code=404, detail="Source image not found") + raise HTTPException(status_code=404, detail=f"Source {source_label} not found") return path, os.path.basename(path), scope @@ -267,6 +296,9 @@ def resolve_source( workspace_dir: Callable[[str], str], uploads_dir: Callable[[], str], asset_finder: Callable[[str], Mapping[str, Any] | None], + expected_kind: str = "image", + allowed_extensions: Iterable[str] = IMAGE_EXTENSIONS, + source_label: str = "image", ) -> tuple[str, str, str]: """Resolve a canonical asset ID or a safe exact source path.""" if payload.asset_id: @@ -276,6 +308,9 @@ def resolve_source( workspace_dir=workspace_dir, uploads_dir=uploads_dir, asset_finder=asset_finder, + expected_kind=expected_kind, + allowed_extensions=allowed_extensions, + source_label=source_label, ) if not payload.source: raise HTTPException(status_code=400, detail="asset_id or source is required") @@ -288,6 +323,8 @@ def resolve_source( destination_workspace=destination_workspace, workspace_dir=workspace_dir, uploads_dir=uploads_dir, + allowed_extensions=allowed_extensions, + source_label=source_label, ) return _named_source( source_without_query, @@ -297,6 +334,8 @@ def resolve_source( destination_workspace=destination_workspace, workspace_dir=workspace_dir, uploads_dir=uploads_dir, + allowed_extensions=allowed_extensions, + source_label=source_label, ) @@ -359,6 +398,7 @@ def job_response( __all__ = [ - "RemoveBackgroundRequest", "child_workspace_name", "destination_context", "job_response", "resolve_source", - "validate_requested_workspace", + "IMAGE_EXTENSIONS", "UPSCALE_IMAGE_EXTENSIONS", "UPSCALE_VIDEO_EXTENSIONS", + "RemoveBackgroundRequest", "child_workspace_name", "destination_context", + "job_response", "resolve_source", "validate_requested_workspace", ] diff --git a/app/wgp.py b/app/wgp.py index 22ae0aaae..1378e79d4 100644 --- a/app/wgp.py +++ b/app/wgp.py @@ -5923,7 +5923,14 @@ def perform_temporal_upsampling(sample, previous_last_frame, temporal_upsampling return sample, previous_last_frame, output_fps -def perform_spatial_upsampling(sample, spatial_upsampling, seed=0, abort_callback=None, progress_callback=None): +def perform_spatial_upsampling( + sample, + spatial_upsampling, + seed=0, + abort_callback=None, + progress_callback=None, + still_image=False, +): from shared.utils.utils import resize_lanczos if spatial_upsampling == "vae2": return sample @@ -5946,7 +5953,7 @@ def perform_spatial_upsampling(sample, spatial_upsampling, seed=0, abort_callbac continue_cache=None, return_continue_cache=False, vae_tile_size=None, process_files=process_files_def, vae_config=vae_config, init_pipe=init_pipe, profile=profile, - still_image=False, abort_callback=abort_callback, + still_image=still_image, abort_callback=abort_callback, progress_callback=progress_callback, ) return sample diff --git a/tests/test_tools_upscale_contract.py b/tests/test_tools_upscale_contract.py new file mode 100644 index 000000000..0b19f3a72 --- /dev/null +++ b/tests/test_tools_upscale_contract.py @@ -0,0 +1,197 @@ +"""Contracts for the shared Tools upscale image/video boundary. + +The real worker is intentionally not imported here: importing the launch +runtime initializes WanGP and model services. Source resolution is exercised +with real temporary files, while the routing assertions inspect the small +launch adapter that owns the expensive pipeline. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +from fastapi import HTTPException + +from shared.tools.background_removal_request import ( + RemoveBackgroundRequest, + UPSCALE_IMAGE_EXTENSIONS, + UPSCALE_VIDEO_EXTENSIONS, + resolve_source, +) + + +ROOT = Path(__file__).resolve().parents[1] + + +def _workspace_functions(workspace: Path, uploads: Path): + return { + "workspace_dir": lambda _workspace: str(workspace), + "uploads_dir": lambda: str(uploads), + } + + +def _resolve(payload: RemoveBackgroundRequest, *, workspace: Path, uploads: Path, kind: str): + functions = _workspace_functions(workspace, uploads) + return resolve_source( + payload, + destination_workspace="default", + workspace_dir=functions["workspace_dir"], + uploads_dir=functions["uploads_dir"], + asset_finder=lambda _asset_id: None, + expected_kind=kind, + allowed_extensions=( + UPSCALE_IMAGE_EXTENSIONS if kind == "image" else UPSCALE_VIDEO_EXTENSIONS + ), + source_label=kind, + ) + + +def test_upscale_source_boundary_accepts_still_and_video_without_crossing_kinds(tmp_path): + workspace = tmp_path / "workspace" + uploads = tmp_path / "uploads" + workspace.mkdir() + uploads.mkdir() + still = workspace / "poster.tiff" + clip = workspace / "shot.webm" + still.write_bytes(b"image fixture") + clip.write_bytes(b"video fixture") + + image_result = _resolve( + RemoveBackgroundRequest(source=still.name, workspace="default"), + workspace=workspace, + uploads=uploads, + kind="image", + ) + video_result = _resolve( + RemoveBackgroundRequest(source=clip.name, workspace="default"), + workspace=workspace, + uploads=uploads, + kind="video", + ) + + assert image_result[:2] == (str(still.resolve()), still.name) + assert video_result[:2] == (str(clip.resolve()), clip.name) + with pytest.raises(HTTPException) as mismatch: + _resolve( + RemoveBackgroundRequest(source=still.name, workspace="default"), + workspace=workspace, + uploads=uploads, + kind="video", + ) + assert mismatch.value.status_code == 404 + + +def test_upscale_asset_boundary_rejects_an_image_asset_for_video_processing(tmp_path): + workspace = tmp_path / "workspace" + uploads = tmp_path / "uploads" + workspace.mkdir() + uploads.mkdir() + source = workspace / "poster.png" + source.write_bytes(b"image fixture") + asset = { + "id": "asset-poster", + "kind": "image", + "locations": [{"workspace_id": "default", "filename": source.name}], + } + + with pytest.raises(HTTPException) as mismatch: + resolve_source( + RemoveBackgroundRequest(asset_id="asset-poster"), + destination_workspace="default", + workspace_dir=lambda _workspace: str(workspace), + uploads_dir=lambda: str(uploads), + asset_finder=lambda _asset_id: asset, + expected_kind="video", + allowed_extensions=UPSCALE_VIDEO_EXTENSIONS, + source_label="video", + ) + assert mismatch.value.status_code == 400 + assert "video" in str(mismatch.value.detail) + + +def _function(tree: ast.AST, name: str) -> ast.AST: + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name: + return node + raise AssertionError(f"Function {name!r} not found") + + +def _called_names(node: ast.AST) -> set[str]: + names = set() + for child in ast.walk(node): + if not isinstance(child, ast.Call): + continue + if isinstance(child.func, ast.Name): + names.add(child.func.id) + elif isinstance(child.func, ast.Attribute): + names.add(child.func.attr) + return names + + +@pytest.fixture(scope="module") +def launch_tree(): + return ast.parse( + (ROOT / "app" / "_launch_runtime.py").read_text(encoding="utf-8"), + filename="app/_launch_runtime.py", + ) + + +def test_image_worker_branch_never_calls_video_decoder_or_writer(launch_tree): + worker = _function(launch_tree, "_run_tool_upscale") + image_branch = next( + node for node in ast.walk(worker) + if isinstance(node, ast.If) + and isinstance(node.test, ast.Compare) + and isinstance(node.test.left, ast.Name) + and node.test.left.id == "source_kind" + and any( + isinstance(op, ast.Eq) for op in node.test.ops + ) + and any( + isinstance(value, ast.Constant) and value.value == "image" + for value in node.test.comparators + ) + ) + image_calls = _called_names(ast.Module(body=image_branch.body, type_ignores=[])) + video_calls = _called_names(ast.Module(body=image_branch.orelse, type_ignores=[])) + + assert "_upscale_tool_image" in image_calls + assert not image_calls.intersection({ + "get_video_info", "extract_audio_tracks", "get_resampled_video", "save_video", + }) + assert {"get_video_info", "extract_audio_tracks", "get_resampled_video", "save_video"} <= video_calls + + +def test_still_adapter_uses_the_existing_upscale_pipeline_in_still_mode(launch_tree): + helper = _function(launch_tree, "_upscale_tool_image") + spatial_calls = [ + node for node in ast.walk(helper) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "perform_spatial_upsampling" + ] + assert len(spatial_calls) == 1 + still_keyword = next( + keyword.value for keyword in spatial_calls[0].keywords + if keyword.arg == "still_image" + ) + assert isinstance(still_keyword, ast.Constant) and still_keyword.value is True + + +def test_shared_upscale_route_accepts_both_source_kinds_and_uses_one_worker(launch_tree): + route = _function(launch_tree, "tools_upscale") + resolver = next( + node for node in ast.walk(route) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "_resolve_tool_source" + ) + expected_kinds = next( + keyword.value for keyword in resolver.keywords if keyword.arg == "expected_kinds" + ) + assert isinstance(expected_kinds, ast.Tuple) + assert [value.value for value in expected_kinds.elts] == ["image", "video"] + assert "_run_tool_upscale" in _called_names(route) + assert "_run_generation" in _called_names(route) diff --git a/ui/e2e/helpers/apiRoutes.ts b/ui/e2e/helpers/apiRoutes.ts index b77f4649e..aa947a0a9 100644 --- a/ui/e2e/helpers/apiRoutes.ts +++ b/ui/e2e/helpers/apiRoutes.ts @@ -12,9 +12,11 @@ export interface ApiRouteSession { } export type BackgroundRemovalE2EMode = 'complete' | 'cancel' | 'fail' +export type UpscaleE2EMode = 'complete' | 'cancel' | 'fail' export interface ApiRouteOptions { backgroundRemovalMode?: BackgroundRemovalE2EMode + upscaleMode?: UpscaleE2EMode } const SYSTEM_STATS = { @@ -160,6 +162,7 @@ const BACKGROUND_SOURCE_ASSET = { execution: { status: 'completed' }, model: { provider: 'local', id: 'flux' }, prompt_preview: 'A hero image used by the Tools E2E flow', + manifest: { technical: { width: 1920, height: 1080 } }, } const BACKGROUND_DERIVED_ASSET = { @@ -232,6 +235,71 @@ const BACKGROUND_DERIVED_OUTPUT = { thumbnail_url: '/api/v1/file/hero-no-background.png?workspace=default', } +const UPSCALE_DERIVED_ASSET = { + id: 'asset-hero-upscaled', + kind: 'image', + filename: 'hero_upscaled.png', + size_bytes: 48, + created_at: 3, + completed_at: 4, + metadata_status: 'canonical', + workspace_ids: ['default'], + locations: [{ + workspace_id: 'default', + filename: 'hero_upscaled.png', + url: '/api/v1/file/hero_upscaled.png?workspace=default', + }], + url: '/api/v1/file/hero_upscaled.png?workspace=default', + origin: { + tool: 'tools', + capability: 'upscale', + actor: 'user', + workspace_id: 'default', + }, + execution: { + status: 'completed', + job_id: 'tool-upscale-e2e', + task_id: 'task-generation-tool-upscale-e2e', + }, + model: { provider: 'local', id: 'post_processing' }, + prompt_preview: '', + manifest: { + schema_version: 1, + origin: { + tool: 'tools', + capability: 'upscale', + actor: 'user', + workspace_id: 'default', + }, + execution: { + status: 'completed', + job_id: 'tool-upscale-e2e', + task_id: 'task-generation-tool-upscale-e2e', + }, + generation: { + model: { provider: 'local', id: 'post_processing' }, + parameters: { method: 'lanczos2', source_kind: 'image' }, + }, + lineage: { + parents: [{ id: 'asset-hero', kind: 'image', uri: 'hero.png', role: 'source' }], + transformations: [{ type: 'upscale', backend: 'lanczos', method: 'lanczos2' }], + }, + }, +} + +const UPSCALE_DERIVED_OUTPUT = { + name: 'hero_upscaled.png', + type: 'image', + mode: 'image', + favorite: false, + size: 48, + created_at: 3, + completed_at: 4, + completion_time_source: 'metadata', + url: '/api/v1/file/hero_upscaled.png?workspace=default', + thumbnail_url: '/api/v1/file/hero_upscaled.png?workspace=default', +} + function json(body: Json, status = 200) { return { status, @@ -388,6 +456,10 @@ export async function installApiRoutes(page: Page, options: ApiRouteOptions = {} let backgroundRemovalSubmitted = false let backgroundRemovalStatusCalls = 0 let backgroundRemovalCancelRequested = false + const upscaleMode = options.upscaleMode || 'complete' + let upscaleSubmitted = false + let upscaleStatusCalls = 0 + let upscaleCancelRequested = false const backgroundAssets = () => [ BACKGROUND_SOURCE_ASSET, @@ -413,6 +485,30 @@ export async function installApiRoutes(page: Page, options: ApiRouteOptions = {} phase: 'completed', message: 'Background removed', output_files: ['hero-no-background.png'], error: null, } } + const toolAssets = () => [ + ...backgroundAssets(), + ...(upscaleSubmitted && upscaleStatusCalls > 1 && !upscaleCancelRequested && upscaleMode === 'complete' + ? [UPSCALE_DERIVED_ASSET] + : []), + ] + const upscaleStatus = () => { + if (upscaleCancelRequested) return { + status: 'cancelled', progress: 0, step: 0, total_steps: 1, + phase: 'cancelled', message: 'Upscale cancelled', output_files: [], error: null, + } + if (upscaleMode === 'fail') return { + status: 'failed', progress: 30, step: 1, total_steps: 3, + phase: 'upscaling', message: 'Upscale failed', output_files: [], error: 'upscale test failure', + } + if (upscaleStatusCalls < 2) return { + status: 'running', progress: 45, step: 1, total_steps: 2, + phase: 'upscaling', message: 'Upscaling image', output_files: [], error: null, + } + return { + status: 'completed', progress: 100, step: 2, total_steps: 2, + phase: 'completed', message: 'Image upscaled', output_files: ['hero_upscaled.png'], error: null, + } + } await page.route('**/api/**', async route => { const request = route.request() @@ -421,13 +517,17 @@ export async function installApiRoutes(page: Page, options: ApiRouteOptions = {} const pathname = url.pathname if (method === 'GET' && pathname === '/api/v1/assets') { - await route.fulfill(json({ assets: backgroundAssets(), total: backgroundAssets().length })) + await route.fulfill(json({ assets: toolAssets(), total: toolAssets().length })) return } if (method === 'GET' && pathname === '/api/v1/assets/asset-hero-cutout') { await route.fulfill(json(BACKGROUND_DERIVED_ASSET)) return } + if (method === 'GET' && pathname === '/api/v1/assets/asset-hero-upscaled') { + await route.fulfill(json(UPSCALE_DERIVED_ASSET)) + return + } if (method === 'POST' && pathname === '/api/v1/tools/remove-background') { backgroundRemovalSubmitted = true backgroundRemovalStatusCalls = 0 @@ -439,6 +539,17 @@ export async function installApiRoutes(page: Page, options: ApiRouteOptions = {} })) return } + if (method === 'POST' && pathname === '/api/v1/tools/upscale') { + upscaleSubmitted = true + upscaleStatusCalls = 0 + upscaleCancelRequested = false + await route.fulfill(json({ + job_id: 'tool-upscale-e2e', + task_id: 'task-generation-tool-upscale-e2e', + root_task_id: 'task-generation-tool-upscale-e2e', + })) + return + } if (method === 'GET' && pathname === '/api/v1/status/tool-bg-e2e' && backgroundRemovalSubmitted) { backgroundRemovalStatusCalls += 1 await route.fulfill(json({ @@ -453,15 +564,42 @@ export async function installApiRoutes(page: Page, options: ApiRouteOptions = {} })) return } + if (method === 'GET' && pathname === '/api/v1/status/tool-upscale-e2e' && upscaleSubmitted) { + upscaleStatusCalls += 1 + await route.fulfill(json({ + job_id: 'tool-upscale-e2e', + task_id: 'task-generation-tool-upscale-e2e', + root_task_id: 'task-generation-tool-upscale-e2e', + created_at: 1, + started_at: 2, + finished_at: upscaleStatusCalls > 1 ? 3 : null, + processing_time_sec: upscaleStatusCalls > 1 ? 1 : null, + generation_details: { + generation_mode: 'image', source_kind: 'image', source_asset_id: 'asset-hero', + capability: 'upscale', method: 'lanczos2', + }, + ...upscaleStatus(), + })) + return + } if (method === 'POST' && pathname === '/api/v1/cancel/tool-bg-e2e') { backgroundRemovalCancelRequested = true await route.fulfill(json({ status: 'cancelling' })) return } + if (method === 'POST' && pathname === '/api/v1/cancel/tool-upscale-e2e') { + upscaleCancelRequested = true + await route.fulfill(json({ status: 'cancelling' })) + return + } if (method === 'GET' && pathname === '/api/v1/outputs' && backgroundRemovalSubmitted && backgroundRemovalStatusCalls > 1 && !backgroundRemovalCancelRequested && backgroundRemovalMode === 'complete') { await route.fulfill(json({ outputs: [BACKGROUND_DERIVED_OUTPUT], total: 1 })) return } + if (method === 'GET' && pathname === '/api/v1/outputs' && upscaleSubmitted && upscaleStatusCalls > 1 && !upscaleCancelRequested && upscaleMode === 'complete') { + await route.fulfill(json({ outputs: [UPSCALE_DERIVED_OUTPUT], total: 1 })) + return + } const exact = catalog[keyFor(method, pathname)] if (exact && 'sse' in exact) { await route.fulfill({ diff --git a/ui/e2e/specs/tools-background-removal.spec.ts b/ui/e2e/specs/tools-background-removal.spec.ts index 9074fcf7d..d7b73501f 100644 --- a/ui/e2e/specs/tools-background-removal.spec.ts +++ b/ui/e2e/specs/tools-background-removal.spec.ts @@ -5,9 +5,21 @@ async function openBackgroundRemovalTools(page: Parameters[0]) { await page.getByRole('button', { name: 'Direct generation', exact: true }).click() await page.getByRole('tab', { name: 'Tools', exact: true }).click() await page.getByRole('button', { name: 'Remove background', exact: true }).click() - const picker = page.getByRole('combobox', { name: 'Source Image', exact: true }) + const picker = page.getByRole('list', { name: 'Source Image', exact: true }) await expect(picker).toBeVisible() - await expect(picker.locator('option[value="asset-hero"]')).toHaveCount(1) + await expect(picker.getByRole('button', { name: 'Select image hero.png', exact: true })).toBeVisible() + return picker +} + +async function openUpscaleTools(page: Parameters[0]) { + await page.getByRole('button', { name: 'Direct generation', exact: true }).click() + await page.getByRole('tab', { name: 'Tools', exact: true }).click() + await page.getByRole('button', { name: 'Upscale', exact: true }).click() + const picker = page.getByRole('list', { name: 'Source Media', exact: true }) + await expect(picker).toBeVisible() + await expect(picker.locator('select')).toHaveCount(0) + await expect(picker.getByRole('button', { name: 'Select image hero.png', exact: true })).toBeVisible() + await expect(picker).toContainText('Image · 1920×1080') return picker } @@ -30,7 +42,7 @@ test('runs Remove Background from direct Tools and exposes the derived asset', a await expect(run).toBeDisabled() await expect(page.getByRole('status')).toContainText('Choose an image from the library') - await picker.selectOption('asset-hero') + await picker.getByRole('button', { name: 'Select image hero.png', exact: true }).click() await expect(page.getByRole('img', { name: 'hero.png', exact: true })).toBeVisible() await expect(run).toBeEnabled() @@ -66,9 +78,9 @@ test('runs Remove Background from direct Tools and exposes the derived asset', a await page.getByRole('tab', { name: 'Tools', exact: true }).click() await page.getByRole('button', { name: 'Remove background', exact: true }).click() await page.getByRole('button', { name: 'Clear', exact: true }).click() - const reusablePicker = page.getByRole('combobox', { name: 'Source Image', exact: true }) - await expect(reusablePicker.locator('option[value="asset-hero-cutout"]')).toHaveCount(1) - await reusablePicker.selectOption('asset-hero-cutout') + const reusablePicker = page.getByRole('list', { name: 'Source Image', exact: true }) + await expect(reusablePicker.getByRole('button', { name: 'Select image hero-no-background.png', exact: true })).toBeVisible() + await reusablePicker.getByRole('button', { name: 'Select image hero-no-background.png', exact: true }).click() await expect(page.locator('aside').getByRole('img', { name: 'hero-no-background.png', exact: true })).toBeVisible() } finally { await closeApp(page, session) @@ -81,7 +93,7 @@ test('shows progress and lets the user cancel a Remove Background run', async ({ try { const picker = await openBackgroundRemovalTools(page) - await picker.selectOption('asset-hero') + await picker.getByRole('button', { name: 'Select image hero.png', exact: true }).click() await page.getByRole('button', { name: 'Remove Background', exact: true }).click() await expect(page.getByRole('button', { name: 'Stop', exact: true })).toBeVisible({ timeout: 7_000 }) @@ -98,7 +110,7 @@ test('keeps a tool failure visible in the activity card', async ({ page }) => { try { const picker = await openBackgroundRemovalTools(page) - await picker.selectOption('asset-hero') + await picker.getByRole('button', { name: 'Select image hero.png', exact: true }).click() await page.getByRole('button', { name: 'Remove Background', exact: true }).click() await expect(page.getByText('Generation Failed', { exact: true })).toBeVisible({ timeout: 7_000 }) await expect(page.getByText('rembg test failure', { exact: true })).toBeVisible() @@ -106,3 +118,47 @@ test('keeps a tool failure visible in the activity card', async ({ page }) => { await closeApp(page, session) } }) + +test('runs the shared Upscale action from an image and publishes a derived asset', async ({ page }) => { + const session = await gotoApp(page, { upscaleMode: 'complete' }) + const submissions = collectRequests(page, '/api/v1/tools/upscale') + const statuses = collectRequests(page, '/api/v1/status/tool-upscale-e2e') + + try { + const picker = await openUpscaleTools(page) + const run = page.getByRole('button', { name: 'Upscale Clip', exact: true }) + await expect(run).toBeDisabled() + + await picker.getByRole('button', { name: 'Select image hero.png', exact: true }).click() + await expect(page.getByRole('img', { name: 'hero.png', exact: true })).toBeVisible() + const imageRun = page.getByRole('button', { name: 'Upscale Image', exact: true }) + await expect(imageRun).toBeEnabled() + await imageRun.click() + + await expect.poll(() => submissions.length).toBe(1) + const payload = JSON.parse(submissions[0].postData() || '{}') as Record + expect(payload).toMatchObject({ + source: 'hero.png', + source_kind: 'image', + asset_id: 'asset-hero', + source_workspace: 'default', + workspace: 'default', + }) + expect(payload.video_path).toBeUndefined() + await expect(page.getByText('Queued...', { exact: true })).toBeVisible() + await expect.poll(() => statuses.length, { timeout: 10_000 }).toBeGreaterThanOrEqual(2) + + await page.getByRole('button', { name: 'Media', exact: true }).click() + await page.getByRole('tab', { name: 'Assets', exact: true }).click() + const upscaled = page.locator('article').filter({ hasText: 'hero_upscaled.png' }) + await expect(upscaled).toBeVisible() + await upscaled.getByRole('button', { name: 'Extra info', exact: true }).click() + const dialog = page.getByRole('dialog', { name: 'Extra info' }) + await expect(dialog).toBeVisible() + await expect(dialog).toContainText('upscale') + await expect(dialog).toContainText('asset-hero') + await dialog.getByRole('button', { name: 'Close', exact: true }).click() + } finally { + await closeApp(page, session) + } +}) diff --git a/ui/src/api/generation.ts b/ui/src/api/generation.ts index c9b8a2348..dda008183 100644 --- a/ui/src/api/generation.ts +++ b/ui/src/api/generation.ts @@ -246,18 +246,37 @@ export async function fetchJobStatus(jobId: string): Promise { return res.json() } -// --- Tools: standalone post-processing on an existing clip --- +// --- Tools: standalone post-processing on existing media --- export async function submitToolUpscale(params: { - video_path: string + source?: string + // Kept for callers that submit the legacy video contract. + video_path?: string + source_kind?: 'image' | 'video' + asset_id?: string + source_workspace?: string method?: string seed?: number workspace?: string + provenance?: { + actor?: 'user' | 'wizard' | 'system' | 'unknown' + capability?: string + workspace_id?: string + command?: { command_id?: string; workflow_id?: string; run_id?: string } + } }): Promise<{ job_id: string }> { + const { source, video_path, source_kind, ...rest } = params + const selectedSource = source ?? video_path + const body = { + ...rest, + ...(source_kind === 'image' + ? { source: selectedSource, source_kind } + : { video_path: selectedSource, ...(source_kind ? { source_kind } : {}) }), + } const res = await fetch(`${BASE}/api/v1/tools/upscale`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(params), + body: JSON.stringify(body), }) if (!res.ok) { const err = await res.json().catch(() => ({ detail: 'Upscale failed' })) diff --git a/ui/src/components/Sidebar/ToolsPanel.tsx b/ui/src/components/Sidebar/ToolsPanel.tsx index c0a27f1cf..68f1565c2 100644 --- a/ui/src/components/Sidebar/ToolsPanel.tsx +++ b/ui/src/components/Sidebar/ToolsPanel.tsx @@ -42,17 +42,23 @@ export function ToolsPanel() { const [vcUploading, setVcUploading] = useState(null) const [imageAssets, setImageAssets] = useState([]) const [imageAssetsLoading, setImageAssetsLoading] = useState(false) + const [imageAssetsError, setImageAssetsError] = useState(false) + const [sourceUploadError, setSourceUploadError] = useState(false) useEffect(() => { - if (tool !== 'remove_background') return + if (tool !== 'upscale' && tool !== 'remove_background') return const controller = new AbortController() setImageAssetsLoading(true) + setImageAssetsError(false) api.fetchAssets({ kind: 'image', limit: 100, signal: controller.signal }) .then(result => { if (!controller.signal.aborted) setImageAssets(result.assets) }) .catch(error => { - if (!controller.signal.aborted) console.error('Image asset catalog failed:', error) + if (!controller.signal.aborted) { + setImageAssetsError(true) + console.error('Image asset catalog failed:', error) + } }) .finally(() => { if (!controller.signal.aborted) setImageAssetsLoading(false) @@ -61,17 +67,26 @@ export function ToolsPanel() { }, [tool, toolsAssetsRevision]) const handleSourceUpload = async (file: File) => { + const image = file.type.startsWith('image/') || /\.(bmp|gif|jpe?g|png|tiff?|webp)$/i.test(file.name) + const video = file.type.startsWith('video/') || /\.(avi|m4v|mkv|mov|mp4|mpeg|mpg|webm|wmv)$/i.test(file.name) + if ((!image && !video) || (tool === 'remove_background' && !image) || (tool === 'revoice' && !video)) { + setSourceUploadError(true) + return + } + setSourceUploadError(false) setUploading(true) try { const r = await api.uploadImage(file) // /api/v1/upload handles video too + const kind = image ? 'image' : 'video' setSource({ path: r.path, name: file.name, url: r.url, - workspace: tool === 'remove_background' ? '__uploads__' : null, - kind: tool === 'remove_background' ? 'image' : 'video', + workspace: '__uploads__', + kind, }) } catch (e) { + setSourceUploadError(true) console.error('Source upload failed:', e) } finally { setUploading(false) @@ -125,7 +140,7 @@ export function ToolsPanel() { const hasVideoSource = !!sourcePath && sourceKind === 'video' const hasImageSource = !!sourcePath && sourceKind === 'image' const canRun = - (tool === 'upscale' && hasVideoSource) || + (tool === 'upscale' && (hasVideoSource || hasImageSource)) || (tool === 'revoice' && hasVideoSource && hasRefs) || (tool === 'remove_background' && hasImageSource) const flashvsrOff = flashvsrMode === 0 && method.startsWith('flashvsr') @@ -170,6 +185,8 @@ export function ToolsPanel() { useCurrentClip={useCurrentClip} imageAssets={imageAssets} imageAssetsLoading={imageAssetsLoading} + imageAssetsError={imageAssetsError} + sourceUploadError={sourceUploadError} selectImageAsset={selectImageAsset} /> {tool === 'remove_background' && !sourcePath && ( @@ -205,7 +222,9 @@ export function ToolsPanel() { }`} > - {tool === 'upscale' ? t('tools.upscaleClip') : tool === 'revoice' ? t('tools.replaceVoiceAction') : t('tools.removeBackgroundAction')} + {tool === 'upscale' + ? sourceKind === 'image' ? t('tools.upscaleImage') : t('tools.upscaleClip') + : tool === 'revoice' ? t('tools.replaceVoiceAction') : t('tools.removeBackgroundAction')} ) diff --git a/ui/src/components/Sidebar/ToolsSourcePanel.tsx b/ui/src/components/Sidebar/ToolsSourcePanel.tsx index 8fdfd975f..7a7f9f680 100644 --- a/ui/src/components/Sidebar/ToolsSourcePanel.tsx +++ b/ui/src/components/Sidebar/ToolsSourcePanel.tsx @@ -31,19 +31,24 @@ type SourceProps = { useCurrentClip: () => void imageAssets: AssetCatalogItem[] imageAssetsLoading: boolean + imageAssetsError: boolean + sourceUploadError: boolean selectImageAsset: (assetId: string) => void } export function ToolsSourcePanel(props: SourceProps) { const { t } = useUiTranslation('studio') return ( -
+
- {props.sourcePath - ? - : } + {props.sourcePath && } + {(!props.sourcePath || props.tool === 'upscale' || props.tool === 'remove_background') && ( + + )}
) } @@ -71,13 +76,37 @@ function SelectedSource({ function SourcePreview({ url, kind, name }: { url: string; kind: ToolSource['kind']; name: string }) { return kind === 'image' - ? {name} - :