diffusers/community-pipelines-mirror

1import types2from typing import List, Optional, Tuple, Union3 4import torch5from transformers import CLIPTextModelWithProjection, CLIPTokenizer6from transformers.models.clip.modeling_clip import CLIPTextModelOutput7 8from diffusers.models import PriorTransformer9from diffusers.pipelines import DiffusionPipeline, StableDiffusionImageVariationPipeline10from diffusers.schedulers import UnCLIPScheduler11from diffusers.utils import logging12from diffusers.utils.torch_utils import randn_tensor13 14 15logger = logging.get_logger(__name__)  # pylint: disable=invalid-name16 17 18def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance):19    image = image.to(device=device)20    image_embeddings = image  # take image as image_embeddings21    image_embeddings = image_embeddings.unsqueeze(1)22 23    # duplicate image embeddings for each generation per prompt, using mps friendly method24    bs_embed, seq_len, _ = image_embeddings.shape25    image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1)26    image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)27 28    if do_classifier_free_guidance:29        uncond_embeddings = torch.zeros_like(image_embeddings)30 31        # For classifier free guidance, we need to do two forward passes.32        # Here we concatenate the unconditional and text embeddings into a single batch33        # to avoid doing two forward passes34        image_embeddings = torch.cat([uncond_embeddings, image_embeddings])35 36    return image_embeddings37 38 39class StableUnCLIPPipeline(DiffusionPipeline):40    def __init__(41        self,42        prior: PriorTransformer,43        tokenizer: CLIPTokenizer,44        text_encoder: CLIPTextModelWithProjection,45        prior_scheduler: UnCLIPScheduler,46        decoder_pipe_kwargs: Optional[dict] = None,47    ):48        super().__init__()49 50        decoder_pipe_kwargs = {"image_encoder": None} if decoder_pipe_kwargs is None else decoder_pipe_kwargs51 52        decoder_pipe_kwargs["torch_dtype"] = decoder_pipe_kwargs.get("torch_dtype", None) or prior.dtype53 54        self.decoder_pipe = StableDiffusionImageVariationPipeline.from_pretrained(55            "lambdalabs/sd-image-variations-diffusers", **decoder_pipe_kwargs56        )57 58        # replace `_encode_image` method59        self.decoder_pipe._encode_image = types.MethodType(_encode_image, self.decoder_pipe)60 61        self.register_modules(62            prior=prior,63            tokenizer=tokenizer,64            text_encoder=text_encoder,65            prior_scheduler=prior_scheduler,66        )67 68    def _encode_prompt(69        self,70        prompt,71        device,72        num_images_per_prompt,73        do_classifier_free_guidance,74        text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,75        text_attention_mask: Optional[torch.Tensor] = None,76    ):77        if text_model_output is None:78            batch_size = len(prompt) if isinstance(prompt, list) else 179            # get prompt text embeddings80            text_inputs = self.tokenizer(81                prompt,82                padding="max_length",83                max_length=self.tokenizer.model_max_length,84                return_tensors="pt",85            )86            text_input_ids = text_inputs.input_ids87            text_mask = text_inputs.attention_mask.bool().to(device)88 89            if text_input_ids.shape[-1] > self.tokenizer.model_max_length:90                removed_text = self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :])91                logger.warning(92                    "The following part of your input was truncated because CLIP can only handle sequences up to"93                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"94                )95                text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]96 97            text_encoder_output = self.text_encoder(text_input_ids.to(device))98 99            text_embeddings = text_encoder_output.text_embeds100            text_encoder_hidden_states = text_encoder_output.last_hidden_state101 102        else:103            batch_size = text_model_output[0].shape[0]104            text_embeddings, text_encoder_hidden_states = text_model_output[0], text_model_output[1]105            text_mask = text_attention_mask106 107        text_embeddings = text_embeddings.repeat_interleave(num_images_per_prompt, dim=0)108        text_encoder_hidden_states = text_encoder_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)109        text_mask = text_mask.repeat_interleave(num_images_per_prompt, dim=0)110 111        if do_classifier_free_guidance:112            uncond_tokens = [""] * batch_size113 114            uncond_input = self.tokenizer(115                uncond_tokens,116                padding="max_length",117                max_length=self.tokenizer.model_max_length,118                truncation=True,119                return_tensors="pt",120            )121            uncond_text_mask = uncond_input.attention_mask.bool().to(device)122            uncond_embeddings_text_encoder_output = self.text_encoder(uncond_input.input_ids.to(device))123 124            uncond_embeddings = uncond_embeddings_text_encoder_output.text_embeds125            uncond_text_encoder_hidden_states = uncond_embeddings_text_encoder_output.last_hidden_state126 127            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method128 129            seq_len = uncond_embeddings.shape[1]130            uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt)131            uncond_embeddings = uncond_embeddings.view(batch_size * num_images_per_prompt, seq_len)132 133            seq_len = uncond_text_encoder_hidden_states.shape[1]134            uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.repeat(1, num_images_per_prompt, 1)135            uncond_text_encoder_hidden_states = uncond_text_encoder_hidden_states.view(136                batch_size * num_images_per_prompt, seq_len, -1137            )138            uncond_text_mask = uncond_text_mask.repeat_interleave(num_images_per_prompt, dim=0)139 140            # done duplicates141 142            # For classifier free guidance, we need to do two forward passes.143            # Here we concatenate the unconditional and text embeddings into a single batch144            # to avoid doing two forward passes145            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])146            text_encoder_hidden_states = torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states])147 148            text_mask = torch.cat([uncond_text_mask, text_mask])149 150        return text_embeddings, text_encoder_hidden_states, text_mask151 152    @property153    def _execution_device(self):154        r"""155        Returns the device on which the pipeline's models will be executed. After calling156        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module157        hooks.158        """159        if self.device != torch.device("meta") or not hasattr(self.prior, "_hf_hook"):160            return self.device161        for module in self.prior.modules():162            if (163                hasattr(module, "_hf_hook")164                and hasattr(module._hf_hook, "execution_device")165                and module._hf_hook.execution_device is not None166            ):167                return torch.device(module._hf_hook.execution_device)168        return self.device169 170    def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):171        if latents is None:172            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)173        else:174            if latents.shape != shape:175                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")176            latents = latents.to(device)177 178        latents = latents * scheduler.init_noise_sigma179        return latents180 181    def to(self, torch_device: Optional[Union[str, torch.device]] = None):182        self.decoder_pipe.to(torch_device)183        super().to(torch_device)184 185    @torch.no_grad()186    def __call__(187        self,188        prompt: Optional[Union[str, List[str]]] = None,189        height: Optional[int] = None,190        width: Optional[int] = None,191        num_images_per_prompt: int = 1,192        prior_num_inference_steps: int = 25,193        generator: Optional[torch.Generator] = None,194        prior_latents: Optional[torch.Tensor] = None,195        text_model_output: Optional[Union[CLIPTextModelOutput, Tuple]] = None,196        text_attention_mask: Optional[torch.Tensor] = None,197        prior_guidance_scale: float = 4.0,198        decoder_guidance_scale: float = 8.0,199        decoder_num_inference_steps: int = 50,200        decoder_num_images_per_prompt: Optional[int] = 1,201        decoder_eta: float = 0.0,202        output_type: Optional[str] = "pil",203        return_dict: bool = True,204    ):205        if prompt is not None:206            if isinstance(prompt, str):207                batch_size = 1208            elif isinstance(prompt, list):209                batch_size = len(prompt)210            else:211                raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")212        else:213            batch_size = text_model_output[0].shape[0]214 215        device = self._execution_device216 217        batch_size = batch_size * num_images_per_prompt218 219        do_classifier_free_guidance = prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0220 221        text_embeddings, text_encoder_hidden_states, text_mask = self._encode_prompt(222            prompt, device, num_images_per_prompt, do_classifier_free_guidance, text_model_output, text_attention_mask223        )224 225        # prior226 227        self.prior_scheduler.set_timesteps(prior_num_inference_steps, device=device)228        prior_timesteps_tensor = self.prior_scheduler.timesteps229 230        embedding_dim = self.prior.config.embedding_dim231 232        prior_latents = self.prepare_latents(233            (batch_size, embedding_dim),234            text_embeddings.dtype,235            device,236            generator,237            prior_latents,238            self.prior_scheduler,239        )240 241        for i, t in enumerate(self.progress_bar(prior_timesteps_tensor)):242            # expand the latents if we are doing classifier free guidance243            latent_model_input = torch.cat([prior_latents] * 2) if do_classifier_free_guidance else prior_latents244 245            predicted_image_embedding = self.prior(246                latent_model_input,247                timestep=t,248                proj_embedding=text_embeddings,249                encoder_hidden_states=text_encoder_hidden_states,250                attention_mask=text_mask,251            ).predicted_image_embedding252 253            if do_classifier_free_guidance:254                predicted_image_embedding_uncond, predicted_image_embedding_text = predicted_image_embedding.chunk(2)255                predicted_image_embedding = predicted_image_embedding_uncond + prior_guidance_scale * (256                    predicted_image_embedding_text - predicted_image_embedding_uncond257                )258 259            if i + 1 == prior_timesteps_tensor.shape[0]:260                prev_timestep = None261            else:262                prev_timestep = prior_timesteps_tensor[i + 1]263 264            prior_latents = self.prior_scheduler.step(265                predicted_image_embedding,266                timestep=t,267                sample=prior_latents,268                generator=generator,269                prev_timestep=prev_timestep,270            ).prev_sample271 272        prior_latents = self.prior.post_process_latents(prior_latents)273 274        image_embeddings = prior_latents275 276        output = self.decoder_pipe(277            image=image_embeddings,278            height=height,279            width=width,280            num_inference_steps=decoder_num_inference_steps,281            guidance_scale=decoder_guidance_scale,282            generator=generator,283            output_type=output_type,284            return_dict=return_dict,285            num_images_per_prompt=decoder_num_images_per_prompt,286            eta=decoder_eta,287        )288        return output289