|
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] |
|
|
|
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 = snapshot_download( |
|
repo_id="RunDiffusion/Juggernaut-XL-v6", allow_patterns="*.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) |
|
for model in Path(MODEL_PATH).glob("*.safetensors"): |
|
model_target_path = Path(f"./ComfyUI/models/checkpoints/{model.name}") |
|
if not model_target_path.exists(): |
|
os.symlink(model, model_target_path) |
|
|
|
|
|
examples = [ |
|
[ |
|
"A very cute monster cat on a glass bottle", |
|
"ugly distorted image, low quality, text, bad, not good ,watermark", |
|
None, |
|
False, |
|
None, |
|
1231231, |
|
5, |
|
], |
|
[ |
|
"A picture from above captures a beautiful, small toucan bird flying in the sky.", |
|
"ugly distorted image, low quality, text, bad, not good ,watermark", |
|
"./examples/bg.png", |
|
False, |
|
"SDXL, Background", |
|
1234144, |
|
8, |
|
], |
|
[ |
|
"a photo a men surrounded by a crowd of people in a circle", |
|
"ugly distorted image, low quality, text, bad, not good ,watermark", |
|
"./examples/lecun.png", |
|
True, |
|
"SDXL, Foreground", |
|
123123, |
|
10, |
|
], |
|
[ |
|
"An image of a galaxy", |
|
"ugly distorted image, low quality, text, bad, not good ,watermark", |
|
"./examples/julien.png", |
|
True, |
|
"SDXL, Foreground", |
|
123123, |
|
10, |
|
], |
|
[ |
|
"a men jumping on swiming pool full of people", |
|
"ugly distorted image, low quality, text, bad, not good ,watermark", |
|
"./examples/old_jump.png", |
|
False, |
|
"SDXL, Foreground", |
|
5350795678007195000, |
|
10, |
|
], |
|
[ |
|
"a cute cat flying over Manhattan time square", |
|
"ugly distorted image, low quality, text, bad, not good ,watermark", |
|
"./examples/cat.png", |
|
True, |
|
"SDXL, Foreground", |
|
123123, |
|
10, |
|
], |
|
] |
|
|
