| 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, |
| ], |
| ] |
|
|
