diff --git a/app/_launch_runtime.py b/app/_launch_runtime.py index aebd8796..3552e482 100644 --- a/app/_launch_runtime.py +++ b/app/_launch_runtime.py @@ -121,6 +121,7 @@ from services.generation import bind_wgp, get_model_def bind_wgp(wgp) from services import model3d_service, minimax_h3_service, minimax_image_service +from services import tools_upscale from services import debug_trace from routers.lan_auth import create_lan_auth_router from services.durable_generation_queue import DurableGenerationQueue @@ -22727,18 +22728,10 @@ def _apply_spatial_upsampling_to_file(video_path: str, method: str, job: dict = # ============================================================================ -_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", - }), -} +# Compatibility aliases keep the existing HTTP validation contract stable while +# the implementation lives in the standalone Tools service. +_TOOL_UPSCALE_METHODS = tools_upscale.TOOL_UPSCALE_METHODS +_TOOL_SOURCE_EXTENSIONS = tools_upscale.TOOL_SOURCE_EXTENSIONS def _tool_asset_roots() -> list[dict[str, str]]: @@ -22937,50 +22930,16 @@ def _upscale_tool_image( 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 + """Compatibility facade for callers that used the old launch symbol.""" + return tools_upscale.upscale_image( + source_path, + output_path, + method, + wgp=wgp, + seed=seed, + abort_callback=abort_callback, + progress_callback=progress_callback, + ) def _write_tool_sidecar( @@ -23056,287 +23015,27 @@ def _write_tool_sidecar( def _run_tool_upscale(job_id: str): - """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: - if not acquired: - return False - try: - if not try_start( - job, message="Preparing upscale...", phase="Preparing", - ): - return False - start_time = float(job.get("started_at") or time.time()) - if not register_abort_state( - job, job_id, _active_gen_states, abort_state, - ): - return False - - params = job["params"] - workspace = job.get("workspace") - out_dir = job.get("out_dir") or wgp.save_path - os.makedirs(out_dir, exist_ok=True) - wgp.save_path = out_dir - - method = params.get("method") or "flashvsr2" - 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 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() - - def _abort(): - return bool(abort_state.get("abort")) or is_cancel_requested(job) - - # FlashVSR's _report_progress always calls back with - # (phase, current_step, total_steps); the latter two may be None. - def _progress(phase, current_step=None, total_steps=None): - changes = {} - if phase: - changes.update(message=str(phase), phase=str(phase)) - try: - if total_steps: - step = int(current_step or 0) - total = int(total_steps) - # Map reported steps onto 5..95% so the bar moves. - changes.update( - step=step, - total_steps=total, - progress=max(5, min(95, int(step / total * 100))), - ) - except (TypeError, ValueError, ZeroDivisionError): - pass - if changes: - update_job(job, **changes) - - 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) - else: - os.replace(tmp_path, final_path) - else: - # 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 - - 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 - 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) - 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=source_filename, - tool="upscale", - 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( - job, - "completed", - progress=100, - phase="", - message="Done", - ) - 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}") - return False - finally: - unregister_abort_state(job_id, _active_gen_states, abort_state) - try: - wgp.cleanup_temp_audio_files(audio_tracks) - except Exception: - pass - try: - wgp.release_flashvsr_vram() - except Exception: - pass - + """Compatibility facade for the standalone Tools upscale service.""" + return tools_upscale.run_tool_upscale( + job_id, + runtime={ + "jobs": _jobs, + "active_gen_states": _active_gen_states, + "wgp": wgp, + "coordinated_generation_slot": _coordinated_generation_slot, + "try_start": try_start, + "register_abort_state": register_abort_state, + "unregister_abort_state": unregister_abort_state, + "is_cancel_requested": is_cancel_requested, + "update_job": update_job, + "finish_job": finish_job, + "acknowledge_cancel": acknowledge_cancel, + "record_job_outputs": record_job_outputs, + "chunked_flashvsr_upscale": _chunked_flashvsr_upscale, + "resolve_tool_clip_path": _resolve_tool_clip_path, + "write_tool_sidecar": _write_tool_sidecar, + }, + ) def _run_tool_revoice(job_id: str): """Background worker: replace the voice(s) in an existing clip via SeedVC. diff --git a/app/services/tools_upscale.py b/app/services/tools_upscale.py new file mode 100644 index 00000000..b77e59ce --- /dev/null +++ b/app/services/tools_upscale.py @@ -0,0 +1,357 @@ +"""Standalone Tools upscale worker. + +The launch runtime owns process state and WanGP's singleton. This module owns +the post-processing algorithm and receives those runtime hooks explicitly so +it can be tested without importing the heavyweight server bootstrap. +""" + +from __future__ import annotations + +import os +import time +import traceback +import uuid +from typing import Any, Mapping + + +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 upscale_image( + source_path: str, + output_path: str, + method: str, + *, + wgp: Any, + seed: int = -1, + abort_callback=None, + progress_callback=None, +) -> tuple[int, int]: + """Upscale one still image through the existing spatial adapter.""" + 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 _upscale_image_job( + *, source_path, out_dir, source_filename, method, params, wgp, abort, progress, update_job +): + if not update_job(message="Upscaling image...", phase="Upscaling", progress=5): + return None, None + final_path = wgp.get_available_filename( + out_dir, source_filename, "_upscaled", force_extension=".png" + ) + image_size = upscale_image( + source_path, + final_path, + method, + wgp=wgp, + seed=int(params.get("seed", -1)), + abort_callback=abort, + progress_callback=progress, + ) + return final_path, image_size + + +def _upscale_video_job( + *, source_path, out_dir, source_filename, method, params, job, wgp, abort, progress, update_job +): + from shared.utils.utils import get_video_info + + fps, _width, _height, _frames = get_video_info(source_path) + 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): + return None, None, audio_tracks + + 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): + tmp_path = progress["chunked"]( + source_path, method, job=job, abort_check=abort, progress_callback=progress["callback"] + ) + if tmp_path is None or abort(): + if tmp_path and os.path.isfile(tmp_path): + try: + os.remove(tmp_path) + except OSError: + pass + return None, None, audio_tracks + 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) + else: + os.replace(tmp_path, final_path) + return final_path, None, audio_tracks + + sample = wgp.get_resampled_video( + source_path, 0, wgp.max_source_video_frames, fps + ) + sample = sample.permute(-1, 0, 1, 2) + sample = wgp.perform_spatial_upsampling( + sample, + method, + seed=int(params.get("seed", -1)), + abort_callback=abort, + progress_callback=progress["callback"], + ) + if abort(): + return None, None, audio_tracks + 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, + ) + return final_path, None, audio_tracks + + +def _publish_upscale_outputs(*, runtime, job_id, job, out_dir, before, source_filename, source_kind, method, image_size): + 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 + ) + if runtime["is_cancel_requested"](job): + for fname in new_files: + try: + os.remove(os.path.join(out_dir, fname)) + except OSError: + pass + return False + runtime["record_job_outputs"](job, new_files) + source_asset_id = job["params"].get("source_asset_id") + source_ref = {"id": source_asset_id, "kind": source_kind, "uri": source_filename, "role": "source"} + for fname in new_files: + runtime["write_tool_sidecar"]( + out_dir, fname, source_name=source_filename, tool="upscale", + params={ + "method": method, "model_type": "post_processing", + "source_asset_id": source_asset_id, "source_kind": source_kind, + "source_filename": source_filename, + }, + elapsed=time.time() - float(job.get("started_at") or time.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"} + ), + ) + return True + + +def _prepare_upscale_source(*, params, workspace, runtime): + """Validate the request and resolve its source before processing.""" + method = params.get("method") or "flashvsr2" + source_kind = str(params.get("source_kind") or "video").casefold() + if method not in TOOL_UPSCALE_METHODS: + raise ValueError("Unsupported upscale method") + if source_kind not in TOOL_SOURCE_EXTENSIONS: + raise ValueError("Unsupported source kind") + 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 = runtime["resolve_tool_clip_path"](source_value, workspace) + if not source_path: + raise FileNotFoundError("Input source not found") + if os.path.splitext(source_path)[1].casefold() not in TOOL_SOURCE_EXTENSIONS[source_kind]: + raise ValueError("Source kind does not match file format") + return method, source_kind, source_path + + +def run_tool_upscale(job_id: str, *, runtime: Mapping[str, Any]) -> bool: + """Run one image/video upscale using explicit launch-runtime hooks.""" + jobs = runtime["jobs"] + active_gen_states = runtime["active_gen_states"] + wgp = runtime["wgp"] + job = jobs[job_id] + abort_state = {"abort": False} + audio_tracks = [] + final_path = None + with runtime["coordinated_generation_slot"]( + job, description="HocusPocus Lab GPU tool · upscale" + ) as acquired: + if not acquired: + return False + try: + if not runtime["try_start"]( + job, message="Preparing upscale...", phase="Preparing" + ): + return False + if not runtime["register_abort_state"]( + job, job_id, active_gen_states, abort_state + ): + return False + params = job["params"] + workspace = job.get("workspace") + out_dir = job.get("out_dir") or wgp.save_path + os.makedirs(out_dir, exist_ok=True) + wgp.save_path = out_dir + method, source_kind, source_path = _prepare_upscale_source( + params=params, workspace=workspace, runtime=runtime + ) + before = set(os.listdir(out_dir)) if os.path.isdir(out_dir) else set() + + def abort(): + return bool(abort_state.get("abort")) or runtime["is_cancel_requested"](job) + + def progress_callback(phase, current_step=None, total_steps=None): + changes = {"message": str(phase), "phase": str(phase)} if phase else {} + if total_steps: + try: + step, total = int(current_step or 0), int(total_steps) + changes.update(step=step, total_steps=total, progress=max(5, min(95, int(step / total * 100)))) + except (TypeError, ValueError, ZeroDivisionError): + pass + if changes: + runtime["update_job"](job, **changes) + + source_filename = str(params.get("source_filename") or os.path.basename(source_path)) + final_path, image_size = (None, None) + if source_kind == "image": + final_path, image_size = _upscale_image_job( + source_path=source_path, out_dir=out_dir, source_filename=source_filename, + method=method, params=params, wgp=wgp, abort=abort, + progress=progress_callback, update_job=lambda **kw: runtime["update_job"](job, **kw), + ) + else: + progress = {"callback": progress_callback, "chunked": runtime["chunked_flashvsr_upscale"]} + final_path, _unused, audio_tracks = _upscale_video_job( + source_path=source_path, out_dir=out_dir, source_filename=source_filename, + method=method, params=params, job=job, wgp=wgp, abort=abort, + progress=progress, update_job=runtime["update_job"], + ) + if not final_path: + return False + if not _publish_upscale_outputs( + runtime=runtime, job_id=job_id, job=job, out_dir=out_dir, before=before, + source_filename=source_filename, source_kind=source_kind, + method=method, image_size=image_size, + ): + return False + completed = runtime["finish_job"]( + job, "completed", progress=100, phase="", message="Done" + ) + print( + f"[Tools/upscale] {source_filename} ({wgp.format_time(time.time() - float(job.get('started_at') or time.time()))})" + ) + return completed + except InterruptedError: + if final_path and os.path.isfile(final_path): + try: + os.remove(final_path) + except OSError: + pass + runtime["acknowledge_cancel"](job) + return False + except Exception as error: + traceback.print_exc() + runtime["finish_job"](job, "failed", error=str(error), message=f"Error: {error}") + return False + finally: + runtime["unregister_abort_state"](job_id, active_gen_states, abort_state) + try: + wgp.cleanup_temp_audio_files(audio_tracks) + except Exception: + pass + try: + wgp.release_flashvsr_vram() + except Exception: + pass diff --git a/tests/test_job_lifecycle_wiring.py b/tests/test_job_lifecycle_wiring.py index 5b1cae5b..97abf4ba 100644 --- a/tests/test_job_lifecycle_wiring.py +++ b/tests/test_job_lifecycle_wiring.py @@ -37,6 +37,22 @@ def _called_names(node: ast.AST) -> set[str]: return names +def _runtime_keys(node: ast.AST) -> set[str]: + """Return dependency names accessed through the extracted runtime map.""" + return { + child.value + for child in ast.walk(node) + if isinstance(child, ast.Constant) + and isinstance(child.value, str) + and child.value in { + "try_start", + "register_abort_state", + "finish_job", + "record_job_outputs", + } + } + + def _load_isolated_function(relative_path: str, name: str, namespace: dict): function = _function(_parse(relative_path), name) module = ast.Module(body=[function], type_ignores=[]) @@ -49,6 +65,7 @@ class TestJobLifecycleWiring(unittest.TestCase): @classmethod def setUpClass(cls): cls.launch = _parse("app/_launch_runtime.py") + cls.tools_upscale = _parse("app/services/tools_upscale.py") def test_each_worker_uses_lifecycle_transitions(self): expected = { @@ -71,7 +88,15 @@ def test_each_worker_uses_lifecycle_transitions(self): } for function_name, required in expected.items(): with self.subTest(function=function_name): - calls = _called_names(_function(self.launch, function_name)) + tree = self.tools_upscale if function_name == "_run_tool_upscale" else self.launch + lookup_name = "run_tool_upscale" if function_name == "_run_tool_upscale" else function_name + function = _function(tree, lookup_name) + functions = [function] + if function_name == "_run_tool_upscale": + functions.append(_function(tree, "_publish_upscale_outputs")) + calls = set().union( + *(_called_names(item) | _runtime_keys(item) for item in functions) + ) self.assertTrue(required <= calls, required - calls) def test_cancel_endpoint_routes_through_shared_helper(self): diff --git a/tests/test_tools_upscale_contract.py b/tests/test_tools_upscale_contract.py index a866556d..55dbaf6f 100644 --- a/tests/test_tools_upscale_contract.py +++ b/tests/test_tools_upscale_contract.py @@ -1,18 +1,21 @@ """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. +The launch runtime is intentionally not imported here: that bootstrap +initializes WanGP and model services. Source resolution is exercised with +real temporary files, and the still-image job helper is run against the +standalone service module with a bound lifecycle hook. """ from __future__ import annotations import ast from pathlib import Path +from types import SimpleNamespace import pytest from fastapi import HTTPException +from services.job_lifecycle import update_job +from services.tools_upscale import _upscale_image_job from shared.tools.background_removal_request import ( RemoveBackgroundRequest, @@ -138,8 +141,16 @@ def launch_tree(): ) -def test_image_worker_branch_never_calls_video_decoder_or_writer(launch_tree): - worker = _function(launch_tree, "_run_tool_upscale") +@pytest.fixture(scope="module") +def service_tree(): + return ast.parse( + (ROOT / "app" / "services" / "tools_upscale.py").read_text(encoding="utf-8"), + filename="app/services/tools_upscale.py", + ) + + +def test_image_worker_branch_never_calls_video_decoder_or_writer(service_tree): + worker = _function(service_tree, "run_tool_upscale") image_branch = next( node for node in ast.walk(worker) if isinstance(node, ast.If) @@ -157,15 +168,56 @@ def test_image_worker_branch_never_calls_video_decoder_or_writer(launch_tree): 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({ + assert "_upscale_image_job" in image_calls + assert "_upscale_video_job" in video_calls + + image_job = _function(service_tree, "_upscale_image_job") + video_job = _function(service_tree, "_upscale_video_job") + assert "upscale_image" in _called_names(image_job) + assert not _called_names(image_job).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 + assert {"get_video_info", "extract_audio_tracks", "get_resampled_video", "save_video"} <= _called_names(video_job) + + +def test_still_image_job_does_not_forward_job_into_the_bound_update_hook(tmp_path, monkeypatch): + """The image branch binds the live job in a lambda; passing job= again TypeErrors.""" + job = {"id": "job-still", "status": "running", "message": "Preparing"} + seen = [] + def fake_upscale_image(source_path, output_path, method, **kwargs): + seen.append((source_path, output_path, method, kwargs.get("seed"))) + return (1920, 1080) -def test_still_adapter_uses_the_existing_upscale_pipeline_in_still_mode(launch_tree): - helper = _function(launch_tree, "_upscale_tool_image") + monkeypatch.setattr("services.tools_upscale.upscale_image", fake_upscale_image) + wgp = SimpleNamespace( + get_available_filename=lambda out_dir, source_filename, suffix, force_extension=".png": str( + tmp_path / f"{Path(source_filename).stem}{suffix}{force_extension}" + ) + ) + + final_path, image_size = _upscale_image_job( + source_path=str(tmp_path / "poster.png"), + out_dir=str(tmp_path), + source_filename="poster.png", + method="lanczos2", + params={"seed": 11}, + wgp=wgp, + abort=lambda: False, + progress=lambda *_args, **_kwargs: None, + update_job=lambda **kw: update_job(job, **kw), + ) + + assert image_size == (1920, 1080) + assert Path(final_path).name == "poster_upscaled.png" + assert seen == [(str(tmp_path / "poster.png"), final_path, "lanczos2", 11)] + assert job["message"] == "Upscaling image..." + assert job["phase"] == "Upscaling" + assert job["progress"] == 5 + + +def test_still_adapter_uses_the_existing_upscale_pipeline_in_still_mode(service_tree): + helper = _function(service_tree, "upscale_image") spatial_calls = [ node for node in ast.walk(helper) if isinstance(node, ast.Call) @@ -180,6 +232,19 @@ def test_still_adapter_uses_the_existing_upscale_pipeline_in_still_mode(launch_t assert isinstance(still_keyword, ast.Constant) and still_keyword.value is True +def test_launch_worker_is_a_thin_facade_over_the_tools_service(launch_tree): + worker = _function(launch_tree, "_run_tool_upscale") + calls = _called_names(worker) + assert "run_tool_upscale" in calls + assert "_coordinated_generation_slot" not in calls + + +def test_upscale_service_does_not_import_the_launch_runtime(): + source = (ROOT / "app" / "services" / "tools_upscale.py").read_text(encoding="utf-8") + assert "_launch_runtime" not in source + assert "from fastapi" not in source + + def test_shared_upscale_route_accepts_both_source_kinds_and_uses_one_worker(launch_tree): route = _function(launch_tree, "tools_upscale") resolver = next(