Skip to content

Fix/z image controlnet logging (supersedes #9289) #9298

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lstein merged 3 commits into from
Jun 24, 2026
Merged

Fix/z image controlnet logging (supersedes #9289) #9298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b1b0ee0
Route Z-Image ControlNet diagnostics through logging
wunianze666-netizen Jun 13, 2026
edd6982
Merge branch 'main' into fix/z-image-controlnet-logging
lstein Jun 24, 2026
38c6a3d
Merge branch 'main' into fix/z-image-controlnet-logging
lstein Jun 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 28 additions & 31 deletions invokeai/backend/z_image/z_image_controlnet_extension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
the base transformer's forward pass.
"""

import logging
from typing import List, Optional, Tuple

import torch
Expand All @@ -16,6 +17,8 @@
from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter
from invokeai.backend.z_image.z_image_patchify_utils import SEQ_MULTI_OF, patchify_control_context

logger = logging.getLogger(__name__)


class ZImageControlNetExtension:
"""Extension for Z-Image ControlNet - computes control hints without duplicating the transformer.
Expand Down Expand Up @@ -57,17 +60,15 @@ def __init__(
# This matches the default configuration in the original implementation
self._control_places = [i * 2 for i in range(self._num_control_blocks)]

# DEBUG: Print control configuration
print(f"[DEBUG] Actual num_control_blocks: {self._num_control_blocks}")
print(f"[DEBUG] control_places: {self._control_places}")
logger.debug("Actual num_control_blocks: %s", self._num_control_blocks)
logger.debug("control_places: %s", self._control_places)

# DEBUG: Check if control_layers have non-zero weights
first_layer = control_adapter.control_layers[0]
if hasattr(first_layer, "after_proj"):
after_proj_norm = first_layer.after_proj.weight.norm().item()
print(f"[DEBUG] First control layer after_proj weight norm: {after_proj_norm}")
logger.debug("First control layer after_proj weight norm: %s", after_proj_norm)
if after_proj_norm < 1e-6:
print("[WARNING] after_proj weights are near-zero! Weights may not be loaded correctly.")
logger.warning("after_proj weights are near-zero! Weights may not be loaded correctly.")

@property
def weight(self) -> float:
Expand Down Expand Up @@ -157,10 +158,9 @@ def prepare_control_state(

control_unified = pad_sequence(control_unified, batch_first=True, padding_value=0.0)

# DEBUG (only once)
if not hasattr(self, "_prepare_printed"):
self._prepare_printed = True
print(f"[DEBUG] Control state prepared: shape {control_unified.shape}")
logger.debug("Control state prepared: shape %s", control_unified.shape)

return control_unified

Expand Down Expand Up @@ -304,26 +304,22 @@ def compute_hints(
control_unified = pad_sequence(control_unified, batch_first=True, padding_value=0.0)
c = control_unified

# Process through control_layers to generate hints
# DEBUG: Print shapes before control_layers (only on first call)
if not hasattr(self, "_debug_printed"):
self._debug_printed = True
print(f"[DEBUG] control_unified shape: {control_unified.shape}")
print(f"[DEBUG] unified_hidden_states shape: {unified_hidden_states.shape}")
print(f"[DEBUG] ctrl_item_seqlens: {ctrl_item_seqlens}, x_item_seqlens: {x_item_seqlens}")
logger.debug("control_unified shape: %s", control_unified.shape)
logger.debug("unified_hidden_states shape: %s", unified_hidden_states.shape)
logger.debug("ctrl_item_seqlens: %s, x_item_seqlens: %s", ctrl_item_seqlens, x_item_seqlens)

# Check weight norms of critical layers
layer0 = self._adapter.control_layers[0]
if hasattr(layer0, "before_proj"):
print(f"[DEBUG] before_proj weight norm: {layer0.before_proj.weight.norm().item():.6f}")
logger.debug("before_proj weight norm: %.6f", layer0.before_proj.weight.norm().item())
if hasattr(layer0, "after_proj"):
print(f"[DEBUG] after_proj weight norm: {layer0.after_proj.weight.norm().item():.6f}")
logger.debug("after_proj weight norm: %.6f", layer0.after_proj.weight.norm().item())

# Check control_noise_refiner weights
if len(self._adapter.control_noise_refiner) > 0:
refiner0 = self._adapter.control_noise_refiner[0]
if hasattr(refiner0, "attn"):
print(f"[DEBUG] noise_refiner[0] attn.wq norm: {refiner0.attn.wq.weight.norm().item():.6f}")
logger.debug("noise_refiner[0] attn.wq norm: %.6f", refiner0.attn.wq.weight.norm().item())

for layer in self._adapter.control_layers:
c = layer(
Expand All @@ -337,16 +333,19 @@ def compute_hints(
# Extract hints (all but the last element which is the running state)
hints = tuple(torch.unbind(c)[:-1])

# DEBUG: Print hint shapes (only on first call)
if not hasattr(self, "_hints_printed"):
self._hints_printed = True
print(f"[DEBUG] Number of hints: {len(hints)}")
logger.debug("Number of hints: %s", len(hints))
if hints:
print(f"[DEBUG] First hint shape: {hints[0].shape}")
# Also check hint statistics for each hint
logger.debug("First hint shape: %s", hints[0].shape)
for i, h in enumerate(hints[:3]): # First 3 hints
print(
f"[DEBUG] Hint[{i}] mean: {h.mean().item():.6f}, std: {h.std().item():.6f}, min: {h.min().item():.6f}, max: {h.max().item():.6f}"
logger.debug(
"Hint[%s] mean: %.6f, std: %.6f, min: %.6f, max: %.6f",
i,
h.mean().item(),
h.std().item(),
h.min().item(),
h.max().item(),
)

return hints
Expand Down Expand Up @@ -476,10 +475,9 @@ def z_image_forward_with_control(
control_weight: float = 1.0
hints: Optional[Tuple[torch.Tensor, ...]] = None

# DEBUG: Print number of transformer layers (only once per session)
if not hasattr(z_image_forward_with_control, "_layers_printed"):
z_image_forward_with_control._layers_printed = True
print(f"[DEBUG] Base transformer has {len(transformer.layers)} layers")
logger.debug("Base transformer has %s layers", len(transformer.layers))

if control_extension is not None:
# Compute ALL hints at once using the INITIAL unified state (before main layers run)
Expand Down Expand Up @@ -511,13 +509,12 @@ def z_image_forward_with_control(
if control_layer_idx >= skip_layers:
hint = hints[control_layer_idx]

# DEBUG: Print on first injection
if not hasattr(z_image_forward_with_control, "_injection_printed"):
z_image_forward_with_control._injection_printed = True
print(f"[DEBUG] Injection at layer {layer_idx} (control_layer {control_layer_idx})")
print(f"[DEBUG] Hint mean: {hint.mean().item():.6f}, std: {hint.std().item():.6f}")
print(f"[DEBUG] Unified mean: {unified.mean().item():.6f}, std: {unified.std().item():.6f}")
print(f"[DEBUG] control_weight: {control_weight}, skip_layers: {skip_layers}")
logger.debug("Injection at layer %s (control_layer %s)", layer_idx, control_layer_idx)
logger.debug("Hint mean: %.6f, std: %.6f", hint.mean().item(), hint.std().item())
logger.debug("Unified mean: %.6f, std: %.6f", unified.mean().item(), unified.std().item())
logger.debug("control_weight: %s, skip_layers: %s", control_weight, skip_layers)

unified = unified + hint * control_weight

Expand Down
24 changes: 24 additions & 0 deletions tests/backend/z_image/test_z_image_controlnet_extension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
from types import SimpleNamespace
from unittest.mock import patch

import torch

from invokeai.backend.z_image.z_image_controlnet_extension import ZImageControlNetExtension


def test_init_logs_control_adapter_diagnostics_without_stdout(capsys):
control_adapter = SimpleNamespace(
control_layers=[
SimpleNamespace(
after_proj=SimpleNamespace(weight=torch.zeros(1)),
)
]
)

with patch("invokeai.backend.z_image.z_image_controlnet_extension.logger") as logger_mock:
ZImageControlNetExtension(control_adapter=control_adapter, control_cond=torch.zeros(1))

captured = capsys.readouterr()
assert captured.out == ""
assert logger_mock.debug.call_count == 3
logger_mock.warning.assert_called_once()
Loading