From 556092d9bcc893264ba9efd5f8da48b3d1502401 Mon Sep 17 00:00:00 2001 From: nuck Date: Mon, 27 Jan 2025 13:19:33 -0500 Subject: [PATCH] Initial commit --- .gitignore | 175 +++++++ __init__.py | 228 +++++++++ noise_warper.py | 1244 ++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 7 + 4 files changed, 1654 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 noise_warper.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1e6919d --- /dev/null +++ b/.gitignore @@ -0,0 +1,175 @@ + + +### Created by https://www.gitignore.io +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +### Python Patch ### +# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration +poetry.toml + +# ruff +.ruff_cache/ + +# LSP config files +pyrightconfig.json + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..1d26226 --- /dev/null +++ b/__init__.py @@ -0,0 +1,228 @@ +import math +import torch +import numpy as np +import cv2 +from .noise_warper import NoiseWarper +import comfy.samplers +import comfy.sample +import latent_preview +import comfy.utils + +############################################################################### +# bgr_to_flow function from your code snippet +############################################################################### +def bgr_to_flow(bgr_image, max_flow=20.0): + """ + Convert a BGR image back to a flow map (H,W,2). + """ + hsv = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2HSV) + angle = hsv[..., 0] * np.pi / 180.0 * 2.0 + magnitude = hsv[..., 2] / 255.0 + flow_map = np.zeros((hsv.shape[0], hsv.shape[1], 2), dtype=np.float32) + flow_map[..., 0], flow_map[..., 1] = cv2.polarToCart(magnitude, angle) + # scale + flow_map *= max_flow + return flow_map + +######################################## +# GLOBAL dict of warpers +######################################## +WARPERS = {} + +######################################## +# ComfyUI Node +######################################## +class NoiseWarperNode: + """ + Node that: + - Accepts an RGB flow image + - Decodes it to (dx, dy) + - Persists a NoiseWarper + - Outputs a LATENT [4, H//8, W//8] by default + - If width/height=0, we use the flow image's size + - Automatically picks a random seed each time we re-init + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "flow_image": ("IMAGE",), + "width": ("INT", {"default":0, "min":0, "max":4096, "step":1}), + "height": ("INT", {"default":0, "min":0, "max":4096, "step":1}), + # "reset": ("BOOL", {"default":False}), + "noise_channels": ("INT", {"default":4, "min":1, "max":8}), + "scale_factor": ("INT", {"default":1, "min":1, "max":8}), + "noise_scale": ("FLOAT", {"default":1.0, "min":0.0, "max":10.0, "step":0.01, "tooltip": "The amount of noise added to the latent image."}), + }, + } + + RETURN_TYPES = ("LATENT","IMAGE") + FUNCTION = "warp_noise" + CATEGORY = "latent" + + @classmethod + def IS_CHANGED(flow_image, width, height, noise_channels, scale_factor, noise_scale): + return float("NaN") + + def warp_noise(self, flow_image, width, height, noise_channels, scale_factor, noise_scale): + # Ensure we only have a single flow image (squeeze batch dimension if needed) + if len(flow_image.shape) == 4: # [B,H,W,C] + assert flow_image.shape[0] == 1, "Only single flow image input is supported" + flow_image = flow_image.squeeze(0) # Remove batch dimension -> [H,W,C] + else: + assert len(flow_image.shape) == 3, "Flow image must have 3 or 4 dimensions" + + node_id = id(self) + flow_h, flow_w, c = flow_image.shape + print(c, flow_h, flow_w) + assert c == 3, "Flow image must have 3 channels (RGB)" + + # 1) If width=0 or height=0, use the flow image's shape + if width == 0: + width = flow_w // 8 + if height == 0: + height = flow_h // 8 + + # 2) Convert ComfyUI tensor -> BGR [uint8, 0..255] + flow_bgr = (flow_image.cpu().numpy() * 255).astype(np.uint8) + # 3) bgr_to_flow => shape [H, W, 2] + flow_map = bgr_to_flow(flow_bgr, max_flow=20.0) + dx_np = flow_map[..., 0] # shape [H, W] + dy_np = flow_map[..., 1] + + # to torch + device = "cuda" if torch.cuda.is_available() else "cpu" + dx = torch.from_numpy(dx_np).float().to(device) + dy = torch.from_numpy(dy_np).float().to(device) + + # 4) Check if we have a warper or need to re-init + warper = WARPERS.get(node_id, None) + do_reset = (warper is None) \ + or (warper.h != height) or (warper.w != width) \ + or (warper.c != noise_channels) + + if do_reset: + # Use a random seed each time we re-init + random_seed = np.random.randint(0, (1 << 31) - 1) + np.random.seed(random_seed & 0xFFFFFFFF) + torch.manual_seed(random_seed) + + warper = NoiseWarper( + c=noise_channels, + h=height, + w=width, + device=device, + scale_factor=scale_factor + ) + WARPERS[node_id] = warper + + + # 5) Warp + random_seed = np.random.randint(0, (1 << 31) - 1) + np.random.seed(random_seed & 0xFFFFFFFF) + torch.manual_seed(random_seed) + warper(dx, dy) + + # 6) Retrieve noise => [noise_channels, H, W] + noise = warper.noise + + # 7) Convert to a 4-channel latent, shape [4, H//8, W//8] + # If noise_channels <4 => pad. If >4 => slice. + latent_channels = 4 + if noise_channels < latent_channels: + pad_c = latent_channels - noise_channels + pad_noise = torch.randn((pad_c, height, width), device=device) + final_noise = torch.cat([noise, pad_noise], dim=0) + else: + final_noise = noise[:latent_channels] + + # downscale to [4, H//8, W//8] + latent_h, latent_w = height, width + latent = final_noise * noise_scale + # latent = torch.nn.functional.interpolate( + # final_noise.unsqueeze(0), + # size=(latent_h, latent_w), + # mode="linear" + # )[0] + + print("RETURNING LATENT NOISE: ", latent.shape) + # Wrap the latent in a dictionary with "samples" key to match SDXL format + # Normalize noise from [-4,4] range to [0,1] range + normalized_noise = (noise + abs(noise.min())) / (noise.max() - noise.min()) + return ({"samples": latent.unsqueeze(0)}, normalized_noise.permute(1, 2, 0).unsqueeze(0),) + +class KSamplerNoiseless: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}), + "disable_noise": ("BOOLEAN", {"default": True, "tooltip": "If True, the noise will not be added to the latent image."}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}), + "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}), + "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}), + "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}), + } + } + + RETURN_TYPES = ("LATENT",) + OUTPUT_TOOLTIPS = ("The denoised latent.",) + FUNCTION = "sample" + + CATEGORY = "sampling" + DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image." + + def sample(self, model, disable_noise, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0): + return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise) + +def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): + latent_image = latent["samples"] + latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image) + + generator = torch.manual_seed(seed) + baseline_noise = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu") + + # Scale latent_image to match baseline_noise range + baseline_min = baseline_noise.min() + baseline_max = baseline_noise.max() + latent_min = latent_image.min() + latent_max = latent_image.max() + + # Apply min-max normalization to match baseline noise range + # latent_image = (latent_image - latent_min) / (latent_max - latent_min) + # latent_image = latent_image * (baseline_max - baseline_min) + baseline_min + + # if not disable_noise: + # batch_inds = latent["batch_index"] if "batch_index" in latent else None + # latent_image = comfy.sample.prepare_noise(latent_image, seed, batch_inds) + + # noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + # print("Noise range: ", latent_image.min(), latent_image.max()) + + + noise_mask = None + if "noise_mask" in latent: + noise_mask = latent["noise_mask"] + + callback = latent_preview.prepare_callback(model, steps) + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + samples = comfy.sample.sample(model, latent_image, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + denoise=denoise, disable_noise=True, start_step=start_step, last_step=last_step, + force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed) + out = latent.copy() + out["samples"] = samples + return (out, ) + + + +# Tell ComfyUI about our node +NODE_CLASS_MAPPINGS = { + "NoiseWarperNode": NoiseWarperNode, + "KSamplerNoiseless": KSamplerNoiseless +} \ No newline at end of file diff --git a/noise_warper.py b/noise_warper.py new file mode 100644 index 0000000..1bc1eab --- /dev/null +++ b/noise_warper.py @@ -0,0 +1,1244 @@ +import glob +import sys + +from einops import rearrange +import numpy as np +import rp +import torch +from tqdm import tqdm + + +sys.path.append(rp.get_path_parent(__file__)) + +def unique_pixels(image): + """ + Find unique pixel values in an image tensor and return their RGB values, counts, and inverse indices. + + Args: + image (torch.Tensor): Image tensor of shape [c, h, w], where c is the number of channels (e.g., 3 for RGB), + h is the height, and w is the width of the image. + + Returns: + tuple: A tuple containing three tensors: + - unique_colors (torch.Tensor): Tensor of shape [u, c] representing the unique RGB values found in the image, + where u is the number of unique colors. + - counts (torch.Tensor): Tensor of shape [u] representing the counts of each unique color. + - index_matrix (torch.Tensor): Tensor of shape [h, w] representing the inverse indices of each pixel, + mapping each pixel to its corresponding unique color index. + """ + c, h, w = image.shape + + # Rearrange the image tensor from [c, h, w] to [h, w, c] using einops + pixels = rearrange(image, "c h w -> h w c") + + # Flatten the image tensor to [h*w, c] + flattened_pixels = rearrange(pixels, "h w c -> (h w) c") + + # Find unique RGB values, counts, and inverse indices + unique_colors, inverse_indices, counts = torch.unique(flattened_pixels, dim=0, return_inverse=True, return_counts=True, sorted=False) + # unique_colors, inverse_indices, counts = torch.unique_consecutive(flattened_pixels, dim=0, return_inverse=True, return_counts=True) + + # Get the number of unique indices + u = unique_colors.shape[0] + + # Reshape the inverse indices back to the original image dimensions [h, w] using einops + index_matrix = rearrange(inverse_indices, "(h w) -> h w", h=h, w=w) + + # Assert the shapes of the output tensors + assert unique_colors.shape == (u, c) + assert counts.shape == (u,) + assert index_matrix.shape == (h, w) + assert index_matrix.min() == 0 + assert index_matrix.max() == u - 1 + + return unique_colors, counts, index_matrix + + +def sum_indexed_values(image, index_matrix): + """ + Sum the values in the CHW image tensor based on the indices specified in the HW index matrix. + + Args: + image (torch.Tensor): Image tensor of shape [C, H, W], where C is the number of channels, + H is the height, and W is the width of the image. + index_matrix (torch.Tensor): Index matrix tensor of shape [H, W] containing indices + specifying the mapping of each pixel to its corresponding + unique value. + Indices range [0, U), where U is the number of unique indices + + Returns: + torch.Tensor: Tensor of shape [U, C] representing the sum of values in the image tensor + based on the indices in the index matrix, where U is the number of unique + indices in the index matrix. + """ + c, h, w = image.shape + u = index_matrix.max() + 1 + + # Rearrange the image tensor from [c, h, w] to [h, w, c] using einops + pixels = rearrange(image, "c h w -> h w c") + + # Flatten the image tensor to [h*w, c] + flattened_pixels = rearrange(pixels, "h w c -> (h w) c") + + # Create an output tensor of shape [u, c] initialized with zeros + output = torch.zeros((u, c), dtype=flattened_pixels.dtype, device=flattened_pixels.device) + + # Scatter sum the flattened pixel values using the index matrix + output.index_add_(0, index_matrix.view(-1), flattened_pixels) + + # Assert the shapes of the input and output tensors + assert image.shape == (c, h, w), f"Expected image shape: ({c}, {h}, {w}), but got: {image.shape}" + assert index_matrix.shape == (h, w), f"Expected index_matrix shape: ({h}, {w}), but got: {index_matrix.shape}" + assert output.shape == (u, c), f"Expected output shape: ({u}, {c}), but got: {output.shape}" + + return output + +def indexed_to_image(index_matrix, unique_colors): + """ + Create a CHW image tensor from an HW index matrix and a UC unique_colors matrix. + + Args: + index_matrix (torch.Tensor): Index matrix tensor of shape [H, W] containing indices + specifying the mapping of each pixel to its corresponding + unique color. + unique_colors (torch.Tensor): Unique colors matrix tensor of shape [U, C] containing + the unique color values, where U is the number of unique + colors and C is the number of channels. + + Returns: + torch.Tensor: Image tensor of shape [C, H, W] representing the reconstructed image + based on the index matrix and unique colors matrix. + """ + h, w = index_matrix.shape + u, c = unique_colors.shape + + # Assert the shapes of the input tensors + assert index_matrix.max() < u, f"Index matrix contains indices ({index_matrix.max()}) greater than the number of unique colors ({u})" + + # Gather the colors based on the index matrix + flattened_image = unique_colors[index_matrix.view(-1)] + + # Reshape the flattened image to [h, w, c] + image = rearrange(flattened_image, "(h w) c -> h w c", h=h, w=w) + + # Rearrange the image tensor from [h, w, c] to [c, h, w] using einops + image = rearrange(image, "h w c -> c h w") + + # Assert the shape of the output tensor + assert image.shape == (c, h, w), f"Expected image shape: ({c}, {h}, {w}), but got: {image.shape}" + + return image + + +def demo_pixellation_via_proxy(): + real_image = rp.as_torch_image( + rp.cv_resize_image( + rp.load_image("https://i.natgeofe.com/n/4f5aaece-3300-41a4-b2a8-ed2708a0a27c/domestic-dog_thumb_square.jpg"), + (512, 512), + ) + ) + + c, h, w = real_image.shape + + noise_image = torch.randn(c, h // 4, w // 4) + + # Resize noise_image using nearest-neighbor interpolation to match the dimensions of real_image + pixelated_noise_image = rp.torch_resize_image(noise_image, 4, "nearest") + assert pixelated_noise_image.shape==(c,h,w) + + # Find unique pixel values, their indices, and counts in the pixelated noise image + unique_colors, counts, index_matrix = unique_pixels(pixelated_noise_image) + + # Sum the color values from real_image based on the indices of the unique noise pixels + summed_colors = sum_indexed_values(real_image, index_matrix) + + # Divide the summed color values by the counts to get the average color for each unique pixel + average_colors = summed_colors / rearrange(counts, "u -> u 1") + + # Create a new pixelated image using the average colors and the index matrix + pixelated_dog_image = indexed_to_image(index_matrix, average_colors) + + rp.display_image(pixelated_dog_image) + +def calculate_wave_pattern(h, w, frame): + # Create a grid of coordinates + y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) + + # Calculate the distance from the center of the image + center_x, center_y = w // 2, h // 2 + dist_from_center = torch.sqrt((x - center_x)**2 + (y - center_y)**2) + + # Calculate the angle from the center of the image + angle_from_center = torch.atan2(y - center_y, x - center_x) + + # Calculate the wave pattern based on the distance and angle + wave_freq = 0.05 # Frequency of the waves + wave_amp = 10.0 # Amplitude of the waves + wave_offset = frame * 0.05 # Offset for animation + + dx = wave_amp * torch.cos(dist_from_center * wave_freq + angle_from_center + wave_offset) + dy = wave_amp * torch.sin(dist_from_center * wave_freq + angle_from_center + wave_offset) + + return dx, dy + +def starfield_zoom(h, w, frame): + # Create a grid of coordinates + y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) + + # Calculate the distance from the center of the image + center_x, center_y = w // 2, h // 2 + dist_from_center = torch.sqrt((x - center_x)**2 + (y - center_y)**2) + + # Calculate the angle from the center of the image + angle_from_center = torch.atan2(y - center_y, x - center_x) + + # Calculate the starfield zoom effect + zoom_speed = 0.01 # Speed of the zoom effect + zoom_scale = 1.0 + frame * zoom_speed # Scale factor for the zoom effect + + # Calculate the displacement based on the distance and angle + dx = dist_from_center * torch.cos(angle_from_center) / zoom_scale + dy = dist_from_center * torch.sin(angle_from_center) / zoom_scale + + return dx, dy + +_arange_cache={} +def _cached_arange(length, device, dtype): + code=hash((length,device,dtype)) + if code in _arange_cache: + return _arange_cache[code] + + + _arange_cache[code]= torch.arange(length , device=device, dtype=dtype) + return _arange_cache[code] + +def fast_nearest_torch_remap_image(image, x, y, *, relative=False, add_alpha_mask=False, use_cached_meshgrid=False): + # assert rp.r.is_torch_image(image), "image must be a torch tensor with shape [C, H, W]" + # assert is_torch_tensor(x) and is_a_matrix(x), "x must be a torch tensor with shape [H_out, W_out]" + # assert is_torch_tensor(y) and is_a_matrix(y), "y must be a torch tensor with shape [H_out, W_out]" + # assert x.shape == y.shape, "x and y must have the same shape, but got x.shape={} and y.shape={}".format(x.shape, y.shape) + # assert image.device==x.device==y.device, "all inputs must be on the same device" + + # pip_import('torch') + + import torch + + in_c, in_height, in_width = image.shape + out_height, out_width = x.shape + + if add_alpha_mask: + alpha_mask = torch.ones_like(image[:1]) + image = torch.cat([image, alpha_mask], dim=0) + + if torch.is_floating_point(x): x = x.round_().long() + if torch.is_floating_point(y): y = y.round_().long() + + if relative: + # assert in_height == out_height, "For relative warping, input and output heights must match, but got in_height={} and out_height={}".format(in_height, out_height) + # assert in_width == out_width , "For relative warping, input and output widths must match, but got in_width={} and out_width={}".format(in_width, out_width) + x += _cached_arange(in_width , device=x.device, dtype=x.dtype) + y += _cached_arange(in_height, device=y.device, dtype=y.dtype)[:,None] + + x.clamp_(0, in_width - 1) + y.clamp_(0,in_height-1) + out = image[:, y, x] + + expected_c = in_c+1 if add_alpha_mask else in_c + assert out.shape == (expected_c, out_height, out_width), "Expected output shape: ({}, {}, {}), but got: {}".format(expected_c, out_height, out_width, out.shape) + + return out + + +def warp_noise(noise, dx, dy, s=1): + #This is *certainly* imperfect. We need to have particle swarm in addition to this. + + dx=dx.round_().int() + dy=dy.round_().int() + + c, h, w = noise.shape + assert dx.shape==(h,w) + assert dy.shape==(h,w) + + #s is scaling factor + hs = h * s + ws = w * s + + #Upscale the warping with linear interpolation. Also scale it appropriately. + if s!=1: + up_dx = rp.torch_resize_image(dx[None], (hs, ws), interp="bilinear")[0] + up_dy = rp.torch_resize_image(dy[None], (hs, ws), interp="bilinear")[0] + up_dx *= s + up_dy *= s + + up_noise = rp.torch_resize_image(noise, (hs, ws), interp="nearest") + else: + up_dx = dx + up_dy = dy + up_noise = noise + assert up_noise.shape == (c, hs, ws) + + # Warp the noise - and put 0 where it lands out-of-bounds + # up_noise = rp.torch_remap_image(up_noise, up_dx, up_dy, relative=True, interp="nearest") + up_noise = fast_nearest_torch_remap_image(up_noise, up_dx, up_dy, relative=True) + assert up_noise.shape == (c, hs, ws) + + # Regaussianize the noise + output, _ = regaussianize(up_noise) + + #Now we resample the noise back down again + if s!=1: + output = rp.torch_resize_image(output, (h, w), interp='area') + output = output * s #Adjust variance by multiplying by sqrt of area, aka sqrt(s*s)=s + + return output + + +def regaussianize(noise): + c, hs, ws = noise.shape + + # Find unique pixel values, their indices, and counts in the pixelated noise image + unique_colors, counts, index_matrix = unique_pixels(noise[:1]) + u = len(unique_colors) + assert unique_colors.shape == (u, 1) + assert counts.shape == (u,) + assert index_matrix.max() == u - 1 + assert index_matrix.min() == 0 + assert index_matrix.shape == (hs, ws) + + foreign_noise = torch.randn_like(noise) + assert foreign_noise.shape == noise.shape == (c, hs, ws) + + summed_foreign_noise_colors = sum_indexed_values(foreign_noise, index_matrix) + assert summed_foreign_noise_colors.shape == (u, c) + + meaned_foreign_noise_colors = summed_foreign_noise_colors / rearrange(counts, "u -> u 1") + assert meaned_foreign_noise_colors.shape == (u, c) + + meaned_foreign_noise = indexed_to_image(index_matrix, meaned_foreign_noise_colors) + assert meaned_foreign_noise.shape == (c, hs, ws) + + zeroed_foreign_noise = foreign_noise - meaned_foreign_noise + assert zeroed_foreign_noise.shape == (c, hs, ws) + + counts_as_colors = rearrange(counts, "u -> u 1") + counts_image = indexed_to_image(index_matrix, counts_as_colors) + assert counts_image.shape == (1, hs, ws) + + #To upsample noise, we must first divide by the area then add zero-sum-noise + output = noise + output = output / counts_image ** .5 + output = output + zeroed_foreign_noise + + assert output.shape == noise.shape == (c, hs, ws) + + return output, counts_image + +def demo_noise_warp(Q=-6,scale_factor=1,num_frames=300): + #Run this in a Jupyter notebook and watch the noise go brrrrrrr + + + d=rp.JupyterDisplayChannel() + d.display() + device='cuda' + h=w=128 + + # rp.cv_imshow(rp.apply_colormap_to_image(output[ω]/output[ω].mean()/4),label='weight') + + warper = NoiseWarper(3,h,w,device=device,scale_factor=scale_factor,) + + #Add some ink + warper._state[4:, int(h/256*10 * scale_factor) : int(h/256*80 * scale_factor), int(h/256*10 * scale_factor) : int(h/256*80 * scale_factor)] = 1 # Make a little black square + warper._state[ + 4:, + h * scale_factor // 2 - 5 * scale_factor : h * scale_factor // 2 + 5 * scale_factor, + w * scale_factor // 2 - 5 * scale_factor : w * scale_factor // 2 + 5 * scale_factor, + ] = 10 # Make a little black square + + noise=torch.randn(3,h,w).to(device) + wdx,wdy=calculate_wave_pattern(h,w,frame=0) + sdx,sdy=starfield_zoom(h,w,frame=1) + + dx=sdx+2*wdx + dy=sdy+2*wdy + + dx/=dx.max() + dy/=dy.max() + # Q=-6 + dy*=Q + dx*=Q + dx=dx.to(device) + dy=dy.to(device) + new_noise=noise + + frames=[] + try: + for _ in range(num_frames): + new_noise=warper(dx,dy).noise + weights = warper._state[2] + weights = rp.torch_resize_image(weights[None],(h,w))[0] + + # rp.display_image(new_noise) + frame=rp.tiled_images( + [ + rp.as_numpy_image(new_noise/4+.5), + rp.apply_colormap_to_image(weights/weights.mean()/4), + ], + border_thickness=0, + ) + frames.append(rp.labeled_image(frame,'Frame %i'%_)) + d.update( + frame + ) + except KeyboardInterrupt: + print("Interrupted demo at frame",_) + return frames + +def demo_webcam_noise_warp(): + import cv2 + from rp import ( + as_numpy_image, + cv_bgr_rgb_swap, + display_image, + get_image_dimensions, + rp, + tiled_images, + ) + + def resize_frame(frame, target_height=64): + aspect_ratio = frame.shape[1] / frame.shape[0] + target_width = int(target_height * aspect_ratio) + resized_frame = cv2.resize(frame, (target_width, target_height)) + # print(resized_frame.shape) + return resized_frame + + def main(): + cap = cv2.VideoCapture(0) + ret, prev_frame = cap.read() + prev_frame = resize_frame(prev_frame) + prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY) + + # Initialize DeepFlow Optical Flow + optical_flow = cv2.optflow.createOptFlow_DeepFlow() + + d = rp.JupyterDisplayChannel() + d.display() + device = "cpu" + h, w = get_image_dimensions(prev_frame) + wdx, wdy = calculate_wave_pattern(h, w, frame=0) + sdx, sdy = starfield_zoom(h, w, frame=1) + + dx = sdx + 2 * wdx + dy = sdy + 2 * wdy + + dx /= dx.max() + dy /= dy.max() + Q = -6 + dy *= Q + dx *= Q + dx = dx.to(device) + dy = dy.to(device) + # new_xyωc = noise_to_xyωc(noise) + + warper=NoiseWarper(3,h,w,device=device,scale_factor=2) + + + while True: + ret, frame = cap.read() + frame = resize_frame(frame) + frame=rp.horizontally_flipped_image(frame) + frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + + + # Compute the optical flow + flow = optical_flow.calc(prev_gray, frame_gray, None) + + x = flow[:, :, 0] + y = flow[:, :, 1] + + dx = torch.Tensor(x) + dy = torch.Tensor(y) + + prev_gray = frame_gray.copy() + + if cv2.waitKey(1) & 0xFF == ord("q"): + break + + new_noise = warper(dx, dy).noise + weights = warper._state[2] + weights = rp.torch_resize_image(weights[None],(h,w))[0] + + display_image( + tiled_images( + [ + as_numpy_image(new_noise / 5 + 0.5), + cv_bgr_rgb_swap(frame), + rp.apply_colormap_to_image(weights / weights.mean() / 4), + ] + ) + ) + + # d.update(rp.as_numpy_image(new_noise/4+.5)) + + cap.release() + cv2.destroyAllWindows() + + main() + + +@rp.memoized +def _xy_meshgrid(h,w,device,dtype): + y, x = torch.meshgrid( + torch.arange(h), + torch.arange(w), + ) + + output = torch.stack( + [x, y], + ).to(device, dtype) + + assert output.shape == (2, h, w) + return output + +def xy_meshgrid_like_image(image): + """ + Example: + >>> image=load_image('https://picsum.photos/id/28/367/267') + ... image=as_torch_image(image) + ... xy=xy_meshgrid_like_image(image) + ... display_image(full_range(as_numpy_array(xy[0]))) + ... display_image(full_range(as_numpy_array(xy[1]))) + """ + assert image.ndim == 3, "image is in CHW form" + c, h, w = image.shape + return _xy_meshgrid(h,w,image.device,image.dtype) + +def noise_to_xyωc(noise): + assert noise.ndim == 3, "noise is in CHW form" + zeros=torch.zeros_like(noise[0][None]) + ones =torch.ones_like (noise[0][None]) + + #Prepend [dx=0, dy=0, weights=1] channels + output=torch.concat([zeros, zeros, ones, noise]) + return output + +def xyωc_to_noise(xyωc): + assert xyωc.ndim == 3, "xyωc is in [ω x y c]·h·w form" + assert xyωc.shape[0]>3, 'xyωc should have at least one noise channel' + noise=xyωc[3:] + return noise + +def warp_xyωc( + I, + F, + xy_mode="none", + # USED FOR ABLATIONS: + expand_only=False, +): + """ + For ablations, set: + - expand_only=True #No contraction + - expand_only='bilinear' #Bilinear Interpolation + - expand_only='nearest' #Nearest Neighbors Warping + """ + #Input assertions + assert F.device==I.device + assert F.ndim==3, str(F.shape)+' F stands for flow, and its in [x y]·h·w form' + assert I.ndim==3, str(I.shape)+' I stands for input, in [ω x y c]·h·w form where ω=weights, x and y are offsets, and c is num noise channels' + xyωc, h, w = I.shape + assert F.shape==(2,h,w) # Should be [x y]·h·w + device=I.device + + #How I'm going to address the different channels: + x = 0 # // index of Δx channel + y = 1 # // index of Δy channel + xy = 2 # I[:xy] + xyω = 3 # I[:xyω] + ω = 2 # I[ω] // index of weight channel + c = xyωc-xyω # I[-c:] // num noise channels + ωc = xyωc-xy # I[-ωc:] + # h_dim = 1 + w_dim = 2 + assert c, 'I has no noise channels. There is nothing to warp.' + assert (I[ω]>0).all(), 'All weights should be greater than 0' + + #Compute the grid of xy indices + grid = xy_meshgrid_like_image(I) + assert grid.shape==(2,h,w) # Shape is [x y]·h·w + + #The default values we initialize to. Todo: cache this. + init = torch.empty_like(I) + init[:xy]=0 + init[ω]=1 + init[-c:]=0 + + #Caluclate initial pre-expand + pre_expand = torch.empty_like(I) + + #The original plan was to use init xy during expand, because the query position is arbitrary.... + #It doesn't actually make deep sense to copy the offsets during this step, but it doesn't seem to hurt either... + #BUT I think I got slightly better results...?...so I'm going to do it anyway. + # pre_expand[:xy] = init[:xy] # <---- Original algorithm I wrote on paper + + #ABLATION STUFF IN THIS PARAGRAPH + #Using F_index instead of F so we can use ablations like bilinear, bicubic etc + interp = 'nearest' if not isinstance(expand_only, str) else expand_only + regauss = not isinstance(expand_only, str) + F_index = F + if interp=='nearest': + #Default behaviour, ablations or not + F_index=F_index.round() + + pre_expand[:xy] = rp.torch_remap_image(I[:xy], * -F, relative=True, interp=interp)# <---- Last minute change + pre_expand[-ωc:] = rp.torch_remap_image(I[-ωc:], * -F, relative=True, interp=interp) + pre_expand[ω][pre_expand[ω]==0]=1 #Give new noise regions a weight of 1 - effectively setting it to init there + + if expand_only: + if regauss: + #This is an ablation option - simple warp + regaussianize + #Enable to preview expansion-only noise warping + #The default behaviour! My algo! + pre_expand[-c:]=regaussianize(pre_expand[-c:])[0] + else: + #Turn zeroes to noise + pre_expand[-c:]=torch.randn_like(pre_expand[-c:]) * (pre_expand[-c:]==0) + pre_expand[-c:] + return pre_expand + + #Calculate initial pre-shrink + pre_shrink = I.clone() + pre_shrink[:xy] += F + + #Pre-Shrink mask - discard out-of-bounds pixels + pos = (grid + pre_shrink[:xy]).round() + in_bounds = (0<= pos[x]) & (pos[x] < w) & (0<= pos[y]) & (pos[y] < h) + in_bounds = in_bounds[None] #Match the shape of the input + out_of_bounds = ~in_bounds + assert out_of_bounds.dtype==torch.bool + assert out_of_bounds.shape==(1,h,w) + assert pre_shrink.shape == init.shape + pre_shrink = torch.where(out_of_bounds, init, pre_shrink) + + #Deal with shrink positions offsets + scat_xy = pre_shrink[:xy].round() + pre_shrink[:xy] -= scat_xy + + #FLOATING POINT POSITIONS: I will disable this for now. It does in fact increase sensitivity! But it also makes it less long-term coherent + assert xy_mode in ['float', 'none'] or isinstance(xy_mode, int) + if xy_mode=='none': + pre_shrink[:xy] = 0 #DEBUG: Uncomment to ablate floating-point swarm positions + + if isinstance(xy_mode, int): + # XY quantization: best to use odd numbers! + quant = xy_mode + pre_shrink[:xy] = ( + pre_shrink[:xy] * quant + ).round() / quant + + #OTHER ways I tried reducing sensitivity to motion. They work - but 0 is best. Let's just use high resolution. + # pre_shrink[:xy][pre_shrink[:xy].abs()<.1] = 0 #DEBUG: Uncomment to ablate floating-point swarm positions + # pre_shrink[:xy] *= -1 #I can't even tell that this is wrong..... + # pre_shrink[:xy] *= .9 + # sensitivity_factor = 4 + + scat = lambda tensor: rp.torch_scatter_add_image(tensor, *scat_xy, relative=True) + + #Where mask==True, we output shrink. Where mask==0, we output expand. + shrink_mask = torch.ones(1,h,w,dtype=bool,device=device) #The purpose is to get zeroes where no element is used + shrink_mask = scat(shrink_mask) + assert shrink_mask.dtype==torch.bool, 'If this fails we gotta convert it with mask.=astype(bool)' + + # rp.cv_imshow(rp.tiled_images([out_of_bounds[0],shrink_mask[0]]),label='OOB') ; return I #DEBUG - uncomment to see the masks + + #Remove the expansion points where we'll use shrink + pre_expand = torch.where(shrink_mask, init, pre_expand) + # rp.cv_imshow(pre_expand[-c:]/5+.5,'preex') + + #Horizontally Concat + concat_dim = w_dim + concat = torch.concat([pre_shrink, pre_expand], dim=concat_dim) + + #Regaussianize + concat[-c:], counts_image = regaussianize(concat[-c:]) + assert counts_image.shape == (1, h, 2*w) + # rp.cv_imshow(concat[-c:]/5+.5,label='regauss') ; return pre_expand #DEBUG - Uncomment to preview regaussianization + + #Distribute Weights + concat[ω] /= counts_image[0] + concat[ω] = concat[ω].nan_to_num() #We shouldn't need this, this is a crutch. Final mask should take care of this. + + pre_shrink, expand = torch.chunk(concat, chunks=2, dim=concat_dim) + assert pre_shrink.shape == expand.shape == (3+c, h, w) + + shrink = torch.empty_like(pre_shrink) + shrink[ω] = scat(pre_shrink[ω][None])[0] + shrink[:xy] = scat(pre_shrink[:xy]*pre_shrink[ω][None]) / shrink[ω][None] + shrink[-c:] = scat(pre_shrink[-c:]*pre_shrink[ω][None]) / scat(pre_shrink[ω][None]**2).sqrt() + + output = torch.where(shrink_mask, shrink, expand) + output[ω] = output[ω] / output[ω].mean() #Don't let them get too big or too small + ε = .00001 + output[ω] += ε #Don't let it go too low + + # rp.debug_comment([output[ω].min(),output[ω].max()])# --> [tensor(0.0010), tensor(2.7004)] + # rp.debug_comment([shrink[ω].min(),shrink[ω].max()])# --> [tensor(0.), tensor(2.7004)] + # rp.debug_comment([expand[ω].min(),expand[ω].max()])# --> [tensor(0.0001), tensor(0.3892)] + # rp.cv_imshow(rp.apply_colormap_to_image(output[ω]/output[ω].mean()/4),label='weight') + # rp.cv_imshow(rp.apply_colormap_to_image(output[ω]/10),label='weight') + assert (output[ω]>0).all() + # print(end='\r%.08f %.08f'%(float(output[ω].min()), float(output[ω].max()))) + + output[ω] **= .9999 #Make it tend towards 1 + + + return output + + +class NoiseWarper: + def __init__( + self, + c, h, w, + device, + dtype=torch.float32, + scale_factor=1, + post_noise_alpha = 0, + progressive_noise_alpha = 0, + warp_kwargs=dict(), + ): + + #Some non-exhaustive input assertions + assert isinstance(c,int) and c>0 + assert isinstance(h,int) and h>0 + assert isinstance(w,int) and w>0 + assert isinstance(scale_factor,int) and w>=1 + + #Record arguments + self.c=c + self.h=h + self.w=w + self.device=device + self.dtype=dtype + self.scale_factor=scale_factor + self.progressive_noise_alpha=progressive_noise_alpha + self.post_noise_alpha=post_noise_alpha + self.warp_kwargs=warp_kwargs + + #Initialize the state + self._state = self._noise_to_state( + noise=torch.randn( + c, + h * scale_factor, + w * scale_factor, + dtype=dtype, + device=device, + ) + ) + + @property + def noise(self): + #TODO: The noise should be downsampled to respect the weights!! + noise = self._state_to_noise(self._state) + weights = self._state[2][None] #xyωc + noise = ( + rp.torch_resize_image(noise * weights, (self.h, self.w), interp="area") + / rp.torch_resize_image(weights**2 , (self.h, self.w), interp="area").sqrt() + ) + noise = noise * self.scale_factor + + if self.post_noise_alpha: + noise = mix_new_noise(noise, self.post_noise_alpha) + + return noise + + def __call__(self, dx, dy): + + if rp.is_numpy_array(dx): dx = torch.tensor(dx).to(self.device, self.dtype) + if rp.is_numpy_array(dy): dy = torch.tensor(dy).to(self.device, self.dtype) + + flow = torch.stack([dx, dy]).to(self.device, self.dtype) + _, oflowh, ofloww = flow.shape #Original height and width of the flow + + assert flow.ndim == 3 and flow.shape[0] == 2, "Flow is in [x y]·h·w form" + flow = rp.torch_resize_image( + flow, + ( + self.h * self.scale_factor, + self.w * self.scale_factor, + ), + ) + + _, flowh, floww = flow.shape + + #Multiply the flow values by the size change + flow[0] *= flowh / oflowh * self.scale_factor + flow[1] *= floww / ofloww * self.scale_factor + + self._state = self._warp_state(self._state, flow) + return self + + #The following three methods can be overridden in subclasses: + + @staticmethod + def _noise_to_state(noise): + return noise_to_xyωc(noise) + + @staticmethod + def _state_to_noise(state): + return xyωc_to_noise(state) + + def _warp_state(self, state, flow): + + if self.progressive_noise_alpha: + state[3:] = mix_new_noise(state[3:], self.progressive_noise_alpha) + + return warp_xyωc(state, flow, **self.warp_kwargs) + +def blend_noise(noise_background, noise_foreground, alpha): + """ Variance-preserving blend """ + return (noise_foreground * alpha + noise_background * (1-alpha))/(alpha ** 2 + (1-alpha) ** 2)**.5 + +def mix_new_noise(noise, alpha): + """As alpha --> 1, noise is destroyed""" + if isinstance(noise, torch.Tensor): return blend_noise(noise, torch.randn_like(noise) , alpha) + elif isinstance(noise, np.ndarray): return blend_noise(noise, np.random.randn(*noise.shape), alpha) + else: raise TypeError(f"Unsupported input type: {type(noise)}. Expected PyTorch Tensor or NumPy array.") + +def resize_noise(noise, size, alpha=None): + """ + Can resize gaussian noise, adjusting for variance and preventing cross-correlation + """ + assert noise.ndim == 3, "resize_noise: noise should be a CHW tensor" + num_channels, old_height, old_width = noise.shape + + if noise.ndim==4: + #If given a batch of noises, do it for each one + return torch.stack([resize_noise(x, new_height, new_width) for x in noise]) + + if rp.is_number(size): + new_height, new_width = int(old_height * size), int(old_width * size) + else: + new_height, new_width = size + + assert new_height<=old_height, 'resize_noise: Only useful for shrinking noise, not growing it' + assert new_width <=old_width , 'resize_noise: Only useful for shrinking noise, not growing it' + + x, y = rp.xy_torch_matrices( + old_height, + old_width, + max_x=new_width, + max_y=new_height, + ) + + if alpha is not None: + #Prepend the alpha + assert alpha.ndim==2,alpha.shape + assert alpha.shape==noise.shape[1:],(alpha.shape,noise.shape) + noise=torch.cat((alpha[None],noise)) + + resized = rp.torch_scatter_add_image( + noise, + x, + y, + height=new_height, + width=new_width, + interp='floor', + prepend_ones=alpha is None + ) + + total, resized = resized[:1], resized[1:] + + adjusted = resized / total**.5 + + return adjusted + +def get_noise_from_video( + video_path: str, + noise_channels: int = 3, + output_folder: str = None, + visualize: bool = True, + resize_frames: tuple = None, + resize_flow: int = 1, + downscale_factor: int = 1, + device=None, + video_preprocessor = None, + save_files=True, + progressive_noise_alpha = 0, + post_noise_alpha = 0, + remove_background=False, + visualize_flow_sensitivity=None, + warp_kwargs=dict(), +): + """ + Extract noise from a video by warping random noise using optical flow between consecutive frames. + + If running this function in a Jupyter notebook, you'll see a live preview of the noise and visualization as it calculates. + + Args: + video_path (str): Path to the input video file (MP4), + a folder containing image frames, + or a glob pattern like "/path/to/images/*.png", + or a URL pointing to a video file online. + You can also pass a numpy video array in THWC form + noise_channels (int, optional): Number of channels in the generated noise. Defaults to 3. + output_folder (str, optional): Folder to save the output noise and visualization. + Defaults to None, in which case the folder name is automatically chosen. + visualize (bool, optional): Whether to generate visualization images and video. Defaults to True. + resize_frames (tuple or float, optional): Size to resize the input frames. + If a tuple (height, width), resizes to the exact dimensions. + If a float, resizes both dimensions relatively and evenly. Defaults to None. + This is applied *before* calculating optical flow, so it will make flow calculation slower. + resize_flow (int): Resize the optical flows calculated after resize_frames. Controls the underlying noise resolution. + Will not make flow calculation slower - effectively uses bilinear interpolation to upsample the flow. + Should be an integer > 0 + downscale_factor (int): Factor by which to downscale the generated noise. + Larger factor --> smaller noise image. + This factor should evenly divide the height and width of the video frames. + video_preprocessor (callable): A function that takes a THW3 RGB uint8 video numpy array, and returns another THW3 numpy array + This can be used for removing watermarks or resizing the video or any other preprocessing steps + This is applied to the input video, directly after loading it from video_path + save_files (bool): If True, will save files to disk. + progressive_noise_alpha: For ryan, don't worry about it + post_noise_alpha: For ryan, don't worry about it + remove_background: If true will attempt to matte out the background and composite the computed noise on that background + visualize_flow_sensitivity (float, optional): If specified, will normalize flow with the assumption that this given value is the max magnitude + warp_kwargs (dict, optional): For experimental features. Don't worry about this if you're not Ryan Burgert. + + Returns: + EasyDict: A dict containing the following keys: + - 'numpy_noises' (np.ndarray): Generated noise with form [T, H, W, C]. + - 'numpy_flows' (np.ndarray): The (dx, dy)'s with form [T-1, 2, H, W] + - 'vis_frames' (np.ndarray): Visualization frames with form [T, H, W, C]. + - 'output_folder' (str): The path to the folder where outputs are saved (if save_files) + + Examples: + # Command line usage + >>> python noise_warp.py --video_path /path/to/video.mp4 --noise_channels 3 --output_folder /path/to/output + >>> python noise_warp.py --video_path /path/to/frames_folder --resize_frames 0.5 --downscale_factor 2 + >>> python noise_warp.py --video_path "/path/to/frames/frame_*.png" --resize_frames (256, 256) + + # Function call + >>> from rp.git.CommonSource import * + >>> numpy_noises, vis_frames = get_noise_from_video( + ... video_path="/path/to/video.mp4", + ... noise_channels=3, + ... output_folder="/path/to/output", + ... visualize=True, + ... resize_frames=0.5, + ... downscale_factor=2, + ... ) + >>> video_demo("/root/CleanCode/Projects/flow_noise_warping/outputs/water/waves_bilinear.mp4", downscale_factor=4) + >>> video_demo("/efs/users/ryan.burgert/public/sharing/KevinSpinnerNoiseWarping/diffuse_images_360", downscale_factor=8, resize_frames=.5) + >>> video_demo("/root/CleanCode/Projects/flow_noise_warping/outputs/kevin_spinner/kevin_vps7.mp4", downscale_factor=4, resize_frames=.5) + + EXAMPLE - Jupyter Notebook Soup-to-nuts: + + >>> #Ryan Burgert 2024 + ... #Run this in a Jupyter notebook code cell for a realtime preview! + ... + ... #Setup: + ... # Run this in a Jupyter Notebook on a computer with at least one GPU + ... # `sudo apt install ffmpeg git` + ... # `pip install rp` + ... # The first time you run this it might be a bit slow (it will download necessary models) + ... # The `rp` package will take care of installing the rest of the python packages for you + ... + ... import rp + ... + ... rp.git_import('CommonSource') #If missing, installs code from https://github.com/RyannDaGreat/CommonSource + ... import rp.git.CommonSource.noise_warp as nw + ... + ... FRAME = 2**-1 #We immediately resize the input frames by this factor, before calculating optical flow + ... #The flow is calulated at (input size) × FRAME resolution. + ... #Higher FLOW values result in slower optical flow calculation and higher intermediate noise resolution + ... #Larger is not always better - watch the preview in Jupyter to see if it looks good! + ... FLOW = 2**4 #Then, we use bilinear interpolation to upscale the flow by this factor + ... #We warp the noise at (input size) × FRAME × FLOW resolution + ... #The noise is then downsampled back to (input size) + ... #Higher FLOW values result in more temporally consistent noise warping at the cost of higher VRAM usage and slower inference time + ... LATENT = 8 #We further downsample the outputs by this amount - because 8 pixels wide corresponds to one latent wide in Stable Diffusion + ... #The final output size is (input size) ÷ LATENT regardless of FRAME and FLOW + ... + ... LATENT = 2 #Uncomment this line for a prettier visualization! But for any Stable-Diffusion based model, use LATENT=8 + ... + ... + ... #You can pass a glob of image files + ... images = "/path_to_images/cam_*_color.png" + ... + ... #You can also use video files or URLs + ... images = "https://www.shutterstock.com/shutterstock/videos/1100085499/preview/stock-footage-bremen-germany-october-old-style-carousel-moving-on-square-in-city-horses-on-traditional.webm" + ... + ... output_folder = rp.get_folder_name(rp.get_parent_folder(images)) + ... + ... #See this function's docstring for more information! + ... output = nw.get_noise_from_video( + ... images, + ... remove_background=False, #Set this to True to matte the foreground - and force the background to have no flow + ... visualize=True, #Generates nice visualization videos and previews in Jupyter notebook + ... save_files=True, #Set this to False if you just want the noises without saving to a numpy file + ... + ... noise_channels=4, + ... output_folder=output_folder, + ... resize_frames=FRAME, + ... resize_flow=FLOW, + ... downscale_factor=round(FRAME * FLOW) * LATENT, + ... ); + ... + ... print("Noise shape:" ,output.numpy_noises.shape) + ... print("Flow shape:" ,output.numpy_flows .shape) + ... print("Output folder:",output.output_folder) + + """ + + #Input assertions + assert isinstance(resize_flow, int) and resize_flow >= 1, resize_flow + + if device is None: + if rp.currently_running_mac(): + device = 'cpu' + else: + device = rp.select_torch_device(prefer_used=True) + + raft_model = raft.RaftOpticalFlow(device, "large") + + # Load video frames into a [T, H, W, C] numpy array, where C=3 and values are between 0 and 1 + # Can be specified as an MP4, a folder that contains images, or a glob like /path/to/*.png + assert rp.is_numpy_array(video_path) or isinstance(video_path, str), type(video_path) + if rp.is_video_file(video_path) or rp.is_valid_url(video_path): + video_frames = rp.load_video(video_path) + elif rp.is_numpy_array(video_path): + #We can also pass a numpy video as an input in THWC form + video_frames = video_path + assert video_frames.ndim==4, video_frames.ndim + video_path = rp.get_unique_copy_path('noisewarp_video.mp4') + else: + if rp.is_a_folder(video_path): + frame_paths = rp.get_all_image_files(video_path, sort_by='number') + else: + frame_paths = glob.glob(video_path) + frame_paths = sorted(sorted(frame_paths),key=len) + if not frame_paths: + raise ValueError(video_path + " is not a video file, a folder of images, or a glob containing images") + video_frames = rp.load_images(frame_paths, show_progress=True) + + if video_preprocessor is not None: + assert callable(video_preprocessor), type(video_preprocessor) + video_frames = rp.as_numpy_array(video_frames) + video_frames = video_preprocessor(video_frames) + + #If resize_frames is specified, resize all frames to that (height, width) + if resize_frames is not None: + rp.fansi_print("Resizing all input frames to size %s"%str(resize_frames), 'yellow') + video_frames=rp.resize_images(video_frames, size=resize_frames, interp='area') + + if remove_background: + alphas = [] + + background_remover = BackgroundRemover(device) + + if visualize and rp.running_in_jupyter_notebook(): + alpha_display_channel = rp.JupyterDisplayChannel() + alpha_display_channel.display() + + for video_frame in rp.eta(video_frames, title='Removing Backgrounds'): + rgba_image = background_remover(video_frame) + alpha = rp.get_alpha_channel(rgba_image) + alpha = rp.as_float_image(alpha) + alphas.append(alpha) + + if visualize and rp.running_in_jupyter_notebook(): + alpha_display_channel.update( + rp.horizontally_concatenated_images( + rp.with_alpha_checkerboard(rgba_image), alpha + ) + ) + + del background_remover #Free GPU usage + + + video_frames = rp.as_rgb_images(video_frames) + video_frames = np.stack(video_frames) + video_frames = video_frames.astype(np.float16)/255 + _, h, w, _ = video_frames.shape + rp.fansi_print(f"Input video shape: {video_frames.shape}", 'yellow') + + if h%downscale_factor or w%downscale_factor: + rp.fansi_print("WARNING: height {h} or width{w} is not divisible by the downscale_factor {downscale_factor}. This will lead to artifacts in the noise.") + + def downscale_noise(noise): + down_noise = rp.torch_resize_image(noise, 1/downscale_factor, interp='area') #Avg pooling + down_noise = down_noise * downscale_factor #Adjust for STD + return down_noise + + # Decide the location of and create the output folder + if save_files: + if output_folder is None: + output_folder = "outputs/" + rp.get_file_name(video_path, include_file_extension=False) + output_folder = rp.make_directory(rp.get_unique_copy_path(output_folder)) + rp.fansi_print("Output folder: " + output_folder, "green") + + with torch.no_grad(): + + if visualize and rp.running_in_jupyter_notebook(): + # For previewing results in Jupyter notebooks, if applicable + display_channel = rp.JupyterDisplayChannel() + display_channel.display() + + warper = NoiseWarper( + c = noise_channels, + h = resize_flow * h, + w = resize_flow * w, + device = device, + post_noise_alpha = post_noise_alpha, + progressive_noise_alpha = progressive_noise_alpha, + warp_kwargs = warp_kwargs, + ) + + prev_video_frame = video_frames[0] + noise = warper.noise + + down_noise = downscale_noise(noise) + numpy_noise = rp.as_numpy_image(down_noise).astype(np.float16) # In HWC form. Using float16 to save RAM, but it might cause problems on come CPU + + numpy_noises = [numpy_noise] + numpy_flows = [] + vis_frames = [] + + try: + for index, video_frame in enumerate(tqdm(video_frames[1:])): + + dx, dy = raft_model(prev_video_frame, video_frame) + noise = warper(dx, dy).noise + prev_video_frame = video_frame + + numpy_flow = np.stack( + [ + rp.as_numpy_array(dx).astype(np.float16), + rp.as_numpy_array(dy).astype(np.float16), + ] + ) + numpy_flows.append(numpy_flow) + + down_noise = downscale_noise(noise) + + numpy_noise = rp.as_numpy_image(down_noise).astype(np.float16) + + if remove_background: + if 'background_noise' not in dir(): + background_noise = np.random.randn(*numpy_noise.shape) + numpy_noise_alpha = alphas[index] + numpy_noise_alpha = rp.cv_resize_image(numpy_noise_alpha, numpy_noise.shape[:2]) + numpy_noise = blend_noise(background_noise, numpy_noise, numpy_noise_alpha[:,:,None]) + + numpy_noises.append(numpy_noise) + + if visualize: + flow_rgb = rp.optical_flow_to_image(dx, dy, sensitivity = visualize_flow_sensitivity) + + #Turn the noise into a numpy HWC RGB array + down_noise_image = np.zeros((*numpy_noise.shape[:2], 3)) + down_noise_image_c = min(noise_channels,3) + down_noise_image[:,:,:down_noise_image_c]=numpy_noise[:,:,:down_noise_image_c] + + + down_size = rp.get_image_dimensions(down_noise_image) + down_video_frame, down_flow_rgb = rp.resize_images(video_frame, flow_rgb, size=down_size) + + optional_images = [] + optional_labels = [] + if remove_background: + alpha = alphas[index] + down_alpha = rp.cv_resize_image(alpha, down_size) + + optional_images.append(down_alpha) + optional_labels.append('Alpha') + + optional_images.append(rp.with_alpha_checkerboard(rp.with_image_alpha(down_video_frame,down_alpha))) + optional_labels.append('RGBA') + + visualization = rp.as_byte_image( + rp.tiled_images( + rp.labeled_images( + [ + down_noise_image / 3 + 0.5, + down_video_frame, + down_flow_rgb, + down_noise_image / 5 + down_video_frame, + ] + optional_images, + [ + "Warped Noise", + "Input Video", + "Optical Flow", + "Overlaid", + ] + optional_labels, + font = 'G:Zilla Slab', + ) + ) + ) + + if rp.running_in_jupyter_notebook(): + display_channel.update(visualization) + + vis_frames.append(visualization) + + except KeyboardInterrupt: + rp.fansi_print("Interrupted! Returning %i noises" % len(numpy_noises), "cyan", "bold") + pass + + numpy_noises = np.stack(numpy_noises).astype(np.float16) + numpy_flows = np.stack(numpy_flows).astype(np.float16) + if vis_frames: vis_frames = np.stack(vis_frames) + + if save_files and len(vis_frames): + vis_img_folder = rp.make_directory(output_folder + "/visualization_images") + vis_img_paths = rp.path_join(vis_img_folder, "visual_%05i.png") + rp.save_images(vis_frames, vis_img_paths, show_progress=True) + + if "ffmpeg" in rp.get_system_commands(): + vis_mp4_path = rp.path_join(output_folder, "visualization_video.mp4") + noise_mp4_path = rp.path_join(output_folder, "noise_video.mp4") + rp.save_video_mp4( + vis_frames, + vis_mp4_path, + video_bitrate="max", + framerate=30, + ) + rp.save_video_mp4( + (numpy_noises / 4 + 0.5)[:,:,:,:3], + noise_mp4_path, + video_bitrate="max", + framerate=30, + ) + if rp.is_video_file(video_path): + try: + #If possible, try to add the original audio and framerate back again + #Only makes sense if the input was an MP4 file and not a folder of images etc + for output_video_path in [vis_mp4_path, noise_mp4_path]: + rp.fansi_print( + "Added audio to output at: " + + rp.add_audio_to_video_file( + rp.printed( + rp.change_video_file_framerate( + output_video_path, + rp.get_video_file_framerate(video_path), + ) + ), + video_path, + ), + "green", + "bold", + ) + except Exception: + rp.print_stack_trace() + else: + rp.fansi_print("Please install ffmpeg! We won't save an MP4 this time - please try again.") + + if save_files: + noises_path = rp.path_join(output_folder, "noises.npy") + flows_path = rp.path_join(output_folder, "flows_dxdy.npy") + np.save(noises_path, numpy_noises) + rp.fansi_print("Saved " + noises_path + " with shape " + str(numpy_noises.shape), "green") + np.save(flows_path, numpy_flows) + rp.fansi_print("Saved " + flows_path + " with shape " + str(numpy_flows.shape), "green") + + rp.fansi_print(rp.get_file_name(__file__)+": Done warping noise, results are at " + rp.get_absolute_path(output_folder), "green", "bold") + + return rp.gather_vars('numpy_noises numpy_flows vis_frames output_folder') + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a81c341 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +torch +numpy +einops +fire +tqdm +opencv-python +rp