268 lines
11 KiB
Python
268 lines
11 KiB
Python
import gc
|
|
from typing import Any, Dict, Optional, Union
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers import DDIMScheduler, StableDiffusionPipeline
|
|
from diffusers.models.unet_2d_condition import UNet2DConditionModel
|
|
from PIL import Image, ImageDraw
|
|
|
|
|
|
class MyUNet2DConditionModel(UNet2DConditionModel):
|
|
def forward(
|
|
self,
|
|
sample: torch.FloatTensor,
|
|
timestep: Union[torch.Tensor, float, int],
|
|
up_ft_indices,
|
|
encoder_hidden_states: torch.Tensor,
|
|
class_labels: Optional[torch.Tensor] = None,
|
|
timestep_cond: Optional[torch.Tensor] = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
cross_attention_kwargs: Optional[Dict[str, Any]] = None
|
|
):
|
|
r"""
|
|
Args:
|
|
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
|
|
timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
|
|
encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
|
|
cross_attention_kwargs (`dict`, *optional*):
|
|
A kwargs dictionary that if specified is passed along to the `AttnProcessor` as defined under
|
|
`self.processor` in
|
|
[diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
|
|
"""
|
|
# By default samples have to be AT least a multiple of the overall upsampling factor.
|
|
# The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
|
|
# However, the upsampling interpolation output size can be forced to fit any upsampling size
|
|
# on the fly if necessary.
|
|
default_overall_up_factor = 2**self.num_upsamplers
|
|
|
|
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
|
forward_upsample_size = False
|
|
upsample_size = None
|
|
|
|
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
|
# logger.info("Forward upsample size to force interpolation output size.")
|
|
forward_upsample_size = True
|
|
|
|
# prepare attention_mask
|
|
if attention_mask is not None:
|
|
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
|
|
attention_mask = attention_mask.unsqueeze(1)
|
|
|
|
# 0. center input if necessary
|
|
if self.config.center_input_sample:
|
|
sample = 2 * sample - 1.0
|
|
|
|
# 1. time
|
|
timesteps = timestep
|
|
if not torch.is_tensor(timesteps):
|
|
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
|
# This would be a good case for the `match` statement (Python 3.10+)
|
|
is_mps = sample.device.type == 'mps'
|
|
if isinstance(timestep, float):
|
|
dtype = torch.float32 if is_mps else torch.float64
|
|
else:
|
|
dtype = torch.int32 if is_mps else torch.int64
|
|
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
|
elif len(timesteps.shape) == 0:
|
|
timesteps = timesteps[None].to(sample.device)
|
|
|
|
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
timesteps = timesteps.expand(sample.shape[0])
|
|
|
|
t_emb = self.time_proj(timesteps)
|
|
|
|
# timesteps does not contain any weights and will always return f32 tensors
|
|
# but time_embedding might actually be running in fp16. so we need to cast here.
|
|
# there might be better ways to encapsulate this.
|
|
t_emb = t_emb.to(dtype=self.dtype)
|
|
|
|
emb = self.time_embedding(t_emb, timestep_cond)
|
|
|
|
if self.class_embedding is not None:
|
|
if class_labels is None:
|
|
raise ValueError('class_labels should be provided when num_class_embeds > 0')
|
|
|
|
if self.config.class_embed_type == 'timestep':
|
|
class_labels = self.time_proj(class_labels)
|
|
|
|
class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
|
|
emb = emb + class_emb
|
|
|
|
# 2. pre-process
|
|
sample = self.conv_in(sample)
|
|
|
|
# 3. down
|
|
down_block_res_samples = (sample,)
|
|
for downsample_block in self.down_blocks:
|
|
if hasattr(downsample_block, 'has_cross_attention') and downsample_block.has_cross_attention:
|
|
sample, res_samples = downsample_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
attention_mask=attention_mask,
|
|
cross_attention_kwargs=cross_attention_kwargs,
|
|
)
|
|
else:
|
|
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
|
|
|
down_block_res_samples += res_samples
|
|
|
|
# 4. mid
|
|
if self.mid_block is not None:
|
|
sample = self.mid_block(
|
|
sample,
|
|
emb,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
attention_mask=attention_mask,
|
|
cross_attention_kwargs=cross_attention_kwargs,
|
|
)
|
|
|
|
# 5. up
|
|
up_ft = {}
|
|
|
|
for i, upsample_block in enumerate(self.up_blocks):
|
|
|
|
if i > np.max(up_ft_indices):
|
|
break
|
|
|
|
is_final_block = i == len(self.up_blocks) - 1
|
|
|
|
res_samples = down_block_res_samples[-len(upsample_block.resnets):]
|
|
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
|
|
|
# if we have not reached the final block and need to forward the
|
|
# upsample size, we do it here
|
|
if not is_final_block and forward_upsample_size:
|
|
upsample_size = down_block_res_samples[-1].shape[2:]
|
|
|
|
if hasattr(upsample_block, 'has_cross_attention') and upsample_block.has_cross_attention:
|
|
sample = upsample_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
res_hidden_states_tuple=res_samples,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
cross_attention_kwargs=cross_attention_kwargs,
|
|
upsample_size=upsample_size,
|
|
attention_mask=attention_mask,
|
|
)
|
|
else:
|
|
sample = upsample_block(
|
|
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
|
|
)
|
|
|
|
if i in up_ft_indices:
|
|
up_ft[i] = sample.detach()
|
|
|
|
output = {}
|
|
output['up_ft'] = up_ft
|
|
|
|
return output
|
|
|
|
|
|
class OneStepSDPipeline(StableDiffusionPipeline):
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
img_tensor,
|
|
t,
|
|
up_ft_indices,
|
|
prompt_embeds: Optional[torch.FloatTensor] = None,
|
|
cross_attention_kwargs: Optional[Dict[str, Any]] = None
|
|
):
|
|
|
|
device = self._execution_device
|
|
latents = self.vae.encode(img_tensor).latent_dist.sample() * self.vae.config.scaling_factor
|
|
t = torch.tensor(t, dtype=torch.long, device=device)
|
|
noise = torch.randn_like(latents).to(device)
|
|
latents_noisy = self.scheduler.add_noise(latents, noise, t)
|
|
unet_output = self.unet(latents_noisy, t, up_ft_indices, encoder_hidden_states=prompt_embeds, cross_attention_kwargs=cross_attention_kwargs)
|
|
return unet_output
|
|
|
|
|
|
class SDFeaturizer:
|
|
def __init__(self, sd_id='pretrained_models/stable-diffusion-v1-4'):
|
|
unet = MyUNet2DConditionModel.from_pretrained(sd_id, subfolder='unet')
|
|
onestep_pipe = OneStepSDPipeline.from_pretrained(sd_id, unet=unet, safety_checker=None)
|
|
onestep_pipe.vae.decoder = None
|
|
onestep_pipe.scheduler = DDIMScheduler.from_pretrained(sd_id, subfolder='scheduler')
|
|
gc.collect()
|
|
onestep_pipe = onestep_pipe.to('cuda')
|
|
onestep_pipe.enable_attention_slicing()
|
|
self.pipe = onestep_pipe
|
|
|
|
@torch.no_grad()
|
|
def forward(self,
|
|
img_tensor,
|
|
prompt,
|
|
t=261,
|
|
up_ft_index=0,
|
|
ensemble_size=8):
|
|
'''
|
|
Args:
|
|
img_tensor: should be a single torch tensor in the shape of [1, C, H, W] or [C, H, W]
|
|
prompt: the prompt to use, a string
|
|
t: the time step to use, should be an int in the range of [0, 1000]
|
|
up_ft_index: which upsampling block of the U-Net to extract feature, you can choose [0, 1, 2, 3]
|
|
ensemble_size: the number of repeated images used in the batch to extract features
|
|
Return:
|
|
unet_ft: a torch tensor in the shape of [1, c, h, w]
|
|
'''
|
|
img_tensor = img_tensor.repeat(ensemble_size, 1, 1, 1).cuda() # ensem, c, h, w
|
|
prompt_embeds = self.pipe._encode_prompt(
|
|
prompt=prompt,
|
|
device='cuda',
|
|
num_images_per_prompt=1,
|
|
do_classifier_free_guidance=False) # [1, 77, dim]
|
|
prompt_embeds = prompt_embeds.repeat(ensemble_size, 1, 1)
|
|
unet_ft_all = self.pipe(
|
|
img_tensor=img_tensor,
|
|
t=t,
|
|
up_ft_indices=[up_ft_index],
|
|
prompt_embeds=prompt_embeds)
|
|
unet_ft = unet_ft_all['up_ft'][up_ft_index] # ensem, c, h, w
|
|
unet_ft = unet_ft.mean(0, keepdim=True) # 1,c,h,w
|
|
return unet_ft
|
|
|
|
|
|
class DIFT_Demo:
|
|
def __init__(self, source_img, source_dift, source_img_size):
|
|
self.source_dift = source_dift # NCHW # torch.Size([1, 1280, 28, 48])
|
|
self.source_img = source_img
|
|
self.source_img_size = source_img_size
|
|
|
|
@torch.no_grad()
|
|
def query(self, target_img, target_dift, target_img_size, query_point, target_point, visualize=False):
|
|
num_channel = self.source_dift.size(1)
|
|
cos = nn.CosineSimilarity(dim=1)
|
|
source_x, source_y = int(np.round(query_point[1])), int(np.round(query_point[0]))
|
|
|
|
src_ft = self.source_dift
|
|
src_ft = nn.Upsample(size=self.source_img_size, mode='bilinear')(src_ft)
|
|
src_vec = src_ft[0, :, source_y, source_x].view(1, num_channel, 1, 1) # 1, C, 1, 1
|
|
|
|
tgt_ft = nn.Upsample(size=target_img_size, mode='bilinear')(target_dift)
|
|
cos_map = cos(src_vec, tgt_ft).cpu().numpy() # N, H, W (1, 448, 768)
|
|
|
|
max_yx = np.unravel_index(cos_map[0].argmax(), cos_map[0].shape)
|
|
target_x, target_y = int(np.round(target_point[1])), int(np.round(target_point[0]))
|
|
|
|
if visualize:
|
|
heatmap = cos_map[0]
|
|
heatmap = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap))
|
|
|
|
cmap = plt.get_cmap('viridis')
|
|
heatmap_color = (cmap(heatmap) * 255)[..., :3].astype(np.uint8)
|
|
|
|
alpha, radius, color = 0.5, 3, (0, 255, 0)
|
|
blended_image = Image.blend(target_img, Image.fromarray(heatmap_color), alpha=alpha)
|
|
draw = ImageDraw.Draw(blended_image)
|
|
draw.ellipse((max_yx[1] - radius, max_yx[0] - radius, max_yx[1] + radius, max_yx[0] + radius), fill=color)
|
|
draw.ellipse((target_x - radius, target_y - radius, target_x + radius, target_y + radius), fill=color)
|
|
else:
|
|
blended_image = None
|
|
dift_feat, confidence = tgt_ft[0, :, target_y, target_x], cos_map[0, target_y, target_x]
|
|
return dift_feat, confidence, blended_image
|