File size: 3,788 Bytes
eadc461 4363538 eadc461 4363538 eadc461 4363538 eadc461 4363538 eadc461 4363538 eadc461 4363538 eadc461 4363538 eadc461 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | from typing import Dict, List, Any
import torch
from torch import autocast
from huggingface_hub import hf_hub_download
from diffusers import DiffusionPipeline
import base64
from io import BytesIO
from safetensors.torch import load_file
from cog_sdxl.dataset_and_utils import TokenEmbeddingsHandler
from cog_sdxl.no_init import no_init_or_tensor
from diffusers.models.attention_processor import LoRAAttnProcessor2_0
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device ~>", device)
class EndpointHandler:
def __init__(self, path=""):
print("path ~>", path)
self.pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16 if device.type == "cuda" else None,
variant="fp16",
).to(device)
lora_path = hf_hub_download(repo_id="SvenN/sdxl-emoji", filename="lora.safetensors", repo_type="model")
embeddings_path = hf_hub_download(repo_id="SvenN/sdxl-emoji", filename="embeddings.pti", repo_type="model")
#Load the LoRA into the UNet
unet = self.pipe.unet
tensors = load_file(lora_path)
unet_lora_attn_procs = {}
name_rank_map = {}
for tk, tv in tensors.items():
# up is N, d
tensors[tk] = tv.half()
if tk.endswith("up.weight"):
proc_name = ".".join(tk.split(".")[:-3])
r = tv.shape[1]
name_rank_map[proc_name] = r
for name, attn_processor in unet.attn_processors.items():
cross_attention_dim = (
None
if name.endswith("attn1.processor")
else unet.config.cross_attention_dim
)
if name.startswith("mid_block"):
hidden_size = unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(unet.config.block_out_channels))[
block_id
]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = unet.config.block_out_channels[block_id]
with no_init_or_tensor():
module = LoRAAttnProcessor2_0(
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
rank=name_rank_map[name],
).half()
unet_lora_attn_procs[name] = module.to("cuda", non_blocking=True)
unet.set_attn_processor(unet_lora_attn_procs)
unet.load_state_dict(tensors, strict=False)
#Load the text embeddings into the text encoder/tokenizer
handler = TokenEmbeddingsHandler(
[self.pipe.text_encoder, self.pipe.text_encoder_2], [self.pipe.tokenizer, self.pipe.tokenizer_2]
)
handler.load_embeddings(embeddings_path)
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
"""
Args:
data (:obj:):
includes the input data and the parameters for the inference.
Return:
A :obj:`dict`:. base64 encoded image
"""
inputs = data.pop("inputs", data)
# Automatically add trigger tokens to the beginning of the prompt
images = self.pipe(
inputs,
cross_attention_kwargs={"scale": 0.6},
**data['parameters']
).images
image = images[0]
return image
if __name__ == "__main__":
handler = EndpointHandler()
print(handler)
output = handler({"inputs": "emoji of a tiger face, white background"})
print(output)
|