hunyuan_image model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate-search status: searched hunyuan_image, HunyuanImagePipeline, AutoencoderKLHunyuanImage, AutoencoderKLHunyuanImageRefiner, plus targeted tiling/shortcut/shape terms. Broad matches were the original integration PR #12333 and HunyuanImage 3.0 request #12412; I did not find duplicates for the issues below. A second gh search batch hit GitHub API rate limits, so I completed targeted duplicate checks with web search.

Issue 1: Base VAE default scaling factor breaks pipeline decode

Affected code:

Problem:
AutoencoderKLHunyuanImage defaults scaling_factor to None, but HunyuanImagePipeline always divides latents by self.vae.config.scaling_factor before decoding.

Impact:
A pipeline assembled from default tiny components can denoise successfully and then fail at decode with TypeError. The refiner VAE already has a numeric default.

Reproduction:

import torch
from diffusers import AutoencoderKLHunyuanImage, FlowMatchEulerDiscreteScheduler, HunyuanImagePipeline, HunyuanImageTransformer2DModel

transformer = HunyuanImageTransformer2DModel(
    in_channels=4, out_channels=4, num_attention_heads=1, attention_head_dim=4,
    num_layers=0, num_single_layers=0, num_refiner_layers=0,
    patch_size=(1, 1), text_embed_dim=4, text_embed_2_dim=4, rope_axes_dim=(2, 2),
)
vae = AutoencoderKLHunyuanImage(
    in_channels=3, out_channels=3, latent_channels=4, block_out_channels=(32,),
    layers_per_block=1, spatial_compression_ratio=8, sample_size=32,
)
pipe = HunyuanImagePipeline(
    scheduler=FlowMatchEulerDiscreteScheduler(), vae=vae, text_encoder=None, tokenizer=None,
    text_encoder_2=None, tokenizer_2=None, transformer=transformer, guider=None, ocr_guider=None,
)
pipe(
    prompt=None,
    prompt_embeds=torch.randn(1, 1, 4),
    prompt_embeds_mask=torch.ones(1, 1, dtype=torch.long),
    prompt_embeds_2=torch.zeros(1, 1, 4),
    prompt_embeds_mask_2=torch.zeros(1, 1, dtype=torch.long),
    height=16, width=16, num_inference_steps=1, output_type="pt",
)

Relevant precedent:

Suggested fix:

scaling_factor: float = 0.476986

Issue 2: HunyuanImagePipelineOutput is not exported from the pipeline package

Affected code:

Problem:
The output class exists and is documented, but diffusers.pipelines.hunyuan_image only lazy-exports the two pipelines.

Impact:
Users cannot import the output type from the package namespace, unlike similar pipeline families.

Reproduction:

from diffusers.pipelines.hunyuan_image import HunyuanImagePipelineOutput

Relevant precedent:

Suggested fix:

_import_structure["pipeline_output"] = ["HunyuanImagePipelineOutput"]
...
from .pipeline_output import HunyuanImagePipelineOutput

Issue 3: Base VAE tiled encode crashes on 4D image tensors

Affected code:

Problem:
AutoencoderKLHunyuanImage is a 2D image VAE, but tiled_encode() unpacks x.shape as if it were 5D and slices tiles with an extra temporal dimension.

Impact:
vae.enable_tiling(); vae.encode(image) fails for image tensors large enough to trigger tiling. Existing tests cover tiled decode through the pipeline, but not tiled encode.

Reproduction:

import torch
from diffusers import AutoencoderKLHunyuanImage

vae = AutoencoderKLHunyuanImage(
    in_channels=3, out_channels=3, latent_channels=4, block_out_channels=(32,),
    layers_per_block=1, spatial_compression_ratio=1, sample_size=8,
)
vae.enable_tiling(tile_sample_min_size=4)
vae.encode(torch.randn(1, 3, 8, 8))

Relevant precedent:
The same file's tiled_decode() correctly treats base VAE latents as 4D:

Suggested fix:

_, _, height, width = x.shape
...
tile = x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
...
result_row.append(tile[:, :, :row_limit, :row_limit])

Issue 4: Refiner VAE tiling mixes sample and latent units

Affected code:

Problem:
tiled_encode() iterates over sample-space height/width using latent-space overlap values. tiled_decode() crops decoded sample tiles using latent-space row limits.

Impact:
Tiled encode can generate invalid edge tiles, and tiled decode returns the wrong spatial size.

Reproduction:

import torch
from diffusers import AutoencoderKLHunyuanImageRefiner

vae = AutoencoderKLHunyuanImageRefiner(
    in_channels=3, out_channels=3, latent_channels=4, block_out_channels=(8, 8),
    layers_per_block=0, spatial_compression_ratio=2, temporal_compression_ratio=1,
)
z = torch.randn(1, 4, 1, 8, 8)
plain = vae.decode(z).sample.shape
vae.enable_tiling(tile_sample_min_height=8, tile_sample_min_width=8, tile_overlap_factor=0.25)
tiled = vae.decode(z).sample.shape
print(plain, tiled)  # plain is [1, 3, 1, 16, 16], tiled is [1, 3, 1, 9, 9]

