mmcm/vision/feature_extractor/vae_extractor.py

import sys
from multiprocessing.pool import Pool
import os
import logging
from typing import Union, List, Tuple, Literal

import torch
import numpy as np
import pandas as pd
import h5py
import diffusers
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
from einops import rearrange
from PIL import Image
from transformers import CLIPTextModel, CLIPTokenizer

from ...data.extract_feature.base_extract_feature import BaseFeatureExtractor
from ...data.emb.h5py_emb import save_value_with_h5py

from ..process.image_process import dynamic_resize_image, dynamic_crop_resize_image
from ..utils.data_type_util import convert_images


class VAEFeatureExtractor(BaseFeatureExtractor):
    def __init__(
        self,
        pretrained_model_name_or_path: str,
        name: str = None,
        device: str = "cpu",
        dtype=torch.float32,
    ):
        super().__init__(device, dtype, name)
        self.pretrained_model_name_or_path = pretrained_model_name_or_path
        vae = AutoencoderKL.from_pretrained(
            pretrained_model_name_or_path, subfolder="vae"
        )
        vae.requires_grad_(False)
        self.vae = vae.to(device=device, dtype=dtype)
        vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

    def extract_images(
        self,
        data: Union[str, List[str], Image.Image, List[Image.Image], np.ndarray],
        target_width: int = None,
        target_height: int = None,
        return_type: str = "numpy",
        input_rgb_order: str = "rgb",
    ) -> Union[np.ndarray, torch.Tensor]:
        data = convert_images(data, return_type="pil", input_rgb_order=input_rgb_order)
        if target_height is not None and target_width is not None:
            data = [
                dynamic_crop_resize_image(
                    image,
                    target_height=target_height,
                    target_width=target_width,
                )
                for image in data
            ]
        batch = self.image_processor.preprocess(data).to(
            device=self.device, dtype=self.dtype
        )
        with torch.no_grad():
            # print("batch", batch.shape, batch.dtype, batch.device, self.vae.device)
            emb = self.vae.encoder(batch)
            quant_emb = self.vae.quant_conv(emb)
        if return_type == "numpy":
            emb = emb.cpu().numpy()
            quant_emb = quant_emb.cpu().numpy()
        return emb, quant_emb

    def extract_video(
        self,
        video_dataset,
        target_width: int = None,
        target_height: int = None,
        return_type: str = "numpy",
        track_performance: bool = False,
        input_rgb_order: str = "rgb",
    ) -> Union[np.ndarray, torch.Tensor]:
        embs = []
        quant_embs = []
        sample_indexs = []
        if track_performance:
            performance = {}
        with torch.no_grad():
            for i, (batch, batch_index) in enumerate(video_dataset):
                # TODO: 现阶段复用hugging face diffusers img2img pipeline中的抽取代码，
                # 由于该代码目前只支持Image的预处理，故先将numpy.ndarray转换成PIL.Image
                batch = [Image.fromarray(batch[b_i]) for b_i in range(len(batch))]
                emb, quant_emb = self.extract_images(
                    data=batch,
                    target_width=target_width,
                    target_height=target_height,
                    return_type=return_type,
                    input_rgb_order=input_rgb_order,
                )
                embs.append(emb)
                quant_embs.append(quant_emb)
                sample_indexs.extend(batch_index)

        sample_indexs = np.array(sample_indexs)
        if return_type == "numpy":
            embs = np.concatenate(embs, axis=0)
            quant_embs = np.concatenate(quant_embs, axis=0)
        elif return_type == "torch":
            embs = torch.concat(embs)
            quant_embs = torch.concat(quant_embs)
            sample_indexs = torch.from_numpy(sample_indexs)
        return sample_indexs, embs, quant_embs

    def extract(
        self,
        data: Union[str, List[str]],
        data_type: Literal["image", "video"],
        return_type: str = "numpy",
        save_emb_path: str = None,
        save_type: str = "h5py",
        emb_key: str = "encoder_emb",
        quant_emb_key: str = "encoder_quant_emb",
        sample_index_key: str = "sample_indexs",
        insert_name_to_key: bool = False,
        overwrite: bool = False,
        save_sample_index: bool = True,
        input_rgb_order: str = "rgb",
        **kwargs,
    ) -> Union[np.ndarray, torch.tensor]:
        if self.name is not None and insert_name_to_key:
            emb_key = f"{self.name}_{emb_key}"
            quant_emb_key = f"{self.name}_{quant_emb_key}"
            sample_index_key = f"{self.name}_{sample_index_key}"
        if save_emb_path is not None and os.path.exists(save_emb_path):
            with h5py.File(save_emb_path, "r") as f:
                if (
                    not overwrite
                    and emb_key in f
                    and quant_emb_key in f
                    and sample_index_key in f
                ):
                    return None

        if data_type == "image":
            emb, quant_emb = self.extract_images(
                data=data,
                return_type=return_type,
                input_rgb_order=input_rgb_order,
                **kwargs,
            )
            if save_emb_path is None:
                return emb, quant_emb
            else:
                raise NotImplementedError("save images emb")
        elif data_type == "video":
            sample_indexs, emb, quant_emb = self.extract_video(
                video_dataset=data,
                return_type=return_type,
                input_rgb_order=input_rgb_order,
                **kwargs,
            )
            if save_emb_path is None:
                return sample_indexs, emb, quant_emb
            else:
                if save_type == "h5py":
                    self.save_video_emb_with_h5py(
                        save_emb_path=save_emb_path,
                        emb=emb,
                        emb_key=emb_key,
                        quant_emb=quant_emb,
                        quant_emb_key=quant_emb_key,
                        sample_indexs=sample_indexs,
                        sample_index_key=sample_index_key,
                        save_sample_index=save_sample_index,
                        overwrite=overwrite,
                    )
                    return sample_indexs, emb, quant_emb
                else:
                    raise ValueError(f"only support save_type={save_type}")

    @staticmethod
    def save_images_emb_with_h5py(
        save_emb_path: str,
        emb: np.ndarray = None,
        emb_key: str = "encoder_emb",
        quant_emb: np.ndarray = None,
        quant_emb_key: str = "encoder_quant_emb",
    ) -> h5py.File:
        save_value_with_h5py(save_emb_path, value=emb, key=emb_key)
        save_value_with_h5py(save_emb_path, value=quant_emb, key=quant_emb_key)

    @staticmethod
    def save_video_emb_with_h5py(
        save_emb_path: str,
        emb: np.ndarray = None,
        emb_key: str = "encoder_emb",
        quant_emb: np.ndarray = None,
        quant_emb_key: str = "encoder_quant_emb",
        sample_indexs: np.ndarray = None,
        sample_index_key: str = "sample_indexs",
        overwrite: bool = False,
        save_sample_index: bool = True,
    ) -> h5py.File:
        # save_value_with_h5py(save_emb_path, value=emb, key=emb_key, overwrite=overwrite)
        if save_sample_index:
            save_value_with_h5py(
                save_emb_path,
                value=quant_emb,
                key=quant_emb_key,
                overwrite=overwrite,
                dtype=np.float16,
            )
            save_value_with_h5py(
                save_emb_path,
                value=sample_indexs,
                key=sample_index_key,
                overwrite=overwrite,
                dtype=np.uint32,
            )