Upload folder using huggingface_hub (part 2)

3afd6d6 verified about 1 month ago

9.44 kB

"""Hidden half-res output restore for Moxie MultiTrack projects.
Single-pass project sampling conditions and samples at half the
editor-selected resolution (see nodes/project.py). This module owns the
other half of that contract: Moxie Project Video Combine restores the exact
editor-selected resolution right before its images output, with no
user-visible wiring anywhere.
Hardware path adapted from NVIDIA's RTX Video Super Resolution ComfyUI node
(Apache License 2.0):
    https://github.com/Comfy-Org/Nvidia_RTX_Nodes_ComfyUI
    Copyright 2025 NVIDIA Corporation
It drives the ``nvidia-vfx`` pip package (RTX Video Super Resolution) per
frame via DLPack. When the package or a supported GPU is unavailable, the
module falls back to ComfyUI's Lanczos resize so the delivered frame size is
always exact and rendering never fails because of the upscaler.
"""

from __future__ import annotations

from typing import Any

__all__ = [
    "rtx_vsr_upscale",
    "restore_project_video_resolution",
    "upscale_images_to_project_resolution",
]


def _log(message: str) -> None:
    try:
        from .log import log_node_info
    except Exception:  # pragma: no cover - standalone tooling without the pack
        print(f"[Moxie] {message}", flush=True)  # noqa: T201
        return
    log_node_info("MultiTrack Output", message)


_FALLBACK_NOTICE_SHOWN = False


def _notice_fallback(message: str) -> None:
    global _FALLBACK_NOTICE_SHOWN
    if _FALLBACK_NOTICE_SHOWN:
        return
    _FALLBACK_NOTICE_SHOWN = True
    _log(message)


def _lanczos_upscale(images: Any, output_width: int, output_height: int) -> Any:
    import comfy.utils

    samples = images[..., :3].movedim(-1, 1)
    upscaled = comfy.utils.common_upscale(
        samples, output_width, output_height, "lanczos", "disabled"
    )
    return upscaled.movedim(1, -1)


def rtx_vsr_upscale(
    images: Any,
    output_width: int,
    output_height: int,
    quality: str = "ULTRA",
) -> Any:
    """Upscale an IMAGE batch to the exact target size, hardware first.
    Frames are processed per frame through RTX Video Super Resolution
    (adapted from Comfy-Org/Nvidia_RTX_Nodes_ComfyUI, Apache 2.0). Any
    failure - missing package, unsupported GPU, VRAM pressure - degrades to
    the Lanczos fallback instead of breaking the render, and the delivered
    size is corrected to the exact target whenever the hardware snap misses.
    """
    import torch

    frame_count, in_height, in_width, channels = images.shape
    if frame_count == 0 or (in_height, in_width) == (output_height, output_width):
        return images

    hardware_ready = channels == 3
    if hardware_ready:
        try:
            import nvvfx
        except Exception:
            hardware_ready = False
            _notice_fallback(
                "Hardware upscale unavailable (install nvidia-vfx for RTX "
                "acceleration); using Lanczos fallback."
            )

    if hardware_ready:
        # /8 snap from the NVIDIA node; editor selections are 8-aligned so
        # this is an identity operation for real project resolutions.
        snapped_width = max(8, round(output_width / 8) * 8)
        snapped_height = max(8, round(output_height / 8) * 8)
        max_batch_pixels = 1024 * 1024 * 16
        batch_size = max(
            1, max_batch_pixels // max(1, snapped_width * snapped_height)
        )
        quality_levels = {
            "LOW": nvvfx.effects.QualityLevel.LOW,
            "MEDIUM": nvvfx.effects.QualityLevel.MEDIUM,
            "HIGH": nvvfx.effects.QualityLevel.HIGH,
            "ULTRA": nvvfx.effects.QualityLevel.ULTRA,
        }
        selected_quality = quality_levels.get(
            str(quality).upper(), nvvfx.effects.QualityLevel.HIGH
        )
        out_tensor = torch.empty(
            (frame_count, snapped_height, snapped_width, channels),
            device=images.device,
            dtype=images.dtype,
        )
        try:
            with nvvfx.VideoSuperRes(selected_quality) as super_res:
                super_res.output_width = snapped_width
                super_res.output_height = snapped_height
                super_res.load()
                for start in range(0, frame_count, batch_size):
                    batch = images[start : start + batch_size]
                    batch_cuda = (
                        batch.cuda().permute(0, 3, 1, 2).float().contiguous()
                    )
                    for offset in range(batch_cuda.shape[0]):
                        result = super_res.run(batch_cuda[offset]).image
                        out_tensor[start + offset : start + offset + 1] = (
                            torch.from_dlpack(result).movedim(0, -1).unsqueeze(0)
                        )
        except Exception as error:
            hardware_ready = False
            _notice_fallback(
                f"Hardware upscale failed ({error}); using Lanczos fallback."
            )
        else:
            if (snapped_height, snapped_width) != (output_height, output_width):
                return _lanczos_upscale(out_tensor, output_width, output_height)
            return out_tensor

    return _lanczos_upscale(images, output_width, output_height)


