File size: 7,365 Bytes
6fecfbe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
import base64
import io
from PIL import Image, ImageOps
from nodes import ImageScale, VAEEncode
import numpy as np
from perlin_noise import perlin_noise
import torch
from torchvision.transforms import ToPILImage
from ..utils import VyroParams
from .perlin import perlin_power_fractal_batch
from PIL import Image
def pil2tensor(image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
class VyroPipeInputV2:
def __init__(self):
self.vae_encoder = VAEEncode()
self.image_scale = ImageScale()
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"user_prompt": ("STRING", {"multiline": True}),
"mode": (VyroParams.MODE, {"default": VyroParams.MODE[0]}),
"vae": ("VAE",),
"init_noise_mode": (["perlin1","perlin2","zeros"], {"default": "perlin1"}),
},
"optional": {
"user_neg_prompt": ("STRING",{"multiline": True}),
"batch_size": ("INT", {"default": 1, "max": 4, "min": 1, "step": 1}),
"cfg": ("FLOAT", {"default": 7.5, "min":1.0, "max":30.0, "step":0.1}),
"steps": ("INT", {"default": 20, "min":10, "max":150, "step":1}),
"width": ("INT", {"default": 1024, "min":64, "max":4096, "step":8}),
"height": ("INT", {"default": 1024, "min":64, "max":4096, "step":8}),
"seed": ("INT", {"default": 1, "min": 1, "max": 2**32 - 1}),
"init_img": ("STRING", {"multiline": True}),
"denoise": ("FLOAT", {"default" : 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"stage1_strength": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
"stage2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"efficiency_multiplier": (VyroParams.MULTIPLIER, {"default": VyroParams.MULTIPLIER[0]}),
}
}
RETURN_TYPES = ("VYRO_PARAMS",)
RETURN_NAMES = ("vyro_params",)
FUNCTION = "prep_input"
#OUTPUT_NODE = False
CATEGORY = "Vyro"
def blank_image(self, width, height, red, green, blue):
# Ensure multiples
width = (width // 8) * 8
height = (height // 8) * 8
# Blend image
blank = Image.new(mode="RGB", size=(width, height),
color=(red, green, blue))
return pil2tensor(blank)
def image_to_mask(self, image, channel):
channels = ["red", "green", "blue", "alpha"]
mask = image[:, :, :, channels.index(channel)]
return (mask,)
def mask_to_image(self, mask):
result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
return (result,)
def upscale_image(self,image, upscale_method='nearest', width=512, height=512, crop='disabled'):
if width == 0 and height == 0:
return image
original_aspect_ratio = image.width / image.height
target_aspect_ratio = width / height
pil_methods = {
'nearest': Image.NEAREST,
'bilinear': Image.BILINEAR,
'area': Image.HAMMING,
'bicubic': Image.BICUBIC,
'lanczos': Image.LANCZOS
}
if crop == 'center':
if original_aspect_ratio > target_aspect_ratio:
new_height = height
new_width = int(round(height * original_aspect_ratio))
else:
new_width = width
new_height = int(round(width / original_aspect_ratio))
resized_image = image.resize((new_width, new_height), pil_methods[upscale_method])
left = (resized_image.width - width) / 2
top = (resized_image.height - height) / 2
right = (resized_image.width + width) / 2
bottom = (resized_image.height + height) / 2
resized_image = resized_image.crop((left, top, right, bottom))
else:
if width == 0:
width = int(round(height * original_aspect_ratio))
elif height == 0:
height = int(round(width / original_aspect_ratio))
resized_image = image.resize((width, height), pil_methods[upscale_method])
return resized_image
def prep_input(self, user_prompt, mode, vae, init_noise_mode, user_neg_prompt='', batch_size=1, cfg=7.5, steps=20, width=1024, height=1024, seed=1, init_img='', denoise=1.0,stage1_strength=0.25, stage2_strength=1.0, efficiency_multiplier=1.0):
width = int(width / efficiency_multiplier)
height = int(height / efficiency_multiplier)
# width = (width // 8) * 8
# height = (height // 8) * 8
if init_img == "" or init_img is None or init_img == 'undefined':
generator = torch.cuda.manual_seed(seed)
if 'perlin2' in init_noise_mode:
init_img = perlin_noise(grid_shape=(2, 8), out_shape=(width // 8, height // 8), batch_size=batch_size*4, generator=generator)
init_img = init_img.reshape(batch_size, 4, height // 8, width // 8)
elif 'perlin1' in init_noise_mode:
init_img = perlin_power_fractal_batch(
batch_size=batch_size,
width=width // 8,
height=height // 8,
X=0,
Y=0,
Z=0,
frame=0,
seed=seed
)[0]
init_img = init_img.reshape(batch_size, 4, height // 8, width // 8)
else:
init_img = torch.zeros([batch_size, 4, height // 8, width // 8])
# init_img = self.blank_image(width, height, 0, 0, 0)
# init_img = init_img.unsqueeze(0)
# init_img = latent = torch.zeros([batch_size, 4, height // 8, width // 8])
else:
# Convert init_img from byte string to PIL image
# Get bytes from string
init_img_bytes = base64.b64decode(init_img)
img = Image.open(io.BytesIO(init_img_bytes))
# Center crop to width/height
init_img = pil2tensor(img)
init_img = self.image_scale.upscale(init_img, "nearest-exact", int(width), int(height), "center")[0]
init_img = init_img.squeeze(0)
init_img = init_img.repeat(batch_size, 1, 1, 1)
init_img = self.vae_encoder.encode(vae, init_img)[0]['samples']
params = VyroParams(latents={"samples":init_img}, user_prompt=user_prompt, mode=mode, cfg=cfg, batch_size=batch_size, steps=steps, width=width, height=height, seed=seed, denoise=denoise, user_neg_prompt=user_neg_prompt, stage1_strength=stage1_strength, stage2_strength=stage2_strength, efficiency_multiplier=efficiency_multiplier)
return (params,)
NODE_CLASS_MAPPINGS = {
"Vyro Pipe Input V2": VyroPipeInputV2,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"VyroPipeInputV2": "Vyro Pipe Input V2",
}
|