Add SMC stuff

src/smc/lora_pipeline.py

@@ -0,0 +1,313 @@
+import os
+from typing import Callable, Dict, List, Optional, Union
+
+import torch
+from huggingface_hub.utils import validate_hf_hub_args
+
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    deprecate,
+    get_submodule_by_name,
+    is_bitsandbytes_available,
+    is_gguf_available,
+    is_peft_available,
+    is_peft_version,
+    is_torch_version,
+    is_transformers_available,
+    is_transformers_version,
+    logging,
+)
+
+from diffusers.loaders.lora_base import (
+    LoraBaseMixin,
+    _fetch_state_dict,
+    _pack_dict_with_prefix
+)
+
+_LOW_CPU_MEM_USAGE_DEFAULT_LORA = False
+if is_torch_version(">=", "1.9.0"):
+    if (
+        is_peft_available()
+        and is_peft_version(">=", "0.13.1")
+        and is_transformers_available()
+        and is_transformers_version(">", "4.45.2")
+    ):
+        _LOW_CPU_MEM_USAGE_DEFAULT_LORA = True
+
+
+logger = logging.get_logger(__name__)
+
+
+TRANSFORMER_NAME = "transformer"
+
+class MeissonicLoraLoaderMixin(LoraBaseMixin):
+    r"""
+    Load LoRA layers into [`Transformer2DModel`]. Specific to [`MeissonicPipeline`].
+    """
+
+    _lora_loadable_modules = ["transformer"]
+    transformer_name = TRANSFORMER_NAME
+
+    @classmethod
+    @validate_hf_hub_args
+    def lora_state_dict(
+        cls,
+        pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+        return_alphas: bool = False,
+        **kwargs,
+    ):
+        r"""
+        Return state dict for lora weights and the network alphas.
+
+        <Tip warning={true}>
+
+        We support loading A1111 formatted LoRA checkpoints in a limited capacity.
+
+        This function is experimental and might change in the future.
+
+        </Tip>
+
+        Parameters:
+            pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
+                Can be either:
+
+                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
+                      the Hub.
+                    - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
+                      with [`ModelMixin.save_pretrained`].
+                    - A [torch state
+                      dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
+
+            cache_dir (`Union[str, os.PathLike]`, *optional*):
+                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
+                is not used.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+
+            proxies (`Dict[str, str]`, *optional*):
+                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            local_files_only (`bool`, *optional*, defaults to `False`):
+                Whether to only load local model weights and configuration files or not. If set to `True`, the model
+                won't be downloaded from the Hub.
+            token (`str` or *bool*, *optional*):
+                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
+                `diffusers-cli login` (stored in `~/.huggingface`) is used.
+            revision (`str`, *optional*, defaults to `"main"`):
+                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
+                allowed by Git.
+            subfolder (`str`, *optional*, defaults to `""`):
+                The subfolder location of a model file within a larger model repository on the Hub or locally.
+            return_lora_metadata (`bool`, *optional*, defaults to False):
+                When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
+        """
+        # Load the main state dict first which has the LoRA layers for either of
+        # transformer and text encoder or both.
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+        weight_name = kwargs.pop("weight_name", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+        return_lora_metadata = kwargs.pop("return_lora_metadata", False)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
+
+        state_dict, metadata = _fetch_state_dict(
+            pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
+            weight_name=weight_name,
+            use_safetensors=use_safetensors,
+            local_files_only=local_files_only,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            allow_pickle=allow_pickle,
+        )
+
+        is_dora_scale_present = any("dora_scale" in k for k in state_dict)
+        if is_dora_scale_present:
+            warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
+            logger.warning(warn_msg)
+            state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
+
+        out = (state_dict, metadata) if return_lora_metadata else state_dict
+        return out
+
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+        adapter_name: Optional[str] = None,
+        hotswap: bool = False,
+        **kwargs,
+    ):
+        """
+        Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
+        `self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
+        [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
+        See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
+        dict is loaded into `self.transformer`.
+
+        Parameters:
+            pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
+                See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
+            adapter_name (`str`, *optional*):
+                Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
+                `default_{i}` where i is the total number of adapters being loaded.
+            low_cpu_mem_usage (`bool`, *optional*):
+                Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
+                weights.
+            hotswap (`bool`, *optional*):
+                See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
+            kwargs (`dict`, *optional*):
+                See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
+        """
+        if not USE_PEFT_BACKEND:
+            raise ValueError("PEFT backend is required for this method.")
+
+        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
+        if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
+            raise ValueError(
+                "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
+            )
+
+        # if a dict is passed, copy it instead of modifying it inplace
+        if isinstance(pretrained_model_name_or_path_or_dict, dict):
+            pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
+
+        # First, ensure that the checkpoint is a compatible one and can be successfully loaded.
+        kwargs["return_lora_metadata"] = True
+        state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
+
+        is_correct_format = all("lora" in key for key in state_dict.keys())
+        if not is_correct_format:
+            raise ValueError("Invalid LoRA checkpoint.")
+        
+        self.load_lora_into_transformer(
+            state_dict,
+            transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
+            adapter_name=adapter_name,
+            metadata=metadata,
+            _pipeline=self,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            hotswap=hotswap,
+        )
+
+    @classmethod
+    def load_lora_into_transformer(
+        cls,
+        state_dict,
+        transformer,
+        adapter_name=None,
+        _pipeline=None,
+        low_cpu_mem_usage=False,
+        hotswap: bool = False,
+        metadata=None,
+    ):
+        """
+        This will load the LoRA layers specified in `state_dict` into `transformer`.
+
+        Parameters:
+            state_dict (`dict`):
+                A standard state dict containing the lora layer parameters. The keys can either be indexed directly
+                into the unet or prefixed with an additional `unet` which can be used to distinguish between text
+                encoder lora layers.
+            transformer (`SD3Transformer2DModel`):
+                The Transformer model to load the LoRA layers into.
+            adapter_name (`str`, *optional*):
+                Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
+                `default_{i}` where i is the total number of adapters being loaded.
+            low_cpu_mem_usage (`bool`, *optional*):
+                Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
+                weights.
+            hotswap (`bool`, *optional*):
+                See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
+            metadata (`dict`):
+                Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
+                from the state dict.
+        """
+        if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
+            raise ValueError(
+                "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
+            )
+
+        # Load the layers corresponding to transformer.
+        logger.info(f"Loading {cls.transformer_name}.")
+        transformer.load_lora_adapter(
+            state_dict,
+            network_alphas=None,
+            adapter_name=adapter_name,
+            metadata=metadata,
+            _pipeline=_pipeline,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            hotswap=hotswap,
+        )
+
+    @classmethod
+    # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
+    def save_lora_weights(
+        cls,
+        save_directory: Union[str, os.PathLike],
+        transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
+        is_main_process: bool = True,
+        weight_name: str = None,
+        save_function: Callable = None,
+        safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: Optional[dict] = None,
+    ):
+        r"""
+        Save the LoRA parameters corresponding to the transformer.
+
+        Arguments:
+            save_directory (`str` or `os.PathLike`):
+                Directory to save LoRA parameters to. Will be created if it doesn't exist.
+            transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
+                State dict of the LoRA layers corresponding to the `transformer`.
+            is_main_process (`bool`, *optional*, defaults to `True`):
+                Whether the process calling this is the main process or not. Useful during distributed training and you
+                need to call this function on all processes. In this case, set `is_main_process=True` only on the main
+                process to avoid race conditions.
+            save_function (`Callable`):
+                The function to use to save the state dictionary. Useful during distributed training when you need to
+                replace `torch.save` with another method. Can be configured with the environment variable
+                `DIFFUSERS_SAVE_MODE`.
+            safe_serialization (`bool`, *optional*, defaults to `True`):
+                Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
+            transformer_lora_adapter_metadata:
+                LoRA adapter metadata associated with the transformer to be serialized with the state dict.
+        """
+        state_dict = {}
+        lora_adapter_metadata = {}
+
+        if not transformer_lora_layers:
+            raise ValueError("You must pass `transformer_lora_layers`.")
+
+        state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
+
+        if transformer_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(
+                _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
+            )
+
+        # Save the model
+        cls.write_lora_layers(
+            state_dict=state_dict,
+            save_directory=save_directory,
+            is_main_process=is_main_process,
+            weight_name=weight_name,
+            save_function=save_function,
+            safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
+        )

src/smc/pipeline.py

