import inspect import math from contextlib import ExitStack from typing import Callable, Iterator, Optional, Tuple import einops import torch import torchvision.transforms as tv_transforms from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image from torchvision.transforms.functional import resize as tv_resize from tqdm import tqdm from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.app.invocations.fields import ( DenoiseMaskField, FieldDescriptions, Input, InputField, LatentsField, ZImageConditioningField, ) from invokeai.app.invocations.latent_noise import validate_noise_tensor_shape from invokeai.app.invocations.model import TransformerField, VAEField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.invocations.z_image_control import ZImageControlField from invokeai.app.invocations.z_image_image_to_latents import ZImageImageToLatentsInvocation from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.schedulers import ZIMAGE_SCHEDULER_LABELS, ZIMAGE_SCHEDULER_MAP, ZIMAGE_SCHEDULER_NAME_VALUES from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ZImageConditioningInfo from invokeai.backend.util.devices import TorchDevice from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension from invokeai.backend.z_image.text_conditioning import ZImageTextConditioning from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter from invokeai.backend.z_image.z_image_controlnet_extension import ( ZImageControlNetExtension, z_image_forward_with_control, ) from invokeai.backend.z_image.z_image_transformer_patch import patch_transformer_for_regional_prompting @invocation( "z_image_denoise", title="Denoise - Z-Image", tags=["image", "z-image"], category="latents", version="1.6.0", classification=Classification.Prototype, ) class ZImageDenoiseInvocation(BaseInvocation): """Run the denoising process with a Z-Image model. Supports regional prompting by connecting multiple conditioning inputs with masks. """ # If latents is provided, this means we are doing image-to-image. latents: Optional[LatentsField] = InputField( default=None, description=FieldDescriptions.latents, input=Input.Connection ) noise: Optional[LatentsField] = InputField( default=None, description=FieldDescriptions.noise, input=Input.Connection ) # denoise_mask is used for image-to-image inpainting. Only the masked region is modified. denoise_mask: Optional[DenoiseMaskField] = InputField( default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection ) denoising_start: float = InputField(default=0.0, ge=0, le=1, description=FieldDescriptions.denoising_start) denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end) add_noise: bool = InputField(default=True, description="Add noise based on denoising start.") transformer: TransformerField = InputField( description=FieldDescriptions.z_image_model, input=Input.Connection, title="Transformer" ) positive_conditioning: ZImageConditioningField | list[ZImageConditioningField] = InputField( description=FieldDescriptions.positive_cond, input=Input.Connection ) negative_conditioning: ZImageConditioningField | list[ZImageConditioningField] | None = InputField( default=None, description=FieldDescriptions.negative_cond, input=Input.Connection ) # Z-Image-Turbo works best without CFG (guidance_scale=1.0) guidance_scale: float = InputField( default=1.0, ge=1.0, description="Guidance scale for classifier-free guidance. 1.0 = no CFG (recommended for Z-Image-Turbo). " "Values > 1.0 amplify guidance.", title="Guidance Scale", ) width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.") height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.") # Z-Image-Turbo uses 8 steps by default steps: int = InputField(default=8, gt=0, description="Number of denoising steps. 8 recommended for Z-Image-Turbo.") seed: int = InputField(default=0, description="Randomness seed for reproducibility.") # Z-Image Control support control: Optional[ZImageControlField] = InputField( default=None, description="Z-Image control conditioning for spatial control (Canny, HED, Depth, Pose, MLSD).", input=Input.Connection, ) # VAE for encoding control images (required when using control) vae: Optional[VAEField] = InputField( default=None, description=FieldDescriptions.vae + " Required for control conditioning.", input=Input.Connection, ) # Shift override for the sigma schedule. If None, shift is auto-calculated from image dimensions. shift: Optional[float] = InputField( default=None, ge=0.0, description="Override the timestep shift (mu) for the sigma schedule. " "Leave blank to auto-calculate based on image dimensions (recommended). " "Lower values (~0.5) produce less noise shifting, higher values (~1.15) produce more.", title="Shift", ) # Scheduler selection for the denoising process scheduler: ZIMAGE_SCHEDULER_NAME_VALUES = InputField( default="euler", description="Scheduler (sampler) for the denoising process. Euler is the default and recommended. " "Heun is 2nd-order (better quality, 2x slower). LCM works with Turbo only (not Base).", ui_choice_labels=ZIMAGE_SCHEDULER_LABELS, ) @torch.no_grad() def invoke(self, context: InvocationContext) -> LatentsOutput: latents = self._run_diffusion(context) latents = latents.detach().to("cpu") name = context.tensors.save(tensor=latents) return LatentsOutput.build(latents_name=name, latents=latents, seed=None) def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> torch.Tensor | None: """Prepare the inpaint mask.""" if self.denoise_mask is None: return None mask = context.tensors.load(self.denoise_mask.mask_name) # Invert mask: 0.0 = regions to denoise, 1.0 = regions to preserve mask = 1.0 - mask _, _, latent_height, latent_width = latents.shape mask = tv_resize( img=mask, size=[latent_height, latent_width], interpolation=tv_transforms.InterpolationMode.BILINEAR, antialias=False, ) mask = mask.to(device=latents.device, dtype=latents.dtype) return mask def _load_text_conditioning( self, context: InvocationContext, cond_field: ZImageConditioningField | list[ZImageConditioningField], img_height: int, img_width: int, dtype: torch.dtype, device: torch.device, ) -> list[ZImageTextConditioning]: """Load Z-Image text conditioning with optional regional masks. Args: context: The invocation context. cond_field: Single conditioning field or list of fields. img_height: Height of the image token grid (H // patch_size). img_width: Width of the image token grid (W // patch_size). dtype: Target dtype. device: Target device. Returns: List of ZImageTextConditioning objects with embeddings and masks. """ # Normalize to a list cond_list = [cond_field] if isinstance(cond_field, ZImageConditioningField) else cond_field text_conditionings: list[ZImageTextConditioning] = [] for cond in cond_list: # Load the text embeddings cond_data = context.conditioning.load(cond.conditioning_name) assert len(cond_data.conditionings) == 1 z_image_conditioning = cond_data.conditionings[0] assert isinstance(z_image_conditioning, ZImageConditioningInfo) z_image_conditioning = z_image_conditioning.to(dtype=dtype, device=device) prompt_embeds = z_image_conditioning.prompt_embeds # Load the mask, if provided mask: torch.Tensor | None = None if cond.mask is not None: mask = context.tensors.load(cond.mask.tensor_name) mask = mask.to(device=device) mask = ZImageRegionalPromptingExtension.preprocess_regional_prompt_mask( mask, img_height, img_width, dtype, device ) text_conditionings.append(ZImageTextConditioning(prompt_embeds=prompt_embeds, mask=mask)) return text_conditionings def _get_noise( self, batch_size: int, num_channels_latents: int, height: int, width: int, dtype: torch.dtype, device: torch.device, seed: int, ) -> torch.Tensor: """Generate initial noise tensor.""" # Generate noise as float32 on CPU for maximum compatibility, # then cast to target dtype/device rand_device = "cpu" rand_dtype = torch.float32 return torch.randn( batch_size, num_channels_latents, int(height) // LATENT_SCALE_FACTOR, int(width) // LATENT_SCALE_FACTOR, device=rand_device, dtype=rand_dtype, generator=torch.Generator(device=rand_device).manual_seed(seed), ).to(device=device, dtype=dtype) def _calculate_shift( self, image_seq_len: int, base_image_seq_len: int = 256, max_image_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ) -> float: """Calculate timestep shift based on image sequence length. Based on diffusers ZImagePipeline.calculate_shift method. Returns a linear shift value (exp(mu) from the original formula). """ import math m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len) b = base_shift - m * base_image_seq_len mu = image_seq_len * m + b # Convert from exponential mu to linear shift value return math.exp(mu) def _get_sigmas(self, shift: float, num_steps: int) -> list[float]: """Generate sigma schedule with linear time shift. Uses linear time shift: shift / (shift + (1/t - 1)). The shift value is used directly as a multiplier. Generates num_steps + 1 sigma values (including terminal 0.0). """ def time_shift(shift: float, t: float) -> float: """Apply linear time shift to a single timestep value.""" if t <= 0: return 0.0 if t >= 1: return 1.0 return shift / (shift + (1 / t - 1)) sigmas = [] for i in range(num_steps + 1): t = 1.0 - i / num_steps # Goes from 1.0 to 0.0 sigma = time_shift(shift, t) sigmas.append(sigma) return sigmas def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: device = TorchDevice.choose_torch_device() inference_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) transformer_info = context.models.load(self.transformer.transformer) # Calculate image token grid dimensions patch_size = 2 # Z-Image uses patch_size=2 latent_height = self.height // LATENT_SCALE_FACTOR latent_width = self.width // LATENT_SCALE_FACTOR img_token_height = latent_height // patch_size img_token_width = latent_width // patch_size img_seq_len = img_token_height * img_token_width # Load positive conditioning with regional masks pos_text_conditionings = self._load_text_conditioning( context=context, cond_field=self.positive_conditioning, img_height=img_token_height, img_width=img_token_width, dtype=inference_dtype, device=device, ) # Create regional prompting extension regional_extension = ZImageRegionalPromptingExtension.from_text_conditionings( text_conditionings=pos_text_conditionings, img_seq_len=img_seq_len, ) # Get the concatenated prompt embeddings for the transformer pos_prompt_embeds = regional_extension.regional_text_conditioning.prompt_embeds # Load negative conditioning if provided and guidance_scale != 1.0 # CFG formula: pred = pred_uncond + cfg_scale * (pred_cond - pred_uncond) # At cfg_scale=1.0: pred = pred_cond (no effect, skip uncond computation) # This matches FLUX's convention where 1.0 means "no CFG" neg_prompt_embeds: torch.Tensor | None = None do_classifier_free_guidance = ( not math.isclose(self.guidance_scale, 1.0) and self.negative_conditioning is not None ) if do_classifier_free_guidance: assert self.negative_conditioning is not None # Load all negative conditionings and concatenate embeddings # Note: We ignore masks for negative conditioning as regional negative prompting is not fully supported neg_text_conditionings = self._load_text_conditioning( context=context, cond_field=self.negative_conditioning, img_height=img_token_height, img_width=img_token_width, dtype=inference_dtype, device=device, ) # Concatenate all negative embeddings neg_prompt_embeds = torch.cat([tc.prompt_embeds for tc in neg_text_conditionings], dim=0) # Calculate shift based on image sequence length, or use override if self.shift is not None: shift = self.shift else: shift = self._calculate_shift(img_seq_len) # Generate sigma schedule with time shift sigmas = self._get_sigmas(shift, self.steps) # Apply denoising_start and denoising_end clipping if self.denoising_start > 0 or self.denoising_end < 1: # Calculate start and end indices based on denoising range total_sigmas = len(sigmas) start_idx = int(self.denoising_start * (total_sigmas - 1)) end_idx = int(self.denoising_end * (total_sigmas - 1)) + 1 sigmas = sigmas[start_idx:end_idx] total_steps = len(sigmas) - 1 # Load input latents if provided (image-to-image) init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None if init_latents is not None: init_latents = init_latents.to(device=device, dtype=inference_dtype) # Generate initial noise. # If noise will never be consumed, avoid validating/loading it. should_ignore_noise = init_latents is not None and not self.add_noise and self.denoise_mask is None noise: torch.Tensor | None if should_ignore_noise: noise = None else: noise = self._prepare_noise_tensor(context, inference_dtype, device) # Prepare input latent image if init_latents is not None: if self.add_noise: assert noise is not None # Noise the init latents using the first sigma from the clipped # InvokeAI schedule. # # Known limitation: if the selected scheduler later starts from a # different first effective sigma/timestep than sigmas[0], the # img2img preblend below may not match that scheduler exactly. # This is an existing pipeline limitation and affects both # internally generated noise and externally supplied noise. s_0 = sigmas[0] latents = s_0 * noise + (1.0 - s_0) * init_latents else: latents = init_latents else: if self.denoising_start > 1e-5: raise ValueError("denoising_start should be 0 when initial latents are not provided.") assert noise is not None latents = noise # Short-circuit if no denoising steps if total_steps <= 0: return latents # Prepare inpaint extension inpaint_mask = self._prep_inpaint_mask(context, latents) inpaint_extension: RectifiedFlowInpaintExtension | None = None if inpaint_mask is not None: if init_latents is None: raise ValueError("Initial latents are required when using an inpaint mask (image-to-image inpainting)") assert noise is not None inpaint_extension = RectifiedFlowInpaintExtension( init_latents=init_latents, inpaint_mask=inpaint_mask, noise=noise, ) step_callback = self._build_step_callback(context) # Initialize the diffusers scheduler if not using built-in Euler scheduler: SchedulerMixin | None = None use_scheduler = self.scheduler != "euler" if use_scheduler: scheduler_class = ZIMAGE_SCHEDULER_MAP[self.scheduler] scheduler = scheduler_class( num_train_timesteps=1000, shift=1.0, ) # Set timesteps - LCM uses its own sigma schedule (num_inference_steps), # while other schedulers can use custom sigmas if supported is_lcm = self.scheduler == "lcm" set_timesteps_sig = inspect.signature(scheduler.set_timesteps) if not is_lcm and "sigmas" in set_timesteps_sig.parameters: scheduler.set_timesteps(sigmas=sigmas, device=device) else: # LCM or a scheduler without custom-sigma support computes its own # schedule from num_inference_steps. That can diverge from sigmas[0] # used in the img2img preblend above. scheduler.set_timesteps(num_inference_steps=total_steps, device=device) # For Heun scheduler, the number of actual steps may differ num_scheduler_steps = len(scheduler.timesteps) else: num_scheduler_steps = total_steps with ExitStack() as exit_stack: # Get transformer config to determine if it's quantized transformer_config = context.models.get_config(self.transformer.transformer) # Determine if the model is quantized. # If the model is quantized, then we need to apply the LoRA weights as sidecar layers. This results in # slower inference than direct patching, but is agnostic to the quantization format. if transformer_config.format in [ModelFormat.Diffusers, ModelFormat.Checkpoint]: model_is_quantized = False elif transformer_config.format in [ModelFormat.GGUFQuantized]: model_is_quantized = True else: raise ValueError(f"Unsupported Z-Image model format: {transformer_config.format}") # Load transformer - always use base transformer, control is handled via extension (cached_weights, transformer) = exit_stack.enter_context(transformer_info.model_on_device()) # Prepare control extension if control is provided control_extension: ZImageControlNetExtension | None = None if self.control is not None: # Load control adapter using context manager (proper GPU memory management) control_model_info = context.models.load(self.control.control_model) (_, control_adapter) = exit_stack.enter_context(control_model_info.model_on_device()) assert isinstance(control_adapter, ZImageControlAdapter) # Get control_in_dim from adapter config (16 for V1, 33 for V2.0) adapter_config = control_adapter.config control_in_dim = adapter_config.get("control_in_dim", 16) num_control_blocks = adapter_config.get("num_control_blocks", 6) # Log control configuration for debugging version = "V2.0" if control_in_dim > 16 else "V1" context.util.signal_progress( f"Using Z-Image ControlNet {version} (Extension): control_in_dim={control_in_dim}, " f"num_blocks={num_control_blocks}, scale={self.control.control_context_scale}" ) # Load and prepare control image - must be VAE-encoded! if self.vae is None: raise ValueError("VAE is required when using Z-Image Control. Connect a VAE to the 'vae' input.") control_image = context.images.get_pil(self.control.image_name) # Resize control image to match output dimensions control_image = control_image.convert("RGB") control_image = control_image.resize((self.width, self.height), Image.Resampling.LANCZOS) # Convert to tensor format for VAE encoding from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor control_image_tensor = image_resized_to_grid_as_tensor(control_image) if control_image_tensor.dim() == 3: control_image_tensor = einops.rearrange(control_image_tensor, "c h w -> 1 c h w") # Encode control image through VAE to get latents vae_info = context.models.load(self.vae.vae) control_latents = ZImageImageToLatentsInvocation.vae_encode( vae_info=vae_info, image_tensor=control_image_tensor, ) # Move to inference device/dtype control_latents = control_latents.to(device=device, dtype=inference_dtype) # Add frame dimension: [B, C, H, W] -> [C, 1, H, W] (single image) control_latents = control_latents.squeeze(0).unsqueeze(1) # Prepare control_cond based on control_in_dim # V1: 16 channels (just control latents) # V2.0: 33 channels = 16 control + 16 reference + 1 mask # - Channels 0-15: control image latents (from VAE encoding) # - Channels 16-31: reference/inpaint image latents (zeros for pure control) # - Channel 32: inpaint mask (1.0 = don't inpaint, 0.0 = inpaint region) # For pure control (no inpainting), we set mask=1 to tell model "use control, don't inpaint" c, f, h, w = control_latents.shape if c < control_in_dim: padding_channels = control_in_dim - c if padding_channels == 17: # V2.0: 16 reference channels (zeros) + 1 mask channel (ones) ref_padding = torch.zeros( (16, f, h, w), device=device, dtype=inference_dtype, ) # Mask channel = 1.0 means "don't inpaint this region, use control signal" mask_channel = torch.ones( (1, f, h, w), device=device, dtype=inference_dtype, ) control_latents = torch.cat([control_latents, ref_padding, mask_channel], dim=0) else: # Generic padding with zeros for other cases zero_padding = torch.zeros( (padding_channels, f, h, w), device=device, dtype=inference_dtype, ) control_latents = torch.cat([control_latents, zero_padding], dim=0) # Create control extension (adapter is already on device from model_on_device) control_extension = ZImageControlNetExtension( control_adapter=control_adapter, control_cond=control_latents, weight=self.control.control_context_scale, begin_step_percent=self.control.begin_step_percent, end_step_percent=self.control.end_step_percent, ) # Apply LoRA models to the transformer. # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching. exit_stack.enter_context( LayerPatcher.apply_smart_model_patches( model=transformer, patches=self._lora_iterator(context), prefix=Z_IMAGE_LORA_TRANSFORMER_PREFIX, dtype=inference_dtype, cached_weights=cached_weights, force_sidecar_patching=model_is_quantized, ) ) # Apply regional prompting patch if we have regional masks exit_stack.enter_context( patch_transformer_for_regional_prompting( transformer=transformer, regional_attn_mask=regional_extension.regional_attn_mask, img_seq_len=img_seq_len, positive_cap_feats=pos_prompt_embeds, ) ) # Denoising loop - supports both built-in Euler and diffusers schedulers # Track user-facing step for progress (accounts for Heun's double steps) user_step = 0 if use_scheduler and scheduler is not None: # Use diffusers scheduler for stepping # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps) # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps pbar = tqdm(total=total_steps, desc="Denoising") for step_index in range(num_scheduler_steps): sched_timestep = scheduler.timesteps[step_index] # Convert scheduler timestep (0-1000) to normalized sigma (0-1) sigma_curr = sched_timestep.item() / scheduler.config.num_train_timesteps # For Heun scheduler, track if we're in first or second order step is_heun = hasattr(scheduler, "state_in_first_order") in_first_order = scheduler.state_in_first_order if is_heun else True # Timestep tensor for Z-Image model # The model expects t=0 at start (noise) and t=1 at end (clean) model_t = 1.0 - sigma_curr timestep = torch.tensor([model_t], device=device, dtype=inference_dtype).expand(latents.shape[0]) # Run transformer for positive prediction latent_model_input = latents.to(transformer.dtype) latent_model_input = latent_model_input.unsqueeze(2) # Add frame dimension latent_model_input_list = list(latent_model_input.unbind(dim=0)) # Determine if control should be applied at this step apply_control = control_extension is not None and control_extension.should_apply( user_step, total_steps ) # Run forward pass if apply_control: model_out_list, _ = z_image_forward_with_control( transformer=transformer, x=latent_model_input_list, t=timestep, cap_feats=[pos_prompt_embeds], control_extension=control_extension, ) else: model_output = transformer( x=latent_model_input_list, t=timestep, cap_feats=[pos_prompt_embeds], ) model_out_list = model_output[0] noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0) noise_pred_cond = noise_pred_cond.squeeze(2) noise_pred_cond = -noise_pred_cond # Z-Image uses v-prediction with negation # Apply CFG if enabled if do_classifier_free_guidance and neg_prompt_embeds is not None: if apply_control: model_out_list_uncond, _ = z_image_forward_with_control( transformer=transformer, x=latent_model_input_list, t=timestep, cap_feats=[neg_prompt_embeds], control_extension=control_extension, ) else: model_output_uncond = transformer( x=latent_model_input_list, t=timestep, cap_feats=[neg_prompt_embeds], ) model_out_list_uncond = model_output_uncond[0] noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0) noise_pred_uncond = noise_pred_uncond.squeeze(2) noise_pred_uncond = -noise_pred_uncond noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) else: noise_pred = noise_pred_cond # Use scheduler.step() for the update step_output = scheduler.step(model_output=noise_pred, timestep=sched_timestep, sample=latents) latents = step_output.prev_sample # Get sigma_prev for inpainting (next sigma value) if step_index + 1 < len(scheduler.sigmas): sigma_prev = scheduler.sigmas[step_index + 1].item() else: sigma_prev = 0.0 if inpaint_extension is not None: latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev) # For Heun, only increment user step after second-order step completes if is_heun: if not in_first_order: user_step += 1 # Only call step_callback if we haven't exceeded total_steps if user_step <= total_steps: pbar.update(1) step_callback( PipelineIntermediateState( step=user_step, order=2, total_steps=total_steps, timestep=int(sigma_curr * 1000), latents=latents, ), ) else: # For first-order schedulers (Euler, LCM) user_step += 1 if user_step <= total_steps: pbar.update(1) step_callback( PipelineIntermediateState( step=user_step, order=1, total_steps=total_steps, timestep=int(sigma_curr * 1000), latents=latents, ), ) pbar.close() else: # Original Euler implementation (default, optimized for Z-Image) for step_idx in tqdm(range(total_steps)): sigma_curr = sigmas[step_idx] sigma_prev = sigmas[step_idx + 1] # Timestep tensor for Z-Image model # The model expects t=0 at start (noise) and t=1 at end (clean) # Sigma goes from 1 (noise) to 0 (clean), so model_t = 1 - sigma model_t = 1.0 - sigma_curr timestep = torch.tensor([model_t], device=device, dtype=inference_dtype).expand(latents.shape[0]) # Run transformer for positive prediction # Z-Image transformer expects: x as list of [C, 1, H, W] tensors, t, cap_feats as list # Prepare latent input: [B, C, H, W] -> [B, C, 1, H, W] -> list of [C, 1, H, W] latent_model_input = latents.to(transformer.dtype) latent_model_input = latent_model_input.unsqueeze(2) # Add frame dimension latent_model_input_list = list(latent_model_input.unbind(dim=0)) # Determine if control should be applied at this step apply_control = control_extension is not None and control_extension.should_apply( step_idx, total_steps ) # Run forward pass - use custom forward with control if extension is active if apply_control: model_out_list, _ = z_image_forward_with_control( transformer=transformer, x=latent_model_input_list, t=timestep, cap_feats=[pos_prompt_embeds], control_extension=control_extension, ) else: model_output = transformer( x=latent_model_input_list, t=timestep, cap_feats=[pos_prompt_embeds], ) model_out_list = model_output[0] # Extract list of tensors from tuple noise_pred_cond = torch.stack([t.float() for t in model_out_list], dim=0) noise_pred_cond = noise_pred_cond.squeeze(2) # Remove frame dimension noise_pred_cond = -noise_pred_cond # Z-Image uses v-prediction with negation # Apply CFG if enabled if do_classifier_free_guidance and neg_prompt_embeds is not None: if apply_control: model_out_list_uncond, _ = z_image_forward_with_control( transformer=transformer, x=latent_model_input_list, t=timestep, cap_feats=[neg_prompt_embeds], control_extension=control_extension, ) else: model_output_uncond = transformer( x=latent_model_input_list, t=timestep, cap_feats=[neg_prompt_embeds], ) model_out_list_uncond = model_output_uncond[0] # Extract list of tensors from tuple noise_pred_uncond = torch.stack([t.float() for t in model_out_list_uncond], dim=0) noise_pred_uncond = noise_pred_uncond.squeeze(2) noise_pred_uncond = -noise_pred_uncond noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) else: noise_pred = noise_pred_cond # Euler step latents_dtype = latents.dtype latents = latents.to(dtype=torch.float32) latents = latents + (sigma_prev - sigma_curr) * noise_pred latents = latents.to(dtype=latents_dtype) if inpaint_extension is not None: latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev) step_callback( PipelineIntermediateState( step=step_idx + 1, order=1, total_steps=total_steps, timestep=int(sigma_curr * 1000), latents=latents, ), ) return latents def _prepare_noise_tensor( self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device ) -> torch.Tensor: if self.noise is not None: noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype) validate_noise_tensor_shape(noise, "Z-Image", self.width, self.height) return noise return self._get_noise( batch_size=1, num_channels_latents=16, height=self.height, width=self.width, dtype=inference_dtype, device=device, seed=self.seed, ) def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]: def step_callback(state: PipelineIntermediateState) -> None: context.util.sd_step_callback(state, BaseModelType.ZImage) return step_callback def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]: """Iterate over LoRA models to apply to the transformer.""" for lora in self.transformer.loras: lora_info = context.models.load(lora.lora) if not isinstance(lora_info.model, ModelPatchRaw): raise TypeError( f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}. " "The LoRA model may be corrupted or incompatible." ) yield (lora_info.model, lora.weight) del lora_info