OmniGen

OmniGen: Unified Image Generation from BAAI, by Shitao Xiao, Yueze Wang, Junjie Zhou, Huaying Yuan, Xingrun Xing, Ruiran Yan, Chaofan Li, Shuting Wang, Tiejun Huang, Zheng Liu.

The abstract from the paper is:

The emergence of Large Language Models (LLMs) has unified language generation tasks and revolutionized human-machine interaction. However, in the realm of image generation, a unified model capable of handling various tasks within a single framework remains largely unexplored. In this work, we introduce OmniGen, a new diffusion model for unified image generation. OmniGen is characterized by the following features: 1) Unification: OmniGen not only demonstrates text-to-image generation capabilities but also inherently supports various downstream tasks, such as image editing, subject-driven generation, and visual conditional generation. 2) Simplicity: The architecture of OmniGen is highly simplified, eliminating the need for additional plugins. Moreover, compared to existing diffusion models, it is more user-friendly and can complete complex tasks end-to-end through instructions without the need for extra intermediate steps, greatly simplifying the image generation workflow. 3) Knowledge Transfer: Benefit from learning in a unified format, OmniGen effectively transfers knowledge across different tasks, manages unseen tasks and domains, and exhibits novel capabilities. We also explore the model’s reasoning capabilities and potential applications of the chain-of-thought mechanism. This work represents the first attempt at a general-purpose image generation model, and we will release our resources at https://github.com/VectorSpaceLab/OmniGen to foster future advancements.

Tip

Make sure to check out the Schedulers guide to learn how to explore the tradeoff between scheduler speed and quality, and see the reuse components across pipelines section to learn how to efficiently load the same components into multiple pipelines.

Tip

This pipeline was contributed by staoxiao. The original codebase can be found here. The original weights can be found under hf.co/shitao.

Inference

First, load the pipeline:

import mindspore as ms
from diffusers import OmniGenPipeline

pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1-diffusers", mindspore_dtype=ms.bfloat16)

For text-to-image, pass a text prompt. By default, OmniGen generates a 1024x1024 image. You can try setting the height and width parameters to generate images with different size.

prompt = "Realistic photo. A young woman sits on a sofa, holding a book and facing the camera. She wears delicate silver hoop earrings adorned with tiny, sparkling diamonds that catch the light, with her long chestnut hair cascading over her shoulders. Her eyes are focused and gentle, framed by long, dark lashes. She is dressed in a cozy cream sweater, which complements her warm, inviting smile. Behind her, there is a table with a cup of water in a sleek, minimalist blue mug. The background is a serene indoor setting with soft natural light filtering through a window, adorned with tasteful art and flowers, creating a cozy and peaceful ambiance. 4K, HD."
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=3,
    generator=np.random.Generator(np.random.PCG64(111)),
).[0][0]
image.save("output.png")

OmniGen supports multimodal inputs. When the input includes an image, you need to add a placeholder <img><|image_1|></img> in the text prompt to represent the image. It is recommended to enable use_input_image_size_as_output to keep the edited image the same size as the original image.

prompt="<img><|image_1|></img> Remove the woman's earrings. Replace the mug with a clear glass filled with sparkling iced cola."
input_images=[load_image("https://raw.githubusercontent.com/VectorSpaceLab/OmniGen/main/imgs/docs_img/t2i_woman_with_book.png")]
image = pipe(
    prompt=prompt,
    input_images=input_images,
    guidance_scale=2,
    img_guidance_scale=1.6,
    use_input_image_size_as_output=True,
    generator=np.random.Generator(np.random.PCG64(222)),).[0][0]
image.save("output.png")

autodoc

mindone.diffusers.OmniGenPipeline

Bases: DiffusionPipeline

The OmniGen pipeline for multimodal-to-image generation.

Reference: https://arxiv.org/pdf/2409.11340

PARAMETER	DESCRIPTION
transformer	Autoregressive Transformer architecture for OmniGen. TYPE: [`OmniGenTransformer2DModel`]
scheduler	A scheduler to be used in combination with transformer to denoise the encoded image latents. TYPE: [`FlowMatchEulerDiscreteScheduler`]
vae	Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. TYPE: [`AutoencoderKL`]
tokenizer	Text tokenizer of class. LlamaTokenizer. TYPE: `LlamaTokenizer`

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

