introstyle.py

import os, gc, random, json, pathlib, glob, pickle
import numpy as np
from PIL import Image, ImageEnhance
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

from torchvision.transforms import PILToTensor
from typing import Any, Callable, Dict, List, Optional, Union, Tuple
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
from diffusers import DDIMScheduler
from diffusers import StableDiffusionPipeline

try:
    from tqdm.notebook import tqdm
except ImportError:
    from tqdm import tqdm

# Optional imports
try:
    import matplotlib.pyplot as plt
    from matplotlib.colors import hsv_to_rgb
except ImportError:
    pass

try:
    import scipy.io
except ImportError:
    pass

try:
    import faiss
except ImportError:
    pass

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,
        skip_ft_indices: Optional[List[int]] = None):
        """
        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 = {}
        skip_ft = {}
        for i, upsample_block in enumerate(self.up_blocks):

            max_up_index = -1 if (up_ft_indices is None or len(up_ft_indices) == 0) else int(np.max(up_ft_indices))
            max_skip_index = -1 if (skip_ft_indices is None or len(skip_ft_indices) == 0) else int(np.max(skip_ft_indices))
            max_index = max(max_up_index, max_skip_index)
            if i > max_index and max_index >= 0:
                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)]

            # Optionally record skip connection features used by this up block
            if skip_ft_indices is not None and i in skip_ft_indices:
                # Choose the highest-resolution skip tensor for this up block
                # res_samples is a tuple/list of tensors; take the last one
                skip_ft[i] = res_samples[-1].detach()

            # 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 up_ft_indices is not None and i in up_ft_indices:
                up_ft[i] = sample.detach()

        output = {}
        output['up_ft'] = up_ft
        if len(skip_ft) > 0:
            output['skip_ft'] = skip_ft
        return output

class OneStepSDPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def __call__(
        self,
        img_tensor,
        t,
        up_ft_indices,
        skip_ft_indices: Optional[List[int]] = None,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        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,
                               skip_ft_indices=skip_ft_indices,
                               encoder_hidden_states=prompt_embeds,
                               cross_attention_kwargs=cross_attention_kwargs)
        return unet_output


class SDFeaturizer:
    def __init__(self, sd_id='stabilityai/stable-diffusion-2-1', device: Optional[str] = 'auto'):
        # Select device
        if device == 'auto' or device is None:
            if torch.cuda.is_available():
                self.device = 'cuda'
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                self.device = 'mps'
            else:
                self.device = 'cpu'
        else:
            self.device = device

        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(self.device)
        onestep_pipe.enable_attention_slicing()
        # Enable xformers only on CUDA
        try:
            if self.device == 'cuda':
                onestep_pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            pass
        self.pipe = onestep_pipe

    @torch.no_grad()
    def forward(self,
                img_tensor, # single image, [1,c,h,w]
                prompt,
                t=261,
                up_ft_index=2,
                ensemble_size=4,
                use_skip=True,
                skip_ft_index=2,
                return_moments=True):
        nums = img_tensor.shape[0]
        nums = img_tensor.shape[0]
        img_tensor = img_tensor.to(self.device)
        if nums > 1:
            img_tensor = torch.cat([img_tensor[i:i+1].repeat(ensemble_size, 1, 1, 1).to(self.device) for i in range(nums)]).to(self.device)
        else:
            img_tensor = img_tensor.repeat(ensemble_size, 1, 1, 1).to(self.device) # ensem, c, h, w

        prompt_embeds = self.pipe.encode_prompt(
            prompt=prompt,
            device=self.device,
            num_images_per_prompt=1,
            do_classifier_free_guidance=False) # [1, 77, dim]

        prompt_embeds = prompt_embeds[0].repeat(ensemble_size*nums, 1, 1)
        if use_skip:
            up_indices = []
            skip_indices = [skip_ft_index]
        else:
            up_indices = [up_ft_index]
            skip_indices = []

        unet_ft_all = self.pipe(
            img_tensor=img_tensor,
            t=t,
            up_ft_indices=up_indices,
            skip_ft_indices=skip_indices,
            prompt_embeds=prompt_embeds)

        if use_skip:
            feat_map = unet_ft_all['skip_ft'][skip_ft_index] # ensem, c, h, w
        else:
            feat_map = unet_ft_all['up_ft'][up_ft_index] # ensem, c, h, w

        # ensemble average back to per-image
        feat_map = torch.cat([feat_map[i*ensemble_size:(i+1)*ensemble_size].mean(0, keepdim=True) for i in range(nums)])

        if return_moments:
            # Per-channel mean and variance over spatial dims
            mu = feat_map.mean(dim=(2,3))
            var = feat_map.var(dim=(2,3), unbiased=False)
            return torch.cat([mu, var], dim=1) # [B, 2C]
        else:
            return feat_map
    
class IntroStyleModel(torch.nn.Module):
    def __init__(self, device: Optional[str] = 'auto'):
        super().__init__()
        self.model = SDFeaturizer(device=device)
        # self.model.eval()
        #self.feat_dim = self.model.num_features
        self.preprocess = transforms.Compose([
            transforms.Resize(512, interpolation=3),
            transforms.CenterCrop(512),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        ])

    def forward(self,
                image,
                t = 25,
                up_ft_index = 1,
                use_skip: bool = True,
                skip_ft_index: int = 2,
                return_moments: bool = True):
        # If use_skip=True, extract from 3rd skip group by default (index 2).
        # If use_skip=False, extract from the selected up block (default index 1 = 2nd up block).
        return self.model.forward(image,
                                  prompt=['a photo of'],
                                  t=t,
                                  up_ft_index=up_ft_index,
                                  use_skip=use_skip,
                                  skip_ft_index=skip_ft_index,
                                  return_moments=return_moments)