|
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): |
|
pad_w, pad_h = ( |
|
np.max(((2, 2), np.ceil(np.array(input_image.size) / 64).astype(int)), axis=0) |
|
* 64 |
|
- input_image.size |
|
) |
|
im_padded = Image.fromarray( |
|
np.pad(np.array(input_image), ((0, pad_h), (0, pad_w), (0, 0)), mode="edge") |
|
) |
|
w, h = im_padded.size |
|
if w == h: |
|
return im_padded |
|
elif w > h: |
|
new_image = Image.new(im_padded.mode, (w, w), (0, 0, 0)) |
|
new_image.paste(im_padded, (0, (w - h) // 2)) |
|
return new_image |
|
else: |
|
new_image = Image.new(im_padded.mode, (h, h), (0, 0, 0)) |
|
new_image.paste(im_padded, ((h - w) // 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 = 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) |
|
|