@@ -0,0 +1,675 @@
+from typing import Optional, Tuple, Callable, List
+import math
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.vq_model import VQModel
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+from src.smc.transformer import Transformer2DModel
+from src.smc.scheduler import BaseScheduler
+from src.smc.resampling import compute_ess_from_log_w, normalize_weights
+from src.smc.lora_pipeline import MeissonicLoraLoaderMixin
+
+
+def logmeanexp(x, dim=None, keepdim=False):
+    """Numerically stable log-mean-exp using torch.logsumexp."""
+    if dim is None:
+        x = x.view(-1)
+        dim = 0
+    # log-sum-exp with or without keeping the reduced dim
+    lse = torch.logsumexp(x, dim=dim, keepdim=keepdim)
+    # subtract log(N) to convert sum into mean (broadcasts correctly)
+    return lse - math.log(x.size(dim))
+
+
+def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+    """
+    Build positional IDs for latent-image tokens.
+
+    Each latent token corresponds to a downsampled image “pixel” in a 2D grid.
+    This function creates a (H//2, W//2, 3) grid where:
+      - channel 0 is reserved (all zeros)
+      - channel 1 stores the row index (0 .. H//2-1)
+      - channel 2 stores the column index (0 .. W//2-1)
+
+    Args:
+        batch_size (int):   Number of images in the batch (unused here, but kept for API consistency).
+        height (int):       Input image height (pre-VAE) or latent height depending on call site.
+        width (int):        Input image width (pre-VAE) or latent width depending on call site.
+        device (torch.device):  Device on which to place the returned tensor.
+        dtype (torch.dtype):    Desired data type of the returned tensor.
+
+    Returns:
+        torch.Tensor of shape ((H//2 * W//2), 3) with dtype and device as specified.
+          Each row is [0, row_index, col_index], flattened in row-major order.
+    """
+    latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+
+    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+    latent_image_ids = latent_image_ids.reshape(
+        latent_image_id_height * latent_image_id_width, latent_image_id_channels
+    )
+
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+
+class Pipeline(
+    DiffusionPipeline,
+    MeissonicLoraLoaderMixin,
+):
+    image_processor: VaeImageProcessor
+    vqvae: VQModel
+    tokenizer: CLIPTokenizer
+    text_encoder: CLIPTextModelWithProjection
+    transformer: Transformer2DModel 
+    scheduler: BaseScheduler
+    
+    model_cpu_offload_seq = "text_encoder->transformer->vqvae"
+    
+    def __init__(
+        self,
+        vqvae: VQModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModelWithProjection,
+        transformer: Transformer2DModel, 
+        scheduler: BaseScheduler,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vqvae=vqvae,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vqvae.config.block_out_channels) - 1) # type: ignore
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False)
+        self.model_dtype = torch.bfloat16
+        
+        self.mask_index = self.scheduler.mask_token_id # type: ignore
+        self.vocab_size = self.transformer.config.vocab_size # type:ignore
+        self.codebook_size = self.transformer.config.codebook_size # type: ignore
+        
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str|List[str],
+        reward_fn: Callable,
+        resample_fn: Callable,
+        resample_frequency: int = 1,
+        kl_weight: float = 1.0,
+        lambdas: Optional[torch.Tensor] = None,
+        height: Optional[int] = 1024,
+        width: Optional[int] = 1024,
+        num_inference_steps: int = 48,
+        guidance_scale: float = 9.0,
+        negative_prompt = None,
+        batches: int = 1, # Number of independent SMCs
+        num_particles: int = 1, # Number of particles per SMC
+        batch_p: int = 1, # Number of parallel particles
+        phi: int = 1, # number of samples for reward approximation
+        tau: float = 1.0, # temperature for taking x0 samples
+        output_type="pil",
+        micro_conditioning_aesthetic_score: int = 6,
+        micro_conditioning_crop_coord: Tuple[int, int] = (0, 0),
+        proposal_type:str = "locally_optimal",
+        ft_model_pipe = None, # needs to supplied if proposal_type is ft_model
+        use_ft_model_for_expected_reward: bool = False, # Whether to use the forward model for expected reward
+        use_continuous_formulation: bool = False, # Whether to use a continuous formulation of carry over unmasking
+        disable_progress_bar: bool = False,
+        final_strategy="argmax_rewards",
+        verbose=True,
+    ):
+        # 0. Set default lambdas
+        if lambdas is None:
+            lambdas = torch.ones(num_inference_steps + 1)
+        assert len(lambdas) == num_inference_steps + 1, f"lambdas must of length {num_inference_steps + 1}"
+        lambdas = lambdas.clamp_min(0.001).to(self._execution_device)
+        
+        # 1. n_particles, batch_size etc
+        total_particles = batches * num_particles
+        batch_p = min(batch_p, total_particles)
+        H, W = height // self.vae_scale_factor, width // self.vae_scale_factor
+        
+        # 2.1. Calculate prompt (and negative prompt) embeddings
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        input_ids = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+        ).input_ids.to(self._execution_device)
+        outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True)
+        prompt_embeds = outputs.text_embeds
+        encoder_hidden_states = outputs.hidden_states[-2]
+        prompt_embeds = prompt_embeds.repeat(batch_p, 1)
+        encoder_hidden_states = encoder_hidden_states.repeat(batch_p, 1, 1)
+        if guidance_scale > 1.0:
+            if negative_prompt is None:
+                negative_prompt = [""]
+            else:
+                negative_prompt = [negative_prompt]
+            input_ids = self.tokenizer(
+                negative_prompt,
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=77,
+            ).input_ids.to(self._execution_device)
+            outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True)
+            negative_prompt_embeds = outputs.text_embeds
+            negative_encoder_hidden_states = outputs.hidden_states[-2]
+            negative_prompt_embeds = negative_prompt_embeds.repeat(batch_p, 1)
+            negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(batch_p, 1, 1)
+            prompt_embeds = torch.concat([negative_prompt_embeds, prompt_embeds])
+            encoder_hidden_states = torch.concat([negative_encoder_hidden_states, encoder_hidden_states])
+        
+        # 2.2. Prepare micro-conditions
+        micro_conds = torch.tensor(
+            [
+                width,
+                height,
+                micro_conditioning_crop_coord[0],
+                micro_conditioning_crop_coord[1],
+                micro_conditioning_aesthetic_score,
+            ],
+            device=self._execution_device,
+            dtype=encoder_hidden_states.dtype,
+        )
+        micro_conds = micro_conds.unsqueeze(0)
+        micro_conds = micro_conds.expand(2 * batch_p if guidance_scale > 1.0 else batch_p, -1)
+        
+        
+        # 3. Intialize latents
+        latents = torch.full(
+            (total_particles, H, W), self.mask_index, dtype=torch.long, device=self._execution_device # type: ignore
+        )
+        
+        # Set some constant vectors
+        ONE = torch.ones(self.vocab_size, device=self._execution_device).float()
+        MASK = F.one_hot(torch.tensor(self.mask_index), num_classes=self.vocab_size).float().to(self._execution_device) # type: ignore
+        
+         # 4. Set scheduler timesteps
+        self.scheduler.set_timesteps(num_inference_steps)
+        
+        # 5. Set SMC variables
+        logits = torch.zeros((*latents.shape, self.vocab_size), device=self._execution_device)
+        logits_ft_model = torch.zeros((*latents.shape, self.vocab_size), device=self._execution_device)
+        rewards = torch.zeros((total_particles,), device=self._execution_device)
+        rewards_grad = torch.zeros((*latents.shape, self.vocab_size), device=self._execution_device)
+        log_twist = torch.zeros((total_particles, ), device=self._execution_device)
+        log_prob_proposal = torch.zeros((total_particles, ), device=self._execution_device)
+        log_prob_diffusion = torch.zeros((total_particles, ), device=self._execution_device)
+        log_w = torch.zeros((total_particles, ), device=self._execution_device)
+        
+        def propagate():
+            if proposal_type == "locally_optimal":
+                propgate_locally_optimal()
+            # elif proposal_type == "straight_through_gradients":
+            #     propagate_straight_through_gradients()
+            elif proposal_type == "reverse":
+                propagate_reverse()
+            elif proposal_type == "without_SMC":
+                propagate_without_SMC()
+            elif proposal_type == "ft_model":
+                propagate_ft_model()
+            else:
+                raise NotImplementedError(f"Proposal type {proposal_type} is not implemented.")
+            
+        def propgate_locally_optimal():
+            nonlocal log_w, latents, log_prob_proposal, log_prob_diffusion, logits, rewards, rewards_grad, log_twist
+            log_twist_prev = log_twist.clone()
+            for j in range(0, total_particles, batch_p):
+                latents_batch = latents[j:j+batch_p]
+                with torch.enable_grad():
+                    latents_one_hot = F.one_hot(latents_batch, num_classes=self.vocab_size).to(dtype=self.model_dtype).requires_grad_(True)
+                    tmp_logits = self.get_logits(latents_one_hot, guidance_scale, height, encoder_hidden_states, micro_conds, prompt_embeds, timestep)
+                    
+                    tmp_rewards = torch.zeros(latents_batch.size(0), phi, device=self._execution_device)
+                    gamma = 1 - ((ONE - MASK) * latents_one_hot).sum(dim=-1, keepdim=True)
+                    for phi_i in range(phi):
+                        sample = F.gumbel_softmax(tmp_logits, tau=tau, hard=True)
+                        if use_continuous_formulation:
+                            sample = gamma * sample + (ONE - MASK) * latents_one_hot
+                        sample = self._decode_one_hot_latents(sample, batch_p, height, width, "pt")
+                        tmp_rewards[:, phi_i] = reward_fn(sample)
+                    tmp_rewards = logmeanexp(tmp_rewards * scale_cur, dim=-1) / scale_cur
+                    
+                    tmp_rewards_grad = torch.autograd.grad(
+                        outputs=tmp_rewards, 
+                        inputs=latents_one_hot,
+                        grad_outputs=torch.ones_like(tmp_rewards)
+                    )[0].detach()
+                
+                logits[j:j+batch_p] = tmp_logits.detach()
+                rewards[j:j+batch_p] = tmp_rewards.detach()
+                rewards_grad[j:j+batch_p] = tmp_rewards_grad.detach()
+                log_twist[j:j+batch_p] = rewards[j:j+batch_p] * scale_cur
+                
+            if verbose:
+                print("Rewards: ", rewards)
+            
+            # Calculate weights
+            incremental_log_w = (log_prob_diffusion - log_prob_proposal) + (log_twist - log_twist_prev)
+            log_w += incremental_log_w
+            
+            # Now reshape log_w to (batches, num_particles)
+            log_w = log_w.reshape(batches, num_particles)
+            
+            if verbose:
+                print("log_prob_diffusion - log_prob_proposal: ", log_prob_diffusion - log_prob_proposal)
+                print("log_twist - log_twist_prev:", log_twist - log_twist_prev)
+                print("Incremental log weights: ", incremental_log_w)
+                print("Log weights: ", log_w)
+                print("Normalized weights: ", normalize_weights(log_w, dim=-1))
+            
+            # Resample particles
+            if verbose:
+                print(f"ESS: ", compute_ess_from_log_w(log_w, dim=-1))
+            
+            if resample_condition:
+                resample_indices = []
+                log_w_new = []
+                is_resampled = False
+                for batch in range(batches):
+                    resample_indices_batch, is_resampled_batch, log_w_batch = resample_fn(log_w[batch])
+                    resample_indices.append(resample_indices_batch + batch * num_particles)
+                    log_w_new.append(log_w_batch)
+                    is_resampled = is_resampled or is_resampled_batch
+                    
+                resample_indices = torch.cat(resample_indices, dim=0)
+                log_w = torch.cat(log_w_new, dim=0)
+                    
+                if is_resampled:
+                    latents = latents[resample_indices]
+                    logits = logits[resample_indices]
+                    rewards = rewards[resample_indices]
+                    rewards_grad = rewards_grad[resample_indices]
+                    log_twist = log_twist[resample_indices]
+                    
+                if verbose:
+                    print("Resample indices: ", resample_indices)
+                
+            if log_w.ndim == 2:
+                log_w = log_w.reshape(total_particles)
+
+            
+            # Propose new particles
+            sched_out = self.scheduler.step_with_approx_guidance(
+                latents=latents,
+                logits=logits,
+                approx_guidance=rewards_grad * scale_next,
+                step=i,
+            )
+            if verbose:
+                print("Approx guidance norm: ", ((rewards_grad * scale_next) ** 2).sum(dim=(1, 2)).sqrt())
+            latents, log_prob_proposal, log_prob_diffusion = (
+                sched_out.new_latents,
+                sched_out.log_prob_proposal,
+                sched_out.log_prob_diffusion,
+            )
+            
+        def propagate_reverse():
+            nonlocal log_w, latents, logits, rewards, log_twist
+            log_twist_prev = log_twist.clone()
+            for j in range(0, total_particles, batch_p):
+                latents_batch = latents[j:j+batch_p]
+                with torch.no_grad():
+                    tmp_logits = self.get_logits(latents_batch, guidance_scale, height, encoder_hidden_states, micro_conds, prompt_embeds, timestep)
+                    
+                    tmp_rewards = torch.zeros(latents_batch.size(0), phi, device=self._execution_device)
+                    tmp_logp_x0 = self.model._subs_parameterization(tmp_logits, latents_batch)
+                    for phi_i in range(phi):
+                        sample = F.gumbel_softmax(tmp_logp_x0, tau=tau, hard=True).argmax(dim=-1)
+                        sample = self._decode_latents(sample, batch_p, height, width, "pt")
+                        tmp_rewards[:, phi_i] = reward_fn(sample)
+                    tmp_rewards = logmeanexp(tmp_rewards * scale_cur, dim=-1) / scale_cur
+                
+                logits[j:j+batch_p] = tmp_logits.detach()
+                rewards[j:j+batch_p] = tmp_rewards.detach()
+                log_twist[j:j+batch_p] = rewards[j:j+batch_p] * scale_cur
+                
+            if verbose:
+                print("Rewards: ", rewards)
+            
+            # Calculate weights
+            incremental_log_w = (log_twist - log_twist_prev)
+            log_w += incremental_log_w
+            
+            # Now reshape log_w to (batches, num_particles)
+            log_w = log_w.reshape(batches, num_particles)
+            
+            if verbose:
+                print("log_twist - log_twist_prev:", log_twist - log_twist_prev)
+                print("Incremental log weights: ", incremental_log_w)
+                print("Log weights: ", log_w)
+                print("Normalized weights: ", normalize_weights(log_w, dim=-1))
+            
+            # Resample particles
+            if verbose:
+                print(f"ESS: ", compute_ess_from_log_w(log_w, dim=-1))
+            
+            if resample_condition:
+                resample_indices = []
+                log_w_new = []
+                is_resampled = False
+                for batch in range(batches):
+                    resample_indices_batch, is_resampled_batch, log_w_batch = resample_fn(log_w[batch])
+                    resample_indices.append(resample_indices_batch + batch * num_particles)
+                    log_w_new.append(log_w_batch)
+                    is_resampled = is_resampled or is_resampled_batch
+                    
+                resample_indices = torch.cat(resample_indices, dim=0)
+                log_w = torch.cat(log_w_new, dim=0)
+                    
+                if is_resampled:
+                    latents = latents[resample_indices]
+                    logits = logits[resample_indices]
+                    rewards = rewards[resample_indices]
+                    log_twist = log_twist[resample_indices]
+                    
+                if verbose:
+                    print("Resample indices: ", resample_indices)
+                
+            if log_w.ndim == 2:
+                log_w = log_w.reshape(total_particles)
+
+            
+            # Propose new particles
+            sched_out = self.scheduler.step(
+                latents=latents,
+                logits=logits,
+                step=i,
+            )
+            latents = sched_out.new_latents
+        
+        def propagate_without_SMC():
+            nonlocal latents, logits
+            for j in range(0, total_particles, batch_p):
+                latents_batch = latents[j:j+batch_p]
+                with torch.no_grad():
+                    tmp_logits = self.get_logits(latents_batch, guidance_scale, height, encoder_hidden_states, micro_conds, prompt_embeds, timestep) 
+                logits[j:j+batch_p] = tmp_logits.detach()
+            
+            # Propose new particles
+            sched_out = self.scheduler.step(
+                latents=latents,
+                logits=logits,
+                step=i,
+            )
+            latents = sched_out.new_latents
+            
+        def propagate_ft_model():
+            assert ft_model_pipe is not None, f"ft_model must be provided for proposal_type={proposal_type}."
+            nonlocal log_w, latents, log_prob_proposal, log_prob_diffusion, logits, logits_ft_model, rewards, log_twist
+            log_twist_prev = log_twist.clone()
+            for j in range(0, total_particles, batch_p):
+                latents_batch = latents[j:j+batch_p]
+                with torch.no_grad():
+                    tmp_logits = self.get_logits(latents_batch, guidance_scale, height, encoder_hidden_states, micro_conds, prompt_embeds, timestep)
+                    tmp_logits_ft_model = ft_model_pipe.get_logits(latents_batch, guidance_scale, height, encoder_hidden_states, micro_conds, prompt_embeds, timestep)
+                    
+                    tmp_rewards = torch.zeros(latents_batch.size(0), phi, device=self._execution_device)
+                    if use_ft_model_for_expected_reward:
+                        tmp_logp_x0 = ft_model_pipe._subs_parameterization(tmp_logits_ft_model, latents_batch)
+                    else:
+                        tmp_logp_x0 = self._subs_parameterization(tmp_logits, latents_batch)
+                    for phi_i in range(phi):
+                        sample = F.gumbel_softmax(tmp_logp_x0, tau=tau, hard=True).argmax(dim=-1)
+                        sample = self._decode_latents(sample, batch_p, height, width, "pt")
+                        tmp_rewards[:, phi_i] = reward_fn(sample)
+                    tmp_rewards = logmeanexp(tmp_rewards * scale_cur, dim=-1) / scale_cur
+                
+                logits[j:j+batch_p] = tmp_logits.detach()
+                logits_ft_model[j:j+batch_p] = tmp_logits_ft_model.detach()
+                rewards[j:j+batch_p] = tmp_rewards.detach()
+                log_twist[j:j+batch_p] = rewards[j:j+batch_p] * scale_cur
+                
+            if verbose:
+                print("Rewards: ", rewards)
+            
+            # Calculate weights
+            incremental_log_w = (log_prob_diffusion - log_prob_proposal) + (log_twist - log_twist_prev)
+            log_w += incremental_log_w
+            
+            # Now reshape log_w to (batches, num_particles)
+            log_w = log_w.reshape(batches, num_particles)
+            
+            if verbose:
+                print("log_prob_diffusion - log_prob_proposal: ", log_prob_diffusion - log_prob_proposal)
+                print("log_twist - log_twist_prev:", log_twist - log_twist_prev)
+                print("Incremental log weights: ", incremental_log_w)
+                print("Log weights: ", log_w)
+                print("Normalized weights: ", normalize_weights(log_w, dim=-1))
+            
+            # Resample particles
+            if verbose:
+                print(f"ESS: ", compute_ess_from_log_w(log_w, dim=-1))
+            
+            if resample_condition:
+                resample_indices = []
+                log_w_new = []
+                is_resampled = False
+                for batch in range(batches):
+                    resample_indices_batch, is_resampled_batch, log_w_batch = resample_fn(log_w[batch])
+                    resample_indices.append(resample_indices_batch + batch * num_particles)
+                    log_w_new.append(log_w_batch)
+                    is_resampled = is_resampled or is_resampled_batch
+                    
+                resample_indices = torch.cat(resample_indices, dim=0)
+                log_w = torch.cat(log_w_new, dim=0)
+                    
+                if is_resampled:
+                    latents = latents[resample_indices]
+                    logits = logits[resample_indices]
+                    logits_ft_model = logits_ft_model[resample_indices]
+                    rewards = rewards[resample_indices]
+                    log_twist = log_twist[resample_indices]
+                    
+                if verbose:
+                    print("Resample indices: ", resample_indices)
+                
+            if log_w.ndim == 2:
+                log_w = log_w.reshape(total_particles)
+
+            
+            # Propose new particles
+            approx_guidance = logits_ft_model - logits # this effectively makes logits_ft_model the proposal distribution
+            approx_guidance[..., self.codebook_size:] = 0.0 # avoid nan due to (inf - inf)
+            sched_out = self.scheduler.step_with_approx_guidance(
+                latents=latents,
+                logits=logits,
+                approx_guidance=approx_guidance, 
+                step=i,
+            )
+            latents, log_prob_proposal, log_prob_diffusion = (
+                sched_out.new_latents,
+                sched_out.log_prob_proposal,
+                sched_out.log_prob_diffusion,
+            )
+        
+        bar = enumerate(reversed(range(num_inference_steps)))
+        if not disable_progress_bar:
+            bar = tqdm(bar, leave=False)
+        for i, timestep in bar:
+            resample_condition = (i + 1) % resample_frequency == 0
+            scale_cur = lambdas[i] / kl_weight
+            scale_next = lambdas[i + 1] / kl_weight
+            if verbose:
+                print(f"scale_cur: {scale_cur}, scale_next: {scale_next}")
+            propagate()
+            print('\n\n')
+        
+        # Final SMC weights
+        scale_cur = lambdas[-1] / kl_weight
+        log_twist_prev = log_twist.clone()
+        for j in range(0, total_particles, batch_p):
+            latents_batch = latents[j:j+batch_p]
+            with torch.no_grad():
+                sample = self._decode_latents(latents_batch, batch_p, height, width, "pt")
+                tmp_rewards = reward_fn(sample)
+                rewards[j:j+batch_p] = tmp_rewards
+                log_twist[j:j+batch_p] = tmp_rewards * scale_cur
+                
+        if verbose:
+            print("Rewards: ", rewards)
+
+        # Calculate weights
+        incremental_log_w = (log_prob_diffusion - log_prob_proposal) + (log_twist - log_twist_prev)
+        log_w += incremental_log_w
+
+        # Now reshape everything to (batches, num_particles) for final strategy
+        log_w = log_w.reshape(batches, num_particles)
+        latents = latents.reshape(batches, num_particles, H, W)
+        rewards = rewards.reshape(batches, num_particles)
+        
+        if verbose:
+            print("log_prob_diffusion - log_prob_proposal: ", log_prob_diffusion - log_prob_proposal)
+            print("log_twist - log_twist_prev:", log_twist - log_twist_prev)
+            print("Incremental log weights: ", incremental_log_w)
+            print("Log weights: ", log_w)
+            print("Normalized weights: ", normalize_weights(log_w, dim=-1))
+        
+        if final_strategy == "multinomial":
+            final_indices = torch.multinomial(normalize_weights(log_w, dim=-1), num_samples=1).squeeze(-1)
+        elif final_strategy == "argmax_rewards":
+            final_indices = rewards.argmax(dim=-1)
+        elif final_strategy == "argmax_weights":
+            final_indices = log_w.argmax(dim=-1)
+        else:
+            raise NotImplementedError(f"Final strategy {final_strategy} is not implemented.")
+        
+        if verbose:
+            print("Final selected indices: ", final_indices)
+        
+        latents = latents[
+            torch.arange(batches, device=latents.device),
+            final_indices
+        ]
+        
+        # Decode latents
+        outputs = []
+        for j in range(0, batches, batch_p):
+            latents_batch = latents[j:j+batch_p]
+            outputs.extend(
+                self._decode_latents(latents_batch, batch_p, height, width, output_type) # type: ignore
+            )
+        if output_type == "pt":
+            outputs = torch.stack(outputs, dim=0)
+        return outputs
+    
+    def get_logits(self, latents, guidance_scale, resolution, encoder_hidden_states, micro_conds, prompt_embeds, timestep):
+        if guidance_scale > 1.0:
+            # Latents are duplicated to get both unconditional and conditional logits
+            model_input = torch.cat([latents] * 2) # type: ignore
+        else:
+            model_input = latents
+        # img_ids, text_ids are used for positional embeddings
+        if resolution == 1024: #args.resolution == 1024:
+            img_ids = _prepare_latent_image_ids(model_input.shape[0], model_input.shape[1],model_input.shape[2],model_input.device,model_input.dtype)
+        else:
+            img_ids = _prepare_latent_image_ids(model_input.shape[0],2*model_input.shape[1],2*model_input.shape[2],model_input.device,model_input.dtype)
+        txt_ids = torch.zeros(encoder_hidden_states.shape[1],3).to(device = encoder_hidden_states.device, dtype = encoder_hidden_states.dtype)
+        model_output = self.transformer(
+            hidden_states = model_input,
+            micro_conds=micro_conds,
+            pooled_projections=prompt_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            img_ids = img_ids,
+            txt_ids = txt_ids,
+            timestep = torch.tensor([timestep], device=model_input.device, dtype=torch.long),
+        )
+        if guidance_scale > 1.0:
+            uncond_logits, cond_logits = model_output.chunk(2)
+            model_output = uncond_logits + guidance_scale * (cond_logits - uncond_logits)
+        tmp_logits = torch.permute(model_output, (0, 2, 3, 1)).float()
+        pad_logits = torch.full(
+            (*tmp_logits.shape[:3], self.vocab_size - self.codebook_size),
+            -torch.inf, 
+            device=tmp_logits.device, dtype=tmp_logits.dtype
+        )
+        tmp_logits = torch.cat([tmp_logits, pad_logits], dim=-1)
+        return tmp_logits
+
+    def _decode_latents(self, latents, batch_size, height, width, output_type):
+        if output_type == "latent":
+            output = latents
+        else:
+            needs_upcasting = self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast # type: ignore
+            if needs_upcasting:
+                self.vqvae.float()
+            output = self.vqvae.decode(
+                latents,
+                force_not_quantize=True,
+                shape=(
+                    batch_size,
+                    height // self.vae_scale_factor,
+                    width // self.vae_scale_factor,
+                    self.vqvae.config.latent_channels, # type: ignore
+                ),
+            ).sample.clip(0, 1) # type: ignore
+            output = self.image_processor.postprocess(output, output_type)
+            if needs_upcasting:
+                self.vqvae.half()
+        return output
+            
+    def _decode_one_hot_latents(self, latents_one_hot, batch_size, height, width, output_type):
+        shape = (
+            batch_size,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+            self.vqvae.config.latent_channels, # type: ignore
+        )
+        codebook_size = self.transformer.config.codebook_size #type: ignore
+        
+        needs_upcasting = self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast # type: ignore
+        if needs_upcasting:
+            self.vqvae.float()
+        
+        # get quantized latent vectors
+        embedding = self.vqvae.quantize.embedding.weight
+        h: torch.Tensor = latents_one_hot[..., :codebook_size].to(embedding.dtype) @ embedding
+        h = h.view(shape)
+        # reshape back to match original input shape
+        h = h.permute(0, 3, 1, 2).contiguous()
+
+        # Setting lookup_from_codebook to False, as we already have the codebook embeddings in h
+        self.vqvae.config.lookup_from_codebook = False # type: ignore
+        output = self.vqvae.decode(
+            h, # type: ignore
+            force_not_quantize=True,
+        ).sample.clip(0, 1) # type: ignore
+        self.vqvae.config.lookup_from_codebook = True # type: ignore
+        
+        output = self.image_processor.postprocess(output, output_type)
+
+        if needs_upcasting:
+            self.vqvae.half()
+            
+        return output
+    
+    def _subs_parameterization(self, logits, latents):
+        B, H, W, C = logits.shape
+        logits = logits.view(B, H * W, C)
+        assert latents.shape == (B, H, W)
+        latents = latents.view(B, H * W)
+        
+        logits = logits - torch.logsumexp(logits, dim=-1,
+                                        keepdim=True)
+        unmasked_indices = (latents != self.mask_index)
+        logits[unmasked_indices] = -torch.inf
+        logits[unmasked_indices, latents[unmasked_indices]] = 0
+        
+        logits = logits.view(B, H, W, C)
+        return logits

