Fix redundant Z-Image terminal timestep

Author: rootonchair

src/diffusers/modular_pipelines/z_image/before_denoise.py

@@ -185,6 +185,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImageTextInputStep(ModularPipelineBlocks):
     model_name = "z-image"
 
@@ -535,13 +540,15 @@ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> P
             base_shift=components.scheduler.config.get("base_shift", 0.5),
             max_shift=components.scheduler.config.get("max_shift", 1.15),
         )
-        components.scheduler.sigma_min = 0.0
+        sigmas = block_state.sigmas
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(block_state.num_inference_steps)
 
         block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
             components.scheduler,
             block_state.num_inference_steps,
             device,
-            sigmas=block_state.sigmas,
+            sigmas=sigmas,
             mu=mu,
         )
 

src/diffusers/pipelines/z_image/pipeline_z_image.py

@@ -134,6 +134,10 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImagePipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
     model_cpu_offload_seq = "text_encoder->transformer->vae"
     _optional_components = []
@@ -474,7 +478,8 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.5),
             self.scheduler.config.get("max_shift", 1.15),
         )
-        self.scheduler.sigma_min = 0.0
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(num_inference_steps)
         scheduler_kwargs = {"mu": mu}
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py

@@ -185,6 +185,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImageControlNetPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
     model_cpu_offload_seq = "text_encoder->transformer->vae"
     _optional_components = []
@@ -593,7 +598,8 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.5),
             self.scheduler.config.get("max_shift", 1.15),
         )
-        self.scheduler.sigma_min = 0.0
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(num_inference_steps)
         scheduler_kwargs = {"mu": mu}
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py

@@ -185,6 +185,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImageControlNetInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
     model_cpu_offload_seq = "text_encoder->transformer->vae"
     _optional_components = []
@@ -615,7 +620,8 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.5),
             self.scheduler.config.get("max_shift", 1.15),
         )
-        self.scheduler.sigma_min = 0.0
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(num_inference_steps)
         scheduler_kwargs = {"mu": mu}
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,

src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py

@@ -146,6 +146,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImageImg2ImgPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
     r"""
     The ZImage pipeline for image-to-image generation.
@@ -563,7 +568,8 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.5),
             self.scheduler.config.get("max_shift", 1.15),
         )
-        self.scheduler.sigma_min = 0.0
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(num_inference_steps)
         scheduler_kwargs = {"mu": mu}
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,

src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py

@@ -162,6 +162,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImageInpaintPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
     r"""
     The ZImage pipeline for inpainting.
@@ -750,7 +755,8 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.5),
             self.scheduler.config.get("max_shift", 1.15),
         )
-        self.scheduler.sigma_min = 0.0
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(num_inference_steps)
         scheduler_kwargs = {"mu": mu}
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,

src/diffusers/pipelines/z_image/pipeline_z_image_omni.py

@@ -135,6 +135,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+# Copied from diffusers.pipelines.z_image.pipeline_z_image.get_default_z_image_sigmas
+def get_default_z_image_sigmas(num_inference_steps: int) -> list[float]:
+    return torch.linspace(1.0, 0.0, num_inference_steps + 1)[:-1].tolist()
+
+
 class ZImageOmniPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
     model_cpu_offload_seq = "text_encoder->transformer->vae"
     _optional_components = []
@@ -604,7 +609,8 @@ def __call__(
             self.scheduler.config.get("base_shift", 0.5),
             self.scheduler.config.get("max_shift", 1.15),
         )
-        self.scheduler.sigma_min = 0.0
+        if sigmas is None:
+            sigmas = get_default_z_image_sigmas(num_inference_steps)
         scheduler_kwargs = {"mu": mu}
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,