Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,502 Bytes
db1df47 |
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 |
import torch
import cv2
import numpy as np
from torchvision.transforms.functional import to_tensor, center_crop, resize
from PIL import Image
from EngageEngine.pipeline import EngagePipeline
from diffusers import (
EulerAncestralDiscreteScheduler,
AutoencoderKL, ControlNetModel,
)
def process_sketch(x, im_size=(1024, 1024), sketch_detail=0.5, sketch_softness=0.5):
x_b = Image.new("RGBA", x.size, "WHITE")
x_b.paste(x, mask=x)
x = to_tensor(x_b.convert('RGB')).unsqueeze(0)
x = center_crop(x, x.shape[-1])
x = resize(x, im_size)
u_th = (1 - sketch_detail) * 190 + 10
l_th = (1 - sketch_detail) ** (sketch_softness * 8 + 1) * 185 + 5
edges = [cv2.Canny(x[i].mul(255).permute(1, 2, 0).numpy().astype(np.uint8),
u_th, l_th, L2gradient=True) for i in range(len(x))]
edges = torch.stack([torch.tensor(e).div(255).unsqueeze(0) for e in edges], dim=0)
edges = torch.concatenate([edges, edges, edges], dim=1)
return edges
def process_mask(x, mask, im_size=(1024, 1024)):
x = to_tensor(x.convert('RGB')).unsqueeze(0)
x = center_crop(x, x.shape[-1])
x = resize(x, im_size)
mask = to_tensor(mask.convert('L')).unsqueeze(0)
mask = center_crop(mask, mask.shape[-1])
mask = resize(mask, im_size)
return x, mask
def fetch_model():
# Load VAE component
vae = AutoencoderKL.from_pretrained(
"madebyollin/sdxl-vae-fp16-fix",
torch_dtype=torch.float16
)
controlnet = ControlNetModel.from_pretrained(
"diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16
)
# Configure the pipeline
pipe = EngagePipeline.from_pretrained(
"dataautogpt3/ProteusV0.4-Lightning",
vae=vae,
controlnet=controlnet,
torch_dtype=torch.float16
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights("EngageEngine/ENGAGE_LORA.safetensors", adapter_name="ENGAGE_LORA")
pipe.load_lora_weights("EngageEngine/FILM_LORA.safetensors", adapter_name="FILM_LORA")
pipe.load_lora_weights("EngageEngine/MJ_LORA.safetensors", adapter_name="MJ_LORA")
pipe.load_lora_weights("EngageEngine/MORE_ART_LORA.safetensors", adapter_name="MORE_ART_LORA")
pipe.set_adapters(["ENGAGE_LORA", "FILM_LORA", "MJ_LORA", "MORE_ART_LORA"], adapter_weights=[0.0, 0.0, 0.0, 0.0])
pipe.to('cuda')
return pipe
|