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.

Beginne mit den zitierten z_image pipeline-, transformer-, ControlNet-, modular block- und test files im Commit 0f1abc4ae8b0eb2a3b40e82a310507281144c423. Führe die bereitgestellten Reproduktionen aus und untersuche bestehende Präzedenzfälle, wobei du die bereits durch PRs 13527 und 13267 abgedeckten Erkenntnisse getrennt behandelst. Als abgeschlossen gilt die Aufgabe, wenn die verbleibenden bestätigten Fehler behoben, nicht unterstützte Eingaben dokumentiert oder entfernt und die fehlende Abdeckung sowie die Prüfungen für die generierte Dokumentation hinzugefügt wurden.

Vom Indexierungsmodell aus dem Issue-Text verfasst.