File size: 5,182 Bytes
bb5540a a389f87 bb5540a a389f87 bb5540a a389f87 bb5540a a389f87 bb5540a a389f87 bb5540a a389f87 |
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 |
import os
import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from PIL import Image, ImageOps, ImageSequence
from typing import List
from pathlib import Path
from huggingface_hub import snapshot_download, hf_hub_download
def tensor_to_pil(images: torch.Tensor | List[torch.Tensor]) -> List[Image.Image]:
if not isinstance(images, list):
images = [images]
imgs = []
for image in images:
i = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(np.squeeze(i), 0, 255).astype(np.uint8))
imgs.append(img)
return imgs
def pad_image(input_image, background_color=(0, 0, 0)):
w, h = input_image.size
pad_w = (64 - w % 64) % 64
pad_h = (64 - h % 64) % 64
new_size = (w + pad_w, h + pad_h)
im_padded = Image.new(input_image.mode, new_size, background_color)
im_padded.paste(input_image, (pad_w // 2, pad_h // 2))
if im_padded.size[0] == im_padded.size[1]:
return im_padded
elif im_padded.size[0] > im_padded.size[1]:
new_size = (im_padded.size[0], im_padded.size[0])
new_image = Image.new(im_padded.mode, new_size, background_color)
new_image.paste(im_padded, (0, (new_size[1] - im_padded.size[1]) // 2))
return new_image
else:
new_size = (im_padded.size[1], im_padded.size[1])
new_image = Image.new(im_padded.mode, new_size, background_color)
new_image.paste(im_padded, ((new_size[0] - im_padded.size[0]) // 2, 0))
return new_image
def pil_to_tensor(image: Image.Image) -> tuple[torch.Tensor, torch.Tensor]:
output_images = []
output_masks = []
for i in ImageSequence.Iterator(image):
i = ImageOps.exif_transpose(i)
if i.mode == "I":
i = i.point(lambda i: i * (1 / 255))
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
if "A" in i.getbands():
mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
mask = 1.0 - torch.from_numpy(mask)
else:
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
output_images.append(image)
output_masks.append(mask.unsqueeze(0))
if len(output_images) > 1:
output_image = torch.cat(output_images, dim=0)
output_mask = torch.cat(output_masks, dim=0)
else:
output_image = output_images[0]
output_mask = output_masks[0]
return (output_image, output_mask)
def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
if len(im.shape) < 3:
im = im[:, :, np.newaxis]
# orig_im_size=im.shape[0:2]
im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
im_tensor = F.interpolate(
torch.unsqueeze(im_tensor, 0), size=model_input_size, mode="bilinear"
).type(torch.uint8)
image = torch.divide(im_tensor, 255.0)
image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
return image
def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
result = torch.squeeze(F.interpolate(result, size=im_size, mode="bilinear"), 0)
ma = torch.max(result)
mi = torch.min(result)
result = (result - mi) / (ma - mi)
im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8)
im_array = np.squeeze(im_array)
return im_array
def downloadModels():
MODEL_PATH = hf_hub_download(
repo_id="RunDiffusion/Juggernaut-XL-v6",
filename="juggernautXL_version6Rundiffusion.safetensors",
)
# MODEL_PATH = hf_hub_download(
# repo_id="lllyasviel/fav_models",
# subfolder="fav",
# filename="juggernautXL_v8Rundiffusion.safetensors",
# )
LAYERS_PATH = snapshot_download(
repo_id="LayerDiffusion/layerdiffusion-v1", allow_patterns="*.safetensors"
)
for file in Path(LAYERS_PATH).glob("*.safetensors"):
target_path = Path(f"./ComfyUI/models/layer_model/{file.name}")
if not target_path.exists():
os.symlink(file, target_path)
model_target_path = Path(
"./ComfyUI/models/checkpoints/juggernautXL_v8Rundiffusion.safetensors"
)
if not model_target_path.exists():
os.symlink(MODEL_PATH, model_target_path)
examples = [
[
"An old men sit on a chair looking at the sky",
"ugly distorted image, low quality, text, bad, not good ,watermark",
None,
False,
None,
],
[
"A beautiful toucan bird flying in the sky",
"ugly distorted image, low quality, text, bad, not good ,watermark",
"./examples/bg.png",
False,
"SDXL, Background",
],
[
"A men watching a concert",
"ugly distorted image, low quality, text, bad, not good ,watermark",
"./examples/lecun.png",
True,
"SDXL, Foreground",
],
[
"A men watching a concert",
"ugly distorted image, low quality, text, bad, not good ,watermark",
"./examples/julien.png",
True,
"SDXL, Foreground",
],
]
|