src/smc/resampling.py

@@ -0,0 +1,149 @@
+from typing import Callable, Tuple
+
+import torch
+
+
+def compute_ess(w, dim=-1):
+    ess = (w.sum(dim=dim))**2 / torch.sum(w**2, dim=dim)
+    return ess
+
+def compute_ess_from_log_w(log_w, dim=-1):
+    return compute_ess(normalize_weights(log_w, dim=dim), dim=dim)
+
+def normalize_weights(log_weights, dim=-1):
+    return torch.exp(normalize_log_weights(log_weights, dim=dim))
+
+def normalize_log_weights(log_weights, dim=-1):
+    log_weights = log_weights - log_weights.max(dim=dim, keepdims=True)[0]
+    log_weights = log_weights - torch.logsumexp(log_weights, dim=dim, keepdims=True) # type: ignore
+    return log_weights
+
+def stratified_resample(log_weights: torch.Tensor):
+    N = log_weights.shape[0]
+    weights = normalize_weights(log_weights)
+    cdf = torch.cumsum(weights, dim=0)
+
+    # Stratified uniform samples
+    u = (torch.arange(N, dtype=torch.float32, device=log_weights.device) + torch.rand(N, device=log_weights.device)) / N
+
+    indices = torch.searchsorted(cdf, u, right=True)
+    return indices
+
+def systematic_resample(log_weights: torch.Tensor, normalized=True):
+    N = log_weights.shape[0]
+    weights = normalize_weights(log_weights)
+    cdf = torch.cumsum(weights, dim=0)
+
+    # Systematic uniform samples
+    u0 = torch.rand(1, device=log_weights.device) / N
+    u = u0 + torch.arange(N, dtype=torch.float32, device=log_weights.device) / N
+
+    indices = torch.searchsorted(cdf, u, right=True)
+    return indices
+
+def multinomial_resample(log_weights: torch.Tensor, normalized=True):
+    N = log_weights.shape[0]
+    weights = normalize_weights(log_weights)
+    resampled_indices = torch.multinomial(weights, N, replacement=True)
+    return resampled_indices
+
+def partial_resample(log_weights: torch.Tensor,
+                     resample_fn: Callable[[torch.Tensor], torch.Tensor],
+                     M: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Perform partial resampling on a set of particles using PyTorch.
+
+    Args:
+        log_weights (torch.Tensor): 1D tensor of shape (K,) containing log-weights.
+        resample_fn (callable): function that takes log_weights and n_samples,
+                                returning a tensor of shape (n_samples,) of sampled indices.
+        M (int): total number of particles to resample.
+
+    Returns:
+        new_indices (torch.Tensor): 1D tensor of shape (K,) mapping each output slot to
+                                    an original particle index.
+        new_log_weights (torch.Tensor): 1D tensor of shape (K,) of updated log-weights.
+    """
+    K = log_weights.numel()
+
+    # Convert log-weights to normalized weights
+    log_weights = normalize_log_weights(log_weights)
+    weights = torch.exp(log_weights)
+
+    # Determine how many high and low weights to resample
+    M_hi = 1 # M // 2
+    M_lo = M - M_hi
+
+    # Get indices of highest and lowest weights
+    _, hi_idx = torch.topk(weights, M_hi, largest=True)
+    _, lo_idx = torch.topk(weights, M_lo, largest=False)
+    I = torch.cat([hi_idx, lo_idx])  # indices selected for resampling
+
+    # Perform multinomial resampling only on selected subset
+    # resample_fn expects log-weights of the subset
+    subset_logw = log_weights[I]
+    local_sampled = resample_fn(subset_logw)  # indices in [0, len(I))
+    # Map back to original indices
+    sampled = I[local_sampled]
+
+    # Build new index mapping: default to identity (retain original)
+    new_indices = torch.arange(K, device=log_weights.device)
+    new_indices[I] = sampled
+
+    # Compute new uniform weight for resampled particles
+    total_I_weight = weights[I].sum()
+    uniform_weight = total_I_weight / M
+
+    # Prepare new log-weights
+    new_log_weight = torch.empty_like(log_weights)
+    # For non-resampled, keep original log-weights
+    mask = torch.ones(K, dtype=torch.bool, device=log_weights.device)
+    mask[I] = False
+    new_log_weight[mask] = log_weights[mask]
+    # For resampled, assign uniform log-weight
+    new_log_weight[I] = torch.log(uniform_weight)
+
+    return new_indices, new_log_weight
+
+
+def resample(log_w, ess_threshold=None, partial=False):
+    """
+    Resample the log weights and return the indices of the resampled particles.
+
+    Parameters
+    ----------
+    log_w : array_like
+        The log weights of the particles.
+    ess_threshold : float, optional
+        The effective sample size (ESS) threshold. If the ESS is below this
+        threshold, resampling is performed. If None, no resampling is
+        performed.
+    partial : bool, optional
+        If True, the resampling is performed on the partial weights. If False,
+        the resampling is performed on the full weights.
+
+    Returns
+    -------
+    array_like
+        The indices of the resampled particles.
+    """
+    base_sampling_fn = systematic_resample
+    N = log_w.size(0)
+    ess = compute_ess_from_log_w(log_w)
+    if ess_threshold is not None and ess >= ess_threshold * N:
+        # Skip resampling as ess is not below the threshold
+        return (
+            torch.arange(N, device=log_w.device),
+            False,
+            log_w
+        )
+    if partial:
+        resample_indices, log_w = partial_resample(log_w, base_sampling_fn, N // 2)
+    else:
+        resample_indices = base_sampling_fn(log_w)
+        log_w = torch.zeros_like(log_w)
+    return (
+        resample_indices,
+        True,
+        log_w
+    )

src/smc/scheduler.py

@@ -0,0 +1,368 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union, List
+
+import math
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from src.meissonic.scheduler import mask_by_random_topk
+
+
+@dataclass
+class SchedulerStepOutput:
+    new_latents: torch.Tensor
+
+
+@dataclass
+class SchedulerApproxGuidanceOutput:
+    new_latents: torch.Tensor
+    log_prob_proposal: torch.Tensor
+    log_prob_diffusion: torch.Tensor
+
+
+class BaseScheduler(ABC):
+    @abstractmethod
+    def step(
+        self,
+        latents: torch.Tensor,
+        step: int,
+        logits: torch.Tensor,
+    ) -> SchedulerStepOutput:
+        pass
+    
+    @abstractmethod
+    def set_timesteps(self, num_inference_steps: int):
+        pass
+
+    @abstractmethod
+    def step_with_approx_guidance(
+        self,
+        latents: torch.Tensor,
+        step: int,
+        logits: torch.Tensor,
+        approx_guidance: torch.Tensor,
+    ) -> SchedulerApproxGuidanceOutput:
+        pass
+
+
+def sum_masked_logits(
+    logits: torch.Tensor,
+    preds: torch.Tensor,
+    mask: torch.Tensor
+) -> torch.Tensor:
+    """
+    Sum logits at `preds` indices, masked by `mask`, handling invalid `preds`.
+
+    Args:
+        logits: Tensor of shape (B, H, W, C) - logits over C classes.
+        preds: Tensor of shape (B, H, W) - predicted class indices.
+        mask: Tensor of shape (B, H, W) - binary mask to include positions.
+
+    Returns:
+        Tensor of shape (B,) - sum of selected logits per batch item.
+    """
+    B, H, W, C = logits.shape
+    # Ensure preds are in valid index range [0, C-1]
+    valid = (preds >= 0) & (preds <= preds[mask].max())
+    # Replace invalid preds with a dummy index (0), which we will mask later
+    safe_preds = preds.masked_fill(~valid, 0)
+    # Gather logits at predicted indices
+    selected = torch.gather(logits, dim=3, index=safe_preds.unsqueeze(-1)).squeeze(-1)
+    # Zero out contributions from invalid preds and masked positions
+    selected = selected * valid * mask
+    # Sum over H, W dimension
+    return selected.sum(dim=(1, 2))
+
+def log1mexp(x: torch.Tensor) -> torch.Tensor:
+    """
+    Numerically stable computation of log(1 - exp(x)) for x < 0.
+    """
+    return torch.where(
+        x > -1,
+        torch.log(-torch.expm1(x)),
+        torch.log1p(-torch.exp(x)),
+    )
+
+
+class MeissonicScheduler(BaseScheduler):
+    def __init__(self, 
+            mask_token_id: int, 
+            masking_schedule: str = "cosine",
+        ):
+        self.mask_token_id = mask_token_id
+        self.masking_schedule = masking_schedule
+    
+    def set_timesteps(self, num_inference_steps: int, temperature: Union[int, Tuple[int, int], List[int]] = (2, 0), device='cuda'):
+        self.num_inference_steps = num_inference_steps
+        self.timesteps = torch.arange(num_inference_steps, device=device).flip(0)
+        if isinstance(temperature, (tuple, list)):
+            self.temperatures = torch.linspace(temperature[0], temperature[1], num_inference_steps, device=device)
+        else:
+            self.temperatures = torch.linspace(temperature, 0.01, num_inference_steps, device=device)
+    
+    def step(
+        self,
+        latents: torch.Tensor,
+        step: int,
+        logits: torch.Tensor,
+    ) -> SchedulerStepOutput:
+        batch_size, height, width, vocab_size = logits.shape
+        sample = latents.reshape(batch_size, height * width)
+        model_output = logits.reshape(batch_size, height * width, vocab_size)
+
+        unknown_map = sample == self.mask_token_id
+
+        probs = model_output.softmax(dim=-1)
+
+        device = probs.device
+        probs_ = probs
+        if probs_.device.type == "cpu" and probs_.dtype != torch.float32:
+            probs_ = probs_.float()  # multinomial is not implemented for cpu half precision
+        probs_ = probs_.reshape(-1, probs.size(-1))
+        pred_original_sample = torch.multinomial(probs_, 1).to(device=device)
+        pred_original_sample = pred_original_sample[:, 0].view(*probs.shape[:-1])
+        pred_original_sample = torch.where(unknown_map, pred_original_sample, sample)
+
+        timestep = self.num_inference_steps - 1 - step
+        if timestep == 0:
+            prev_sample = pred_original_sample
+        else:
+            seq_len = sample.shape[1]
+            step_idx = (self.timesteps == timestep).nonzero()
+            ratio = (step_idx + 1) / len(self.timesteps)
+
+            if self.masking_schedule == "cosine":
+                mask_ratio = torch.cos(ratio * math.pi / 2)
+            elif self.masking_schedule == "linear":
+                mask_ratio = 1 - ratio
+            else:
+                raise ValueError(f"unknown masking schedule {self.masking_schedule}")
+
+            mask_len = (seq_len * mask_ratio).floor()
+            # do not mask more than amount previously masked
+            mask_len = torch.min(unknown_map.sum(dim=-1, keepdim=True) - 1, mask_len)
+            # mask at least one
+            mask_len = torch.max(torch.tensor([1], device=model_output.device), mask_len)
+
+            selected_probs = torch.gather(probs, -1, pred_original_sample[:, :, None])[:, :, 0]
+            # Ignores the tokens given in the input by overwriting their confidence.
+            selected_probs = torch.where(unknown_map, selected_probs, torch.finfo(selected_probs.dtype).max)
+
+            masking = mask_by_random_topk(mask_len, selected_probs, self.temperatures[step_idx].item())
+
+            # Masks tokens with lower confidence.
+            prev_sample = torch.where(masking, self.mask_token_id, pred_original_sample)
+
+        print("Unmasked:", (prev_sample != self.mask_token_id).sum(dim=1))
+        prev_sample = prev_sample.reshape(batch_size, height, width)
+        pred_original_sample = pred_original_sample.reshape(batch_size, height, width)
+        
+        return SchedulerStepOutput(new_latents=prev_sample)
+        
+    
+    def step_with_approx_guidance(
+        self,
+        latents: torch.Tensor,
+        step: int,
+        logits: torch.Tensor,
+        approx_guidance: torch.Tensor,
+    ) -> SchedulerApproxGuidanceOutput:
+        proposal_logits = logits + approx_guidance
+        sched_out = self.step(latents, step, proposal_logits)
+        new_latents = sched_out.new_latents
+        
+        newly_filled_positions = (latents != new_latents)
+        print("Newly filled positions:", newly_filled_positions.sum(dim=(1, 2)))
+        
+        log_prob_proposal = sum_masked_logits(
+            logits=proposal_logits.log_softmax(dim=-1),
+            preds=new_latents,
+            mask=newly_filled_positions,
+        )
+        log_prob_diffusion = sum_masked_logits(
+            logits=logits.log_softmax(dim=-1),
+            preds=new_latents,
+            mask=newly_filled_positions,
+        )
+        print("log prob proposal:", log_prob_proposal)
+        print("log prob diffusion:", log_prob_diffusion)
+        return SchedulerApproxGuidanceOutput(
+            new_latents,
+            log_prob_proposal,
+            log_prob_diffusion,
+        )
+
+
+class ReMDMScheduler(BaseScheduler):
+    def __init__(
+        self,
+        schedule,
+        remask_strategy,
+        eta,
+        mask_token_id,
+        temperature=1.0,
+    ):
+        self.schedule = schedule
+        self.remask_strategy = remask_strategy
+        self.eta = eta 
+        self.temperature = temperature
+        self.mask_token_id = mask_token_id
+    
+    def set_timesteps(self, num_inference_steps: int):
+        self.num_inference_steps = num_inference_steps
+        if self.schedule == "linear":
+            self.alphas = 1 - torch.linspace(0, 1, num_inference_steps + 1)
+        elif self.schedule == "cosine":
+            self.alphas = 1 - torch.cos((math.pi/2) * (1 - torch.linspace(0, 1, num_inference_steps + 1)))
+        else:
+            raise ValueError(f"unknown masking schedule {self.schedule}")
+    
+    def step(
+        self,
+        latents: torch.Tensor,
+        step: int,
+        logits: torch.Tensor,
+    ) -> SchedulerStepOutput:
+        B, H, W, C = logits.shape
+        assert latents.shape == (B, H, W)
+        
+        latents = latents.reshape(B, H*W)
+        logits = logits.reshape(B, H*W, C)
+        
+        t = self.num_inference_steps - step
+        s = t - 1
+        
+        alpha_t = self.alphas[t]
+        alpha_s = self.alphas[s]
+        sigma_t_max = torch.clamp_max((1 - alpha_s) / alpha_t, 1.0)
+        if self.remask_strategy == "max_cap":
+            sigma_t = torch.clamp_max(sigma_t_max, self.eta)
+        elif self.remask_strategy == "rescale":
+            sigma_t = sigma_t_max * self.eta
+        else:
+            raise ValueError(f"unknown masking schedule {self.remask_strategy}")
+        
+        # z_t != m
+        x_theta = F.one_hot(latents, num_classes=C).float()
+        logits_z_t_neq_m = (
+            torch.log(x_theta) +
+            torch.log(1 - sigma_t)
+        )
+        logits_z_t_neq_m[..., self.mask_token_id] = (
+            torch.log(sigma_t)
+        )
+        
+        # z_t = m
+        log_x_theta = (logits / self.temperature).log_softmax(dim=-1)
+        logits_z_t_eq_m = (
+            log_x_theta + 
+            torch.log((alpha_s - (1 - sigma_t) * alpha_t) / (1 - alpha_t))
+        )
+        logits_z_t_eq_m[..., self.mask_token_id] = (
+            torch.log((1 - alpha_s - sigma_t * alpha_t) / (1 - alpha_t))
+        )
+        
+        z_t_neq_m = (latents != self.mask_token_id)
+        p_theta_logits = torch.where(
+            z_t_neq_m.unsqueeze(-1).expand(-1, -1, C),
+            logits_z_t_neq_m,
+            logits_z_t_eq_m,
+        )
+        assert torch.allclose(torch.exp(p_theta_logits).sum(dim=-1), torch.ones(B, H*W, device=logits.device)), (torch.exp(p_theta_logits).sum(dim=-1) - torch.ones(B, H*W, device=logits.device)).abs().max()
+        diffusion_dist = torch.distributions.Categorical(logits=p_theta_logits) # type: ignore
+        new_latents = diffusion_dist.sample()
+        print("Unmasked:", (new_latents != self.mask_token_id).sum(dim=1))
+        return SchedulerStepOutput(new_latents.reshape(B, H, W))
+    
+    def step_with_approx_guidance(
+        self,
+        latents: torch.Tensor,
+        step: int,
+        logits: torch.Tensor,
+        approx_guidance: torch.Tensor,
+    ) -> SchedulerApproxGuidanceOutput:
+        B, H, W, C = logits.shape
+        assert latents.shape == (B, H, W)
+        assert approx_guidance.shape == (B, H, W, C)
+        
+        latents = latents.reshape(B, H*W)
+        logits = logits.reshape(B, H*W, C)
+        approx_guidance = approx_guidance.reshape(B, H*W, C)
+        
+        t = self.num_inference_steps - step
+        s = t - 1
+        
+        alpha_t = self.alphas[t]
+        alpha_s = self.alphas[s]
+        sigma_t_max = torch.clamp_max((1 - alpha_s) / alpha_t, 1.0)
+        if self.remask_strategy == "max_cap":
+            sigma_t = torch.clamp_max(sigma_t_max, self.eta)
+        elif self.remask_strategy == "rescale":
+            sigma_t = sigma_t_max * self.eta
+        else:
+            raise ValueError(f"unknown masking schedule {self.remask_strategy}")
+        
+        # z_t != m
+        x_theta = F.one_hot(latents, num_classes=C).float()
+        logits_z_t_neq_m = (
+            torch.log(x_theta) +
+            torch.log(1 - sigma_t)
+        )
+        logits_z_t_neq_m[..., self.mask_token_id] = (
+            torch.log(sigma_t)
+        )
+        
+        # z_t = m
+        log_x_theta = (logits / self.temperature).log_softmax(dim=-1)
+        logits_z_t_eq_m = (
+            log_x_theta + 
+            torch.log((alpha_s - (1 - sigma_t) * alpha_t) / (1 - alpha_t))
+        )
+        logits_z_t_eq_m[..., self.mask_token_id] = (
+            torch.log((1 - alpha_s - sigma_t * alpha_t) / (1 - alpha_t))
+        )
+        
+        z_t_neq_m = (latents != self.mask_token_id)
+        p_theta_logits = torch.where(
+            z_t_neq_m.unsqueeze(-1).expand(-1, -1, C),
+            logits_z_t_neq_m,
+            logits_z_t_eq_m,
+        )
+        assert torch.allclose(torch.exp(p_theta_logits).sum(dim=-1), torch.ones(B, H*W, device=logits.device))
+        
+        proposal_logits = (p_theta_logits + approx_guidance).log_softmax(dim=-1)
+        assert torch.allclose(torch.exp(proposal_logits).sum(dim=-1), torch.ones(B, H*W, device=logits.device))
+        
+        # modify proposal logits to have the same mask schedule as the original logits
+        proposal_logits[..., :self.mask_token_id] += (
+            torch.logsumexp(p_theta_logits[..., :self.mask_token_id], dim=(1, 2), keepdim=True) - 
+            torch.logsumexp(proposal_logits[..., :self.mask_token_id], dim=(1, 2), keepdim=True)
+        )
+        proposal_logits[..., :self.mask_token_id] = torch.where(
+            proposal_logits[..., :self.mask_token_id].logsumexp(dim=-1, keepdim=True) >= 0,
+            proposal_logits[..., :self.mask_token_id].log_softmax(dim=-1),
+            proposal_logits[..., :self.mask_token_id]
+        )
+        assert not (proposal_logits[..., :self.mask_token_id].logsumexp(dim=-1) > 1e-6).any(), proposal_logits[..., :self.mask_token_id].logsumexp(dim=-1).max()
+        proposal_logits[..., self.mask_token_id] = (
+            log1mexp(proposal_logits[..., :self.mask_token_id].logsumexp(dim=-1).clamp_max(0))
+        )
+        assert torch.allclose(torch.exp(proposal_logits).sum(dim=-1), torch.ones(B, H*W, device=logits.device)), (torch.exp(proposal_logits).sum(dim=-1) - torch.ones(B, H*W, device=logits.device)).abs().max()
+        # modify proposal logits to have the same mask schedule as the original logits
+        
+        proposal_dist = torch.distributions.Categorical(logits=proposal_logits) # type: ignore
+        diffusion_dist = torch.distributions.Categorical(logits=p_theta_logits) # type: ignore
+        
+        new_latents = proposal_dist.sample()
+        
+        log_prob_proposal = proposal_dist.log_prob(new_latents).sum(dim=1)
+        log_prob_diffusion = diffusion_dist.log_prob(new_latents).sum(dim=1)
+        
+        print("Unmasked:", (new_latents != self.mask_token_id).sum(dim=1))
+        return SchedulerApproxGuidanceOutput(
+            new_latents.reshape(B, H, W),
+            log_prob_proposal,
+            log_prob_diffusion,
+        )

src/smc/transformer.py

@@ -0,0 +1,1119 @@
+# Copyright 2024 Black Forest Labs, The HuggingFace Team, The InstantX Team and The MeissonFlow Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward, BasicTransformerBlock, SkipFFTransformerBlock
+from diffusers.models.attention_processor import (
+    Attention,
+    AttentionProcessor,
+    FluxAttnProcessor2_0,
+    # FusedFluxAttnProcessor2_0,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, GlobalResponseNorm, RMSNorm
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings,TimestepEmbedding, get_timestep_embedding #,FluxPosEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput 
+from diffusers.models.resnet import Downsample2D, Upsample2D
+
+from typing import List
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+
+def get_3d_rotary_pos_embed(
+    embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    RoPE for video tokens with 3D structure.
+
+    Args:
+    embed_dim: (`int`):
+        The embedding dimension size, corresponding to hidden_size_head.
+    crops_coords (`Tuple[int]`):
+        The top-left and bottom-right coordinates of the crop.
+    grid_size (`Tuple[int]`):
+        The grid size of the spatial positional embedding (height, width).
+    temporal_size (`int`):
+        The size of the temporal dimension.
+    theta (`float`):
+        Scaling factor for frequency computation.
+    use_real (`bool`):
+        If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+
+    Returns:
+        `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
+    """
+    start, stop = crops_coords
+    grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)
+    grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)
+    grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)
+
+    # Compute dimensions for each axis
+    dim_t = embed_dim // 4
+    dim_h = embed_dim // 8 * 3
+    dim_w = embed_dim // 8 * 3
+
+    # Temporal frequencies
+    freqs_t = 1.0 / (theta ** (torch.arange(0, dim_t, 2).float() / dim_t))
+    grid_t = torch.from_numpy(grid_t).float()
+    freqs_t = torch.einsum("n , f -> n f", grid_t, freqs_t)
+    freqs_t = freqs_t.repeat_interleave(2, dim=-1)
+
+    # Spatial frequencies for height and width
+    freqs_h = 1.0 / (theta ** (torch.arange(0, dim_h, 2).float() / dim_h))
+    freqs_w = 1.0 / (theta ** (torch.arange(0, dim_w, 2).float() / dim_w))
+    grid_h = torch.from_numpy(grid_h).float()
+    grid_w = torch.from_numpy(grid_w).float()
+    freqs_h = torch.einsum("n , f -> n f", grid_h, freqs_h)
+    freqs_w = torch.einsum("n , f -> n f", grid_w, freqs_w)
+    freqs_h = freqs_h.repeat_interleave(2, dim=-1)
+    freqs_w = freqs_w.repeat_interleave(2, dim=-1)
+
+    # Broadcast and concatenate tensors along specified dimension
+    def broadcast(tensors, dim=-1):
+        num_tensors = len(tensors)
+        shape_lens = {len(t.shape) for t in tensors}
+        assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+        shape_len = list(shape_lens)[0]
+        dim = (dim + shape_len) if dim < 0 else dim
+        dims = list(zip(*(list(t.shape) for t in tensors)))
+        expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+        assert all(
+            [*(len(set(t[1])) <= 2 for t in expandable_dims)]
+        ), "invalid dimensions for broadcastable concatenation"
+        max_dims = [(t[0], max(t[1])) for t in expandable_dims]
+        expanded_dims = [(t[0], (t[1],) * num_tensors) for t in max_dims]
+        expanded_dims.insert(dim, (dim, dims[dim]))
+        expandable_shapes = list(zip(*(t[1] for t in expanded_dims)))
+        tensors = [t[0].expand(*t[1]) for t in zip(tensors, expandable_shapes)]
+        return torch.cat(tensors, dim=dim)
+
+    freqs = broadcast((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1)
+
+    t, h, w, d = freqs.shape
+    freqs = freqs.view(t * h * w, d)
+
+    # Generate sine and cosine components
+    sin = freqs.sin()
+    cos = freqs.cos()
+
+    if use_real:
+        return cos, sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs_cis
+
+
+def get_2d_rotary_pos_embed(embed_dim, crops_coords, grid_size, use_real=True):
+    """
+    RoPE for image tokens with 2d structure.
+
+    Args:
+    embed_dim: (`int`):
+        The embedding dimension size
+    crops_coords (`Tuple[int]`)
+        The top-left and bottom-right coordinates of the crop.
+    grid_size (`Tuple[int]`):
+        The grid size of the positional embedding.
+    use_real (`bool`):
+        If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+
+    Returns:
+        `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`.
+    """
+    start, stop = crops_coords
+    grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)
+    grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)  # [2, W, H]
+
+    grid = grid.reshape([2, 1, *grid.shape[1:]])
+    pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real)
+    return pos_embed
+
+
+def get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=False):
+    assert embed_dim % 4 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_rotary_pos_embed(
+        embed_dim // 2, grid[0].reshape(-1), use_real=use_real
+    )  # (H*W, D/2) if use_real else (H*W, D/4)
+    emb_w = get_1d_rotary_pos_embed(
+        embed_dim // 2, grid[1].reshape(-1), use_real=use_real
+    )  # (H*W, D/2) if use_real else (H*W, D/4)
+
+    if use_real:
+        cos = torch.cat([emb_h[0], emb_w[0]], dim=1)  # (H*W, D)
+        sin = torch.cat([emb_h[1], emb_w[1]], dim=1)  # (H*W, D)
+        return cos, sin
+    else:
+        emb = torch.cat([emb_h, emb_w], dim=1)  # (H*W, D/2)
+        return emb
+
+
+def get_2d_rotary_pos_embed_lumina(embed_dim, len_h, len_w, linear_factor=1.0, ntk_factor=1.0):
+    assert embed_dim % 4 == 0
+
+    emb_h = get_1d_rotary_pos_embed(
+        embed_dim // 2, len_h, linear_factor=linear_factor, ntk_factor=ntk_factor
+    )  # (H, D/4)
+    emb_w = get_1d_rotary_pos_embed(
+        embed_dim // 2, len_w, linear_factor=linear_factor, ntk_factor=ntk_factor
+    )  # (W, D/4)
+    emb_h = emb_h.view(len_h, 1, embed_dim // 4, 1).repeat(1, len_w, 1, 1)  # (H, W, D/4, 1)
+    emb_w = emb_w.view(1, len_w, embed_dim // 4, 1).repeat(len_h, 1, 1, 1)  # (H, W, D/4, 1)
+
+    emb = torch.cat([emb_h, emb_w], dim=-1).flatten(2)  # (H, W, D/2)
+    return emb
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[np.ndarray, int],
+    theta: float = 10000.0,
+    use_real=False,
+    linear_factor=1.0,
+    ntk_factor=1.0,
+    repeat_interleave_real=True,
+    freqs_dtype=torch.float32,  # torch.float32 (hunyuan, stable audio), torch.float64 (flux)
+):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+    index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+    data type.
+
+    Args:
+        dim (`int`): Dimension of the frequency tensor.
+        pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+        theta (`float`, *optional*, defaults to 10000.0):
+            Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (`bool`, *optional*):
+            If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+        linear_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor for the context extrapolation. Defaults to 1.0.
+        ntk_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.
+        repeat_interleave_real (`bool`, *optional*, defaults to `True`):
+            If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.
+            Otherwise, they are concateanted with themselves.
+        freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):
+            the dtype of the frequency tensor.
+    Returns:
+        `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+    """
+    assert dim % 2 == 0
+
+    if isinstance(pos, int):
+        pos = np.arange(pos)
+    theta = theta * ntk_factor
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype)[: (dim // 2)] / dim)) / linear_factor  # [D/2]
+    t = torch.from_numpy(pos).to(freqs.device)  # type: ignore  # [S]
+    freqs = torch.outer(t, freqs)  # type: ignore   # [S, D/2]
+    if use_real and repeat_interleave_real:
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+    elif use_real:
+        freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float()  # [S, D]
+        freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs).float()  # complex64     # [S, D/2]
+        return freqs_cis
+
+
+class FluxPosEmbed(nn.Module):
+    # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
+    def __init__(self, theta: int, axes_dim: List[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        n_axes = ids.shape[-1]
+        cos_out = []
+        sin_out = []
+        pos = ids.squeeze().float().cpu().numpy()
+        is_mps = ids.device.type == "mps"
+        freqs_dtype = torch.float32 if is_mps else torch.float64
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i], pos[:, i], repeat_interleave_real=True, use_real=True, freqs_dtype=freqs_dtype
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+
+class FusedFluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "FusedFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        # `context` projections.
+        if encoder_hidden_states is not None:
+            encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+            split_size = encoder_qkv.shape[-1] // 3
+            (
+                encoder_hidden_states_query_proj,
+                encoder_hidden_states_key_proj,
+                encoder_hidden_states_value_proj,
+            ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+
+@maybe_allow_in_graph
+class  SingleTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+        processor = FluxAttnProcessor2_0()
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=processor,
+            qk_norm="rms_norm",
+            eps=1e-6,
+            pre_only=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+    ):
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return hidden_states
+
+@maybe_allow_in_graph
+class TransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+
+        self.norm1_context = AdaLayerNormZero(dim)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            processor = FluxAttnProcessor2_0()
+        else:
+            raise ValueError(
+                "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+            )
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=processor,
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        # self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="swiglu")
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        # self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="swiglu")
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+    ):
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+        # Attention.
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = hidden_states + ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class UVit2DConvEmbed(nn.Module):
+    def __init__(self, in_channels, block_out_channels, vocab_size, elementwise_affine, eps, bias):
+        super().__init__()
+        self.embeddings = nn.Embedding(vocab_size, in_channels)
+        self.layer_norm = RMSNorm(in_channels, eps, elementwise_affine)
+        self.conv = nn.Conv2d(in_channels, block_out_channels, kernel_size=1, bias=bias)
+
+    def forward(self, input_ids):
+        if input_ids.is_floating_point():
+            embeddings = input_ids @ self.embeddings.weight
+        else:
+            embeddings = self.embeddings(input_ids)
+        embeddings = self.layer_norm(embeddings)
+        embeddings = embeddings.permute(0, 3, 1, 2)
+        embeddings = self.conv(embeddings)
+        return embeddings
+
+class ConvMlmLayer(nn.Module):
+    def __init__(
+        self,
+        block_out_channels: int,
+        in_channels: int,
+        use_bias: bool,
+        ln_elementwise_affine: bool,
+        layer_norm_eps: float,
+        codebook_size: int,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(block_out_channels, in_channels, kernel_size=1, bias=use_bias)
+        self.layer_norm = RMSNorm(in_channels, layer_norm_eps, ln_elementwise_affine)
+        self.conv2 = nn.Conv2d(in_channels, codebook_size, kernel_size=1, bias=use_bias)
+
+    def forward(self, hidden_states):
+        hidden_states = self.conv1(hidden_states)
+        hidden_states = self.layer_norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        logits = self.conv2(hidden_states)
+        return logits
+
+class SwiGLU(nn.Module):
+    r"""
+    A [variant](https://arxiv.org/abs/2002.05202) of the gated linear unit activation function. It's similar to `GEGLU`
+    but uses SiLU / Swish instead of GeLU.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)
+        self.activation = nn.SiLU()
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_states, gate = hidden_states.chunk(2, dim=-1)
+        return hidden_states * self.activation(gate)
+
+class ConvNextBlock(nn.Module):
+    def __init__(
+        self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4
+    ):
+        super().__init__()
+        self.depthwise = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            padding=1,
+            groups=channels,
+            bias=use_bias,
+        )
+        self.norm = RMSNorm(channels, layer_norm_eps, ln_elementwise_affine)
+        self.channelwise_linear_1 = nn.Linear(channels, int(channels * res_ffn_factor), bias=use_bias)
+        self.channelwise_act = nn.GELU()
+        self.channelwise_norm = GlobalResponseNorm(int(channels * res_ffn_factor))
+        self.channelwise_linear_2 = nn.Linear(int(channels * res_ffn_factor), channels, bias=use_bias)
+        self.channelwise_dropout = nn.Dropout(hidden_dropout)
+        self.cond_embeds_mapper = nn.Linear(hidden_size, channels * 2, use_bias)
+
+    def forward(self, x, cond_embeds):
+        x_res = x
+
+        x = self.depthwise(x)
+
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+
+        x = self.channelwise_linear_1(x)
+        x = self.channelwise_act(x)
+        x = self.channelwise_norm(x)
+        x = self.channelwise_linear_2(x)
+        x = self.channelwise_dropout(x)
+
+        x = x.permute(0, 3, 1, 2)
+
+        x = x + x_res
+
+        scale, shift = self.cond_embeds_mapper(F.silu(cond_embeds)).chunk(2, dim=1)
+        x = x * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+
+        return x
+
+class Simple_UVitBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        ln_elementwise_affine,
+        layer_norm_eps,
+        use_bias,
+        downsample: bool,
+        upsample: bool,
+    ):
+        super().__init__()
+
+        if downsample:
+            self.downsample = Downsample2D(
+                channels,
+                use_conv=True,
+                padding=0,
+                name="Conv2d_0",
+                kernel_size=2,
+                norm_type="rms_norm",
+                eps=layer_norm_eps,
+                elementwise_affine=ln_elementwise_affine,
+                bias=use_bias,
+            )
+        else:
+            self.downsample = None
+
+        if upsample:
+            self.upsample = Upsample2D(
+                channels,
+                use_conv_transpose=True,
+                kernel_size=2,
+                padding=0,
+                name="conv",
+                norm_type="rms_norm",
+                eps=layer_norm_eps,
+                elementwise_affine=ln_elementwise_affine,
+                bias=use_bias,
+                interpolate=False,
+            )
+        else:
+            self.upsample = None
+
+    def forward(self, x):
+        # print("before,", x.shape)
+        if self.downsample is not None:
+            # print('downsample')
+            x = self.downsample(x)
+
+        if self.upsample is not None:
+            # print('upsample')
+            x = self.upsample(x)
+        # print("after,", x.shape)
+        return x
+
+class Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    """
+    The Transformer model introduced in Flux.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Parameters:
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
+        num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
+        guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
+    """
+
+    _supports_gradient_checkpointing = False #True 
+    # Due to NotImplementedError: DDPOptimizer backend: Found a higher order op in the graph. This is not supported. Please turn off DDP optimizer using torch._dynamo.config.optimize_ddp=False. Note that this can cause performance degradation because there will be one bucket for the entire Dynamo graph. 
+    # Please refer to this issue - https://github.com/pytorch/pytorch/issues/104674.
+    _no_split_modules = ["TransformerBlock", "SingleTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = 768,
+        guidance_embeds: bool = False, # unused in our implementation
+        axes_dims_rope: Tuple[int] = (16, 56, 56),
+        vocab_size: int = 8256,
+        codebook_size: int = 8192,
+        downsample: bool = False,
+        upsample: bool = False,
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim 
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+     
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+
+        self.gradient_checkpointing = False
+
+        in_channels_embed = self.inner_dim 
+        ln_elementwise_affine = True
+        layer_norm_eps = 1e-06
+        use_bias = False
+        micro_cond_embed_dim = 1280
+        self.embed = UVit2DConvEmbed(
+            in_channels_embed, self.inner_dim, self.config.vocab_size, ln_elementwise_affine, layer_norm_eps, use_bias
+        )
+        self.mlm_layer = ConvMlmLayer(
+            self.inner_dim, in_channels_embed, use_bias, ln_elementwise_affine, layer_norm_eps, self.config.codebook_size
+        )
+        self.cond_embed = TimestepEmbedding( 
+            micro_cond_embed_dim + self.config.pooled_projection_dim, self.inner_dim, sample_proj_bias=use_bias
+        )
+        self.encoder_proj_layer_norm = RMSNorm(self.inner_dim, layer_norm_eps, ln_elementwise_affine)
+        self.project_to_hidden_norm = RMSNorm(in_channels_embed, layer_norm_eps, ln_elementwise_affine)
+        self.project_to_hidden = nn.Linear(in_channels_embed, self.inner_dim, bias=use_bias)
+        self.project_from_hidden_norm = RMSNorm(self.inner_dim, layer_norm_eps, ln_elementwise_affine)
+        self.project_from_hidden = nn.Linear(self.inner_dim, in_channels_embed, bias=use_bias)
+        
+        self.down_block = Simple_UVitBlock(
+            self.inner_dim, 
+            ln_elementwise_affine,
+            layer_norm_eps,
+            use_bias,
+            downsample,
+            False,
+        )
+        self.up_block = Simple_UVitBlock(
+            self.inner_dim, #block_out_channels,
+            ln_elementwise_affine,
+            layer_norm_eps,
+            use_bias,
+            False,
+            upsample=upsample,
+        )
+       
+        # self.fuse_qkv_projections()
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedFluxAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples= None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        micro_conds: torch.Tensor = None,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        micro_cond_encode_dim = 256 # same as self.config.micro_cond_encode_dim = 256 from amused
+        micro_cond_embeds = get_timestep_embedding(
+            micro_conds.flatten(), micro_cond_encode_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        ) 
+        micro_cond_embeds = micro_cond_embeds.reshape((hidden_states.shape[0], -1)) 
+
+        pooled_projections = torch.cat([pooled_projections, micro_cond_embeds], dim=1)
+        pooled_projections = pooled_projections.to(dtype=self.dtype)
+        pooled_projections = self.cond_embed(pooled_projections).to(encoder_hidden_states.dtype)    
+       
+
+        hidden_states = self.embed(hidden_states) 
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states) 
+        encoder_hidden_states = self.encoder_proj_layer_norm(encoder_hidden_states)
+        hidden_states = self.down_block(hidden_states)
+
+        batch_size, channels, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+        hidden_states = self.project_to_hidden_norm(hidden_states) 
+        hidden_states = self.project_to_hidden(hidden_states)
+
+       
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+        temb = (      
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        ) 
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+       
+        image_rotary_emb = self.pos_embed(ids) 
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,  
+                    image_rotary_emb=image_rotary_emb,
+                )
+                
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] 
+
+       
+        hidden_states = self.project_from_hidden_norm(hidden_states)
+        hidden_states = self.project_from_hidden(hidden_states)
+       
+
+        hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+
+        hidden_states = self.up_block(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+        
+        output = self.mlm_layer(hidden_states)
+        # self.unfuse_qkv_projections()
+        if not return_dict:
+            return (output,)
+
+    
+        return output