def _project_output_resolution(project_name: Any) -> tuple[int, int] | None:
    """Editor-selected resolution stored in the project manifest, or ``None``.
    The manifest is written from the editor selection at the start of every
    run (nodes/project.py), so reading it fresh here can never see a stale
    frontend widget value.
    """
    try:
        from .h3_project import load_h3_project_data, safe_h3_project_name

        data = load_h3_project_data(safe_h3_project_name(project_name))
        width = int(data.get("width") or 0)
        height = int(data.get("height") or 0)
    except Exception:
        return None
    if width <= 0 or height <= 0 or (width, height) == (32, 32):
        return None
    return width, height


def upscale_images_to_project_resolution(
    images: Any,
    project_name: Any,
) -> Any:
    """Restore the project's editor resolution on combined frames.
    Used by Moxie Project Video Combine right before the ``images`` output.
    Frames already at the project resolution - full-res renders, dual-pass,
    SelfLift and passthrough outputs - pass through untouched; only frames
    whose size differs from the manifest (the hidden half-res sampling
    signature) are restored to the exact editor-selected resolution.
    """
    try:
        import torch
    except ImportError:  # pragma: no cover - torch is required in ComfyUI
        return images
    if (
        not isinstance(images, torch.Tensor)
        or images.ndim != 4
        or images.shape[0] == 0
    ):
        return images
    target = _project_output_resolution(project_name)
    if target is None:
        return images
    output_width, output_height = target
    if (int(images.shape[2]), int(images.shape[1])) == (
        output_width,
        output_height,
    ):
        return images
    return rtx_vsr_upscale(images, output_width, output_height)


