224 lines
8.0 KiB
Python
224 lines
8.0 KiB
Python
#!/usr/bin/env python
|
|
|
|
from . import py3d_tools as p3dT
|
|
from . import disco_xform_utils as dxf
|
|
import torchvision.transforms as T
|
|
import cv2
|
|
import pandas as pd
|
|
import gc
|
|
import io
|
|
import math
|
|
import lpips
|
|
from PIL import Image, ImageOps
|
|
import requests
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
import torchvision.transforms.functional as TF
|
|
import subprocess
|
|
from importlib import util as importlibutil
|
|
import numpy as np
|
|
import os
|
|
import requests
|
|
import tqdm
|
|
from urllib.parse import urlparse
|
|
|
|
import comfy.model_management
|
|
|
|
def split_prompts(prompts, max_frames):
|
|
prompt_series = pd.Series([np.nan for a in range(max_frames)])
|
|
for i, prompt in prompts.items():
|
|
prompt_series[i] = prompt
|
|
# prompt_series = prompt_series.astype(str)
|
|
prompt_series = prompt_series.ffill().bfill()
|
|
return prompt_series
|
|
|
|
# https://gist.github.com/adefossez/0646dbe9ed4005480a2407c62aac8869
|
|
|
|
def interp(t):
|
|
return 3 * t**2 - 2 * t ** 3
|
|
|
|
def perlin(width, height, scale=10, device=None):
|
|
gx, gy = torch.randn(2, width + 1, height + 1, 1, 1, device=device)
|
|
xs = torch.linspace(0, 1, scale + 1)[:-1, None].to(device)
|
|
ys = torch.linspace(0, 1, scale + 1)[None, :-1].to(device)
|
|
wx = 1 - interp(xs)
|
|
wy = 1 - interp(ys)
|
|
dots = 0
|
|
dots += wx * wy * (gx[:-1, :-1] * xs + gy[:-1, :-1] * ys)
|
|
dots += (1 - wx) * wy * (-gx[1:, :-1] * (1 - xs) + gy[1:, :-1] * ys)
|
|
dots += wx * (1 - wy) * (gx[:-1, 1:] * xs - gy[:-1, 1:] * (1 - ys))
|
|
dots += (1 - wx) * (1 - wy) * (-gx[1:, 1:] * (1 - xs) - gy[1:, 1:] * (1 - ys))
|
|
return dots.permute(0, 2, 1, 3).contiguous().view(width * scale, height * scale)
|
|
|
|
def perlin_ms(octaves, width, height, grayscale, device=None):
|
|
if not device:
|
|
device = comfy.model_management.get_torch_device()
|
|
out_array = [0.5] if grayscale else [0.5, 0.5, 0.5]
|
|
# out_array = [0.0] if grayscale else [0.0, 0.0, 0.0]
|
|
for i in range(1 if grayscale else 3):
|
|
scale = 2 ** len(octaves)
|
|
oct_width = width
|
|
oct_height = height
|
|
for oct in octaves:
|
|
p = perlin(oct_width, oct_height, scale, device)
|
|
out_array[i] += p * oct
|
|
scale //= 2
|
|
oct_width *= 2
|
|
oct_height *= 2
|
|
return torch.cat(out_array)
|
|
|
|
def create_perlin_noise(side_x, side_y, octaves=[1, 1, 1, 1], width=2, height=2, grayscale=True):
|
|
out = perlin_ms(octaves, width, height, grayscale)
|
|
if grayscale:
|
|
out = TF.resize(size=(side_y, side_x), img=out.unsqueeze(0))
|
|
out = TF.to_pil_image(out.clamp(0, 1)).convert('RGB')
|
|
else:
|
|
out = out.reshape(-1, 3, out.shape[0]//3, out.shape[1])
|
|
out = TF.resize(size=(side_y, side_x), img=out)
|
|
out = TF.to_pil_image(out.clamp(0, 1).squeeze())
|
|
|
|
out = ImageOps.autocontrast(out)
|
|
return out
|
|
|
|
def regen_perlin(perlin_mode, batch_size, expand=False):
|
|
if perlin_mode == 'color':
|
|
init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)
|
|
init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, False)
|
|
elif perlin_mode == 'gray':
|
|
init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, True)
|
|
init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)
|
|
else:
|
|
init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)
|
|
init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)
|
|
|
|
device = comfy.model_management.get_torch_device()
|
|
init = TF.to_tensor(init).add(TF.to_tensor(init2)).div(2).to(device).unsqueeze(0).mul(2).sub(1)
|
|
del init2
|
|
if expand:
|
|
return init.expand(batch_size, -1, -1, -1)
|
|
return init
|
|
|
|
def fetch(url_or_path):
|
|
if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):
|
|
r = requests.get(url_or_path)
|
|
r.raise_for_status()
|
|
fd = io.BytesIO()
|
|
fd.write(r.content)
|
|
fd.seek(0)
|
|
return fd
|
|
return open(url_or_path, 'rb')
|
|
|
|
def read_image_workaround(path):
|
|
"""OpenCV reads images as BGR, Pillow saves them as RGB. Work around
|
|
this incompatibility to avoid colour inversions."""
|
|
im_tmp = cv2.imread(path)
|
|
return cv2.cvtColor(im_tmp, cv2.COLOR_BGR2RGB)
|
|
|
|
def parse_prompt(prompt):
|
|
if prompt.startswith('http://') or prompt.startswith('https://'):
|
|
vals = prompt.rsplit(':', 2)
|
|
vals = [vals[0] + ':' + vals[1], *vals[2:]]
|
|
else:
|
|
vals = prompt.rsplit(':', 1)
|
|
vals = vals + ['', '1'][len(vals):]
|
|
return vals[0], float(vals[1])
|
|
|
|
def sinc(x):
|
|
return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))
|
|
|
|
def lanczos(x, a):
|
|
cond = torch.logical_and(-a < x, x < a)
|
|
out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))
|
|
return out / out.sum()
|
|
|
|
def ramp(ratio, width):
|
|
n = math.ceil(width / ratio + 1)
|
|
out = torch.empty([n])
|
|
cur = 0
|
|
for i in range(out.shape[0]):
|
|
out[i] = cur
|
|
cur += ratio
|
|
return torch.cat([-out[1:].flip([0]), out])[1:-1]
|
|
|
|
def resample(input, size, align_corners=True):
|
|
n, c, h, w = input.shape
|
|
dh, dw = size
|
|
|
|
input = input.reshape([n * c, 1, h, w])
|
|
|
|
if dh < h:
|
|
kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)
|
|
pad_h = (kernel_h.shape[0] - 1) // 2
|
|
input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')
|
|
input = F.conv2d(input, kernel_h[None, None, :, None])
|
|
|
|
if dw < w:
|
|
kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)
|
|
pad_w = (kernel_w.shape[0] - 1) // 2
|
|
input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')
|
|
input = F.conv2d(input, kernel_w[None, None, None, :])
|
|
|
|
input = input.reshape([n, c, h, w])
|
|
return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)
|
|
|
|
def spherical_dist_loss(x, y):
|
|
x = F.normalize(x, dim=-1)
|
|
y = F.normalize(y, dim=-1)
|
|
return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)
|
|
|
|
def tv_loss(input):
|
|
"""L2 total variation loss, as in Mahendran et al."""
|
|
input = F.pad(input, (0, 1, 0, 1), 'replicate')
|
|
x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]
|
|
y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]
|
|
return (x_diff**2 + y_diff**2).mean([1, 2, 3])
|
|
|
|
|
|
def range_loss(input):
|
|
return (input - input.clamp(-1, 1)).pow(2).mean([1, 2, 3])
|
|
|
|
|
|
def alpha_sigma_to_t(alpha, sigma):
|
|
return torch.atan2(sigma, alpha) * 2 / math.pi
|
|
|
|
|
|
normalize = T.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])
|
|
|
|
|
|
def pyget(url, path=None, filename=None, progress=True):
|
|
try:
|
|
response = requests.get(url, stream=True)
|
|
response.raise_for_status()
|
|
parsed_url = urlparse(url)
|
|
filename = filename if filename else os.path.basename(parsed_url.path)
|
|
path = os.path.join(path, filename) if path else filename
|
|
total_size_in_bytes = int(response.headers.get('content-length', 0))
|
|
pbar = None
|
|
if progress:
|
|
pbar = comfy.utils.ProgressBar(total_size_in_bytes)
|
|
tqdm_bar = tqdm.tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
chunk_size = 10 * 1024 * 1024
|
|
with open(path, 'wb') as file:
|
|
for chunk in response.iter_content(chunk_size):
|
|
if chunk:
|
|
if progress:
|
|
chunk_length = len(chunk)
|
|
tqdm_bar.update(chunk_length)
|
|
pbar.update(chunk_length)
|
|
file.write(chunk)
|
|
if os.path.exists(path):
|
|
return True
|
|
else:
|
|
print(f"Unable to save file to: {path}")
|
|
except requests.exceptions.HTTPError as errh:
|
|
print(f"HTTP Error: ({url}): {errh}")
|
|
except requests.exceptions.ConnectionError as errc:
|
|
print(f"Connection Error: ({url}): {errc}")
|
|
except requests.exceptions.Timeout as errt:
|
|
print(f"Timeout Error: ({url}): {errt}")
|
|
except requests.exceptions.RequestException as err:
|
|
print(f"Request Exception: ({url}): {err}")
|
|
return False
|