sd_pipeline.py

from diffusers import StableDiffusionPipeline 
import torch 
from transformers import CLIPModel, CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
from aesthetic_scorer import SinusoidalTimeMLP
from typing import Callable, List, Optional, Union, Dict, Any
import torchvision
import numpy as np

from diffusers.configuration_utils import FrozenDict

import os
import sys
sys.path.append(os.getcwd())
from diffusers_patch.ddim_with_kl import ddim_step_KL, predict_x0_from_xt, ddim_step_KL_modified

#from diffusers.loaders import FromCkptMixin, LoraLoaderMixin, TextualInversionLoaderMixin
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import (
        deprecate,
        is_accelerate_available,
        is_accelerate_version,
        logging,
        randn_tensor,
        replace_example_docstring,

)
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker

from compressibility_scorer import condition_CompressibilityScorerDiff

class SMC_SDPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        r"""
        Function invoked when calling the pipeline for generation.
        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
        Examples:
        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )
        #weights = self.calculate_initial_weight(latents)



        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        kl_loss = 0
        
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order


        new_noise_pred = torch.zeros([batch_size, 4, 64, 64]).to(device)

        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    old_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

               

                noise_pred = old_noise_pred 

                ############################################################
                ############################################################


                # compute the previous noisy sample x_t -> x_t-1
                # latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample

                # Decoding: I Changed this part: We get multiple samples for each particle [Batch_size * Duplicates , 4, 64, 64  ]
                # The following generates [Batch_size * Duplicates , 4, 64, 64  ]
                with torch.no_grad():
                    previous_latents = torch.clone(latents) 
                    latents, kl_terms = ddim_step_KL(
                                        self.scheduler,
                                        noise_pred,   # (2,4,64,64),
                                        old_noise_pred, # (2,4,64,64),
                                        t,
                                        latents,
                                        eta=eta,  # 1.0
                                    )
                    kl_loss += torch.mean(kl_terms)
                ##### Start Calcuate Next Weight (Value functions)
                if i < len(timesteps) - 1: 
                    latent_model_input = torch.cat([latents] * 2) 
                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, timesteps[i+1])
                    noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample 
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    ggg_noise_pred = torch.clone(new_noise_pred) 
                    new_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # Get noises corresponding to latents 
            
                    weights =  self.calculate_weight_SMC(
                                latents, previous_latents, 
                                new_noise_pred, ggg_noise_pred, timesteps[i+1], self.alpha
                            )
                    #weights = weights.cpu().detach().numpy()
                    final_sample_indices = np.random.choice(weights.shape[0], weights.shape[0], p =  weights/weights.sum() ) 
                    latents = latents[final_sample_indices]
                    
                else:  #If we are in the last step 
                    pass

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
            has_nsfw_concept = None
        elif output_type == "pil":
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

            # 10. Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

        ##############
        has_nsfw_concept = False
        ##############

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return image, kl_loss

    def setup_oracle(self, scorer):
        self.scorer = scorer

    def set_reward(self, reward):
        self.reward = reward

    def setup_scorer(self, scorer):
        self.scorer = scorer
        self.scorer.requires_grad_(False)
        self.scorer.eval()

    def set_target(self, target):
        self.target = target

    def set_guidance(self, guidance):
        self.target_guidance = guidance
        
    def set_variant(self, variant):
        self.variant = variant
    
    def set_parameters(self, batch_size, duplicate_size, alpha = 10.0):
        self.batch_size = batch_size
        self.duplicate = duplicate_size
        self.alpha = alpha

   
    @torch.no_grad()
    def calculate_weight_SMC(self, latents, old_latents, new_noise_pred, old_noise_pred, t, alpha = 10.0): # t = 981, 961, 941 ..
        
        pred_original_sample = predict_x0_from_xt(
                                        self.scheduler,
                                        new_noise_pred,   
                                        t,
                                        latents
                                    )
        old_pred_original_sample = predict_x0_from_xt(
                                        self.scheduler,
                                        old_noise_pred,   
                                        t,
                                        old_latents
                                    )

        if self.reward == 'compressibility':
            ### Calculate new scores
            images = self.decode_latents(pred_original_sample)
            images = (images * 255).round().astype("uint8")
            images_scores = self.scorer(images)

            ### Calcualte old scores
            old_images = self.decode_latents(old_pred_original_sample)
            old_images = (old_images * 255).round().astype("uint8")
            old_images_scores = self.scorer(old_images)

            return np.exp( (images_scores-old_images_scores)/alpha)

        else: 
            ### Calculate new scores
            im_pix_un = self.vae.decode(pred_original_sample.to(self.vae.dtype) / 0.18215).sample
            im_pix = ((im_pix_un / 2) + 0.5).clamp(0, 1) 
            resize = torchvision.transforms.Resize(224, antialias=False)
            im_pix = resize(im_pix)
            normalize = torchvision.transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                    std=[0.26862954, 0.26130258, 0.27577711])
            im_pix = normalize(im_pix).to(im_pix_un.dtype)
            images_scores, _ = self.scorer(im_pix)
            images_scores = images_scores.cpu().detach().numpy()
            ### Calulate old scores 
            im_pix_un = self.vae.decode(old_pred_original_sample.to(self.vae.dtype) / 0.18215).sample
            im_pix = ((im_pix_un / 2) + 0.5).clamp(0, 1) 
            resize = torchvision.transforms.Resize(224, antialias=False)
            im_pix = resize(im_pix)
            normalize = torchvision.transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                    std=[0.26862954, 0.26130258, 0.27577711])
            im_pix = normalize(im_pix).to(im_pix_un.dtype)
            old_images_scores, _ = self.scorer(im_pix)

            old_images_scores = old_images_scores.cpu().detach().numpy()

            return np.exp( (images_scores-old_images_scores)/alpha) 



class Decoding_nonbatch_SDPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        r"""
        Function invoked when calling the pipeline for generation.
        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
        Examples:
        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )
        #weights = self.calculate_initial_weight(latents)



        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        kl_loss = 0
        
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    old_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

               

                noise_pred = old_noise_pred 

                ############################################################
                ############################################################


                # compute the previous noisy sample x_t -> x_t-1
                # latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample

                # Decoding: I Changed this part: We get multiple samples for each particle [Batch_size * Duplicates , 4, 64, 64  ]
                # The following generates [Batch_size * Duplicates , 4, 64, 64  ]
                if i < len(timesteps) - 1: 
                    weights_list = [ ]
                    latents_list = [ ]
                    for ttt in range(self.duplicate):
                        with torch.no_grad():
                            latents_duplicate, kl_terms = ddim_step_KL(
                                                self.scheduler,
                                                noise_pred,   # (2,4,64,64),
                                                old_noise_pred, # (2,4,64,64),
                                                t,
                                                latents,
                                                eta=eta,  # 1.0
                                            )
                            kl_loss += torch.mean(kl_terms)
                    ##### Start Calcuate Next Weight (Value functions)

                        latent_model_input = torch.cat([latents_duplicate] * 2) 
                        latent_model_input = self.scheduler.scale_model_input(latent_model_input, timesteps[i+1])
                        noise_pred_duplicate = self.unet(latent_model_input, timesteps[i+1], encoder_hidden_states= prompt_embeds, cross_attention_kwargs=cross_attention_kwargs).sample             
                        noise_pred_uncond, noise_pred_text = noise_pred_duplicate.chunk(2)
                        new_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # Get noises corresponding to latents 
                        weights =  self.calculate_weight(
                                    latents_duplicate, 
                                    new_noise_pred, timesteps[i+1] 
                                )
                        if  self.variant == 'PM' and self.reward == 'compressibility':
                            pass
                        else: 
                            weights = weights.cpu().detach().numpy()
                        weights_list.append(weights)
                        latents_list.append(latents_duplicate.cpu().detach().numpy()) 

                    weights_list = np.array(weights_list)
                    latents_list = np.array(latents_list) 
                    #index_chosen = [ ]
                    #for i in range(self.batch_size):
                    #    index_chosen.append(i *self.duplicate + np.argmax(weights_list[i*self.duplicate : (i+1)*self.duplicate]))
                    #latents = latents_list[index_chosen, :, :, :]
                    index_chosen = np.argmax(weights_list,0)
                    latents = np.array([latents_list[ index_chosen[i], i, :, :,:]  for i in range(self.batch_size)])
                    latents = torch.tensor(latents).to(device)
                else:  #If we are in the last step 
                    with torch.no_grad():
                        latents, kl_terms = ddim_step_KL(
                                            self.scheduler,
                                            noise_pred,   # (2,4,64,64),
                                            old_noise_pred, # (2,4,64,64),
                                            t,
                                            latents,
                                            eta=eta,  # 1.0
                                        ) 
                  

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
            has_nsfw_concept = None
        elif output_type == "pil":
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

            # 10. Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

        ##############
        has_nsfw_concept = False
        ##############

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return image, kl_loss

    def setup_oracle(self, scorer):
        self.scorer = scorer

    def set_reward(self, reward):
        self.reward = reward

    def setup_scorer(self, scorer):
        self.scorer = scorer
        self.scorer.requires_grad_(False)
        self.scorer.eval()

    def set_target(self, target):
        self.target = target

    def set_guidance(self, guidance):
        self.target_guidance = guidance
        
    def set_variant(self, variant):
        self.variant = variant
    
    def set_parameters(self, batch_size, duplicate_size):
        self.batch_size = batch_size
        self.duplicate = duplicate_size

   
    @torch.no_grad()
    def calculate_weight(self, latents, new_noise_pred, t): # t = 981, 961, 941 ..

        if  self.variant == 'PM' and self.reward == 'compressibility':
            pred_original_sample = predict_x0_from_xt(
                                self.scheduler,
                                new_noise_pred,   
                                t,
                                latents
                            )
            images = self.decode_latents(pred_original_sample)
            images = (images * 255).round().astype("uint8")
            weights = self.scorer(images)
            return weights           
        ## Calculate E[x_0|x_t]
        if self.variant == 'PM':
            pred_original_sample = predict_x0_from_xt(
                                        self.scheduler,
                                        new_noise_pred,   
                                        t,
                                        latents
                                    )
            im_pix_un = self.vae.decode(pred_original_sample.to(self.vae.dtype) / 0.18215).sample
        elif self.variant == 'MC':
            im_pix_un = self.vae.decode(latents.to(self.vae.dtype) / 0.18215).sample
            
        im_pix = ((im_pix_un / 2) + 0.5).clamp(0, 1) 

        # resize = torchvision.transforms.Resize(224, antialias=False)
        if self.reward == 'compressibility':
            resize = torchvision.transforms.Resize(512, antialias=False)
        elif self.reward == 'aesthetic':
            resize = torchvision.transforms.Resize(224, antialias=False)
        else:
            raise ValueError('Invalid reward type')
        
        im_pix = resize(im_pix)
        normalize = torchvision.transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                std=[0.26862954, 0.26130258, 0.27577711])
        im_pix = normalize(im_pix).to(im_pix_un.dtype)
        
        if self.variant == 'PM':
            weights, _ = self.scorer(im_pix)
        elif self.variant == 'MC':
            if self.reward == 'compressibility':
                weights, _ = self.scorer(latents, timesteps=t.repeat(im_pix.shape[0]))
            else:
                weights, _ = self.scorer(im_pix, timesteps=t.repeat(im_pix.shape[0]))
            
        return weights


