Views
No views yet
1# You can use the following code to call our trained style encoder. Hope it helps.
2import torchvision.transforms.functional as F
3from torchvision import transforms
4from transformers import (AutoModel, AutoProcessor, AutoTokenizer, AutoConfig,
5 CLIPImageProcessor, CLIPVisionModelWithProjection)
6class SEStyleEmbedding:
7 def __init__(self, pretrained_path: str = "xingpng/OneIG-StyleEncoder", device: str = "cuda", dtype=torch.bfloat16):
8 self.device = torch.device(device)
9 self.dtype = dtype
10 self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(pretrained_path)
11 self.image_encoder.to(self.device, dtype=self.dtype)
12 self.image_encoder.eval()
13 self.processor = CLIPImageProcessor()
14
15 def _l2_normalize(self, x):
16 return torch.nn.functional.normalize(x, p=2, dim=-1)
17
18 def get_style_embedding(self, image_path: str):
19 image = Image.open(image_path).convert('RGB')
20 inputs = self.processor(images=image, return_tensors="pt").pixel_values.to(self.device, dtype=self.dtype)
21
22 with torch.no_grad():
23 outputs = self.image_encoder(inputs)
24 image_embeds = outputs.image_embeds
25 image_embeds_norm = self._l2_normalize(image_embeds)
26 return image_embeds_norm