python You can use the following code to call our trained style encoder. Hope it helps. import torchvision.transforms.functional as F from torchvision import transforms from transformers import (AutoModel, AutoProcessor, AutoTokenizer, AutoConfig, CLIPImageProcessor, CLIPVisionModelWithProjection) class SEStyleEmbedding: def init (self, pretrained path: str = "xingpng/OneIG StyleEncoder", device: str = "cuda", dtype=torch.bfloat16): self.device = torch.device(device) self.dtype = dtype self.image encoder = CLIPVisionModelWithProjection.from pretrained(pretrained path) self.image encoder.to(self.device, dtype=self.dtype) self.image encoder.eval() self.processor = CLIPImageProcessor() def l2 normalize(self, x): return torch.nn.functional.normalize(x, p=2, dim= 1) def get style embedding(self, image path: str): image = Image.open(image path).convert('RGB') inputs = self.processor(images=image, return tensors="pt").pixel values.to(self.device, dtype=self.dtype) with torch.no grad(): outputs = self.image encoder(inputs) image embeds = outputs.image embeds image embeds norm = self. l2 normalize(image embeds) return image embeds norm
We use cookies for essential functionality and analytics. You can accept or reject analytics cookies.Cookie policy