SunderAli17 commited on
Commit
f458bd8
1 Parent(s): be0d2f2

Upload 3 files

Browse files
Files changed (3) hide show
  1. toonmage/fluxpipeline.py +188 -0
  2. toonmage/pipeline.py +232 -0
  3. toonmage/utils.py +76 -0
toonmage/fluxpipeline.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gc
2
+
3
+ import cv2
4
+ import insightface
5
+ import torch
6
+ import torch.nn as nn
7
+ from basicsr.utils import img2tensor, tensor2img
8
+ from facexlib.parsing import init_parsing_model
9
+ from facexlib.utils.face_restoration_helper import FaceRestoreHelper
10
+ from huggingface_hub import hf_hub_download, snapshot_download
11
+ from insightface.app import FaceAnalysis
12
+ from safetensors.torch import load_file
13
+ from torchvision.transforms import InterpolationMode
14
+ from torchvision.transforms.functional import normalize, resize
15
+
16
+ from eva_clip import create_model_and_transforms
17
+ from eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
18
+ from toonmage.encoders_flux import IDFormer, PerceiverAttentionCA
19
+
20
+
21
+ class ToonMagePipeline(nn.Module):
22
+ def __init__(self, dit, device, weight_dtype=torch.bfloat16, *args, **kwargs):
23
+ super().__init__()
24
+ self.device = device
25
+ self.weight_dtype = weight_dtype
26
+ double_interval = 2
27
+ single_interval = 4
28
+
29
+ # init encoder
30
+ self.toonmage_encoder = IDFormer().to(self.device, self.weight_dtype)
31
+
32
+ num_ca = 19 // double_interval + 38 // single_interval
33
+ if 19 % double_interval != 0:
34
+ num_ca += 1
35
+ if 38 % single_interval != 0:
36
+ num_ca += 1
37
+ self.toonmage_ca = nn.ModuleList([
38
+ PerceiverAttentionCA().to(self.device, self.weight_dtype) for _ in range(num_ca)
39
+ ])
40
+
41
+ dit.toonmage_ca = self.toonmage_ca
42
+ dit.toonmage_double_interval = double_interval
43
+ dit.toonmage_single_interval = single_interval
44
+
45
+ # preprocessors
46
+ # face align and parsing
47
+ self.face_helper = FaceRestoreHelper(
48
+ upscale_factor=1,
49
+ face_size=512,
50
+ crop_ratio=(1, 1),
51
+ det_model='retinaface_resnet50',
52
+ save_ext='png',
53
+ device=self.device,
54
+ )
55
+ self.face_helper.face_parse = None
56
+ self.face_helper.face_parse = init_parsing_model(model_name='bisenet', device=self.device)
57
+ # clip-vit backbone
58
+ model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', 'eva_clip', force_custom_clip=True)
59
+ model = model.visual
60
+ self.clip_vision_model = model.to(self.device, dtype=self.weight_dtype)
61
+ eva_transform_mean = getattr(self.clip_vision_model, 'image_mean', OPENAI_DATASET_MEAN)
62
+ eva_transform_std = getattr(self.clip_vision_model, 'image_std', OPENAI_DATASET_STD)
63
+ if not isinstance(eva_transform_mean, (list, tuple)):
64
+ eva_transform_mean = (eva_transform_mean,) * 3
65
+ if not isinstance(eva_transform_std, (list, tuple)):
66
+ eva_transform_std = (eva_transform_std,) * 3
67
+ self.eva_transform_mean = eva_transform_mean
68
+ self.eva_transform_std = eva_transform_std
69
+ # antelopev2
70
+ snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2')
71
+ self.app = FaceAnalysis(
72
+ name='antelopev2', root='.', providers=['CPUExecutionProvider']
73
+ )
74
+ self.app.prepare(ctx_id=0, det_size=(640, 640))
75
+ self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx', providers=['CPUExecutionProvider'])
76
+ self.handler_ante.prepare(ctx_id=0)
77
+
78
+ gc.collect()
79
+ torch.cuda.empty_cache()
80
+
81
+ # self.load_pretrain()
82
+
83
+ # other configs
84
+ self.debug_img_list = []
85
+
86
+ def load_pretrain(self, pretrain_path=None):
87
+ hf_hub_download('SunderAli17/SAK', 'toonmage_flux_v2.safetensors', local_dir='models')
88
+ ckpt_path = 'models/toonmage_flux_v2.safetensors'
89
+ if pretrain_path is not None:
90
+ ckpt_path = pretrain_path
91
+ state_dict = load_file(ckpt_path)
92
+ state_dict_dict = {}
93
+ for k, v in state_dict.items():
94
+ module = k.split('.')[0]
95
+ state_dict_dict.setdefault(module, {})
96
+ new_k = k[len(module) + 1:]
97
+ state_dict_dict[module][new_k] = v
98
+
99
+ for module in state_dict_dict:
100
+ print(f'loading from {module}')
101
+ getattr(self, module).load_state_dict(state_dict_dict[module], strict=True)
102
+
103
+ del state_dict
104
+ del state_dict_dict
105
+
106
+ def to_gray(self, img):
107
+ x = 0.299 * img[:, 0:1] + 0.587 * img[:, 1:2] + 0.114 * img[:, 2:3]
108
+ x = x.repeat(1, 3, 1, 1)
109
+ return x
110
+
111
+ def get_id_embedding(self, image, cal_uncond=False):
112
+ """
113
+ Args:
114
+ image: numpy rgb image, range [0, 255]
115
+ """
116
+ self.face_helper.clean_all()
117
+ self.debug_img_list = []
118
+ image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
119
+ # get antelopev2 embedding
120
+ # for k in self.app.models.keys():
121
+ # self.app.models[k].session.set_providers(['CUDAExecutionProvider'])
122
+ face_info = self.app.get(image_bgr)
123
+ if len(face_info) > 0:
124
+ face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[
125
+ -1
126
+ ] # only use the maximum face
127
+ id_ante_embedding = face_info['embedding']
128
+ self.debug_img_list.append(
129
+ image[
130
+ int(face_info['bbox'][1]) : int(face_info['bbox'][3]),
131
+ int(face_info['bbox'][0]) : int(face_info['bbox'][2]),
132
+ ]
133
+ )
134
+ else:
135
+ id_ante_embedding = None
136
+
137
+ # using facexlib to detect and align face
138
+ self.face_helper.read_image(image_bgr)
139
+ self.face_helper.get_face_landmarks_5(only_center_face=True)
140
+ self.face_helper.align_warp_face()
141
+ if len(self.face_helper.cropped_faces) == 0:
142
+ raise RuntimeError('facexlib align face fail')
143
+ align_face = self.face_helper.cropped_faces[0]
144
+ # incase insightface didn't detect face
145
+ if id_ante_embedding is None:
146
+ print('fail to detect face using insightface, extract embedding on align face')
147
+ # self.handler_ante.session.set_providers(['CUDAExecutionProvider'])
148
+ id_ante_embedding = self.handler_ante.get_feat(align_face)
149
+
150
+ id_ante_embedding = torch.from_numpy(id_ante_embedding).to(self.device, self.weight_dtype)
151
+ if id_ante_embedding.ndim == 1:
152
+ id_ante_embedding = id_ante_embedding.unsqueeze(0)
153
+
154
+ # parsing
155
+ input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
156
+ input = input.to(self.device)
157
+ parsing_out = self.face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
158
+ parsing_out = parsing_out.argmax(dim=1, keepdim=True)
159
+ bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
160
+ bg = sum(parsing_out == i for i in bg_label).bool()
161
+ white_image = torch.ones_like(input)
162
+ # only keep the face features
163
+ face_features_image = torch.where(bg, white_image, self.to_gray(input))
164
+ self.debug_img_list.append(tensor2img(face_features_image, rgb2bgr=False))
165
+
166
+ # transform img before sending to eva-clip-vit
167
+ face_features_image = resize(face_features_image, self.clip_vision_model.image_size, InterpolationMode.BICUBIC)
168
+ face_features_image = normalize(face_features_image, self.eva_transform_mean, self.eva_transform_std)
169
+ id_cond_vit, id_vit_hidden = self.clip_vision_model(
170
+ face_features_image.to(self.weight_dtype), return_all_features=False, return_hidden=True, shuffle=False
171
+ )
172
+ id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True)
173
+ id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm)
174
+
175
+ id_cond = torch.cat([id_ante_embedding, id_cond_vit], dim=-1)
176
+
177
+ id_embedding = self.toonmage_encoder(id_cond, id_vit_hidden)
178
+
179
+ if not cal_uncond:
180
+ return id_embedding, None
181
+
182
+ id_uncond = torch.zeros_like(id_cond)
183
+ id_vit_hidden_uncond = []
184
+ for layer_idx in range(0, len(id_vit_hidden)):
185
+ id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden[layer_idx]))
186
+ uncond_id_embedding = self.toonmage_encoder(id_uncond, id_vit_hidden_uncond)
187
+
188
+ return id_embedding, uncond_id_embedding
toonmage/pipeline.py ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gc
2
+
3
+ import cv2
4
+ import insightface
5
+ import torch
6
+ import torch.nn as nn
7
+ from basicsr.utils import img2tensor, tensor2img
8
+ from diffusers import (
9
+ DPMSolverMultistepScheduler,
10
+ StableDiffusionXLPipeline,
11
+ UNet2DConditionModel,
12
+ )
13
+ from facexlib.parsing import init_parsing_model
14
+ from facexlib.utils.face_restoration_helper import FaceRestoreHelper
15
+ from huggingface_hub import hf_hub_download, snapshot_download
16
+ from insightface.app import FaceAnalysis
17
+ from safetensors.torch import load_file
18
+ from torchvision.transforms import InterpolationMode
19
+ from torchvision.transforms.functional import normalize, resize
20
+
21
+ from eva_clip import create_model_and_transforms
22
+ from eva_clip.constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
23
+ from toonmage.encoders import IDEncoder
24
+ from toonmage.utils import is_torch2_available
25
+
26
+ if is_torch2_available():
27
+ from toonmage.attention_processor import AttnProcessor2_0 as AttnProcessor
28
+ from toonmage.attention_processor import IDAttnProcessor2_0 as IDAttnProcessor
29
+ else:
30
+ from toonmage.attention_processor import AttnProcessor, IDAttnProcessor
31
+
32
+
33
+ class ToonMagePipeline:
34
+ def __init__(self, *args, **kwargs):
35
+ super().__init__()
36
+ self.device = 'cuda'
37
+ sdxl_base_repo = 'stabilityai/stable-diffusion-xl-base-1.0'
38
+ sdxl_lightning_repo = 'ByteDance/SDXL-Lightning'
39
+ self.sdxl_base_repo = sdxl_base_repo
40
+
41
+ # load base model
42
+ unet = UNet2DConditionModel.from_config(sdxl_base_repo, subfolder='unet').to(self.device, torch.float16)
43
+ unet.load_state_dict(
44
+ load_file(
45
+ hf_hub_download(sdxl_lightning_repo, 'sdxl_lightning_4step_unet.safetensors'), device=self.device
46
+ )
47
+ )
48
+ unet.half()
49
+ self.hack_unet_attn_layers(unet)
50
+ self.pipe = StableDiffusionXLPipeline.from_pretrained(
51
+ sdxl_base_repo, unet=unet, torch_dtype=torch.float16, variant="fp16"
52
+ ).to(self.device)
53
+ self.pipe.watermark = None
54
+
55
+ # scheduler
56
+ self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
57
+ self.pipe.scheduler.config, timestep_spacing="trailing"
58
+ )
59
+
60
+ # ID adapters
61
+ self.id_adapter = IDEncoder().to(self.device)
62
+
63
+ # preprocessors
64
+ # face align and parsing
65
+ self.face_helper = FaceRestoreHelper(
66
+ upscale_factor=1,
67
+ face_size=512,
68
+ crop_ratio=(1, 1),
69
+ det_model='retinaface_resnet50',
70
+ save_ext='png',
71
+ device=self.device,
72
+ )
73
+ self.face_helper.face_parse = None
74
+ self.face_helper.face_parse = init_parsing_model(model_name='bisenet', device=self.device)
75
+ # clip-vit backbone
76
+ model, _, _ = create_model_and_transforms('EVA02-CLIP-L-14-336', 'eva_clip', force_custom_clip=True)
77
+ model = model.visual
78
+ self.clip_vision_model = model.to(self.device)
79
+ eva_transform_mean = getattr(self.clip_vision_model, 'image_mean', OPENAI_DATASET_MEAN)
80
+ eva_transform_std = getattr(self.clip_vision_model, 'image_std', OPENAI_DATASET_STD)
81
+ if not isinstance(eva_transform_mean, (list, tuple)):
82
+ eva_transform_mean = (eva_transform_mean,) * 3
83
+ if not isinstance(eva_transform_std, (list, tuple)):
84
+ eva_transform_std = (eva_transform_std,) * 3
85
+ self.eva_transform_mean = eva_transform_mean
86
+ self.eva_transform_std = eva_transform_std
87
+ # antelopev2
88
+ snapshot_download('DIAMONIK7777/antelopev2', local_dir='models/antelopev2')
89
+ self.app = FaceAnalysis(
90
+ name='antelopev2', root='.', providers=['CPUExecutionProvider']
91
+ )
92
+ self.app.prepare(ctx_id=0, det_size=(640, 640))
93
+ self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx', providers=['CPUExecutionProvider'])
94
+ self.handler_ante.prepare(ctx_id=0)
95
+
96
+ print('load done')
97
+
98
+ gc.collect()
99
+ torch.cuda.empty_cache()
100
+
101
+ self.load_pretrain()
102
+
103
+ # other configs
104
+ self.debug_img_list = []
105
+
106
+ def hack_unet_attn_layers(self, unet):
107
+ id_adapter_attn_procs = {}
108
+ for name, _ in unet.attn_processors.items():
109
+ cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
110
+ if name.startswith("mid_block"):
111
+ hidden_size = unet.config.block_out_channels[-1]
112
+ elif name.startswith("up_blocks"):
113
+ block_id = int(name[len("up_blocks.")])
114
+ hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
115
+ elif name.startswith("down_blocks"):
116
+ block_id = int(name[len("down_blocks.")])
117
+ hidden_size = unet.config.block_out_channels[block_id]
118
+ if cross_attention_dim is not None:
119
+ id_adapter_attn_procs[name] = IDAttnProcessor(
120
+ hidden_size=hidden_size,
121
+ cross_attention_dim=cross_attention_dim,
122
+ ).to(unet.device)
123
+ else:
124
+ id_adapter_attn_procs[name] = AttnProcessor()
125
+ unet.set_attn_processor(id_adapter_attn_procs)
126
+ self.id_adapter_attn_layers = nn.ModuleList(unet.attn_processors.values())
127
+
128
+ def load_pretrain(self):
129
+ hf_hub_download('SunderAli17/SAK', 'toonmage_v2.bin', local_dir='models')
130
+ ckpt_path = 'models/toonmage_v2.bin'
131
+ state_dict = torch.load(ckpt_path, map_location='cpu')
132
+ state_dict_dict = {}
133
+ for k, v in state_dict.items():
134
+ module = k.split('.')[0]
135
+ state_dict_dict.setdefault(module, {})
136
+ new_k = k[len(module) + 1 :]
137
+ state_dict_dict[module][new_k] = v
138
+
139
+ for module in state_dict_dict:
140
+ print(f'loading from {module}')
141
+ getattr(self, module).load_state_dict(state_dict_dict[module], strict=True)
142
+
143
+ def to_gray(self, img):
144
+ x = 0.299 * img[:, 0:1] + 0.587 * img[:, 1:2] + 0.114 * img[:, 2:3]
145
+ x = x.repeat(1, 3, 1, 1)
146
+ return x
147
+
148
+ def get_id_embedding(self, image):
149
+ """
150
+ Args:
151
+ image: numpy rgb image, range [0, 255]
152
+ """
153
+ self.face_helper.clean_all()
154
+ image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
155
+ # get antelopev2 embedding
156
+ face_info = self.app.get(image_bgr)
157
+ if len(face_info) > 0:
158
+ face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * x['bbox'][3] - x['bbox'][1])[
159
+ -1
160
+ ] # only use the maximum face
161
+ id_ante_embedding = face_info['embedding']
162
+ self.debug_img_list.append(
163
+ image[
164
+ int(face_info['bbox'][1]) : int(face_info['bbox'][3]),
165
+ int(face_info['bbox'][0]) : int(face_info['bbox'][2]),
166
+ ]
167
+ )
168
+ else:
169
+ id_ante_embedding = None
170
+
171
+ # using facexlib to detect and align face
172
+ self.face_helper.read_image(image_bgr)
173
+ self.face_helper.get_face_landmarks_5(only_center_face=True)
174
+ self.face_helper.align_warp_face()
175
+ if len(self.face_helper.cropped_faces) == 0:
176
+ raise RuntimeError('facexlib align face fail')
177
+ align_face = self.face_helper.cropped_faces[0]
178
+ # incase insightface didn't detect face
179
+ if id_ante_embedding is None:
180
+ print('fail to detect face using insightface, extract embedding on align face')
181
+ id_ante_embedding = self.handler_ante.get_feat(align_face)
182
+
183
+ id_ante_embedding = torch.from_numpy(id_ante_embedding).to(self.device)
184
+ if id_ante_embedding.ndim == 1:
185
+ id_ante_embedding = id_ante_embedding.unsqueeze(0)
186
+
187
+ # parsing
188
+ input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
189
+ input = input.to(self.device)
190
+ parsing_out = self.face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
191
+ parsing_out = parsing_out.argmax(dim=1, keepdim=True)
192
+ bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
193
+ bg = sum(parsing_out == i for i in bg_label).bool()
194
+ white_image = torch.ones_like(input)
195
+ # only keep the face features
196
+ face_features_image = torch.where(bg, white_image, self.to_gray(input))
197
+ self.debug_img_list.append(tensor2img(face_features_image, rgb2bgr=False))
198
+
199
+ # transform img before sending to eva-clip-vit
200
+ face_features_image = resize(face_features_image, self.clip_vision_model.image_size, InterpolationMode.BICUBIC)
201
+ face_features_image = normalize(face_features_image, self.eva_transform_mean, self.eva_transform_std)
202
+ id_cond_vit, id_vit_hidden = self.clip_vision_model(
203
+ face_features_image, return_all_features=False, return_hidden=True, shuffle=False
204
+ )
205
+ id_cond_vit_norm = torch.norm(id_cond_vit, 2, 1, True)
206
+ id_cond_vit = torch.div(id_cond_vit, id_cond_vit_norm)
207
+
208
+ id_cond = torch.cat([id_ante_embedding, id_cond_vit], dim=-1)
209
+ id_uncond = torch.zeros_like(id_cond)
210
+ id_vit_hidden_uncond = []
211
+ for layer_idx in range(0, len(id_vit_hidden)):
212
+ id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden[layer_idx]))
213
+
214
+ id_embedding = self.id_adapter(id_cond, id_vit_hidden)
215
+ uncond_id_embedding = self.id_adapter(id_uncond, id_vit_hidden_uncond)
216
+
217
+ # return id_embedding
218
+ return torch.cat((uncond_id_embedding, id_embedding), dim=0)
219
+
220
+ def inference(self, prompt, size, prompt_n='', image_embedding=None, id_scale=1.0, guidance_scale=1.2, steps=4):
221
+ images = self.pipe(
222
+ prompt=prompt,
223
+ negative_prompt=prompt_n,
224
+ num_images_per_prompt=size[0],
225
+ height=size[1],
226
+ width=size[2],
227
+ num_inference_steps=steps,
228
+ guidance_scale=guidance_scale,
229
+ cross_attention_kwargs={'id_embedding': image_embedding, 'id_scale': id_scale},
230
+ ).images
231
+
232
+ return images
toonmage/utils.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import importlib
2
+ import os
3
+ import random
4
+
5
+ import cv2
6
+ import numpy as np
7
+ import torch
8
+ import torch.nn.functional as F
9
+ from transformers import PretrainedConfig
10
+
11
+
12
+ def seed_everything(seed):
13
+ os.environ["PL_GLOBAL_SEED"] = str(seed)
14
+ random.seed(seed)
15
+ np.random.seed(seed)
16
+ torch.manual_seed(seed)
17
+ torch.cuda.manual_seed_all(seed)
18
+
19
+
20
+ def is_torch2_available():
21
+ return hasattr(F, "scaled_dot_product_attention")
22
+
23
+
24
+ def instantiate_from_config(config):
25
+ if "target" not in config:
26
+ if config == '__is_first_stage__' or config == "__is_unconditional__":
27
+ return None
28
+ raise KeyError("Expected key `target` to instantiate.")
29
+ return get_obj_from_str(config["target"])(**config.get("params", {}))
30
+
31
+
32
+ def get_obj_from_str(string, reload=False):
33
+ module, cls = string.rsplit(".", 1)
34
+ if reload:
35
+ module_imp = importlib.import_module(module)
36
+ importlib.reload(module_imp)
37
+ return getattr(importlib.import_module(module, package=None), cls)
38
+
39
+
40
+ def drop_seq_token(seq, drop_rate=0.5):
41
+ idx = torch.randperm(seq.size(1))
42
+ num_keep_tokens = int(len(idx) * (1 - drop_rate))
43
+ idx = idx[:num_keep_tokens]
44
+ seq = seq[:, idx]
45
+ return seq
46
+
47
+
48
+ def import_model_class_from_model_name_or_path(
49
+ pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder"
50
+ ):
51
+ text_encoder_config = PretrainedConfig.from_pretrained(
52
+ pretrained_model_name_or_path, subfolder=subfolder, revision=revision
53
+ )
54
+ model_class = text_encoder_config.architectures[0]
55
+
56
+ if model_class == "CLIPTextModel":
57
+ from transformers import CLIPTextModel
58
+
59
+ return CLIPTextModel
60
+ elif model_class == "CLIPTextModelWithProjection": # noqa RET505
61
+ from transformers import CLIPTextModelWithProjection
62
+
63
+ return CLIPTextModelWithProjection
64
+ else:
65
+ raise ValueError(f"{model_class} is not supported.")
66
+
67
+
68
+ def resize_numpy_image_long(image, resize_long_edge=768):
69
+ h, w = image.shape[:2]
70
+ if max(h, w) <= resize_long_edge:
71
+ return image
72
+ k = resize_long_edge / max(h, w)
73
+ h = int(h * k)
74
+ w = int(w * k)
75
+ image = cv2.resize(image, (w, h), interpolation=cv2.INTER_LANCZOS4)
76
+ return image