class OmniGenPipeline(
    DiffusionPipeline,
):
    r"""
    The OmniGen pipeline for multimodal-to-image generation.

    Reference: https://arxiv.org/pdf/2409.11340

    Args:
        transformer ([`OmniGenTransformer2DModel`]):
            Autoregressive Transformer architecture for OmniGen.
        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
        vae ([`AutoencoderKL`]):
            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
        tokenizer (`LlamaTokenizer`):
            Text tokenizer of class.
            [LlamaTokenizer](https://huggingface.co/docs/transformers/main/model_doc/llama#transformers.LlamaTokenizer).
    """

    model_cpu_offload_seq = "transformer->vae"
    _optional_components = []
    _callback_tensor_inputs = ["latents"]

    def __init__(
        self,
        transformer: OmniGenTransformer2DModel,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: AutoencoderKL,
        tokenizer: LlamaTokenizer,
    ):
        super().__init__()

        self.register_modules(
            vae=vae,
            tokenizer=tokenizer,
            transformer=transformer,
            scheduler=scheduler,
        )
        self.vae_scale_factor = (
            2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) is not None else 8
        )
        # OmniGen latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)

        self.multimodal_processor = OmniGenMultiModalProcessor(tokenizer, max_image_size=1024)
        self.tokenizer_max_length = (
            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 120000
        )
        self.default_sample_size = 128

    def encode_input_images(
        self,
        input_pixel_values: List[ms.Tensor],
        dtype: Optional[ms.Type] = None,
    ):
        """
        get the continue embedding of input images by VAE

        Args:
            input_pixel_values: normalized pixel of input images
        Returns: ms.Tensor
        """
        dtype = dtype or self.vae.dtype

        input_img_latents = []
        for img in input_pixel_values:
            img = self.vae.diag_gauss_dist.sample(self.vae.encode(img.to(dtype))[0]).mul_(
                self.vae.config.scaling_factor
            )
            input_img_latents.append(img)
        return input_img_latents

    def check_inputs(
        self,
        prompt,
        input_images,
        height,
        width,
        use_input_image_size_as_output,
        callback_on_step_end_tensor_inputs=None,
    ):
        if input_images is not None:
            if len(input_images) != len(prompt):
                raise ValueError(
                    f"The number of prompts: {len(prompt)} does not match the number of input images: {len(input_images)}."
                )
            for i in range(len(input_images)):
                if input_images[i] is not None:
                    if not all(f"<img><|image_{k + 1}|></img>" in prompt[i] for k in range(len(input_images[i]))):
                        raise ValueError(
                            f"prompt `{prompt[i]}` doesn't have enough placeholders for the input images `{input_images[i]}`"
                        )

        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
            logger.warning(
                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
            )

        if use_input_image_size_as_output:
            if input_images is None or input_images[0] is None:
                raise ValueError(
                    "`use_input_image_size_as_output` is set to True, but no input image was found. \
                        If you are performing a text-to-image task, please set it to False."
                )

        if callback_on_step_end_tensor_inputs is not None and not all(
            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
        ):
            raise ValueError(
                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, \
                    but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
            )

    def enable_vae_slicing(self):
        r"""
        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
        """
        self.vae.enable_slicing()

    def disable_vae_slicing(self):
        r"""
        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
        computing decoding in one step.
        """
        self.vae.disable_slicing()

    def enable_vae_tiling(self):
        r"""
        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
        processing larger images.
        """
        self.vae.enable_tiling()

    def disable_vae_tiling(self):
        r"""
        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
        computing decoding in one step.
        """
        self.vae.disable_tiling()

    # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents
    def prepare_latents(
        self,
        batch_size,
        num_channels_latents,
        height,
        width,
        dtype,
        generator,
        latents=None,
    ):
        if latents is not None:
            return latents.to(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, dtype=dtype)

        return latents

    @property
    def guidance_scale(self):
        return self._guidance_scale

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

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

    def __call__(
        self,
        prompt: Union[str, List[str]],
        input_images: Union[PipelineImageInput, List[PipelineImageInput]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        max_input_image_size: int = 1024,
        timesteps: List[int] = None,
        guidance_scale: float = 2.5,
        img_guidance_scale: float = 1.6,
        use_input_image_size_as_output: bool = False,
        num_images_per_prompt: Optional[int] = 1,
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
        latents: Optional[ms.Tensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = False,
        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
    ):
        r"""
        Function invoked when calling the pipeline for generation.

        Args:
            prompt (`str` or `List[str]`, *optional*):
                The prompt or prompts to guide the image generation. If the input includes images, need to add
                placeholders `<img><|image_i|></img>` in the prompt to indicate the position of the i-th images.
            input_images (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
                The list of input images. We will replace the "<|image_i|>" in prompt with the i-th image in list.
            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The height in pixels of the generated image. This is set to 1024 by default for the best results.
            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
                The width in pixels of the generated image. This is set to 1024 by default for the best results.
            num_inference_steps (`int`, *optional*, defaults to 50):
                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                expense of slower inference.
            max_input_image_size (`int`, *optional*, defaults to 1024):
                the maximum size of input image, which will be used to crop the input image to the maximum size
            timesteps (`List[int]`, *optional*):
                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
                passed will be used. Must be in descending order.
            guidance_scale (`float`, *optional*, defaults to 2.5):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            img_guidance_scale (`float`, *optional*, defaults to 1.6):
                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800).
            use_input_image_size_as_output (bool, defaults to False):
                whether to use the input image size as the output image size, which can be used for single-image input,
                e.g., image editing task
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            generator (`np.random.Generator` or `List[np.random.Generator]`, *optional*):
                One or a list of [`np.random.Generator`](https://numpy.org/doc/stable/reference/random/generator.html)
                to make generation deterministic.
            latents (`ms.Tensor`, *optional*):
                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
                tensor will ge generated by sampling using the supplied random `generator`.
            output_type (`str`, *optional*, defaults to `"pil"`):
                The output format of the generate image. Choose between
                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
            callback_on_step_end (`Callable`, *optional*):
                A function that calls at the end of each denoising steps during the inference. The function is called
                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
                `callback_on_step_end_tensor_inputs`.
            callback_on_step_end_tensor_inputs (`List`, *optional*):
                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                `._callback_tensor_inputs` attribute of your pipeline class.

        Examples:

        Returns: [`~pipelines.ImagePipelineOutput`] or `tuple`:
            If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is returned
            where the first element is a list with the generated images.
        """

        height = height or self.default_sample_size * self.vae_scale_factor
        width = width or self.default_sample_size * self.vae_scale_factor
        num_cfg = 2 if input_images is not None else 1
        use_img_cfg = True if input_images is not None else False
        if isinstance(prompt, str):
            prompt = [prompt]
            input_images = [input_images]

        # 1. Check inputs. Raise error if not correct
        self.check_inputs(
            prompt,
            input_images,
            height,
            width,
            use_input_image_size_as_output,
            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
        )

        self._guidance_scale = guidance_scale
        self._interrupt = False

        # 2. Define call parameters
        batch_size = len(prompt)

        # 3. process multi-modal instructions
        if max_input_image_size != self.multimodal_processor.max_image_size:
            self.multimodal_processor.reset_max_image_size(max_image_size=max_input_image_size)
        processed_data = self.multimodal_processor(
            prompt,
            input_images,
            height=height,
            width=width,
            use_img_cfg=use_img_cfg,
            use_input_image_size_as_output=use_input_image_size_as_output,
            num_images_per_prompt=num_images_per_prompt,
        )

        # 4. Encode input images
        input_img_latents = self.encode_input_images(processed_data["input_pixel_values"])

        # 5. Prepare timesteps
        sigmas = np.linspace(1, 0, num_inference_steps + 1)[:num_inference_steps]
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler, num_inference_steps, timesteps, sigmas=sigmas
        )
        self._num_timesteps = len(timesteps)

        # 6. Prepare latents
        transformer_dtype = self.transformer.dtype
        if use_input_image_size_as_output:
            height, width = processed_data["input_pixel_values"][0].shape[-2:]
        latent_channels = self.transformer.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            latent_channels,
            height,
            width,
            ms.float32,
            generator,
            latents,
        )

        # 8. Denoising loop
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = mint.cat([latents] * (num_cfg + 1))
                latent_model_input = latent_model_input.to(transformer_dtype)

                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                timestep = t.broadcast_to((latent_model_input.shape[0],))

                noise_pred = self.transformer(
                    hidden_states=latent_model_input,
                    timestep=timestep,
                    input_ids=processed_data["input_ids"],
                    input_img_latents=input_img_latents,
                    input_image_sizes=processed_data["input_image_sizes"],
                    attention_mask=processed_data["attention_mask"],
                    position_ids=processed_data["position_ids"],
                    return_dict=False,
                )[0]

                if num_cfg == 2:
                    cond, uncond, img_cond = mint.split(noise_pred, len(noise_pred) // 3, dim=0)
                    noise_pred = uncond + img_guidance_scale * (img_cond - uncond) + guidance_scale * (cond - img_cond)
                else:
                    cond, uncond = mint.split(noise_pred, len(noise_pred) // 2, dim=0)
                    noise_pred = uncond + guidance_scale * (cond - uncond)

                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(noise_pred, t, latents, 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)

                progress_bar.update()

        if not output_type == "latent":
            latents = latents.to(self.vae.dtype)
            latents = latents / self.vae.config.scaling_factor
            image = self.vae.decode(latents, return_dict=False)[0]
            image = self.image_processor.postprocess(image, output_type=output_type)
        else:
            image = latents

        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image)

mindone.diffusers.OmniGenPipeline.__call__(prompt, input_images=None, height=None, width=None, num_inference_steps=50, max_input_image_size=1024, timesteps=None, guidance_scale=2.5, img_guidance_scale=1.6, use_input_image_size_as_output=False, num_images_per_prompt=1, generator=None, latents=None, output_type='pil', return_dict=False, callback_on_step_end=None, callback_on_step_end_tensor_inputs=['latents'])

Function invoked when calling the pipeline for generation.

PARAMETER	DESCRIPTION
prompt	The prompt or prompts to guide the image generation. If the input includes images, need to add placeholders <img><|image_i|></img> in the prompt to indicate the position of the i-th images. TYPE: `str` or `List[str]`, *optional*
input_images	The list of input images. We will replace the "<|image_i|>" in prompt with the i-th image in list. TYPE: `PipelineImageInput` or `List[PipelineImageInput]`, *optional* DEFAULT: None
height	The height in pixels of the generated image. This is set to 1024 by default for the best results. TYPE: `int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor DEFAULT: None
width	The width in pixels of the generated image. This is set to 1024 by default for the best results. TYPE: `int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor DEFAULT: None
num_inference_steps	The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. TYPE: `int`, *optional*, defaults to 50 DEFAULT: 50
max_input_image_size	the maximum size of input image, which will be used to crop the input image to the maximum size TYPE: `int`, *optional*, defaults to 1024 DEFAULT: 1024
timesteps	Custom timesteps to use for the denoising process with schedulers which support a timesteps argument in their set_timesteps method. If not defined, the default behavior when num_inference_steps is passed will be used. Must be in descending order. TYPE: `List[int]`, *optional* DEFAULT: None
guidance_scale	Guidance scale as defined in Classifier-Free Diffusion Guidance. guidance_scale is defined as w of equation 2. of Imagen Paper. Guidance scale is enabled by setting guidance_scale > 1. Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality. TYPE: `float`, *optional*, defaults to 2.5 DEFAULT: 2.5
img_guidance_scale	Defined as equation 3 in Instrucpix2pix. TYPE: `float`, *optional*, defaults to 1.6 DEFAULT: 1.6
use_input_image_size_as_output	whether to use the input image size as the output image size, which can be used for single-image input, e.g., image editing task TYPE: bool, defaults to False DEFAULT: False
num_images_per_prompt	The number of images to generate per prompt. TYPE: `int`, *optional*, defaults to 1 DEFAULT: 1
generator	One or a list of np.random.Generator to make generation deterministic. TYPE: `np.random.Generator` or `List[np.random.Generator]`, *optional* DEFAULT: None
latents	Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will ge generated by sampling using the supplied random generator. TYPE: `ms.Tensor`, *optional* DEFAULT: None
output_type	The output format of the generate image. Choose between PIL: PIL.Image.Image or np.array. TYPE: `str`, *optional*, defaults to `"pil"` DEFAULT: 'pil'
return_dict	Whether or not to return a [~pipelines.ImagePipelineOutput] instead of a plain tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
callback_on_step_end	A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict). callback_kwargs will include a list of all tensors as specified by callback_on_step_end_tensor_inputs. TYPE: `Callable`, *optional* DEFAULT: None
callback_on_step_end_tensor_inputs	The list of tensor inputs for the callback_on_step_end function. The tensors specified in the list will be passed as callback_kwargs argument. You will only be able to include variables listed in the ._callback_tensor_inputs attribute of your pipeline class. TYPE: `List`, *optional* DEFAULT: ['latents']

