olivv-cs-ComfyUI-FunPack/modules/output/render/upscale.py

"""Post-render upscale: a finished video through a models/upscale_models model.

Frames stream through ffmpeg in small chunks (never the whole clip in RAM), each chunk
goes through ComfyUI's own ImageUpscaleWithModel (tiling, OOM back-off, device handling),
and the result is encoded with the source's frame rate and sound.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess

from . import files, rtx
from ..._core import addons

CHUNK = 8                # frames per upscale call, at most
CHUNK_PIXELS = 2_000_000   # ...and at most this many input pixels per call: a big upscale of 1080p is gigabytes per frame
SUBFOLDER = "funpack_upscaled"


def probe(path):
    """-> (width, height, fps string, has sound)."""
    ffprobe = shutil.which("ffprobe")
    if not ffprobe:
        raise files.ClipError("ffprobe was not found on PATH: install ffmpeg to upscale videos.")
    out = subprocess.run([ffprobe, "-v", "error", "-show_entries",
                          "stream=codec_type,width,height,r_frame_rate", "-of", "json", path],
                         capture_output=True, text=True, check=True)
    streams = json.loads(out.stdout).get("streams") or []
    video = next((s for s in streams if s.get("codec_type") == "video"), None)
    if not video:
        raise files.ClipError(f"no video stream in {os.path.basename(path)}")
    return (int(video["width"]), int(video["height"]), video.get("r_frame_rate") or "25",
            any(s.get("codec_type") == "audio" for s in streams))


def _read(stream, n):
    buf = bytearray()
    while len(buf) < n:
        part = stream.read(n - len(buf))
        if not part:
            break
        buf += part
    return bytes(buf)


def upscale_file(src, dst, upscale, interrupted=lambda: None):
    """Stream `src` through `upscale` ([B,H,W,3] float 0-1 -> same, larger) into `dst`.
    Returns the frames written. `interrupted()` raises to stop between chunks."""
    import numpy as np
    import torch
    w, h, fps, has_sound = probe(src)
    frame = w * h * 3
    chunk = max(1, min(CHUNK, CHUNK_PIXELS // max(1, w * h)))
    ffmpeg = files.ffmpeg()
    dec = subprocess.Popen([ffmpeg, "-v", "error", "-i", src, "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
                           stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    enc, frames = None, 0
    try:
        while True:
            interrupted()
            raw = _read(dec.stdout, frame * chunk)
            n = len(raw) // frame
            if n == 0:
                break
            x = torch.from_numpy(np.frombuffer(raw[:n * frame], dtype=np.uint8)
                                 .reshape(n, h, w, 3).copy()).float() / 255.0
            y = (upscale(x).clamp(0, 1) * 255.0).round().to(torch.uint8).cpu().numpy()
            if enc is None:
                oh, ow = y.shape[1], y.shape[2]
                cmd = [ffmpeg, "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
                       "-s", f"{ow}x{oh}", "-r", fps, "-i", "-"]
                if has_sound:
                    cmd += ["-i", src, "-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac", "-b:a", "192k"]
                cmd += ["-vf", "scale=trunc(iw/2)*2:trunc(ih/2)*2", "-c:v", "libx264", "-crf", "16",
                        "-pix_fmt", "yuv420p", "-movflags", "+faststart", dst]
                enc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
            enc.stdin.write(y.tobytes())
            frames += n
        if enc is None:
            raise files.ClipError(f"no frames could be read from {os.path.basename(src)}")
        enc.stdin.close()
        err = enc.stderr.read().decode(errors="replace")
        if enc.wait() != 0:
            raise files.ClipError(f"ffmpeg could not encode the upscaled video: {err[-500:]}")
        return frames
    except BaseException:
        for p in (dec, enc):
            if p is not None and p.poll() is None:
                p.kill()
            if p is not None:
                p.wait()
                for pipe in (p.stdin, p.stdout, p.stderr):
                    try:
                        if pipe:
                            pipe.close()
                    except OSError:                        # a killed encoder's unflushed pipe
                        pass
        if os.path.exists(dst):
            os.remove(dst)
        raise
    finally:
        if dec.poll() is None:
            dec.kill()


def free_name(out_dir, filename, model):
    stem = os.path.splitext(os.path.basename(filename))[0]
    tag = os.path.splitext(os.path.basename(model))[0]
    n, name = 0, f"{stem}_{tag}.mp4"
    while os.path.exists(os.path.join(out_dir, name)):
        n += 1
        name = f"{stem}_{tag}_{n}.mp4"
    return name


def models() -> list:
    try:
        import folder_paths
        return list(folder_paths.get_filename_list("upscale_models"))
    except Exception:                                  # noqa: BLE001 -- headless
        return []


try:
    from comfy_api.latest import io
except Exception:                                      # noqa: BLE001 -- tests without ComfyUI
    io = None

if io is not None:
    class FunPackUpscaleVideo(io.ComfyNode):
        @classmethod
        def define_schema(cls) -> io.Schema:
            return io.Schema(
                node_id="FunPackUpscaleVideo",
                display_name="FunPack Upscale Video",
                category="FunPack/Output",
                description="Upscale a finished video file with an upscale model or NVIDIA RTX Super Resolution; keeps its fps and sound.",
                inputs=[
                    io.String.Input("filename", default=""),
                    io.String.Input("subfolder", default=""),
                    io.Combo.Input("type", options=["output", "temp"], default="output"),
                    io.Combo.Input("upscale_model", options=[*([rtx.NAME] if addons.enabled("nvidia_vsr") else []), *(models() or [""])],
                                   tooltip="From ComfyUI/models/upscale_models (use .safetensors), or NVIDIA RTX Super Resolution."),
                    io.String.Input("rtx_options", default="{}", optional=True,
                                    tooltip="NVIDIA RTX Super Resolution options as JSON: quality, resize, scale, width, height, strength, tenbit."),
                ],
                outputs=[],
                is_output_node=True,
            )

        @classmethod
        def execute(cls, filename, subfolder, type, upscale_model, rtx_options="{}") -> io.NodeOutput:
            import comfy.model_management as mm
            import folder_paths

            src = files.comfy_path(filename, subfolder, type)
            if not src or not os.path.isfile(src):
                raise RuntimeError(f"{filename} is not in the {type} folder (anymore).")
            if upscale_model == rtx.NAME:
                opts = rtx.options(rtx_options)
                fn, tag = rtx.make(opts), "rtx_" + opts["quality"].lower()
            else:
                from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel, UpscaleModelLoader
                model = UpscaleModelLoader.execute(upscale_model).result[0]
                fn, tag = (lambda x: ImageUpscaleWithModel.execute(model, x).result[0]), upscale_model
            out_dir = os.path.join(folder_paths.get_output_directory(), SUBFOLDER)
            os.makedirs(out_dir, exist_ok=True)
            name = free_name(out_dir, filename, tag)
            frames = upscale_file(src, os.path.join(out_dir, name), fn, mm.throw_exception_if_processing_interrupted)
            from ..._core import log
            log.info("FunPack Upscale", f"{filename} with {tag}: {frames} frames -> {SUBFOLDER}/{name}")
            return io.NodeOutput(ui={"videos": [{"filename": name, "subfolder": SUBFOLDER,
                                                 "type": "output", "format": "video/h264-mp4"}]})