class Decoding_SDPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        r"""
        Function invoked when calling the pipeline for generation.
        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
        Examples:
        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )
        #weights = self.calculate_initial_weight(latents)



        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        kl_loss = 0
        
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    old_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

               

                noise_pred = old_noise_pred 

                ############################################################
                ############################################################


                # compute the previous noisy sample x_t -> x_t-1
                # latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample

                # Decoding: I Changed this part: We get multiple samples for each particle [Batch_size * Duplicates , 4, 64, 64  ]
                # The following generates [Batch_size * Duplicates , 4, 64, 64  ]
                with torch.no_grad():
                    latents,  kl_terms = ddim_step_KL_modified(
                                        self.scheduler,
                                        noise_pred,   # (2,4,64,64),
                                        old_noise_pred, # (2,4,64,64),
                                        t,
                                        latents,  # (2,4,64,64)
                                        eta=eta,  # 1.0
                                        batch_size = self.batch_size, 
                                        duplicate = self.duplicate
                                    )
                    kl_loss += torch.mean(kl_terms)
                ##### Start Calcuate Next Weight (Value functions)
                if i < len(timesteps) - 1: 
                    latent_model_input = torch.cat([latents] * 2) 
                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, timesteps[i+1])
                    noise_pred = self.unet(latent_model_input, timesteps[i+1], encoder_hidden_states= torch.cat([prompt_embeds] * self.duplicate), cross_attention_kwargs=cross_attention_kwargs).sample             
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    new_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # Get noises corresponding to latents 
                    weights =  self.calculate_weight(
                                latents, 
                                new_noise_pred, timesteps[i+1] 
                            )
                    if  self.variant == 'PM' and self.reward == 'compressibility':
                        pass
                    else: 
                        weights = weights.cpu().detach().numpy()
                    index_chosen = [ ]
                    for i in range(self.batch_size):
                        index_chosen.append(i *self.duplicate + np.argmax(weights[i*self.duplicate : (i+1)*self.duplicate]))
                    latents = latents[index_chosen, :, :, :]
                else:  #If we are in the last step 
                    index_chosen = [i*self.duplicate for i in range(self.batch_size)]
                    latents = latents[index_chosen, :, :, :]

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
            has_nsfw_concept = None
        elif output_type == "pil":
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

            # 10. Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

        ##############
        has_nsfw_concept = False
        ##############

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return image, kl_loss

    def setup_oracle(self, scorer):
        self.scorer = scorer

    def set_reward(self, reward):
        self.reward = reward

    def setup_scorer(self, scorer):
        self.scorer = scorer
        self.scorer.requires_grad_(False)
        self.scorer.eval()

    def set_target(self, target):
        self.target = target

    def set_guidance(self, guidance):
        self.target_guidance = guidance
        
    def set_variant(self, variant):
        self.variant = variant
    
    def set_parameters(self, batch_size, duplicate_size):
        self.batch_size = batch_size
        self.duplicate = duplicate_size

   
    @torch.no_grad()
    def calculate_weight(self, latents, new_noise_pred, t): # t = 981, 961, 941 ..

        if  self.variant == 'PM' and self.reward == 'compressibility':
            pred_original_sample = predict_x0_from_xt(
                                self.scheduler,
                                new_noise_pred,   
                                t,
                                latents
                            )
            images = self.decode_latents(pred_original_sample)
            images = (images * 255).round().astype("uint8")
            weights = self.scorer(images)
            return weights           
        ## Calculate E[x_0|x_t]
        if self.variant == 'PM':
            pred_original_sample = predict_x0_from_xt(
                                        self.scheduler,
                                        new_noise_pred,   
                                        t,
                                        latents
                                    )
            im_pix_un = self.vae.decode(pred_original_sample.to(self.vae.dtype) / 0.18215).sample
        elif self.variant == 'MC':
            im_pix_un = self.vae.decode(latents.to(self.vae.dtype) / 0.18215).sample
            
        im_pix = ((im_pix_un / 2) + 0.5).clamp(0, 1) 

        # resize = torchvision.transforms.Resize(224, antialias=False)
        if self.reward == 'compressibility':
            resize = torchvision.transforms.Resize(512, antialias=False)
        elif self.reward == 'aesthetic':
            resize = torchvision.transforms.Resize(224, antialias=False)
        else:
            raise ValueError('Invalid reward type')
        
        im_pix = resize(im_pix)
        normalize = torchvision.transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                std=[0.26862954, 0.26130258, 0.27577711])
        im_pix = normalize(im_pix).to(im_pix_un.dtype)
        
        if self.variant == 'PM':
            weights, _ = self.scorer(im_pix)
        elif self.variant == 'MC':
            if self.reward == 'compressibility':
                weights, _ = self.scorer(latents, timesteps=t.repeat(im_pix.shape[0]))
            else:
                weights, _ = self.scorer(im_pix, timesteps=t.repeat(im_pix.shape[0]))
                #weights, _ = self.scorer(im_pix)
            
        return weights 



class DPS_continuous_SDPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        r"""
        Function invoked when calling the pipeline for generation.
        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
        Examples:
        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        kl_loss = 0
        
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    old_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)



                ############################################################
                ############################################################
                ## Guided Diffusion Modification ##

                ## grad = nabla_x 0.5 * || y - mu(x) ||^2
                ## nabla_x log p_t (y|x_t) = - [1/sigma^2] * grad

                ## For DDIM scheduler,
                ## modified noise = original noise - sqrt( 1-alpha_t ) * (nabla_x log p_t (y|x_t)) ,
                ## see eq(14) of http://arxiv.org/abs/2105.05233


                ## self.target_guidance <---> 1 / sigma^2
                ## self.target  <---> y

                # target = torch.tensor([self.target], dtype=torch.int64).to(latents.device)
                # target = target.repeat(batch_size * num_images_per_prompt)
                sqrt_1minus_alpha_t = (1 - self.scheduler.alphas_cumprod[t] ) **0.5
            
                # timestep_list = torch.tensor([t]).to(latents.device).repeat(batch_size * num_images_per_prompt)
                grad =  self.compute_gradient(
                            latents, 
                            old_noise_pred, 
                            t,
                            eta
                        )
                noise_pred = old_noise_pred - self.target_guidance * sqrt_1minus_alpha_t * grad


                ############################################################
                ############################################################


                # compute the previous noisy sample x_t -> x_t-1
                # latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
                with torch.no_grad():
                    latents, kl_terms = ddim_step_KL(
                                        self.scheduler,
                                        noise_pred,   # (2,4,64,64),
                                        old_noise_pred, # (2,4,64,64),
                                        t,
                                        latents,
                                        eta=eta,  # 1.0
                                    )
                    kl_loss += torch.mean(kl_terms)

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
            has_nsfw_concept = None
        elif output_type == "pil":
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

            # 10. Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

        ##############
        has_nsfw_concept = False
        ##############

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return image, kl_loss


    def setup_scorer(self, scorer):
        self.scorer = scorer
        self.scorer.requires_grad_(False)
        self.scorer.eval()

    # def set_target(self, target):
    #     self.target = target
    
    def set_reward(self, reward):
        self.reward = reward

    def set_guidance(self, guidance):
        self.target_guidance = guidance


    @torch.enable_grad()
    def compute_gradient(self, latent, old_noise_pred, t, eta):
        latent.requires_grad_(True)
        pred_original_sample = predict_x0_from_xt(
                            self.scheduler,
                            old_noise_pred,   # (2,4,64,64),
                            t,
                            latent,
                            eta=eta,  # 1.0
                        )
        
        im_pix_un = self.vae.decode(pred_original_sample.to(self.vae.dtype) / 0.18215).sample
        im_pix = ((im_pix_un / 2) + 0.5).clamp(0, 1) 
        
        if self.reward == 'compressibility':
            resize = torchvision.transforms.Resize(512, antialias=False)
        
        elif self.reward == 'aesthetic':
            resize = torchvision.transforms.Resize(224, antialias=False)
        else:
            raise ValueError('Invalid reward type')

        
        im_pix = resize(im_pix)
        normalize = torchvision.transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                std=[0.26862954, 0.26130258, 0.27577711])
        im_pix = normalize(im_pix).to(im_pix_un.dtype)
            
        rewards, _ = self.scorer(im_pix)
        reward = rewards.mean()
        
        self.scorer.zero_grad()
        reward.backward()
        return latent.grad.clone()