[`~PIPELINES.IMAGEPIPELINEOUTPUT`] OR `TUPLE`	DESCRIPTION
	If return_dict is True, [~pipelines.ImagePipelineOutput] is returned, otherwise a tuple is returned
	where the first element is a list with the generated images.

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

def __call__(
    self,
    prompt: Union[str, List[str]],
    input_images: Union[PipelineImageInput, List[PipelineImageInput]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 50,
    max_input_image_size: int = 1024,
    timesteps: List[int] = None,
    guidance_scale: float = 2.5,
    img_guidance_scale: float = 1.6,
    use_input_image_size_as_output: bool = False,
    num_images_per_prompt: Optional[int] = 1,
    generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
    latents: Optional[ms.Tensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = False,
    callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
):
    r"""
    Function invoked when calling the pipeline for generation.

    Args:
        prompt (`str` or `List[str]`, *optional*):
            The prompt or prompts to guide the image generation. If the input includes images, need to add
            placeholders `<img><|image_i|></img>` in the prompt to indicate the position of the i-th images.
        input_images (`PipelineImageInput` or `List[PipelineImageInput]`, *optional*):
            The list of input images. We will replace the "<|image_i|>" in prompt with the i-th image in list.
        height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
            The height in pixels of the generated image. This is set to 1024 by default for the best results.
        width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
            The width in pixels of the generated image. This is set to 1024 by default for the best results.
        num_inference_steps (`int`, *optional*, defaults to 50):
            The number of denoising steps. More denoising steps usually lead to a higher quality image at the
            expense of slower inference.
        max_input_image_size (`int`, *optional*, defaults to 1024):
            the maximum size of input image, which will be used to crop the input image to the maximum size
        timesteps (`List[int]`, *optional*):
            Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
            in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
            passed will be used. Must be in descending order.
        guidance_scale (`float`, *optional*, defaults to 2.5):
            Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
            `guidance_scale` is defined as `w` of equation 2. of [Imagen
            Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
            1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
            usually at the expense of lower image quality.
        img_guidance_scale (`float`, *optional*, defaults to 1.6):
            Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800).
        use_input_image_size_as_output (bool, defaults to False):
            whether to use the input image size as the output image size, which can be used for single-image input,
            e.g., image editing task
        num_images_per_prompt (`int`, *optional*, defaults to 1):
            The number of images to generate per prompt.
        generator (`np.random.Generator` or `List[np.random.Generator]`, *optional*):
            One or a list of [`np.random.Generator`](https://numpy.org/doc/stable/reference/random/generator.html)
            to make generation deterministic.
        latents (`ms.Tensor`, *optional*):
            Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
            generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
            tensor will ge generated by sampling using the supplied random `generator`.
        output_type (`str`, *optional*, defaults to `"pil"`):
            The output format of the generate image. Choose between
            [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~pipelines.ImagePipelineOutput`] instead of a plain tuple.
        callback_on_step_end (`Callable`, *optional*):
            A function that calls at the end of each denoising steps during the inference. The function is called
            with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
            callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
            `callback_on_step_end_tensor_inputs`.
        callback_on_step_end_tensor_inputs (`List`, *optional*):
            The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
            will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
            `._callback_tensor_inputs` attribute of your pipeline class.

    Examples:

    Returns: [`~pipelines.ImagePipelineOutput`] or `tuple`:
        If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is returned
        where the first element is a list with the generated images.
    """

    height = height or self.default_sample_size * self.vae_scale_factor
    width = width or self.default_sample_size * self.vae_scale_factor
    num_cfg = 2 if input_images is not None else 1
    use_img_cfg = True if input_images is not None else False
    if isinstance(prompt, str):
        prompt = [prompt]
        input_images = [input_images]

    # 1. Check inputs. Raise error if not correct
    self.check_inputs(
        prompt,
        input_images,
        height,
        width,
        use_input_image_size_as_output,
        callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
    )

    self._guidance_scale = guidance_scale
    self._interrupt = False

    # 2. Define call parameters
    batch_size = len(prompt)

    # 3. process multi-modal instructions
    if max_input_image_size != self.multimodal_processor.max_image_size:
        self.multimodal_processor.reset_max_image_size(max_image_size=max_input_image_size)
    processed_data = self.multimodal_processor(
        prompt,
        input_images,
        height=height,
        width=width,
        use_img_cfg=use_img_cfg,
        use_input_image_size_as_output=use_input_image_size_as_output,
        num_images_per_prompt=num_images_per_prompt,
    )

    # 4. Encode input images
    input_img_latents = self.encode_input_images(processed_data["input_pixel_values"])

    # 5. Prepare timesteps
    sigmas = np.linspace(1, 0, num_inference_steps + 1)[:num_inference_steps]
    timesteps, num_inference_steps = retrieve_timesteps(
        self.scheduler, num_inference_steps, timesteps, sigmas=sigmas
    )
    self._num_timesteps = len(timesteps)

    # 6. Prepare latents
    transformer_dtype = self.transformer.dtype
    if use_input_image_size_as_output:
        height, width = processed_data["input_pixel_values"][0].shape[-2:]
    latent_channels = self.transformer.config.in_channels
    latents = self.prepare_latents(
        batch_size * num_images_per_prompt,
        latent_channels,
        height,
        width,
        ms.float32,
        generator,
        latents,
    )

    # 8. Denoising loop
    with self.progress_bar(total=num_inference_steps) as progress_bar:
        for i, t in enumerate(timesteps):
            # expand the latents if we are doing classifier free guidance
            latent_model_input = mint.cat([latents] * (num_cfg + 1))
            latent_model_input = latent_model_input.to(transformer_dtype)

            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
            timestep = t.broadcast_to((latent_model_input.shape[0],))

            noise_pred = self.transformer(
                hidden_states=latent_model_input,
                timestep=timestep,
                input_ids=processed_data["input_ids"],
                input_img_latents=input_img_latents,
                input_image_sizes=processed_data["input_image_sizes"],
                attention_mask=processed_data["attention_mask"],
                position_ids=processed_data["position_ids"],
                return_dict=False,
            )[0]

            if num_cfg == 2:
                cond, uncond, img_cond = mint.split(noise_pred, len(noise_pred) // 3, dim=0)
                noise_pred = uncond + img_guidance_scale * (img_cond - uncond) + guidance_scale * (cond - img_cond)
            else:
                cond, uncond = mint.split(noise_pred, len(noise_pred) // 2, dim=0)
                noise_pred = uncond + guidance_scale * (cond - uncond)

            # compute the previous noisy sample x_t -> x_t-1
            latents = self.scheduler.step(noise_pred, t, latents, 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)

            progress_bar.update()

    if not output_type == "latent":
        latents = latents.to(self.vae.dtype)
        latents = latents / self.vae.config.scaling_factor
        image = self.vae.decode(latents, return_dict=False)[0]
        image = self.image_processor.postprocess(image, output_type=output_type)
    else:
        image = latents

    if not return_dict:
        return (image,)

    return ImagePipelineOutput(images=image)

mindone.diffusers.OmniGenPipeline.disable_vae_slicing()

Disable sliced VAE decoding. If enable_vae_slicing was previously enabled, this method will go back to computing decoding in one step.

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

def disable_vae_slicing(self):
    r"""
    Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
    computing decoding in one step.
    """
    self.vae.disable_slicing()

mindone.diffusers.OmniGenPipeline.disable_vae_tiling()

Disable tiled VAE decoding. If enable_vae_tiling was previously enabled, this method will go back to computing decoding in one step.

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

def disable_vae_tiling(self):
    r"""
    Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
    computing decoding in one step.
    """
    self.vae.disable_tiling()

mindone.diffusers.OmniGenPipeline.enable_vae_slicing()

Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

def enable_vae_slicing(self):
    r"""
    Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
    compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
    """
    self.vae.enable_slicing()

mindone.diffusers.OmniGenPipeline.enable_vae_tiling()

Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images.

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

def enable_vae_tiling(self):
    r"""
    Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
    compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
    processing larger images.
    """
    self.vae.enable_tiling()

mindone.diffusers.OmniGenPipeline.encode_input_images(input_pixel_values, dtype=None)

get the continue embedding of input images by VAE

PARAMETER	DESCRIPTION
input_pixel_values	normalized pixel of input images TYPE: List[Tensor]

Source code in mindone/diffusers/pipelines/omnigen/pipeline_omnigen.py

def encode_input_images(
    self,
    input_pixel_values: List[ms.Tensor],
    dtype: Optional[ms.Type] = None,
):
    """
    get the continue embedding of input images by VAE

    Args:
        input_pixel_values: normalized pixel of input images
    Returns: ms.Tensor
    """
    dtype = dtype or self.vae.dtype

    input_img_latents = []
    for img in input_pixel_values:
        img = self.vae.diag_gauss_dist.sample(self.vae.encode(img.to(dtype))[0]).mul_(
            self.vae.config.scaling_factor
        )
        input_img_latents.append(img)
    return input_img_latents