HunyuanImage-2.1

[yjy415]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @yjy415, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces the full HunyuanImage 2.1 model architecture and its associated inference pipeline. The changes encompass a novel Diffusion Transformer (DiT) for robust image generation, a Variational Autoencoder (VAE) for efficient latent space manipulation, and a dual text encoding system utilizing both Qwen and ByT5 models to enhance prompt understanding, including support for glyphs and OCR-related features. The new pipeline provides a streamlined and modular approach to text-to-image synthesis, complete with custom scheduling and adaptive guidance mechanisms.

Highlights

• HunyuanImage 2.1 Model Integration: Introduced the complete HunyuanImage 2.1 model suite, including its Diffusion Transformer (DiT), Variational Autoencoder (VAE), and specialized text encoders (Qwen and ByT5).
• Core Diffusion Transformer (DiT): Added HYImageDiffusionTransformer with advanced features such as multimodal attention, rotary positional embeddings, FlashAttention for efficiency, and various modulation and normalization layers.
• Variational Autoencoder (VAE): Implemented HunyuanImageVAE2D for efficient encoding and decoding of images to and from a latent space, featuring Resnet blocks, attention mechanisms, and spatial down/upsampling.
• Dual Text Encoders for Richer Prompts: Integrated HunyuanImage_TextEncoder_Qwen (based on Qwen2.5-VL-7B-Instruct) for general text understanding and HunyuanImage_ByT5 (based on ByT5) with specialized glyph/OCR support for handling font and color tokens in prompts.
• Comprehensive Inference Pipeline: Developed HunyuanImagePipeline to orchestrate the entire image generation process, incorporating a custom flow-matching scheduler, adaptive projected guidance, and modular pipeline units for preprocessing and embedding.
• State Dictionary Converters: Provided dedicated state dictionary converters for the DiT, VAE, Qwen, and ByT5 models to ensure compatibility and proper loading of pre-trained weights.
• Example Usage: Included an example script (HunyuanImage.py) demonstrating how to download the necessary models and run an inference pipeline for text-to-image generation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for the HunyuanImage-2.1 model, adding new model implementations for the Diffusion Transformer (DiT), VAE, and text encoders (Qwen and T5), along with a new inference pipeline and corresponding state dictionary converters. The changes are extensive and well-structured. My review identified a few critical issues that could lead to runtime errors, such as improper handling of optional parameters and incorrect attribute access in the model implementations. I have also provided suggestions to improve code clarity, maintainability, and correctness by removing dead code, simplifying redundant logic, and correcting type hints.

[gemini-code-assist]

The call attn_mask.bool() will raise an AttributeError if attn_mask is None. This can happen if use_attention_mask is set to False in HYImageDiffusionTransformer, as the None value for the mask is passed down to this function.

You should handle the case where attn_mask is None by providing a default mask. Also, the .bool() call is redundant because the mask is already converted to a boolean tensor in IndividualTokenRefiner.forward.

Suggested change

	key_padding_mask=attn_mask.bool(),
	key_padding_mask=attn_mask if attn_mask is not None else torch.ones(x.shape[:2], dtype=torch.bool, device=x.device),

[gemini-code-assist]

The encode method returns a DiagonalGaussianDistribution object directly, which does not have a latent_dist attribute. Accessing .latent_dist will cause an AttributeError. You should use the returned posterior object directly.

Suggested change

	posterior = self.encode(sample).latent_dist
	posterior = self.encode(sample)

[gemini-code-assist]

The return type hint for this function is Union[torch.Tensor, Dict[str, torch.Tensor]], but the function actually returns a tuple (img, features_list, shape). This inconsistency can be misleading. Please update the type hint to match the actual return type for better code clarity and correctness.

Suggested change

	) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], tuple]:

[gemini-code-assist]

Using assert for control flow, especially for parameter validation, is not recommended. Asserts can be disabled globally (e.g., with Python's -O flag), which would cause this check to be skipped. It's better to raise a NotImplementedError to make it clear that this functionality is not supported.

Suggested change

	assert not return_dict, "return_dict is not supported."
	if return_dict: raise NotImplementedError("return_dict is not supported.")

[gemini-code-assist]

The conditional logic to handle empty tensors is redundant. Indexing a tensor with an empty index tensor (like true_idx or false_idx when they are empty) will correctly return an empty tensor of the appropriate shape. You can simplify this code by removing the if len(...) > 0 checks and the explicit creation of empty tensors.

Suggested change

	noise_pred_text_true = noise_pred_text[true_idx] if len(true_idx) > 0 else \
	torch.empty((0, noise_pred_text.size(1)), dtype=noise_pred_text.dtype, device=noise_pred_text.device)
	noise_pred_text_false = noise_pred_text[false_idx] if len(false_idx) > 0 else \
	torch.empty((0, noise_pred_text.size(1)), dtype=noise_pred_text.dtype, device=noise_pred_text.device)
	noise_pred_uncond_true = noise_pred_uncond[true_idx] if len(true_idx) > 0 else \
	torch.empty((0, noise_pred_uncond.size(1)), dtype=noise_pred_uncond.dtype, device=noise_pred_uncond.device)
	noise_pred_uncond_false = noise_pred_uncond[false_idx] if len(false_idx) > 0 else \
	torch.empty((0, noise_pred_uncond.size(1)), dtype=noise_pred_uncond.dtype, device=noise_pred_uncond.device)
	noise_pred_text_true = noise_pred_text[true_idx]
	noise_pred_text_false = noise_pred_text[false_idx]
	noise_pred_uncond_true = noise_pred_uncond[true_idx]
	noise_pred_uncond_false = noise_pred_uncond[false_idx]

[gemini-code-assist]

The hidden_size_map dictionary is initialized and populated but is never used anywhere in the function. This is dead code and should be removed to improve code clarity and maintainability.