class Normal_continuous_SDPipeline(StableDiffusionPipeline):
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ):
        r"""
        Function invoked when calling the pipeline for generation.
        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                instead.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            guidance_scale (`float`, *optional*, defaults to 7.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            negative_prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
                less than `1`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            eta (`float`, *optional*, defaults to 0.0):
                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
                [`schedulers.DDIMScheduler`], will be ignored for others.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
                to make generation deterministic.
            latents (`torch.FloatTensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                provided, text embeddings will be generated from `prompt` input argument.
            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
                argument.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
                plain tuple.
            callback (`Callable`, *optional*):
                A function that will be called every `callback_steps` steps during inference. The function will be
                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
            callback_steps (`int`, *optional*, defaults to 1):
                The frequency at which the `callback` function will be called. If not specified, the callback will be
                called at every step.
            cross_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
        Examples:
        Returns:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
            When returning a tuple, the first element is a list with the generated images, and the second element is a
            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
            (nsfw) content, according to the `safety_checker`.
        """
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
        )

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0

        # 3. Encode input prompt
        prompt_embeds = self._encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
        )

        # 4. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        # 5. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )

        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)

        # 7. Denoising loop
        kl_loss = 0
        
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    cross_attention_kwargs=cross_attention_kwargs,
                ).sample

                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    old_noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)



                ############################################################
                ############################################################
                ## Guided Diffusion Modification ##

                ## grad = nabla_x 0.5 * || y - mu(x) ||^2
                ## nabla_x log p_t (y|x_t) = - [1/sigma^2] * grad

                ## For DDIM scheduler,
                ## modified noise = original noise - sqrt( 1-alpha_t ) * (nabla_x log p_t (y|x_t)) ,
                ## see eq(14) of http://arxiv.org/abs/2105.05233


                ## self.target_guidance <---> 1 / sigma^2
                ## self.target  <---> y

                # target = torch.tensor([self.target], dtype=torch.int64).to(latents.device)
                # target = target.repeat(batch_size * num_images_per_prompt)
                sqrt_1minus_alpha_t = (1 - self.scheduler.alphas_cumprod[t] ) **0.5
            
                noise_pred = old_noise_pred 

                ############################################################
                ############################################################


                # compute the previous noisy sample x_t -> x_t-1
                # latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
                with torch.no_grad():
                    latents, kl_terms = ddim_step_KL(
                                        self.scheduler,
                                        noise_pred,   # (2,4,64,64),
                                        old_noise_pred, # (2,4,64,64),
                                        t,
                                        latents,
                                        eta=eta,  # 1.0
                                    )
                    kl_loss += torch.mean(kl_terms)

                # call the callback, if provided
                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                    progress_bar.update()
                    if callback is not None and i % callback_steps == 0:
                        callback(i, t, latents)

        if output_type == "latent":
            image = latents
            has_nsfw_concept = None
        elif output_type == "pil":
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

            # 10. Convert to PIL
            image = self.numpy_to_pil(image)
        else:
            # 8. Post-processing
            image = self.decode_latents(latents)

            # 9. Run safety checker
            #############################################
            ## Disabled for correct evaluation of the reward
            #############################################
            #image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)

        ##############
        has_nsfw_concept = False
        ##############

        # Offload last model to CPU
        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
            self.final_offload_hook.offload()

        if not return_dict:
            return (image, has_nsfw_concept)

        return image, kl_loss


    def setup_scorer(self, scorer):
        self.scorer = scorer
        self.scorer.requires_grad_(False)
        self.scorer.eval()

    # def set_target(self, target):
    #     self.target = target
    
    def set_reward(self, reward):
        self.reward = reward

    def set_guidance(self, guidance):
        self.target_guidance = guidance


    @torch.enable_grad()
    def compute_gradient(self, latent, old_noise_pred, t, eta):
        latent.requires_grad_(True)
        pred_original_sample = predict_x0_from_xt(
                            self.scheduler,
                            old_noise_pred,   # (2,4,64,64),
                            t,
                            latent,
                            eta=eta,  # 1.0
                        )
        
        im_pix_un = self.vae.decode(pred_original_sample.to(self.vae.dtype) / 0.18215).sample
        im_pix = ((im_pix_un / 2) + 0.5).clamp(0, 1) 
        
        if self.reward == 'compressibility':
            resize = torchvision.transforms.Resize(512, antialias=False)
        
        elif self.reward == 'aesthetic':
            resize = torchvision.transforms.Resize(224, antialias=False)
        else:
            raise ValueError('Invalid reward type')

        
        im_pix = resize(im_pix)
        normalize = torchvision.transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                                std=[0.26862954, 0.26130258, 0.27577711])
        im_pix = normalize(im_pix).to(im_pix_un.dtype)
            
        rewards, _ = self.scorer(im_pix)
        reward = rewards.mean()
        
        self.scorer.zero_grad()
        reward.backward()
        return latent.grad.clone()