Relevant precedent:

Suggested fix:
Use sample-space stride/overlap for encode iteration, latent-space crop sizes for encoded tiles, latent-space stride for decode iteration, and sample-space crop sizes for decoded tiles. The existing tile_sample_stride_height / tile_sample_stride_width fields should either be used or removed.

Issue 5: Base VAE shortcut projection is applied to the wrong tensor

Affected code:

Problem:
When in_channels != out_channels, HunyuanImageResnetBlock.forward() applies conv_shortcut to x after conv2, not to the saved residual.

Impact:
Non-default serialized configs with downsample_match_channel=False or upsample_match_channel=False crash when a resnet block changes channels.

Reproduction:

import torch
from diffusers import AutoencoderKLHunyuanImage

vae = AutoencoderKLHunyuanImage(
    in_channels=3, out_channels=3, latent_channels=4, block_out_channels=(32, 64),
    layers_per_block=1, spatial_compression_ratio=2, sample_size=16,
    downsample_match_channel=False,
)
vae.encode(torch.randn(1, 3, 16, 16))

Relevant precedent:

Suggested fix:

if self.conv_shortcut is not None:
    residual = self.conv_shortcut(residual)
return x + residual

Issue 6: Transformer silently truncates latent sizes not divisible by patch_size

Affected code:

Problem:
The transformer computes post_patch_sizes = d // p and unpatchifies to post_patch * patch, silently dropping remainder pixels. The pipelines warn about divisibility but do not round latent sizes to the transformer patch grid.

Impact:
Custom configs with patch_size > 1 can produce a model output smaller than the input latents, leading to scheduler shape errors or silent shape loss when the model is used directly.

Reproduction:

import torch
from diffusers import HunyuanImageTransformer2DModel

model = HunyuanImageTransformer2DModel(
    in_channels=4, out_channels=4, num_attention_heads=1, attention_head_dim=4,
    num_layers=0, num_single_layers=0, num_refiner_layers=0,
    patch_size=(2, 2), text_embed_dim=4, rope_axes_dim=(2, 2),
)
out = model(
    hidden_states=torch.randn(1, 4, 3, 3),
    timestep=torch.ones(1),
    encoder_hidden_states=torch.randn(1, 1, 4),
    encoder_attention_mask=torch.ones(1, 1, dtype=torch.long),
).sample
print(out.shape)  # torch.Size([1, 4, 2, 2])

Relevant precedent:

Suggested fix:

if any(size % patch != 0 for size, patch in zip(sizes, self.config.patch_size)):
    raise ValueError(f"`hidden_states` spatial/temporal sizes {sizes} must be divisible by patch_size {self.config.patch_size}.")

Issue 7: Base VAE has checkpointing code but disables the public capability flag

Affected code:

Problem:
The encoder and decoder contain self._gradient_checkpointing_func(...) branches, but AutoencoderKLHunyuanImage._supports_gradient_checkpointing is False.

Impact:
Training or fine-tuning code cannot enable gradient checkpointing on the base VAE even though the implementation paths exist.

Reproduction:

from diffusers import AutoencoderKLHunyuanImage

vae = AutoencoderKLHunyuanImage(
    in_channels=3, out_channels=3, latent_channels=4, block_out_channels=(32,),
    layers_per_block=1, spatial_compression_ratio=1, sample_size=8,
)
vae.enable_gradient_checkpointing()

Relevant precedent:

Suggested fix:

_supports_gradient_checkpointing = True

Issue 8: Test coverage is incomplete, including missing slow tests

Affected code:

Problem:
There is one fast pipeline test class for HunyuanImagePipeline. I found no model fast tests for HunyuanImageTransformer2DModel, AutoencoderKLHunyuanImage, or AutoencoderKLHunyuanImageRefiner, no fast tests for HunyuanImageRefinerPipeline, and no @slow tests for this family.

Impact:
The broken VAE encode tiling, refiner tiling, refiner pipeline paths, model patch-size behavior, and default scaling-factor failure are not covered.

Reproduction:

from pathlib import Path

files = list(Path("tests").rglob("*.py"))
hunyuan_files = [p for p in files if "hunyuan_image" in str(p).lower() or "hunyuanimage" in p.read_text(encoding="utf-8", errors="ignore")]
print([str(p) for p in hunyuan_files])
print("has_refiner_pipeline_test", any("HunyuanImageRefinerPipeline" in p.read_text(encoding="utf-8", errors="ignore") for p in hunyuan_files))
print("has_slow_test", any("@slow" in p.read_text(encoding="utf-8", errors="ignore") for p in hunyuan_files))

Relevant precedent:
Other families keep slow pipeline tests alongside fast tests, for example:

Suggested fix:
Add focused model tests under tests/models/ for the transformer and both VAEs, add fast refiner pipeline tests with tiny components, and add at least one slow test for the base and refiner pretrained checkpoints. I attempted the existing fast Hunyuan test file locally with .venv, but collection failed before running tests because this environment's PyTorch build lacks torch._C._distributed_c10d.