BLIP3o/BLIP3o-Model-4B

1# -*- coding: utf-8 -*-2 3# ===========================================================================================4#5#    Copyright (c) Beijing Academy of Artificial Intelligence (BAAI). All rights reserved.6#7#    Author        : Fan Zhang8#    Email         : zhangfan@baai.ac.cn9#    Institute     : Beijing Academy of Artificial Intelligence (BAAI)10#    Create On     : 2023-12-19 10:4511#    Last Modified : 2023-12-25 07:5912#    File Name     : pipeline_emu2_gen.py13#    Description   :14#15# ===========================================================================================16 17from dataclasses import dataclass18from typing import List, Optional19 20from PIL import Image21import numpy as np22import torch23from torchvision import transforms as TF24from tqdm import tqdm25 26from diffusers import DiffusionPipeline27from diffusers.utils import BaseOutput28 29from diffusers import UNet2DConditionModel, EulerDiscreteScheduler, AutoencoderKL30from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker31from transformers import CLIPImageProcessor32from transformers import AutoModelForCausalLM, AutoTokenizer33 34EVA_IMAGE_SIZE = 44835OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)36OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)37DEFAULT_IMG_PLACEHOLDER = "[<IMG_PLH>]"38 39@dataclass40class EmuVisualGenerationPipelineOutput(BaseOutput):41    image: Image.Image42    nsfw_content_detected: Optional[bool]43 44 45class EmuVisualGenerationPipeline(DiffusionPipeline):46 47    def __init__(48        self,49        tokenizer: AutoTokenizer,50        multimodal_encoder: AutoModelForCausalLM,51        scheduler: EulerDiscreteScheduler,52        unet: UNet2DConditionModel,53        vae: AutoencoderKL,54        feature_extractor: CLIPImageProcessor,55        safety_checker: StableDiffusionSafetyChecker,56        eva_size=EVA_IMAGE_SIZE,57        eva_mean=OPENAI_DATASET_MEAN,58        eva_std=OPENAI_DATASET_STD,59    ):60        super().__init__()61        self.register_modules(62            tokenizer=tokenizer,63            multimodal_encoder=multimodal_encoder,64            scheduler=scheduler,65            unet=unet,66            vae=vae,67            feature_extractor=feature_extractor,68            safety_checker=safety_checker,69        )70 71        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)72 73        self.transform = TF.Compose([74            TF.Resize((eva_size, eva_size), interpolation=TF.InterpolationMode.BICUBIC),75            TF.ToTensor(),76            TF.Normalize(mean=eva_mean, std=eva_std),77        ])78 79        self.negative_prompt = {}80 81    def device(self, module):82        return next(module.parameters()).device83 84    def dtype(self, module):85        return next(module.parameters()).dtype86 87    @torch.no_grad()88    def __call__(89        self,90        inputs: List[Image.Image | str] | str | Image.Image,91        height: int = 1024,92        width: int = 1024,93        num_inference_steps: int = 50,94        guidance_scale: float = 3.,95        crop_info: List[int] = [0, 0],96        original_size: List[int] = [1024, 1024],97    ):98        if not isinstance(inputs, list):99            inputs = [inputs]100 101        # 0. Default height and width to unet102        height = height or self.unet.config.sample_size * self.vae_scale_factor103        width = width or self.unet.config.sample_size * self.vae_scale_factor104 105        device = self.device(self.unet)106        dtype = self.dtype(self.unet)107 108        do_classifier_free_guidance = guidance_scale > 1.0109 110        # 1. Encode input prompt111        prompt_embeds = self._prepare_and_encode_inputs(112            inputs,113            do_classifier_free_guidance,114        ).to(dtype).to(device)115        batch_size = prompt_embeds.shape[0] // 2 if do_classifier_free_guidance else prompt_embeds.shape[0]116 117        unet_added_conditions = {}118        time_ids = torch.LongTensor(original_size + crop_info + [height, width]).to(device)119        if do_classifier_free_guidance:120            unet_added_conditions["time_ids"] = torch.cat([time_ids, time_ids], dim=0)121        else:122            unet_added_conditions["time_ids"] = time_ids123        unet_added_conditions["text_embeds"] = torch.mean(prompt_embeds, dim=1)124 125        # 2. Prepare timesteps126        self.scheduler.set_timesteps(num_inference_steps, device=device)127        timesteps = self.scheduler.timesteps128 129        # 3. Prepare latent variables130        shape = (131            batch_size,132            self.unet.config.in_channels,133            height // self.vae_scale_factor,134            width // self.vae_scale_factor,135        )136        latents = torch.randn(shape, device=device, dtype=dtype)137        latents = latents * self.scheduler.init_noise_sigma138 139        # 4. Denoising loop140        for t in tqdm(timesteps):141            # expand the latents if we are doing classifier free guidance142            # 2B x 4 x H x W143            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents144            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)145 146            noise_pred = self.unet(147                latent_model_input,148                t,149                encoder_hidden_states=prompt_embeds,150                added_cond_kwargs=unet_added_conditions,151            ).sample152 153            # perform guidance154            if do_classifier_free_guidance:155                noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)156                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)157 158            # compute the previous noisy sample x_t -> x_t-1159            latents = self.scheduler.step(noise_pred, t, latents).prev_sample160 161        # 5. Post-processing162        images = self.decode_latents(latents)163 164        # 6. Run safety checker165        images, has_nsfw_concept = self.run_safety_checker(images)166 167        # 7. Convert to PIL168        images = self.numpy_to_pil(images)169        return EmuVisualGenerationPipelineOutput(170            image=images[0],171            nsfw_content_detected=None if has_nsfw_concept is None else has_nsfw_concept[0],172        )173 174    def _prepare_and_encode_inputs(175        self,176        inputs: List[str | Image.Image],177        do_classifier_free_guidance: bool = False,178        placeholder: str = DEFAULT_IMG_PLACEHOLDER,179    ):180        device = self.device(self.multimodal_encoder.model.visual)181        dtype = self.dtype(self.multimodal_encoder.model.visual)182 183        has_image, has_text = False, False184        text_prompt, image_prompt = "", []185        for x in inputs:186            if isinstance(x, str):187                has_text = True188                text_prompt += x189            else:190                has_image = True191                text_prompt += placeholder192                image_prompt.append(self.transform(x))193 194        if len(image_prompt) == 0:195            image_prompt = None196        else:197            image_prompt = torch.stack(image_prompt)198            image_prompt = image_prompt.type(dtype).to(device)199 200        if has_image and not has_text:201            prompt = self.multimodal_encoder.model.encode_image(image=image_prompt)202            if do_classifier_free_guidance:203                key = "[NULL_IMAGE]"204                if key not in self.negative_prompt:205                    negative_image = torch.zeros_like(image_prompt)206                    self.negative_prompt[key] = self.multimodal_encoder.model.encode_image(image=negative_image)207                prompt = torch.cat([prompt, self.negative_prompt[key]], dim=0)208        else:209            prompt = self.multimodal_encoder.generate_image(text=[text_prompt], image=image_prompt, tokenizer=self.tokenizer)210            if do_classifier_free_guidance:211                key = ""212                if key not in self.negative_prompt:213                    self.negative_prompt[key] = self.multimodal_encoder.generate_image(text=[""], tokenizer=self.tokenizer)214                prompt = torch.cat([prompt, self.negative_prompt[key]], dim=0)215 216        return prompt217 218    def decode_latents(self, latents: torch.Tensor) -> np.ndarray:219        latents = 1 / self.vae.config.scaling_factor * latents220        image = self.vae.decode(latents).sample221        image = (image / 2 + 0.5).clamp(0, 1)222        image = image.cpu().permute(0, 2, 3, 1).float().numpy()223        return image224 225    def numpy_to_pil(self, images: np.ndarray) -> List[Image.Image]:226        """227        Convert a numpy image or a batch of images to a PIL image.228        """229        if images.ndim == 3:230            images = images[None, ...]231        images = (images * 255).round().astype("uint8")232        if images.shape[-1] == 1:233            # special case for grayscale (single channel) images234            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]235        else:236            pil_images = [Image.fromarray(image) for image in images]237 238        return pil_images239 240    def run_safety_checker(self, images: np.ndarray):241        if self.safety_checker is not None:242            device = self.device(self.safety_checker)243            dtype = self.dtype(self.safety_checker)244            safety_checker_input = self.feature_extractor(self.numpy_to_pil(images), return_tensors="pt").to(device)245            images, has_nsfw_concept = self.safety_checker(246                images=images, clip_input=safety_checker_input.pixel_values.to(dtype)247            )248        else:249            has_nsfw_concept = None250        return images, has_nsfw_concept251