Upload 2 files

ui_model.py

@@ -0,0 +1,70 @@
+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