OmniGen (original) (raw)

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.

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.

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 torch from diffusers import OmniGenPipeline

pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1-diffusers", torch_dtype=torch.bfloat16) pipe.to("cuda")

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=torch.Generator(device="cpu").manual_seed(111), ).images[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="<|image_1|> 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=torch.Generator(device="cpu").manual_seed(222)).images[0] image.save("output.png")

OmniGenPipeline

class diffusers.OmniGenPipeline

< source >

( transformer: OmniGenTransformer2DModel scheduler: FlowMatchEulerDiscreteScheduler vae: AutoencoderKL tokenizer: LlamaTokenizer )

Parameters

• 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.

The OmniGen pipeline for multimodal-to-image generation.

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

__call__

< source >

( prompt: typing.Union[str, typing.List[str]] input_images: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor], typing.List[typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, typing.List[PIL.Image.Image], typing.List[numpy.ndarray], typing.List[torch.Tensor]]]] = None height: typing.Optional[int] = None width: typing.Optional[int] = None num_inference_steps: int = 50 max_input_image_size: int = 1024 timesteps: typing.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: typing.Optional[int] = 1 generator: typing.Union[torch._C.Generator, typing.List[torch._C.Generator], NoneType] = None latents: typing.Optional[torch.Tensor] = None output_type: typing.Optional[str] = 'pil' return_dict: bool = True callback_on_step_end: typing.Optional[typing.Callable[[int, int, typing.Dict], NoneType]] = None callback_on_step_end_tensor_inputs: typing.List[str] = ['latents'] )

Parameters

• 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.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.
• img_guidance_scale (float, optional, defaults to 1.6) — Defined as equation 3 in Instrucpix2pix.
• 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 (torch.Generator or List[torch.Generator], optional) — One or a list of torch generator(s)to make generation deterministic.
• latents (torch.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 betweenPIL: PIL.Image.Image or np.array.
• return_dict (bool, optional, defaults to True) — Whether or not to return a ~pipelines.flux.FluxPipelineOutput 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 bycallback_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.

Function invoked when calling the pipeline for generation.

Examples:

import torch from diffusers import OmniGenPipeline

pipe = OmniGenPipeline.from_pretrained("Shitao/OmniGen-v1-diffusers", torch_dtype=torch.bfloat16) pipe.to("cuda")

prompt = "A cat holding a sign that says hello world"

image = pipe(prompt, num_inference_steps=50, guidance_scale=2.5).images[0] image.save("output.png")

Returns: ImagePipelineOutput or tuple: If return_dict is True, ImagePipelineOutput is returned, otherwise a tuple is returned where the first element is a list with the generated images.

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

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

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.

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.

encode_input_images

< source >

( input_pixel_values: typing.List[torch.Tensor] device: typing.Optional[torch.device] = None dtype: typing.Optional[torch.dtype] = None )

Parameters

• input_pixel_values — normalized pixel of input images
• device —

get the continue embedding of input images by VAE

Returns: torch.Tensor

< > Update on GitHub