import math import random import torch from diffusers import DiffusionPipeline, DDPMScheduler from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker, StableDiffusionPipelineOutput from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput from diffusers.image_processor import VaeImageProcessor from huggingface_hub import PyTorchModelHubMixin from PIL import Image from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor class CombinedStableDiffusion( DiffusionPipeline, PyTorchModelHubMixin ): """ A Stable Diffusion model wrapper that provides functionality for text-to-image synthesis, noise scheduling, latent space manipulation, and image decoding. """ def __init__( self, original_unet: torch.nn.Module, fine_tuned_unet: torch.nn.Module, scheduler: DDPMScheduler, vae: torch.nn.Module, tokenizer: CLIPTextModel, safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, text_encoder: CLIPTokenizer, ) -> None: super().__init__() self.register_modules( tokenizer=tokenizer, text_encoder=text_encoder, original_unet=original_unet, fine_tuned_unet=fine_tuned_unet, scheduler=scheduler, vae=vae, safety_checker=safety_checker, feature_extractor=feature_extractor, ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) self.image_processor = VaeImageProcessor( vae_scale_factor=self.vae_scale_factor ) def _get_negative_prompts(self, batch_size: int) -> torch.Tensor: return self.tokenizer( [""] * batch_size, max_length=self.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt", ).input_ids def _get_encoder_hidden_states( self, tokenized_prompts: torch.Tensor, do_classifier_free_guidance: bool = False ) -> torch.Tensor: if do_classifier_free_guidance: tokenized_prompts = torch.cat( [ self._get_negative_prompts(tokenized_prompts.shape[0]).to( tokenized_prompts.device ), tokenized_prompts, ] ) return self.text_encoder(tokenized_prompts)[0] def _get_unet_prediction( self, latent_model_input: torch.Tensor, timestep: int, encoder_hidden_states: torch.Tensor, ) -> torch.Tensor: """ Return unet noise prediction Args: latent_model_input (torch.Tensor): Unet latents input timestep (int): noise scheduler timestep encoder_hidden_states (torch.Tensor): Text encoder hidden states Returns: torch.Tensor: noise prediction """ unet = self.original_unet if self._use_original_unet else self.fine_tuned_unet return unet( latent_model_input, timestep=timestep, encoder_hidden_states=encoder_hidden_states, ).sample def get_noise_prediction( self, latents: torch.Tensor, timestep_index: int, encoder_hidden_states: torch.Tensor, do_classifier_free_guidance: bool = False, detach_main_path: bool = False, ): """ Return noise prediction Args: latents (torch.Tensor): Image latents timestep_index (int): noise scheduler timestep index encoder_hidden_states (torch.Tensor): Text encoder hidden states do_classifier_free_guidance (bool) Whether to do classifier free guidance detach_main_path (bool): Detach gradient Returns: torch.Tensor: noise prediction """ timestep = self.scheduler.timesteps[timestep_index] latent_model_input = self.scheduler.scale_model_input( sample=torch.cat([latents] * 2) if do_classifier_free_guidance else latents, timestep=timestep, ) noise_pred = self._get_unet_prediction( latent_model_input=latent_model_input, timestep=timestep, encoder_hidden_states=encoder_hidden_states, ) if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) if detach_main_path: noise_pred_text = noise_pred_text.detach() noise_pred = noise_pred_uncond + self.guidance_scale * ( noise_pred_text - noise_pred_uncond ) return noise_pred def sample_next_latents( self, latents: torch.Tensor, timestep_index: int, noise_pred: torch.Tensor, return_pred_original: bool = False, ) -> torch.Tensor: """ Return next latents prediction Args: latents (torch.Tensor): Image latents timestep_index (int): noise scheduler timestep index noise_pred (torch.Tensor): noise prediction return_pred_original (bool) Whether to sample original sample Returns: torch.Tensor: latent prediction """ timestep = self.scheduler.timesteps[timestep_index] sample = self.scheduler.step( model_output=noise_pred, timestep=timestep, sample=latents ) return ( sample.pred_original_sample if return_pred_original else sample.prev_sample ) def predict_next_latents( self, latents: torch.Tensor, timestep_index: int, encoder_hidden_states: torch.Tensor, return_pred_original: bool = False, do_classifier_free_guidance: bool = False, detach_main_path: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """ Predicts the next latent states during the diffusion process. Args: latents (torch.Tensor): Current latent states. timestep_index (int): Index of the current timestep. encoder_hidden_states (torch.Tensor): Encoder hidden states from the text encoder. return_pred_original (bool): Whether to return the predicted original sample. do_classifier_free_guidance (bool) Whether to do classifier free guidance detach_main_path (bool): Detach gradient Returns: tuple: Next latents and predicted noise tensor. """ noise_pred = self.get_noise_prediction( latents=latents, timestep_index=timestep_index, encoder_hidden_states=encoder_hidden_states, do_classifier_free_guidance=do_classifier_free_guidance, detach_main_path=detach_main_path, ) latents = self.sample_next_latents( latents=latents, noise_pred=noise_pred, timestep_index=timestep_index, return_pred_original=return_pred_original, ) return latents, noise_pred def get_latents(self, batch_size: int, device: torch.device) -> torch.Tensor: latent_resolution = int(self.resolution) // self.vae_scale_factor return torch.randn( ( batch_size, self.original_unet.config.in_channels, latent_resolution, latent_resolution, ), device=device, ) def do_k_diffusion_steps( self, start_timestep_index: int, end_timestep_index: int, latents: torch.Tensor, encoder_hidden_states: torch.Tensor, return_pred_original: bool = False, do_classifier_free_guidance: bool = False, detach_main_path: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """ Performs multiple diffusion steps between specified timesteps. Args: start_timestep_index (int): Starting timestep index. end_timestep_index (int): Ending timestep index. latents (torch.Tensor): Initial latents. encoder_hidden_states (torch.Tensor): Encoder hidden states. return_pred_original (bool): Whether to return the predicted original sample. do_classifier_free_guidance (bool) Whether to do classifier free guidance detach_main_path (bool): Detach gradient Returns: tuple: Resulting latents and encoder hidden states. """ assert start_timestep_index <= end_timestep_index for timestep_index in range(start_timestep_index, end_timestep_index - 1): latents, _ = self.predict_next_latents( latents=latents, timestep_index=timestep_index, encoder_hidden_states=encoder_hidden_states, return_pred_original=False, do_classifier_free_guidance=do_classifier_free_guidance, detach_main_path=detach_main_path, ) res, _ = self.predict_next_latents( latents=latents, timestep_index=end_timestep_index - 1, encoder_hidden_states=encoder_hidden_states, return_pred_original=return_pred_original, do_classifier_free_guidance=do_classifier_free_guidance, ) return res, encoder_hidden_states def get_pil_image(self, raw_images: torch.Tensor) -> list[Image]: do_denormalize = [True] * raw_images.shape[0] images = self.inference_image_processor.postprocess( raw_images, output_type="pil", do_denormalize=do_denormalize ) return images def get_reward_image(self, raw_images: torch.Tensor) -> torch.Tensor: reward_images = (raw_images / 2 + 0.5).clamp(0, 1) if self.use_image_shifting: self._shift_tensor_batch( reward_images, dx=random.randint(0, math.ceil(self.resolution / 224)), dy=random.randint(0, math.ceil(self.resolution / 224)), ) return self.reward_image_processor(reward_images) def run_safety_checker(self, image, device, dtype): if self.safety_checker is None: has_nsfw_concept = None else: if torch.is_tensor(image): feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") else: feature_extractor_input = self.image_processor.numpy_to_pil(image) safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device) image, has_nsfw_concept = self.safety_checker( images=image, clip_input=safety_checker_input.pixel_values.to(dtype) ) return image, has_nsfw_concept @torch.no_grad() def __call__( self, prompt: str | list[str], num_inference_steps=40, original_unet_steps=30, resolution=512, guidance_scale=7.5, output_type: str = "pil", return_dict: bool = True, generator=None, ): self.guidance_scale = guidance_scale batch_size = 1 if isinstance(prompt, str) else len(prompt) tokenized_prompts = self.tokenizer( prompt, return_tensors="pt", padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True ).input_ids.to(self.device) original_encoder_hidden_states = self._get_encoder_hidden_states( tokenized_prompts=tokenized_prompts, do_classifier_free_guidance=True ) fine_tuned_encoder_hidden_states = self._get_encoder_hidden_states( tokenized_prompts=tokenized_prompts, do_classifier_free_guidance=False ) latent_resolution = int(resolution) // self.vae_scale_factor latents = torch.randn( ( batch_size, self.original_unet.config.in_channels, latent_resolution, latent_resolution, ), device=self.device, ) self.scheduler.set_timesteps( num_inference_steps, device=self.device ) self._use_original_unet = True latents, _ = self.do_k_diffusion_steps( start_timestep_index=0, end_timestep_index=original_unet_steps, latents=latents, encoder_hidden_states=original_encoder_hidden_states, return_pred_original=False, do_classifier_free_guidance=True, ) self._use_original_unet = False latents, _ = self.do_k_diffusion_steps( start_timestep_index=original_unet_steps, end_timestep_index=num_inference_steps, latents=latents, encoder_hidden_states=fine_tuned_encoder_hidden_states, return_pred_original=False, do_classifier_free_guidance=False, ) if not output_type == "latent": image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ 0 ] image, has_nsfw_concept = self.run_safety_checker( image, self.device, original_encoder_hidden_states.dtype) else: image = latents has_nsfw_concept = None if has_nsfw_concept is None: do_denormalize = [True] * image.shape[0] else: do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] image = self.image_processor.postprocess( image, output_type=output_type, do_denormalize=do_denormalize ) # Offload all models self.maybe_free_model_hooks() if not return_dict: return image, has_nsfw_concept return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)