def restore_project_video_resolution(
    video: Any,
    project_name: Any,
    source_path: str | None = None,
) -> Any:
    """Restore the manifest resolution on the composed project video.
    The combine composes the timeline at the segments' own size, which the
    hidden half-res sampling contract makes smaller than the manifest. This
    helper returns the video object untouched when the composed file already
    matches the manifest (full-res renders, dual-pass, SelfLift, passthrough);
    otherwise it decodes the frames once, upscales them to the exact manifest
    resolution, and rebuilds the video so every combine output (VIDEO,
    images, audio, fps) carries the editor-selected size. Any failure keeps
    the original video - rendering must never break because of the restore.
    """
    try:
        target = _project_output_resolution(project_name)
        if target is None:
            return video
        output_width, output_height = target
        if source_path:
            # Cheap ffprobe first: only pay the frame decode when the composed
            # file is genuinely smaller than the project resolution.
            from .video import ffprobe_info

            info = ffprobe_info(str(source_path))
            probed_width = info.get("width") if isinstance(info, dict) else None
            probed_height = info.get("height") if isinstance(info, dict) else None
            if (
                isinstance(probed_width, int)
                and isinstance(probed_height, int)
                and (probed_width, probed_height) == (output_width, output_height)
            ):
                return video
        components = video.get_components()
        images = upscale_images_to_project_resolution(
            components.images, project_name
        )
        if images is components.images:
            return video
        from comfy_api.latest import InputImpl, Types

        return InputImpl.VideoFromComponents(
            Types.VideoComponents(
                images=images,
                audio=components.audio,
                frame_rate=components.frame_rate,
            )
        )
    except Exception as error:
        _log(
            f"Output restore skipped ({error}); the composed video keeps its "
            "original resolution."
        )
        return video

	"""Hidden half-res output restore for Moxie MultiTrack projects.

	Single-pass project sampling conditions and samples at half the
	editor-selected resolution (see nodes/project.py). This module owns the
	other half of that contract: Moxie Project Video Combine restores the exact
	editor-selected resolution right before its images output, with no
	user-visible wiring anywhere.

	Hardware path adapted from NVIDIA's RTX Video Super Resolution ComfyUI node
	(Apache License 2.0):

	https://github.com/Comfy-Org/Nvidia_RTX_Nodes_ComfyUI
	Copyright 2025 NVIDIA Corporation

	It drives the ``nvidia-vfx`` pip package (RTX Video Super Resolution) per
	frame via DLPack. When the package or a supported GPU is unavailable, the
	module falls back to ComfyUI's Lanczos resize so the delivered frame size is
	always exact and rendering never fails because of the upscaler.
	"""

	from __future__ import annotations

	from typing import Any

	__all__ = [
	"rtx_vsr_upscale",
	"restore_project_video_resolution",
	"upscale_images_to_project_resolution",
	]


	def _log(message: str) -> None:
	try:
	from .log import log_node_info
	except Exception: # pragma: no cover - standalone tooling without the pack
	print(f"[Moxie] {message}", flush=True) # noqa: T201
	return
	log_node_info("MultiTrack Output", message)


	_FALLBACK_NOTICE_SHOWN = False


	def _notice_fallback(message: str) -> None:
	global _FALLBACK_NOTICE_SHOWN
	if _FALLBACK_NOTICE_SHOWN:
	return
	_FALLBACK_NOTICE_SHOWN = True
	_log(message)


	def _lanczos_upscale(images: Any, output_width: int, output_height: int) -> Any:
	import comfy.utils

	samples = images[..., :3].movedim(-1, 1)
	upscaled = comfy.utils.common_upscale(
	samples, output_width, output_height, "lanczos", "disabled"
	)
	return upscaled.movedim(1, -1)


	def rtx_vsr_upscale(
	images: Any,
	output_width: int,
	output_height: int,
	quality: str = "ULTRA",
	) -> Any:
	"""Upscale an IMAGE batch to the exact target size, hardware first.

	Frames are processed per frame through RTX Video Super Resolution
	(adapted from Comfy-Org/Nvidia_RTX_Nodes_ComfyUI, Apache 2.0). Any
	failure - missing package, unsupported GPU, VRAM pressure - degrades to
	the Lanczos fallback instead of breaking the render, and the delivered
	size is corrected to the exact target whenever the hardware snap misses.
	"""
	import torch

	frame_count, in_height, in_width, channels = images.shape
	if frame_count == 0 or (in_height, in_width) == (output_height, output_width):
	return images

	hardware_ready = channels == 3
	if hardware_ready:
	try:
	import nvvfx
	except Exception:
	hardware_ready = False
	_notice_fallback(
	"Hardware upscale unavailable (install nvidia-vfx for RTX "
	"acceleration); using Lanczos fallback."
	)

	if hardware_ready:
	# /8 snap from the NVIDIA node; editor selections are 8-aligned so
	# this is an identity operation for real project resolutions.
	snapped_width = max(8, round(output_width / 8) * 8)
	snapped_height = max(8, round(output_height / 8) * 8)
	max_batch_pixels = 1024 * 1024 * 16
	batch_size = max(
	1, max_batch_pixels // max(1, snapped_width * snapped_height)
	)
	quality_levels = {
	"LOW": nvvfx.effects.QualityLevel.LOW,
	"MEDIUM": nvvfx.effects.QualityLevel.MEDIUM,
	"HIGH": nvvfx.effects.QualityLevel.HIGH,
	"ULTRA": nvvfx.effects.QualityLevel.ULTRA,
	}
	selected_quality = quality_levels.get(
	str(quality).upper(), nvvfx.effects.QualityLevel.HIGH
	)
	out_tensor = torch.empty(
	(frame_count, snapped_height, snapped_width, channels),
	device=images.device,
	dtype=images.dtype,
	)
	try:
	with nvvfx.VideoSuperRes(selected_quality) as super_res:
	super_res.output_width = snapped_width
	super_res.output_height = snapped_height
	super_res.load()
	for start in range(0, frame_count, batch_size):
	batch = images[start : start + batch_size]
	batch_cuda = (
	batch.cuda().permute(0, 3, 1, 2).float().contiguous()
	)
	for offset in range(batch_cuda.shape[0]):
	result = super_res.run(batch_cuda[offset]).image
	out_tensor[start + offset : start + offset + 1] = (
	torch.from_dlpack(result).movedim(0, -1).unsqueeze(0)
	)
	except Exception as error:
	hardware_ready = False
	_notice_fallback(
	f"Hardware upscale failed ({error}); using Lanczos fallback."
	)
	else:
	if (snapped_height, snapped_width) != (output_height, output_width):
	return _lanczos_upscale(out_tensor, output_width, output_height)
	return out_tensor

	return _lanczos_upscale(images, output_width, output_height)


	def _project_output_resolution(project_name: Any) -> tuple[int, int] \| None:
	"""Editor-selected resolution stored in the project manifest, or ``None``.

	The manifest is written from the editor selection at the start of every
	run (nodes/project.py), so reading it fresh here can never see a stale
	frontend widget value.
	"""
	try:
	from .h3_project import load_h3_project_data, safe_h3_project_name

	data = load_h3_project_data(safe_h3_project_name(project_name))
	width = int(data.get("width") or 0)
	height = int(data.get("height") or 0)
	except Exception:
	return None
	if width <= 0 or height <= 0 or (width, height) == (32, 32):
	return None
	return width, height


	def upscale_images_to_project_resolution(
	images: Any,
	project_name: Any,
	) -> Any:
	"""Restore the project's editor resolution on combined frames.

	Used by Moxie Project Video Combine right before the ``images`` output.
	Frames already at the project resolution - full-res renders, dual-pass,
	SelfLift and passthrough outputs - pass through untouched; only frames
	whose size differs from the manifest (the hidden half-res sampling
	signature) are restored to the exact editor-selected resolution.
	"""
	try:
	import torch
	except ImportError: # pragma: no cover - torch is required in ComfyUI
	return images
	if (
	not isinstance(images, torch.Tensor)
	or images.ndim != 4
	or images.shape[0] == 0
	):
	return images
	target = _project_output_resolution(project_name)
	if target is None:
	return images
	output_width, output_height = target
	if (int(images.shape[2]), int(images.shape[1])) == (
	output_width,
	output_height,
	):
	return images
	return rtx_vsr_upscale(images, output_width, output_height)


	def restore_project_video_resolution(
	video: Any,
	project_name: Any,
	source_path: str \| None = None,
	) -> Any:
	"""Restore the manifest resolution on the composed project video.

	The combine composes the timeline at the segments' own size, which the
	hidden half-res sampling contract makes smaller than the manifest. This
	helper returns the video object untouched when the composed file already
	matches the manifest (full-res renders, dual-pass, SelfLift, passthrough);
	otherwise it decodes the frames once, upscales them to the exact manifest
	resolution, and rebuilds the video so every combine output (VIDEO,
	images, audio, fps) carries the editor-selected size. Any failure keeps
	the original video - rendering must never break because of the restore.
	"""
	try:
	target = _project_output_resolution(project_name)
	if target is None:
	return video
	output_width, output_height = target
	if source_path:
	# Cheap ffprobe first: only pay the frame decode when the composed
	# file is genuinely smaller than the project resolution.
	from .video import ffprobe_info

	info = ffprobe_info(str(source_path))
	probed_width = info.get("width") if isinstance(info, dict) else None
	probed_height = info.get("height") if isinstance(info, dict) else None
	if (
	isinstance(probed_width, int)
	and isinstance(probed_height, int)
	and (probed_width, probed_height) == (output_width, output_height)
	):
	return video
	components = video.get_components()
	images = upscale_images_to_project_resolution(
	components.images, project_name
	)
	if images is components.images:
	return video
	from comfy_api.latest import InputImpl, Types

	return InputImpl.VideoFromComponents(
	Types.VideoComponents(
	images=images,
	audio=components.audio,
	frame_rate=components.frame_rate,
	)
	)
	except Exception as error:
	_log(
	f"Output restore skipped ({error}); the composed video keeps its "
	"original resolution."
	)
	return video