z_image model/pipeline review

# `z_image` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Duplicate search status: checked GitHub Issues/PRs for `z_image`, affected class names, and failure modes. Duplicates found for Issue 1 and Issue 3; noted below.

## Issue 1: `ZImageOmniPipeline` crashes when `guidance_scale=0`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py#L579-L592

Problem:
`negative_condition_siglip_embeds` is only assigned inside `if self.do_classifier_free_guidance`, but it is normalized unconditionally immediately after. The public example uses `guidance_scale=0.0`, so the documented Omni path raises before denoising.

Duplicate:
Already covered by open PR https://github.com/huggingface/diffusers/pull/13527. This is not a new finding.

Impact:
`ZImageOmniPipeline(..., guidance_scale=0.0)` fails for the documented turbo/no-CFG usage.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, ZImageOmniPipeline, ZImageTransformer2DModel

transformer = ZImageTransformer2DModel(
all_patch_size=(2,), all_f_patch_size=(1,), in_channels=4, dim=16,
n_layers=1, n_refiner_layers=1, n_heads=2, n_kv_heads=2,
cap_feat_dim=8, axes_dims=[4, 2, 2], axes_lens=[32, 32, 32],
)
vae = AutoencoderKL(
in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"],
block_out_channels=[16], layers_per_block=1, latent_channels=4, norm_num_groups=4, sample_size=32,
scaling_factor=0.3611, shift_factor=0.1159,
)
pipe = ZImageOmniPipeline(FlowMatchEulerDiscreteScheduler(), vae, None, None, transformer, None, None)
pipe(prompt_embeds=[[torch.randn(3, 8)]], height=32, width=32, num_inference_steps=1, guidance_scale=0.0, output_type="latent")
```

Relevant precedent:
Open duplicate PR: https://github.com/huggingface/diffusers/pull/13527

Suggested fix:
```python
condition_siglip_embeds = [None if sels == [] else sels + [None] for sels in condition_siglip_embeds]
if self.do_classifier_free_guidance:
negative_condition_siglip_embeds = [
None if sels == [] else sels + [None] for sels in negative_condition_siglip_embeds
]
else:
negative_condition_siglip_embeds = None
```

## Issue 2: Omni condition-image encoding hard-casts VAE input to `bfloat16`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py#L293-L306

Problem:
`prepare_image_latents()` calls `self.vae.encode(image.bfloat16())` regardless of the VAE dtype. A float32 VAE on CPU receives bf16 inputs with float32 weights and raises a dtype mismatch.

Impact:
Omni image-conditioned generation fails outside the exact bf16 VAE setup. It also violates the dtype/device handling rule by hardcoding a dtype in pipeline runtime code.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, ZImageOmniPipeline

vae = AutoencoderKL(
in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"],
block_out_channels=[16], layers_per_block=1, latent_channels=4, norm_num_groups=4, sample_size=32,
scaling_factor=0.3611, shift_factor=0.1159,
)
pipe = object.__new__(ZImageOmniPipeline)
pipe.vae = vae
pipe.prepare_image_latents([torch.rand(1, 3, 32, 32)], 1, torch.device("cpu"), torch.float32)
```

Relevant precedent:
Other image-encoding paths convert to the requested/vae dtype before `vae.encode`, not a fixed bf16 dtype.

Suggested fix:
```python
vae_dtype = self.vae.dtype
image = image.to(device=device, dtype=vae_dtype)
image_latent = (
self.vae.encode(image).latent_dist.mode()[0] - self.vae.config.shift_factor
) * self.vae.config.scaling_factor
image_latent = image_latent.unsqueeze(1).to(dtype)
```

## Issue 3: `ZImageControlNetModel` has gradient-checkpointing flag but never initializes it

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_z_image.py#L433-L517
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_z_image.py#L753-L833

Problem:
The model sets `_supports_gradient_checkpointing = True` and branches on `self.gradient_checkpointing`, but `__init__` never sets `self.gradient_checkpointing = False`.

Duplicate:
Already covered by open PR https://github.com/huggingface/diffusers/pull/13267. This is not a new finding.

