From e7ec13f209258703491160d128f286d3f9b835fa Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Wed, 28 Feb 2024 12:15:39 -0500 Subject: [PATCH] Remove scheduler_args from ConditioningData structure. --- invokeai/app/invocations/latent.py | 25 +++++++++++-------- .../stable_diffusion/diffusers_pipeline.py | 7 +++++- .../diffusion/conditioning_data.py | 20 ++------------- 3 files changed, 22 insertions(+), 30 deletions(-) diff --git a/invokeai/app/invocations/latent.py b/invokeai/app/invocations/latent.py index 9b140910c7..1511c1ba15 100644 --- a/invokeai/app/invocations/latent.py +++ b/invokeai/app/invocations/latent.py @@ -1,5 +1,6 @@ # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) +import inspect import math from contextlib import ExitStack from functools import singledispatchmethod @@ -375,14 +376,6 @@ class DenoiseLatentsInvocation(BaseInvocation): guidance_rescale_multiplier=self.cfg_rescale_multiplier, ) - conditioning_data = conditioning_data.add_scheduler_args_if_applicable( - scheduler, - # for ddim scheduler - eta=0.0, # ddim_eta - # for ancestral and sde schedulers - # flip all bits to have noise different from initial - generator=torch.Generator(device=unet.device).manual_seed(seed ^ 0xFFFFFFFF), - ) return conditioning_data def create_pipeline( @@ -642,7 +635,7 @@ class DenoiseLatentsInvocation(BaseInvocation): # original idea by https://github.com/AmericanPresidentJimmyCarter # TODO: research more for second order schedulers timesteps - def init_scheduler(self, scheduler, device, steps, denoising_start, denoising_end): + def init_scheduler(self, scheduler, device, steps, denoising_start, denoising_end, seed: int): if scheduler.config.get("cpu_only", False): scheduler.set_timesteps(steps, device="cpu") timesteps = scheduler.timesteps.to(device=device) @@ -669,7 +662,15 @@ class DenoiseLatentsInvocation(BaseInvocation): timesteps = timesteps[t_start_idx : t_start_idx + t_end_idx] num_inference_steps = len(timesteps) // scheduler.order - return num_inference_steps, timesteps, init_timestep + scheduler_step_kwargs = {} + scheduler_step_signature = inspect.signature(scheduler.step) + if "generator" in scheduler_step_signature.parameters: + # At some point, someone decided that schedulers that accept a generator should use the original seed with + # all bits flipped. I don't know the original rationale for this, but now we must keep it like this for + # reproducibility. + scheduler_step_kwargs = {"generator": torch.Generator(device=device).manual_seed(seed ^ 0xFFFFFFFF)} + + return num_inference_steps, timesteps, init_timestep, scheduler_step_kwargs def prep_inpaint_mask(self, context, latents): if self.denoise_mask is None: @@ -783,12 +784,13 @@ class DenoiseLatentsInvocation(BaseInvocation): exit_stack=exit_stack, ) - num_inference_steps, timesteps, init_timestep = self.init_scheduler( + num_inference_steps, timesteps, init_timestep, scheduler_step_kwargs = self.init_scheduler( scheduler, device=unet.device, steps=self.steps, denoising_start=self.denoising_start, denoising_end=self.denoising_end, + seed=seed, ) result_latents = pipeline.latents_from_embeddings( @@ -800,6 +802,7 @@ class DenoiseLatentsInvocation(BaseInvocation): mask=mask, masked_latents=masked_latents, num_inference_steps=num_inference_steps, + scheduler_step_kwargs=scheduler_step_kwargs, conditioning_data=conditioning_data, control_data=controlnet_data, ip_adapter_data=ip_adapter_data, diff --git a/invokeai/backend/stable_diffusion/diffusers_pipeline.py b/invokeai/backend/stable_diffusion/diffusers_pipeline.py index 073d6b1747..3b54136469 100644 --- a/invokeai/backend/stable_diffusion/diffusers_pipeline.py +++ b/invokeai/backend/stable_diffusion/diffusers_pipeline.py @@ -309,6 +309,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): self, latents: torch.Tensor, num_inference_steps: int, + scheduler_step_kwargs: dict[str, Any], conditioning_data: ConditioningData, *, noise: Optional[torch.Tensor], @@ -368,6 +369,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): latents, timesteps, conditioning_data, + scheduler_step_kwargs=scheduler_step_kwargs, additional_guidance=additional_guidance, control_data=control_data, ip_adapter_data=ip_adapter_data, @@ -388,6 +390,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): latents: torch.Tensor, timesteps, conditioning_data: ConditioningData, + scheduler_step_kwargs: dict[str, Any], *, additional_guidance: List[Callable] = None, control_data: List[ControlNetData] = None, @@ -454,6 +457,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): conditioning_data, step_index=i, total_step_count=len(timesteps), + scheduler_step_kwargs=scheduler_step_kwargs, additional_guidance=additional_guidance, control_data=control_data, ip_adapter_data=ip_adapter_data, @@ -485,6 +489,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): conditioning_data: ConditioningData, step_index: int, total_step_count: int, + scheduler_step_kwargs: dict[str, Any], additional_guidance: List[Callable] = None, control_data: List[ControlNetData] = None, ip_adapter_data: Optional[list[IPAdapterData]] = None, @@ -584,7 +589,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): ) # compute the previous noisy sample x_t -> x_t-1 - step_output = self.scheduler.step(noise_pred, timestep, latents, **conditioning_data.scheduler_args) + step_output = self.scheduler.step(noise_pred, timestep, latents, **scheduler_step_kwargs) # TODO: issue to diffusers? # undo internal counter increment done by scheduler.step, so timestep can be resolved as before call diff --git a/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py b/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py index 2e816b982b..267a17e917 100644 --- a/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py @@ -1,7 +1,5 @@ -import dataclasses -import inspect -from dataclasses import dataclass, field -from typing import Any, List, Optional, Union +from dataclasses import dataclass +from typing import List, Optional, Union import torch @@ -73,19 +71,5 @@ class ConditioningData: ref [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf) """ guidance_rescale_multiplier: float = 0 - scheduler_args: dict[str, Any] = field(default_factory=dict) ip_adapter_conditioning: Optional[list[IPAdapterConditioningInfo]] = None - - def add_scheduler_args_if_applicable(self, scheduler, **kwargs): - scheduler_args = dict(self.scheduler_args) - step_method = inspect.signature(scheduler.step) - for name, value in kwargs.items(): - try: - step_method.bind_partial(**{name: value}) - except TypeError: - # FIXME: don't silently discard arguments - pass # debug("%s does not accept argument named %r", scheduler, name) - else: - scheduler_args[name] = value - return dataclasses.replace(self, scheduler_args=scheduler_args)