From 89446e2a5800f4879047a20775ccf935b5619189 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Tue, 23 Jun 2026 18:35:00 +0000 Subject: [PATCH 1/8] Add OmniDreams postprocessing output path Signed-off-by: Gangzheng Tong --- .../tests/test_runner_postprocess.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 integrations/omnidreams/tests/test_runner_postprocess.py diff --git a/integrations/omnidreams/tests/test_runner_postprocess.py b/integrations/omnidreams/tests/test_runner_postprocess.py new file mode 100644 index 00000000..8c4ddd98 --- /dev/null +++ b/integrations/omnidreams/tests/test_runner_postprocess.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU checks for Omnidreams runner output post-processing wiring.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from omnidreams.config import OMNIDREAMS_RUNNERS +from omnidreams.runner import OmnidreamsRunner + +pytestmark = pytest.mark.ci_cpu + + +def _runner() -> OmnidreamsRunner: + runner = object.__new__(OmnidreamsRunner) + runner.config = SimpleNamespace() + return runner + + +def test_default_output_keeps_hdmap_and_generated_canvas() -> None: + runner = _runner() + condition = torch.full((1, 1, 2, 3, 2, 3), -1.0) + video = torch.full((1, 1, 2, 3, 2, 3), 1.0) + + canvas, description = runner._prepare_canvas_for_write( + condition=condition, + video=video, + ) + + assert description == "HDMap/RGB canvas" + assert canvas.shape == (2, 4, 3, 3) + assert torch.equal(canvas[:, :2], torch.full((2, 2, 3, 3), -1.0)) + assert torch.equal(canvas[:, 2:], torch.full((2, 2, 3, 3), 1.0)) + + +def test_omnidreams_runners_process_generated_views_independently() -> None: + assert OMNIDREAMS_RUNNERS + for runner_config in OMNIDREAMS_RUNNERS.values(): + assert runner_config.postprocess_output_layout == "bvtchw" + assert runner_config.postprocess_per_view is True From 08cb9d08e41afd4f3348ba56fad99e3bf9552694 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 26 Jun 2026 16:05:20 -0700 Subject: [PATCH 2/8] refactor(postprocess): register presets via entry points Replace the bespoke CLI postprocess mode selector with VideoPostprocessChainConfig.preset discovery, lazy FlashVSR pipeline building, finalize-on-failure handling, and runner-aligned defaults. --- integrations/flashvsr/flashvsr/postprocess.py | 20 +++++++++++-------- .../flashvsr/tests/test_postprocess.py | 4 ++-- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/integrations/flashvsr/flashvsr/postprocess.py b/integrations/flashvsr/flashvsr/postprocess.py index 15915012..51b8b38f 100644 --- a/integrations/flashvsr/flashvsr/postprocess.py +++ b/integrations/flashvsr/flashvsr/postprocess.py @@ -321,14 +321,18 @@ def _next_target_size(self) -> int: def _run_flashvsr_chunk(self, clip: Tensor) -> Tensor: assert self._pipeline is not None assert self._cache is not None - output = self._pipeline.generate( - autoregressive_index=self._ar_idx, - cache=self._cache, - input=clip, - ) - self._pipeline.finalize(autoregressive_index=self._ar_idx, cache=self._cache) - self._ar_idx += 1 - return output + try: + return self._pipeline.generate( + autoregressive_index=self._ar_idx, + cache=self._cache, + input=clip, + ) + finally: + self._pipeline.finalize( + autoregressive_index=self._ar_idx, + cache=self._cache, + ) + self._ar_idx += 1 def _consume_metadata(self, frames: int, *, source: str) -> dict[str, Any]: """Consume metadata spans covering ``frames`` buffered frames.""" diff --git a/integrations/flashvsr/tests/test_postprocess.py b/integrations/flashvsr/tests/test_postprocess.py index af4a1be2..902ea334 100644 --- a/integrations/flashvsr/tests/test_postprocess.py +++ b/integrations/flashvsr/tests/test_postprocess.py @@ -157,7 +157,7 @@ def test_flashvsr_postprocess_can_drop_short_tail( assert created[0].inputs == [] -def test_flashvsr_postprocess_does_not_finalize_when_generate_raises( +def test_flashvsr_postprocess_finalizes_when_generate_raises( monkeypatch: pytest.MonkeyPatch, ) -> None: created = _install_fake_builder(monkeypatch) @@ -181,7 +181,7 @@ def failing_generate( session.process(VideoChunk(tensor=video, layout="tchw")) pipeline = created[0] - assert pipeline.finalized == [] + assert pipeline.finalized == [0] assert len(pipeline.inputs) == 1 From 38c6a0f083d70e03a6a4a3e81eb0f9e9d1169789 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Fri, 26 Jun 2026 17:28:48 -0700 Subject: [PATCH 3/8] Remove unnecessary test --- .../tests/test_runner_postprocess.py | 56 ------------------- 1 file changed, 56 deletions(-) delete mode 100644 integrations/omnidreams/tests/test_runner_postprocess.py diff --git a/integrations/omnidreams/tests/test_runner_postprocess.py b/integrations/omnidreams/tests/test_runner_postprocess.py deleted file mode 100644 index 8c4ddd98..00000000 --- a/integrations/omnidreams/tests/test_runner_postprocess.py +++ /dev/null @@ -1,56 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CPU checks for Omnidreams runner output post-processing wiring.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch -from omnidreams.config import OMNIDREAMS_RUNNERS -from omnidreams.runner import OmnidreamsRunner - -pytestmark = pytest.mark.ci_cpu - - -def _runner() -> OmnidreamsRunner: - runner = object.__new__(OmnidreamsRunner) - runner.config = SimpleNamespace() - return runner - - -def test_default_output_keeps_hdmap_and_generated_canvas() -> None: - runner = _runner() - condition = torch.full((1, 1, 2, 3, 2, 3), -1.0) - video = torch.full((1, 1, 2, 3, 2, 3), 1.0) - - canvas, description = runner._prepare_canvas_for_write( - condition=condition, - video=video, - ) - - assert description == "HDMap/RGB canvas" - assert canvas.shape == (2, 4, 3, 3) - assert torch.equal(canvas[:, :2], torch.full((2, 2, 3, 3), -1.0)) - assert torch.equal(canvas[:, 2:], torch.full((2, 2, 3, 3), 1.0)) - - -def test_omnidreams_runners_process_generated_views_independently() -> None: - assert OMNIDREAMS_RUNNERS - for runner_config in OMNIDREAMS_RUNNERS.values(): - assert runner_config.postprocess_output_layout == "bvtchw" - assert runner_config.postprocess_per_view is True From 0b4c88a17f65e092416a51e17edfb7afd10f323c Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Tue, 23 Jun 2026 19:59:27 +0000 Subject: [PATCH 4/8] feat(realesrgan): add video upsampler integration Signed-off-by: Gangzheng Tong --- THIRD-PARTY-NOTICES | 13 + integrations/realesrgan/README.md | 108 +++++ integrations/realesrgan/pyproject.toml | 55 +++ .../realesrgan/realesrgan/__init__.py | 32 ++ integrations/realesrgan/realesrgan/archs.py | 191 ++++++++ integrations/realesrgan/realesrgan/cli.py | 206 +++++++++ .../realesrgan/realesrgan/postprocess.py | 142 ++++++ .../realesrgan/realesrgan/upsampler.py | 429 ++++++++++++++++++ integrations/realesrgan/scripts/gpu_smoke.py | 100 ++++ .../realesrgan/tests/test_realesrgan.py | 180 ++++++++ pyproject.toml | 2 + uv.lock | 23 + 12 files changed, 1481 insertions(+) create mode 100644 integrations/realesrgan/README.md create mode 100644 integrations/realesrgan/pyproject.toml create mode 100644 integrations/realesrgan/realesrgan/__init__.py create mode 100644 integrations/realesrgan/realesrgan/archs.py create mode 100644 integrations/realesrgan/realesrgan/cli.py create mode 100644 integrations/realesrgan/realesrgan/postprocess.py create mode 100644 integrations/realesrgan/realesrgan/upsampler.py create mode 100644 integrations/realesrgan/scripts/gpu_smoke.py create mode 100644 integrations/realesrgan/tests/test_realesrgan.py diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index bc8f3ff4..b46924a5 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -50,6 +50,13 @@ Wan 2.1 / Wan 2.2 Apache-2.0 https://github.com/Wan-Video/Wan2.1 Wan-Video team. Model weights are downloaded from upstream at runtime and are not redistributed in this repository. +Real-ESRGAN BSD-3-Clause https://github.com/xinntao/Real-ESRGAN +BasicSR Apache-2.0 https://github.com/XPixelGroup/BasicSR + The optional Real-ESRGAN integration implements RRDBNet and compact SRVGG + architecture definitions compatible with public Real-ESRGAN checkpoints. + Model weights are downloaded from upstream at runtime and are not + redistributed in this repository. + ================================================================================ Optional integration: integrations/omnidreams ================================================================================ @@ -100,6 +107,12 @@ aiohttp Apache-2.0 https://github.com/aio-libs/aiohttp aiortc BSD-3-Clause https://github.com/aiortc/aiortc opencv-python-headless Apache-2.0 https://github.com/opencv/opencv-python +================================================================================ +Optional integration: integrations/realesrgan +================================================================================ + +opencv-python-headless Apache-2.0 https://github.com/opencv/opencv-python + ================================================================================ Source-level redistributions ================================================================================ diff --git a/integrations/realesrgan/README.md b/integrations/realesrgan/README.md new file mode 100644 index 00000000..c7f0d052 --- /dev/null +++ b/integrations/realesrgan/README.md @@ -0,0 +1,108 @@ + + +# flashdreams-realesrgan + +Real-ESRGAN frame and video upsampling packaged as a FlashDreams workspace +integration. It provides: + +- a reusable `RealESRGANUpsampler` Python API, +- a `RealESRGANPostProcessorConfig` for FlashDreams generated RGB video chunks, +- a `realesrgan-upsample` CLI for image and video files. + +The architecture definitions follow the public Real-ESRGAN / BasicSR RRDBNet +and compact SRVGG network layouts so public checkpoints can be loaded directly. +Weights are downloaded at first use into +`$FLASHDREAMS_CACHE_DIR/realesrgan` or `~/.cache/flashdreams/realesrgan`. + +## Install + +From the repository root: + +```bash +uv sync --package flashdreams-realesrgan --extra dev +``` + +The root workspace glob includes `integrations/realesrgan`, so editable +workspace usage works without extra path setup. + +## CLI + +Upsample one image: + +```bash +uv run --package flashdreams-realesrgan realesrgan-upsample \ + --input /path/to/input.png \ + --output /path/to/input_2x.png \ + --scale 2 +``` + +Upsample a video: + +```bash +uv run --package flashdreams-realesrgan realesrgan-upsample \ + --input /path/to/input.mp4 \ + --output /path/to/input_2x.mp4 \ + --scale 2 \ + --tile 256 \ + --compile +``` + +Useful flags: + +| flag | description | +| --- | --- | +| `--scale {2,4}` | Output scale. Defaults to `2`. | +| `--model-name` | Public checkpoint name such as `RealESRGAN_x2plus` or `RealESRGAN_x4plus`. | +| `--model-path` | Local checkpoint path. Skips public download. | +| `--tile` | Tile size for lower peak VRAM. `0` processes whole frames. | +| `--fp32` | Disable fp16 CUDA inference. | +| `--compile` | Enable `torch.compile` with `mode="reduce-overhead"`. | +| `--compile-mode` | Override the `torch.compile` mode. | +| `--profile-warmup-frames` | Frames excluded from steady FPS metrics. Defaults to `10`. | +| `--max-frames` | Video-only frame cap for quick smoke tests. | + +## FlashDreams Postprocess API + +```python +from flashdreams.infra.postprocess import VideoPostprocessChainConfig +from realesrgan import RealESRGANPostProcessorConfig + +postprocess = VideoPostprocessChainConfig( + processors=( + RealESRGANPostProcessorConfig( + scale=2, + model_name="RealESRGAN_x2plus", + tile=256, + compile_model=True, + device="cuda", + ), + ) +) +``` + +The postprocessor accepts generated RGB chunks in any FlashDreams video layout, +runs frame-local Real-ESRGAN inference, and returns `bvtchw` chunks in +`[-1, 1]`. + +## Tests + +CPU-safe tests use random weights and do not download checkpoints: + +```bash +uv run --package flashdreams-realesrgan pytest integrations/realesrgan/tests/test_realesrgan.py -m ci_cpu +``` + +Checkpoint-backed GPU smoke test: + +```bash +uv run --package flashdreams-realesrgan python integrations/realesrgan/scripts/gpu_smoke.py \ + --scale 2 \ + --device cuda +``` + +The smoke script downloads `RealESRGAN_x2plus` once, runs one CUDA fp16 +forward pass on a `16x16` RGB frame, checks the `32x32` output, and prints +model load time, inference time, and peak allocated VRAM. diff --git a/integrations/realesrgan/pyproject.toml b/integrations/realesrgan/pyproject.toml new file mode 100644 index 00000000..44034a2e --- /dev/null +++ b/integrations/realesrgan/pyproject.toml @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-realesrgan" +version = "0.1.0" +description = "Real-ESRGAN frame and video upsampler integration for flashdreams." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "flashdreams", + "opencv-python-headless>=4.8", +] + +[tool.uv] +managed = true + +[tool.uv.sources] +flashdreams = { workspace = true } + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", +] + +[project.scripts] +realesrgan-upsample = "realesrgan.cli:main" + +[tool.setuptools.packages.find] +include = ["realesrgan*"] +exclude = ["tests"] + +[tool.pytest.ini_options] +addopts = "--import-mode=importlib -p flashdreams._pytest_plugins.marker_enforcement" +markers = [ + "ci_cpu: CPU-safe test, runs on the CPU CI runner", + "ci_gpu: requires GPU or libGL (cv2), runs on the GPU CI runner", + "manual: heavy or environment-specific test, opt-in only", +] diff --git a/integrations/realesrgan/realesrgan/__init__.py b/integrations/realesrgan/realesrgan/__init__.py new file mode 100644 index 00000000..f0efe4e6 --- /dev/null +++ b/integrations/realesrgan/realesrgan/__init__.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-ESRGAN upsampling integration for FlashDreams.""" + +from realesrgan.postprocess import RealESRGANPostProcessorConfig +from realesrgan.upsampler import ( + MODEL_CONFIGS, + MODEL_URLS, + RealESRGANUpsampler, + create_model, +) + +__all__ = [ + "MODEL_CONFIGS", + "MODEL_URLS", + "RealESRGANPostProcessorConfig", + "RealESRGANUpsampler", + "create_model", +] diff --git a/integrations/realesrgan/realesrgan/archs.py b/integrations/realesrgan/realesrgan/archs.py new file mode 100644 index 00000000..7b94ee1b --- /dev/null +++ b/integrations/realesrgan/realesrgan/archs.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-ESRGAN network architectures. + +The RRDBNet and compact SRVGG definitions mirror the public Real-ESRGAN +and BasicSR architectures so upstream checkpoints can be loaded directly. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torch.nn import functional as F + + +def pixel_unshuffle(x: torch.Tensor, scale: int) -> torch.Tensor: + """Return the inverse pixel-shuffle transform for ``x``.""" + batch, channels, height, width = x.shape + assert height % scale == 0 and width % scale == 0 + out_channels = channels * scale * scale + height_out = height // scale + width_out = width // scale + x = x.view(batch, channels, height_out, scale, width_out, scale) + return x.permute(0, 1, 3, 5, 2, 4).reshape( + batch, out_channels, height_out, width_out + ) + + +def default_init_weights(module_list: list[nn.Module], scale: float = 1.0) -> None: + """Initialize convolutional and linear weights in ``module_list``.""" + for module in module_list: + for child in module.modules(): + if isinstance(child, nn.Conv2d): + nn.init.kaiming_normal_(child.weight, a=0, mode="fan_in") + child.weight.data *= scale + if child.bias is not None: + child.bias.data.zero_() + elif isinstance(child, nn.Linear): + nn.init.kaiming_normal_(child.weight, a=0, mode="fan_in") + child.weight.data *= scale + if child.bias is not None: + child.bias.data.zero_() + + +class ResidualDenseBlock(nn.Module): + """Residual dense block used by RRDBNet.""" + + def __init__(self, num_feat: int = 64, num_grow_ch: int = 32) -> None: + super().__init__() + self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) + self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1) + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + default_init_weights( + [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], + 0.1, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the residual dense block.""" + x1 = self.lrelu(self.conv1(x)) + x2 = self.lrelu(self.conv2(torch.cat((x, x1), dim=1))) + x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), dim=1))) + x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), dim=1))) + x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), dim=1)) + return x5 * 0.2 + x + + +class RRDB(nn.Module): + """Residual-in-residual dense block used by RRDBNet.""" + + def __init__(self, num_feat: int, num_grow_ch: int = 32) -> None: + super().__init__() + self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch) + self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch) + self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the residual-in-residual block.""" + out = self.rdb1(x) + out = self.rdb2(out) + out = self.rdb3(out) + return out * 0.2 + x + + +class RRDBNet(nn.Module): + """RRDBNet used by RealESRGAN_x2plus and RealESRGAN_x4plus.""" + + def __init__( + self, + num_in_ch: int, + num_out_ch: int, + scale: int = 4, + num_feat: int = 64, + num_block: int = 23, + num_grow_ch: int = 32, + ) -> None: + super().__init__() + self.scale = scale + if scale == 2: + num_in_ch *= 4 + elif scale == 1: + num_in_ch *= 16 + + self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) + self.body = nn.Sequential( + *[ + RRDB(num_feat=num_feat, num_grow_ch=num_grow_ch) + for _ in range(num_block) + ] + ) + self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run RRDBNet on an RGB image batch in ``[0, 1]``.""" + if self.scale == 2: + feat = pixel_unshuffle(x, scale=2) + elif self.scale == 1: + feat = pixel_unshuffle(x, scale=4) + else: + feat = x + + feat = self.conv_first(feat) + body_feat = self.conv_body(self.body(feat)) + feat = feat + body_feat + feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2))) + feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2))) + return self.conv_last(self.lrelu(self.conv_hr(feat))) + + +class SRVGGNetCompact(nn.Module): + """Compact VGG-style super-resolution network used by Real-ESRGAN v3.""" + + def __init__( + self, + num_in_ch: int = 3, + num_out_ch: int = 3, + num_feat: int = 64, + num_conv: int = 16, + upscale: int = 4, + act_type: str = "prelu", + ) -> None: + super().__init__() + self.upscale = upscale + self.body = nn.ModuleList() + self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)) + self.body.append(self._make_activation(act_type, num_feat)) + for _ in range(num_conv): + self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1)) + self.body.append(self._make_activation(act_type, num_feat)) + self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1)) + self.upsampler = nn.PixelShuffle(upscale) + + @staticmethod + def _make_activation(act_type: str, num_feat: int) -> nn.Module: + if act_type == "relu": + return nn.ReLU(inplace=True) + if act_type == "prelu": + return nn.PReLU(num_parameters=num_feat) + if act_type == "leakyrelu": + return nn.LeakyReLU(negative_slope=0.1, inplace=True) + raise ValueError(f"Unknown activation type: {act_type}") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run SRVGG on an RGB image batch in ``[0, 1]``.""" + out = x + for layer in self.body: + out = layer(out) + out = self.upsampler(out) + base = F.interpolate(x, scale_factor=self.upscale, mode="nearest") + return out + base diff --git a/integrations/realesrgan/realesrgan/cli.py b/integrations/realesrgan/realesrgan/cli.py new file mode 100644 index 00000000..21442191 --- /dev/null +++ b/integrations/realesrgan/realesrgan/cli.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Command-line Real-ESRGAN image/video upsampler.""" + +from __future__ import annotations + +import argparse +import time +from pathlib import Path +from typing import Literal + +import cv2 +import torch + +from realesrgan.upsampler import ( + RealESRGANUpsampler, + default_model_name, + write_bgr_image, +) + +IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"} +VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv", ".wmv"} + + +def main() -> None: + """Run the Real-ESRGAN upsampler CLI.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", "-i", type=Path, required=True) + parser.add_argument("--output", "-o", type=Path, required=True) + parser.add_argument("--scale", "-s", type=int, choices=(2, 4), default=2) + parser.add_argument("--model-name", type=str, default=None) + parser.add_argument("--model-path", type=Path, default=None) + parser.add_argument("--tile", type=int, default=0) + parser.add_argument("--tile-pad", type=int, default=10) + parser.add_argument("--pre-pad", type=int, default=10) + parser.add_argument("--fp32", action="store_true") + parser.add_argument( + "--compile", + action="store_true", + help="Compile the model with torch.compile().", + ) + parser.add_argument( + "--compile-mode", + type=str, + default="reduce-overhead", + help="Mode passed to torch.compile() when --compile is used.", + ) + parser.add_argument("--device", type=str, default=None) + parser.add_argument( + "--profile-warmup-frames", + type=int, + default=10, + help="Frames to exclude from steady FPS profiling.", + ) + parser.add_argument( + "--max-frames", + type=int, + default=None, + help="Optional cap for video smoke tests.", + ) + args = parser.parse_args() + + scale: Literal[2, 4] = 2 if args.scale == 2 else 4 + model_name = args.model_name or default_model_name(scale) + upsampler = RealESRGANUpsampler( + scale=scale, + model_name=model_name, + model_path=args.model_path, + tile=args.tile, + tile_pad=args.tile_pad, + pre_pad=args.pre_pad, + half=not args.fp32, + compile_model=args.compile, + compile_mode=args.compile_mode, + device=args.device, + ) + + suffix = args.input.suffix.lower() + if suffix in IMAGE_EXTENSIONS: + _upsample_image(args.input, args.output, upsampler) + return + if suffix in VIDEO_EXTENSIONS: + _upsample_video( + args.input, + args.output, + upsampler, + max_frames=args.max_frames, + profile_warmup_frames=args.profile_warmup_frames, + ) + return + raise SystemExit(f"Unsupported input extension {suffix!r}.") + + +def _upsample_image( + input_path: Path, + output_path: Path, + upsampler: RealESRGANUpsampler, +) -> None: + image = cv2.imread(str(input_path), cv2.IMREAD_UNCHANGED) + if image is None: + raise RuntimeError(f"Failed to read image {input_path}") + output, mode = upsampler.upsample_bgr_image(image) + if mode == "RGBA": + output_path = output_path.with_suffix(".png") + write_bgr_image(output_path, output) + print(f"wrote {output.shape[1]}x{output.shape[0]} {output_path}") + + +def _upsample_video( + input_path: Path, + output_path: Path, + upsampler: RealESRGANUpsampler, + *, + max_frames: int | None, + profile_warmup_frames: int, +) -> None: + cap = cv2.VideoCapture(str(input_path)) + if not cap.isOpened(): + raise RuntimeError(f"Failed to open video {input_path}") + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = float(cap.get(cv2.CAP_PROP_FPS) or 30.0) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + output_path.parent.mkdir(parents=True, exist_ok=True) + fourcc_factory = getattr(cv2, "VideoWriter_fourcc") + writer = cv2.VideoWriter( + str(output_path), + fourcc_factory(*"mp4v"), + fps, + (width * upsampler.scale, height * upsampler.scale), + ) + if not writer.isOpened(): + cap.release() + raise RuntimeError(f"Failed to create video writer {output_path}") + + start = time.perf_counter() + frame_idx = 0 + process_durations: list[float] = [] + infer_durations: list[float] = [] + try: + while True: + ok, frame = cap.read() + if not ok: + break + process_start = time.perf_counter() + infer_start = time.perf_counter() + output, _mode = upsampler.upsample_bgr_image(frame) + _synchronize(upsampler) + infer_durations.append(time.perf_counter() - infer_start) + writer.write(output) + process_durations.append(time.perf_counter() - process_start) + frame_idx += 1 + if max_frames is not None and frame_idx >= max_frames: + break + finally: + cap.release() + writer.release() + + elapsed = time.perf_counter() - start + total_note = f"/{total_frames}" if total_frames else "" + warmup_frames = min(max(profile_warmup_frames, 0), len(process_durations)) + steady_process = process_durations[warmup_frames:] + steady_infer = infer_durations[warmup_frames:] + print(f"wrote {frame_idx}{total_note} frames -> {output_path}") + print( + f"profile total_s={elapsed:.3f} total_fps={_fps(frame_idx, elapsed):.2f} " + f"process_fps={_duration_fps(process_durations):.2f} " + f"infer_fps={_duration_fps(infer_durations):.2f}" + ) + print( + f"profile_steady warmup_frames={warmup_frames} " + f"steady_frames={len(steady_process)} " + f"steady_process_fps={_duration_fps(steady_process):.2f} " + f"steady_infer_fps={_duration_fps(steady_infer):.2f}" + ) + + +def _synchronize(upsampler: RealESRGANUpsampler) -> None: + if upsampler.device.type == "cuda": + torch.cuda.synchronize(upsampler.device) + + +def _duration_fps(durations: list[float]) -> float: + return _fps(len(durations), sum(durations)) + + +def _fps(frames: int, elapsed: float) -> float: + return frames / elapsed if elapsed > 0 else 0.0 + + +if __name__ == "__main__": + main() diff --git a/integrations/realesrgan/realesrgan/postprocess.py b/integrations/realesrgan/realesrgan/postprocess.py new file mode 100644 index 00000000..8f87a03d --- /dev/null +++ b/integrations/realesrgan/realesrgan/postprocess.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FlashDreams post-processing wrapper for Real-ESRGAN.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +import torch + +from flashdreams.infra.postprocess import ( + VideoChunk, + VideoPostProcessor, + VideoPostProcessorConfig, + VideoPostProcessorSession, + VideoSpec, + to_bvtchw, + to_minus_one_one, +) +from realesrgan.upsampler import RealESRGANUpsampler, default_model_name + + +@dataclass(kw_only=True) +class RealESRGANPostProcessorConfig(VideoPostProcessorConfig): + """Post-process generated RGB video with Real-ESRGAN.""" + + _target: type["RealESRGANPostProcessor"] = field( + default_factory=lambda: RealESRGANPostProcessor + ) + + scale: Literal[2, 4] = 2 + """Spatial upsample factor.""" + + model_name: str | None = None + """Real-ESRGAN checkpoint name. Defaults to the general model for scale.""" + + model_path: str | Path | None = None + """Optional local checkpoint path.""" + + tile: int = 0 + """Input tile size. ``0`` processes each frame as a single tensor.""" + + tile_pad: int = 10 + """Input tile overlap in pixels.""" + + pre_pad: int = 10 + """Reflection padding applied before model inference.""" + + half: bool = True + """Use fp16 on CUDA.""" + + compile_model: bool = False + """Compile the Real-ESRGAN model with ``torch.compile``.""" + + compile_mode: str = "reduce-overhead" + """Mode passed to ``torch.compile`` when :attr:`compile_model` is enabled.""" + + device: str = "cuda" + """Torch device used by the Real-ESRGAN model.""" + + +class RealESRGANPostProcessor(VideoPostProcessor[RealESRGANPostProcessorConfig]): + """Factory for Real-ESRGAN post-processing sessions.""" + + def start(self, spec: VideoSpec) -> VideoPostProcessorSession: + """Start one Real-ESRGAN post-processing stream.""" + return _RealESRGANPostProcessorSession(self.config) + + +class _RealESRGANPostProcessorSession(VideoPostProcessorSession): + def __init__(self, config: RealESRGANPostProcessorConfig) -> None: + self._config = config + self._upsampler: RealESRGANUpsampler | None = None + + @torch.no_grad() + def process(self, chunk: VideoChunk) -> list[VideoChunk]: + canonical = to_bvtchw( + to_minus_one_one(chunk.tensor, value_range=chunk.value_range), + layout=chunk.layout, + ) + batch, views, _, channels, _, _ = canonical.shape + if channels != 3: + raise ValueError( + f"Real-ESRGAN expects RGB chunks; got {channels} channels." + ) + upsampler = self._ensure_upsampler() + outputs = [] + for batch_idx in range(batch): + view_outputs = [] + for view_idx in range(views): + view_outputs.append( + upsampler.upsample_video_tensor(canonical[batch_idx, view_idx]) + ) + outputs.append(torch.stack(view_outputs, dim=0)) + output = torch.stack(outputs, dim=0) + return [ + VideoChunk( + tensor=output, + layout="bvtchw", + value_range="minus_one_one", + is_final=chunk.is_final, + metadata={**chunk.metadata, "source": "realesrgan"}, + ) + ] + + def flush(self) -> list[VideoChunk]: + """Real-ESRGAN is frame-local and keeps no buffered tail.""" + return [] + + def _ensure_upsampler(self) -> RealESRGANUpsampler: + if self._upsampler is None: + model_name = self._config.model_name or default_model_name( + self._config.scale + ) + self._upsampler = RealESRGANUpsampler( + scale=self._config.scale, + model_name=model_name, + model_path=self._config.model_path, + tile=self._config.tile, + tile_pad=self._config.tile_pad, + pre_pad=self._config.pre_pad, + half=self._config.half, + compile_model=self._config.compile_model, + compile_mode=self._config.compile_mode, + device=self._config.device, + ) + return self._upsampler diff --git a/integrations/realesrgan/realesrgan/upsampler.py b/integrations/realesrgan/realesrgan/upsampler.py new file mode 100644 index 00000000..e8dfc376 --- /dev/null +++ b/integrations/realesrgan/realesrgan/upsampler.py @@ -0,0 +1,429 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Real-ESRGAN checkpoint loading and tensor upsampling.""" + +from __future__ import annotations + +import math +import os +from pathlib import Path +from typing import Literal + +import cv2 +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F + +from flashdreams.core.io.download import download_to_cache +from realesrgan.archs import RRDBNet, SRVGGNetCompact + +RealESRGANModelName = Literal[ + "RealESRGAN_x4plus", + "RealESRNet_x4plus", + "RealESRGAN_x4plus_anime_6B", + "RealESRGAN_x2plus", + "realesr-animevideov3", + "realesr-general-x4v3", +] + +MODEL_URLS: dict[str, str] = { + "RealESRGAN_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", + "RealESRNet_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth", + "RealESRGAN_x4plus_anime_6B": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth", + "RealESRGAN_x2plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth", + "realesr-animevideov3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth", + "realesr-general-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", +} +"""Public Real-ESRGAN checkpoint URLs. Files are cached at runtime.""" + +MODEL_CONFIGS: dict[str, dict[str, int | str]] = { + "RealESRGAN_x4plus": { + "arch": "RRDBNet", + "num_in_ch": 3, + "num_out_ch": 3, + "num_feat": 64, + "num_block": 23, + "num_grow_ch": 32, + "scale": 4, + }, + "RealESRNet_x4plus": { + "arch": "RRDBNet", + "num_in_ch": 3, + "num_out_ch": 3, + "num_feat": 64, + "num_block": 23, + "num_grow_ch": 32, + "scale": 4, + }, + "RealESRGAN_x4plus_anime_6B": { + "arch": "RRDBNet", + "num_in_ch": 3, + "num_out_ch": 3, + "num_feat": 64, + "num_block": 6, + "num_grow_ch": 32, + "scale": 4, + }, + "RealESRGAN_x2plus": { + "arch": "RRDBNet", + "num_in_ch": 3, + "num_out_ch": 3, + "num_feat": 64, + "num_block": 23, + "num_grow_ch": 32, + "scale": 2, + }, + "realesr-animevideov3": { + "arch": "SRVGGNetCompact", + "num_in_ch": 3, + "num_out_ch": 3, + "num_feat": 64, + "num_conv": 16, + "upscale": 4, + "act_type": "prelu", + "scale": 4, + }, + "realesr-general-x4v3": { + "arch": "SRVGGNetCompact", + "num_in_ch": 3, + "num_out_ch": 3, + "num_feat": 64, + "num_conv": 32, + "upscale": 4, + "act_type": "prelu", + "scale": 4, + }, +} +"""Model architecture table for public Real-ESRGAN checkpoints.""" + +CACHE_DIR = ( + Path(os.path.expanduser(os.getenv("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams"))) + / "realesrgan" +) +"""User-writable cache directory for Real-ESRGAN checkpoint downloads.""" + + +def default_model_name(scale: Literal[2, 4]) -> RealESRGANModelName: + """Return the default public Real-ESRGAN model for ``scale``.""" + return "RealESRGAN_x2plus" if scale == 2 else "RealESRGAN_x4plus" + + +def create_model(model_name: str) -> tuple[nn.Module, int]: + """Create the architecture for a named Real-ESRGAN checkpoint.""" + if model_name not in MODEL_CONFIGS: + raise ValueError( + f"Unknown Real-ESRGAN model {model_name!r}; " + f"available models: {sorted(MODEL_CONFIGS)}" + ) + config = MODEL_CONFIGS[model_name] + arch = str(config["arch"]) + scale = int(config["scale"]) + if arch == "RRDBNet": + return ( + RRDBNet( + num_in_ch=int(config["num_in_ch"]), + num_out_ch=int(config["num_out_ch"]), + num_feat=int(config["num_feat"]), + num_block=int(config["num_block"]), + num_grow_ch=int(config["num_grow_ch"]), + scale=scale, + ), + scale, + ) + if arch == "SRVGGNetCompact": + return ( + SRVGGNetCompact( + num_in_ch=int(config["num_in_ch"]), + num_out_ch=int(config["num_out_ch"]), + num_feat=int(config["num_feat"]), + num_conv=int(config["num_conv"]), + upscale=int(config["upscale"]), + act_type=str(config["act_type"]), + ), + scale, + ) + raise ValueError(f"Unknown Real-ESRGAN architecture {arch!r}.") + + +def resolve_model_path(model_name: str, model_path: str | Path | None = None) -> Path: + """Return a local checkpoint path, downloading public weights if needed.""" + if model_path is not None: + return Path(model_path).expanduser() + if model_name not in MODEL_URLS: + raise ValueError(f"No public checkpoint URL for model {model_name!r}.") + return download_to_cache(MODEL_URLS[model_name], cache_dir=CACHE_DIR) + + +def load_weights(model: nn.Module, checkpoint_path: str | Path) -> None: + """Load Real-ESRGAN weights from ``checkpoint_path`` into ``model``.""" + checkpoint = torch.load( + Path(checkpoint_path).expanduser(), + map_location="cpu", + weights_only=False, + ) + if isinstance(checkpoint, dict) and "params_ema" in checkpoint: + state = checkpoint["params_ema"] + elif isinstance(checkpoint, dict) and "params" in checkpoint: + state = checkpoint["params"] + else: + state = checkpoint + model.load_state_dict(state, strict=True) + + +class RealESRGANUpsampler: + """Real-ESRGAN RGB tensor upsampler. + + Args: + scale: Output upsample factor. + model_name: Public Real-ESRGAN model name. Defaults to the general + x2 or x4 RRDBNet model matching ``scale``. + model_path: Optional local checkpoint path. + tile: Optional input tile size. ``0`` runs each frame as one tensor. + tile_pad: Input overlap on every tile edge. + pre_pad: Reflection padding applied before model inference. + half: Use fp16 on CUDA. + compile_model: Compile the model with ``torch.compile`` after loading. + compile_mode: Mode passed to ``torch.compile``. + device: Torch device string or object. + """ + + def __init__( + self, + *, + scale: Literal[2, 4] = 2, + model_name: str | None = None, + model_path: str | Path | None = None, + tile: int = 0, + tile_pad: int = 10, + pre_pad: int = 10, + half: bool = True, + compile_model: bool = False, + compile_mode: str = "reduce-overhead", + device: str | torch.device | None = None, + load_checkpoint: bool = True, + ) -> None: + self.scale = scale + self.model_name = model_name or default_model_name(scale) + self.tile = tile + self.tile_pad = tile_pad + self.pre_pad = pre_pad + self.half = half + self.compile_model = compile_model + self.compile_mode = compile_mode + self.device = torch.device( + device + if device is not None + else ("cuda" if torch.cuda.is_available() else "cpu") + ) + self.model, self.model_scale = create_model(self.model_name) + if self.model_scale != scale: + raise ValueError( + f"Model {self.model_name!r} has native scale {self.model_scale}; " + f"requested scale {scale}." + ) + if load_checkpoint: + load_weights(self.model, resolve_model_path(self.model_name, model_path)) + self.model.eval().to(self.device) + if self.half and self.device.type == "cuda": + self.model.half() + if self.compile_model: + self.model = torch.compile(self.model, mode=self.compile_mode) + self.dtype = ( + torch.float16 if self.half and self.device.type == "cuda" else torch.float32 + ) + + @torch.no_grad() + def upsample_video_tensor(self, video: torch.Tensor) -> torch.Tensor: + """Upsample ``video`` shaped ``[T, C, H, W]`` in ``[-1, 1]``.""" + if video.ndim != 4 or video.shape[1] != 3: + raise ValueError( + f"Expected [T, 3, H, W] video tensor; got {tuple(video.shape)}" + ) + frames = [self.upsample_frame_tensor(frame) for frame in video] + return torch.stack(frames, dim=0) + + @torch.no_grad() + def upsample_frame_tensor(self, frame: torch.Tensor) -> torch.Tensor: + """Upsample one RGB frame shaped ``[3, H, W]`` in ``[-1, 1]``.""" + if frame.ndim != 3 or frame.shape[0] != 3: + raise ValueError( + f"Expected [3, H, W] frame tensor; got {tuple(frame.shape)}" + ) + image = ((frame.to(self.device, dtype=self.dtype) + 1.0) * 0.5).clamp(0, 1) + output = self._process_image(image.unsqueeze(0)) + return (output.squeeze(0).float().cpu().clamp(0, 1) * 2.0) - 1.0 + + @torch.no_grad() + def upsample_bgr_image(self, image: np.ndarray) -> tuple[np.ndarray, str]: + """Upsample an OpenCV image array and return ``(image, mode)``.""" + mode = "RGB" + alpha: np.ndarray | None = None + if image.ndim == 2: + mode = "L" + image_rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) + elif image.shape[2] == 4: + mode = "RGBA" + alpha = image[:, :, 3] + image_rgb = cv2.cvtColor(image[:, :, :3], cv2.COLOR_BGR2RGB) + else: + image_rgb = cv2.cvtColor(image[:, :, :3], cv2.COLOR_BGR2RGB) + + max_range = 65535.0 if image.dtype == np.uint16 else 255.0 + tensor = torch.from_numpy( + (image_rgb.astype(np.float32) / max_range).transpose(2, 0, 1) + ) + output = (self.upsample_frame_tensor(tensor * 2.0 - 1.0) + 1.0) * 0.5 + if mode == "L": + return _rgb_tensor_to_gray_image(output, image.dtype), mode + output_bgr = _rgb_tensor_to_bgr_image(output, image.dtype) + if mode == "RGBA" and alpha is not None: + alpha_out = cv2.resize( + alpha, + (output_bgr.shape[1], output_bgr.shape[0]), + interpolation=cv2.INTER_LINEAR, + ) + output_bgr = _bgr_to_bgra(output_bgr, alpha_out) + + return output_bgr, mode + + def _process_image(self, image: torch.Tensor) -> torch.Tensor: + padded, crop = self._pad_input(image) + output = ( + self._process_tiled(padded) + if self.tile > 0 + else self.model(padded.to(dtype=self.dtype)) + ) + return self._crop_output(output, crop) + + def _pad_input(self, image: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]: + if self.pre_pad: + image = F.pad(image, (0, self.pre_pad, 0, self.pre_pad), "reflect") + mod_scale = 2 if self.model_scale == 2 else None + mod_pad_h = 0 + mod_pad_w = 0 + if mod_scale is not None: + _, _, height, width = image.shape + mod_pad_h = (mod_scale - height % mod_scale) % mod_scale + mod_pad_w = (mod_scale - width % mod_scale) % mod_scale + if mod_pad_h or mod_pad_w: + image = F.pad(image, (0, mod_pad_w, 0, mod_pad_h), "reflect") + return image, ( + (self.pre_pad + mod_pad_h) * self.model_scale, + (self.pre_pad + mod_pad_w) * self.model_scale, + ) + + def _crop_output(self, output: torch.Tensor, crop: tuple[int, int]) -> torch.Tensor: + crop_h, crop_w = crop + if crop_h: + output = output[:, :, :-crop_h, :] + if crop_w: + output = output[:, :, :, :-crop_w] + return output + + def _process_tiled(self, image: torch.Tensor) -> torch.Tensor: + batch, channels, height, width = image.shape + output = image.new_zeros( + batch, + channels, + height * self.model_scale, + width * self.model_scale, + ) + tiles_x = math.ceil(width / self.tile) + tiles_y = math.ceil(height / self.tile) + for y in range(tiles_y): + for x in range(tiles_x): + in_x0 = x * self.tile + in_y0 = y * self.tile + in_x1 = min(in_x0 + self.tile, width) + in_y1 = min(in_y0 + self.tile, height) + pad_x0 = max(in_x0 - self.tile_pad, 0) + pad_y0 = max(in_y0 - self.tile_pad, 0) + pad_x1 = min(in_x1 + self.tile_pad, width) + pad_y1 = min(in_y1 + self.tile_pad, height) + tile = image[:, :, pad_y0:pad_y1, pad_x0:pad_x1] + output_tile = self.model(tile.to(dtype=self.dtype)) + + out_x0 = in_x0 * self.model_scale + out_y0 = in_y0 * self.model_scale + out_x1 = in_x1 * self.model_scale + out_y1 = in_y1 * self.model_scale + tile_x0 = (in_x0 - pad_x0) * self.model_scale + tile_y0 = (in_y0 - pad_y0) * self.model_scale + tile_x1 = tile_x0 + (in_x1 - in_x0) * self.model_scale + tile_y1 = tile_y0 + (in_y1 - in_y0) * self.model_scale + output[:, :, out_y0:out_y1, out_x0:out_x1] = output_tile[ + :, :, tile_y0:tile_y1, tile_x0:tile_x1 + ] + return output + + +def write_bgr_image(path: str | Path, image: np.ndarray) -> None: + """Write ``image`` to ``path`` and raise if OpenCV rejects it.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + if not cv2.imwrite(str(path), image): + raise RuntimeError(f"Failed to write image to {path}") + + +def _rgb_tensor_to_bgr_image(tensor: torch.Tensor, dtype: np.dtype) -> np.ndarray: + """Convert a CPU ``[3, H, W]`` RGB float tensor in ``[0, 1]`` to BGR image.""" + image_dtype = _output_image_dtype(dtype) + scale = _dtype_max(image_dtype) + torch_dtype = _torch_dtype(image_dtype) + return ( + tensor.clamp(0, 1) + .mul(scale) + .round() + .to(torch_dtype)[[2, 1, 0]] + .permute(1, 2, 0) + .contiguous() + .numpy() + ) + + +def _rgb_tensor_to_gray_image(tensor: torch.Tensor, dtype: np.dtype) -> np.ndarray: + """Convert a CPU ``[3, H, W]`` RGB float tensor in ``[0, 1]`` to grayscale.""" + image_dtype = _output_image_dtype(dtype) + scale = _dtype_max(image_dtype) + torch_dtype = _torch_dtype(image_dtype) + tensor = tensor.clamp(0, 1) + gray = tensor[0] * 0.299 + tensor[1] * 0.587 + tensor[2] * 0.114 + return gray.mul(scale).round().to(torch_dtype).contiguous().numpy() + + +def _bgr_to_bgra(bgr: np.ndarray, alpha: np.ndarray) -> np.ndarray: + """Append ``alpha`` to a BGR image without routing through float color convert.""" + output = np.empty((*bgr.shape[:2], 4), dtype=bgr.dtype) + output[:, :, :3] = bgr + output[:, :, 3] = alpha.astype(bgr.dtype, copy=False) + return output + + +def _output_image_dtype(dtype: np.dtype) -> np.dtype: + return ( + np.dtype(np.uint16) + if np.dtype(dtype) == np.dtype(np.uint16) + else np.dtype(np.uint8) + ) + + +def _dtype_max(dtype: np.dtype) -> float: + return 65535.0 if np.dtype(dtype) == np.dtype(np.uint16) else 255.0 + + +def _torch_dtype(dtype: np.dtype) -> torch.dtype: + return torch.uint16 if np.dtype(dtype) == np.dtype(np.uint16) else torch.uint8 diff --git a/integrations/realesrgan/scripts/gpu_smoke.py b/integrations/realesrgan/scripts/gpu_smoke.py new file mode 100644 index 00000000..69a32d2a --- /dev/null +++ b/integrations/realesrgan/scripts/gpu_smoke.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run a checkpoint-backed Real-ESRGAN GPU smoke test.""" + +from __future__ import annotations + +import argparse +import time +from typing import Literal, cast + +import torch + +from realesrgan.upsampler import RealESRGANUpsampler, default_model_name + + +def main() -> None: + """Load a Real-ESRGAN checkpoint and run one synthetic RGB frame.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--scale", type=int, choices=(2, 4), default=2) + parser.add_argument("--model-name", type=str, default=None) + parser.add_argument("--size", type=int, default=16) + parser.add_argument("--tile", type=int, default=0) + parser.add_argument("--fp32", action="store_true") + parser.add_argument("--device", type=str, default="cuda") + args = parser.parse_args() + + if args.device.startswith("cuda") and not torch.cuda.is_available(): + raise SystemExit("CUDA is not available.") + + scale = cast(Literal[2, 4], args.scale) + model_name = args.model_name or default_model_name(scale) + model_start = time.perf_counter() + upsampler = RealESRGANUpsampler( + scale=scale, + model_name=model_name, + tile=args.tile, + pre_pad=0, + half=not args.fp32, + device=args.device, + ) + _synchronize(args.device) + model_elapsed = time.perf_counter() - model_start + + if args.device.startswith("cuda"): + torch.cuda.reset_peak_memory_stats(args.device) + frame = torch.zeros(3, args.size, args.size) + infer_start = time.perf_counter() + output = upsampler.upsample_frame_tensor(frame) + _synchronize(args.device) + infer_elapsed = time.perf_counter() - infer_start + + expected_shape = (3, args.size * scale, args.size * scale) + if output.shape != expected_shape: + raise RuntimeError( + f"Expected output shape {expected_shape}, got {output.shape}" + ) + if not torch.isfinite(output).all(): + raise RuntimeError("Real-ESRGAN output contains non-finite values.") + + peak_mb = ( + torch.cuda.max_memory_allocated(args.device) / 1024 / 1024 + if args.device.startswith("cuda") + else 0.0 + ) + device_name = ( + torch.cuda.get_device_name(args.device) + if args.device.startswith("cuda") + else args.device + ) + print(f"device={device_name}") + print( + f"output_shape={tuple(output.shape)} dtype={output.dtype} " + f"range=({float(output.min()):.4f}, {float(output.max()):.4f})" + ) + print( + f"model_load_s={model_elapsed:.3f} " + f"infer_s={infer_elapsed:.3f} peak_allocated_mb={peak_mb:.1f}" + ) + + +def _synchronize(device: str) -> None: + if device.startswith("cuda"): + torch.cuda.synchronize(device) + + +if __name__ == "__main__": + main() diff --git a/integrations/realesrgan/tests/test_realesrgan.py b/integrations/realesrgan/tests/test_realesrgan.py new file mode 100644 index 00000000..4de925b5 --- /dev/null +++ b/integrations/realesrgan/tests/test_realesrgan.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Real-ESRGAN upsampler integration.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pytest +import torch + +import realesrgan.postprocess as postprocess_mod +from realesrgan.postprocess import RealESRGANPostProcessorConfig +from realesrgan.upsampler import RealESRGANUpsampler, create_model + +from flashdreams.infra.postprocess import VideoChunk + +pytestmark = pytest.mark.ci_cpu + + +def test_create_x2_model_matches_expected_shape() -> None: + model, scale = create_model("RealESRGAN_x2plus") + model.eval() + x = torch.zeros(1, 3, 8, 8) + + with torch.no_grad(): + y = model(x) + + assert scale == 2 + assert y.shape == (1, 3, 16, 16) + + +def test_random_weight_upsampler_scales_tensor_without_checkpoint() -> None: + upsampler = RealESRGANUpsampler( + scale=2, + model_name="RealESRGAN_x2plus", + pre_pad=0, + half=False, + device="cpu", + load_checkpoint=False, + ) + frame = torch.zeros(3, 8, 8) + + output = upsampler.upsample_frame_tensor(frame) + + assert output.shape == (3, 16, 16) + assert output.dtype == torch.float32 + + +def test_upsampler_can_compile_model(monkeypatch: pytest.MonkeyPatch) -> None: + compile_calls = [] + + def fake_compile(model: torch.nn.Module, *, mode: str) -> torch.nn.Module: + compile_calls.append(mode) + return model + + monkeypatch.setattr(torch, "compile", fake_compile) + + RealESRGANUpsampler( + scale=2, + model_name="RealESRGAN_x2plus", + pre_pad=0, + half=False, + compile_model=True, + compile_mode="reduce-overhead", + device="cpu", + load_checkpoint=False, + ) + + assert compile_calls == ["reduce-overhead"] + + +def test_bgr_image_output_uses_opencv_channel_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + upsampler = RealESRGANUpsampler( + scale=2, + model_name="RealESRGAN_x2plus", + pre_pad=0, + half=False, + device="cpu", + load_checkpoint=False, + ) + + def fake_upsample_frame_tensor(frame: torch.Tensor) -> torch.Tensor: + rgb = torch.empty(3, 2, 2) + rgb[0].fill_(0.2) + rgb[1].fill_(0.4) + rgb[2].fill_(0.6) + return rgb * 2.0 - 1.0 + + monkeypatch.setattr(upsampler, "upsample_frame_tensor", fake_upsample_frame_tensor) + + image = np.zeros((1, 1, 3), dtype=np.uint8) + output, mode = upsampler.upsample_bgr_image(image) + + assert mode == "RGB" + assert output.dtype == np.uint8 + assert output.shape == (2, 2, 3) + assert output[0, 0].tolist() == [153, 102, 51] + + +def test_bgra_image_output_preserves_alpha(monkeypatch: pytest.MonkeyPatch) -> None: + upsampler = RealESRGANUpsampler( + scale=2, + model_name="RealESRGAN_x2plus", + pre_pad=0, + half=False, + device="cpu", + load_checkpoint=False, + ) + + def fake_upsample_frame_tensor(frame: torch.Tensor) -> torch.Tensor: + return torch.zeros(3, 2, 2) + + monkeypatch.setattr(upsampler, "upsample_frame_tensor", fake_upsample_frame_tensor) + + image = np.array([[[10, 20, 30, 77]]], dtype=np.uint8) + output, mode = upsampler.upsample_bgr_image(image) + + assert mode == "RGBA" + assert output.dtype == np.uint8 + assert output.shape == (2, 2, 4) + assert np.all(output[:, :, 3] == 77) + + +@dataclass +class _FakeUpsampler: + scale: int + model_name: str | None = None + model_path: str | None = None + tile: int = 0 + tile_pad: int = 10 + pre_pad: int = 10 + half: bool = True + compile_model: bool = False + compile_mode: str = "reduce-overhead" + device: str = "cuda" + + def upsample_video_tensor(self, video: torch.Tensor) -> torch.Tensor: + return video.repeat_interleave(self.scale, dim=-2).repeat_interleave( + self.scale, + dim=-1, + ) + + +def test_postprocess_scales_batched_views(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(postprocess_mod, "RealESRGANUpsampler", _FakeUpsampler) + session = ( + RealESRGANPostProcessorConfig(scale=2) + .setup() + .start(postprocess_mod.VideoSpec(height=2, width=3, channels=3)) + ) + chunk = VideoChunk( + tensor=torch.zeros(1, 2, 4, 3, 2, 3), + layout="bvtchw", + value_range="minus_one_one", + is_final=True, + ) + + outputs = session.process(chunk) + + assert len(outputs) == 1 + assert outputs[0].layout == "bvtchw" + assert outputs[0].metadata["source"] == "realesrgan" + assert outputs[0].tensor.shape == (1, 2, 4, 3, 4, 6) diff --git a/pyproject.toml b/pyproject.toml index 250048b8..40624cac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ extraPaths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/realesrgan", "integrations/self_forcing", "integrations/wan21", "integrations/wan22", @@ -74,6 +75,7 @@ extra-paths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/realesrgan", "integrations/self_forcing", "integrations/wan21", "integrations/wan22", diff --git a/uv.lock b/uv.lock index 2b0d0321..32ca7b09 100644 --- a/uv.lock +++ b/uv.lock @@ -27,6 +27,7 @@ members = [ "flashdreams-hy-worldplay", "flashdreams-lingbot", "flashdreams-omnidreams", + "flashdreams-realesrgan", "flashdreams-self-forcing", "flashdreams-wan21", "flashdreams-wan22", @@ -1567,6 +1568,28 @@ requires-dist = [ ] provides-extras = ["interactive-drive", "dev"] +[[package]] +name = "flashdreams-realesrgan" +version = "0.1.0" +source = { editable = "integrations/realesrgan" } +dependencies = [ + { name = "flashdreams" }, + { name = "opencv-python-headless" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", editable = "flashdreams" }, + { name = "opencv-python-headless", specifier = ">=4.8" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, +] +provides-extras = ["dev"] + [[package]] name = "flashdreams-self-forcing" version = "0.1.0" From 3c5e2f0d98fdd8631c63db08ae08c979d2d804da Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Wed, 24 Jun 2026 20:37:15 +0000 Subject: [PATCH 5/8] Add postprocess profiling and simplify Real-ESRGAN mode Signed-off-by: Gangzheng Tong --- flashdreams/flashdreams/scripts/cli.py | 1 + integrations/flashvsr/README.md | 6 ++ integrations/flashvsr/flashvsr/runner.py | 83 ++++++++++++++++++- integrations/realesrgan/README.md | 20 +++++ integrations/realesrgan/pyproject.toml | 3 + .../realesrgan/realesrgan/__init__.py | 8 +- integrations/realesrgan/realesrgan/cli.py | 33 +++++--- .../realesrgan/realesrgan/postprocess.py | 4 + .../realesrgan/realesrgan/upsampler.py | 60 ++++++++++++-- .../realesrgan/tests/test_realesrgan.py | 58 +++++++++++-- 10 files changed, 248 insertions(+), 28 deletions(-) diff --git a/flashdreams/flashdreams/scripts/cli.py b/flashdreams/flashdreams/scripts/cli.py index e6c69ae7..70dc2ae6 100644 --- a/flashdreams/flashdreams/scripts/cli.py +++ b/flashdreams/flashdreams/scripts/cli.py @@ -28,6 +28,7 @@ flashdreams-run wan21-i2v-14b-480p --prompt "..." --image-path frame.png flashdreams-run --no-instantiate template-offline # resolve config only flashdreams-run wan21-t2v-1.3b-480p --postprocess.preset flashvsr-v1.1-sparse-2.0 + flashdreams-run wan21-t2v-1.3b-480p --postprocess.preset realesrgan # Multi-GPU via context-parallelism (integration transformers auto-detect # CP size from the launcher's WORLD group). ``--no-python`` tells diff --git a/integrations/flashvsr/README.md b/integrations/flashvsr/README.md index 9a0ea644..54adb043 100644 --- a/integrations/flashvsr/README.md +++ b/integrations/flashvsr/README.md @@ -288,6 +288,12 @@ for i, (start, size) in enumerate(chunks): hand-rolled kernel) or `"torch"` (pure-torch wavelet + AdaIN reference). - `enable_sync_and_profile`: per-AR-step CUDA-event profiling. Adds one `cuda.synchronize()` per step. +- Runner `profile_warmup_frames`: output frames excluded from the steady FPS + summary. Warmup is applied at chunk granularity. Summary rows use the same + names as `realesrgan-upsample`: `model_fps` is CUDA-event timing for + `generate`/`finalize`, `pipeline_fps` is the synchronized FlashVSR pipeline + chunk path, `video_loop_fps` also includes the chunk CPU copy, and + `end_to_end_fps` covers the video pass through output write. ## Files diff --git a/integrations/flashvsr/flashvsr/runner.py b/integrations/flashvsr/flashvsr/runner.py index fcd6be57..12e0d3bf 100644 --- a/integrations/flashvsr/flashvsr/runner.py +++ b/integrations/flashvsr/flashvsr/runner.py @@ -19,6 +19,7 @@ import json import os +import time from dataclasses import dataclass, field from pathlib import Path from typing import Literal @@ -226,6 +227,11 @@ class FlashVSRRunnerConfig(RunnerConfig): """Sparse-attention budget multiplier used to re-derive ``transformer.topk_ratio`` from each video's post-crop target area.""" + profile_warmup_frames: int = 0 + """Output frames to exclude from steady-state FPS summaries. Warmup is + applied at chunk granularity, so any chunk touched by this threshold is + excluded from steady metrics.""" + def _ensure_mgpu_config_supported( config: FlashVSRRunnerConfig, world_size: int @@ -259,6 +265,27 @@ def _ensure_mgpu_config_supported( ) +def _fps(frames: int, elapsed_s: float) -> float: + return frames / elapsed_s if elapsed_s > 0 else 0.0 + + +def _duration_fps(frames: list[int], durations_s: list[float]) -> float: + return _fps(sum(frames), sum(durations_s)) + + +def _duration_ms_fps(frames: list[int], durations_ms: list[float]) -> float: + return _fps(sum(frames), sum(durations_ms) / 1000.0) + + +def _steady_start_index(frames: list[int], warmup_frames: int) -> int: + remaining = max(warmup_frames, 0) + index = 0 + while index < len(frames) and remaining > 0: + remaining -= frames[index] + index += 1 + return index + + ## Runner @@ -433,10 +460,23 @@ def run(self) -> None: cache = self._initialize_cache() + run_started_at = time.perf_counter() postprocess_stream = self.create_postprocess_stream(fps=fps) stats_history: list[dict[str, object]] = [] + profile_frames: list[int] = [] + video_loop_durations_s: list[float] = [] + pipeline_durations_s: list[float] = [] + model_durations_ms: list[float] = [] for chunk_idx, (start, size) in enumerate(chunks): + video_loop_started_at = time.perf_counter() clip = video_t[:, :, start : start + size] + pipeline_started_at = time.perf_counter() + start_event: torch.cuda.Event | None = None + end_event: torch.cuda.Event | None = None + if self.pipeline.device.type == "cuda": + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() video_chunk = self.pipeline.generate( autoregressive_index=chunk_idx, cache=cache, @@ -444,13 +484,25 @@ def run(self) -> None: ) pipeline_frames = int(video_chunk.shape[2]) stats = self.pipeline.finalize(autoregressive_index=chunk_idx, cache=cache) + if end_event is not None: + end_event.record() + torch.cuda.synchronize(self.pipeline.device) + assert start_event is not None + model_ms = float(start_event.elapsed_time(end_event)) + else: + model_ms = (time.perf_counter() - pipeline_started_at) * 1000.0 + pipeline_elapsed_s = time.perf_counter() - pipeline_started_at + chunk_frames = int(video_chunk.shape[2]) + profile_frames.append(chunk_frames) + pipeline_durations_s.append(pipeline_elapsed_s) + model_durations_ms.append(model_ms) postprocess_stream.process( video_chunk, autoregressive_index=chunk_idx, ) - if postprocess_stream.collect_output and stats is not None: - # Pipeline throughput is based on this AR step's direct output. - # Postprocess emission/buffering is reported separately. + if stats is not None: + # Report throughput against the visible output frames for this + # chunk. ``fps`` is reserved for the output video's frame rate. chunk_total_ms = stats["total_ms"] chunk_fps = ( pipeline_frames / chunk_total_ms * 1000.0 @@ -465,6 +517,7 @@ def run(self) -> None: "fps": chunk_fps, } ) + video_loop_durations_s.append(time.perf_counter() - video_loop_started_at) generated = postprocess_stream.finish() if generated is None: @@ -482,6 +535,30 @@ def run(self) -> None: f"-> {video_path.resolve()}" ) + elapsed_s = time.perf_counter() - run_started_at + output_frames = int(generated.shape[2]) + warmup_start = _steady_start_index(profile_frames, config.profile_warmup_frames) + steady_frames = profile_frames[warmup_start:] + steady_video_loop_s = video_loop_durations_s[warmup_start:] + steady_pipeline_s = pipeline_durations_s[warmup_start:] + steady_model_ms = model_durations_ms[warmup_start:] + logger.info( + "profile " + f"total_s={elapsed_s:.3f} " + f"end_to_end_fps={_fps(output_frames, elapsed_s):.2f} " + f"video_loop_fps={_duration_fps(profile_frames, video_loop_durations_s):.2f} " + f"pipeline_fps={_duration_fps(profile_frames, pipeline_durations_s):.2f} " + f"model_fps={_duration_ms_fps(profile_frames, model_durations_ms):.2f}" + ) + logger.info( + "profile_steady " + f"warmup_frames={config.profile_warmup_frames} " + f"steady_frames={sum(steady_frames)} " + f"steady_video_loop_fps={_duration_fps(steady_frames, steady_video_loop_s):.2f} " + f"steady_pipeline_fps={_duration_fps(steady_frames, steady_pipeline_s):.2f} " + f"steady_model_fps={_duration_ms_fps(steady_frames, steady_model_ms):.2f}" + ) + if stats_history: stats_path = config.output_dir / f"stats_{config.runner_name}.json" stats_path.write_text(json.dumps(stats_history, indent=2)) diff --git a/integrations/realesrgan/README.md b/integrations/realesrgan/README.md index c7f0d052..220af52f 100644 --- a/integrations/realesrgan/README.md +++ b/integrations/realesrgan/README.md @@ -50,6 +50,14 @@ uv run --package flashdreams-realesrgan realesrgan-upsample \ --compile ``` +Use Real-ESRGAN as a `flashdreams-run` postprocessor: + +```bash +uv run --package flashdreams-realesrgan flashdreams-run \ + --postprocess.preset realesrgan \ + wan21-t2v-1.3b-480p +``` + Useful flags: | flag | description | @@ -64,6 +72,18 @@ Useful flags: | `--profile-warmup-frames` | Frames excluded from steady FPS metrics. Defaults to `10`. | | `--max-frames` | Video-only frame cap for quick smoke tests. | +For `flashdreams-run`, Real-ESRGAN is selected with the registered +`--postprocess.preset realesrgan` preset. The preset uses the default 2x public +checkpoint on CUDA. Use the standalone `realesrgan-upsample` CLI or configure a +`RealESRGANPostProcessorConfig` directly when you need custom scale, tiling, +precision, checkpoint, or compile settings. + +Video runs print two profile rows. `model_fps` is CUDA-event timing for the +Real-ESRGAN model call only. `pipeline_fps` includes image conversion, +CPU/GPU copies, padding/cropping, the model call, and output conversion. +`video_loop_fps` adds video frame read/write around the pipeline call. +`end_to_end_fps` is the whole video pass wall-clock throughput. + ## FlashDreams Postprocess API ```python diff --git a/integrations/realesrgan/pyproject.toml b/integrations/realesrgan/pyproject.toml index 44034a2e..4f979090 100644 --- a/integrations/realesrgan/pyproject.toml +++ b/integrations/realesrgan/pyproject.toml @@ -42,6 +42,9 @@ dev = [ [project.scripts] realesrgan-upsample = "realesrgan.cli:main" +[project.entry-points."flashdreams.postprocess_presets"] +realesrgan = "realesrgan.postprocess:POSTPROCESS_PRESET_REALESRGAN" + [tool.setuptools.packages.find] include = ["realesrgan*"] exclude = ["tests"] diff --git a/integrations/realesrgan/realesrgan/__init__.py b/integrations/realesrgan/realesrgan/__init__.py index f0efe4e6..fb64357f 100644 --- a/integrations/realesrgan/realesrgan/__init__.py +++ b/integrations/realesrgan/realesrgan/__init__.py @@ -15,10 +15,14 @@ """Real-ESRGAN upsampling integration for FlashDreams.""" -from realesrgan.postprocess import RealESRGANPostProcessorConfig +from realesrgan.postprocess import ( + POSTPROCESS_PRESET_REALESRGAN, + RealESRGANPostProcessorConfig, +) from realesrgan.upsampler import ( MODEL_CONFIGS, MODEL_URLS, + RealESRGANFrameProfile, RealESRGANUpsampler, create_model, ) @@ -26,6 +30,8 @@ __all__ = [ "MODEL_CONFIGS", "MODEL_URLS", + "POSTPROCESS_PRESET_REALESRGAN", + "RealESRGANFrameProfile", "RealESRGANPostProcessorConfig", "RealESRGANUpsampler", "create_model", diff --git a/integrations/realesrgan/realesrgan/cli.py b/integrations/realesrgan/realesrgan/cli.py index 21442191..b7394421 100644 --- a/integrations/realesrgan/realesrgan/cli.py +++ b/integrations/realesrgan/realesrgan/cli.py @@ -150,17 +150,19 @@ def _upsample_video( start = time.perf_counter() frame_idx = 0 process_durations: list[float] = [] - infer_durations: list[float] = [] + pipeline_durations: list[float] = [] + model_durations_ms: list[float | None] = [] try: while True: + process_start = time.perf_counter() ok, frame = cap.read() if not ok: break - process_start = time.perf_counter() - infer_start = time.perf_counter() - output, _mode = upsampler.upsample_bgr_image(frame) + pipeline_start = time.perf_counter() + output, _mode, profile = upsampler.upsample_bgr_image_profiled(frame) _synchronize(upsampler) - infer_durations.append(time.perf_counter() - infer_start) + pipeline_durations.append(time.perf_counter() - pipeline_start) + model_durations_ms.append(profile.model_ms) writer.write(output) process_durations.append(time.perf_counter() - process_start) frame_idx += 1 @@ -174,18 +176,22 @@ def _upsample_video( total_note = f"/{total_frames}" if total_frames else "" warmup_frames = min(max(profile_warmup_frames, 0), len(process_durations)) steady_process = process_durations[warmup_frames:] - steady_infer = infer_durations[warmup_frames:] + steady_pipeline = pipeline_durations[warmup_frames:] + steady_model_ms = model_durations_ms[warmup_frames:] print(f"wrote {frame_idx}{total_note} frames -> {output_path}") print( - f"profile total_s={elapsed:.3f} total_fps={_fps(frame_idx, elapsed):.2f} " - f"process_fps={_duration_fps(process_durations):.2f} " - f"infer_fps={_duration_fps(infer_durations):.2f}" + f"profile total_s={elapsed:.3f} " + f"end_to_end_fps={_fps(frame_idx, elapsed):.2f} " + f"video_loop_fps={_duration_fps(process_durations):.2f} " + f"pipeline_fps={_duration_fps(pipeline_durations):.2f} " + f"model_fps={_duration_ms_fps(model_durations_ms):.2f}" ) print( f"profile_steady warmup_frames={warmup_frames} " f"steady_frames={len(steady_process)} " - f"steady_process_fps={_duration_fps(steady_process):.2f} " - f"steady_infer_fps={_duration_fps(steady_infer):.2f}" + f"steady_video_loop_fps={_duration_fps(steady_process):.2f} " + f"steady_pipeline_fps={_duration_fps(steady_pipeline):.2f} " + f"steady_model_fps={_duration_ms_fps(steady_model_ms):.2f}" ) @@ -198,6 +204,11 @@ def _duration_fps(durations: list[float]) -> float: return _fps(len(durations), sum(durations)) +def _duration_ms_fps(durations_ms: list[float | None]) -> float: + durations = [duration for duration in durations_ms if duration is not None] + return _fps(len(durations), sum(durations) / 1000.0) + + def _fps(frames: int, elapsed: float) -> float: return frames / elapsed if elapsed > 0 else 0.0 diff --git a/integrations/realesrgan/realesrgan/postprocess.py b/integrations/realesrgan/realesrgan/postprocess.py index 8f87a03d..a1e7a92d 100644 --- a/integrations/realesrgan/realesrgan/postprocess.py +++ b/integrations/realesrgan/realesrgan/postprocess.py @@ -82,6 +82,10 @@ def start(self, spec: VideoSpec) -> VideoPostProcessorSession: return _RealESRGANPostProcessorSession(self.config) +POSTPROCESS_PRESET_REALESRGAN = RealESRGANPostProcessorConfig() +"""Default Real-ESRGAN post-processing preset.""" + + class _RealESRGANPostProcessorSession(VideoPostProcessorSession): def __init__(self, config: RealESRGANPostProcessorConfig) -> None: self._config = config diff --git a/integrations/realesrgan/realesrgan/upsampler.py b/integrations/realesrgan/realesrgan/upsampler.py index e8dfc376..787dd899 100644 --- a/integrations/realesrgan/realesrgan/upsampler.py +++ b/integrations/realesrgan/realesrgan/upsampler.py @@ -19,6 +19,7 @@ import math import os +from dataclasses import dataclass from pathlib import Path from typing import Literal @@ -117,6 +118,14 @@ """User-writable cache directory for Real-ESRGAN checkpoint downloads.""" +@dataclass(frozen=True, kw_only=True) +class RealESRGANFrameProfile: + """Per-frame timing captured while producing a Real-ESRGAN output.""" + + model_ms: float | None + """CUDA-event time for the model/tiled-model call, excluding CPU I/O.""" + + def default_model_name(scale: Literal[2, 4]) -> RealESRGANModelName: """Return the default public Real-ESRGAN model for ``scale``.""" return "RealESRGAN_x2plus" if scale == 2 else "RealESRGAN_x4plus" @@ -259,17 +268,34 @@ def upsample_video_tensor(self, video: torch.Tensor) -> torch.Tensor: @torch.no_grad() def upsample_frame_tensor(self, frame: torch.Tensor) -> torch.Tensor: """Upsample one RGB frame shaped ``[3, H, W]`` in ``[-1, 1]``.""" + output, _profile = self.upsample_frame_tensor_profiled(frame) + return output + + @torch.no_grad() + def upsample_frame_tensor_profiled( + self, frame: torch.Tensor + ) -> tuple[torch.Tensor, RealESRGANFrameProfile]: + """Upsample one RGB frame and return model-only timing metadata.""" if frame.ndim != 3 or frame.shape[0] != 3: raise ValueError( f"Expected [3, H, W] frame tensor; got {tuple(frame.shape)}" ) image = ((frame.to(self.device, dtype=self.dtype) + 1.0) * 0.5).clamp(0, 1) - output = self._process_image(image.unsqueeze(0)) - return (output.squeeze(0).float().cpu().clamp(0, 1) * 2.0) - 1.0 + output, profile = self._process_image_profiled(image.unsqueeze(0)) + output = (output.squeeze(0).float().cpu().clamp(0, 1) * 2.0) - 1.0 + return output, profile @torch.no_grad() def upsample_bgr_image(self, image: np.ndarray) -> tuple[np.ndarray, str]: """Upsample an OpenCV image array and return ``(image, mode)``.""" + output, mode, _profile = self.upsample_bgr_image_profiled(image) + return output, mode + + @torch.no_grad() + def upsample_bgr_image_profiled( + self, image: np.ndarray + ) -> tuple[np.ndarray, str, RealESRGANFrameProfile]: + """Upsample an OpenCV image array and return model timing metadata.""" mode = "RGB" alpha: np.ndarray | None = None if image.ndim == 2: @@ -286,9 +312,10 @@ def upsample_bgr_image(self, image: np.ndarray) -> tuple[np.ndarray, str]: tensor = torch.from_numpy( (image_rgb.astype(np.float32) / max_range).transpose(2, 0, 1) ) - output = (self.upsample_frame_tensor(tensor * 2.0 - 1.0) + 1.0) * 0.5 + output, profile = self.upsample_frame_tensor_profiled(tensor * 2.0 - 1.0) + output = (output + 1.0) * 0.5 if mode == "L": - return _rgb_tensor_to_gray_image(output, image.dtype), mode + return _rgb_tensor_to_gray_image(output, image.dtype), mode, profile output_bgr = _rgb_tensor_to_bgr_image(output, image.dtype) if mode == "RGBA" and alpha is not None: alpha_out = cv2.resize( @@ -298,16 +325,37 @@ def upsample_bgr_image(self, image: np.ndarray) -> tuple[np.ndarray, str]: ) output_bgr = _bgr_to_bgra(output_bgr, alpha_out) - return output_bgr, mode + return output_bgr, mode, profile def _process_image(self, image: torch.Tensor) -> torch.Tensor: + output, _profile = self._process_image_profiled(image) + return output + + def _process_image_profiled( + self, image: torch.Tensor + ) -> tuple[torch.Tensor, RealESRGANFrameProfile]: padded, crop = self._pad_input(image) + start: torch.cuda.Event | None = None + end: torch.cuda.Event | None = None + if self.device.type == "cuda": + if self.compile_model: + torch.compiler.cudagraph_mark_step_begin() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() output = ( self._process_tiled(padded) if self.tile > 0 else self.model(padded.to(dtype=self.dtype)) ) - return self._crop_output(output, crop) + if end is not None: + end.record() + output = self._crop_output(output, crop) + model_ms: float | None = None + if start is not None and end is not None: + torch.cuda.synchronize(self.device) + model_ms = float(start.elapsed_time(end)) + return output, RealESRGANFrameProfile(model_ms=model_ms) def _pad_input(self, image: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]: if self.pre_pad: diff --git a/integrations/realesrgan/tests/test_realesrgan.py b/integrations/realesrgan/tests/test_realesrgan.py index 4de925b5..379c45c0 100644 --- a/integrations/realesrgan/tests/test_realesrgan.py +++ b/integrations/realesrgan/tests/test_realesrgan.py @@ -25,7 +25,11 @@ import realesrgan.postprocess as postprocess_mod from realesrgan.postprocess import RealESRGANPostProcessorConfig -from realesrgan.upsampler import RealESRGANUpsampler, create_model +from realesrgan.upsampler import ( + RealESRGANFrameProfile, + RealESRGANUpsampler, + create_model, +) from flashdreams.infra.postprocess import VideoChunk @@ -61,6 +65,23 @@ def test_random_weight_upsampler_scales_tensor_without_checkpoint() -> None: assert output.dtype == torch.float32 +def test_random_weight_upsampler_profiled_tensor_reports_cpu_profile() -> None: + upsampler = RealESRGANUpsampler( + scale=2, + model_name="RealESRGAN_x2plus", + pre_pad=0, + half=False, + device="cpu", + load_checkpoint=False, + ) + frame = torch.zeros(3, 8, 8) + + output, profile = upsampler.upsample_frame_tensor_profiled(frame) + + assert output.shape == (3, 16, 16) + assert profile == RealESRGANFrameProfile(model_ms=None) + + def test_upsampler_can_compile_model(monkeypatch: pytest.MonkeyPatch) -> None: compile_calls = [] @@ -96,14 +117,20 @@ def test_bgr_image_output_uses_opencv_channel_order( load_checkpoint=False, ) - def fake_upsample_frame_tensor(frame: torch.Tensor) -> torch.Tensor: + def fake_upsample_frame_tensor_profiled( + frame: torch.Tensor, + ) -> tuple[torch.Tensor, RealESRGANFrameProfile]: rgb = torch.empty(3, 2, 2) rgb[0].fill_(0.2) rgb[1].fill_(0.4) rgb[2].fill_(0.6) - return rgb * 2.0 - 1.0 + return rgb * 2.0 - 1.0, RealESRGANFrameProfile(model_ms=None) - monkeypatch.setattr(upsampler, "upsample_frame_tensor", fake_upsample_frame_tensor) + monkeypatch.setattr( + upsampler, + "upsample_frame_tensor_profiled", + fake_upsample_frame_tensor_profiled, + ) image = np.zeros((1, 1, 3), dtype=np.uint8) output, mode = upsampler.upsample_bgr_image(image) @@ -124,10 +151,16 @@ def test_bgra_image_output_preserves_alpha(monkeypatch: pytest.MonkeyPatch) -> N load_checkpoint=False, ) - def fake_upsample_frame_tensor(frame: torch.Tensor) -> torch.Tensor: - return torch.zeros(3, 2, 2) + def fake_upsample_frame_tensor_profiled( + frame: torch.Tensor, + ) -> tuple[torch.Tensor, RealESRGANFrameProfile]: + return torch.zeros(3, 2, 2), RealESRGANFrameProfile(model_ms=None) - monkeypatch.setattr(upsampler, "upsample_frame_tensor", fake_upsample_frame_tensor) + monkeypatch.setattr( + upsampler, + "upsample_frame_tensor_profiled", + fake_upsample_frame_tensor_profiled, + ) image = np.array([[[10, 20, 30, 77]]], dtype=np.uint8) output, mode = upsampler.upsample_bgr_image(image) @@ -178,3 +211,14 @@ def test_postprocess_scales_batched_views(monkeypatch: pytest.MonkeyPatch) -> No assert outputs[0].layout == "bvtchw" assert outputs[0].metadata["source"] == "realesrgan" assert outputs[0].tensor.shape == (1, 2, 4, 3, 4, 6) + + +def test_realesrgan_postprocess_preset_is_registered() -> None: + from flashdreams.plugins.registry import resolve_postprocess_preset + + processor = resolve_postprocess_preset("realesrgan") + + assert isinstance(processor, RealESRGANPostProcessorConfig) + assert processor.scale == 2 + assert processor.device == "cuda" + assert processor.half is True From 93022a784d85d69c0cd50c242eadf04614461224 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Wed, 1 Jul 2026 18:39:50 +0000 Subject: [PATCH 6/8] refactor(serving): share uplift gRPC protocol --- .../flashdreams/serving/uplift/__init__.py | 4 + .../serving/uplift}/protos/__init__.py | 3 +- .../serving/uplift/protos/uplift.proto | 11 +- .../serving/uplift/protos/uplift_pb2.py | 56 ++ .../serving/uplift/protos/uplift_pb2.pyi | 0 .../serving/uplift/protos/uplift_pb2_grpc.py | 111 +-- .../serving/uplift}/streaming_view.py | 47 +- .../serving/uplift}/uplift_client.py | 105 +- flashdreams/pyproject.toml | 6 + integrations/flashvsr/README.md | 20 +- .../flashvsr/grpc/protos/flashvsr_pb2.py | 56 -- .../flashvsr/flashvsr/grpc/uplift_server.py | 29 +- integrations/flashvsr/pyproject.toml | 11 +- .../flashvsr/tests/test_grpc_client.py | 32 +- .../flashvsr/tests/test_streaming_view.py | 75 ++ integrations/realesrgan/pyproject.toml | 3 +- .../realesrgan/realesrgan/grpc/__init__.py | 4 + .../realesrgan/grpc/uplift_server.py | 917 ++++++++++++++++++ .../scripts/feed_realesrgan_uplift_video.sh | 31 + .../scripts/run_realesrgan_uplift_server.sh | 19 + .../realesrgan/tests/test_uplift_server.py | 127 +++ uv.lock | 22 +- 22 files changed, 1501 insertions(+), 188 deletions(-) create mode 100644 flashdreams/flashdreams/serving/uplift/__init__.py rename {integrations/flashvsr/flashvsr/grpc => flashdreams/flashdreams/serving/uplift}/protos/__init__.py (90%) rename integrations/flashvsr/flashvsr/grpc/protos/flashvsr.proto => flashdreams/flashdreams/serving/uplift/protos/uplift.proto (95%) create mode 100644 flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.py rename integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2.pyi => flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.pyi (100%) rename integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2_grpc.py => flashdreams/flashdreams/serving/uplift/protos/uplift_pb2_grpc.py (63%) rename {integrations/flashvsr/flashvsr/grpc => flashdreams/flashdreams/serving/uplift}/streaming_view.py (92%) rename {integrations/flashvsr/flashvsr/grpc => flashdreams/flashdreams/serving/uplift}/uplift_client.py (85%) delete mode 100644 integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2.py create mode 100644 integrations/flashvsr/tests/test_streaming_view.py create mode 100644 integrations/realesrgan/realesrgan/grpc/__init__.py create mode 100644 integrations/realesrgan/realesrgan/grpc/uplift_server.py create mode 100755 integrations/realesrgan/scripts/feed_realesrgan_uplift_video.sh create mode 100755 integrations/realesrgan/scripts/run_realesrgan_uplift_server.sh create mode 100644 integrations/realesrgan/tests/test_uplift_server.py diff --git a/flashdreams/flashdreams/serving/uplift/__init__.py b/flashdreams/flashdreams/serving/uplift/__init__.py new file mode 100644 index 00000000..559cc7ac --- /dev/null +++ b/flashdreams/flashdreams/serving/uplift/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared gRPC uplift protocol, client, and browser-viewer utilities.""" diff --git a/integrations/flashvsr/flashvsr/grpc/protos/__init__.py b/flashdreams/flashdreams/serving/uplift/protos/__init__.py similarity index 90% rename from integrations/flashvsr/flashvsr/grpc/protos/__init__.py rename to flashdreams/flashdreams/serving/uplift/protos/__init__.py index e8ae6167..b33e22d7 100644 --- a/integrations/flashvsr/flashvsr/grpc/protos/__init__.py +++ b/flashdreams/flashdreams/serving/uplift/protos/__init__.py @@ -13,5 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Generated protobuf modules for FlashVSR gRPC serving.""" - +"""Generated protobuf modules for FlashDreams video uplift serving.""" diff --git a/integrations/flashvsr/flashvsr/grpc/protos/flashvsr.proto b/flashdreams/flashdreams/serving/uplift/protos/uplift.proto similarity index 95% rename from integrations/flashvsr/flashvsr/grpc/protos/flashvsr.proto rename to flashdreams/flashdreams/serving/uplift/protos/uplift.proto index 4efab6b2..4ff2ab45 100644 --- a/integrations/flashvsr/flashvsr/grpc/protos/flashvsr.proto +++ b/flashdreams/flashdreams/serving/uplift/protos/uplift.proto @@ -15,12 +15,13 @@ syntax = "proto3"; -package flashvsr; +package flashdreams.uplift; // --------------------------------------------------------------------------- -// FlashVSR gRPC Service +// FlashDreams video uplift gRPC service // -// Streaming video super-resolution (2× or 4×) using a diffusion-based model. +// Streaming video super-resolution (2× or 4×) using any server-side model +// implementation that speaks this protocol. // Two usage patterns: // // Unary chunk-by-chunk flow (best for integration into existing pipelines): @@ -36,7 +37,7 @@ package flashvsr; // get_status — liveness/readiness check // --------------------------------------------------------------------------- -service FlashVSR { +service VideoUplift { rpc get_status (StatusRequest) returns (StatusResponse); rpc start_session (StartSessionRequest) returns (StartSessionResponse); rpc end_session (EndSessionRequest) returns (EndSessionResponse); @@ -53,7 +54,7 @@ message StatusRequest {} message StatusResponse { bool ready = 1; // true once the model is loaded string device = 2; // e.g. "cuda:0" - string model_name = 3; // e.g. "FlashVSR-v1.1" + string model_name = 3; // e.g. "FlashVSR-v1.1" or "RealESRGAN/x2plus" repeated string active_sessions = 4; // currently open session IDs } diff --git a/flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.py b/flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.py new file mode 100644 index 00000000..936b5a3a --- /dev/null +++ b/flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.py @@ -0,0 +1,56 @@ +# -*- coding: utf-8 -*- +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: flashdreams/serving/uplift/protos/uplift.proto +# Protobuf Python Version: 6.31.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 31, + 1, + '', + 'flashdreams/serving/uplift/protos/uplift.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n.flashdreams/serving/uplift/protos/uplift.proto\x12\x12\x66lashdreams.uplift\"\x0f\n\rStatusRequest\"\\\n\x0eStatusResponse\x12\r\n\x05ready\x18\x01 \x01(\x08\x12\x0e\n\x06\x64\x65vice\x18\x02 \x01(\t\x12\x12\n\nmodel_name\x18\x03 \x01(\t\x12\x17\n\x0f\x61\x63tive_sessions\x18\x04 \x03(\t\"y\n\x13StartSessionRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x14\n\x0cinput_height\x18\x02 \x01(\x05\x12\x13\n\x0binput_width\x18\x03 \x01(\x05\x12\r\n\x05scale\x18\x04 \x01(\x05\x12\x14\n\x0csparse_ratio\x18\x05 \x01(\x02\"J\n\x14StartSessionResponse\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x0f\n\x07success\x18\x02 \x01(\x08\x12\r\n\x05\x65rror\x18\x03 \x01(\t\"\'\n\x11\x45ndSessionRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\"%\n\x12\x45ndSessionResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\"\xbb\x02\n\x13UpscaleChunkRequest\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x14\n\x0cinput_height\x18\x02 \x01(\x05\x12\x13\n\x0binput_width\x18\x03 \x01(\x05\x12\r\n\x05scale\x18\x04 \x01(\x05\x12\x14\n\x0csparse_ratio\x18\x05 \x01(\x02\x12\x12\n\nframes_rgb\x18\x06 \x01(\x0c\x12\x12\n\nnum_frames\x18\x07 \x01(\x05\x12\x0e\n\x06height\x18\x08 \x01(\x05\x12\r\n\x05width\x18\t \x01(\x05\x12\x13\n\x0b\x63hunk_index\x18\n \x01(\x05\x12\x39\n\x0e\x66rame_encoding\x18\x0b \x01(\x0e\x32!.flashdreams.uplift.FrameEncoding\x12\x13\n\x0b\x66rames_jpeg\x18\x0c \x03(\x0c\x12\x14\n\x0c\x64isplay_only\x18\r \x01(\x08\"\xc1\x01\n\x14UpscaleChunkResponse\x12\x12\n\nsession_id\x18\x01 \x01(\t\x12\x12\n\nframes_rgb\x18\x02 \x01(\x0c\x12\x12\n\nnum_frames\x18\x03 \x01(\x05\x12\x0e\n\x06height\x18\x04 \x01(\x05\x12\r\n\x05width\x18\x05 \x01(\x05\x12\x13\n\x0b\x63hunk_index\x18\x06 \x01(\x05\x12\x12\n\nelapsed_ms\x18\x07 \x01(\x02\x12\r\n\x05\x65rror\x18\x08 \x01(\t\x12\x16\n\x0e\x66rames_omitted\x18\t \x01(\x08*D\n\rFrameEncoding\x12\x1a\n\x16\x46RAME_ENCODING_RAW_RGB\x10\x00\x12\x17\n\x13\x46RAME_ENCODING_JPEG\x10\x01\x32\xf0\x03\n\x0bVideoUplift\x12S\n\nget_status\x12!.flashdreams.uplift.StatusRequest\x1a\".flashdreams.uplift.StatusResponse\x12\x62\n\rstart_session\x12\'.flashdreams.uplift.StartSessionRequest\x1a(.flashdreams.uplift.StartSessionResponse\x12\\\n\x0b\x65nd_session\x12%.flashdreams.uplift.EndSessionRequest\x1a&.flashdreams.uplift.EndSessionResponse\x12\x62\n\rupscale_chunk\x12\'.flashdreams.uplift.UpscaleChunkRequest\x1a(.flashdreams.uplift.UpscaleChunkResponse\x12\x66\n\rupscale_video\x12\'.flashdreams.uplift.UpscaleChunkRequest\x1a(.flashdreams.uplift.UpscaleChunkResponse(\x01\x30\x01\x62\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flashdreams.serving.uplift.protos.uplift_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_FRAMEENCODING']._serialized_start=974 + _globals['_FRAMEENCODING']._serialized_end=1042 + _globals['_STATUSREQUEST']._serialized_start=70 + _globals['_STATUSREQUEST']._serialized_end=85 + _globals['_STATUSRESPONSE']._serialized_start=87 + _globals['_STATUSRESPONSE']._serialized_end=179 + _globals['_STARTSESSIONREQUEST']._serialized_start=181 + _globals['_STARTSESSIONREQUEST']._serialized_end=302 + _globals['_STARTSESSIONRESPONSE']._serialized_start=304 + _globals['_STARTSESSIONRESPONSE']._serialized_end=378 + _globals['_ENDSESSIONREQUEST']._serialized_start=380 + _globals['_ENDSESSIONREQUEST']._serialized_end=419 + _globals['_ENDSESSIONRESPONSE']._serialized_start=421 + _globals['_ENDSESSIONRESPONSE']._serialized_end=458 + _globals['_UPSCALECHUNKREQUEST']._serialized_start=461 + _globals['_UPSCALECHUNKREQUEST']._serialized_end=776 + _globals['_UPSCALECHUNKRESPONSE']._serialized_start=779 + _globals['_UPSCALECHUNKRESPONSE']._serialized_end=972 + _globals['_VIDEOUPLIFT']._serialized_start=1045 + _globals['_VIDEOUPLIFT']._serialized_end=1541 +# @@protoc_insertion_point(module_scope) diff --git a/integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2.pyi b/flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.pyi similarity index 100% rename from integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2.pyi rename to flashdreams/flashdreams/serving/uplift/protos/uplift_pb2.pyi diff --git a/integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2_grpc.py b/flashdreams/flashdreams/serving/uplift/protos/uplift_pb2_grpc.py similarity index 63% rename from integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2_grpc.py rename to flashdreams/flashdreams/serving/uplift/protos/uplift_pb2_grpc.py index d1d1f52d..246b7777 100644 --- a/integrations/flashvsr/flashvsr/grpc/protos/flashvsr_pb2_grpc.py +++ b/flashdreams/flashdreams/serving/uplift/protos/uplift_pb2_grpc.py @@ -5,7 +5,7 @@ import grpc import warnings -from flashvsr.grpc.protos import flashvsr_pb2 as flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2 +from flashdreams.serving.uplift.protos import uplift_pb2 as flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2 GRPC_GENERATED_VERSION = '1.80.0' GRPC_VERSION = grpc.__version__ @@ -20,18 +20,19 @@ if _version_not_supported: raise RuntimeError( f'The grpc package installed is at version {GRPC_VERSION},' - + ' but the generated code in flashvsr/grpc/protos/flashvsr_pb2_grpc.py depends on' + + ' but the generated code in flashdreams/serving/uplift/protos/uplift_pb2_grpc.py depends on' + f' grpcio>={GRPC_GENERATED_VERSION}.' + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' ) -class FlashVSRStub(object): +class VideoUpliftStub(object): """--------------------------------------------------------------------------- - FlashVSR gRPC Service + FlashDreams video uplift gRPC service - Streaming video super-resolution (2× or 4×) using a diffusion-based model. + Streaming video super-resolution (2× or 4×) using any server-side model + implementation that speaks this protocol. Two usage patterns: Unary chunk-by-chunk flow (best for integration into existing pipelines): @@ -56,37 +57,38 @@ def __init__(self, channel): channel: A grpc.Channel. """ self.get_status = channel.unary_unary( - '/flashvsr.FlashVSR/get_status', - request_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StatusRequest.SerializeToString, - response_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StatusResponse.FromString, + '/flashdreams.uplift.VideoUplift/get_status', + request_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StatusRequest.SerializeToString, + response_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StatusResponse.FromString, _registered_method=True) self.start_session = channel.unary_unary( - '/flashvsr.FlashVSR/start_session', - request_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StartSessionRequest.SerializeToString, - response_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StartSessionResponse.FromString, + '/flashdreams.uplift.VideoUplift/start_session', + request_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StartSessionRequest.SerializeToString, + response_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StartSessionResponse.FromString, _registered_method=True) self.end_session = channel.unary_unary( - '/flashvsr.FlashVSR/end_session', - request_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.EndSessionRequest.SerializeToString, - response_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.EndSessionResponse.FromString, + '/flashdreams.uplift.VideoUplift/end_session', + request_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.EndSessionRequest.SerializeToString, + response_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.EndSessionResponse.FromString, _registered_method=True) self.upscale_chunk = channel.unary_unary( - '/flashvsr.FlashVSR/upscale_chunk', - request_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkRequest.SerializeToString, - response_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkResponse.FromString, + '/flashdreams.uplift.VideoUplift/upscale_chunk', + request_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkRequest.SerializeToString, + response_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkResponse.FromString, _registered_method=True) self.upscale_video = channel.stream_stream( - '/flashvsr.FlashVSR/upscale_video', - request_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkRequest.SerializeToString, - response_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkResponse.FromString, + '/flashdreams.uplift.VideoUplift/upscale_video', + request_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkRequest.SerializeToString, + response_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkResponse.FromString, _registered_method=True) -class FlashVSRServicer(object): +class VideoUpliftServicer(object): """--------------------------------------------------------------------------- - FlashVSR gRPC Service + FlashDreams video uplift gRPC service - Streaming video super-resolution (2× or 4×) using a diffusion-based model. + Streaming video super-resolution (2× or 4×) using any server-side model + implementation that speaks this protocol. Two usage patterns: Unary chunk-by-chunk flow (best for integration into existing pipelines): @@ -135,46 +137,47 @@ def upscale_video(self, request_iterator, context): raise NotImplementedError('Method not implemented!') -def add_FlashVSRServicer_to_server(servicer, server): +def add_VideoUpliftServicer_to_server(servicer, server): rpc_method_handlers = { 'get_status': grpc.unary_unary_rpc_method_handler( servicer.get_status, - request_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StatusRequest.FromString, - response_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StatusResponse.SerializeToString, + request_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StatusRequest.FromString, + response_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StatusResponse.SerializeToString, ), 'start_session': grpc.unary_unary_rpc_method_handler( servicer.start_session, - request_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StartSessionRequest.FromString, - response_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StartSessionResponse.SerializeToString, + request_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StartSessionRequest.FromString, + response_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StartSessionResponse.SerializeToString, ), 'end_session': grpc.unary_unary_rpc_method_handler( servicer.end_session, - request_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.EndSessionRequest.FromString, - response_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.EndSessionResponse.SerializeToString, + request_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.EndSessionRequest.FromString, + response_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.EndSessionResponse.SerializeToString, ), 'upscale_chunk': grpc.unary_unary_rpc_method_handler( servicer.upscale_chunk, - request_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkRequest.FromString, - response_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkResponse.SerializeToString, + request_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkRequest.FromString, + response_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkResponse.SerializeToString, ), 'upscale_video': grpc.stream_stream_rpc_method_handler( servicer.upscale_video, - request_deserializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkRequest.FromString, - response_serializer=flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkResponse.SerializeToString, + request_deserializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkRequest.FromString, + response_serializer=flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( - 'flashvsr.FlashVSR', rpc_method_handlers) + 'flashdreams.uplift.VideoUplift', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) - server.add_registered_method_handlers('flashvsr.FlashVSR', rpc_method_handlers) + server.add_registered_method_handlers('flashdreams.uplift.VideoUplift', rpc_method_handlers) # This class is part of an EXPERIMENTAL API. -class FlashVSR(object): +class VideoUplift(object): """--------------------------------------------------------------------------- - FlashVSR gRPC Service + FlashDreams video uplift gRPC service - Streaming video super-resolution (2× or 4×) using a diffusion-based model. + Streaming video super-resolution (2× or 4×) using any server-side model + implementation that speaks this protocol. Two usage patterns: Unary chunk-by-chunk flow (best for integration into existing pipelines): @@ -206,9 +209,9 @@ def get_status(request, return grpc.experimental.unary_unary( request, target, - '/flashvsr.FlashVSR/get_status', - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StatusRequest.SerializeToString, - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StatusResponse.FromString, + '/flashdreams.uplift.VideoUplift/get_status', + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StatusRequest.SerializeToString, + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StatusResponse.FromString, options, channel_credentials, insecure, @@ -233,9 +236,9 @@ def start_session(request, return grpc.experimental.unary_unary( request, target, - '/flashvsr.FlashVSR/start_session', - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StartSessionRequest.SerializeToString, - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.StartSessionResponse.FromString, + '/flashdreams.uplift.VideoUplift/start_session', + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StartSessionRequest.SerializeToString, + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.StartSessionResponse.FromString, options, channel_credentials, insecure, @@ -260,9 +263,9 @@ def end_session(request, return grpc.experimental.unary_unary( request, target, - '/flashvsr.FlashVSR/end_session', - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.EndSessionRequest.SerializeToString, - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.EndSessionResponse.FromString, + '/flashdreams.uplift.VideoUplift/end_session', + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.EndSessionRequest.SerializeToString, + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.EndSessionResponse.FromString, options, channel_credentials, insecure, @@ -287,9 +290,9 @@ def upscale_chunk(request, return grpc.experimental.unary_unary( request, target, - '/flashvsr.FlashVSR/upscale_chunk', - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkRequest.SerializeToString, - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkResponse.FromString, + '/flashdreams.uplift.VideoUplift/upscale_chunk', + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkRequest.SerializeToString, + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkResponse.FromString, options, channel_credentials, insecure, @@ -314,9 +317,9 @@ def upscale_video(request_iterator, return grpc.experimental.stream_stream( request_iterator, target, - '/flashvsr.FlashVSR/upscale_video', - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkRequest.SerializeToString, - flashvsr_dot_grpc_dot_protos_dot_flashvsr__pb2.UpscaleChunkResponse.FromString, + '/flashdreams.uplift.VideoUplift/upscale_video', + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkRequest.SerializeToString, + flashdreams_dot_serving_dot_uplift_dot_protos_dot_uplift__pb2.UpscaleChunkResponse.FromString, options, channel_credentials, insecure, diff --git a/integrations/flashvsr/flashvsr/grpc/streaming_view.py b/flashdreams/flashdreams/serving/uplift/streaming_view.py similarity index 92% rename from integrations/flashvsr/flashvsr/grpc/streaming_view.py rename to flashdreams/flashdreams/serving/uplift/streaming_view.py index eec54f48..0b4964ad 100644 --- a/integrations/flashvsr/flashvsr/grpc/streaming_view.py +++ b/flashdreams/flashdreams/serving/uplift/streaming_view.py @@ -13,13 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HTTP MJPEG viewer support for the FlashVSR gRPC server.""" +"""HTTP MJPEG viewer support for video uplift servers.""" import io import queue import threading import time -from dataclasses import dataclass +from dataclasses import dataclass, field from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -31,6 +31,7 @@ DEFAULT_VIEWER_JPEG_QUALITY = 90 DEFAULT_VIEWER_JPEG_BACKEND = "auto" DEFAULT_VIEWER_MAX_FPS = 60.0 +DEFAULT_VIEWER_PLAYBACK_FPS = 30.0 DEFAULT_VIEWER_FRAME_STRIDE = 1 _VIEWER_STOP = object() @@ -40,6 +41,7 @@ @dataclass class _ViewerPlaybackChunk: elapsed_ms: float + ready_at: float = field(default_factory=time.perf_counter) frames: np.ndarray | None = None jpegs: list[bytes] | None = None @@ -154,6 +156,7 @@ def __init__( jpeg_backend: str, chunk_queue_depth: int, max_fps: float, + playback_fps: float, frame_stride: int, ) -> None: self.host = host @@ -162,6 +165,7 @@ def __init__( self.jpeg_backend = jpeg_backend self.chunk_queue_depth = max(1, int(chunk_queue_depth)) self.max_fps = float(max_fps) + self.playback_fps = float(playback_fps) self.frame_stride = max(1, int(frame_stride)) self.original_hub = _MjpegFrameHub() self.upscaled_hub = _MjpegFrameHub() @@ -207,7 +211,7 @@ def _send_index(self) -> None: - FlashVSR Stream + FlashDreams Uplift Stream