if __name__ == "__main__":
    
    torch.cuda.empty_cache()
    
    device = 'cuda'
    pipeline = Normal_continuous_SDPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", local_files_only=True)
    pipeline.to(device)
    pipeline.vae.requires_grad_(False)
    pipeline.text_encoder.requires_grad_(False)
    pipeline.unet.requires_grad_(False)

    pipeline.vae.eval()
    pipeline.text_encoder.eval()
    pipeline.unet.eval()
    
    pipeline.scheduler.set_timesteps(50, device=device)
    print(pipeline.scheduler.timesteps)

    reward_model = torch.load('model/reward_predictor_epoch_3.pth').to(device) # get the sinusoidalTime MLP
    reward_model.eval()
    reward_model.requires_grad_(False)
    
    import prompts as prompts_file
    eval_prompt_fn = getattr(prompts_file, 'eval_simple_animals')
    
    pipeline.setup_reward_model(reward_model)
    pipeline.set_target(7)
    pipeline.set_guidance(100)
    
    image = []
    batch_size = 2
    num_images = 64
    init_latents = None
    
    for i in range(num_images // batch_size):
        eval_prompts, eval_prompt_metadata = zip(
            *[eval_prompt_fn() for _ in range(batch_size)]
        )   
        eval_prompts = list(eval_prompts)
        
        if init_latents is None:
            init_i = None
        else:
            init_i = init_latents[i]
        image_ = pipeline(eval_prompts, num_images_per_prompt=batch_size, latents=init_i).images # List of PIL.Image objects
        image.extend(image_)