diff --git a/inspect_zimage.py b/inspect_zimage.py new file mode 100644 index 0000000..761dc86 --- /dev/null +++ b/inspect_zimage.py @@ -0,0 +1,39 @@ +import inspect +import torch +import sys + +try: + # Try to import the class directly if possible + # Based on PR link, it might be in diffusers.models.transformers.z_image_transformer_2d + # But let's try to find it in diffusers.models + from diffusers import DiffusionPipeline + + # We can try to load the pipeline (might be slow) or just inspect the module structure + # Let's try to find the class in diffusers.models + import diffusers.models + + found = False + for name, obj in inspect.getmembers(diffusers.models): + if "ZImage" in name and "Transformer" in name: + print(f"Found class: {name}") + sig = inspect.signature(obj.forward) + print(f"Signature of {name}.forward:") + print(sig) + found = True + + if not found: + print("Could not find ZImageTransformer class in diffusers.models directly.") + print("Trying to load pipeline to get the object...") + # Load a small dummy or just the class if we can + # Let's try to import specifically from the likely module path + try: + from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel + print("Found ZImageTransformer2DModel in diffusers.models.transformers.transformer_z_image") + sig = inspect.signature(ZImageTransformer2DModel.forward) + print(f"Signature of ZImageTransformer2DModel.forward:") + print(sig) + except ImportError: + print("Could not import from diffusers.models.transformers.transformer_z_image") + +except Exception as e: + print(f"Error: {e}") diff --git a/pipeline_z_image.py b/pipeline_z_image.py new file mode 100644 index 0000000..a4fcacb --- /dev/null +++ b/pipeline_z_image.py @@ -0,0 +1,596 @@ +# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from transformers import AutoTokenizer, PreTrainedModel + +from ...image_processor import VaeImageProcessor +from ...loaders import FromSingleFileMixin +from ...models.autoencoders import AutoencoderKL +from ...models.transformers import ZImageTransformer2DModel +from ...pipelines.pipeline_utils import DiffusionPipeline +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from .pipeline_output import ZImagePipelineOutput + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import ZImagePipeline + + >>> pipe = ZImagePipeline.from_pretrained("Z-a-o/Z-Image-Turbo", torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> # Optionally, set the attention backend to flash-attn 2 or 3, default is SDPA in PyTorch. + >>> # (1) Use flash attention 2 + >>> # pipe.transformer.set_attention_backend("flash") + >>> # (2) Use flash attention 3 + >>> # pipe.transformer.set_attention_backend("_flash_3") + + >>> prompt = "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。画面巧妙地将文字概念视觉化：一辆复古蒸汽小火车化身为巨大的拉链头，正拉开厚厚的冬日积雪，展露出一个生机盎然的春天。" + >>> image = pipe( + ... prompt, + ... height=1024, + ... width=1024, + ... num_inference_steps=9, + ... guidance_scale=0.0, + ... generator=torch.Generator("cuda").manual_seed(42), + ... ).images[0] + >>> image.save("zimage.png") + ``` +""" + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class ZImagePipeline(DiffusionPipeline, FromSingleFileMixin): + model_cpu_offload_seq = "text_encoder->transformer->vae" + _optional_components = [] + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: PreTrainedModel, + tokenizer: AutoTokenizer, + transformer: ZImageTransformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer, + ) + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 + ) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + + def encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[Union[str, List[str]]] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + max_sequence_length: int = 512, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt_embeds = self._encode_prompt( + prompt=prompt, + device=device, + prompt_embeds=prompt_embeds, + max_sequence_length=max_sequence_length, + ) + + if do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = ["" for _ in prompt] + else: + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + assert len(prompt) == len(negative_prompt) + negative_prompt_embeds = self._encode_prompt( + prompt=negative_prompt, + device=device, + prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, + ) + else: + negative_prompt_embeds = [] + return prompt_embeds, negative_prompt_embeds + + def _encode_prompt( + self, + prompt: Union[str, List[str]], + device: Optional[torch.device] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + max_sequence_length: int = 512, + ) -> List[torch.FloatTensor]: + device = device or self._execution_device + + if prompt_embeds is not None: + return prompt_embeds + + if isinstance(prompt, str): + prompt = [prompt] + + for i, prompt_item in enumerate(prompt): + messages = [ + {"role": "user", "content": prompt_item}, + ] + prompt_item = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + prompt[i] = prompt_item + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids.to(device) + prompt_masks = text_inputs.attention_mask.to(device).bool() + + prompt_embeds = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_masks, + output_hidden_states=True, + ).hidden_states[-2] + + embeddings_list = [] + + for i in range(len(prompt_embeds)): + embeddings_list.append(prompt_embeds[i][prompt_masks[i]]) + + return embeddings_list + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def joint_attention_kwargs(self): + return self._joint_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + sigmas: Optional[List[float]] = None, + guidance_scale: float = 5.0, + cfg_normalization: bool = False, + cfg_truncation: float = 1.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds: Optional[List[torch.FloatTensor]] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 512, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + cfg_normalization (`bool`, *optional*, defaults to False): + Whether to apply configuration normalization. + cfg_truncation (`float`, *optional*, defaults to 1.0): + The truncation value for configuration. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`List[torch.FloatTensor]`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain + tuple. + joint_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if + `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the + generated images. + """ + height = height or 1024 + width = width or 1024 + + vae_scale = self.vae_scale_factor * 2 + if height % vae_scale != 0: + raise ValueError( + f"Height must be divisible by {vae_scale} (got {height}). " + f"Please adjust the height to a multiple of {vae_scale}." + ) + if width % vae_scale != 0: + raise ValueError( + f"Width must be divisible by {vae_scale} (got {width}). " + f"Please adjust the width to a multiple of {vae_scale}." + ) + + device = self._execution_device + + self._guidance_scale = guidance_scale + self._joint_attention_kwargs = joint_attention_kwargs + self._interrupt = False + self._cfg_normalization = cfg_normalization + self._cfg_truncation = cfg_truncation + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = len(prompt_embeds) + + # If prompt_embeds is provided and prompt is None, skip encoding + if prompt_embeds is not None and prompt is None: + if self.do_classifier_free_guidance and negative_prompt_embeds is None: + raise ValueError( + "When `prompt_embeds` is provided without `prompt`, " + "`negative_prompt_embeds` must also be provided for classifier-free guidance." + ) + else: + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.in_channels + + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + torch.float32, + device, + generator, + latents, + ) + + # Repeat prompt_embeds for num_images_per_prompt + if num_images_per_prompt > 1: + prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] + if self.do_classifier_free_guidance and negative_prompt_embeds: + negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] + + actual_batch_size = batch_size * num_images_per_prompt + image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) + + # 5. Prepare timesteps + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + self.scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + # 6. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latents.shape[0]) + timestep = (1000 - timestep) / 1000 + # Normalized time for time-aware config (0 at start, 1 at end) + t_norm = timestep[0].item() + + # Handle cfg truncation + current_guidance_scale = self.guidance_scale + if ( + self.do_classifier_free_guidance + and self._cfg_truncation is not None + and float(self._cfg_truncation) <= 1 + ): + if t_norm > self._cfg_truncation: + current_guidance_scale = 0.0 + + # Run CFG only if configured AND scale is non-zero + apply_cfg = self.do_classifier_free_guidance and current_guidance_scale > 0 + + if apply_cfg: + latents_typed = latents.to(self.transformer.dtype) + latent_model_input = latents_typed.repeat(2, 1, 1, 1) + prompt_embeds_model_input = prompt_embeds + negative_prompt_embeds + timestep_model_input = timestep.repeat(2) + else: + latent_model_input = latents.to(self.transformer.dtype) + prompt_embeds_model_input = prompt_embeds + timestep_model_input = timestep + + latent_model_input = latent_model_input.unsqueeze(2) + latent_model_input_list = list(latent_model_input.unbind(dim=0)) + + model_out_list = self.transformer( + latent_model_input_list, + timestep_model_input, + prompt_embeds_model_input, + )[0] + + if apply_cfg: + # Perform CFG + pos_out = model_out_list[:actual_batch_size] + neg_out = model_out_list[actual_batch_size:] + + noise_pred = [] + for j in range(actual_batch_size): + pos = pos_out[j].float() + neg = neg_out[j].float() + + pred = pos + current_guidance_scale * (pos - neg) + + # Renormalization + if self._cfg_normalization and float(self._cfg_normalization) > 0.0: + ori_pos_norm = torch.linalg.vector_norm(pos) + new_pos_norm = torch.linalg.vector_norm(pred) + max_new_norm = ori_pos_norm * float(self._cfg_normalization) + if new_pos_norm > max_new_norm: + pred = pred * (max_new_norm / new_pos_norm) + + noise_pred.append(pred) + + noise_pred = torch.stack(noise_pred, dim=0) + else: + noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) + + noise_pred = noise_pred.squeeze(2) + noise_pred = -noise_pred + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0] + assert latents.dtype == torch.float32 + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if output_type == "latent": + image = latents + + else: + latents = latents.to(self.vae.dtype) + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ZImagePipelineOutput(images=image) diff --git a/sdnq_nodes.py b/sdnq_nodes.py index 119c6ed..ae1208e 100644 --- a/sdnq_nodes.py +++ b/sdnq_nodes.py @@ -36,143 +36,6 @@ except ImportError: except ImportError: raise ImportError("SDNQ installation failed. Please install manually.") -# Custom Pipeline for Img2Img support -class ZImageImg2ImgPipeline(DiffusionPipeline): - # Define offload sequence for CPU offloading - model_cpu_offload_seq = "text_encoder->transformer->vae" - - def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler): - super().__init__() - self.register_modules( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - scheduler=scheduler, - ) - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self.vae, "config") else 8 - - def __call__( - self, - prompt=None, - height=1024, - width=1024, - num_inference_steps=50, - guidance_scale=5.0, - negative_prompt=None, - num_images_per_prompt=1, - generator=None, - latents=None, - prompt_embeds=None, - negative_prompt_embeds=None, - output_type="pil", - return_dict=True, - max_sequence_length=512, - # Img2Img specific parameters - image=None, - strength=1.0, - noise_scale=1.0, - ): - # 0. Default height and width to 1024 - height = height or 1024 - width = width or 1024 - - # 1. Check inputs - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self._execution_device - - # 2. Encode prompt - # Note: We assume the original pipeline has an encode_prompt method or similar logic - # Since we are inheriting/wrapping, we might need to access the original method if available - # But ZImagePipeline structure is specific. Let's try to reuse the components directly. - - # Simplified prompt encoding for Z-Image (based on observation of original pipeline) - if prompt_embeds is None: - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - prompt_embeds = self.text_encoder(text_input_ids)[0] - - # 3. Prepare latents - # If image is provided, encode it - if image is not None: - # Image is already a tensor [B, C, H, W] in range [-1, 1] on correct device/dtype - if latents is None: - latents = self.vae.encode(image).latent_dist.sample(generator) - latents = latents * self.vae.config.scaling_factor - - # If no latents (text-to-image), generate random noise - if latents is None: - shape = (batch_size, self.transformer.config.in_channels, height // self.vae_scale_factor, width // self.vae_scale_factor) - latents = torch.randn(shape, generator=generator, device=device, dtype=prompt_embeds.dtype) - - # 4. Prepare noise and mix for Img2Img - if image is not None and strength < 1.0: - # Generate noise - noise = torch.randn(latents.shape, generator=generator, device=device, dtype=latents.dtype) - - # Apply noise scale if requested - if noise_scale != 1.0: - noise = noise * noise_scale - - # Rectified Flow Interpolation: latents = (1 - t) * image_latents + t * noise - # t = strength (0.0 = image, 1.0 = noise) - t = strength - latents = (1 - t) * latents + t * noise - - # We don't need to skip steps in Rectified Flow the same way as DDIM - # The ODE solver integrates from t=0 (noise) to t=1 (data) or vice versa - # But standard pipelines usually go from Noise -> Data - # Here we are starting from an intermediate state - - # However, standard schedulers might expect to start from pure noise - # For simplicity in this custom implementation, we'll just pass the mixed latents - # as the starting point. - - # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - - # 6. Denoising loop - for i, t in enumerate(timesteps): - # Expand latents if needed - latent_model_input = latents - - # Predict noise/velocity - noise_pred = self.transformer( - latent_model_input, - timestep=t, - encoder_hidden_states=prompt_embeds, - return_dict=False, - )[0] - - # Compute previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - # 7. Decode latents - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - else: - image = latents - - # Convert to PIL if requested - if output_type == "pil": - image = self.image_processor.postprocess(image, output_type=output_type) - - from diffusers.pipelines.pipeline_utils import ImagePipelineOutput - return ImagePipelineOutput(images=image) - class LoadZImageSDNQ: @classmethod def INPUT_TYPES(cls): @@ -224,24 +87,13 @@ class LoadZImageSDNQ: else: dev = device - # Load original pipeline first - original_pipeline = DiffusionPipeline.from_pretrained( + # Load the original ZImagePipeline directly + pipeline = DiffusionPipeline.from_pretrained( model_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=low_cpu_mem_usage ) - # Wrap in our custom pipeline - # Note: We need to extract components. Z-Image pipeline structure: - # vae, text_encoder, tokenizer, transformer, scheduler - pipeline = ZImageImg2ImgPipeline( - vae=original_pipeline.vae, - text_encoder=original_pipeline.text_encoder, - tokenizer=original_pipeline.tokenizer, - transformer=original_pipeline.transformer, - scheduler=original_pipeline.scheduler - ) - # Apply attention backend if attention_backend != "default": print(f"Setting attention backend to: {attention_backend}") @@ -278,7 +130,7 @@ class LoadZImageSDNQ: except Exception as e: print(f"Warning: Compilation failed: {e}") - print("SDNQ Pipeline (Custom Img2Img) loaded successfully.") + print("SDNQ Pipeline loaded successfully.") return (pipeline,) # Resolution presets from official Gradio app @@ -450,20 +302,18 @@ class ZImageSDNQGenerate: generator = torch.manual_seed(seed) - # Prepare input image if provided - image_tensor = None + # Prepare latents for img2img if input_image is provided + latents = None if input_image is not None: - print(f"Preparing input image for img2img (strength={strength})...") - # Convert ComfyUI image format [B, H, W, C] to torch [B, C, H, W] - # Match VAE dtype and device - vae_dtype = pipeline.vae.dtype - - # Use _execution_device if available (handles CPU offloading correctly) + print(f"Preparing img2img latents (strength={strength}, noise_scale={noise_scale})...") + # Get device if hasattr(pipeline, "_execution_device"): device = pipeline._execution_device else: device = pipeline.device + # Convert ComfyUI image format [B, H, W, C] to torch [B, C, H, W] + vae_dtype = pipeline.vae.dtype image_tensor = input_image.permute(0, 3, 1, 2).to(device=device, dtype=vae_dtype) # Resize if needed @@ -473,6 +323,24 @@ class ZImageSDNQGenerate: # Normalize to [-1, 1] (ComfyUI images are [0, 1]) image_tensor = 2.0 * image_tensor - 1.0 + + # Encode to latents using VAE + with torch.no_grad(): + latents = pipeline.vae.encode(image_tensor).latent_dist.sample() + latents = latents * pipeline.vae.config.scaling_factor + + # Apply noise mixing for img2img effect + # Generate noise on CPU (compatible with CPU generator), then move to device + noise = torch.randn(latents.shape, generator=generator, device="cpu", dtype=latents.dtype).to(device) + + # Apply noise scale + if noise_scale != 1.0: + noise = noise * noise_scale + + # Rectified Flow interpolation: latents = (1 - strength) * image_latents + strength * noise + # strength=0.0 → pure image, strength=1.0 → pure noise + latents = (1 - strength) * latents + strength * noise + print(f"Mixed latents: {100*(1-strength):.1f}% image + {100*strength:.1f}% noise") # Configure shift on the scheduler before generation original_shift = None @@ -493,16 +361,14 @@ class ZImageSDNQGenerate: "guidance_scale": guidance_scale, "max_sequence_length": max_sequence_length, "generator": generator, - "noise_scale": noise_scale, } - # Add img2img parameters if image is provided - if image_tensor is not None: - gen_kwargs["image"] = image_tensor - gen_kwargs["strength"] = strength + # Add pre-processed latents if we have them (img2img) + if latents is not None: + gen_kwargs["latents"] = latents - # Call the custom pipeline - # The custom pipeline handles encoding, noise mixing, and denoising loop + # Call the original ZImagePipeline + # It handles all the transformer calls and denoising internally image = pipeline(**gen_kwargs).images[0] # Restore original shift value if we changed it diff --git a/test_device.py b/test_device.py new file mode 100644 index 0000000..5c1ac24 --- /dev/null +++ b/test_device.py @@ -0,0 +1,31 @@ +import torch +from diffusers import DiffusionPipeline + +# Load a small test to see how pipeline handles latents +print("Testing latents device handling...") + +# Simulate what we're doing +class MockPipeline: + def __init__(self): + self.device = torch.device("cuda") + self._execution_device = torch.device("cuda") + +# Our current approach +pipeline_mock = MockPipeline() + +# Get device (our current code) +if hasattr(pipeline_mock, "_execution_device"): + device = pipeline_mock._execution_device +else: + device = pipeline_mock.device + +print(f"Device selected: {device}") + +# Create latents (simulating our VAE encode result) +latents = torch.randn(1, 4, 128, 128, device=device, dtype=torch.bfloat16) +print(f"Latents device: {latents.device}") +print(f"Latents dtype: {latents.dtype}") + +# When we pass to pipeline, it should already be on the correct device +print(f"\n✓ Latents are already on the execution device: {device}") +print("✓ Pipeline will accept them directly without device transfers")