| import math |
| import cv2 |
| import numpy as np |
| import torch |
| from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
| from transformers import CLIPVisionModelWithProjection |
| from preservation import PreservedPipeline, preservation_mask |
|
|
| BASE_MODEL = "OnomaAIResearch/Illustrious-XL-v2.0" |
| BASE_FILE = "Illustrious-XL-v2.0.safetensors" |
| BASE_CONFIG = "stabilityai/stable-diffusion-xl-base-1.0" |
| CONTROL_MODEL = "diffusers/controlnet-canny-sdxl-1.0" |
| ADAPTER_MODEL = "h94/IP-Adapter" |
| VAE_MODEL = "madebyollin/sdxl-vae-fp16-fix" |
| SIZE = 1024 |
|
|
|
|
| def prepare_illustrious(input_path, boxes=(), text_mask=None): |
| with Image.open(input_path) as source: |
| original = source.convert("RGB") |
| width, height = original.size |
| scale = min(SIZE / width, SIZE / height) |
| size = (round(width * scale), round(height * scale)) |
| x, y = (SIZE - size[0]) // 2, (SIZE - size[1]) // 2 |
| image = Image.new("RGB", (SIZE, SIZE)) |
| image.paste(original.resize(size, Image.Resampling.LANCZOS), (x, y)) |
|
|
| edges = cv2.Canny(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY), 100, 150) |
| control = Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2)) |
| sx, sy = size[0] / width, size[1] / height |
| padded_mask = Image.new("L", (SIZE, SIZE)) |
| if text_mask is not None: |
| padded_mask.paste(text_mask.resize(size, Image.Resampling.NEAREST), (x, y)) |
| mask = preservation_mask((SIZE, SIZE), [ |
| (x + math.floor(x1 * sx), y + math.floor(y1 * sy), |
| x + math.ceil(x2 * sx), y + math.ceil(y2 * sy)) |
| for x1, y1, x2, y2 in boxes |
| ], padded_mask) |
| return image, control, mask, (x, y, x + size[0], y + size[1]), (width, height) |
|
|
|
|
| def load_illustrious(device): |
| encoder = CLIPVisionModelWithProjection.from_pretrained( |
| ADAPTER_MODEL, subfolder="models/image_encoder", torch_dtype=torch.float16, |
| ) |
| pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_single_file( |
| hf_hub_download(repo_id=BASE_MODEL, filename=BASE_FILE), |
| config=BASE_CONFIG, |
| controlnet=ControlNetModel.from_pretrained(CONTROL_MODEL, torch_dtype=torch.float16), |
| image_encoder=encoder, |
| vae=AutoencoderKL.from_pretrained(VAE_MODEL, torch_dtype=torch.float16), |
| torch_dtype=torch.float16, |
| ) |
| pipe.load_ip_adapter( |
| ADAPTER_MODEL, subfolder="sdxl_models", |
| weight_name="ip-adapter-plus_sdxl_vit-h.safetensors", |
| ) |
| pipe.set_ip_adapter_scale(0.5) |
| return pipe.to(device) |
|
|
|
|
| def infer_illustrious(pipe, image, control, mask, strength, steps): |
| return PreservedPipeline(pipe, [mask])( |
| prompt="masterpiece, best quality, anime illustration", |
| negative_prompt="photorealistic, low quality, blurry", |
| image=image, |
| control_image=control, |
| ip_adapter_image=image, |
| strength=strength, |
| num_inference_steps=steps, |
| guidance_scale=5.0, |
| ).images[0] |
|
|
|
|
| def save_illustrious(result, crop, size, output_path): |
| result.crop(crop).resize(size, Image.Resampling.LANCZOS).save(output_path) |
|
|
