| import os,sys |
| import PIL |
| from PIL import Image |
| import numpy as np |
|
|
| import torch |
| import torchvision.transforms as tvtrans |
| from .lib.cfg_helper import model_cfg_bank |
| from .lib.model_zoo import get_model |
| from random import randint |
| from .lib.model_zoo.ddim import DDIMSampler |
| |
|
|
| def highlight_print(info): |
| print('') |
| print(''.join(['#']*(len(info)+4))) |
| print('# '+info+' #') |
| print(''.join(['#']*(len(info)+4))) |
| print('') |
|
|
| def decompose(x, q=20, niter=100): |
| x_mean = x.mean(-1, keepdim=True) |
| x_input = x - x_mean |
| u, s, v = torch.pca_lowrank(x_input, q=q, center=False, niter=niter) |
| ss = torch.stack([torch.diag(si) for si in s]) |
| x_lowrank = torch.bmm(torch.bmm(u, ss), torch.permute(v, [0, 2, 1])) |
| x_remain = x_input - x_lowrank |
| return u, s, v, x_mean, x_remain |
|
|
| class adjust_rank(object): |
| def __init__(self, max_drop_rank=[1, 5], q=20): |
| self.max_semantic_drop_rank = max_drop_rank[0] |
| self.max_style_drop_rank = max_drop_rank[1] |
| self.q = q |
|
|
| def t2y0_semf_wrapper(t0, y00, t1, y01): |
| return lambda t: (np.exp((t-0.5)*2)-t0)/(t1-t0)*(y01-y00)+y00 |
| t0, y00 = np.exp((0 -0.5)*2), -self.max_semantic_drop_rank |
| t1, y01 = np.exp((0.5-0.5)*2), 1 |
| self.t2y0_semf = t2y0_semf_wrapper(t0, y00, t1, y01) |
|
|
| def x2y_semf_wrapper(x0, x1, y1): |
| return lambda x, y0: (x-x0)/(x1-x0)*(y1-y0)+y0 |
| x0 = 0 |
| x1, y1 = self.max_semantic_drop_rank+1, 1 |
| self.x2y_semf = x2y_semf_wrapper(x0, x1, y1) |
| |
| def t2y0_styf_wrapper(t0, y00, t1, y01): |
| return lambda t: (np.exp((t-0.5)*2)-t0)/(t1-t0)*(y01-y00)+y00 |
| t0, y00 = np.exp((1 -0.5)*2), -(q-self.max_style_drop_rank) |
| t1, y01 = np.exp((0.5-0.5)*2), 1 |
| self.t2y0_styf = t2y0_styf_wrapper(t0, y00, t1, y01) |
|
|
| def x2y_styf_wrapper(x0, x1, y1): |
| return lambda x, y0: (x-x0)/(x1-x0)*(y1-y0)+y0 |
| x0 = q-1 |
| x1, y1 = self.max_style_drop_rank-1, 1 |
| self.x2y_styf = x2y_styf_wrapper(x0, x1, y1) |
|
|
| def __call__(self, x, lvl): |
| if lvl == 0.5: |
| return x |
|
|
| if x.dtype == torch.float16: |
| fp16 = True |
| x = x.float() |
| else: |
| fp16 = False |
| std_save = x.std(axis=[-2, -1]) |
|
|
| u, s, v, x_mean, x_remain = decompose(x, q=self.q) |
|
|
| if lvl < 0.5: |
| assert lvl>=0 |
| for xi in range(0, self.max_semantic_drop_rank+1): |
| y0 = self.t2y0_semf(lvl) |
| yi = self.x2y_semf(xi, y0) |
| yi = 0 if yi<0 else yi |
| s[:, xi] *= yi |
|
|
| elif lvl > 0.5: |
| assert lvl <= 1 |
| for xi in range(self.max_style_drop_rank, self.q): |
| y0 = self.t2y0_styf(lvl) |
| yi = self.x2y_styf(xi, y0) |
| yi = 0 if yi<0 else yi |
| s[:, xi] *= yi |
| x_remain = 0 |
|
|
| ss = torch.stack([torch.diag(si) for si in s]) |
| x_lowrank = torch.bmm(torch.bmm(u, ss), torch.permute(v, [0, 2, 1])) |
| x_new = x_lowrank + x_mean + x_remain |
|
|
| std_new = x_new.std(axis=[-2, -1]) |
| x_new = x_new / std_new * std_save |
|
|
| if fp16: |
| x_new = x_new.half() |
|
|
| return x_new |
|
|
| class Reconstructor(object): |
| def __init__(self, fp16=True, device="cuda:0", cache_dir="../cache", ddim_steps=50, deprecated=False): |
| print(f"Reconstructor: Loading model... fp16: {fp16}") |
| print("Taking new code 2.") |
| if deprecated: |
| cfgm_name = 'vd_noema' |
| else: |
| cfgm_name = 'vd_four_flow_v1-0' |
| |
| cfgm = model_cfg_bank()(cfgm_name) |
| cfgm['args']['vae_cfg_list'][0][1]['pth'] = f'{cache_dir}/kl-f8.pth' |
| cfgm['args']['vae_cfg_list'][1][1]['pth'] =f'{cache_dir}/optimus-vae.pth' |
| net = get_model()(cfgm) |
|
|
| if fp16: |
| net.ctx['text'].fp16 = True |
| net.ctx['image'].fp16 = True |
| net = net.half() |
| self.dtype = torch.float16 |
| if deprecated: |
| sd = torch.load(f'{cache_dir}/vd-four-flow-v1-0-fp16-deprecated.pth', map_location='cpu') |
| else: |
| sd = torch.load(f'{cache_dir}/vd-four-flow-v1-0-fp16.pth', map_location='cpu') |
| else: |
| self.dtype = torch.float32 |
| sd = torch.load(f'{cache_dir}/vd-four-flow-v1-0.pth', map_location='cpu') |
| |
| self.device=device |
| self.output_dim = [512, 512] |
| self.ddim_steps = ddim_steps |
| self.ddim_eta = 0.0 |
| self.image_latent_dim = 4 |
| |
| net.load_state_dict(sd, strict=False) |
| self.sampler = DDIMSampler(net) |
| self.sampler.make_schedule(ddim_num_steps=self.ddim_steps, ddim_eta=self.ddim_eta, verbose=False) |
|
|
| net.to(self.device) |
| self.net = net |
| self.adjust_rank_f = adjust_rank(max_drop_rank=[1, 5], q=20) |
| self.scale = 3.5 |
| self.disentanglement_noglobal = True |
| |
| def embed_text(self, prompt): |
| if isinstance(prompt, str): |
| prompt = [prompt] |
| text_encoding = self.net.ctx_encode(prompt, which='text') |
| return text_encoding |
| |
| def embed_image(self, image): |
| if isinstance(image, PIL.Image.Image): |
| image = tvtrans.ToTensor()(image) |
| |
| if image.ndim == 3: |
| image = image.unsqueeze(0) |
| |
| image_encoding = self.net.ctx_encode(image, which='image') |
| return image_encoding |
| |
| def project_clip(self, expanded_clip): |
| reduced_clip = expanded_clip[:, 0, :] |
| reduced_clip = reduced_clip * torch.norm(reduced_clip, dim=-1, keepdim=True) |
| print("RECONSTRUCTOR REDUCED CLIP SHAPE: ", reduced_clip.shape) |
| projected_clip = self.net.ctx["image"].model.visual_projection(reduced_clip) |
| return projected_clip |
| |
| def reconstruct(self, |
| image=None, |
| c_i=None, |
| c_t=None, |
| n_samples=1, |
| textstrength=0.5, |
| strength=1.0, |
| color_adjust=False, |
| fcs_lvl=0.5, |
| seed=None |
| ): |
| |
| numClips =0 |
| h, w = 512, 512 |
| BICUBIC = PIL.Image.Resampling.BICUBIC |
| |
| if strength == 0: |
| return [image]*n_samples |
| else: |
| assert (c_t is not None) or (c_i is not None) |
| c_info_list = [] |
| scale = self.scale |
| if c_t is not None and textstrength != 0: |
| c_t = c_t.reshape((77,768)).to(dtype=torch.float16, device=self.device) |
| ut = self.net.ctx_encode([""], which='text').repeat(n_samples, 1, 1) |
| ct = c_t.repeat(n_samples, 1, 1) |
| print(ct.shape) |
| c_info_list.append({ |
| 'type':'text', |
| 'conditioning':ct.to(torch.float16), |
| 'unconditional_conditioning':ut, |
| 'unconditional_guidance_scale':scale, |
| 'ratio': textstrength, }) |
| numClips +=1 |
| else: |
| textstrength=0 |
|
|
| if c_i is not None and textstrength != 1: |
| c_i = c_i.reshape((257,768)).to(dtype=torch.float16, device=self.device) |
| ci = c_i |
|
|
| if self.disentanglement_noglobal: |
| ci_glb = ci[:, 0:1] |
| ci_loc = ci[:, 1: ] |
| ci_loc = self.adjust_rank_f(ci_loc, fcs_lvl) |
| ci = torch.cat([ci_glb, ci_loc], dim=1).repeat(n_samples, 1, 1) |
| else: |
| ci = self.adjust_rank_f(ci, fcs_lvl).repeat(n_samples, 1, 1) |
|
|
| c_info_list.append({ |
| 'type':'image', |
| 'conditioning':ci.to(torch.float16), |
| 'unconditional_conditioning':torch.zeros_like(ci), |
| 'unconditional_guidance_scale':scale, |
| 'ratio': (1-textstrength), }) |
| numClips +=1 |
| else: |
| textstrength=1 |
| if(image is not None): |
| image_tensor = tvtrans.Compose([ |
| tvtrans.ToTensor(), |
| tvtrans.Resize((w, h)) |
| ])(image).to(self.device).to(self.dtype) |
| if image_tensor.ndim == 3: |
| image_tensor = image_tensor.unsqueeze(0) |
| shape = [n_samples, self.image_latent_dim, h//8, w//8] |
| if(seed): |
| np.random.seed(seed) |
| torch.manual_seed(seed + 100) |
| else: |
| seed = randint(0,1000) |
| np.random.seed(seed) |
| torch.manual_seed(seed + 100) |
| if strength!=1 and image: |
| x0 = self.net.vae_encode(image_tensor, which='image').repeat(n_samples, 1, 1, 1) |
| step = int(self.ddim_steps * (strength)) |
| if numClips==2: |
| x, _ = self.sampler.sample_multicontext( |
| steps=self.ddim_steps, |
| x_info={'type':'image', 'x0':x0, 'x0_forward_timesteps':step}, |
| c_info_list=c_info_list, |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| else: |
| x, _ = self.sampler.sample( |
| steps=self.ddim_steps, |
| x_info={'type':'image', 'x0':x0, 'x0_forward_timesteps':step}, |
| c_info=c_info_list[0], |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| else: |
| if numClips ==2: |
| x, _ = self.sampler.sample_multicontext( |
| steps=self.ddim_steps, |
| x_info={'type':'image',}, |
| c_info_list=c_info_list, |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| else: |
| x, _ = self.sampler.sample( |
| steps=self.ddim_steps, |
| x_info={'type':'image',}, |
| c_info=c_info_list[0], |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| imout = self.net.vae_decode(x, which='image') |
| if color_adjust: |
| cx_mean = image_tensor.view(3, -1).mean(-1)[:, None, None] |
| cx_std = image_tensor.view(3, -1).std(-1)[:, None, None] |
| imout_mean = [imouti.view(3, -1).mean(-1)[:, None, None] for imouti in imout] |
| imout_std = [imouti.view(3, -1).std(-1)[:, None, None] for imouti in imout] |
| imout = [(ii-mi)/si*cx_std+cx_mean for ii, mi, si in zip(imout, imout_mean, imout_std)] |
| imout = [torch.clamp(ii, 0, 1) for ii in imout] |
| imout = [tvtrans.ToPILImage()(i) for i in imout] |
| if len(imout)==1: |
| return imout[0] |
| else: |
| return imout |
|
|
| def reconstruct_batch(self, |
| image=None, |
| c_i=None, |
| c_t=None, |
| textstrength=0.5, |
| strength=1.0, |
| color_adjust=False, |
| fcs_lvl=0.5, |
| seed=None |
| ): |
| n_samples = c_i.shape[0] if c_i is not None else c_t.shape[0] |
| if (c_i is not None) and (c_t is not None): |
| assert (len(c_i) == len(c_t)), "Make sure the batch size of your clip text and clip image are the same" |
| numClips =0 |
| h, w = 512, 512 |
| BICUBIC = PIL.Image.Resampling.BICUBIC |
| |
| if strength == 0: |
| return [image]*n_samples |
| else: |
| assert (c_t is not None) or (c_i is not None) |
| c_info_list = [] |
| scale = self.scale |
| if c_t is not None and textstrength != 0: |
| c_t = c_t.to(dtype=torch.float16, device=self.device) |
| ut = self.net.ctx_encode([""], which='text').repeat(n_samples, 1, 1) |
| ct = c_t |
| c_info_list.append({ |
| 'type':'text', |
| 'conditioning':ct.to(torch.float16), |
| 'unconditional_conditioning':ut, |
| 'unconditional_guidance_scale':scale, |
| 'ratio': textstrength, }) |
| numClips +=1 |
| else: |
| textstrength=0 |
|
|
| if c_i is not None and textstrength != 1: |
| c_i = c_i.to(dtype=torch.float16, device=self.device) |
| |
| if self.disentanglement_noglobal: |
| ci_final = torch.empty(c_i.shape, dtype=torch.float16, device=self.device) |
| for i in range(len(c_i)): |
| ci = c_i[i] |
| ci_glb = ci[:, 0:1] |
| ci_loc = ci[:, 1: ] |
| ci_loc = self.adjust_rank_f(ci_loc, fcs_lvl) |
| ci = torch.cat([ci_glb, ci_loc], dim=1) |
| ci_final[i,:,:] = ci |
| else: |
| ci_final = torch.empty(c_i.shape, dtype=torch.float16, device=self.device) |
| for i in range(len(c_i)): |
| ci = c_i[i] |
| ci = self.adjust_rank_f(ci, fcs_lvl) |
| ci_final[i,:,:] = ci |
|
|
| c_info_list.append({ |
| 'type':'image', |
| 'conditioning':ci_final.to(torch.float16), |
| 'unconditional_conditioning':torch.zeros_like(ci_final), |
| 'unconditional_guidance_scale':scale, |
| 'ratio': (1-textstrength), }) |
| numClips +=1 |
| else: |
| textstrength=1 |
| if(image is not None): |
| image_tensor = tvtrans.Compose([ |
| tvtrans.ToTensor(), |
| tvtrans.Resize((w, h)) |
| ])(image).to(self.device).to(self.dtype) |
| if image_tensor.ndim == 3: |
| image_tensor = image_tensor.unsqueeze(0) |
| shape = [n_samples, self.image_latent_dim, h//8, w//8] |
| if(seed): |
| np.random.seed(seed) |
| torch.manual_seed(seed + 100) |
| else: |
| seed = randint(0,1000) |
| np.random.seed(seed) |
| torch.manual_seed(seed + 100) |
| if strength!=1 and image: |
| x0 = self.net.vae_encode(image_tensor, which='image').repeat(n_samples, 1, 1, 1) |
| step = int(self.ddim_steps * (strength)) |
| if numClips==2: |
| x, _ = self.sampler.sample_multicontext( |
| steps=self.ddim_steps, |
| x_info={'type':'image', 'x0':x0, 'x0_forward_timesteps':step}, |
| c_info_list=c_info_list, |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| else: |
| x, _ = self.sampler.sample( |
| steps=self.ddim_steps, |
| x_info={'type':'image', 'x0':x0, 'x0_forward_timesteps':step}, |
| c_info=c_info_list[0], |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| else: |
| if numClips ==2: |
| x, _ = self.sampler.sample_multicontext( |
| steps=self.ddim_steps, |
| x_info={'type':'image',}, |
| c_info_list=c_info_list, |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| else: |
| x, _ = self.sampler.sample( |
| steps=self.ddim_steps, |
| x_info={'type':'image',}, |
| c_info=c_info_list[0], |
| shape=shape, |
| verbose=False, |
| eta=self.ddim_eta) |
| imout = self.net.vae_decode(x, which='image') |
| if color_adjust: |
| cx_mean = image_tensor.view(3, -1).mean(-1)[:, None, None] |
| cx_std = image_tensor.view(3, -1).std(-1)[:, None, None] |
| imout_mean = [imouti.view(3, -1).mean(-1)[:, None, None] for imouti in imout] |
| imout_std = [imouti.view(3, -1).std(-1)[:, None, None] for imouti in imout] |
| imout = [(ii-mi)/si*cx_std+cx_mean for ii, mi, si in zip(imout, imout_mean, imout_std)] |
| imout = [torch.clamp(ii, 0, 1) for ii in imout] |
| imout = [tvtrans.ToPILImage()(i) for i in imout] |
| if len(imout)==1: |
| return imout[0] |
| else: |
| return imout |
