diff --git a/src/diffusers/modular_pipelines/z_image/before_denoise.py b/src/diffusers/modular_pipelines/z_image/before_denoise.py index 8558f2c67f65..820de47d6335 100644 --- a/src/diffusers/modular_pipelines/z_image/before_denoise.py +++ b/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, ) diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image.py b/src/diffusers/pipelines/z_image/pipeline_z_image.py index 46403a0719cd..67df30425acb 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image.py +++ b/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, diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py index 1e49737bb5b0..38ef711dd09f 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py +++ b/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, diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py index 09f9b2395458..c2613bede08b 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_controlnet_inpaint.py +++ b/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, diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py b/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py index ee57f51dd957..4ef0e60677bc 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_img2img.py +++ b/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, diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py b/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py index e740a48e65ec..4374dd98d217 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py +++ b/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, diff --git a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py b/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py index 9199e176a1f6..e5c11cf01e11 100644 --- a/src/diffusers/pipelines/z_image/pipeline_z_image_omni.py +++ b/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,