|
|
from typing import List |
|
|
import torch |
|
|
from torchvision import transforms |
|
|
from transformers import CLIPImageProcessor |
|
|
from transformers import CLIPVisionModel as OriginalCLIPVisionModel |
|
|
from ._clip import CLIPVisionModel |
|
|
from PIL import Image |
|
|
import torch.nn.functional as F |
|
|
import torch.nn as nn |
|
|
import os |
|
|
|
|
|
def is_torch2_available(): |
|
|
return hasattr(F, "scaled_dot_product_attention") |
|
|
if is_torch2_available(): |
|
|
from .attention_processor import SSRAttnProcessor2_0 as SSRAttnProcessor, AttnProcessor2_0 as AttnProcessor |
|
|
else: |
|
|
from .attention_processor import SSRAttnProcessor, AttnProcessor |
|
|
from .resampler import Resampler |
|
|
|
|
|
class detail_encoder(torch.nn.Module): |
|
|
"""from SSR-encoder""" |
|
|
def __init__(self, unet, image_encoder_path, device="cuda", dtype=torch.float32): |
|
|
super().__init__() |
|
|
self.device = device |
|
|
self.dtype = dtype |
|
|
|
|
|
|
|
|
clip_encoder = OriginalCLIPVisionModel.from_pretrained(image_encoder_path) |
|
|
self.image_encoder = CLIPVisionModel(clip_encoder.config) |
|
|
state_dict = clip_encoder.state_dict() |
|
|
self.image_encoder.load_state_dict(state_dict, strict=False) |
|
|
self.image_encoder.to(self.device, self.dtype) |
|
|
del clip_encoder |
|
|
self.clip_image_processor = CLIPImageProcessor() |
|
|
|
|
|
|
|
|
attn_procs = {} |
|
|
for name in unet.attn_processors.keys(): |
|
|
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] |
|
|
if cross_attention_dim is None: |
|
|
attn_procs[name] = AttnProcessor() |
|
|
else: |
|
|
attn_procs[name] = SSRAttnProcessor(hidden_size=hidden_size, cross_attention_dim=1024, scale=1).to(self.device, dtype=self.dtype) |
|
|
unet.set_attn_processor(attn_procs) |
|
|
adapter_modules = torch.nn.ModuleList(unet.attn_processors.values()) |
|
|
self.SSR_layers = adapter_modules |
|
|
self.SSR_layers.to(self.device, dtype=self.dtype) |
|
|
self.resampler = self.init_proj() |
|
|
|
|
|
def init_proj(self): |
|
|
resampler = Resampler().to(self.device, dtype=self.dtype) |
|
|
return resampler |
|
|
|
|
|
def forward(self, img): |
|
|
image_embeds = self.image_encoder(img, output_hidden_states=True)['hidden_states'][2::2] |
|
|
image_embeds = torch.cat(image_embeds, dim=1) |
|
|
image_embeds = self.resampler(image_embeds) |
|
|
return image_embeds |
|
|
|
|
|
@torch.inference_mode() |
|
|
def get_image_embeds(self, pil_image): |
|
|
if isinstance(pil_image, Image.Image): |
|
|
pil_image = [pil_image] |
|
|
clip_image = [] |
|
|
for pil in pil_image: |
|
|
tensor_image = self.clip_image_processor(images=pil, return_tensors="pt").pixel_values.to(self.device, dtype=self.dtype) |
|
|
clip_image.append(tensor_image) |
|
|
clip_image = torch.cat(clip_image, dim=0) |
|
|
|
|
|
|
|
|
clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True)['hidden_states'][2::2] |
|
|
clip_image_embeds = torch.cat(clip_image_embeds, dim=1) |
|
|
uncond_clip_image_embeds = self.image_encoder(torch.zeros_like(clip_image), output_hidden_states=True)['hidden_states'][2::2] |
|
|
uncond_clip_image_embeds = torch.cat(uncond_clip_image_embeds, dim=1) |
|
|
clip_image_embeds = self.resampler(clip_image_embeds) |
|
|
uncond_clip_image_embeds = self.resampler(uncond_clip_image_embeds) |
|
|
return clip_image_embeds, uncond_clip_image_embeds |
|
|
|
|
|
def generate( |
|
|
self, |
|
|
id_image, |
|
|
makeup_image, |
|
|
seed=None, |
|
|
guidance_scale=2, |
|
|
num_inference_steps=30, |
|
|
pipe=None, |
|
|
**kwargs, |
|
|
): |
|
|
image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(makeup_image) |
|
|
|
|
|
prompt_embeds = image_prompt_embeds |
|
|
negative_prompt_embeds = uncond_image_prompt_embeds |
|
|
|
|
|
generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None |
|
|
image = pipe( |
|
|
image=id_image, |
|
|
prompt_embeds=prompt_embeds, |
|
|
negative_prompt_embeds=negative_prompt_embeds, |
|
|
guidance_scale=guidance_scale, |
|
|
num_inference_steps=num_inference_steps, |
|
|
generator=generator, |
|
|
**kwargs, |
|
|
).images[0] |
|
|
|
|
|
return image |
