internvlu_image_dit / intern_native /pipeline_internvlu_generation_decoder.py
Publish verified native DiT using pinned official LFS copy and numerical validation
a130d59 verified 1 day ago
26.9 kB
# --------------------------------------------------------
# InternVL-U
# Modifications Copyright (c) 2026 OpenGVLab
# This file includes code from PixArt-Sigma and HuggingFace,
# licensed under the Apache License, Version 2.0.
# --------------------------------------------------------
# Copyright 2025 PixArt-Sigma Authors and The HuggingFace 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.

import math
import inspect

from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import PIL
import numpy as np
import torch

from einops import rearrange
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import PixArtImageProcessor
from diffusers.models import AutoencoderDC
from diffusers.schedulers import DPMSolverMultistepScheduler
from diffusers.utils import (
    is_torch_xla_available,
    logging,
)
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha import (
    ASPECT_RATIO_512_BIN,
    ASPECT_RATIO_1024_BIN,
)
from diffusers.pipelines.pixart_alpha.pipeline_pixart_sigma import ASPECT_RATIO_2048_BIN
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import (
    retrieve_timesteps,
)
from diffusers.pipelines.pipeline_utils import StableDiffusionMixin
from diffusers.utils import BaseOutput

from .modeling_internvlu_generation_decoder import InternVLUGenerationDecoder

if is_torch_xla_available():
    import torch_xla.core.xla_model as xm

    XLA_AVAILABLE = True
else:
    XLA_AVAILABLE = False


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


ASPECT_RATIO_4096_BIN = {
    "0.25": [2048.0, 8192.0],
    "0.26": [2048.0, 7936.0],
    "0.27": [2048.0, 7680.0],
    "0.28": [2048.0, 7424.0],
    "0.32": [2304.0, 7168.0],
    "0.33": [2304.0, 6912.0],
    "0.35": [2304.0, 6656.0],
    "0.4": [2560.0, 6400.0],
    "0.42": [2560.0, 6144.0],
    "0.48": [2816.0, 5888.0],
    "0.5": [2816.0, 5632.0],
    "0.52": [2816.0, 5376.0],
    "0.57": [3072.0, 5376.0],
    "0.6": [3072.0, 5120.0],
    "0.68": [3328.0, 4864.0],
    "0.72": [3328.0, 4608.0],
    "0.78": [3584.0, 4608.0],
    "0.82": [3584.0, 4352.0],
    "0.88": [3840.0, 4352.0],
    "0.94": [3840.0, 4096.0],
    "1.0": [4096.0, 4096.0],
    "1.07": [4096.0, 3840.0],
    "1.13": [4352.0, 3840.0],
    "1.21": [4352.0, 3584.0],
    "1.29": [4608.0, 3584.0],
    "1.38": [4608.0, 3328.0],
    "1.46": [4864.0, 3328.0],
    "1.67": [5120.0, 3072.0],
    "1.75": [5376.0, 3072.0],
    "2.0": [5632.0, 2816.0],
    "2.09": [5888.0, 2816.0],
    "2.4": [6144.0, 2560.0],
    "2.5": [6400.0, 2560.0],
    "2.89": [6656.0, 2304.0],
    "3.0": [6912.0, 2304.0],
    "3.11": [7168.0, 2304.0],
    "3.62": [7424.0, 2048.0],
    "3.75": [7680.0, 2048.0],
    "3.88": [7936.0, 2048.0],
    "4.0": [8192.0, 2048.0],
}


@dataclass
class InternVLUDiffusionPipelineOutput(BaseOutput):
    """
    Output class for Stable Diffusion pipelines.
    Args:
        images (`List[PIL.Image.Image]` or `np.ndarray`)
            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
    """

    images: Union[List[PIL.Image.Image], np.ndarray]


def retrieve_latents(
    encoder_output: torch.Tensor,
    generator: Optional[torch.Generator] = None,
    sample_mode: str = "sample",
):
    """Retrieve latent samples from an encoder output object.
    Args:
        encoder_output (`torch.Tensor` or object): Output that may expose `latent_dist` or `latents`.
        generator (`torch.Generator`, *optional*): RNG to use when sampling.
        sample_mode (`str`, *optional*, defaults to `"sample"`): `"sample"` to sample, `"argmax"` to use mode.
    Returns:
        `torch.Tensor`: Latent tensor extracted from the encoder output.
    """
    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
        return encoder_output.latent_dist.sample(generator)
    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
        return encoder_output.latent_dist.mode()
    elif hasattr(encoder_output, "latents"):
        return encoder_output.latents
    else:
        raise AttributeError("Could not access latents of provided encoder_output")