Impact:
Direct grad-enabled forward, training, or checkpointing setup fails with `AttributeError`.

Reproduction:
```python
import torch
from diffusers import ZImageControlNetModel, ZImageTransformer2DModel

transformer = ZImageTransformer2DModel(
all_patch_size=(2,), all_f_patch_size=(1,), in_channels=4, dim=16,
n_layers=1, n_refiner_layers=1, n_heads=2, n_kv_heads=2,
cap_feat_dim=8, axes_dims=[4, 2, 2], axes_lens=[64, 64, 64],
)
controlnet = ZImageControlNetModel(
control_layers_places=[0], control_refiner_layers_places=[0], control_in_dim=4,
all_patch_size=(2,), all_f_patch_size=(1,), dim=16, n_refiner_layers=1,
n_heads=2, n_kv_heads=2,
)
controlnet = ZImageControlNetModel.from_transformer(controlnet, transformer)
controlnet([torch.randn(4, 1, 32, 32)], torch.tensor([0.5]), [torch.randn(3, 8)], [torch.randn(4, 1, 32, 32)])
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_qwenimage.py#L101

Suggested fix:
```python
self.gradient_checkpointing = False
```

## Issue 4: `ZImageInpaintPipeline.masked_image_latents` is accepted but ignored

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py#L537-L576
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py#L788-L797

Problem:
The public argument says precomputed masked latents skip encoding, but the value is never passed into `prepare_mask_latents()` and never affects denoising. Different supplied `masked_image_latents` produce identical outputs.

Impact:
Users cannot actually provide precomputed masked latents, and the callback tensor implies a state value that does not participate in generation.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, ZImageInpaintPipeline, ZImageTransformer2DModel

torch.manual_seed(0)
transformer = ZImageTransformer2DModel(
all_patch_size=(2,), all_f_patch_size=(1,), in_channels=4, dim=16,
n_layers=1, n_refiner_layers=1, n_heads=2, n_kv_heads=2,
cap_feat_dim=8, axes_dims=[4, 2, 2], axes_lens=[64, 64, 64],
)
vae = AutoencoderKL(
in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"],
block_out_channels=[16], layers_per_block=1, latent_channels=4, norm_num_groups=4, sample_size=32,
scaling_factor=0.3611, shift_factor=0.1159,
)
pipe = ZImageInpaintPipeline(FlowMatchEulerDiscreteScheduler(), vae, None, None, transformer)
pipe.set_progress_bar_config(disable=True)
kwargs = dict(
prompt_embeds=[torch.randn(3, 8)], image=torch.rand(1, 3, 32, 32), mask_image=torch.ones(1, 1, 32, 32),
height=32, width=32, num_inference_steps=1, guidance_scale=0.0, output_type="latent",
latents=torch.randn(1, 4, 32, 32),
)
a = pipe(**kwargs, masked_image_latents=torch.zeros(1, 4, 32, 32), generator=torch.Generator().manual_seed(123)).images
b = pipe(**kwargs, masked_image_latents=torch.randn(1, 4, 32, 32), generator=torch.Generator().manual_seed(123)).images
print(torch.equal(a, b))
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py#L1204-L1209

Suggested fix:
If Z-Image inpaint is intended to be latent-blending only, remove or deprecate `masked_image_latents` and the callback tensor. If it is intended to match SD-style inpaint conditioning, thread the provided tensor through `prepare_mask_latents()` and into the model input path.

## Issue 5: Model dtype rules are violated in shared transformer/controlnet helpers

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_z_image.py#L65
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_z_image.py#L327-L332
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_z_image.py#L67
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_z_image.py#L304-L309

Problem:
`TimestepEmbedder.forward()` casts by reading `self.mlp[0].weight.dtype`, and `RopeEmbedder.precompute_freqs_cis()` unconditionally constructs float64 tensors. Both patterns are explicitly called out in the model review rules.

Impact:
This is fragile for quantized/GGUF/layerwise-casting loads and violates backend portability expectations for MPS/NPU-style environments.

Reproduction:
```python
from pathlib import Path

