diff --git a/src/diffusers/pipelines/modular_pipeline.py b/src/diffusers/pipelines/modular_pipeline.py index 683b3940c7..4961d158e1 100644 --- a/src/diffusers/pipelines/modular_pipeline.py +++ b/src/diffusers/pipelines/modular_pipeline.py @@ -138,13 +138,31 @@ class BlockState: return f"BlockState(\n{attributes}\n)" +@dataclass +class ComponentSpec: + """Specification for a pipeline component.""" + name: str + type_hint: Type + description: Optional[str] = None + default: Any = None # you can create a default component if it is a stateless class like scheduler, guider or image processor + default_class_name: Union[str, List[str], Tuple[str, str]] = None # Either "class_name" or ["module", "class_name"] + default_repo: Optional[Union[str, List[str]]] = None # either "repo" or ["repo", "subfolder"] + +@dataclass +class ConfigSpec: + """Specification for a pipeline configuration parameter.""" + name: str + default: Any + description: Optional[str] = None + + @dataclass class InputParam: name: str + type_hint: Any = None default: Any = None required: bool = False description: str = "" - type_hint: Any = Any def __repr__(self): return f"<{self.name}: {'required' if self.required else 'optional'}, default={self.default}>" @@ -152,8 +170,8 @@ class InputParam: @dataclass class OutputParam: name: str + type_hint: Any description: str = "" - type_hint: Any = Any def __repr__(self): return f"<{self.name}: {self.type_hint.__name__ if hasattr(self.type_hint, '__name__') else str(self.type_hint)}>" @@ -338,49 +356,40 @@ def make_doc_string(inputs, intermediates_inputs, outputs, description=""): return output -@dataclass -class ComponentSpec: - """Specification for a pipeline component.""" - name: str - type_hint: Optional[Type] = None - description: Optional[str] = None - default: Any = None # you can create a default component if it is a stateless class like scheduler, guider or image processor - default_class_name: Union[str, List[str], Tuple[str, str]] # Either "class_name" or ["module", "class_name"] - default_repo: Optional[Union[str, List[str]]] = None # either "repo" or ["repo", "subfolder"] - -@dataclass -class ConfigSpec: - """Specification for a pipeline configuration parameter.""" - name: str - default: Any - description: Optional[str] = None - type_hint: Optional[Type] = None class PipelineBlock: - - component_specs: List[ComponentSpec] = [] - config_specs: List[ConfigSpec] = [] + model_name = None @property def description(self) -> str: """Description of the block. Must be implemented by subclasses.""" raise NotImplementedError("description method must be implemented in subclasses") + + @property + def components(self) -> List[ComponentSpec]: + return [] + @property + def configs(self) -> List[ConfigSpec]: + return [] + + + # YiYi TODO: can we combine inputs and intermediates_inputs? the difference is inputs are immutable @property def inputs(self) -> List[InputParam]: """List of input parameters. Must be implemented by subclasses.""" - raise NotImplementedError("inputs method must be implemented in subclasses") + return [] @property def intermediates_inputs(self) -> List[InputParam]: """List of intermediate input parameters. Must be implemented by subclasses.""" - raise NotImplementedError("intermediates_inputs method must be implemented in subclasses") + return [] @property def intermediates_outputs(self) -> List[OutputParam]: """List of intermediate output parameters. Must be implemented by subclasses.""" - raise NotImplementedError("intermediates_outputs method must be implemented in subclasses") + return [] # Adding outputs attributes here for consistency between PipelineBlock/AutoPipelineBlocks/SequentialPipelineBlocks @property @@ -403,10 +412,6 @@ class PipelineBlock: input_names.append(input_param.name) return input_names - def __init__(self): - self.components: Dict[str, Any] = {} - self.auxiliaries: Dict[str, Any] = {} - self.configs: Dict[str, Any] = {} def __call__(self, pipeline, state: PipelineState) -> PipelineState: raise NotImplementedError("__call__ method must be implemented in subclasses") diff --git a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_modular.py b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_modular.py index 9c7ad5ec19..5e2b8ae779 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_modular.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_modular.py @@ -141,23 +141,7 @@ def retrieve_latents( class StableDiffusionXLLoraStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="text_encoder", - type_hint=CLIPTextModel, - default_class_name=["transformers", "CLIPTextModel"], - ), - ComponentSpec( - name="text_encoder_2", - type_hint=CLIPTextModelWithProjection, - default_class_name=["transformers", "CLIPTextModelWithProjection"], - ), - ComponentSpec( - name="unet", - type_hint=UNet2DConditionModel, - default_class_name=["diffusers", "UNet2DConditionModel"], - ) - ] + model_name = "stable-diffusion-xl" @property @@ -167,49 +151,22 @@ class StableDiffusionXLLoraStep(PipelineBlock): " See [StableDiffusionXLLoraLoaderMixin](https://huggingface.co/docs/diffusers/api/loaders/lora#diffusers.loaders.StableDiffusionXLLoraLoaderMixin)" " for more details" ) - @property - def inputs(self) -> List[InputParam]: - return [] - - @property - def intermediates_inputs(self) -> List[InputParam]: - return [] + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("text_encoder", CLIPTextModel), + ComponentSpec("text_encoder_2", CLIPTextModelWithProjection), + ComponentSpec("unet", UNet2DConditionModel), + ] - @property - def intermediates_outputs(self) -> List[OutputParam]: - return [] - - def __init__(self): - super().__init__() - self.components["text_encoder"] = None - self.components["text_encoder_2"] = None - self.components["unet"] = None @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: raise EnvironmentError("StableDiffusionXLLoraStep is desgined to be used to load lora weights, __call__ is not implemented") -class StableDiffusionXLIPAdapterStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="image_encoder", - type_hint=CLIPVisionModelWithProjection, - default_class_name=["transformers", "CLIPVisionModelWithProjection"], - ), - ComponentSpec( - name="feature_extractor", - type_hint=CLIPImageProcessor, - default_class_name=["transformers", "CLIPImageProcessor"], - ), - ComponentSpec( - name="unet", - type_hint=UNet2DConditionModel, - default_class_name=["diffusers", "UNet2DConditionModel"], - ) - ] +class StableDiffusionXLIPAdapterStep(PipelineBlock, ModularIPAdapterMixin): model_name = "stable-diffusion-xl" @@ -220,26 +177,30 @@ class StableDiffusionXLIPAdapterStep(PipelineBlock): " See [ModularIPAdapterMixin](https://huggingface.co/docs/diffusers/api/loaders/ip_adapter#diffusers.loaders.ModularIPAdapterMixin)" " for more details" ) + + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("image_encoder", CLIPVisionModelWithProjection), + ComponentSpec("feature_extractor", CLIPImageProcessor), + ComponentSpec("unet", UNet2DConditionModel), + ] @property def inputs(self) -> List[InputParam]: return [ InputParam( "ip_adapter_image", - required=True, - type_hint=PipelineImageInput, + PipelineImageInput, + required=True, description="The image(s) to be used as ip adapter" ), InputParam( "guidance_scale", default=5.0, - description="Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). Guidance scale is enabled by setting `guidance_scale > 1`." ), ] - @property - def intermediates_inputs(self) -> List[InputParam]: - return [] @property def intermediates_outputs(self) -> List[OutputParam]: @@ -248,12 +209,6 @@ class StableDiffusionXLIPAdapterStep(PipelineBlock): OutputParam("negative_ip_adapter_embeds", type_hint=torch.Tensor, description="Negative IP adapter image embeddings") ] - def __init__(self): - super().__init__() - self.components["image_encoder"] = None - self.components["feature_extractor"] = None - self.components["unet"] = None - @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: @@ -281,36 +236,6 @@ class StableDiffusionXLIPAdapterStep(PipelineBlock): class StableDiffusionXLTextEncoderStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="text_encoder", - type_hint=CLIPTextModel, - default_class_name=["transformers", "CLIPTextModel"], - ), - ComponentSpec( - name="text_encoder_2", - type_hint=CLIPTextModelWithProjection, - default_class_name=["transformers", "CLIPTextModelWithProjection"], - ), - ComponentSpec( - name="tokenizer", - type_hint=CLIPTokenizer, - default_class_name=["transformers", "CLIPTokenizer"], - ), - ComponentSpec( - name="tokenizer_2", - type_hint=CLIPTokenizer, - default_class_name=["transformers", "CLIPTokenizer"], - ) - ] - - expected_configs = [ - ConfigSpec( - name="force_zeros_for_empty_prompt", - default=True, - type_hint=bool - ) - ] model_name = "stable-diffusion-xl" @@ -319,52 +244,32 @@ class StableDiffusionXLTextEncoderStep(PipelineBlock): return( "Text Encoder step that generate text_embeddings to guide the image generation" ) + + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("text_encoder", CLIPTextModel), + ComponentSpec("text_encoder_2", CLIPTextModelWithProjection), + ComponentSpec("tokenizer", CLIPTokenizer), + ComponentSpec("tokenizer_2", CLIPTokenizer), + ] + @property + def configs(self) -> List[ConfigSpec]: + return [ConfigSpec("force_zeros_for_empty_prompt", True)] @property def inputs(self) -> List[InputParam]: return [ - InputParam( - name="prompt", - type_hint=Union[str, List[str]], - description="The prompt or prompts to guide the image generation.", - ), - InputParam( - name="prompt_2", - type_hint=Union[str, List[str]], - description="The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is used in both text-encoders", - ), - InputParam( - name="negative_prompt", - type_hint=Union[str, List[str]], - description="The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).", - ), - InputParam( - name="negative_prompt_2", - type_hint=Union[str, List[str]], - description="The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders", - ), - InputParam( - name="cross_attention_kwargs", - type_hint=Optional[dict], - description="A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor]", - ), - InputParam( - name="guidance_scale", - type_hint=float, - default=5.0, - description="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.", - ), - InputParam( - name="clip_skip", - type_hint=Optional[int], - ), + InputParam("prompt"), + InputParam("prompt_2"), + InputParam("negative_prompt"), + InputParam("negative_prompt_2"), + InputParam("cross_attention_kwargs"), + InputParam("guidance_scale",default=5.0), + InputParam("clip_skip"), ] - - @property - def intermediates_inputs(self) -> List[InputParam]: - return [] @property def intermediates_outputs(self) -> List[OutputParam]: @@ -375,14 +280,6 @@ class StableDiffusionXLTextEncoderStep(PipelineBlock): OutputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor, description="negative pooled text embeddings used to guide the image generation"), ] - def __init__(self): - super().__init__() - self.configs["force_zeros_for_empty_prompt"] = True - self.components["text_encoder"] = None - self.components["text_encoder_2"] = None - self.components["tokenizer"] = None - self.components["tokenizer_2"] = None - def check_inputs(self, pipeline, data): if data.prompt is not None and (not isinstance(data.prompt, str) and not isinstance(data.prompt, list)): @@ -431,18 +328,7 @@ class StableDiffusionXLTextEncoderStep(PipelineBlock): class StableDiffusionXLVaeEncoderStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="vae", - type_hint=AutoencoderKL, - default_class_name=["diffusers", "AutoencoderKL"], - ), - ComponentSpec( - name="image_processor", - type_hint=VaeImageProcessor, - default=VaeImageProcessor() - ) - ] + model_name = "stable-diffusion-xl" @@ -452,37 +338,20 @@ class StableDiffusionXLVaeEncoderStep(PipelineBlock): "Vae Encoder step that encode the input image into a latent representation" ) + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec("image_processor", VaeImageProcessor, default=VaeImageProcessor()), + ] + @property def inputs(self) -> List[InputParam]: return [ - InputParam( - name="image", - type_hint=PipelineImageInput, - required=True, - description="The image(s) to modify with the pipeline, for img2img or inpainting task. When using for inpainting task, parts of the image will be masked out with `mask_image` and repainted according to `prompt`." - ), - InputParam( - name="generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)" - "to make generation deterministic." - ), - InputParam( - name="height", - type_hint=Optional[int], - description="The height in pixels of the generated image. This is set to 1024 by default for the best results. " - "Anything below 512 pixels won't work well for [stabilityai/stable-diffusion-xl-base-1.0]" - "(https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) and checkpoints that are not " - "specifically fine-tuned on low resolutions.", - ), - InputParam( - name="width", - type_hint=Optional[int], - description="The width in pixels of the generated image. This is set to 1024 by default for the best results. " - "Anything below 512 pixels won't work well for [stabilityai/stable-diffusion-xl-base-1.0]" - "(https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) and checkpoints that are not " - "specifically fine-tuned on low resolutions.", - ), + InputParam("image", required=True), + InputParam("generator"), + InputParam("height"), + InputParam("width"), ] @property @@ -495,10 +364,6 @@ class StableDiffusionXLVaeEncoderStep(PipelineBlock): def intermediates_outputs(self) -> List[OutputParam]: return [OutputParam("image_latents", type_hint=torch.Tensor, description="The latents representing the reference image for image-to-image/inpainting generation")] - def __init__(self): - super().__init__() - self.components["vae"] = None - self.auxiliaries["image_processor"] = VaeImageProcessor() @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: @@ -528,24 +393,16 @@ class StableDiffusionXLVaeEncoderStep(PipelineBlock): class StableDiffusionXLInpaintVaeEncoderStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="vae", - type_hint=AutoencoderKL, - default_class_name=["diffusers", "AutoencoderKL"], - ), - ComponentSpec( - name="image_processor", - type_hint=VaeImageProcessor, - default=VaeImageProcessor() - ), - ComponentSpec( - name="mask_processor", - type_hint=VaeImageProcessor, - default=VaeImageProcessor(do_normalize=False, do_binarize=True, do_convert_grayscale=True) - ) - ] model_name = "stable-diffusion-xl" + + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec("image_processor", VaeImageProcessor, default=VaeImageProcessor()), + ComponentSpec("mask_processor", VaeImageProcessor, default=VaeImageProcessor(do_normalize=False, do_binarize=True, do_convert_grayscale=True)), + ] + @property def description(self) -> str: @@ -556,53 +413,12 @@ class StableDiffusionXLInpaintVaeEncoderStep(PipelineBlock): @property def inputs(self) -> List[InputParam]: return [ - InputParam( - "height", - type_hint=Optional[int], - description="The height in pixels of the generated image. This is set to 1024 by default for the best results. " - "Anything below 512 pixels won't work well for [stabilityai/stable-diffusion-xl-base-1.0]" - "(https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) and checkpoints that are not " - "specifically fine-tuned on low resolutions.", - ), - InputParam( - "width", - type_hint=Optional[int], - description="The width in pixels of the generated image. This is set to 1024 by default for the best results. " - "Anything below 512 pixels won't work well for [stabilityai/stable-diffusion-xl-base-1.0]" - "(https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) and checkpoints that are not " - "specifically fine-tuned on low resolutions.", - ), - InputParam( - "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) " - "to make generation deterministic." - ), - InputParam( - "image", - required=True, - type_hint=PipelineImageInput, - description="The image(s) to modify with the pipeline, for img2img or inpainting task. When using for inpainting task, parts of the image will be masked out with `mask_image` and repainted according to `prompt`." - ), - InputParam( - "mask_image", - required=True, - type_hint=PipelineImageInput, - description="`Image`, or tensor representing an image batch, to mask `image`. White pixels in the mask will be " - "repainted, while black pixels will be preserved. If `mask_image` is a PIL image, it will be converted " - "to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) " - "instead of 3, so the expected shape would be `(B, H, W, 1)`." - ), - InputParam( - "padding_mask_crop", - type_hint=Optional[Tuple[int, int]], - description="The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied to " - "image and mask_image. If `padding_mask_crop` is not `None`, it will first find a rectangular region " - "with the same aspect ratio of the image and contains all masked area, and then expand that area based " - "on `padding_mask_crop`. The image and mask_image will then be cropped based on the expanded area before " - "resizing to the original image size for inpainting. This is useful when the masked area is small while " - "the image is large and contain information irrelevant for inpainting, such as background." - ), + InputParam("height"), + InputParam("width"), + InputParam("generator"), + InputParam("image", required=True), + InputParam("mask_image", required=True), + InputParam("padding_mask_crop"), ] @property @@ -615,12 +431,7 @@ class StableDiffusionXLInpaintVaeEncoderStep(PipelineBlock): OutputParam("mask", type_hint=torch.Tensor, description="The mask to use for the inpainting process"), OutputParam("masked_image_latents", type_hint=torch.Tensor, description="The masked image latents to use for the inpainting process (only for inpainting-specifid unet)"), OutputParam("crops_coords", type_hint=Optional[Tuple[int, int]], description="The crop coordinates to use for the preprocess/postprocess of the image and mask")] - - def __init__(self): - super().__init__() - self.auxiliaries["image_processor"] = VaeImageProcessor() - self.auxiliaries["mask_processor"] = VaeImageProcessor(do_normalize=False, do_binarize=True, do_convert_grayscale=True) - self.components["vae"] = None + @torch.no_grad() def __call__(self, pipeline: DiffusionPipeline, state: PipelineState) -> PipelineState: @@ -682,12 +493,7 @@ class StableDiffusionXLInputStep(PipelineBlock): @property def inputs(self) -> List[InputParam]: return [ - InputParam( - name="num_images_per_prompt", - type_hint=int, - default=1, - description="The number of images to generate per prompt.", - ), + InputParam("num_images_per_prompt", default=1), ] @property @@ -788,15 +594,14 @@ class StableDiffusionXLInputStep(PipelineBlock): class StableDiffusionXLImg2ImgSetTimestepsStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ) - ] model_name = "stable-diffusion-xl" + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ] + @property def description(self) -> str: return ( @@ -807,50 +612,14 @@ class StableDiffusionXLImg2ImgSetTimestepsStep(PipelineBlock): @property def inputs(self) -> List[InputParam]: return [ - InputParam( - "num_inference_steps", - default=50, - type_hint=int, - description="The number of denoising steps. More denoising steps usually lead to a higher quality image at the" - " expense of slower inference." - ), - InputParam( - "timesteps", - type_hint=Optional[torch.Tensor], - description="Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order." - ), - InputParam( - "sigmas", - type_hint=Optional[torch.Tensor], - description="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." - ), - InputParam( - "denoising_end", - type_hint=Optional[float], - description="When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be completed before it is intentionally prematurely terminated. As a result, the returned sample will still retain a substantial amount of noise as determined by the discrete timesteps selected by the scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a 'Mixture of Denoisers' multi-pipeline setup." - ), - InputParam( - "strength", - default=0.3, - type_hint=float, - description="Conceptually, indicates how much to transform the reference `image` (the masked portion of image for inpainting). Must be between 0 and 1. `image` " - "will be used as a starting point, adding more noise to it the larger the `strength`. The number of " - "denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will " - "be maximum and the denoising process will run for the full number of iterations specified in " - "`num_inference_steps`. A value of 1, therefore, essentially ignores `image`. Note that in the case of " - "`denoising_start` being declared as an integer, the value of `strength` will be ignored." - ), - InputParam( - "denoising_start", - type_hint=Optional[float], - description="The denoising start value to use for the scheduler. Determines the starting point of the denoising process." - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt. Defaults to 1." - ), + InputParam("num_inference_steps", default=50), + InputParam("timesteps"), + InputParam("sigmas"), + InputParam("denoising_end"), + InputParam("strength", default=0.3), + InputParam("denoising_start"), + # YiYi TODO: do we need num_images_per_prompt here? + InputParam("num_images_per_prompt", default=1), ] @property @@ -867,9 +636,6 @@ class StableDiffusionXLImg2ImgSetTimestepsStep(PipelineBlock): OutputParam("latent_timestep", type_hint=torch.Tensor, description="The timestep that represents the initial noise level for image-to-image generation") ] - def __init__(self): - super().__init__() - self.components["scheduler"] = None @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: @@ -908,14 +674,14 @@ class StableDiffusionXLImg2ImgSetTimestepsStep(PipelineBlock): class StableDiffusionXLSetTimestepsStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ) - ] + model_name = "stable-diffusion-xl" + + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ] @property def description(self) -> str: @@ -926,27 +692,10 @@ class StableDiffusionXLSetTimestepsStep(PipelineBlock): @property def inputs(self) -> List[InputParam]: return [ - InputParam( - "num_inference_steps", - default=50, - type_hint=int, - description="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference." - ), - InputParam( - "timesteps", - type_hint=Optional[torch.Tensor], - description="Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order." - ), - InputParam( - "sigmas", - type_hint=Optional[torch.Tensor], - description="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." - ), - InputParam( - "denoising_end", - type_hint=Optional[float], - description="When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be completed before it is intentionally prematurely terminated. As a result, the returned sample will still retain a substantial amount of noise as determined by the discrete timesteps selected by the scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a 'Mixture of Denoisers' multi-pipeline setup." - ), + InputParam("num_inference_steps", default=50), + InputParam("timesteps"), + InputParam("sigmas"), + InputParam("denoising_end"), ] @property @@ -954,13 +703,6 @@ class StableDiffusionXLSetTimestepsStep(PipelineBlock): return [OutputParam("timesteps", type_hint=torch.Tensor, description="The timesteps to use for inference"), OutputParam("num_inference_steps", type_hint=int, description="The number of denoising steps to perform at inference time")] - @property - def intermediates_inputs(self) -> List[InputParam]: - return [] - - def __init__(self): - super().__init__() - self.components["scheduler"] = None @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: @@ -987,15 +729,15 @@ class StableDiffusionXLSetTimestepsStep(PipelineBlock): class StableDiffusionXLInpaintPrepareLatentsStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ) - ] + model_name = "stable-diffusion-xl" + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ] + @property def description(self) -> str: return ( @@ -1005,31 +747,13 @@ class StableDiffusionXLInpaintPrepareLatentsStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) " - "to make generation deterministic."), - InputParam( - "latents", - type_hint=Optional[torch.Tensor], - description="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 ge generated by sampling using the supplied random `generator`." - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt" - ), - InputParam( - "denoising_start", - type_hint=Optional[float], - description="When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be bypassed before it is initiated. The initial part of the denoising process is skipped and it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, strength will be ignored. Useful for 'Mixture of Denoisers' multi-pipeline setups." - ), + InputParam("generator"), + InputParam("latents"), + InputParam("num_images_per_prompt", default=1), + InputParam("denoising_start"), InputParam( "strength", - default=0.9999, - type_hint=float, + default=0.9999, description="Conceptually, indicates how much to transform the reference `image` (the masked portion of image for inpainting). Must be between 0 and 1. `image` " "will be used as a starting point, adding more noise to it the larger the `strength`. The number of " "denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will " @@ -1085,9 +809,6 @@ class StableDiffusionXLInpaintPrepareLatentsStep(PipelineBlock): OutputParam("masked_image_latents", type_hint=torch.Tensor, description="The masked image latents to use for the inpainting generation (only for inpainting-specific unet)"), OutputParam("noise", type_hint=torch.Tensor, description="The noise added to the image latents, used for inpainting generation")] - def __init__(self): - super().__init__() - self.components["scheduler"] = None @torch.no_grad() def __call__(self, pipeline: DiffusionPipeline, state: PipelineState) -> PipelineState: @@ -1143,20 +864,15 @@ class StableDiffusionXLInpaintPrepareLatentsStep(PipelineBlock): class StableDiffusionXLImg2ImgPrepareLatentsStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="vae", - type_hint=AutoencoderKL, - default_class_name=["diffusers", "AutoencoderKL"], - ), - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ) - ] model_name = "stable-diffusion-xl" + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ] + @property def description(self) -> str: return ( @@ -1166,28 +882,10 @@ class StableDiffusionXLImg2ImgPrepareLatentsStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) " - "to make generation deterministic." - ), - InputParam( - "latents", - type_hint=Optional[torch.Tensor], - description="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 ge generated by sampling using the supplied random `generator`." - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt" - ), - InputParam( - "denoising_start", - type_hint=Optional[float], - description="When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be bypassed before it is initiated. The initial part of the denoising process is skipped and it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, strength will be ignored. Useful for 'Mixture of Denoisers' multi-pipeline setups." - ), + InputParam("generator"), + InputParam("latents"), + InputParam("num_images_per_prompt", default=1), + InputParam("denoising_start"), ] @property @@ -1202,9 +900,6 @@ class StableDiffusionXLImg2ImgPrepareLatentsStep(PipelineBlock): def intermediates_outputs(self) -> List[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process")] - def __init__(self): - super().__init__() - self.components["scheduler"] = None @torch.no_grad() def __call__(self, pipeline: DiffusionPipeline, state: PipelineState) -> PipelineState: @@ -1231,15 +926,14 @@ class StableDiffusionXLImg2ImgPrepareLatentsStep(PipelineBlock): class StableDiffusionXLPrepareLatentsStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ) - ] model_name = "stable-diffusion-xl" + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ] + @property def description(self) -> str: return ( @@ -1249,37 +943,11 @@ class StableDiffusionXLPrepareLatentsStep(PipelineBlock): @property def inputs(self) -> List[InputParam]: return [ - InputParam( - "height", - type_hint=Optional[int], - description="The height in pixels of the generated image. This is set to 1024 by default for the best results. " - "Anything below 512 pixels won't work well for [stabilityai/stable-diffusion-xl-base-1.0]" - "(https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) and checkpoints that are not " - "specifically fine-tuned on low resolutions."), - InputParam( - "width", - type_hint=Optional[int], - description="The width in pixels of the generated image. This is set to 1024 by default for the best results. " - "Anything below 512 pixels won't work well for [stabilityai/stable-diffusion-xl-base-1.0]" - "(https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) and checkpoints that are not " - "specifically fine-tuned on low resolutions."), - InputParam( - "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) " - "to make generation deterministic." - ), - InputParam( - "latents", - type_hint=Optional[torch.Tensor], - description="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 ge generated by sampling using the supplied random `generator`." - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt" - ), + InputParam("height"), + InputParam("width"), + InputParam("generator"), + InputParam("latents"), + InputParam("num_images_per_prompt", default=1), ] @property @@ -1308,9 +976,6 @@ class StableDiffusionXLPrepareLatentsStep(PipelineBlock): ) ] - def __init__(self): - super().__init__() - self.components["scheduler"] = None @staticmethod def check_inputs(pipeline, data): @@ -1355,14 +1020,12 @@ class StableDiffusionXLPrepareLatentsStep(PipelineBlock): class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(PipelineBlock): - expected_configs = [ - ConfigSpec( - name="requires_aesthetics_score", - default=False, - type_hint=bool, - ) - ] + model_name = "stable-diffusion-xl" + + @property + def configs(self) -> List[ConfigSpec]: + return [ConfigSpec("requires_aesthetics_score", default=False),] @property def description(self) -> str: @@ -1373,75 +1036,16 @@ class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "original_size", - type_hint=Optional[Tuple[int]], - description="If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. " - "`original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as " - "explained in section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "target_size", - type_hint=Optional[Tuple[int]], - description="For most cases, `target_size` should be set to the desired height and width of the generated image. If " - "not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in " - "section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "negative_original_size", - type_hint=Optional[Tuple[int]], - description="To negatively condition the generation process based on a specific image resolution. Part of SDXL's " - "micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "negative_target_size", - type_hint=Optional[Tuple[int]], - description="To negatively condition the generation process based on a target image resolution. It should be as same " - "as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of " - "https://huggingface.co/papers/2307.01952" - ), - InputParam( - "crops_coords_top_left", - default=(0, 0), - type_hint=Tuple[int], - description="`crops_coords_top_left` can be used to generate an image that appears to be \"cropped\" from the position " - "`crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning" - ), - InputParam( - "negative_crops_coords_top_left", - default=(0, 0), - type_hint=Tuple[int], - description="To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's " - "micro-conditioning" - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt." - ), - InputParam( - "guidance_scale", - default=5.0, - type_hint=float, - description="Guidance scale as defined in Classifier-Free Diffusion Guidance. `guidance_scale` is defined as `w` of equation 2. " - "Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, " - "usually at the expense of lower image quality." - ), - InputParam( - "aesthetic_score", - default=6.0, - type_hint=float, - description="Used to simulate an aesthetic score of the generated image by influencing the positive text condition. " - "Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "negative_aesthetic_score", - default=2.0, - type_hint=float, - description="Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952. " - "Can be used to simulate an aesthetic score of the generated image by influencing the negative text condition." - ), + InputParam("original_size"), + InputParam("target_size"), + InputParam("negative_original_size"), + InputParam("negative_target_size"), + InputParam("crops_coords_top_left", default=(0, 0)), + InputParam("negative_crops_coords_top_left", default=(0, 0)), + InputParam("num_images_per_prompt", default=1), + InputParam("guidance_scale", required=True), + InputParam("aesthetic_score", default=6.0), + InputParam("negative_aesthetic_score", default=2.0), ] @property @@ -1458,9 +1062,57 @@ class StableDiffusionXLImg2ImgPrepareAdditionalConditioningStep(PipelineBlock): OutputParam("negative_add_time_ids", type_hint=torch.Tensor, description="The negative time ids to condition the denoising process"), OutputParam("timestep_cond", type_hint=torch.Tensor, description="The timestep cond to use for LCM")] - def __init__(self): - super().__init__() - self.configs["requires_aesthetics_score"] = False + # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids with self -> components + def _get_add_time_ids_img2img( + components, + original_size, + crops_coords_top_left, + target_size, + aesthetic_score, + negative_aesthetic_score, + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype, + text_encoder_projection_dim=None, + ): + if components.config.requires_aesthetics_score: + add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) + add_neg_time_ids = list( + negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) + ) + else: + add_time_ids = list(original_size + crops_coords_top_left + target_size) + add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size) + + passed_add_embed_dim = ( + components.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim + ) + expected_add_embed_dim = components.unet.add_embedding.linear_1.in_features + + if ( + expected_add_embed_dim > passed_add_embed_dim + and (expected_add_embed_dim - passed_add_embed_dim) == components.unet.config.addition_time_embed_dim + ): + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." + ) + elif ( + expected_add_embed_dim < passed_add_embed_dim + and (passed_add_embed_dim - expected_add_embed_dim) == components.unet.config.addition_time_embed_dim + ): + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." + ) + elif expected_add_embed_dim != passed_add_embed_dim: + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." + ) + + add_time_ids = torch.tensor([add_time_ids], dtype=dtype) + add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) + + return add_time_ids, add_neg_time_ids @torch.no_grad() def __call__(self, pipeline: DiffusionPipeline, state: PipelineState) -> PipelineState: @@ -1526,55 +1178,14 @@ class StableDiffusionXLPrepareAdditionalConditioningStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "original_size", - type_hint=Tuple[int, int], - default=(1024, 1024), - description="The original size (height, width) of the image that conditions the generation process. If different from target_size, the image will appear to be down- or upsampled. Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "target_size", - type_hint=Tuple[int, int], - default=(1024, 1024), - description="The target size (height, width) of the generated image. For most cases, this should be set to the desired output dimensions. Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "negative_original_size", - type_hint=Tuple[int, int], - default=(1024, 1024), - description="The negative original size to condition against during generation. Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952. See: https://github.com/huggingface/diffusers/issues/4208" - ), - InputParam( - "negative_target_size", - type_hint=Tuple[int, int], - default=(1024, 1024), - description="The negative target size to condition against during generation. Should typically match target_size. Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952. See: https://github.com/huggingface/diffusers/issues/4208" - ), - InputParam( - "crops_coords_top_left", - default=(0, 0), - type_hint=Tuple[int, int], - description="The top-left coordinates (x, y) used to condition the generation process. Setting this to (0, 0) typically produces well-centered images. Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952" - ), - InputParam( - "negative_crops_coords_top_left", - default=(0, 0), - type_hint=Tuple[int, int], - description="The top-left coordinates (x, y) used to negatively condition the generation process. Part of SDXL's micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952. For more information, see: https://github.com/huggingface/diffusers/issues/4208" - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt" - ), - InputParam( - "guidance_scale", - default=5.0, - type_hint=float, - description="Guidance scale as defined in Classifier-Free Diffusion Guidance. `guidance_scale` is defined as `w` of equation 2. " - "Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, " - "usually at the expense of lower image quality."), + InputParam("original_size"), + InputParam("target_size"), + InputParam("negative_original_size"), + InputParam("negative_target_size"), + InputParam("crops_coords_top_left", default=(0, 0)), + InputParam("negative_crops_coords_top_left", default=(0, 0)), + InputParam("num_images_per_prompt", default=1), + InputParam("guidance_scale", default=5.0), ] @property @@ -1606,6 +1217,58 @@ class StableDiffusionXLPrepareAdditionalConditioningStep(PipelineBlock): OutputParam("negative_add_time_ids", type_hint=torch.Tensor, description="The negative time ids to condition the denoising process"), OutputParam("timestep_cond", type_hint=torch.Tensor, description="The timestep cond to use for LCM")] + # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids with self -> components + def _get_add_time_ids_img2img( + components, + original_size, + crops_coords_top_left, + target_size, + aesthetic_score, + negative_aesthetic_score, + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype, + text_encoder_projection_dim=None, + ): + if components.config.requires_aesthetics_score: + add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) + add_neg_time_ids = list( + negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) + ) + else: + add_time_ids = list(original_size + crops_coords_top_left + target_size) + add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size) + + passed_add_embed_dim = ( + components.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim + ) + expected_add_embed_dim = components.unet.add_embedding.linear_1.in_features + + if ( + expected_add_embed_dim > passed_add_embed_dim + and (expected_add_embed_dim - passed_add_embed_dim) == components.unet.config.addition_time_embed_dim + ): + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." + ) + elif ( + expected_add_embed_dim < passed_add_embed_dim + and (passed_add_embed_dim - expected_add_embed_dim) == components.unet.config.addition_time_embed_dim + ): + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." + ) + elif expected_add_embed_dim != passed_add_embed_dim: + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." + ) + + add_time_ids = torch.tensor([add_time_ids], dtype=dtype) + add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) + + return add_time_ids, add_neg_time_ids + @torch.no_grad() def __call__(self, pipeline: DiffusionPipeline, state: PipelineState) -> PipelineState: data = self.get_block_state(state) @@ -1658,25 +1321,17 @@ class StableDiffusionXLPrepareAdditionalConditioningStep(PipelineBlock): class StableDiffusionXLDenoiseStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="unet", - type_hint=UNet2DConditionModel, - default_class_name=["diffusers", "UNet2DConditionModel"], - ), - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ), - ComponentSpec( - name="guider", - type_hint=Guiders, - default_class_name=["diffusers", "CFGGuider"], - ) - ] + model_name = "stable-diffusion-xl" + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("guider", CFGGuider), + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ComponentSpec("unet", UNet2DConditionModel), + ] + @property def description(self) -> str: return ( @@ -1686,47 +1341,13 @@ class StableDiffusionXLDenoiseStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "guidance_scale", - type_hint=float, - default=5.0, - description="Guidance scale as defined in Classifier-Free Diffusion Guidance. Higher values encourage images closely linked to the text prompt, potentially at the expense of image quality. Enabled when > 1." - ), - InputParam( - "guidance_rescale", - type_hint=float, - default=0.0, - description="Guidance rescale factor (φ) to fix overexposure when using zero terminal SNR, as proposed in 'Common Diffusion Noise Schedules and Sample Steps are Flawed'." - ), - InputParam( - "cross_attention_kwargs", - type_hint=Optional[Dict[str, Any]], - default=None, - description="Optional kwargs dictionary passed to the AttentionProcessor." - ), - InputParam( - "generator", - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of torch generator(s) to make generation deterministic." - ), - InputParam( - "eta", - type_hint=float, - default=0.0, - description="Parameter η in the DDIM paper. Only applies to DDIMScheduler, ignored for others." - ), - InputParam( - "guider_kwargs", - type_hint=Optional[Dict[str, Any]], - default=None, - description="Optional kwargs dictionary passed to the Guider." - ), - InputParam( - "num_images_per_prompt", - type_hint=int, - default=1, - description="The number of images to generate per prompt." - ), + InputParam("guidance_scale", default=5.0), + InputParam("guidance_rescale", default=0.0), + InputParam("cross_attention_kwargs"), + InputParam("generator"), + InputParam("eta", default=0.0), + InputParam("guider_kwargs"), + InputParam("num_images_per_prompt", default=1), ] @property @@ -1830,11 +1451,6 @@ class StableDiffusionXLDenoiseStep(PipelineBlock): def intermediates_outputs(self) -> List[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] - def __init__(self): - super().__init__() - self.components["guider"] = CFGGuider() - self.components["scheduler"] = None - self.components["unet"] = None def check_inputs(self, pipeline, data): @@ -1961,41 +1577,20 @@ class StableDiffusionXLDenoiseStep(PipelineBlock): class StableDiffusionXLControlNetDenoiseStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="unet", - type_hint=UNet2DConditionModel, - default_class_name=["diffusers", "UNet2DConditionModel"], - ), - ComponentSpec( - name="controlnet", - type_hint=ControlNetModel, - default_class_name=["diffusers", "ControlNetModel"], - ), - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ), - ComponentSpec( - name="guider", - type_hint=Guiders, - default_class_name=["diffusers", "CFGGuider"], - ), - ComponentSpec( - name="controlnet_guider", - type_hint=Guiders, - default_class_name=["diffusers", "CFGGuider"], - ), - ComponentSpec( - name="control_image_processor", - type_hint=VaeImageProcessor, - default=VaeImageProcessor(do_convert_rgb=True, do_normalize=False) - ) - ] model_name = "stable-diffusion-xl" + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("guider", CFGGuider), + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ComponentSpec("unet", UNet2DConditionModel), + ComponentSpec("controlnet", ControlNetModel), + ComponentSpec("control_image_processor", VaeImageProcessor, default=VaeImageProcessor(do_convert_rgb=True, do_normalize=False)), + ComponentSpec("controlnet_guider", CFGGuider), + ] + @property def description(self) -> str: return "step that iteratively denoise the latents for the text-to-image/image-to-image/inpainting generation process. Using ControlNet to condition the denoising process" @@ -2003,78 +1598,18 @@ class StableDiffusionXLControlNetDenoiseStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "control_image", - required=True, - type_hint=PipelineImageInput, - description="The ControlNet input condition to provide guidance to the unet for generation. If passed as torch.Tensor, it is used as-is. PIL.Image.Image inputs are accepted and default to image dimensions. For multiple ControlNets, pass images as a list for proper batching." - ), - InputParam( - "control_guidance_start", - default=0.0, - type_hint=Union[float, List[float]], - description="The percentage of total steps at which the ControlNet starts applying." - ), - InputParam( - "control_guidance_end", - default=1.0, - type_hint=Union[float, List[float]], - description="The percentage of total steps at which the ControlNet stops applying." - ), - InputParam( - "controlnet_conditioning_scale", - default=1.0, - type_hint=Union[float, List[float]], - description="Scale factor for ControlNet outputs before adding to unet residual. For multiple ControlNets, can be set as a list of scales." - ), - InputParam( - "guess_mode", - default=False, - type_hint=bool, - description="Enables ControlNet encoder to recognize input image content without prompts. Recommended guidance_scale: 3.0-5.0." - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt." - ), - InputParam( - "guidance_scale", - default=5.0, - type_hint=float, - description="Guidance scale as defined in Classifier-Free Diffusion Guidance. Higher values encourage images closely linked to the text prompt, potentially at the expense of image quality. Enabled when > 1." - ), - InputParam( - "guidance_rescale", - default=0.0, - type_hint=float, - description="Guidance rescale factor (φ) to fix overexposure when using zero terminal SNR, as proposed in 'Common Diffusion Noise Schedules and Sample Steps are Flawed'." - ), - InputParam( - "cross_attention_kwargs", - default=None, - type_hint=Optional[Dict[str, Any]], - description="Optional kwargs dictionary passed to the AttentionProcessor." - ), - InputParam( - "generator", - default=None, - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of torch generator(s) to make generation deterministic." - ), - InputParam( - "eta", - default=0.0, - type_hint=float, - description="Parameter η in the DDIM paper. Only applies to DDIMScheduler, ignored for others." - ), - InputParam( - "guider_kwargs", - default=None, - type_hint=Optional[Dict[str, Any]], - description="Optional kwargs dictionary passed to the Guider." - ), + InputParam("control_image", required=True), + InputParam("control_guidance_start", default=0.0), + InputParam("control_guidance_end", default=1.0), + InputParam("controlnet_conditioning_scale", default=1.0), + InputParam("guess_mode", default=False), + InputParam("num_images_per_prompt", default=1), + InputParam("guidance_scale", default=5.0), + InputParam("guidance_rescale", default=0.0), + InputParam("cross_attention_kwargs"), + InputParam("generator"), + InputParam("eta", default=0.0), + InputParam("guider_kwargs"), ] @property @@ -2183,16 +1718,6 @@ class StableDiffusionXLControlNetDenoiseStep(PipelineBlock): def intermediates_outputs(self) -> List[OutputParam]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] - def __init__(self): - super().__init__() - self.components["guider"] = CFGGuider() - self.components["controlnet_guider"] = CFGGuider() - self.components["scheduler"] = None - self.components["unet"] = None - self.components["controlnet"] = None - control_image_processor = VaeImageProcessor(do_convert_rgb=True, do_normalize=False) - self.auxiliaries["control_image_processor"] = control_image_processor - def check_inputs(self, pipeline, data): num_channels_unet = pipeline.unet.config.in_channels @@ -2447,39 +1972,18 @@ class StableDiffusionXLControlNetDenoiseStep(PipelineBlock): class StableDiffusionXLControlNetUnionDenoiseStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="unet", - type_hint=UNet2DConditionModel, - default_class_name=["diffusers", "UNet2DConditionModel"], - ), - ComponentSpec( - name="controlnet", - type_hint=ControlNetUnionModel, - default_class_name=["diffusers", "ControlNetUnionModel"], - ), - ComponentSpec( - name="scheduler", - type_hint=KarrasDiffusionSchedulers, - default_class_name=["diffusers", "EulerDiscreteScheduler"], - ), - ComponentSpec( - name="guider", - type_hint=Guiders, - default_class_name=["diffusers", "CFGGuider"], - ), - ComponentSpec( - name="controlnet_guider", - type_hint=Guiders, - default_class_name=["diffusers", "CFGGuider"], - ), - ComponentSpec( - name="control_image_processor", - type_hint=VaeImageProcessor, - default=VaeImageProcessor(do_convert_rgb=True, do_normalize=False) - ) - ] model_name = "stable-diffusion-xl" + + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("unet", UNet2DConditionModel), + ComponentSpec("controlnet", ControlNetUnionModel), + ComponentSpec("scheduler", KarrasDiffusionSchedulers), + ComponentSpec("guider", CFGGuider), + ComponentSpec("controlnet_guider", CFGGuider), + ComponentSpec("control_image_processor", VaeImageProcessor, default=VaeImageProcessor(do_convert_rgb=True, do_normalize=False)), + ] @property def description(self) -> str: @@ -2487,75 +1991,19 @@ class StableDiffusionXLControlNetUnionDenoiseStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "control_image", - required=True, - type_hint=PipelineImageInput, - description="The ControlNet input condition to provide guidance to the unet for generation. If passed as torch.Tensor, it is used as-is. PIL.Image.Image inputs are accepted and default to image dimensions. For multiple ControlNets, pass images as a list for proper batching."), - InputParam( - "control_guidance_start", - default=0.0, - type_hint=Union[float, List[float]], - description="The percentage of total steps at which the ControlNet starts applying."), - InputParam( - "control_guidance_end", - default=1.0, - type_hint=Union[float, List[float]], - description="The percentage of total steps at which the ControlNet stops applying."), - InputParam( - "control_mode", - required=True, - type_hint=List[int], - description="The control mode for union controlnet, 0 for openpose, 1 for depth, 2 for hed/pidi/scribble/ted, 3 for canny/lineart/anime_lineart/mlsd, 4 for normal and 5 for segment" - ), - InputParam( - "controlnet_conditioning_scale", - default=1.0, - type_hint=Union[float, List[float]], - description="Scale factor for ControlNet outputs before adding to unet residual. For multiple ControlNets, can be set as a list of scales." - ), - InputParam( - "guess_mode", - default=False, - type_hint=bool, - description="Enables ControlNet encoder to recognize input image content without prompts. Recommended guidance_scale: 3.0-5.0." - ), - InputParam( - "num_images_per_prompt", - default=1, - type_hint=int, - description="The number of images to generate per prompt." - ), - InputParam( - "guidance_scale", - default=5.0, - type_hint=float, - description="Guidance scale as defined in Classifier-Free Diffusion Guidance. Higher values encourage images closely linked to the text prompt, potentially at the expense of image quality. Enabled when > 1."), - InputParam( - "guidance_rescale", - default=0.0, - type_hint=float, - description="Guidance rescale factor (φ) to fix overexposure when using zero terminal SNR, as proposed in 'Common Diffusion Noise Schedules and Sample Steps are Flawed'."), - InputParam( - "cross_attention_kwargs", - default=None, - type_hint=Optional[Dict[str, Any]], - description="Optional kwargs dictionary passed to the AttentionProcessor."), - InputParam( - "generator", - default=None, - type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], - description="One or a list of torch generator(s) to make generation deterministic."), - InputParam( - "eta", - default=0.0, - type_hint=float, - description="Parameter η in the DDIM paper. Only applies to DDIMScheduler, ignored for others."), - InputParam( - "guider_kwargs", - default=None, - type_hint=Optional[Dict[str, Any]], - description="Optional kwargs dictionary passed to the Guider."), + InputParam("control_image", required=True), + InputParam("control_guidance_start", default=0.0), + InputParam("control_guidance_end", default=1.0), + InputParam("control_mode", required=True), + InputParam("controlnet_conditioning_scale", default=1.0), + InputParam("guess_mode", default=False), + InputParam("num_images_per_prompt", default=1), + InputParam("guidance_scale", default=5.0), + InputParam("guidance_rescale", default=0.0), + InputParam("cross_attention_kwargs"), + InputParam("generator"), + InputParam("eta", default=0.0), + InputParam("guider_kwargs") ] @property @@ -2664,16 +2112,6 @@ class StableDiffusionXLControlNetUnionDenoiseStep(PipelineBlock): def intermediates_outputs(self) -> List[str]: return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")] - def __init__(self): - super().__init__() - self.components["guider"] = CFGGuider() - self.components["controlnet_guider"] = CFGGuider() - self.components["scheduler"] = None - self.components["unet"] = None - self.components["controlnet"] = None - control_image_processor = VaeImageProcessor(do_convert_rgb=True, do_normalize=False) - self.auxiliaries["control_image_processor"] = control_image_processor - def check_inputs(self, pipeline, data): num_channels_unet = pipeline.unet.config.in_channels @@ -2921,19 +2359,15 @@ class StableDiffusionXLControlNetUnionDenoiseStep(PipelineBlock): class StableDiffusionXLDecodeLatentsStep(PipelineBlock): - expected_components = [ - ComponentSpec( - name="vae", - type_hint=AutoencoderKL, - default_class_name=["diffusers", "AutoencoderKL"], - ), - ComponentSpec( - name="image_processor", - type_hint=VaeImageProcessor, - default=VaeImageProcessor() - ) - ] + model_name = "stable-diffusion-xl" + + @property + def components(self) -> List[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec("image_processor", VaeImageProcessor, default=VaeImageProcessor()) + ] @property def description(self) -> str: @@ -2942,12 +2376,7 @@ class StableDiffusionXLDecodeLatentsStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "output_type", - type_hint=str, - default="pil", - description="The output format of the generated image. Choose between PIL (PIL.Image.Image), torch.Tensor or np.array." - ), + InputParam("output_type", default="pil"), ] @property @@ -2958,11 +2387,6 @@ class StableDiffusionXLDecodeLatentsStep(PipelineBlock): def intermediates_outputs(self) -> List[str]: return [OutputParam("images", type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], description="The generated images, can be a PIL.Image.Image, torch.Tensor or a numpy array")] - def __init__(self): - super().__init__() - self.components["vae"] = None - self.auxiliaries["image_processor"] = VaeImageProcessor(vae_scale_factor=8) - @torch.no_grad() def __call__(self, pipeline, state: PipelineState) -> PipelineState: data = self.get_block_state(state) @@ -3028,24 +2452,9 @@ class StableDiffusionXLInpaintOverlayMaskStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: return [ - InputParam( - "image", - type_hint=PipelineImageInput, - required=True, - description="The image(s) to modify with the pipeline, for img2img or inpainting task. When using for inpainting task, parts of the image will be masked out with `mask_image` and repainted according to `prompt`." - ), - InputParam( - "mask_image", - type_hint=PipelineImageInput, - required=True, - description="The mask image(s) to use for inpainting, white pixels in the mask will be repainted, while black pixels will be preserved. If mask_image is a PIL image, it will be converted to a single channel (luminance) before use. If it's a tensor, it should contain one color channel (L) instead of 3, so the expected shape would be (B, H, W, 1). Must be a `PIL.Image.Image`" - ), - InputParam( - "padding_mask_crop", - type_hint=Optional[Tuple[int, int]], - default=None, - description="The size of margin in the crop to be applied to the image and masking. If `None`, no crop is applied. If set, it will find a rectangular region with the same aspect ratio as the image that contains all masked areas, then expand that area by this margin. The image and mask_image are cropped to this expanded area before resizing to the original size for inpainting. Useful when the masked area is small in a large image with irrelevant background information." - ), + InputParam("image", required=True), + InputParam("mask_image", required=True), + InputParam("padding_mask_crop"), ] @property @@ -3080,7 +2489,7 @@ class StableDiffusionXLOutputStep(PipelineBlock): @property def inputs(self) -> List[Tuple[str, Any]]: - return [(InputParam("return_dict", type_hint=bool, default=True, description="Whether or not to return a StableDiffusionXLPipelineOutput instead of a plain tuple."))] + return [InputParam("return_dict", default=True)] @property def intermediates_inputs(self) -> List[str]: @@ -3227,7 +2636,7 @@ class StableDiffusionXLAutoDecodeStep(AutoPipelineBlocks): " - `StableDiffusionXLDecodeStep` (non-inpaint) is used when `padding_mask_crop` is not provided." -class StableDiffusionXLAutoIPAdapterStep(AutoPipelineBlocks): +class StableDiffusionXLAutoIPAdapterStep(AutoPipelineBlocks, ModularIPAdapterMixin): block_classes = [StableDiffusionXLIPAdapterStep] block_names = ["ip_adapter"] block_trigger_inputs = ["ip_adapter_image"] @@ -3325,8 +2734,7 @@ SDXL_SUPPORTED_BLOCKS = { } -class StableDiffusionXLModularPipeline( - ModularPipeline, +class StableDiffusionXLComponents( StableDiffusionMixin, TextualInversionLoaderMixin, StableDiffusionXLLoraLoaderMixin, @@ -3360,24 +2768,6 @@ class StableDiffusionXLModularPipeline( num_channels_latents = self.vae.config.latent_channels return num_channels_latents - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids - def _get_add_time_ids( - self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None - ): - add_time_ids = list(original_size + crops_coords_top_left + target_size) - - passed_add_embed_dim = ( - self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim - ) - expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features - - if expected_add_embed_dim != passed_add_embed_dim: - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." - ) - - add_time_ids = torch.tensor([add_time_ids], dtype=dtype) - return add_time_ids # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids def _get_add_time_ids_img2img( @@ -4142,824 +3532,107 @@ class StableDiffusionXLModularPipeline( assert emb.shape == (w.shape[0], embedding_dim) return emb - ) - if data.do_classifier_free_guidance: - data.negative_ip_adapter_embeds = [] - for i, image_embeds in enumerate(data.ip_adapter_embeds): - negative_image_embeds, image_embeds = image_embeds.chunk(2) - data.negative_ip_adapter_embeds.append(negative_image_embeds) - data.ip_adapter_embeds[i] = image_embeds - self.add_block_state(state, data) - return pipeline, state - @property - def default_sample_size(self): - default_sample_size = 128 - if hasattr(self, "unet") and self.unet is not None: - default_sample_size = self.unet.config.sample_size - return default_sample_size - - @property - def vae_scale_factor(self): - vae_scale_factor = 8 - if hasattr(self, "vae") and self.vae is not None: - vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) - return vae_scale_factor - - @property - def num_channels_unet(self): - num_channels_unet = 4 - if hasattr(self, "unet") and self.unet is not None: - num_channels_unet = self.unet.config.in_channels - return num_channels_unet - - @property - def num_channels_latents(self): - num_channels_latents = 4 - if hasattr(self, "vae") and self.vae is not None: - num_channels_latents = self.vae.config.latent_channels - return num_channels_latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._get_add_time_ids - def _get_add_time_ids( - self, original_size, crops_coords_top_left, target_size, dtype, text_encoder_projection_dim=None - ): - add_time_ids = list(original_size + crops_coords_top_left + target_size) - - passed_add_embed_dim = ( - self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim - ) - expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features - - if expected_add_embed_dim != passed_add_embed_dim: - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." - ) - - add_time_ids = torch.tensor([add_time_ids], dtype=dtype) - return add_time_ids - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids - def _get_add_time_ids_img2img( - self, - original_size, - crops_coords_top_left, - target_size, - aesthetic_score, - negative_aesthetic_score, - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype, - text_encoder_projection_dim=None, - ): - if self.config.requires_aesthetics_score: - add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) - add_neg_time_ids = list( - negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) - ) - else: - add_time_ids = list(original_size + crops_coords_top_left + target_size) - add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size) - - passed_add_embed_dim = ( - self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim - ) - expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features - - if ( - expected_add_embed_dim > passed_add_embed_dim - and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_embed_dim - ): - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." - ) - elif ( - expected_add_embed_dim < passed_add_embed_dim - and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_embed_dim - ): - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." - ) - elif expected_add_embed_dim != passed_add_embed_dim: - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." - ) - - add_time_ids = torch.tensor([add_time_ids], dtype=dtype) - add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) - - return add_time_ids, add_neg_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None): - dtype = next(self.image_encoder.parameters()).dtype - - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=dtype) - if output_hidden_states: - image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2] - image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_enc_hidden_states = self.image_encoder( - torch.zeros_like(image), output_hidden_states=True - ).hidden_states[-2] - uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave( - num_images_per_prompt, dim=0 - ) - return image_enc_hidden_states, uncond_image_enc_hidden_states - else: - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0) - uncond_image_embeds = torch.zeros_like(image_embeds) - - return image_embeds, uncond_image_embeds - - # Modified from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image - # 1. return image without apply any guidance - # 2. add crops_coords and resize_mode to preprocess() - def prepare_control_image( - self, - image, - width, - height, - batch_size, - num_images_per_prompt, - device, - dtype, - crops_coords=None, - ): - if crops_coords is not None: - image = self.control_image_processor.preprocess(image, height=height, width=width, crops_coords=crops_coords, resize_mode="fill").to(dtype=torch.float32) - else: - image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) - image_batch_size = image.shape[0] - - if image_batch_size == 1: - repeat_by = batch_size - else: - # image batch size is the same as prompt batch size - repeat_by = num_images_per_prompt - - image = image.repeat_interleave(repeat_by, dim=0) - - image = image.to(device=device, dtype=dtype) - - return image - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt - def encode_prompt( - self, - prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.Tensor] = None, - negative_prompt_embeds: Optional[torch.Tensor] = None, - pooled_prompt_embeds: Optional[torch.Tensor] = None, - negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - used in both text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - do_classifier_free_guidance (`bool`): - whether to use classifier free guidance or not - 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`). - negative_prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and - `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders - prompt_embeds (`torch.Tensor`, *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 (`torch.Tensor`, *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. - pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided, pooled text embeddings will be generated from `prompt` input argument. - negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` - input argument. - lora_scale (`float`, *optional*): - A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. - clip_skip (`int`, *optional*): - Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that - the output of the pre-final layer will be used for computing the prompt embeddings. - """ - device = device or self._execution_device - - # set lora scale so that monkey patched LoRA - # function of text encoder can correctly access it - if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin): - self._lora_scale = lora_scale - - # dynamically adjust the LoRA scale - if self.text_encoder is not None: - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) - else: - scale_lora_layers(self.text_encoder, lora_scale) - - if self.text_encoder_2 is not None: - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale) - else: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] - text_encoders = ( - [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2] - ) - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # textual inversion: process multi-vector tokens if necessary - prompt_embeds_list = [] - prompts = [prompt, prompt_2] - for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - prompt = self.maybe_convert_prompt(prompt, tokenizer) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( - text_input_ids, untruncated_ids - ): - removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {tokenizer.model_max_length} tokens: {removed_text}" - ) - - prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) - - # We are only ALWAYS interested in the pooled output of the final text encoder - pooled_prompt_embeds = prompt_embeds[0] - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - # "2" because SDXL always indexes from the penultimate layer. - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)] - - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) - - # get unconditional embeddings for classifier free guidance - zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt - if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) - elif do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - - # normalize str to list - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - negative_prompt_2 = ( - batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 - ) - - uncond_tokens: List[str] - if prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - else: - uncond_tokens = [negative_prompt, negative_prompt_2] - - negative_prompt_embeds_list = [] - for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer) - - max_length = prompt_embeds.shape[1] - uncond_input = tokenizer( - negative_prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), - output_hidden_states=True, - ) - # We are only ALWAYS interested in the pooled output of the final text encoder - negative_pooled_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] - - negative_prompt_embeds_list.append(negative_prompt_embeds) - - negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1) - - if self.text_encoder_2 is not None: - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - else: - prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device) - - bs_embed, seq_len, _ = prompt_embeds.shape - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) - - if do_classifier_free_guidance: - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - seq_len = negative_prompt_embeds.shape[1] - - if self.text_encoder_2 is not None: - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) - else: - negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed * num_images_per_prompt, -1 - ) - if do_classifier_free_guidance: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( - bs_embed * num_images_per_prompt, -1 - ) - - if self.text_encoder is not None: - if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder, lora_scale) - - if self.text_encoder_2 is not None: - if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder_2, lora_scale) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds - def prepare_ip_adapter_image_embeds( - self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance - ): - image_embeds = [] - if do_classifier_free_guidance: - negative_image_embeds = [] - if ip_adapter_image_embeds is None: - if not isinstance(ip_adapter_image, list): - ip_adapter_image = [ip_adapter_image] - - if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers): - raise ValueError( - f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters." - ) - - for single_ip_adapter_image, image_proj_layer in zip( - ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers - ): - output_hidden_state = not isinstance(image_proj_layer, ImageProjection) - single_image_embeds, single_negative_image_embeds = self.encode_image( - single_ip_adapter_image, device, 1, output_hidden_state - ) - - image_embeds.append(single_image_embeds[None, :]) - if do_classifier_free_guidance: - negative_image_embeds.append(single_negative_image_embeds[None, :]) - else: - for single_image_embeds in ip_adapter_image_embeds: - if do_classifier_free_guidance: - single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2) - negative_image_embeds.append(single_negative_image_embeds) - image_embeds.append(single_image_embeds) - - ip_adapter_image_embeds = [] - for i, single_image_embeds in enumerate(image_embeds): - single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0) - if do_classifier_free_guidance: - single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0) - single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0) - - single_image_embeds = single_image_embeds.to(device=device) - ip_adapter_image_embeds.append(single_image_embeds) - - return ip_adapter_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps - def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): - # get the original timestep using init_timestep - if denoising_start is None: - init_timestep = min(int(num_inference_steps * strength), num_inference_steps) - t_start = max(num_inference_steps - init_timestep, 0) - - timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start * self.scheduler.order) - - return timesteps, num_inference_steps - t_start - - else: - # Strength is irrelevant if we directly request a timestep to start at; - # that is, strength is determined by the denoising_start instead. - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_start * self.scheduler.config.num_train_timesteps) - ) - ) - - num_inference_steps = (self.scheduler.timesteps < discrete_timestep_cutoff).sum().item() - if self.scheduler.order == 2 and num_inference_steps % 2 == 0: - # if the scheduler is a 2nd order scheduler we might have to do +1 - # because `num_inference_steps` might be even given that every timestep - # (except the highest one) is duplicated. If `num_inference_steps` is even it would - # mean that we cut the timesteps in the middle of the denoising step - # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1 - # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler - num_inference_steps = num_inference_steps + 1 - - # because t_n+1 >= t_n, we slice the timesteps starting from the end - t_start = len(self.scheduler.timesteps) - num_inference_steps - timesteps = self.scheduler.timesteps[t_start:] - if hasattr(self.scheduler, "set_begin_index"): - self.scheduler.set_begin_index(t_start) - return timesteps, num_inference_steps - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents - # YiYi TODO: refactor using _encode_vae_image - def prepare_latents_img2img( - self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True - ): - if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): - raise ValueError( - f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" - ) - - # Offload text encoder if `enable_model_cpu_offload` was enabled - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.text_encoder_2.to("cpu") - torch.cuda.empty_cache() - - image = image.to(device=device, dtype=dtype) - - batch_size = batch_size * num_images_per_prompt - - if image.shape[1] == 4: - init_latents = image - - else: - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - # make sure the VAE is in float32 mode, as it overflows in float16 - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - elif isinstance(generator, list): - if image.shape[0] < batch_size and batch_size % image.shape[0] == 0: - image = torch.cat([image] * (batch_size // image.shape[0]), dim=0) - elif image.shape[0] < batch_size and batch_size % image.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {image.shape[0]} to effective batch_size {batch_size} " - ) - - init_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(batch_size) - ] - init_latents = torch.cat(init_latents, dim=0) - else: - init_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - init_latents = init_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=device, dtype=dtype) - latents_std = latents_std.to(device=device, dtype=dtype) - init_latents = (init_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - init_latents = self.vae.config.scaling_factor * init_latents - - if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: - # expand init_latents for batch_size - additional_image_per_prompt = batch_size // init_latents.shape[0] - init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) - elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: - raise ValueError( - f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." - ) - else: - init_latents = torch.cat([init_latents], dim=0) - - if add_noise: - shape = init_latents.shape - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # get latents - init_latents = self.scheduler.add_noise(init_latents, noise, timestep) - - latents = init_latents - - return latents - - # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_latents - def prepare_latents_inpaint( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None, - image=None, - timestep=None, - is_strength_max=True, - add_noise=True, - return_noise=False, - return_image_latents=False, - ): - shape = ( - batch_size, - num_channels_latents, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - if (image is None or timestep is None) and not is_strength_max: - raise ValueError( - "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise." - "However, either the image or the noise timestep has not been provided." - ) - - if image.shape[1] == 4: - image_latents = image.to(device=device, dtype=dtype) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - elif return_image_latents or (latents is None and not is_strength_max): - image = image.to(device=device, dtype=dtype) - image_latents = self._encode_vae_image(image=image, generator=generator) - image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1) - - if latents is None and add_noise: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - # if strength is 1. then initialise the latents to noise, else initial to image + noise - latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep) - # if pure noise then scale the initial latents by the Scheduler's init sigma - latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents - elif add_noise: - noise = latents.to(device) - latents = noise * self.scheduler.init_noise_sigma - else: - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = image_latents.to(device) - - outputs = (latents,) - - if return_noise: - outputs += (noise,) - - if return_image_latents: - outputs += (image_latents,) - - return outputs - - - # Modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline._encode_vae_image - # YiYi TODO: update the _encode_vae_image so that we can use #Coped from - def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): - - latents_mean = latents_std = None - if hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None: - latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1) - if hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None: - latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1) - - dtype = image.dtype - if self.vae.config.force_upcast: - image = image.float() - self.vae.to(dtype=torch.float32) - - if isinstance(generator, list): - image_latents = [ - retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) - for i in range(image.shape[0]) - ] - image_latents = torch.cat(image_latents, dim=0) - else: - image_latents = retrieve_latents(self.vae.encode(image), generator=generator) - - if self.vae.config.force_upcast: - self.vae.to(dtype) - - image_latents = image_latents.to(dtype) - if latents_mean is not None and latents_std is not None: - latents_mean = latents_mean.to(device=image_latents.device, dtype=dtype) - latents_std = latents_std.to(device=image_latents.device, dtype=dtype) - image_latents = (image_latents - latents_mean) * self.vae.config.scaling_factor / latents_std - else: - image_latents = self.vae.config.scaling_factor * image_latents - - return image_latents - - - # modified from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint.StableDiffusionXLInpaintPipeline.prepare_mask_latents - # do not accept do_classifier_free_guidance - def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator - ): - # resize the mask to latents shape as we concatenate the mask to the latents - # we do that before converting to dtype to avoid breaking in case we're using cpu_offload - # and half precision - mask = torch.nn.functional.interpolate( - mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor) - ) - mask = mask.to(device=device, dtype=dtype) - - # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method - if mask.shape[0] < batch_size: - if not batch_size % mask.shape[0] == 0: - raise ValueError( - "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to" - f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number" - " of masks that you pass is divisible by the total requested batch size." - ) - mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1) - - if masked_image is not None and masked_image.shape[1] == 4: - masked_image_latents = masked_image - else: - masked_image_latents = None - - if masked_image is not None: - if masked_image_latents is None: - masked_image = masked_image.to(device=device, dtype=dtype) - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) - - if masked_image_latents.shape[0] < batch_size: - if not batch_size % masked_image_latents.shape[0] == 0: - raise ValueError( - "The passed images and the required batch size don't match. Images are supposed to be duplicated" - f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed." - " Make sure the number of images that you pass is divisible by the total requested batch size." - ) - masked_image_latents = masked_image_latents.repeat( - batch_size // masked_image_latents.shape[0], 1, 1, 1 - ) - - # aligning device to prevent device errors when concating it with the latent model input - masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) - - return mask, masked_image_latents - - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding( - self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32 - ) -> torch.Tensor: - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - w (`torch.Tensor`): - Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. - embedding_dim (`int`, *optional*, defaults to 512): - Dimension of the embeddings to generate. - dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): - Data type of the generated embeddings. - - Returns: - `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb +# YiYi Notes: not used yet, maintain a list of schema that can be used across all pipeline blocks +sdxl_inputs_map = { + "prompt": InputParam("prompt", type_hint=Union[str, List[str]], description="The prompt or prompts to guide the image generation"), + "prompt_2": InputParam("prompt_2", type_hint=Union[str, List[str]], description="The prompt or prompts to be sent to the tokenizer_2 and text_encoder_2"), + "negative_prompt": InputParam("negative_prompt", type_hint=Union[str, List[str]], description="The prompt or prompts not to guide the image generation"), + "negative_prompt_2": InputParam("negative_prompt_2", type_hint=Union[str, List[str]], description="The negative prompt or prompts for text_encoder_2"), + "cross_attention_kwargs": InputParam("cross_attention_kwargs", type_hint=Optional[dict], description="Kwargs dictionary passed to the AttentionProcessor"), + "guidance_scale": InputParam("guidance_scale", type_hint=float, default=5.0, description="Classifier-Free Diffusion Guidance scale"), + "clip_skip": InputParam("clip_skip", type_hint=Optional[int], description="Number of layers to skip in CLIP text encoder"), + "image": InputParam("image", type_hint=PipelineImageInput, required=True, description="The image(s) to modify for img2img or inpainting"), + "mask_image": InputParam("mask_image", type_hint=PipelineImageInput, required=True, description="Mask image for inpainting, white pixels will be repainted"), + "generator": InputParam("generator", type_hint=Optional[Union[torch.Generator, List[torch.Generator]]], description="Generator(s) for deterministic generation"), + "height": InputParam("height", type_hint=Optional[int], description="Height in pixels of the generated image"), + "width": InputParam("width", type_hint=Optional[int], description="Width in pixels of the generated image"), + "num_images_per_prompt": InputParam("num_images_per_prompt", type_hint=int, default=1, description="Number of images to generate per prompt"), + "num_inference_steps": InputParam("num_inference_steps", type_hint=int, default=50, description="Number of denoising steps"), + "timesteps": InputParam("timesteps", type_hint=Optional[torch.Tensor], description="Custom timesteps for the denoising process"), + "sigmas": InputParam("sigmas", type_hint=Optional[torch.Tensor], description="Custom sigmas for the denoising process"), + "denoising_end": InputParam("denoising_end", type_hint=Optional[float], description="Fraction of denoising process to complete before termination"), + # YiYi Notes: img2img defaults to 0.3, inpainting defaults to 0.9999 + "strength": InputParam("strength", type_hint=float, default=0.3, description="How much to transform the reference image"), + "denoising_start": InputParam("denoising_start", type_hint=Optional[float], description="Starting point of the denoising process"), + "latents": InputParam("latents", type_hint=Optional[torch.Tensor], description="Pre-generated noisy latents for image generation"), + "padding_mask_crop": InputParam("padding_mask_crop", type_hint=Optional[Tuple[int, int]], description="Size of margin in crop for image and mask"), + "original_size": InputParam("original_size", type_hint=Optional[Tuple[int, int]], description="Original size of the image for SDXL's micro-conditioning"), + "target_size": InputParam("target_size", type_hint=Optional[Tuple[int, int]], description="Target size for SDXL's micro-conditioning"), + "negative_original_size": InputParam("negative_original_size", type_hint=Optional[Tuple[int, int]], description="Negative conditioning based on image resolution"), + "negative_target_size": InputParam("negative_target_size", type_hint=Optional[Tuple[int, int]], description="Negative conditioning based on target resolution"), + "crops_coords_top_left": InputParam("crops_coords_top_left", type_hint=Tuple[int, int], default=(0, 0), description="Top-left coordinates for SDXL's micro-conditioning"), + "negative_crops_coords_top_left": InputParam("negative_crops_coords_top_left", type_hint=Tuple[int, int], default=(0, 0), description="Negative conditioning crop coordinates"), + "aesthetic_score": InputParam("aesthetic_score", type_hint=float, default=6.0, description="Simulates aesthetic score of generated image"), + "negative_aesthetic_score": InputParam("negative_aesthetic_score", type_hint=float, default=2.0, description="Simulates negative aesthetic score"), + "guidance_rescale": InputParam("guidance_rescale", type_hint=float, default=0.0, description="Guidance rescale factor to fix overexposure"), + "eta": InputParam("eta", type_hint=float, default=0.0, description="Parameter η in the DDIM paper"), + "guider_kwargs": InputParam("guider_kwargs", type_hint=Optional[Dict[str, Any]], description="Kwargs dictionary passed to the Guider"), + "output_type": InputParam("output_type", type_hint=str, default="pil", description="Output format (pil/tensor/np.array)"), + "return_dict": InputParam("return_dict", type_hint=bool, default=True, description="Whether to return a StableDiffusionXLPipelineOutput"), + "ip_adapter_image": InputParam("ip_adapter_image", type_hint=PipelineImageInput, required=True, description="Image(s) to be used as IP adapter"), + "control_image": InputParam("control_image", type_hint=PipelineImageInput, required=True, description="ControlNet input condition"), + "control_guidance_start": InputParam("control_guidance_start", type_hint=Union[float, List[float]], default=0.0, description="When ControlNet starts applying"), + "control_guidance_end": InputParam("control_guidance_end", type_hint=Union[float, List[float]], default=1.0, description="When ControlNet stops applying"), + "controlnet_conditioning_scale": InputParam("controlnet_conditioning_scale", type_hint=Union[float, List[float]], default=1.0, description="Scale factor for ControlNet outputs"), + "guess_mode": InputParam("guess_mode", type_hint=bool, default=False, description="Enables ControlNet encoder to recognize input without prompts"), + "control_mode": InputParam("control_mode", type_hint=List[int], required=True, description="Control mode for union controlnet") +} + + +sdxl_intermediate_inputs_map = { + "prompt_embeds": InputParam("prompt_embeds", type_hint=torch.Tensor, required=True, description="Text embeddings used to guide image generation"), + "negative_prompt_embeds": InputParam("negative_prompt_embeds", type_hint=torch.Tensor, description="Negative text embeddings"), + "pooled_prompt_embeds": InputParam("pooled_prompt_embeds", type_hint=torch.Tensor, required=True, description="Pooled text embeddings"), + "negative_pooled_prompt_embeds": InputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor, description="Negative pooled text embeddings"), + "batch_size": InputParam("batch_size", type_hint=int, required=True, description="Number of prompts"), + "dtype": InputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"), + "preprocess_kwargs": InputParam("preprocess_kwargs", type_hint=Optional[dict], description="Kwargs for ImageProcessor"), + "latents": InputParam("latents", type_hint=torch.Tensor, required=True, description="Initial latents for denoising process"), + "timesteps": InputParam("timesteps", type_hint=torch.Tensor, required=True, description="Timesteps for inference"), + "num_inference_steps": InputParam("num_inference_steps", type_hint=int, required=True, description="Number of denoising steps"), + "latent_timestep": InputParam("latent_timestep", type_hint=torch.Tensor, required=True, description="Initial noise level timestep"), + "image_latents": InputParam("image_latents", type_hint=torch.Tensor, required=True, description="Latents representing reference image"), + "mask": InputParam("mask", type_hint=torch.Tensor, required=True, description="Mask for inpainting"), + "masked_image_latents": InputParam("masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting"), + "add_time_ids": InputParam("add_time_ids", type_hint=torch.Tensor, required=True, description="Time ids for conditioning"), + "negative_add_time_ids": InputParam("negative_add_time_ids", type_hint=torch.Tensor, description="Negative time ids"), + "timestep_cond": InputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"), + "noise": InputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"), + "crops_coords": InputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"), + "ip_adapter_embeds": InputParam("ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter"), + "negative_ip_adapter_embeds": InputParam("negative_ip_adapter_embeds", type_hint=List[torch.Tensor], description="Negative image embeddings for IP-Adapter"), + "images": InputParam("images", type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], required=True, description="Generated images") +} + + +sdxl_intermediate_outputs_map = { + "prompt_embeds": OutputParam("prompt_embeds", type_hint=torch.Tensor, description="Text embeddings used to guide image generation"), + "negative_prompt_embeds": OutputParam("negative_prompt_embeds", type_hint=torch.Tensor, description="Negative text embeddings"), + "pooled_prompt_embeds": OutputParam("pooled_prompt_embeds", type_hint=torch.Tensor, description="Pooled text embeddings"), + "negative_pooled_prompt_embeds": OutputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor, description="Negative pooled text embeddings"), + "batch_size": OutputParam("batch_size", type_hint=int, description="Number of prompts"), + "dtype": OutputParam("dtype", type_hint=torch.dtype, description="Data type of model tensor inputs"), + "image_latents": OutputParam("image_latents", type_hint=torch.Tensor, description="Latents representing reference image"), + "mask": OutputParam("mask", type_hint=torch.Tensor, description="Mask for inpainting"), + "masked_image_latents": OutputParam("masked_image_latents", type_hint=torch.Tensor, description="Masked image latents for inpainting"), + "crops_coords": OutputParam("crops_coords", type_hint=Optional[Tuple[int]], description="Crop coordinates"), + "timesteps": OutputParam("timesteps", type_hint=torch.Tensor, description="Timesteps for inference"), + "num_inference_steps": OutputParam("num_inference_steps", type_hint=int, description="Number of denoising steps"), + "latent_timestep": OutputParam("latent_timestep", type_hint=torch.Tensor, description="Initial noise level timestep"), + "add_time_ids": OutputParam("add_time_ids", type_hint=torch.Tensor, description="Time ids for conditioning"), + "negative_add_time_ids": OutputParam("negative_add_time_ids", type_hint=torch.Tensor, description="Negative time ids"), + "timestep_cond": OutputParam("timestep_cond", type_hint=torch.Tensor, description="Timestep conditioning for LCM"), + "latents": OutputParam("latents", type_hint=torch.Tensor, description="Denoised latents"), + "noise": OutputParam("noise", type_hint=torch.Tensor, description="Noise added to image latents"), + "ip_adapter_embeds": OutputParam("ip_adapter_embeds", type_hint=List[torch.Tensor], description="Image embeddings for IP-Adapter"), + "negative_ip_adapter_embeds": OutputParam("negative_ip_adapter_embeds", type_hint=List[torch.Tensor], description="Negative image embeddings for IP-Adapter"), + "images": OutputParam("images", type_hint=Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]], description="Generated images") +} + + +sdxl_outputs_map = { + "images": OutputParam("images", type_hint=Union[Tuple[Union[List[PIL.Image.Image], List[torch.Tensor], List[np.array]]], StableDiffusionXLPipelineOutput], description="The final generated images") +} \ No newline at end of file