class InternVLUDiffusionPipeline(DiffusionPipeline, StableDiffusionMixin):
    """Diffusion pipeline wrapper around the InternVLU generation decoder."""

    def __init__(
        self,
        vae: AutoencoderDC,
        generation_decoder: InternVLUGenerationDecoder,
        scheduler: DPMSolverMultistepScheduler,
    ):
        super().__init__()

        self.register_modules(
            vae=vae, generation_decoder=generation_decoder, scheduler=scheduler
        )
        self.scheduler.config["flow_shift"] = self.generation_decoder.config.flow_shift

        self.transformer = self.generation_decoder.decoder

        if hasattr(self, "vae") and self.vae is not None:
            if hasattr(self.vae.config, "encoder_block_out_channels"):
                self.vae_scale_factor = 2 ** (
                    len(self.vae.config.encoder_block_out_channels) - 1
                )
            elif hasattr(self.vae.config, "temperal_downsample"):
                self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
            else:
                self.vae_scale_factor = 32

            self.latent_channels = (
                self.vae.config.z_dim if getattr(self, "vae", None) else 16
            )
        self.image_processor = PixArtImageProcessor(
            vae_scale_factor=self.vae_scale_factor
        )

    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
    def prepare_extra_step_kwargs(self, generator, eta):
        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
        # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502
        # and should be between [0, 1]

        accepts_eta = "eta" in set(
            inspect.signature(self.scheduler.step).parameters.keys()
        )
        extra_step_kwargs = {}
        if accepts_eta:
            extra_step_kwargs["eta"] = eta

        # check if the scheduler accepts generator
        accepts_generator = "generator" in set(
            inspect.signature(self.scheduler.step).parameters.keys()
        )
        if accepts_generator:
            extra_step_kwargs["generator"] = generator
        return extra_step_kwargs

    def prepare_latents(
        self,
        batch_size,
        num_channels_latents,
        height,
        width,
        dtype,
        device,
        generator,
        latents=None,
    ):
        if latents is not None:
            return latents.to(device=device, dtype=dtype)

        shape = (
            batch_size,
            num_channels_latents,
            int(height) // self.vae_scale_factor,
            int(width) // self.vae_scale_factor,
        )
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        return latents

    @property
    def attention_kwargs(self):
        return self._attention_kwargs

    @property
    def do_classifier_free_guidance(self):
        return True

    @property
    def num_timesteps(self):
        return self._num_timesteps

    @property
    def interrupt(self):
        return self._interrupt

    @torch.no_grad()
    def pixels_to_latents(self, x, max_bs=32):
        x = x.to(dtype=self.dtype)
        x = x.unsqueeze(2)

        max_bs = math.floor(max_bs * 3 * 512**2 / x[0].numel())
        num_iter = math.ceil(x.shape[0] / max_bs)
        z_list = []
        for i in range(num_iter):
            x_bs = x[i * max_bs : (i + 1) * max_bs]
            image_latents = retrieve_latents(
                self.vae.encode(x_bs), generator=None, sample_mode="argmax"
            )
            z_list.append(image_latents)
        image_latents = torch.cat(z_list, dim=0)
        # image_latents = retrieve_latents(self.vae.encode(x), generator=None, sample_mode="argmax")
        latents_mean = (
            torch.tensor(self.vae.config.latents_mean)
            .view(1, self.latent_channels, 1, 1, 1)
            .to(image_latents.device, image_latents.dtype)
        )
        latents_std = (
            torch.tensor(self.vae.config.latents_std)
            .view(1, self.latent_channels, 1, 1, 1)
            .to(image_latents.device, image_latents.dtype)
        )
        z = (image_latents - latents_mean) / latents_std
        z = z.squeeze(2)
        return z

    @torch.no_grad()
    def latents_to_pixels(self, z, max_bs=8):
        latents = z.to(self.vae.dtype).unsqueeze(2)
        latents_mean = (
            torch.tensor(self.vae.config.latents_mean)
            .view(1, self.vae.config.z_dim, 1, 1, 1)
            .to(latents.device, latents.dtype)
        )
        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
            1, self.vae.config.z_dim, 1, 1, 1
        ).to(latents.device, latents.dtype)
        latents = latents / latents_std + latents_mean

        max_bs = math.floor(max_bs * 16 * 64**2 / latents[0].numel())
        num_iter = math.ceil(latents.shape[0] / max_bs)
        x_list = []
        for i in range(num_iter):
            latents_bs = latents[i * max_bs : (i + 1) * max_bs]
            x_rec = self.vae.decode(latents_bs, return_dict=False)[0][:, :, 0]
            x_list.append(x_rec)
        x_rec = torch.cat(x_list)
        return x_rec

    @torch.no_grad()
    def __call__(
        self,
        encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
        encoder_image_token_mask: Union[torch.Tensor, List[torch.Tensor]] = None,
        image_fhw_cond: Union[torch.Tensor, List[torch.Tensor]] = None,
        image_hidden_states: Union[torch.Tensor, List[torch.Tensor]] = None,
        image_rel_idxs: Union[torch.Tensor, List[torch.Tensor]] = None,
        conditional_image: Union[torch.Tensor, List[torch.Tensor]] = None,
        num_images_per_prompt: Optional[int] = 1,
        num_inference_steps: int = 20,
        timesteps: List[int] = None,
        sigmas: List[float] = None,
        all_cfg_scale: float = 4.5,
        part_cfg_scale: float = 2.0,
        height: int = 256,
        width: int = 256,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        use_resolution_binning: bool = True,
        attention_kwargs: Optional[Dict[str, Any]] = None,
        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
        conditional_input: Optional[List[torch.Tensor]] = None,
        timestep_trunc=930,
        timestep_shift=3.0,
        image_grid_thw_gen: torch.Tensor = None,
        image_grid_thw_gen_cond: torch.Tensor = None,
    ) -> Union[InternVLUDiffusionPipelineOutput, Tuple]:
        """
        Generate images from VLM conditioning and optional image guidance.
        Args:
            encoder_hidden_states (`torch.Tensor` or `List[torch.Tensor]`):
                Conditioning hidden states from the VLM.
            encoder_image_token_mask (`torch.Tensor` or `List[torch.Tensor]`, *optional*):
                Mask indicating image token positions in the encoder states.
            image_fhw_cond (`torch.Tensor` or `List[torch.Tensor]`, *optional*):
                Frame/height/width metadata for conditioning images.
            image_hidden_states (`torch.Tensor` or `List[torch.Tensor]`, *optional*):
                Optional image hidden states used for conditioning.
            image_rel_idxs (`torch.Tensor` or `List[torch.Tensor]`, *optional*):
                Relative indices for channel-add conditioning.
            conditional_image (`torch.Tensor` or `List[torch.Tensor]`, *optional*):
                Conditional images or latents for image-based guidance.
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                Number of images to generate per prompt.
            num_inference_steps (`int`, *optional*, defaults to 20):
                Number of denoising steps.
            timesteps (`List[int]`, *optional*):
                Custom timesteps to use for the denoising process.
            sigmas (`List[float]`, *optional*):
                Custom sigmas to use for the denoising process.
            all_cfg_scale (`float`, *optional*, defaults to 4.5):
                Classifier-free guidance scale for full conditioning.
            part_cfg_scale (`float`, *optional*, defaults to 2.0):
                Classifier-free guidance scale for partial conditioning.
            height (`int`, *optional*, defaults to 256):
                Height in pixels of the generated image.
            width (`int`, *optional*, defaults to 256):
                Width in pixels of the generated image.
            eta (`float`, *optional*, defaults to 0.0):
                Eta (η) for DDIM-like schedulers.
            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
                Random generator(s) for deterministic output.
            latents (`torch.Tensor`, *optional*):
                Pre-generated noisy latents. If not provided, latents are sampled.
            output_type (`str`, *optional*, defaults to `"pil"`):
                Output format: `"pil"` or `"np"`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a pipeline output dataclass or a tuple.
            use_resolution_binning (`bool`, *optional*, defaults to `True`):
                Map requested size to the closest aspect ratio bin before decoding.
            attention_kwargs (`Dict[str, Any]`, *optional*):
                Extra kwargs passed to attention processors.
            callback_on_step_end (`Callable`, *optional*):
                Callback invoked at the end of each denoising step.
            callback_on_step_end_tensor_inputs (`List[str]`, *optional*):
                Tensor names to pass to the callback.
            conditional_input (`List[torch.Tensor]`, *optional*):
                Optional conditional inputs for the decoder.
            timestep_trunc (`int`, *optional*, defaults to 930):
                Maximum timestep used when truncating the schedule.
            timestep_shift (`float`, *optional*, defaults to 3.0):
                Shift value applied to timesteps.
            image_grid_thw_gen (`torch.Tensor`, *optional*):
                Generation grid metadata for image tokens.
            image_grid_thw_gen_cond (`torch.Tensor`, *optional*):
                Conditioning grid metadata for image tokens.
        Returns:
            `InternVLUDiffusionPipelineOutput` or `tuple`: Generated images.
        """

        # 0. perpare inputs

        assert (
            len(encoder_hidden_states) % 3 == 0
        ), "encoder_hidden_states should be mod by 3, as it is split into three parts for dual cfg guidance."

        conditional_input = None

        image_hidden_states = None
        if conditional_image is not None:
            split_sizes = None
            if isinstance(conditional_image, list):
                split_sizes = [len(c) for c in conditional_image]
                conditional_image = torch.cat(conditional_image)
            if conditional_image.shape[0] > 0:
                image_hidden_states = self.pixels_to_latents(conditional_image)
                if split_sizes is not None:
                    image_hidden_states = list(
                        torch.split(image_hidden_states, split_sizes)
                    )
                conditional_input = image_hidden_states
                image_hidden_states = None

        encoder_hidden_states, attention_masks, encoder_image_token_mask = (
            self.generation_decoder.prepare_forward_input(
                encoder_hidden_states,
                encoder_image_token_mask=encoder_image_token_mask,
            )
        )

        prompt_embeds = encoder_hidden_states
        prompt_attention_mask = attention_masks

        # duplicate text embeddings for each generation per prompt, using mps friendly method
        def _duplicate(embeds: torch.Tensor, attention_mask: torch.Tensor):
            bs_embed, seq_len, _ = embeds.shape
            # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
            embeds = embeds.repeat(1, num_images_per_prompt, 1)
            embeds = embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
            attention_mask = attention_mask.view(bs_embed, -1)
            attention_mask = attention_mask.repeat_interleave(num_images_per_prompt, 0)
            # attention_mask = attention_mask.repeat(num_images_per_prompt, 1)

            return embeds, attention_mask

        prompt_embeds, prompt_attention_mask = _duplicate(
            prompt_embeds, prompt_attention_mask
        )

        if conditional_input is not None:
            conditional_input_new = []
            for c in conditional_input:
                conditional_input_new.extend([c] * num_images_per_prompt)
            conditional_input = conditional_input_new
        if encoder_image_token_mask is not None:
            encoder_image_token_mask = encoder_image_token_mask.repeat_interleave(
                num_images_per_prompt, 0
            )
        if image_fhw_cond is not None:
            image_fhw_cond_new = []
            for c in image_fhw_cond:
                image_fhw_cond_new.extend([c] * num_images_per_prompt)
            image_fhw_cond = image_fhw_cond_new
        if image_grid_thw_gen is not None:
            image_grid_thw_gen = image_grid_thw_gen.repeat_interleave(
                num_images_per_prompt, 0
            )
        if image_grid_thw_gen_cond is not None:
            image_grid_thw_gen_cond_new = []
            for c in image_grid_thw_gen_cond:
                image_grid_thw_gen_cond_new.extend([c] * num_images_per_prompt)
            image_grid_thw_gen_cond = image_grid_thw_gen_cond_new

        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs

        # 1. Check inputs. Raise error if not correct
        if use_resolution_binning:
            if self.transformer.config.sample_size == 128:
                aspect_ratio_bin = ASPECT_RATIO_4096_BIN
            elif self.transformer.config.sample_size == 64:
                aspect_ratio_bin = ASPECT_RATIO_2048_BIN
            elif self.transformer.config.sample_size == 32:
                aspect_ratio_bin = ASPECT_RATIO_1024_BIN
            elif self.transformer.config.sample_size == 16:
                aspect_ratio_bin = ASPECT_RATIO_512_BIN
            else:
                raise ValueError("Invalid sample size")
            height, width = self.image_processor.classify_height_width_bin(
                height, width, ratios=aspect_ratio_bin
            )

        self._attention_kwargs = attention_kwargs
        self._interrupt = False

        # 2. Default height and width to transformer
        batch_size = prompt_embeds.shape[0]

        device = self._execution_device
        lora_scale = (
            self.attention_kwargs.get("scale", None)
            if self.attention_kwargs is not None
            else None
        )

        # print(self.scheduler, self.scheduler.config)
        # 4. Prepare timesteps
        if timestep_shift is not None:
            self.scheduler.config.flow_shift = timestep_shift
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler, num_inference_steps, device, timesteps, sigmas
        )

        # 5. Prepare latents.
        latent_channels = self.transformer.config.in_channels // (
            self.transformer.config.patch_size**2
        )
        if image_grid_thw_gen is not None:
            height = max(image_grid_thw_gen[:, 1]) * self.vae_scale_factor
            width = max(image_grid_thw_gen[:, 2]) * self.vae_scale_factor
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt // 3,
            latent_channels,
            height,
            width,
            torch.float32,
            device,
            generator,
            latents,
        )

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

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

        transformer_dtype = self.transformer.dtype
        self.set_progress_bar_config(disable=True)
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
                    continue

                # latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                latent_model_input = torch.cat([latents] * 3)

                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                timestep = t.expand(latent_model_input.shape[0])
                timestep = timestep * getattr(
                    self.transformer.config, "timestep_scale", 1.0
                )

                # predict noise model_output
                noise_pred = self.transformer(
                    latent_model_input.to(dtype=transformer_dtype),
                    encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype),
                    encoder_attention_mask=prompt_attention_mask,
                    encoder_image_token_mask=encoder_image_token_mask,
                    image_fhw_cond=image_fhw_cond,
                    timestep=timestep,
                    return_dict=False,
                    attention_kwargs=self.attention_kwargs,
                    conditional_input=conditional_input,
                    image_grid_thw_gen=image_grid_thw_gen,
                    image_grid_thw_gen_cond=image_grid_thw_gen_cond,
                )[0]
                noise_pred = noise_pred.float()

                noise_pred_cond, noise_pred_part_cond, noise_pred_uncond = (
                    noise_pred.chunk(3)
                )

                if timestep_trunc > 0 and t > timestep_trunc:
                    part_diff_norm = torch.norm(
                        noise_pred_part_cond - noise_pred_uncond, dim=1, keepdim=True
                    )
                    diff_norm = torch.norm(
                        noise_pred_cond - noise_pred_part_cond, dim=1, keepdim=True
                    )
                    noise_pred = (
                        noise_pred_uncond
                        + part_cfg_scale
                        * (noise_pred_part_cond - noise_pred_uncond)
                        / self.process_diff_norm(part_diff_norm, k=0.4)
                        + all_cfg_scale
                        * (noise_pred_cond - noise_pred_part_cond)
                        / self.process_diff_norm(diff_norm, k=0.4)
                    )
                else:
                    noise_pred = (
                        noise_pred_uncond
                        + part_cfg_scale * (noise_pred_part_cond - noise_pred_uncond)
                        + all_cfg_scale * (noise_pred_cond - noise_pred_part_cond)
                    )

                # learned sigma
                if self.transformer.config.out_channels // 2 == latent_channels:
                    noise_pred = noise_pred.chunk(2, dim=1)[0]

                # compute previous image: x_t -> x_t-1
                latents = self.scheduler.step(
                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
                )[0]

                if callback_on_step_end is not None:
                    callback_kwargs = {}
                    for k in callback_on_step_end_tensor_inputs:
                        callback_kwargs[k] = locals()[k]
                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)

                    latents = callback_outputs.pop("latents", latents)
                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                    negative_prompt_embeds = callback_outputs.pop(
                        "negative_prompt_embeds", negative_prompt_embeds
                    )

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

                if XLA_AVAILABLE:
                    xm.mark_step()

        if output_type == "latent":
            image = latents
        else:
            image = self.latents_to_pixels(latents.to(self.generation_decoder.dtype))

        # Offload all models
        self.maybe_free_model_hooks()

        if not return_dict:
            return (image,)

        return InternVLUDiffusionPipelineOutput(images=image)

    @staticmethod
    def process_diff_norm(diff_norm, k):
        pow_result = torch.pow(diff_norm, k)

        result = torch.where(
            diff_norm > 1.0,
            pow_result,
            torch.where(diff_norm < 1.0, torch.ones_like(diff_norm), diff_norm),
        )
        return result