for path in [
"src/diffusers/models/transformers/transformer_z_image.py",
"src/diffusers/models/controlnets/controlnet_z_image.py",
]:
for i, line in enumerate(Path(path).read_text().splitlines(), 1):
if "weight.dtype" in line or "torch.float64" in line:
print(path, i, line.strip())
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L508-L516

Suggested fix:
Use float32 for RoPE precompute unless there is measured need for gated float64, and pass the desired activation dtype from the caller into `TimestepEmbedder` instead of reading parameter storage dtype.

## Issue 6: Modular pipeline generated docs still contain TODO placeholders

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py#L46-L75

Problem:
`modular_blocks_z_image.py` contains many generated `TODO: Add description.` entries. The modular review rules explicitly require generated modular docstrings to be fixed after running auto-docstring generation.

Impact:
Public modular pipeline docs/API metadata are incomplete for several inputs, including `height`, `width`, `latents`, `generator`, `sigmas`, and workflow-specific inputs.

Reproduction:
```python
from pathlib import Path

path = Path("src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py")
print(sum("TODO: Add description." in line for line in path.read_text().splitlines()))
```

Relevant precedent:
`.ai/modular.md` conversion checklist requires running `utils/modular_auto_docstring.py --fix_and_overwrite` and resolving TODO placeholders.

Suggested fix:
Add accurate `InputParam` descriptions/types for the missing fields, rerun `python utils/modular_auto_docstring.py --fix_and_overwrite`, and verify no generated TODOs remain.

## Issue 7: Coverage gaps: no slow tests, no ControlNet/Omni pipeline tests, and docs omit public variants

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/z_image/__init__.py#L24-L30
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/z_image/test_z_image.py#L45
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/api/pipelines/z_image.md#L91-L105

Problem:
The package exports `ZImageControlNetPipeline`, `ZImageControlNetInpaintPipeline`, and `ZImageOmniPipeline`, but `tests/pipelines/z_image/` only has fast tests for text2img/img2img/inpaint. There are no z_image slow tests. The pipeline docs only autodoc the three non-ControlNet/non-Omni pipelines.

Impact:
The exact Omni no-CFG crash and ControlNet checkpointing issue above are not covered by pipeline tests. Real-checkpoint regressions are also unguarded.

Reproduction:
```python
from pathlib import Path

test_text = "\n".join(p.read_text() for p in Path("tests").rglob("*z_image*.py"))
docs = Path("docs/source/en/api/pipelines/z_image.md").read_text()

print("@slow" in test_text or "slow(" in test_text)
print("ZImageOmniPipeline" in test_text, "ZImageControlNetPipeline" in test_text)
print("ZImageOmniPipeline" in docs, "ZImageControlNetPipeline" in docs)
```

Relevant precedent:
Most mature pipeline families include at least one `@slow` real-checkpoint smoke test for public pipelines, plus fast tests for every exported variant.

Suggested fix:
Add fast tests for Omni and both ControlNet pipelines using tiny fixtures, add at least one slow real-checkpoint smoke test for the z_image family, and add autodoc sections for the public ControlNet and Omni pipelines.

Verification performed:
- Minimal `.venv` snippets confirmed Issues 1-4.
- `tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py -q`: `14 passed`.
- Pipeline fast tests could not be collected in this `.venv` because the installed PyTorch build lacks `torch._C._distributed_c10d`, imported via shared training test utilities.

Inizia dai file citati z_image pipeline, transformer, ControlNet, modular block e test al commit 0f1abc4ae8b0eb2a3b40e82a310507281144c423. Esegui le riproduzioni fornite e analizza i precedenti esistenti, separando i risultati già coperti dalle PR 13527 e 13267. Il lavoro è completo quando i difetti confermati rimanenti sono stati risolti, gli input non supportati sono stati documentati o rimossi e sono stati aggiunti la coverage mancante e i controlli sulla documentazione generata.

Scritto dal modello di indicizzazione a partire dal testo della issue.